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())?;
|
||||
|
||||
Reference in New Issue
Block a user