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:
@@ -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) => {
|
||||
|
||||
Reference in New Issue
Block a user