feat(tagger): add JoyTag provider behind a tagger-model abstraction
Introduces a selectable tagging model so Phokus isn't locked to the anime-leaning WD tagger. JoyTag uses the Danbooru schema but generalizes to photographic content and is strong on NSFW concepts — a better fit for a photo manager while keeping the explicitness range. - `TaggerModel` enum (wd | joytag) persisted as a `tagger_model` setting; `model_dir`, status, and download now follow the active model (per-model repo/dir/file list). WD stays the default. - `Tagger` trait implemented by `WdTagger` and the new `JoyTagger`; `create_active_tagger` builds the selected one. The worker holds `Box<dyn Tagger>` and rebuilds on model change (TAGGER_SESSION_DIRTY). - Shared `assemble_batch` skeleton (pack -> one forward pass -> per-image fallback, results in input order); both providers and the shared decode/pad/resize are de-duped onto common helpers. - JoyTag specifics: NCHW + RGB + CLIP-normalized input (vs WD's NHWC/BGR/raw), flat top_tags.txt labels, logits -> sigmoid -> threshold (default 0.4). It has no native rating, so the explicitness rating is derived from its NSFW tags. - Tags are attributed to the model that produced them (ai_tagger_model), via Tagger::model_name, instead of a hardcoded WD constant. - New get/set_tagger_model commands. Backend only; the Settings model picker and end-to-end testing against the JoyTag model files come next.
This commit is contained in:
@@ -8,7 +8,7 @@ use crate::db::{
|
||||
use crate::embedder;
|
||||
use crate::hnsw_index;
|
||||
use crate::indexer::{self, WatcherHandle};
|
||||
use crate::tagger::{self, TaggerAcceleration, TaggerModelStatus, TaggerRuntimeProbe};
|
||||
use crate::tagger::{self, TaggerAcceleration, TaggerModel, TaggerModelStatus, TaggerRuntimeProbe};
|
||||
use crate::vector;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashSet;
|
||||
@@ -1970,6 +1970,11 @@ pub struct SetTaggerAccelerationParams {
|
||||
pub acceleration: TaggerAcceleration,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub struct SetTaggerModelParams {
|
||||
pub model: TaggerModel,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub struct SetTaggerThresholdParams {
|
||||
pub threshold: f32,
|
||||
@@ -2030,6 +2035,21 @@ pub async fn set_tagger_acceleration(
|
||||
tagger::set_tagger_acceleration(&app_dir, params.acceleration).map_err(|e| e.to_string())
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
pub async fn get_tagger_model(app: AppHandle) -> Result<TaggerModel, String> {
|
||||
let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?;
|
||||
Ok(tagger::tagger_model(&app_dir))
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
pub async fn set_tagger_model(
|
||||
app: AppHandle,
|
||||
params: SetTaggerModelParams,
|
||||
) -> Result<TaggerModel, String> {
|
||||
let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?;
|
||||
tagger::set_tagger_model(&app_dir, params.model).map_err(|e| e.to_string())
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
pub async fn probe_tagger_runtime(app: AppHandle) -> Result<TaggerRuntimeProbe, String> {
|
||||
let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?;
|
||||
|
||||
@@ -3,7 +3,7 @@ use crate::db::{self, DbPool, EmbeddingJob, FolderJobProgress, ImageRecord, Inde
|
||||
use crate::embedder::{embedding_source_path, ClipImageEmbedder};
|
||||
use crate::media::{probe_video_metadata, MediaTools};
|
||||
use crate::storage::{detect_storage_profile, RuntimeAdaptiveProfile, StorageProfile};
|
||||
use crate::tagger::{self, WdTagger};
|
||||
use crate::tagger::{self, Tagger};
|
||||
use crate::thumbnail;
|
||||
use crate::vector;
|
||||
use anyhow::Result;
|
||||
@@ -330,7 +330,7 @@ pub fn start_caption_worker(app: AppHandle, pool: DbPool, app_data_dir: PathBuf)
|
||||
|
||||
pub fn start_tagging_worker(app: AppHandle, pool: DbPool, app_data_dir: PathBuf) {
|
||||
std::thread::spawn(move || {
|
||||
let mut tagger_instance: Option<WdTagger> = None;
|
||||
let mut tagger_instance: Option<Box<dyn Tagger>> = None;
|
||||
log::info!("Tagging worker started.");
|
||||
loop {
|
||||
// If the acceleration setting changed, drop the cached session so
|
||||
@@ -1236,7 +1236,7 @@ fn process_tagging_batch(
|
||||
app: &AppHandle,
|
||||
pool: &DbPool,
|
||||
app_data_dir: &Path,
|
||||
tagger_instance: &mut Option<WdTagger>,
|
||||
tagger_instance: &mut Option<Box<dyn Tagger>>,
|
||||
) -> Result<bool> {
|
||||
if !tagger::tagger_model_status(app_data_dir).ready {
|
||||
return Ok(false);
|
||||
@@ -1264,7 +1264,7 @@ fn process_tagging_batch(
|
||||
}
|
||||
|
||||
if tagger_instance.is_none() {
|
||||
match WdTagger::new(app_data_dir) {
|
||||
match tagger::create_active_tagger(app_data_dir) {
|
||||
Ok(model) => *tagger_instance = Some(model),
|
||||
Err(error) => {
|
||||
with_db_write_lock(|| {
|
||||
@@ -1323,6 +1323,10 @@ fn process_tagging_batch(
|
||||
}
|
||||
let infer_elapsed = infer_started_at.elapsed();
|
||||
|
||||
// Attribute the tags to the model that actually produced them, not a
|
||||
// hardcoded one (WD vs JoyTag are both possible).
|
||||
let tagger_model_name = tagger_ref.model_name();
|
||||
|
||||
let tag_results = jobs.iter().cloned().zip(outputs).collect::<Vec<_>>();
|
||||
|
||||
let write_started_at = Instant::now();
|
||||
@@ -1355,7 +1359,7 @@ fn process_tagging_batch(
|
||||
job.image_id,
|
||||
&tag_pairs,
|
||||
&output.rating,
|
||||
tagger::WD_TAGGER_MODEL_NAME,
|
||||
tagger_model_name,
|
||||
)?;
|
||||
}
|
||||
Err(error) => {
|
||||
|
||||
@@ -202,6 +202,8 @@ pub fn run() {
|
||||
commands::get_tagger_model_status,
|
||||
commands::get_tagger_acceleration,
|
||||
commands::set_tagger_acceleration,
|
||||
commands::get_tagger_model,
|
||||
commands::set_tagger_model,
|
||||
commands::probe_tagger_runtime,
|
||||
commands::get_tagger_threshold,
|
||||
commands::set_tagger_threshold,
|
||||
|
||||
+525
-121
@@ -16,15 +16,27 @@ pub const WD_TAGGER_MODEL_NAME: &str = "wd-swinv2-tagger-v3";
|
||||
const TAGGER_ACCELERATION_FILE: &str = "settings/tagger_acceleration.txt";
|
||||
const TAGGER_THRESHOLD_FILE: &str = "settings/tagger_threshold.txt";
|
||||
const TAGGER_BATCH_SIZE_FILE: &str = "settings/tagger_batch_size.txt";
|
||||
const TAGGER_MODEL_FILE: &str = "settings/tagger_model.txt";
|
||||
|
||||
// Files required on disk before the tagger can run. The ONNX runtime DLLs
|
||||
// are shared with the captioner and live in the same `onnxruntime/` directory.
|
||||
const TAGGER_REQUIRED_FILES: &[&str] = &[
|
||||
pub const JOYTAG_MODEL_ID: &str = "fancyfeast/joytag";
|
||||
pub const JOYTAG_MODEL_NAME: &str = "joytag";
|
||||
|
||||
// JoyTag preprocessing differs from the WD tagger: it expects RGB (not BGR),
|
||||
// CLIP-style mean/std normalization on [0,1] values (not raw [0,255]), and an
|
||||
// NCHW layout (not NHWC). These are the OpenAI CLIP normalization constants.
|
||||
const JOYTAG_MEAN: [f32; 3] = [0.481_454_66, 0.457_827_5, 0.408_210_73];
|
||||
const JOYTAG_STD: [f32; 3] = [0.268_629_54, 0.261_302_58, 0.275_777_11];
|
||||
// JoyTag's recommended detection threshold (used as the default; the user's
|
||||
// tagger_threshold setting still overrides it).
|
||||
const JOYTAG_DEFAULT_THRESHOLD: f32 = 0.4;
|
||||
|
||||
// Shared ONNX Runtime DLLs (live in the caption model's `onnxruntime/` dir and
|
||||
// are reused by the tagger). The per-model weight + label files are listed by
|
||||
// `TaggerModel::download_files`.
|
||||
const TAGGER_RUNTIME_DLLS: &[&str] = &[
|
||||
"onnxruntime/onnxruntime.dll",
|
||||
"onnxruntime/onnxruntime_providers_shared.dll",
|
||||
"onnxruntime/DirectML.dll",
|
||||
"model.onnx",
|
||||
"selected_tags.csv",
|
||||
];
|
||||
|
||||
// Tags in these Danbooru categories are kept in the output.
|
||||
@@ -78,6 +90,61 @@ impl TaggerAcceleration {
|
||||
}
|
||||
}
|
||||
|
||||
/// Which tagging model is active. Both produce a `TaggerOutput` (tags + an
|
||||
/// explicitness rating); they differ in vocabulary, preprocessing, and how the
|
||||
/// rating is derived. WD is Danbooru-trained (anime-leaning); JoyTag uses the
|
||||
/// Danbooru schema but generalizes to photographic content and is stronger on
|
||||
/// NSFW concepts.
|
||||
#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum TaggerModel {
|
||||
#[default]
|
||||
Wd,
|
||||
JoyTag,
|
||||
}
|
||||
|
||||
impl TaggerModel {
|
||||
fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::Wd => "wd",
|
||||
Self::JoyTag => "joytag",
|
||||
}
|
||||
}
|
||||
|
||||
/// Hugging Face repo the model files are fetched from.
|
||||
fn repo_id(self) -> &'static str {
|
||||
match self {
|
||||
Self::Wd => WD_TAGGER_MODEL_ID,
|
||||
Self::JoyTag => JOYTAG_MODEL_ID,
|
||||
}
|
||||
}
|
||||
|
||||
/// Stable display/identifier name.
|
||||
fn model_name(self) -> &'static str {
|
||||
match self {
|
||||
Self::Wd => WD_TAGGER_MODEL_NAME,
|
||||
Self::JoyTag => JOYTAG_MODEL_NAME,
|
||||
}
|
||||
}
|
||||
|
||||
/// Subdirectory under `models/` where this model's files live.
|
||||
fn dir_name(self) -> &'static str {
|
||||
match self {
|
||||
Self::Wd => "wd-swinv2-tagger-v3",
|
||||
Self::JoyTag => "joytag",
|
||||
}
|
||||
}
|
||||
|
||||
/// Files fetched from `repo_id` (the shared ONNX Runtime DLLs are handled
|
||||
/// separately). `model.onnx` is the weights; the second is the label list.
|
||||
fn download_files(self) -> &'static [&'static str] {
|
||||
match self {
|
||||
Self::Wd => &["model.onnx", "selected_tags.csv"],
|
||||
Self::JoyTag => &["model.onnx", "top_tags.txt"],
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Status / probe types exposed to the frontend
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -152,14 +219,41 @@ struct TagEntry {
|
||||
// Path helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Directory of the *active* tagger model's files (driven by the
|
||||
/// `tagger_model` setting), e.g. `…/models/wd-swinv2-tagger-v3`.
|
||||
pub fn model_dir(app_data_dir: &Path) -> PathBuf {
|
||||
app_data_dir.join("models").join("wd-swinv2-tagger-v3")
|
||||
app_data_dir
|
||||
.join("models")
|
||||
.join(tagger_model(app_data_dir).dir_name())
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Settings persistence
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
pub fn tagger_model(app_data_dir: &Path) -> TaggerModel {
|
||||
let path = app_data_dir.join(TAGGER_MODEL_FILE);
|
||||
let Ok(value) = std::fs::read_to_string(path) else {
|
||||
return TaggerModel::default();
|
||||
};
|
||||
match value.trim().to_ascii_lowercase().as_str() {
|
||||
"joytag" => TaggerModel::JoyTag,
|
||||
_ => TaggerModel::Wd,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn set_tagger_model(app_data_dir: &Path, model: TaggerModel) -> Result<TaggerModel> {
|
||||
let path = app_data_dir.join(TAGGER_MODEL_FILE);
|
||||
if let Some(parent) = path.parent() {
|
||||
std::fs::create_dir_all(parent)?;
|
||||
}
|
||||
std::fs::write(path, model.as_str())?;
|
||||
// Switching models means the cached session is for the wrong model; the
|
||||
// worker drops and rebuilds it on the next batch (same flag as an EP change).
|
||||
TAGGER_SESSION_DIRTY.store(true, Ordering::Relaxed);
|
||||
Ok(model)
|
||||
}
|
||||
|
||||
pub fn tagger_acceleration(app_data_dir: &Path) -> TaggerAcceleration {
|
||||
let path = app_data_dir.join(TAGGER_ACCELERATION_FILE);
|
||||
let Ok(value) = std::fs::read_to_string(path) else {
|
||||
@@ -231,25 +325,27 @@ pub fn set_tagger_batch_size(app_data_dir: &Path, batch_size: usize) -> Result<u
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
pub fn tagger_model_status(app_data_dir: &Path) -> TaggerModelStatus {
|
||||
let model = tagger_model(app_data_dir);
|
||||
let local_dir = model_dir(app_data_dir);
|
||||
// The ONNX runtime DLLs live in the caption model dir; reuse them.
|
||||
let caption_model_dir = app_data_dir.join("models").join("florence-2-base-ft");
|
||||
let missing_files = TAGGER_REQUIRED_FILES
|
||||
|
||||
let mut missing_files: Vec<String> = TAGGER_RUNTIME_DLLS
|
||||
.iter()
|
||||
.filter(|file| {
|
||||
let path = if file.starts_with("onnxruntime/") {
|
||||
caption_model_dir.join(file)
|
||||
} else {
|
||||
local_dir.join(file)
|
||||
};
|
||||
!path.exists()
|
||||
})
|
||||
.filter(|file| !caption_model_dir.join(file).exists())
|
||||
.map(|file| (*file).to_string())
|
||||
.collect::<Vec<_>>();
|
||||
.collect();
|
||||
missing_files.extend(
|
||||
model
|
||||
.download_files()
|
||||
.iter()
|
||||
.filter(|file| !local_dir.join(file).exists())
|
||||
.map(|file| (*file).to_string()),
|
||||
);
|
||||
|
||||
TaggerModelStatus {
|
||||
model_id: WD_TAGGER_MODEL_ID,
|
||||
model_name: WD_TAGGER_MODEL_NAME,
|
||||
model_id: model.repo_id(),
|
||||
model_name: model.model_name(),
|
||||
local_dir: local_dir.to_string_lossy().to_string(),
|
||||
ready: missing_files.is_empty(),
|
||||
missing_files,
|
||||
@@ -260,19 +356,20 @@ pub fn prepare_tagger_model_with_progress(
|
||||
app_data_dir: &Path,
|
||||
emit_progress: impl Fn(TaggerModelProgress),
|
||||
) -> Result<TaggerModelStatus> {
|
||||
let model = tagger_model(app_data_dir);
|
||||
let local_dir = model_dir(app_data_dir);
|
||||
std::fs::create_dir_all(&local_dir)?;
|
||||
|
||||
// The tagger shares the ONNX runtime DLLs with the captioner; download
|
||||
// them here so the tagger works even on a clean install where the caption
|
||||
// model has never been fetched (ensure_onnx_runtime only initializes).
|
||||
const DOWNLOAD_FILES: &[&str] = &["model.onnx", "selected_tags.csv"];
|
||||
let download_files = model.download_files();
|
||||
let caption_model_dir = crate::captioner::model_dir(app_data_dir);
|
||||
std::fs::create_dir_all(&caption_model_dir)?;
|
||||
|
||||
// Unified step count across DLLs and tagger files so the bar is coherent.
|
||||
let dll_count = crate::captioner::missing_onnx_runtime_count(&caption_model_dir);
|
||||
let model_pending = DOWNLOAD_FILES
|
||||
let model_pending = download_files
|
||||
.iter()
|
||||
.filter(|file| !local_dir.join(file).exists())
|
||||
.count();
|
||||
@@ -318,9 +415,9 @@ pub fn prepare_tagger_model_with_progress(
|
||||
// (timeout + resume), rather than hf-hub's download_with_progress, whose
|
||||
// agent has no read timeout and would hang on a stalled connection.
|
||||
let api = Api::new()?;
|
||||
let repo = api.repo(Repo::new(WD_TAGGER_MODEL_ID.to_string(), RepoType::Model));
|
||||
let repo = api.repo(Repo::new(model.repo_id().to_string(), RepoType::Model));
|
||||
|
||||
for file in DOWNLOAD_FILES {
|
||||
for file in download_files {
|
||||
let destination = local_dir.join(file);
|
||||
if destination.exists() {
|
||||
continue;
|
||||
@@ -424,11 +521,96 @@ pub fn probe_tagger_runtime(app_data_dir: &Path) -> Result<TaggerRuntimeProbe> {
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Top-level inference entry point
|
||||
// Tagger trait + shared batch skeleton
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// A loaded tagging model. Implementations differ in vocabulary, preprocessing,
|
||||
/// and how the explicitness rating is derived, but all turn a batch of image
|
||||
/// paths into one `TaggerOutput` per path (in order), with per-image failures
|
||||
/// reflected in the individual `Result`s. Built for the active model by
|
||||
/// [`create_active_tagger`].
|
||||
pub trait Tagger {
|
||||
fn run_batch(&mut self, image_paths: &[PathBuf], max_tags: usize) -> Vec<Result<TaggerOutput>>;
|
||||
|
||||
/// Stable name of this model, written to the DB as `ai_tagger_model` so
|
||||
/// tags can be attributed to (and re-tagged across) models.
|
||||
fn model_name(&self) -> &'static str;
|
||||
}
|
||||
|
||||
/// Build the tagger for the currently-selected model.
|
||||
pub fn create_active_tagger(app_data_dir: &Path) -> Result<Box<dyn Tagger>> {
|
||||
match tagger_model(app_data_dir) {
|
||||
TaggerModel::Wd => Ok(Box::new(WdTagger::new(app_data_dir)?)),
|
||||
TaggerModel::JoyTag => Ok(Box::new(JoyTagger::new(app_data_dir)?)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Shared batch skeleton: pack the successfully-preprocessed images into one
|
||||
/// contiguous buffer, run a single batched forward pass via `infer`, and fall
|
||||
/// back to per-image inference if the batch fails (e.g. a model pinned to
|
||||
/// batch=1). Returns one result per input slot, in order — a decode failure
|
||||
/// stays attached to its own slot. `infer(pixels, count)` runs `count`
|
||||
/// contiguous images (`count * stride` floats) and returns `count` outputs.
|
||||
fn assemble_batch(
|
||||
preprocessed: Vec<Result<Vec<f32>>>,
|
||||
stride: usize,
|
||||
model_label: &str,
|
||||
mut infer: impl FnMut(&[f32], usize) -> Result<Vec<TaggerOutput>>,
|
||||
) -> Vec<Result<TaggerOutput>> {
|
||||
let mut batch_slots: Vec<usize> = Vec::new();
|
||||
let mut batch_pixels: Vec<f32> = Vec::with_capacity(preprocessed.len() * stride);
|
||||
for (i, result) in preprocessed.iter().enumerate() {
|
||||
if let Ok(pixels) = result {
|
||||
batch_slots.push(i);
|
||||
batch_pixels.extend_from_slice(pixels);
|
||||
}
|
||||
}
|
||||
|
||||
// Seed each slot with its decode error; inference overwrites decoded slots.
|
||||
let mut results: Vec<Result<TaggerOutput>> = preprocessed
|
||||
.into_iter()
|
||||
.map(|r| match r {
|
||||
Ok(_) => Err(anyhow::anyhow!(
|
||||
"tagging inference did not run for this image"
|
||||
)),
|
||||
Err(error) => Err(error),
|
||||
})
|
||||
.collect();
|
||||
|
||||
if batch_slots.is_empty() {
|
||||
return results; // nothing decoded
|
||||
}
|
||||
|
||||
match infer(&batch_pixels, batch_slots.len()) {
|
||||
Ok(outputs) => {
|
||||
// `infer` must return exactly one output per packed image; the zip
|
||||
// below would otherwise silently leave trailing slots as errors.
|
||||
debug_assert_eq!(outputs.len(), batch_slots.len());
|
||||
for (&slot, output) in batch_slots.iter().zip(outputs) {
|
||||
results[slot] = Ok(output);
|
||||
}
|
||||
}
|
||||
Err(batch_error) => {
|
||||
log::warn!(
|
||||
"{model_label} batch inference failed for {} images, falling back to per-image: {batch_error}",
|
||||
batch_slots.len()
|
||||
);
|
||||
for (k, &slot) in batch_slots.iter().enumerate() {
|
||||
let one = &batch_pixels[k * stride..(k + 1) * stride];
|
||||
results[slot] = infer(one, 1).and_then(|mut out| {
|
||||
out.drain(..)
|
||||
.next()
|
||||
.ok_or_else(|| anyhow::anyhow!("tagger produced no output for image"))
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
results
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tagger implementation
|
||||
// WD tagger implementation
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
pub struct WdTagger {
|
||||
@@ -509,85 +691,6 @@ impl WdTagger {
|
||||
})
|
||||
}
|
||||
|
||||
/// Tag a batch of images in a single forward pass. Returns one result per
|
||||
/// input path, in the same order. Per-image decode failures — and, if the
|
||||
/// batched inference itself fails (e.g. a model pinned to batch=1), a
|
||||
/// per-image inference fallback — are reflected in the individual `Result`s.
|
||||
pub fn run_batch(
|
||||
&mut self,
|
||||
image_paths: &[PathBuf],
|
||||
max_tags: usize,
|
||||
) -> Vec<Result<TaggerOutput>> {
|
||||
if image_paths.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
let started_at = Instant::now();
|
||||
let input_size = self.input_size;
|
||||
let stride = input_size * input_size * 3;
|
||||
|
||||
// Decode + preprocess every image in parallel, keeping results aligned to
|
||||
// the inputs so a failure stays attached to its own path.
|
||||
let preprocessed: Vec<Result<Vec<f32>>> = image_paths
|
||||
.par_iter()
|
||||
.map(|path| preprocess_image(path, input_size))
|
||||
.collect();
|
||||
|
||||
// Pack the successfully-decoded images into one contiguous buffer for a
|
||||
// single batched inference, remembering which input slot each came from.
|
||||
let mut batch_slots: Vec<usize> = Vec::new();
|
||||
let mut batch_pixels: Vec<f32> = Vec::with_capacity(image_paths.len() * stride);
|
||||
for (i, result) in preprocessed.iter().enumerate() {
|
||||
if let Ok(pixels) = result {
|
||||
batch_slots.push(i);
|
||||
batch_pixels.extend_from_slice(pixels);
|
||||
}
|
||||
}
|
||||
|
||||
// Seed each slot with its decode error; inference overwrites the ones
|
||||
// that decoded. The placeholder error never survives a decoded slot.
|
||||
let mut results: Vec<Result<TaggerOutput>> = preprocessed
|
||||
.into_iter()
|
||||
.map(|r| match r {
|
||||
Ok(_) => Err(anyhow::anyhow!(
|
||||
"tagging inference did not run for this image"
|
||||
)),
|
||||
Err(error) => Err(error),
|
||||
})
|
||||
.collect();
|
||||
|
||||
if batch_slots.is_empty() {
|
||||
return results; // nothing decoded
|
||||
}
|
||||
|
||||
match self.infer_batch(&batch_pixels, batch_slots.len(), max_tags) {
|
||||
Ok(outputs) => {
|
||||
for (&slot, output) in batch_slots.iter().zip(outputs) {
|
||||
results[slot] = Ok(output);
|
||||
}
|
||||
}
|
||||
Err(batch_error) => {
|
||||
log::warn!(
|
||||
"WD tagger batch inference failed for {} images, falling back to per-image: {batch_error}",
|
||||
batch_slots.len()
|
||||
);
|
||||
for (k, &slot) in batch_slots.iter().enumerate() {
|
||||
let pixels = batch_pixels[k * stride..(k + 1) * stride].to_vec();
|
||||
results[slot] = self.infer_one(pixels, max_tags);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
log::debug!(
|
||||
"WD tagger batch: {} images ({} decoded) in {:?}",
|
||||
image_paths.len(),
|
||||
batch_slots.len(),
|
||||
started_at.elapsed(),
|
||||
);
|
||||
|
||||
results
|
||||
}
|
||||
|
||||
/// Run `count` already-preprocessed images (contiguous `[count, H, W, 3]`
|
||||
/// pixels) through the model in one forward pass.
|
||||
fn infer_batch(
|
||||
@@ -632,14 +735,6 @@ impl WdTagger {
|
||||
.collect())
|
||||
}
|
||||
|
||||
/// Run a single preprocessed image through the model.
|
||||
fn infer_one(&mut self, image_array: Vec<f32>, max_tags: usize) -> Result<TaggerOutput> {
|
||||
self.infer_batch(&image_array, 1, max_tags)?
|
||||
.into_iter()
|
||||
.next()
|
||||
.ok_or_else(|| anyhow::anyhow!("tagger produced no output for image"))
|
||||
}
|
||||
|
||||
/// Convert one image's class probabilities into its rating + sorted tags.
|
||||
/// Associated (not `&self`) so callers can hold a disjoint session borrow.
|
||||
fn tags_from_probs(
|
||||
@@ -678,6 +773,177 @@ impl WdTagger {
|
||||
}
|
||||
}
|
||||
|
||||
impl Tagger for WdTagger {
|
||||
fn run_batch(&mut self, image_paths: &[PathBuf], max_tags: usize) -> Vec<Result<TaggerOutput>> {
|
||||
if image_paths.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
let input_size = self.input_size;
|
||||
let preprocessed: Vec<Result<Vec<f32>>> = image_paths
|
||||
.par_iter()
|
||||
.map(|path| preprocess_image(path, input_size))
|
||||
.collect();
|
||||
assemble_batch(
|
||||
preprocessed,
|
||||
input_size * input_size * 3,
|
||||
"WD tagger",
|
||||
|pixels, count| self.infer_batch(pixels, count, max_tags),
|
||||
)
|
||||
}
|
||||
|
||||
fn model_name(&self) -> &'static str {
|
||||
WD_TAGGER_MODEL_NAME
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// JoyTag implementation
|
||||
// JoyTag uses the Danbooru tag schema but generalizes to photographic content
|
||||
// and is strong on NSFW concepts. It has no rating output, so the explicitness
|
||||
// rating is derived from its NSFW tags (see `joytag_rating`). Input is NCHW,
|
||||
// RGB, CLIP-normalized — see `preprocess_joytag`.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
pub struct JoyTagger {
|
||||
session: Session,
|
||||
labels: Vec<String>,
|
||||
threshold: f32,
|
||||
input_size: usize,
|
||||
}
|
||||
|
||||
impl JoyTagger {
|
||||
pub fn new(app_data_dir: &Path) -> Result<Self> {
|
||||
let started_at = Instant::now();
|
||||
|
||||
let status = tagger_model_status(app_data_dir);
|
||||
if !status.ready {
|
||||
anyhow::bail!(
|
||||
"JoyTag model is missing {} required file(s): {}",
|
||||
status.missing_files.len(),
|
||||
status.missing_files.join(", ")
|
||||
);
|
||||
}
|
||||
|
||||
let local_dir = model_dir(app_data_dir);
|
||||
|
||||
// Shared ONNX runtime DLLs (see WdTagger::new).
|
||||
let caption_model_dir = app_data_dir.join("models").join("florence-2-base-ft");
|
||||
crate::captioner::ensure_onnx_runtime(&caption_model_dir).map_err(|e| {
|
||||
anyhow::anyhow!(
|
||||
"ONNX Runtime not initialised — download the Florence-2 caption model first \
|
||||
to get the shared runtime DLLs. Original error: {e}"
|
||||
)
|
||||
})?;
|
||||
|
||||
let acceleration = tagger_acceleration(app_data_dir);
|
||||
let threshold = joytag_threshold(app_data_dir);
|
||||
let model_path = local_dir.join("model.onnx");
|
||||
let labels_path = local_dir.join("top_tags.txt");
|
||||
|
||||
let session = create_tagger_session(&model_path, acceleration)?;
|
||||
|
||||
// JoyTag uses NCHW (N, 3, H, W); the spatial size lives at axis 2.
|
||||
let (input_size, batch_axis) = {
|
||||
let inputs = session.inputs();
|
||||
inputs
|
||||
.first()
|
||||
.and_then(|inp| {
|
||||
if let ort::value::ValueType::Tensor { shape, .. } = inp.dtype() {
|
||||
let size = shape.get(2).copied().filter(|&d| d > 0).map(|d| d as usize);
|
||||
Some((size.unwrap_or(448), shape.first().copied()))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.unwrap_or((448, None))
|
||||
};
|
||||
|
||||
let labels = load_joytag_labels(&labels_path)?;
|
||||
|
||||
log::info!(
|
||||
"JoyTag loaded in {:?} ({} tags, input {}x{}, batch axis {:?}, {:?} acceleration)",
|
||||
started_at.elapsed(),
|
||||
labels.len(),
|
||||
input_size,
|
||||
input_size,
|
||||
batch_axis,
|
||||
acceleration,
|
||||
);
|
||||
|
||||
Ok(Self {
|
||||
session,
|
||||
labels,
|
||||
threshold,
|
||||
input_size,
|
||||
})
|
||||
}
|
||||
|
||||
/// Run `count` already-preprocessed images (contiguous `[count, 3, H, W]`
|
||||
/// pixels) through the model in one forward pass.
|
||||
fn infer_batch(
|
||||
&mut self,
|
||||
pixels: &[f32],
|
||||
count: usize,
|
||||
max_tags: usize,
|
||||
) -> Result<Vec<TaggerOutput>> {
|
||||
let input = Tensor::from_array((
|
||||
[count, 3usize, self.input_size, self.input_size],
|
||||
pixels.to_vec().into_boxed_slice(),
|
||||
))
|
||||
.map_err(|error| anyhow::anyhow!("{error}"))?;
|
||||
|
||||
let input_name: String = self.session.inputs()[0].name().to_string();
|
||||
let outputs = self
|
||||
.session
|
||||
.run(ort::inputs! { input_name.as_str() => input })
|
||||
.map_err(|error| anyhow::anyhow!("{error}"))?;
|
||||
|
||||
let (_, logits) = outputs[0]
|
||||
.try_extract_tensor::<f32>()
|
||||
.map_err(|error| anyhow::anyhow!("{error}"))?;
|
||||
|
||||
let n_labels = self.labels.len();
|
||||
if logits.len() != count * n_labels {
|
||||
anyhow::bail!(
|
||||
"Model output length {} does not match {} images x {} labels",
|
||||
logits.len(),
|
||||
count,
|
||||
n_labels
|
||||
);
|
||||
}
|
||||
|
||||
let labels = &self.labels;
|
||||
let threshold = self.threshold;
|
||||
Ok(logits
|
||||
.chunks_exact(n_labels)
|
||||
.map(|row| joytag_tags_from_logits(labels, threshold, row, max_tags))
|
||||
.collect())
|
||||
}
|
||||
}
|
||||
|
||||
impl Tagger for JoyTagger {
|
||||
fn run_batch(&mut self, image_paths: &[PathBuf], max_tags: usize) -> Vec<Result<TaggerOutput>> {
|
||||
if image_paths.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
let input_size = self.input_size;
|
||||
let preprocessed: Vec<Result<Vec<f32>>> = image_paths
|
||||
.par_iter()
|
||||
.map(|path| preprocess_joytag(path, input_size))
|
||||
.collect();
|
||||
assemble_batch(
|
||||
preprocessed,
|
||||
input_size * input_size * 3,
|
||||
"JoyTag",
|
||||
|pixels, count| self.infer_batch(pixels, count, max_tags),
|
||||
)
|
||||
}
|
||||
|
||||
fn model_name(&self) -> &'static str {
|
||||
JOYTAG_MODEL_NAME
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Session creation – mirrors captioner's `create_session`
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -771,12 +1037,132 @@ fn load_labels(path: &Path) -> Result<Vec<TagEntry>> {
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Image preprocessing
|
||||
// WD tagger expects: (1, H, W, 3) float32, raw [0,255] values, BGR channel
|
||||
// order, padded to square with white (255,255,255).
|
||||
// JoyTag: label loading, threshold, rating derivation
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
fn preprocess_image(image_path: &Path, target_size: usize) -> Result<Vec<f32>> {
|
||||
/// JoyTag labels: one tag per line, index == output position. Underscores are
|
||||
/// replaced with spaces to match the display style used elsewhere.
|
||||
fn load_joytag_labels(path: &Path) -> Result<Vec<String>> {
|
||||
let content = std::fs::read_to_string(path)?;
|
||||
let labels: Vec<String> = content
|
||||
.lines()
|
||||
.map(str::trim)
|
||||
.filter(|line| !line.is_empty())
|
||||
.map(|line| line.replace('_', " "))
|
||||
.collect();
|
||||
if labels.is_empty() {
|
||||
anyhow::bail!("top_tags.txt is empty or could not be parsed");
|
||||
}
|
||||
Ok(labels)
|
||||
}
|
||||
|
||||
/// JoyTag detection threshold. Honours the shared `tagger_threshold` setting if
|
||||
/// the user has set it, otherwise uses JoyTag's recommended default (0.4).
|
||||
fn joytag_threshold(app_data_dir: &Path) -> f32 {
|
||||
match std::fs::read_to_string(app_data_dir.join(TAGGER_THRESHOLD_FILE)) {
|
||||
Ok(value) => value
|
||||
.trim()
|
||||
.parse::<f32>()
|
||||
.unwrap_or(JOYTAG_DEFAULT_THRESHOLD)
|
||||
.clamp(0.01, 1.0),
|
||||
Err(_) => JOYTAG_DEFAULT_THRESHOLD,
|
||||
}
|
||||
}
|
||||
|
||||
fn sigmoid(x: f32) -> f32 {
|
||||
1.0 / (1.0 + (-x).exp())
|
||||
}
|
||||
|
||||
// Explicitness buckets for deriving a rating from JoyTag's tags (highest match
|
||||
// wins). Names use spaces, since underscores are stripped on load. Tunable.
|
||||
const JOYTAG_EXPLICIT_TAGS: &[&str] = &[
|
||||
"sex",
|
||||
"vaginal",
|
||||
"anal",
|
||||
"oral",
|
||||
"fellatio",
|
||||
"cunnilingus",
|
||||
"penis",
|
||||
"pussy",
|
||||
"cum",
|
||||
"ejaculation",
|
||||
"erection",
|
||||
"handjob",
|
||||
"paizuri",
|
||||
"masturbation",
|
||||
];
|
||||
const JOYTAG_QUESTIONABLE_TAGS: &[&str] = &[
|
||||
"nude",
|
||||
"completely nude",
|
||||
"nipples",
|
||||
"topless",
|
||||
"bottomless",
|
||||
"pubic hair",
|
||||
"areola",
|
||||
"areolae",
|
||||
];
|
||||
const JOYTAG_SENSITIVE_TAGS: &[&str] = &[
|
||||
"lingerie",
|
||||
"underwear",
|
||||
"panties",
|
||||
"swimsuit",
|
||||
"bikini",
|
||||
"cleavage",
|
||||
"bra",
|
||||
"midriff",
|
||||
];
|
||||
|
||||
/// Derive an explicitness rating from JoyTag's tags. JoyTag has no rating
|
||||
/// output, so this maps its NSFW-content tags onto the WD-style buckets.
|
||||
fn joytag_rating(tags: &[TagResult]) -> String {
|
||||
let has = |bucket: &[&str]| tags.iter().any(|t| bucket.contains(&t.tag.as_str()));
|
||||
if has(JOYTAG_EXPLICIT_TAGS) {
|
||||
"explicit".to_string()
|
||||
} else if has(JOYTAG_QUESTIONABLE_TAGS) {
|
||||
"questionable".to_string()
|
||||
} else if has(JOYTAG_SENSITIVE_TAGS) {
|
||||
"sensitive".to_string()
|
||||
} else {
|
||||
"general".to_string()
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert one image's JoyTag logits into sorted tags + a derived rating.
|
||||
fn joytag_tags_from_logits(
|
||||
labels: &[String],
|
||||
threshold: f32,
|
||||
logits: &[f32],
|
||||
max_tags: usize,
|
||||
) -> TaggerOutput {
|
||||
let mut tags: Vec<TagResult> = labels
|
||||
.iter()
|
||||
.zip(logits.iter())
|
||||
.filter_map(|(name, logit)| {
|
||||
let confidence = sigmoid(*logit);
|
||||
(confidence >= threshold).then(|| TagResult {
|
||||
tag: name.clone(),
|
||||
confidence,
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
tags.sort_by(|a, b| b.confidence.total_cmp(&a.confidence));
|
||||
|
||||
// Derive the rating from all kept tags before truncating to max_tags.
|
||||
let rating = joytag_rating(&tags);
|
||||
tags.truncate(max_tags);
|
||||
|
||||
TaggerOutput { tags, rating }
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Image preprocessing
|
||||
// Both taggers pad to square with white and resize to the model input size;
|
||||
// they differ only in channel order, value range, and tensor layout.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Decode an image, composite any alpha onto white, pad to a centered square,
|
||||
/// and resize to `target_size`. Shared by both taggers.
|
||||
fn decode_pad_resize(image_path: &Path, target_size: usize) -> Result<image::RgbImage> {
|
||||
let image = ImageReader::open(image_path)?.decode()?;
|
||||
|
||||
// Composite any alpha channel onto a white background.
|
||||
@@ -787,22 +1173,25 @@ fn preprocess_image(image_path: &Path, target_size: usize) -> Result<Vec<f32>> {
|
||||
image::imageops::overlay(&mut canvas_rgba, &image_rgba, 0, 0);
|
||||
let image_rgb = DynamicImage::ImageRgba8(canvas_rgba).to_rgb8();
|
||||
|
||||
// Pad to square.
|
||||
// Pad to a centered square.
|
||||
let max_dim = width.max(height);
|
||||
let pad_left = (max_dim - width) / 2;
|
||||
let pad_top = (max_dim - height) / 2;
|
||||
let mut square = image::RgbImage::from_pixel(max_dim, max_dim, image::Rgb([255, 255, 255]));
|
||||
image::imageops::overlay(&mut square, &image_rgb, pad_left as i64, pad_top as i64);
|
||||
|
||||
// Resize to model input size.
|
||||
let resized = image::imageops::resize(
|
||||
// Resize to model input size (CatmullRom ≈ PIL BICUBIC).
|
||||
Ok(image::imageops::resize(
|
||||
&square,
|
||||
target_size as u32,
|
||||
target_size as u32,
|
||||
FilterType::CatmullRom,
|
||||
);
|
||||
))
|
||||
}
|
||||
|
||||
// Flatten to (H, W, 3) float32 in BGR order, values in [0, 255].
|
||||
/// WD tagger input: (N, H, W, 3) float32, raw [0,255] values, BGR order.
|
||||
fn preprocess_image(image_path: &Path, target_size: usize) -> Result<Vec<f32>> {
|
||||
let resized = decode_pad_resize(image_path, target_size)?;
|
||||
let mut pixel_values = vec![0.0f32; target_size * target_size * 3];
|
||||
for (x, y, pixel) in resized.enumerate_pixels() {
|
||||
let base = (y as usize * target_size + x as usize) * 3;
|
||||
@@ -811,6 +1200,21 @@ fn preprocess_image(image_path: &Path, target_size: usize) -> Result<Vec<f32>> {
|
||||
pixel_values[base + 1] = f32::from(pixel[1]); // G
|
||||
pixel_values[base + 2] = f32::from(pixel[0]); // R
|
||||
}
|
||||
|
||||
Ok(pixel_values)
|
||||
}
|
||||
|
||||
/// JoyTag input: (N, 3, H, W) float32, RGB, CLIP-normalized ((x/255 − mean)/std).
|
||||
fn preprocess_joytag(image_path: &Path, target_size: usize) -> Result<Vec<f32>> {
|
||||
let resized = decode_pad_resize(image_path, target_size)?;
|
||||
let plane = target_size * target_size;
|
||||
let mut pixel_values = vec![0.0f32; 3 * plane];
|
||||
for (x, y, pixel) in resized.enumerate_pixels() {
|
||||
let idx = y as usize * target_size + x as usize;
|
||||
// Channel-major (NCHW): R plane, then G, then B.
|
||||
for c in 0..3 {
|
||||
pixel_values[c * plane + idx] =
|
||||
(f32::from(pixel[c]) / 255.0 - JOYTAG_MEAN[c]) / JOYTAG_STD[c];
|
||||
}
|
||||
}
|
||||
Ok(pixel_values)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user