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:
2026-06-29 09:40:12 +01:00
parent 3a2b134103
commit 52d54d2404
4 changed files with 557 additions and 127 deletions
+9 -5
View File
@@ -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) => {