diff --git a/src-tauri/src/commands.rs b/src-tauri/src/commands.rs index b591786..48595ca 100644 --- a/src-tauri/src/commands.rs +++ b/src-tauri/src/commands.rs @@ -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 { + 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 { + 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 { let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?; diff --git a/src-tauri/src/indexer.rs b/src-tauri/src/indexer.rs index ba09c05..28dab25 100644 --- a/src-tauri/src/indexer.rs +++ b/src-tauri/src/indexer.rs @@ -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 = None; + let mut tagger_instance: Option> = 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, + tagger_instance: &mut Option>, ) -> Result { 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::>(); 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) => { diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 6c5fc8d..18a7eac 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -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, diff --git a/src-tauri/src/tagger.rs b/src-tauri/src/tagger.rs index 82518c7..45e5150 100644 --- a/src-tauri/src/tagger.rs +++ b/src-tauri/src/tagger.rs @@ -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 { + 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 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 = 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::>(); + .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 { + 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 { } // --------------------------------------------------------------------------- -// 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>; + + /// 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> { + 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>>, + stride: usize, + model_label: &str, + mut infer: impl FnMut(&[f32], usize) -> Result>, +) -> Vec> { + let mut batch_slots: Vec = Vec::new(); + let mut batch_pixels: Vec = 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> = 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> { - 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>> = 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 = Vec::new(); - let mut batch_pixels: Vec = 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> = 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, max_tags: usize) -> Result { - 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> { + if image_paths.is_empty() { + return Vec::new(); + } + let input_size = self.input_size; + let preprocessed: Vec>> = 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, + threshold: f32, + input_size: usize, +} + +impl JoyTagger { + pub fn new(app_data_dir: &Path) -> Result { + 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> { + 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::() + .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> { + if image_paths.is_empty() { + return Vec::new(); + } + let input_size = self.input_size; + let preprocessed: Vec>> = 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> { } // --------------------------------------------------------------------------- -// 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> { +/// 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> { + let content = std::fs::read_to_string(path)?; + let labels: Vec = 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::() + .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 = 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 { 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> { 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> { + 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> { 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> { + 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) }