.gitignore

@@ -29,3 +29,9 @@ dist-ssr
 
 # Local staging area
 /staging
+
+# Misc
+*.py
+*.json
+
+*.pyc

src-tauri/Cargo.lock

@@ -874,6 +874,27 @@ dependencies = [
  "syn 2.0.117",
 ]
 
+[[package]]
+name = "csv"
+version = "1.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "52cd9d68cf7efc6ddfaaee42e7288d3a99d613d4b50f76ce9827ae0c6e14f938"
+dependencies = [
+ "csv-core",
+ "itoa",
+ "ryu",
+ "serde_core",
+]
+
+[[package]]
+name = "csv-core"
+version = "0.1.13"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "704a3c26996a80471189265814dbc2c257598b96b8a7feae2d31ace646bb9782"
+dependencies = [
+ "memchr",
+]
+
 [[package]]
 name = "ctor"
 version = "0.2.9"
@@ -4009,6 +4030,7 @@ dependencies = [
  "candle-nn",
  "candle-transformers",
  "chrono",
+ "csv",
  "fast_image_resize",
  "ffmpeg-sidecar",
  "hf-hub",
@@ -4030,9 +4052,11 @@ dependencies = [
  "tauri-plugin-opener",
  "tokenizers",
  "tokio",
+ "ureq",
  "uuid",
  "walkdir",
  "xxhash-rust",
+ "zip 4.6.1",
 ]
 
 [[package]]

src-tauri/Cargo.toml

@@ -44,4 +44,7 @@ candle-nn = { version = "0.10.2", features = ["cuda"] }
 candle-transformers = { version = "0.10.2", features = ["cuda"] }
 hf-hub = { version = "0.5.0", default-features = false, features = ["ureq", "native-tls"] }
 tokenizers = "0.22.1"
-ort = { version = "2.0.0-rc.12", default-features = false, features = ["std", "ndarray", "download-binaries", "copy-dylibs", "load-dynamic", "api-24", "tls-native"] }
+ort = { version = "2.0.0-rc.12", default-features = false, features = ["std", "ndarray", "download-binaries", "copy-dylibs", "load-dynamic", "directml", "api-24", "tls-native"] }
+ureq = { version = "3.3.0", default-features = false, features = ["native-tls"] }
+zip = { version = "4.6.1", default-features = false, features = ["deflate"] }
+csv = "1"

src-tauri/src/captioner.rs

@@ -1,18 +1,54 @@
 use anyhow::Result;
 use hf_hub::{api::sync::Api, Repo, RepoType};
 use image::{imageops::FilterType, ImageReader};
+use ort::ep;
 use ort::session::SessionInputValue;
+use ort::session::SessionOutputs;
 use ort::session::{builder::GraphOptimizationLevel, Session};
 use ort::value::{Shape, Tensor};
-use serde::Serialize;
+use serde::{Deserialize, Serialize};
 use std::borrow::Cow;
+use std::io::{Cursor, Read};
 use std::path::{Path, PathBuf};
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::OnceLock;
+use std::time::Instant;
 use tokenizers::Tokenizer;
 
 pub const FLORENCE_MODEL_ID: &str = "onnx-community/Florence-2-base-ft";
 pub const FLORENCE_CAPTION_MODEL_NAME: &str = "florence-2-base-ft-onnx-q4";
+const ONNX_RUNTIME_NUGET_URL: &str =
+    "https://www.nuget.org/api/v2/package/Microsoft.ML.OnnxRuntime.DirectML/1.24.2";
+const DIRECTML_NUGET_URL: &str =
+    "https://www.nuget.org/api/v2/package/Microsoft.AI.DirectML/1.15.4";
+const ONNX_RUNTIME_DLL_FILE: &str = "onnxruntime/onnxruntime.dll";
+const ONNX_RUNTIME_PROVIDERS_DLL_FILE: &str = "onnxruntime/onnxruntime_providers_shared.dll";
+const DIRECTML_DLL_FILE: &str = "onnxruntime/DirectML.dll";
+const CAPTION_ACCELERATION_FILE: &str = "settings/caption_acceleration.txt";
+const CAPTION_DETAIL_FILE: &str = "settings/caption_detail.txt";
+
+const ONNX_RUNTIME_FILES: &[(&str, &str, &str)] = &[
+    (
+        ONNX_RUNTIME_DLL_FILE,
+        ONNX_RUNTIME_NUGET_URL,
+        "runtimes/win-x64/native/onnxruntime.dll",
+    ),
+    (
+        ONNX_RUNTIME_PROVIDERS_DLL_FILE,
+        ONNX_RUNTIME_NUGET_URL,
+        "runtimes/win-x64/native/onnxruntime_providers_shared.dll",
+    ),
+    (
+        DIRECTML_DLL_FILE,
+        DIRECTML_NUGET_URL,
+        "bin/x64-win/DirectML.dll",
+    ),
+];
 
 const REQUIRED_FILES: &[&str] = &[
+    ONNX_RUNTIME_DLL_FILE,
+    ONNX_RUNTIME_PROVIDERS_DLL_FILE,
+    DIRECTML_DLL_FILE,
     "config.json",
     "generation_config.json",
     "preprocessor_config.json",
@@ -21,10 +57,81 @@ const REQUIRED_FILES: &[&str] = &[
     "special_tokens_map.json",
     "onnx/vision_encoder_fp16.onnx",
     "onnx/encoder_model_q4.onnx",
+    "onnx/decoder_model_q4.onnx",
     "onnx/decoder_model_merged_q4.onnx",
     "onnx/embed_tokens_fp16.onnx",
 ];
 
+static ORT_RUNTIME_INIT: OnceLock<std::result::Result<(), String>> = OnceLock::new();
+
+/// Set to `true` by `set_caption_acceleration` so the caption worker loop
+/// knows to drop its cached `FlorenceCaptioner` and reload with the new EP.
+pub static CAPTION_SESSION_DIRTY: AtomicBool = AtomicBool::new(false);
+
+#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
+#[serde(rename_all = "lowercase")]
+pub enum CaptionAcceleration {
+    Auto,
+    Cpu,
+    Directml,
+}
+
+impl CaptionAcceleration {
+    fn as_str(self) -> &'static str {
+        match self {
+            Self::Auto => "auto",
+            Self::Cpu => "cpu",
+            Self::Directml => "directml",
+        }
+    }
+}
+
+impl Default for CaptionAcceleration {
+    fn default() -> Self {
+        Self::Auto
+    }
+}
+
+#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
+#[serde(rename_all = "lowercase")]
+pub enum CaptionDetail {
+    Short,
+    Detailed,
+    Paragraph,
+}
+
+impl CaptionDetail {
+    fn as_str(self) -> &'static str {
+        match self {
+            Self::Short => "short",
+            Self::Detailed => "detailed",
+            Self::Paragraph => "paragraph",
+        }
+    }
+
+    fn prompt(self) -> &'static str {
+        match self {
+            Self::Short => "What does the image describe?",
+            Self::Detailed => "Describe in detail what is shown in the image.",
+            Self::Paragraph => "Describe with a paragraph what is shown in the image.",
+        }
+    }
+
+    fn max_new_tokens(self) -> usize {
+        match self {
+            Self::Short => 32,
+            Self::Detailed => 56,
+            Self::Paragraph => 96,
+        }
+    }
+}
+
+impl Default for CaptionDetail {
+    fn default() -> Self {
+        Self::Paragraph
+    }
+}
+
 #[derive(Serialize)]
 pub struct CaptionModelStatus {
     pub model_id: &'static str,
@@ -45,6 +152,8 @@ pub struct CaptionModelProgress {
 #[derive(Serialize)]
 pub struct CaptionRuntimeProbe {
     pub ready: bool,
+    pub acceleration: CaptionAcceleration,
+    pub detail: CaptionDetail,
     pub tokenizer_vocab_size: usize,
     pub sessions: Vec<CaptionRuntimeSessionProbe>,
 }
@@ -61,6 +170,7 @@ pub struct CaptionVisionProbe {
     pub input_shape: Vec<i64>,
     pub output_shape: Vec<i64>,
     pub output_values: usize,
+    pub acceleration: CaptionAcceleration,
 }
 
 #[derive(Clone)]
@@ -71,9 +181,11 @@ struct TensorData {
 
 pub struct FlorenceCaptioner {
     tokenizer: Tokenizer,
+    caption_detail: CaptionDetail,
     vision_session: Session,
     embed_session: Session,
     encoder_session: Session,
+    decoder_prefill_session: Session,
     decoder_session: Session,
 }
 
@@ -81,6 +193,54 @@ pub fn model_dir(app_data_dir: &Path) -> PathBuf {
     app_data_dir.join("models").join("florence-2-base-ft")
 }
 
+pub fn caption_acceleration(app_data_dir: &Path) -> CaptionAcceleration {
+    let path = app_data_dir.join(CAPTION_ACCELERATION_FILE);
+    let Ok(value) = std::fs::read_to_string(path) else {
+        return CaptionAcceleration::default();
+    };
+
+    match value.trim().to_ascii_lowercase().as_str() {
+        "cpu" => CaptionAcceleration::Cpu,
+        "directml" => CaptionAcceleration::Directml,
+        _ => CaptionAcceleration::Auto,
+    }
+}
+
+pub fn set_caption_acceleration(
+    app_data_dir: &Path,
+    acceleration: CaptionAcceleration,
+) -> Result<CaptionAcceleration> {
+    let path = app_data_dir.join(CAPTION_ACCELERATION_FILE);
+    if let Some(parent) = path.parent() {
+        std::fs::create_dir_all(parent)?;
+    }
+    std::fs::write(path, acceleration.as_str())?;
+    CAPTION_SESSION_DIRTY.store(true, Ordering::Relaxed);
+    Ok(acceleration)
+}
+
+pub fn caption_detail(app_data_dir: &Path) -> CaptionDetail {
+    let path = app_data_dir.join(CAPTION_DETAIL_FILE);
+    let Ok(value) = std::fs::read_to_string(path) else {
+        return CaptionDetail::default();
+    };
+
+    match value.trim().to_ascii_lowercase().as_str() {
+        "short" => CaptionDetail::Short,
+        "paragraph" => CaptionDetail::Paragraph,
+        _ => CaptionDetail::Detailed,
+    }
+}
+
+pub fn set_caption_detail(app_data_dir: &Path, detail: CaptionDetail) -> Result<CaptionDetail> {
+    let path = app_data_dir.join(CAPTION_DETAIL_FILE);
+    if let Some(parent) = path.parent() {
+        std::fs::create_dir_all(parent)?;
+    }
+    std::fs::write(path, detail.as_str())?;
+    Ok(detail)
+}
+
 pub fn caption_model_status(app_data_dir: &Path) -> CaptionModelStatus {
     let local_dir = model_dir(app_data_dir);
     let missing_files = REQUIRED_FILES
@@ -133,9 +293,20 @@ pub fn prepare_caption_model_with_progress(
         if let Some(parent) = destination.parent() {
             std::fs::create_dir_all(parent)?;
         }
-        let cached = repo.get(file)?;
-        std::fs::copy(cached, destination)?;
-        completed_files += 1;
+        if ONNX_RUNTIME_FILES
+            .iter()
+            .any(|(runtime_file, _, _)| runtime_file == file)
+        {
+            download_onnx_runtime_files(&local_dir)?;
+            completed_files = REQUIRED_FILES
+                .iter()
+                .filter(|file| local_dir.join(file).exists())
+                .count();
+        } else {
+            let cached = repo.get(file)?;
+            std::fs::copy(cached, destination)?;
+            completed_files += 1;
+        }
         emit_progress(CaptionModelProgress {
             total_files: REQUIRED_FILES.len(),
             completed_files,
@@ -172,21 +343,40 @@ pub fn probe_caption_runtime(app_data_dir: &Path) -> Result<CaptionRuntimeProbe>
     }
 
     let local_dir = model_dir(app_data_dir);
+    ensure_onnx_runtime(&local_dir)?;
     let tokenizer =
         Tokenizer::from_file(local_dir.join("tokenizer.json")).map_err(anyhow::Error::msg)?;
 
-    let sessions = [
-        "onnx/vision_encoder_fp16.onnx",
+    let acceleration = caption_acceleration(app_data_dir);
+
+    // For the vision encoder (the only session that runs on the GPU EP) we
+    // create an actual ORT session so we can verify which EP was actually
+    // loaded, rather than just echoing the settings file.
+    let vision_session = probe_vision_session(
+        &local_dir.join("onnx/vision_encoder_fp16.onnx"),
+        acceleration,
+    )?;
+
+    // The remaining entries are DLLs and CPU-only text/decoder models — probe
+    // them with the lightweight file-size check as before.
+    let other_files = [
+        "onnxruntime/onnxruntime.dll",
+        "onnxruntime/onnxruntime_providers_shared.dll",
+        "onnxruntime/DirectML.dll",
         "onnx/embed_tokens_fp16.onnx",
         "onnx/encoder_model_q4.onnx",
+        "onnx/decoder_model_q4.onnx",
         "onnx/decoder_model_merged_q4.onnx",
-    ]
-    .into_iter()
-    .map(|file| probe_session(file, &local_dir.join(file)))
-    .collect::<Result<Vec<_>>>()?;
+    ];
+    let mut sessions = vec![vision_session];
+    for file in other_files {
+        sessions.push(probe_session(file, &local_dir.join(file))?);
+    }
 
     Ok(CaptionRuntimeProbe {
         ready: true,
+        acceleration,
+        detail: caption_detail(app_data_dir),
         tokenizer_vocab_size: tokenizer.get_vocab_size(false),
         sessions,
     })
@@ -202,11 +392,17 @@ pub fn probe_caption_vision(app_data_dir: &Path, image_path: &Path) -> Result<Ca
     }
 
     let local_dir = model_dir(app_data_dir);
+    ensure_onnx_runtime(&local_dir)?;
     let pixels = preprocess_image(image_path)?;
     let input_shape = vec![1, 3, 768, 768];
     let input = Tensor::from_array(([1usize, 3, 768, 768], pixels.into_boxed_slice()))
         .map_err(|error| anyhow::anyhow!("{error}"))?;
-    let mut session = create_session(&local_dir.join("onnx/vision_encoder_fp16.onnx"))?;
+    let acceleration = caption_acceleration(app_data_dir);
+    let mut session = create_session(
+        &local_dir.join("onnx/vision_encoder_fp16.onnx"),
+        acceleration,
+        true,
+    )?;
     let outputs = session
         .run(ort::inputs! {
             "pixel_values" => input
@@ -220,6 +416,7 @@ pub fn probe_caption_vision(app_data_dir: &Path, image_path: &Path) -> Result<Ca
         input_shape,
         output_shape: output_shape.to_vec(),
         output_values: output_values.len(),
+        acceleration,
     })
 }
 
@@ -230,6 +427,7 @@ pub fn generate_caption(app_data_dir: &Path, image_path: &Path) -> Result<String
 
 impl FlorenceCaptioner {
     pub fn new(app_data_dir: &Path) -> Result<Self> {
+        let started_at = Instant::now();
         let status = caption_model_status(app_data_dir);
         if !status.ready {
             anyhow::bail!(
@@ -239,48 +437,92 @@ impl FlorenceCaptioner {
         }
 
         let local_dir = model_dir(app_data_dir);
+        ensure_onnx_runtime(&local_dir)?;
         let tokenizer =
             Tokenizer::from_file(local_dir.join("tokenizer.json")).map_err(anyhow::Error::msg)?;
+        let caption_detail = caption_detail(app_data_dir);
 
-        let vision_session = create_session(&local_dir.join("onnx/vision_encoder_fp16.onnx"))?;
-        let embed_session = create_session(&local_dir.join("onnx/embed_tokens_fp16.onnx"))?;
-        let encoder_session = create_session(&local_dir.join("onnx/encoder_model_q4.onnx"))?;
-        let decoder_session = create_session(&local_dir.join("onnx/decoder_model_merged_q4.onnx"))?;
+        let sessions_started_at = Instant::now();
+        let acceleration = caption_acceleration(app_data_dir);
+        let vision_session = create_session(
+            &local_dir.join("onnx/vision_encoder_fp16.onnx"),
+            acceleration,
+            true,
+        )?;
+        let embed_session = create_session(
+            &local_dir.join("onnx/embed_tokens_fp16.onnx"),
+            acceleration,
+            false,
+        )?;
+        let encoder_session = create_session(
+            &local_dir.join("onnx/encoder_model_q4.onnx"),
+            acceleration,
+            false,
+        )?;
+        let decoder_prefill_session = create_session(
+            &local_dir.join("onnx/decoder_model_q4.onnx"),
+            acceleration,
+            false,
+        )?;
+        let decoder_session = create_session(
+            &local_dir.join("onnx/decoder_model_merged_q4.onnx"),
+            acceleration,
+            false,
+        )?;
+        println!(
+            "Florence sessions loaded in {:?} with {:?} acceleration (total init {:?})",
+            sessions_started_at.elapsed(),
+            acceleration,
+            started_at.elapsed()
+        );
 
         Ok(Self {
             tokenizer,
+            caption_detail,
             vision_session,
             embed_session,
             encoder_session,
+            decoder_prefill_session,
             decoder_session,
         })
     }
 
     pub fn generate(&mut self, image_path: &Path) -> Result<String> {
+        let started_at = Instant::now();
+        println!("Florence caption started: {}", image_path.display());
         let image_features = run_vision_encoder(&mut self.vision_session, image_path)?;
+        println!("Florence vision encoder done in {:?}", started_at.elapsed());
         let prompt_ids = self
             .tokenizer
-            .encode("What does the image describe?", false)
+            .encode(self.caption_detail.prompt(), false)
             .map_err(anyhow::Error::msg)?
             .get_ids()
             .iter()
             .map(|id| i64::from(*id))
             .collect::<Vec<_>>();
         let prompt_embeds = run_token_embedder(&mut self.embed_session, &prompt_ids)?;
-        let encoder_embeds = concatenate_sequence_embeddings(&prompt_embeds, &image_features)?;
+        println!(
+            "Florence token embeddings done in {:?}",
+            started_at.elapsed()
+        );
+        let encoder_embeds = concatenate_sequence_embeddings(&image_features, &prompt_embeds)?;
         let encoder_attention_mask = vec![1_i64; encoder_embeds.shape[1] as usize];
         let encoder_hidden_states = run_encoder(
             &mut self.encoder_session,
             &encoder_embeds,
             &encoder_attention_mask,
         )?;
+        println!("Florence encoder done in {:?}", started_at.elapsed());
 
         let generated_ids = run_decoder(
+            &mut self.decoder_prefill_session,
             &mut self.decoder_session,
             &mut self.embed_session,
             &encoder_hidden_states,
             &encoder_attention_mask,
+            self.caption_detail.max_new_tokens(),
         )?;
+        println!("Florence decoder done in {:?}", started_at.elapsed());
 
         let generated_u32 = generated_ids
             .into_iter()
@@ -330,6 +572,26 @@ fn probe_session(file: &'static str, path: &Path) -> Result<CaptionRuntimeSessio
             ],
             vec!["logits".to_string(), "present_key_values".to_string()],
         ),
+        "onnxruntime/onnxruntime.dll" => (
+            vec!["OrtGetApiBase".to_string()],
+            vec!["ORT DirectML 1.24.2".to_string()],
+        ),
+        "onnxruntime/onnxruntime_providers_shared.dll" => (
+            vec!["providers".to_string()],
+            vec!["shared provider bridge".to_string()],
+        ),
+        "onnxruntime/DirectML.dll" => (
+            vec!["DirectX 12".to_string()],
+            vec!["DirectML 1.15.4".to_string()],
+        ),
+        "onnx/decoder_model_q4.onnx" => (
+            vec![
+                "encoder_attention_mask".to_string(),
+                "encoder_hidden_states".to_string(),
+                "inputs_embeds".to_string(),
+            ],
+            vec!["logits".to_string(), "present_key_values".to_string()],
+        ),
         _ => (Vec::new(), Vec::new()),
     };
 
@@ -340,6 +602,111 @@ fn probe_session(file: &'static str, path: &Path) -> Result<CaptionRuntimeSessio
     })
 }
 
+/// Create an actual ORT session for the vision encoder and return a probe
+/// entry that reflects which execution provider was successfully loaded.
+/// This is the only session that can legitimately run on DirectML; all
+/// text/decoder sessions intentionally use CPU only.
+fn probe_vision_session(
+    path: &Path,
+    acceleration: CaptionAcceleration,
+) -> Result<CaptionRuntimeSessionProbe> {
+    let metadata = std::fs::metadata(path)?;
+    if metadata.len() == 0 {
+        anyhow::bail!("{} is empty", path.display());
+    }
+
+    // Try to create the session with the requested acceleration.  If DirectML
+    // was requested but fails we report what actually loaded.
+    let loaded_acceleration = match acceleration {
+        CaptionAcceleration::Cpu => {
+            create_session(path, CaptionAcceleration::Cpu, false)?;
+            CaptionAcceleration::Cpu
+        }
+        CaptionAcceleration::Auto => {
+            // `fail_silently` — session will fall back to CPU if DirectML
+            // is unavailable; we detect this by trying Directml explicitly.
+            let directml_ok = create_session(path, CaptionAcceleration::Directml, true).is_ok();
+            if directml_ok {
+                CaptionAcceleration::Directml
+            } else {
+                create_session(path, CaptionAcceleration::Cpu, false)?;
+                CaptionAcceleration::Cpu
+            }
+        }
+        CaptionAcceleration::Directml => {
+            create_session(path, CaptionAcceleration::Directml, true)?;
+            CaptionAcceleration::Directml
+        }
+    };
+
+    Ok(CaptionRuntimeSessionProbe {
+        file: "onnx/vision_encoder_fp16.onnx",
+        inputs: vec!["pixel_values".to_string()],
+        outputs: vec![format!("image_features [EP: {:?}]", loaded_acceleration)],
+    })
+}
+
+pub fn ensure_onnx_runtime(local_dir: &Path) -> Result<()> {
+    let dll_path = local_dir.join(ONNX_RUNTIME_DLL_FILE);
+    ORT_RUNTIME_INIT
+        .get_or_init(|| {
+            if !dll_path.exists() {
+                return Err(format!(
+                    "ONNX Runtime DLL is missing: {}",
+                    dll_path.display()
+                ));
+            }
+            ort::environment::init_from(&dll_path)
+                .map_err(|error| error.to_string())?
+                .with_name("phokus-florence")
+                .commit();
+            Ok(())
+        })
+        .clone()
+        .map_err(anyhow::Error::msg)
+}
+
+fn download_onnx_runtime_files(local_dir: &Path) -> Result<()> {
+    if !cfg!(target_os = "windows") {
+        anyhow::bail!(
+            "Florence-2 ONNX Runtime download is currently configured for Windows builds"
+        );
+    }
+
+    for (destination_file, source_url, archive_path) in ONNX_RUNTIME_FILES {
+        let destination = local_dir.join(destination_file);
+        download_nuget_file(source_url, archive_path, &destination)?;
+    }
+
+    Ok(())
+}
+
+fn download_nuget_file(source_url: &str, archive_path: &str, destination: &Path) -> Result<()> {
+    let mut response = ureq::get(source_url)
+        .call()
+        .map_err(|error| anyhow::anyhow!("{error}"))?;
+    let mut bytes = Vec::new();
+    response
+        .body_mut()
+        .as_reader()
+        .read_to_end(&mut bytes)
+        .map_err(|error| anyhow::anyhow!("{error}"))?;
+
+    let mut archive = zip::ZipArchive::new(Cursor::new(bytes))?;
+    let mut dll = archive.by_name(archive_path)?;
+    if let Some(parent) = destination.parent() {
+        std::fs::create_dir_all(parent)?;
+    }
+
+    let temp_destination = destination.with_extension("tmp");
+    {
+        let mut file = std::fs::File::create(&temp_destination)?;
+        std::io::copy(&mut dll, &mut file)?;
+    }
+    std::fs::rename(temp_destination, destination)?;
+    Ok(())
+}
+
 fn run_vision_encoder(session: &mut Session, image_path: &Path) -> Result<TensorData> {
     let pixels = preprocess_image(image_path)?;
     let input = Tensor::from_array(([1usize, 3, 768, 768], pixels.into_boxed_slice()))
@@ -387,24 +754,50 @@ fn run_encoder(
 }
 
 fn run_decoder(
+    decoder_prefill_session: &mut Session,
     decoder_session: &mut Session,
     embed_session: &mut Session,
     encoder_hidden_states: &TensorData,
     encoder_attention_mask: &[i64],
+    max_new_tokens: usize,
 ) -> Result<Vec<i64>> {
     const DECODER_LAYERS: usize = 6;
-    const DECODER_HEADS: usize = 12;
-    const HEAD_DIM: usize = 64;
     const DECODER_START_TOKEN_ID: i64 = 2;
+    const FORCED_BOS_TOKEN_ID: i64 = 0;
     const EOS_TOKEN_ID: i64 = 2;
-    const MAX_NEW_TOKENS: usize = 32;
 
     let mut generated = Vec::new();
-    let mut next_input_id = DECODER_START_TOKEN_ID;
-    let mut past: Vec<TensorData> = Vec::new();
-    let mut use_cache_branch = false;
+    let encoder_attention_mask_tensor = Tensor::from_array((
+        [1usize, encoder_attention_mask.len()],
+        encoder_attention_mask.to_vec().into_boxed_slice(),
+    ))
+    .map_err(|error| anyhow::anyhow!("{error}"))?;
+    let encoder_hidden_states_tensor = tensor_from_data(encoder_hidden_states)?;
+    let prefill_inputs_embeds = run_token_embedder(embed_session, &[DECODER_START_TOKEN_ID])?;
+    let prefill_inputs_embeds_tensor = tensor_from_data(&prefill_inputs_embeds)?;
+    let prefill_started_at = Instant::now();
+    let prefill_outputs = decoder_prefill_session
+        .run(ort::inputs! {
+            "encoder_attention_mask" => encoder_attention_mask_tensor,
+            "encoder_hidden_states" => encoder_hidden_states_tensor,
+            "inputs_embeds" => prefill_inputs_embeds_tensor
+        })
+        .map_err(|error| anyhow::anyhow!("{error}"))?;
+    println!(
+        "Florence decoder prefill done in {:?}",
+        prefill_started_at.elapsed()
+    );
+    let _prefill_logits = tensor_data(&prefill_outputs[0])?;
+    let mut next_input_id = FORCED_BOS_TOKEN_ID;
+    let encoder_kv = collect_present_key_values(&prefill_outputs, DECODER_LAYERS)?;
+    let mut decoder_kv = encoder_kv.clone();
+
+    for _ in 0..max_new_tokens {
+        if next_input_id == EOS_TOKEN_ID {
+            break;
+        }
+        generated.push(next_input_id);
 
-    for _ in 0..MAX_NEW_TOKENS {
         let inputs_embeds = run_token_embedder(embed_session, &[next_input_id])?;
         let encoder_attention_mask_tensor = Tensor::from_array((
             [1usize, encoder_attention_mask.len()],
@@ -413,9 +806,8 @@ fn run_decoder(
         .map_err(|error| anyhow::anyhow!("{error}"))?;
         let encoder_hidden_states_tensor = tensor_from_data(encoder_hidden_states)?;
         let inputs_embeds_tensor = tensor_from_data(&inputs_embeds)?;
-        let use_cache_branch_tensor =
-            Tensor::from_array(([1usize], vec![use_cache_branch].into_boxed_slice()))
-                .map_err(|error| anyhow::anyhow!("{error}"))?;
+        let use_cache_branch_tensor = Tensor::from_array(([1usize], vec![true].into_boxed_slice()))
+            .map_err(|error| anyhow::anyhow!("{error}"))?;
 
         let mut inputs = ort::inputs! {
             "encoder_attention_mask" => encoder_attention_mask_tensor,
@@ -424,92 +816,82 @@ fn run_decoder(
             "use_cache_branch" => use_cache_branch_tensor
         };
 
-        if past.is_empty() {
-            for layer in 0..DECODER_LAYERS {
-                push_tensor_input(
-                    &mut inputs,
-                    format!("past_key_values.{layer}.decoder.key"),
-                    TensorData::zeros(vec![1, DECODER_HEADS as i64, 0, HEAD_DIM as i64]),
-                )?;
-                push_tensor_input(
-                    &mut inputs,
-                    format!("past_key_values.{layer}.decoder.value"),
-                    TensorData::zeros(vec![1, DECODER_HEADS as i64, 0, HEAD_DIM as i64]),
-                )?;
-                push_tensor_input(
-                    &mut inputs,
-                    format!("past_key_values.{layer}.encoder.key"),
-                    TensorData::zeros(vec![1, DECODER_HEADS as i64, 0, HEAD_DIM as i64]),
-                )?;
-                push_tensor_input(
-                    &mut inputs,
-                    format!("past_key_values.{layer}.encoder.value"),
-                    TensorData::zeros(vec![1, DECODER_HEADS as i64, 0, HEAD_DIM as i64]),
-                )?;
-            }
-        } else {
-            for layer in 0..DECODER_LAYERS {
-                for cache_name in [
-                    "decoder.key",
-                    "decoder.value",
-                    "encoder.key",
-                    "encoder.value",
-                ] {
-                    let past_index = layer * 4
-                        + match cache_name {
-                            "decoder.key" => 0,
-                            "decoder.value" => 1,
-                            "encoder.key" => 2,
-                            "encoder.value" => 3,
-                            _ => unreachable!(),
-                        };
-                    push_tensor_input(
-                        &mut inputs,
-                        format!("past_key_values.{layer}.{cache_name}"),
-                        past[past_index].clone(),
-                    )?;
-                }
-            }
+        for layer in 0..DECODER_LAYERS {
+            push_tensor_input(
+                &mut inputs,
+                format!("past_key_values.{layer}.decoder.key"),
+                decoder_kv[layer * 4].clone(),
+            )?;
+            push_tensor_input(
+                &mut inputs,
+                format!("past_key_values.{layer}.decoder.value"),
+                decoder_kv[layer * 4 + 1].clone(),
+            )?;
+            push_tensor_input(
+                &mut inputs,
+                format!("past_key_values.{layer}.encoder.key"),
+                encoder_kv[layer * 4 + 2].clone(),
+            )?;
+            push_tensor_input(
+                &mut inputs,
+                format!("past_key_values.{layer}.encoder.value"),
+                encoder_kv[layer * 4 + 3].clone(),
+            )?;
         }
 
         let outputs = decoder_session
             .run(inputs)
             .map_err(|error| anyhow::anyhow!("{error}"))?;
         let logits = tensor_data(&outputs["logits"])?;
-        let token_id = argmax_last_token(&logits)?;
-        if token_id == EOS_TOKEN_ID {
-            break;
-        }
-        generated.push(token_id);
-        next_input_id = token_id;
-
-        past.clear();
-        for layer in 0..DECODER_LAYERS {
-            for cache_name in [
-                "decoder.key",
-                "decoder.value",
-                "encoder.key",
-                "encoder.value",
-            ] {
-                past.push(tensor_data(
-                    &outputs[format!("present.{layer}.{cache_name}").as_str()],
-                )?);
-            }
-        }
-        use_cache_branch = true;
+        next_input_id = argmax_last_token(&logits)?;
+        decoder_kv = collect_present_key_values(&outputs, DECODER_LAYERS)?;
     }
 
+    println!("Florence decoder produced {} token(s)", generated.len());
     Ok(generated)
 }
 
-fn create_session(path: &Path) -> Result<Session> {
+fn create_session(
+    path: &Path,
+    acceleration: CaptionAcceleration,
+    allow_directml: bool,
+) -> Result<Session> {
     let builder = Session::builder().map_err(|error| anyhow::anyhow!("{error}"))?;
     let builder = builder
         .with_optimization_level(GraphOptimizationLevel::Level3)
         .map_err(|error| anyhow::anyhow!("{error}"))?;
-    let mut builder = builder
+    let use_directml = allow_directml
+        && matches!(
+            acceleration,
+            CaptionAcceleration::Auto | CaptionAcceleration::Directml
+        );
+    let builder = builder
+        .with_memory_pattern(!use_directml)
+        .map_err(|error| anyhow::anyhow!("{error}"))?;
+    let builder = builder
+        .with_parallel_execution(false)
+        .map_err(|error| anyhow::anyhow!("{error}"))?;
+    let builder = builder
         .with_intra_threads(1)
         .map_err(|error| anyhow::anyhow!("{error}"))?;
+    let mut builder = match acceleration {
+        CaptionAcceleration::Cpu => builder,
+        CaptionAcceleration::Auto if use_directml => builder
+            .with_execution_providers([ep::DirectML::default().build().fail_silently()])
+            .unwrap_or_else(|error| error.recover()),
+        CaptionAcceleration::Directml if use_directml => builder
+            .with_execution_providers([ep::DirectML::default().build().error_on_failure()])
+            .map_err(|error| anyhow::anyhow!("{error}"))?,
+        CaptionAcceleration::Auto | CaptionAcceleration::Directml => {
+            // `allow_directml` is false for the 4 text/decoder sessions —
+            // they intentionally run on CPU regardless of the setting.
+            println!(
+                "Florence: using CPU for {} (DirectML disabled for this session type)",
+                path.display()
+            );
+            builder
+        }
+    };
     let session = builder
         .commit_from_file(path)
         .map_err(|error| anyhow::anyhow!("{error}"))?;
@@ -537,6 +919,24 @@ fn concatenate_sequence_embeddings(text: &TensorData, image: &TensorData) -> Res
     })
 }
 
+fn collect_present_key_values(
+    outputs: &SessionOutputs<'_>,
+    layers: usize,
+) -> Result<Vec<TensorData>> {
+    let expected = layers * 4;
+    if outputs.len() < expected + 1 {
+        anyhow::bail!(
+            "Decoder returned {} output(s), expected at least {}",
+            outputs.len(),
+            expected + 1
+        );
+    }
+
+    (1..=expected)
+        .map(|index| tensor_data(&outputs[index]))
+        .collect::<Result<Vec<_>>>()
+}
+
 fn tensor_data(value: &ort::value::DynValue) -> Result<TensorData> {
     let (shape, values) = value
         .try_extract_tensor::<f32>()
@@ -608,13 +1008,3 @@ fn preprocess_image(image_path: &Path) -> Result<Vec<f32>> {
 
     Ok(pixel_values)
 }
-
-impl TensorData {
-    fn zeros(shape: Vec<i64>) -> Self {
-        let values_len = shape.iter().map(|value| (*value).max(0) as usize).product();
-        Self {
-            shape,
-            values: vec![0.0; values_len],
-        }
-    }
-}

src-tauri/src/commands.rs

@@ -1,7 +1,11 @@
-use crate::captioner::{self, CaptionModelStatus, CaptionRuntimeProbe, CaptionVisionProbe};
-use crate::db::{self, DbPool, Folder, FolderJobProgress, ImageRecord};
+use crate::captioner::{
+    self, CaptionAcceleration, CaptionDetail, CaptionModelStatus, CaptionRuntimeProbe,
+    CaptionVisionProbe,
+};
+use crate::db::{self, DbPool, Folder, FolderJobProgress, ImageRecord, ImageTag};
 use crate::embedder;
 use crate::indexer;
+use crate::tagger::{self, TaggerAcceleration, TaggerModelStatus, TaggerRuntimeProbe};
 use crate::vector;
 use serde::{Deserialize, Serialize};
 use std::path::PathBuf;
@@ -39,12 +43,14 @@ pub struct UpdateImageDetailsParams {
 #[derive(Deserialize)]
 pub struct FindSimilarImagesParams {
     pub image_id: i64,
+    pub folder_id: Option<i64>,
     pub limit: Option<usize>,
 }
 
 #[derive(Deserialize)]
 pub struct DebugSimilarImagesParams {
     pub image_id: i64,
+    pub folder_id: Option<i64>,
     pub limit: Option<usize>,
 }
 
@@ -72,6 +78,26 @@ pub struct QueueCaptionJobsParams {
     pub image_id: Option<i64>,
 }
 
+#[derive(Deserialize)]
+pub struct ClearCaptionJobsParams {
+    pub folder_id: Option<i64>,
+}
+
+#[derive(Deserialize)]
+pub struct ResetCaptionsParams {
+    pub folder_id: Option<i64>,
+}
+
+#[derive(Deserialize)]
+pub struct SetCaptionAccelerationParams {
+    pub acceleration: CaptionAcceleration,
+}
+
+#[derive(Deserialize)]
+pub struct SetCaptionDetailParams {
+    pub detail: CaptionDetail,
+}
+
 #[derive(Deserialize)]
 pub struct ProbeCaptionImageParams {
     pub image_id: i64,
@@ -235,8 +261,8 @@ pub async fn find_similar_images(
         db::repair_embedding_consistency(&conn).map_err(|e| e.to_string())?;
         return Ok(Vec::new());
     }
-    let image_ids =
-        vector::find_similar_image_ids(&conn, params.image_id, limit).map_err(|e| e.to_string())?;
+    let image_ids = vector::find_similar_image_ids(&conn, params.image_id, limit, params.folder_id)
+        .map_err(|e| e.to_string())?;
     db::get_images_by_ids(&conn, &image_ids).map_err(|e| e.to_string())
 }
 
@@ -258,7 +284,8 @@ pub async fn debug_similar_images(
     let vector_count = vector::count_image_vectors(&conn).map_err(|e| e.to_string())?;
     let has_vector = vector::has_image_vector(&conn, params.image_id).map_err(|e| e.to_string())?;
     let similar_ids = if has_vector {
-        vector::find_similar_image_ids(&conn, params.image_id, limit).map_err(|e| e.to_string())?
+        vector::find_similar_image_ids(&conn, params.image_id, limit, params.folder_id)
+            .map_err(|e| e.to_string())?
     } else {
         Vec::new()
     };
@@ -332,6 +359,36 @@ pub async fn get_caption_model_status(app: AppHandle) -> Result<CaptionModelStat
     Ok(captioner::caption_model_status(&app_dir))
 }
 
+#[tauri::command]
+pub async fn get_caption_acceleration(app: AppHandle) -> Result<CaptionAcceleration, String> {
+    let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?;
+    Ok(captioner::caption_acceleration(&app_dir))
+}
+
+#[tauri::command]
+pub async fn set_caption_acceleration(
+    app: AppHandle,
+    params: SetCaptionAccelerationParams,
+) -> Result<CaptionAcceleration, String> {
+    let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?;
+    captioner::set_caption_acceleration(&app_dir, params.acceleration).map_err(|e| e.to_string())
+}
+
+#[tauri::command]
+pub async fn get_caption_detail(app: AppHandle) -> Result<CaptionDetail, String> {
+    let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?;
+    Ok(captioner::caption_detail(&app_dir))
+}
+
+#[tauri::command]
+pub async fn set_caption_detail(
+    app: AppHandle,
+    params: SetCaptionDetailParams,
+) -> Result<CaptionDetail, String> {
+    let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?;
+    captioner::set_caption_detail(&app_dir, params.detail).map_err(|e| e.to_string())
+}
+
 #[tauri::command]
 pub async fn prepare_caption_model(app: AppHandle) -> Result<CaptionModelStatus, String> {
     let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?;
@@ -441,6 +498,24 @@ pub async fn queue_caption_jobs(
     }
 }
 
+#[tauri::command]
+pub async fn clear_caption_jobs(
+    db: State<'_, DbState>,
+    params: ClearCaptionJobsParams,
+) -> Result<usize, String> {
+    let conn = db.get().map_err(|e| e.to_string())?;
+    db::clear_caption_jobs(&conn, params.folder_id).map_err(|e| e.to_string())
+}
+
+#[tauri::command]
+pub async fn reset_generated_captions(
+    db: State<'_, DbState>,
+    params: ResetCaptionsParams,
+) -> Result<usize, String> {
+    let conn = db.get().map_err(|e| e.to_string())?;
+    db::reset_generated_captions(&conn, params.folder_id).map_err(|e| e.to_string())
+}
+
 #[derive(Serialize, Deserialize)]
 pub struct TagCloudEntry {
     pub count: usize,
@@ -670,11 +745,27 @@ pub struct FolderWorkerStates {
     pub metadata_paused: bool,
     pub embedding_paused: bool,
     pub caption_paused: bool,
+    pub tagging_paused: bool,
 }
 
 #[tauri::command]
-pub async fn set_worker_paused(worker: String, folder_id: i64, paused: bool) -> Result<(), String> {
+pub async fn set_worker_paused(
+    db: State<'_, DbState>,
+    worker: String,
+    folder_id: i64,
+    paused: bool,
+) -> Result<(), String> {
     indexer::set_worker_paused(&worker, folder_id, paused);
+    if worker == "caption" && paused {
+        let conn = db.get().map_err(|e| e.to_string())?;
+        db::requeue_processing_caption_jobs_for_folder(&conn, folder_id)
+            .map_err(|e| e.to_string())?;
+    }
+    if worker == "tagging" && paused {
+        let conn = db.get().map_err(|e| e.to_string())?;
+        db::requeue_processing_tagging_jobs_for_folder(&conn, folder_id)
+            .map_err(|e| e.to_string())?;
+    }
     Ok(())
 }
 
@@ -684,23 +775,213 @@ pub async fn get_worker_states(folder_ids: Vec<i64>) -> Result<Vec<FolderWorkerS
     Ok(folder_ids
         .into_iter()
         .map(|folder_id| {
-            let state =
-                states
-                    .get(&folder_id)
-                    .copied()
-                    .unwrap_or(indexer::FolderWorkerPausedState {
-                        thumbnail: false,
-                        metadata: false,
-                        embedding: false,
-                        caption: false,
-                    });
+            let state = states
+                .get(&folder_id)
+                .copied()
+                .unwrap_or(indexer::FolderWorkerPausedState {
+                    thumbnail: false,
+                    metadata: false,
+                    embedding: false,
+                    caption: false,
+                    tagging: false,
+                });
             FolderWorkerStates {
                 folder_id,
                 thumbnail_paused: state.thumbnail,
                 metadata_paused: state.metadata,
                 embedding_paused: state.embedding,
                 caption_paused: state.caption,
+                tagging_paused: state.tagging,
             }
         })
         .collect())
 }
+
+// ---------------------------------------------------------------------------
+// Tagger commands
+// ---------------------------------------------------------------------------
+
+#[derive(Deserialize)]
+pub struct SetTaggerAccelerationParams {
+    pub acceleration: TaggerAcceleration,
+}
+
+#[derive(Deserialize)]
+pub struct SetTaggerThresholdParams {
+    pub threshold: f32,
+}
+
+#[derive(Deserialize)]
+pub struct QueueTaggingJobsParams {
+    pub folder_id: Option<i64>,
+    pub image_id: Option<i64>,
+}
+
+#[derive(Deserialize)]
+pub struct ClearTaggingJobsParams {
+    pub folder_id: Option<i64>,
+}
+
+#[derive(Deserialize)]
+pub struct GetImageTagsParams {
+    pub image_id: i64,
+}
+
+#[derive(Deserialize)]
+pub struct AddUserTagParams {
+    pub image_id: i64,
+    pub tag: String,
+}
+
+#[derive(Deserialize)]
+pub struct RemoveTagParams {
+    pub tag_id: i64,
+}
+
+#[tauri::command]
+pub async fn get_tagger_model_status(app: AppHandle) -> Result<TaggerModelStatus, String> {
+    let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?;
+    Ok(tagger::tagger_model_status(&app_dir))
+}
+
+#[tauri::command]
+pub async fn get_tagger_acceleration(app: AppHandle) -> Result<TaggerAcceleration, String> {
+    let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?;
+    Ok(tagger::tagger_acceleration(&app_dir))
+}
+
+#[tauri::command]
+pub async fn set_tagger_acceleration(
+    app: AppHandle,
+    params: SetTaggerAccelerationParams,
+) -> Result<TaggerAcceleration, String> {
+    let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?;
+    tagger::set_tagger_acceleration(&app_dir, params.acceleration).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())?;
+    tauri::async_runtime::spawn_blocking(move || tagger::probe_tagger_runtime(&app_dir))
+        .await
+        .map_err(|e| e.to_string())?
+        .map_err(|e| e.to_string())
+}
+
+#[tauri::command]
+pub async fn get_tagger_threshold(app: AppHandle) -> Result<f32, String> {
+    let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?;
+    Ok(tagger::tagger_threshold(&app_dir))
+}
+
+#[tauri::command]
+pub async fn set_tagger_threshold(
+    app: AppHandle,
+    params: SetTaggerThresholdParams,
+) -> Result<f32, String> {
+    let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?;
+    tagger::set_tagger_threshold(&app_dir, params.threshold).map_err(|e| e.to_string())
+}
+
+#[tauri::command]
+pub async fn prepare_tagger_model(app: AppHandle) -> Result<TaggerModelStatus, String> {
+    let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?;
+    tauri::async_runtime::spawn_blocking(move || {
+        let app = app.clone();
+        tagger::prepare_tagger_model_with_progress(&app_dir, move |progress| {
+            let _ = app.emit("tagger-model-progress", progress);
+        })
+    })
+    .await
+    .map_err(|e| e.to_string())?
+    .map_err(|e| e.to_string())
+}
+
+#[tauri::command]
+pub async fn delete_tagger_model(app: AppHandle) -> Result<TaggerModelStatus, String> {
+    let app_dir = app.path().app_data_dir().map_err(|e| e.to_string())?;
+    tauri::async_runtime::spawn_blocking(move || tagger::delete_tagger_model(&app_dir))
+        .await
+        .map_err(|e| e.to_string())?
+        .map_err(|e| e.to_string())
+}
+
+#[tauri::command]
+pub async fn queue_tagging_jobs(
+    app: AppHandle,
+    db: State<'_, DbState>,
+    params: QueueTaggingJobsParams,
+) -> Result<usize, String> {
+    let conn = db.get().map_err(|e| e.to_string())?;
+    let (total, folder_ids) = match (params.folder_id, params.image_id) {
+        (_, Some(image_id)) => {
+            db::enqueue_tagging_job(&conn, image_id).map_err(|e| e.to_string())?;
+            // Look up just this image's folder_id rather than fetching all folders
+            let image = db::get_image_by_id(&conn, image_id).map_err(|e| e.to_string())?;
+            (1usize, vec![image.folder_id])
+        }
+        (Some(folder_id), None) => {
+            let n = db::enqueue_missing_tagging_jobs_for_folder(&conn, folder_id)
+                .map_err(|e| e.to_string())?;
+            (n, vec![folder_id])
+        }
+        (None, None) => {
+            let folders = db::get_folders(&conn).map_err(|e| e.to_string())?;
+            let folder_ids: Vec<i64> = folders.iter().map(|f| f.id).collect();
+            let mut total = 0usize;
+            for &folder_id in &folder_ids {
+                total += db::enqueue_missing_tagging_jobs_for_folder(&conn, folder_id)
+                    .map_err(|e| e.to_string())?;
+            }
+            (total, folder_ids)
+        }
+    };
+    drop(conn);
+    indexer::emit_folder_job_progress(&app, db.inner(), &folder_ids, true);
+    Ok(total)
+}
+
+#[tauri::command]
+pub async fn clear_tagging_jobs(
+    app: AppHandle,
+    db: State<'_, DbState>,
+    params: ClearTaggingJobsParams,
+) -> Result<usize, String> {
+    let conn = db.get().map_err(|e| e.to_string())?;
+    let n = db::clear_tagging_jobs(&conn, params.folder_id).map_err(|e| e.to_string())?;
+    let folder_ids: Vec<i64> = match params.folder_id {
+        Some(id) => vec![id],
+        None => db::get_folders(&conn)
+            .map_err(|e| e.to_string())?
+            .into_iter()
+            .map(|f| f.id)
+            .collect(),
+    };
+    drop(conn);
+    indexer::emit_folder_job_progress(&app, db.inner(), &folder_ids, true);
+    Ok(n)
+}
+
+#[tauri::command]
+pub async fn get_image_tags(
+    db: State<'_, DbState>,
+    params: GetImageTagsParams,
+) -> Result<Vec<ImageTag>, String> {
+    let conn = db.get().map_err(|e| e.to_string())?;
+    db::get_image_tags(&conn, params.image_id).map_err(|e| e.to_string())
+}
+
+#[tauri::command]
+pub async fn add_user_tag(
+    db: State<'_, DbState>,
+    params: AddUserTagParams,
+) -> Result<ImageTag, String> {
+    let conn = db.get().map_err(|e| e.to_string())?;
+    db::add_user_tag(&conn, params.image_id, &params.tag).map_err(|e| e.to_string())
+}
+
+#[tauri::command]
+pub async fn remove_tag(db: State<'_, DbState>, params: RemoveTagParams) -> Result<(), String> {
+    let conn = db.get().map_err(|e| e.to_string())?;
+    db::remove_tag(&conn, params.tag_id).map_err(|e| e.to_string())
+}

src-tauri/src/db.rs

@@ -61,6 +61,10 @@ pub struct ImageRecord {
     pub caption_model: Option<String>,
     pub caption_updated_at: Option<String>,
     pub caption_error: Option<String>,
+    pub ai_rating: Option<String>,
+    pub ai_tagger_model: Option<String>,
+    pub ai_tagged_at: Option<String>,
+    pub ai_tagger_error: Option<String>,
 }
 
 #[allow(dead_code)]
@@ -100,6 +104,24 @@ pub struct CaptionJob {
     pub path: String,
 }
 
+#[derive(Debug, Clone)]
+pub struct TaggingJob {
+    pub image_id: i64,
+    pub folder_id: i64,
+    pub path: String,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ImageTag {
+    pub id: i64,
+    pub image_id: i64,
+    pub tag: String,
+    pub source: String,
+    pub ai_model: Option<String>,
+    pub confidence: Option<f64>,
+    pub created_at: String,
+}
+
 #[derive(Debug, Clone)]
 pub struct IndexedMediaEntry {
     pub id: i64,
@@ -120,6 +142,9 @@ pub struct FolderJobProgress {
     pub caption_pending: i64,
     pub caption_ready: i64,
     pub caption_failed: i64,
+    pub tagging_pending: i64,
+    pub tagging_ready: i64,
+    pub tagging_failed: i64,
 }
 
 pub fn create_pool(db_path: &Path) -> Result<DbPool> {
@@ -195,6 +220,26 @@ pub fn migrate(conn: &Connection) -> Result<()> {
             updated_at  TEXT NOT NULL DEFAULT (datetime('now'))
         );
 
+        CREATE TABLE IF NOT EXISTS tagging_jobs (
+            image_id    INTEGER PRIMARY KEY REFERENCES images(id) ON DELETE CASCADE,
+            status      TEXT NOT NULL DEFAULT 'pending',
+            attempts    INTEGER NOT NULL DEFAULT 0,
+            last_error  TEXT,
+            created_at  TEXT NOT NULL DEFAULT (datetime('now')),
+            updated_at  TEXT NOT NULL DEFAULT (datetime('now'))
+        );
+
+        CREATE TABLE IF NOT EXISTS image_tags (
+            id         INTEGER PRIMARY KEY AUTOINCREMENT,
+            image_id   INTEGER NOT NULL REFERENCES images(id) ON DELETE CASCADE,
+            tag        TEXT NOT NULL,
+            source     TEXT NOT NULL DEFAULT 'user',
+            ai_model   TEXT,
+            confidence REAL,
+            created_at TEXT NOT NULL DEFAULT (datetime('now')),
+            UNIQUE(image_id, tag)
+        );
+
         CREATE TABLE IF NOT EXISTS tag_cloud_cache (
             folder_scope    TEXT PRIMARY KEY,
             image_ids_hash  INTEGER NOT NULL,
@@ -208,6 +253,10 @@ pub fn migrate(conn: &Connection) -> Result<()> {
         CREATE INDEX IF NOT EXISTS idx_thumbnail_jobs_status ON thumbnail_jobs(status);
         CREATE INDEX IF NOT EXISTS idx_metadata_jobs_status ON metadata_jobs(status);
         CREATE INDEX IF NOT EXISTS idx_caption_jobs_status ON caption_jobs(status);
+        CREATE INDEX IF NOT EXISTS idx_tagging_jobs_status ON tagging_jobs(status);
+        CREATE INDEX IF NOT EXISTS idx_image_tags_image_id ON image_tags(image_id);
+        CREATE INDEX IF NOT EXISTS idx_image_tags_source ON image_tags(source);
+        CREATE INDEX IF NOT EXISTS idx_image_tags_tag ON image_tags(tag);
         ",
     )?;
 
@@ -237,6 +286,10 @@ pub fn migrate(conn: &Connection) -> Result<()> {
     ensure_column(conn, "images", "caption_model", "TEXT")?;
     ensure_column(conn, "images", "caption_updated_at", "TEXT")?;
     ensure_column(conn, "images", "caption_error", "TEXT")?;
+    ensure_column(conn, "images", "ai_rating", "TEXT")?;
+    ensure_column(conn, "images", "ai_tagger_model", "TEXT")?;
+    ensure_column(conn, "images", "ai_tagged_at", "TEXT")?;
+    ensure_column(conn, "images", "ai_tagger_error", "TEXT")?;
 
     vector::migrate(conn)?;
     Ok(())
@@ -257,8 +310,8 @@ pub fn insert_folder(conn: &Connection, path: &str, name: &str) -> Result<i64> {
 
 pub fn upsert_image(conn: &Connection, img: &ImageRecord) -> Result<i64> {
     let id = conn.query_row(
-        "INSERT INTO images (folder_id, path, filename, thumbnail_path, width, height, file_size, created_at, modified_at, mime_type, media_kind, duration_ms, video_codec, audio_codec, metadata_updated_at, metadata_error, favorite, rating, embedding_status, embedding_model, embedding_updated_at, embedding_error, generated_caption, caption_model, caption_updated_at, caption_error)
-         VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18, ?19, ?20, ?21, ?22, ?23, ?24, ?25, ?26)
+        "INSERT INTO images (folder_id, path, filename, thumbnail_path, width, height, file_size, created_at, modified_at, mime_type, media_kind, duration_ms, video_codec, audio_codec, metadata_updated_at, metadata_error, favorite, rating, embedding_status, embedding_model, embedding_updated_at, embedding_error, generated_caption, caption_model, caption_updated_at, caption_error, ai_rating, ai_tagger_model, ai_tagged_at, ai_tagger_error)
+         VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18, ?19, ?20, ?21, ?22, ?23, ?24, ?25, ?26, ?27, ?28, ?29, ?30)
          ON CONFLICT(path) DO UPDATE SET
             folder_id = excluded.folder_id,
             filename = excluded.filename,
@@ -311,6 +364,10 @@ pub fn upsert_image(conn: &Connection, img: &ImageRecord) -> Result<i64> {
             img.caption_model,
             img.caption_updated_at,
             img.caption_error,
+            img.ai_rating,
+            img.ai_tagger_model,
+            img.ai_tagged_at,
+            img.ai_tagger_error,
         ],
         |row| row.get(0),
     )?;
@@ -433,6 +490,13 @@ pub fn reset_inflight_jobs(conn: &Connection) -> Result<()> {
         "UPDATE caption_jobs SET status = 'pending' WHERE status = 'processing'",
         [],
     )?;
+    conn.execute(
+        "UPDATE tagging_jobs SET status = 'pending' WHERE status = 'processing'",
+        [],
+    )?;
+    // Delete any rows that were cancelled before the previous shutdown so
+    // they don't silently linger in the DB across restarts.
+    conn.execute("DELETE FROM tagging_jobs WHERE status = 'cancelled'", [])?;
     Ok(())
 }
 
@@ -493,6 +557,110 @@ pub fn requeue_caption_jobs(conn: &Connection, image_ids: &[i64]) -> Result<()>
     Ok(())
 }
 
+pub fn requeue_processing_caption_jobs_for_folder(
+    conn: &Connection,
+    folder_id: i64,
+) -> Result<usize> {
+    conn.execute(
+        "UPDATE caption_jobs
+         SET status = 'pending', updated_at = datetime('now')
+         WHERE status = 'processing'
+           AND image_id IN (
+               SELECT id FROM images WHERE folder_id = ?1
+           )",
+        [folder_id],
+    )
+    .map_err(Into::into)
+}
+
+pub fn clear_caption_jobs(conn: &Connection, folder_id: Option<i64>) -> Result<usize> {
+    match folder_id {
+        Some(folder_id) => {
+            let deleted = conn.execute(
+                "DELETE FROM caption_jobs
+                 WHERE image_id IN (
+                    SELECT id FROM images WHERE folder_id = ?1
+                 )",
+                [folder_id],
+            )?;
+            conn.execute(
+                "UPDATE images
+                 SET caption_error = NULL
+                 WHERE folder_id = ?1
+                   AND generated_caption IS NULL",
+                [folder_id],
+            )?;
+            Ok(deleted)
+        }
+        None => {
+            let deleted = conn.execute("DELETE FROM caption_jobs", [])?;
+            conn.execute(
+                "UPDATE images
+                 SET caption_error = NULL
+                 WHERE generated_caption IS NULL",
+                [],
+            )?;
+            Ok(deleted)
+        }
+    }
+}
+
+pub fn reset_generated_captions(conn: &Connection, folder_id: Option<i64>) -> Result<usize> {
+    let image_ids = match folder_id {
+        Some(folder_id) => {
+            let mut stmt = conn.prepare(
+                "SELECT id FROM images WHERE folder_id = ?1 AND generated_caption IS NOT NULL",
+            )?;
+            let rows = stmt.query_map([folder_id], |row| row.get::<_, i64>(0))?;
+            rows.collect::<rusqlite::Result<Vec<_>>>()?
+        }
+        None => {
+            let mut stmt =
+                conn.prepare("SELECT id FROM images WHERE generated_caption IS NOT NULL")?;
+            let rows = stmt.query_map([], |row| row.get::<_, i64>(0))?;
+            rows.collect::<rusqlite::Result<Vec<_>>>()?
+        }
+    };
+
+    let tx = conn.unchecked_transaction()?;
+    for image_id in &image_ids {
+        vector::delete_caption_embedding(&tx, *image_id)?;
+    }
+
+    match folder_id {
+        Some(folder_id) => {
+            tx.execute(
+                "UPDATE images
+                 SET generated_caption = NULL,
+                     caption_model = NULL,
+                     caption_updated_at = NULL,
+                     caption_error = NULL
+                 WHERE folder_id = ?1",
+                [folder_id],
+            )?;
+            tx.execute(
+                "DELETE FROM caption_jobs
+                 WHERE image_id IN (SELECT id FROM images WHERE folder_id = ?1)",
+                [folder_id],
+            )?;
+        }
+        None => {
+            tx.execute(
+                "UPDATE images
+                 SET generated_caption = NULL,
+                     caption_model = NULL,
+                     caption_updated_at = NULL,
+                     caption_error = NULL",
+                [],
+            )?;
+            tx.execute("DELETE FROM caption_jobs", [])?;
+        }
+    }
+
+    tx.commit()?;
+    Ok(image_ids.len())
+}
+
 pub fn enqueue_missing_caption_jobs_for_folder(conn: &Connection, folder_id: i64) -> Result<usize> {
     let inserted = conn.execute(
         "INSERT INTO caption_jobs (image_id, status, attempts, last_error, created_at, updated_at)
@@ -791,9 +959,36 @@ pub fn get_folder_job_progress(conn: &Connection, folder_id: i64) -> Result<Fold
     )?;
 
     let caption_failed = conn.query_row(
+        "SELECT COUNT(*)
+         FROM caption_jobs j
+         JOIN images i ON i.id = j.image_id
+         WHERE i.folder_id = ?1 AND j.status = 'failed'",
+        [folder_id],
+        |row| row.get(0),
+    )?;
+
+    let tagging_pending = conn.query_row(
+        "SELECT COUNT(*)
+         FROM tagging_jobs j
+         JOIN images i ON i.id = j.image_id
+         WHERE i.folder_id = ?1 AND j.status IN ('pending', 'processing')",
+        [folder_id],
+        |row| row.get(0),
+    )?;
+
+    let tagging_ready = conn.query_row(
         "SELECT COUNT(*)
          FROM images
-         WHERE folder_id = ?1 AND caption_error IS NOT NULL",
+         WHERE folder_id = ?1 AND ai_tagged_at IS NOT NULL",
+        [folder_id],
+        |row| row.get(0),
+    )?;
+
+    let tagging_failed = conn.query_row(
+        "SELECT COUNT(*)
+         FROM tagging_jobs j
+         JOIN images i ON i.id = j.image_id
+         WHERE i.folder_id = ?1 AND j.status = 'failed'",
         [folder_id],
         |row| row.get(0),
     )?;
@@ -808,6 +1003,9 @@ pub fn get_folder_job_progress(conn: &Connection, folder_id: i64) -> Result<Fold
         caption_pending,
         caption_ready,
         caption_failed,
+        tagging_pending,
+        tagging_ready,
+        tagging_failed,
     })
 }
 
@@ -1046,7 +1244,8 @@ pub fn update_image_details(
         "SELECT id, folder_id, path, filename, thumbnail_path, width, height, file_size, created_at, modified_at, mime_type,
                 media_kind, duration_ms, video_codec, audio_codec, metadata_updated_at, metadata_error,
                 favorite, rating, embedding_status, embedding_model, embedding_updated_at, embedding_error,
-                generated_caption, caption_model, caption_updated_at, caption_error
+                generated_caption, caption_model, caption_updated_at, caption_error,
+                ai_rating, ai_tagger_model, ai_tagged_at, ai_tagger_error
          FROM images
          WHERE id = ?1",
         [image_id],
@@ -1060,7 +1259,8 @@ pub fn get_image_by_id(conn: &Connection, image_id: i64) -> Result<ImageRecord>
         "SELECT id, folder_id, path, filename, thumbnail_path, width, height, file_size, created_at, modified_at, mime_type,
                 media_kind, duration_ms, video_codec, audio_codec, metadata_updated_at, metadata_error,
                 favorite, rating, embedding_status, embedding_model, embedding_updated_at, embedding_error,
-                generated_caption, caption_model, caption_updated_at, caption_error
+                generated_caption, caption_model, caption_updated_at, caption_error,
+                ai_rating, ai_tagger_model, ai_tagged_at, ai_tagger_error
          FROM images
          WHERE id = ?1",
         [image_id],
@@ -1126,7 +1326,8 @@ pub fn get_images(
         "SELECT id, folder_id, path, filename, thumbnail_path, width, height, file_size, created_at, modified_at, mime_type,
                 media_kind, duration_ms, video_codec, audio_codec, metadata_updated_at, metadata_error,
                 favorite, rating, embedding_status, embedding_model, embedding_updated_at, embedding_error,
-                generated_caption, caption_model, caption_updated_at, caption_error
+                generated_caption, caption_model, caption_updated_at, caption_error,
+                ai_rating, ai_tagger_model, ai_tagged_at, ai_tagger_error
          FROM images
          WHERE (?1 IS NULL OR folder_id = ?1)
            AND (?2 IS NULL OR filename LIKE ?2)
@@ -1240,6 +1441,299 @@ pub fn mark_caption_failed(conn: &Connection, image_id: i64, error: &str) -> Res
     Ok(())
 }
 
+// ---------------------------------------------------------------------------
+// Tagging jobs
+// ---------------------------------------------------------------------------
+
+pub fn enqueue_tagging_job(conn: &Connection, image_id: i64) -> Result<()> {
+    conn.execute(
+        "INSERT INTO tagging_jobs (image_id, status, attempts, last_error, created_at, updated_at)
+         VALUES (?1, 'pending', 0, NULL, datetime('now'), datetime('now'))
+         ON CONFLICT(image_id) DO UPDATE SET
+            status = CASE WHEN status != 'processing' THEN 'pending' ELSE status END,
+            last_error = CASE WHEN status != 'processing' THEN NULL ELSE last_error END,
+            updated_at = datetime('now')",
+        [image_id],
+    )?;
+    Ok(())
+}
+
+pub fn claim_tagging_jobs(
+    conn: &mut Connection,
+    paused_folder_ids: &std::collections::HashSet<i64>,
+    limit: usize,
+) -> Result<Vec<TaggingJob>> {
+    let tx = conn.transaction()?;
+    let candidates = get_pending_tagging_jobs_excluding(&tx, paused_folder_ids, limit * 2)?;
+    let mut claimed = Vec::new();
+    for job in candidates {
+        let updated = tx.execute(
+            "UPDATE tagging_jobs
+             SET status = 'processing', attempts = attempts + 1, updated_at = datetime('now')
+             WHERE image_id = ?1 AND status = 'pending'",
+            [job.image_id],
+        )?;
+        if updated > 0 {
+            claimed.push(job);
+            if claimed.len() >= limit {
+                break;
+            }
+        }
+    }
+    tx.commit()?;
+    Ok(claimed)
+}
+
+fn get_pending_tagging_jobs_excluding(
+    conn: &Connection,
+    paused_folder_ids: &std::collections::HashSet<i64>,
+    limit: usize,
+) -> Result<Vec<TaggingJob>> {
+    let mut stmt = conn.prepare(
+        "SELECT j.image_id, i.folder_id, i.path
+         FROM tagging_jobs j
+         JOIN images i ON i.id = j.image_id
+         WHERE j.status = 'pending'
+         ORDER BY j.created_at ASC
+         LIMIT ?1",
+    )?;
+    let rows = stmt
+        .query_map([limit as i64], |row| {
+            Ok(TaggingJob {
+                image_id: row.get(0)?,
+                folder_id: row.get(1)?,
+                path: row.get(2)?,
+            })
+        })?
+        .collect::<rusqlite::Result<Vec<_>>>()?;
+    Ok(rows
+        .into_iter()
+        .filter(|job| !paused_folder_ids.contains(&job.folder_id))
+        .collect())
+}
+
+pub fn update_ai_tags(
+    conn: &Connection,
+    image_id: i64,
+    tags: &[(String, f64)],
+    rating: &str,
+    model: &str,
+) -> Result<()> {
+    // NOTE: callers are responsible for wrapping this in a transaction.
+    // Do NOT open a nested transaction here — rusqlite/SQLite do not support
+    // nested transactions and will error with "cannot start a transaction
+    // within a transaction".
+
+    // Remove previous AI tags for this image before inserting fresh ones
+    conn.execute(
+        "DELETE FROM image_tags WHERE image_id = ?1 AND source = 'ai'",
+        [image_id],
+    )?;
+
+    let mut stmt = conn.prepare(
+        "INSERT INTO image_tags (image_id, tag, source, ai_model, confidence, created_at)
+         VALUES (?1, ?2, 'ai', ?3, ?4, datetime('now'))
+         ON CONFLICT(image_id, tag) DO UPDATE SET
+            source = 'ai',
+            ai_model = excluded.ai_model,
+            confidence = excluded.confidence",
+    )?;
+    for (tag, confidence) in tags {
+        stmt.execute(params![image_id, tag, model, confidence])?;
+    }
+    drop(stmt);
+
+    conn.execute(
+        "UPDATE images
+         SET ai_rating = ?2,
+             ai_tagger_model = ?3,
+             ai_tagged_at = datetime('now'),
+             ai_tagger_error = NULL
+         WHERE id = ?1",
+        params![image_id, rating, model],
+    )?;
+    conn.execute("DELETE FROM tagging_jobs WHERE image_id = ?1", [image_id])?;
+    Ok(())
+}
+
+pub fn mark_tagging_failed(conn: &Connection, image_id: i64, error: &str) -> Result<()> {
+    conn.execute(
+        "UPDATE images SET ai_tagger_error = ?2 WHERE id = ?1",
+        params![image_id, error],
+    )?;
+    conn.execute(
+        "UPDATE tagging_jobs
+         SET status = 'failed', last_error = ?2, updated_at = datetime('now')
+         WHERE image_id = ?1",
+        params![image_id, error],
+    )?;
+    Ok(())
+}
+
+pub fn clear_tagging_jobs(conn: &Connection, folder_id: Option<i64>) -> Result<usize> {
+    // Rows currently being processed by the worker must not be deleted mid-flight
+    // — the worker holds a reference and will write results back after inference.
+    // Mark them 'cancelled' so the worker discards the result instead of saving it,
+    // then delete every non-processing row.  On the next poll the worker will find
+    // no pending work and the queue will appear empty.
+    let deleted = match folder_id {
+        Some(fid) => {
+            conn.execute(
+                "UPDATE tagging_jobs
+                 SET status = 'cancelled', last_error = NULL, updated_at = datetime('now')
+                 WHERE status = 'processing'
+                   AND image_id IN (SELECT id FROM images WHERE folder_id = ?1)",
+                [fid],
+            )?;
+            let n = conn.execute(
+                "DELETE FROM tagging_jobs
+                 WHERE status != 'processing'
+                   AND image_id IN (SELECT id FROM images WHERE folder_id = ?1)",
+                [fid],
+            )?;
+            conn.execute(
+                "UPDATE images
+                 SET ai_tagger_error = NULL
+                 WHERE folder_id = ?1 AND ai_tagged_at IS NULL",
+                [fid],
+            )?;
+            n
+        }
+        None => {
+            conn.execute(
+                "UPDATE tagging_jobs
+                 SET status = 'cancelled', last_error = NULL, updated_at = datetime('now')
+                 WHERE status = 'processing'",
+                [],
+            )?;
+            let n = conn.execute("DELETE FROM tagging_jobs WHERE status != 'processing'", [])?;
+            conn.execute(
+                "UPDATE images SET ai_tagger_error = NULL WHERE ai_tagged_at IS NULL",
+                [],
+            )?;
+            n
+        }
+    };
+    Ok(deleted)
+}
+
+pub fn get_image_tags(conn: &Connection, image_id: i64) -> Result<Vec<ImageTag>> {
+    let mut stmt = conn.prepare(
+        "SELECT id, image_id, tag, source, ai_model, confidence, created_at
+         FROM image_tags
+         WHERE image_id = ?1
+         ORDER BY source DESC, confidence DESC NULLS LAST, tag ASC",
+    )?;
+    let rows = stmt
+        .query_map([image_id], |row| {
+            Ok(ImageTag {
+                id: row.get(0)?,
+                image_id: row.get(1)?,
+                tag: row.get(2)?,
+                source: row.get(3)?,
+                ai_model: row.get(4)?,
+                confidence: row.get(5)?,
+                created_at: row.get(6)?,
+            })
+        })?
+        .collect::<rusqlite::Result<Vec<_>>>()?;
+    Ok(rows)
+}
+
+pub fn add_user_tag(conn: &Connection, image_id: i64, tag: &str) -> Result<ImageTag> {
+    conn.execute(
+        "INSERT INTO image_tags (image_id, tag, source, ai_model, confidence, created_at)
+         VALUES (?1, ?2, 'user', NULL, NULL, datetime('now'))
+         ON CONFLICT(image_id, tag) DO NOTHING",
+        params![image_id, tag],
+    )?;
+    let row = conn.query_row(
+        "SELECT id, image_id, tag, source, ai_model, confidence, created_at
+         FROM image_tags WHERE image_id = ?1 AND tag = ?2",
+        params![image_id, tag],
+        |row| {
+            Ok(ImageTag {
+                id: row.get(0)?,
+                image_id: row.get(1)?,
+                tag: row.get(2)?,
+                source: row.get(3)?,
+                ai_model: row.get(4)?,
+                confidence: row.get(5)?,
+                created_at: row.get(6)?,
+            })
+        },
+    )?;
+    Ok(row)
+}
+
+pub fn remove_tag(conn: &Connection, tag_id: i64) -> Result<()> {
+    conn.execute("DELETE FROM image_tags WHERE id = ?1", [tag_id])?;
+    Ok(())
+}
+
+pub fn enqueue_missing_tagging_jobs_for_folder(conn: &Connection, folder_id: i64) -> Result<usize> {
+    let inserted = conn.execute(
+        "INSERT INTO tagging_jobs (image_id, status, attempts, last_error, created_at, updated_at)
+         SELECT id, 'pending', 0, NULL, datetime('now'), datetime('now')
+         FROM images
+         WHERE folder_id = ?1
+           AND media_kind = 'image'
+           AND ai_tagged_at IS NULL
+         ON CONFLICT(image_id) DO UPDATE SET
+            -- Only reset to 'pending' if the row is not currently being
+            -- processed or cancelled; leave 'processing' rows untouched so
+            -- the worker can complete them and clean up normally.
+            status = CASE WHEN status NOT IN ('processing', 'cancelled') THEN 'pending' ELSE status END,
+            last_error = CASE WHEN status != 'processing' THEN NULL ELSE last_error END,
+            updated_at = datetime('now')",
+        [folder_id],
+    )?;
+    conn.execute(
+        "UPDATE images
+         SET ai_tagger_error = NULL
+         WHERE folder_id = ?1
+           AND media_kind = 'image'
+           AND ai_tagged_at IS NULL",
+        [folder_id],
+    )?;
+    Ok(inserted)
+}
+
+pub fn requeue_processing_tagging_jobs_for_folder(conn: &Connection, folder_id: i64) -> Result<()> {
+    conn.execute(
+        "UPDATE tagging_jobs
+         SET status = 'pending', updated_at = datetime('now')
+         WHERE status = 'processing'
+           AND image_id IN (SELECT id FROM images WHERE folder_id = ?1)",
+        [folder_id],
+    )?;
+    Ok(())
+}
+
+/// Returns `true` when the job row for `image_id` currently has status = 'cancelled'.
+/// Used by the worker to discard inference results for jobs that were cancelled
+/// while inference was running.
+pub fn is_tagging_job_cancelled(conn: &Connection, image_id: i64) -> Result<bool> {
+    let count: i64 = conn.query_row(
+        "SELECT COUNT(*) FROM tagging_jobs WHERE image_id = ?1 AND status = 'cancelled'",
+        [image_id],
+        |row| row.get(0),
+    )?;
+    Ok(count > 0)
+}
+
+pub fn requeue_tagging_jobs(conn: &Connection, image_ids: &[i64]) -> Result<()> {
+    for image_id in image_ids {
+        conn.execute(
+            "UPDATE tagging_jobs
+             SET status = 'pending', updated_at = datetime('now')
+             WHERE image_id = ?1 AND status = 'processing'",
+            [image_id],
+        )?;
+    }
+    Ok(())
+}
+
 pub fn suggest_tags_from_caption(
     conn: &Connection,
     image_id: i64,
@@ -1332,6 +1826,10 @@ fn map_image_row(row: &Row<'_>) -> rusqlite::Result<ImageRecord> {
         caption_model: row.get(24)?,
         caption_updated_at: row.get(25)?,
         caption_error: row.get(26)?,
+        ai_rating: row.get(27)?,
+        ai_tagger_model: row.get(28)?,
+        ai_tagged_at: row.get(29)?,
+        ai_tagger_error: row.get(30)?,
     })
 }
 

src-tauri/src/indexer.rs

@@ -3,6 +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::thumbnail;
 use crate::vector;
 use anyhow::Result;
@@ -35,6 +36,7 @@ struct PausedWorkerFolders {
     metadata: HashSet<i64>,
     embedding: HashSet<i64>,
     caption: HashSet<i64>,
+    tagging: HashSet<i64>,
 }
 
 #[derive(Clone, Copy)]
@@ -43,6 +45,7 @@ pub struct FolderWorkerPausedState {
     pub metadata: bool,
     pub embedding: bool,
     pub caption: bool,
+    pub tagging: bool,
 }
 
 pub fn set_worker_paused(worker: &str, folder_id: i64, paused: bool) {
@@ -55,6 +58,7 @@ pub fn set_worker_paused(worker: &str, folder_id: i64, paused: bool) {
             "metadata" => Some(&mut paused_folders.metadata),
             "embedding" => Some(&mut paused_folders.embedding),
             "caption" => Some(&mut paused_folders.caption),
+            "tagging" => Some(&mut paused_folders.tagging),
             _ => None,
         };
 
@@ -87,6 +91,7 @@ pub fn get_worker_paused_states(folder_ids: &[i64]) -> HashMap<i64, FolderWorker
                     metadata: paused_folders.metadata.contains(&folder_id),
                     embedding: paused_folders.embedding.contains(&folder_id),
                     caption: paused_folders.caption.contains(&folder_id),
+                    tagging: paused_folders.tagging.contains(&folder_id),
                 },
             )
         })
@@ -105,6 +110,8 @@ fn paused_folder_ids(worker: &str) -> HashSet<i64> {
         "thumbnail" => paused_folders.thumbnail.clone(),
         "metadata" => paused_folders.metadata.clone(),
         "embedding" => paused_folders.embedding.clone(),
+        "caption" => paused_folders.caption.clone(),
+        "tagging" => paused_folders.tagging.clone(),
         _ => HashSet::new(),
     }
 }
@@ -191,6 +198,12 @@ pub fn start_caption_worker(app: AppHandle, pool: DbPool, app_data_dir: PathBuf)
         let mut captioner: Option<FlorenceCaptioner> = None;
         println!("Caption worker started.");
         loop {
+            // If the acceleration setting changed, drop the cached session so
+            // the next batch picks it up with the new execution provider.
+            if captioner::CAPTION_SESSION_DIRTY.swap(false, std::sync::atomic::Ordering::Relaxed) {
+                println!("Caption worker: acceleration setting changed — resetting session.");
+                captioner = None;
+            }
             if let Err(error) = process_caption_batch(&app, &pool, &app_data_dir, &mut captioner) {
                 eprintln!("Caption worker error: {}", error);
                 captioner = None;
@@ -200,6 +213,28 @@ 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;
+        println!("Tagging worker started.");
+        loop {
+            // If the acceleration setting changed, drop the cached session so
+            // the next batch picks it up with the new execution provider.
+            if tagger::TAGGER_SESSION_DIRTY.swap(false, std::sync::atomic::Ordering::Relaxed) {
+                println!("Tagging worker: acceleration setting changed — resetting session.");
+                tagger_instance = None;
+            }
+            if let Err(error) =
+                process_tagging_batch(&app, &pool, &app_data_dir, &mut tagger_instance)
+            {
+                eprintln!("Tagging worker error: {}", error);
+                tagger_instance = None;
+            }
+            std::thread::sleep(std::time::Duration::from_millis(750));
+        }
+    });
+}
+
 fn do_index(app: AppHandle, pool: DbPool, folder_id: i64, folder_path: PathBuf) -> Result<()> {
     let existing_entries = {
         let conn = pool.get()?;
@@ -370,6 +405,10 @@ fn build_record(
         caption_model: None,
         caption_updated_at: None,
         caption_error: None,
+        ai_rating: None,
+        ai_tagger_model: None,
+        ai_tagged_at: None,
+        ai_tagger_error: None,
     })
 }
 
@@ -825,6 +864,134 @@ fn process_caption_batch(
     Ok(())
 }
 
+const TAGGING_BATCH_SIZE: usize = 1;
+
+fn process_tagging_batch(
+    app: &AppHandle,
+    pool: &DbPool,
+    app_data_dir: &Path,
+    tagger_instance: &mut Option<WdTagger>,
+) -> Result<()> {
+    if !tagger::tagger_model_status(app_data_dir).ready {
+        return Ok(());
+    }
+
+    let paused_folders = paused_folder_ids("tagging");
+    let jobs = with_db_write_lock(|| {
+        let mut conn = pool.get()?;
+        db::claim_tagging_jobs(&mut conn, &paused_folders, TAGGING_BATCH_SIZE)
+    })?;
+
+    if jobs.is_empty() {
+        return Ok(());
+    }
+
+    if tagger_instance.is_none() {
+        match WdTagger::new(app_data_dir) {
+            Ok(model) => *tagger_instance = Some(model),
+            Err(error) => {
+                with_db_write_lock(|| {
+                    let conn = pool.get()?;
+                    db::requeue_tagging_jobs(
+                        &conn,
+                        &jobs.iter().map(|job| job.image_id).collect::<Vec<_>>(),
+                    )
+                })?;
+                return Err(error);
+            }
+        }
+    }
+
+    let folder_ids = jobs.iter().map(|job| job.folder_id).collect::<HashSet<_>>();
+    emit_folder_job_progress(
+        app,
+        pool,
+        &folder_ids.iter().copied().collect::<Vec<_>>(),
+        false,
+    );
+
+    let tagger_ref = tagger_instance
+        .as_mut()
+        .expect("tagger should be initialized before tagging batch processing");
+
+    let tag_results = jobs
+        .iter()
+        .map(|job| {
+            (
+                job.clone(),
+                tagger_ref.run(Path::new(&job.path), tagger::DEFAULT_MAX_TAGS),
+            )
+        })
+        .collect::<Vec<_>>();
+
+    let updated_images = with_db_write_lock(|| {
+        let mut conn = pool.get()?;
+        let tx = conn.transaction()?;
+        let mut updated_images = Vec::with_capacity(tag_results.len());
+
+        for (job, tag_result) in &tag_results {
+            // If the job was cancelled while inference was running, discard
+            // the result and delete the row — don't save tags or mark failed.
+            if db::is_tagging_job_cancelled(&tx, job.image_id)? {
+                tx.execute(
+                    "DELETE FROM tagging_jobs WHERE image_id = ?1",
+                    [job.image_id],
+                )?;
+                continue;
+            }
+            match tag_result {
+                Ok(output) => {
+                    let tag_pairs: Vec<(String, f64)> = output
+                        .tags
+                        .iter()
+                        .map(|t| (t.tag.clone(), t.confidence as f64))
+                        .collect();
+                    db::update_ai_tags(
+                        &tx,
+                        job.image_id,
+                        &tag_pairs,
+                        &output.rating,
+                        tagger::WD_TAGGER_MODEL_NAME,
+                    )?;
+                }
+                Err(error) => {
+                    db::mark_tagging_failed(&tx, job.image_id, &error.to_string())?;
+                }
+            }
+            updated_images.push(db::get_image_by_id(&tx, job.image_id)?);
+        }
+
+        tx.commit()?;
+        Ok(updated_images)
+    })
+    .or_else(|db_err| {
+        // The DB write failed.  Try to requeue the claimed jobs so they aren't
+        // left stuck in 'processing' until the next app restart.
+        let image_ids: Vec<i64> = jobs.iter().map(|job| job.image_id).collect();
+        let _ = with_db_write_lock(|| {
+            let conn = pool.get()?;
+            db::requeue_tagging_jobs(&conn, &image_ids)
+        });
+        Err(db_err)
+    })?;
+
+    if !updated_images.is_empty() {
+        let folder_ids = updated_images
+            .iter()
+            .map(|image| image.folder_id)
+            .collect::<HashSet<_>>();
+        emit_media_updates(
+            app,
+            &MediaUpdateBatch {
+                images: updated_images,
+            },
+        );
+        emit_folder_job_progress(app, pool, &folder_ids.into_iter().collect::<Vec<_>>(), true);
+    }
+
+    Ok(())
+}
+
 fn active_indexing_folders() -> HashSet<i64> {
     ACTIVE_INDEXING_FOLDERS
         .get_or_init(|| Mutex::new(HashSet::new()))
@@ -910,7 +1077,7 @@ fn emit_media_updates(app: &AppHandle, batch: &MediaUpdateBatch) {
     let _ = app.emit("media-updated", batch);
 }
 
-fn emit_folder_job_progress(app: &AppHandle, pool: &DbPool, folder_ids: &[i64], force: bool) {
+pub fn emit_folder_job_progress(app: &AppHandle, pool: &DbPool, folder_ids: &[i64], force: bool) {
     let mut unique_folder_ids = folder_ids.iter().copied().collect::<Vec<_>>();
     unique_folder_ids.sort_unstable();
     unique_folder_ids.dedup();

src-tauri/src/tagger.rs

@@ -0,0 +1,625 @@
+use anyhow::Result;
+use hf_hub::{api::sync::Api, Repo, RepoType};
+use image::{imageops::FilterType, DynamicImage, ImageReader};
+use ort::ep;
+use ort::session::{builder::GraphOptimizationLevel, Session};
+use ort::value::Tensor;
+use serde::{Deserialize, Serialize};
+use std::path::{Path, PathBuf};
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::time::Instant;
+
+pub const WD_TAGGER_MODEL_ID: &str = "SmilingWolf/wd-swinv2-tagger-v3";
+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";
+
+// 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] = &[
+    "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.
+// Category 0 = general, category 4 = character.
+// Category 9 = rating (explicit/questionable/sensitive/general) – used for
+// `ai_rating` but NOT emitted as individual tags.
+const GENERAL_CATEGORY: u8 = 0;
+const CHARACTER_CATEGORY: u8 = 4;
+const RATING_CATEGORY: u8 = 9;
+
+pub const DEFAULT_THRESHOLD: f32 = 0.35;
+pub const DEFAULT_MAX_TAGS: usize = 30;
+
+/// Set to `true` by `set_tagger_acceleration` so the tagging worker loop
+/// knows to drop its cached `WdTagger` and reload with the new EP.
+pub static TAGGER_SESSION_DIRTY: AtomicBool = AtomicBool::new(false);
+
+// ---------------------------------------------------------------------------
+// Settings types
+// ---------------------------------------------------------------------------
+
+#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
+#[serde(rename_all = "lowercase")]
+pub enum TaggerAcceleration {
+    Auto,
+    Cpu,
+    Directml,
+}
+
+impl TaggerAcceleration {
+    fn as_str(self) -> &'static str {
+        match self {
+            Self::Auto => "auto",
+            Self::Cpu => "cpu",
+            Self::Directml => "directml",
+        }
+    }
+}
+
+impl Default for TaggerAcceleration {
+    fn default() -> Self {
+        Self::Auto
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Status / probe types exposed to the frontend
+// ---------------------------------------------------------------------------
+
+#[derive(Serialize)]
+pub struct TaggerModelStatus {
+    pub model_id: &'static str,
+    pub model_name: &'static str,
+    pub local_dir: String,
+    pub ready: bool,
+    pub missing_files: Vec<String>,
+}
+
+#[derive(Clone, Serialize)]
+pub struct TaggerModelProgress {
+    pub total_files: usize,
+    pub completed_files: usize,
+    pub current_file: Option<String>,
+    pub done: bool,
+}
+
+// ---------------------------------------------------------------------------
+// Runtime probe types exposed to the frontend
+// ---------------------------------------------------------------------------
+
+#[derive(Serialize)]
+pub struct TaggerRuntimeProbe {
+    pub ready: bool,
+    pub acceleration: TaggerAcceleration,
+    pub session: TaggerRuntimeSessionProbe,
+}
+
+#[derive(Serialize)]
+pub struct TaggerRuntimeSessionProbe {
+    pub file: &'static str,
+    pub inputs: Vec<String>,
+    pub outputs: Vec<String>,
+}
+
+// ---------------------------------------------------------------------------
+// Tag record returned to callers
+// ---------------------------------------------------------------------------
+
+#[derive(Debug, Clone, Serialize)]
+pub struct TagResult {
+    pub tag: String,
+    pub confidence: f32,
+}
+
+#[derive(Debug, Clone, Serialize)]
+pub struct TaggerOutput {
+    pub tags: Vec<TagResult>,
+    /// Highest-scoring rating label: "general" | "sensitive" | "questionable" | "explicit"
+    pub rating: String,
+}
+
+// ---------------------------------------------------------------------------
+// Internal label table built from selected_tags.csv
+// ---------------------------------------------------------------------------
+
+#[derive(Debug, Clone)]
+struct TagEntry {
+    name: String,
+    category: u8,
+}
+
+// ---------------------------------------------------------------------------
+// Path helpers
+// ---------------------------------------------------------------------------
+
+pub fn model_dir(app_data_dir: &Path) -> PathBuf {
+    app_data_dir.join("models").join("wd-swinv2-tagger-v3")
+}
+
+// ---------------------------------------------------------------------------
+// Settings persistence
+// ---------------------------------------------------------------------------
+
+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 {
+        return TaggerAcceleration::default();
+    };
+    match value.trim().to_ascii_lowercase().as_str() {
+        "cpu" => TaggerAcceleration::Cpu,
+        "directml" => TaggerAcceleration::Directml,
+        _ => TaggerAcceleration::Auto,
+    }
+}
+
+pub fn set_tagger_acceleration(
+    app_data_dir: &Path,
+    acceleration: TaggerAcceleration,
+) -> Result<TaggerAcceleration> {
+    let path = app_data_dir.join(TAGGER_ACCELERATION_FILE);
+    if let Some(parent) = path.parent() {
+        std::fs::create_dir_all(parent)?;
+    }
+    std::fs::write(path, acceleration.as_str())?;
+    TAGGER_SESSION_DIRTY.store(true, Ordering::Relaxed);
+    Ok(acceleration)
+}
+
+pub fn tagger_threshold(app_data_dir: &Path) -> f32 {
+    let path = app_data_dir.join(TAGGER_THRESHOLD_FILE);
+    let Ok(value) = std::fs::read_to_string(path) else {
+        return DEFAULT_THRESHOLD;
+    };
+    value
+        .trim()
+        .parse::<f32>()
+        .unwrap_or(DEFAULT_THRESHOLD)
+        .clamp(0.01, 1.0)
+}
+
+pub fn set_tagger_threshold(app_data_dir: &Path, threshold: f32) -> Result<f32> {
+    let clamped = threshold.clamp(0.01, 1.0);
+    let path = app_data_dir.join(TAGGER_THRESHOLD_FILE);
+    if let Some(parent) = path.parent() {
+        std::fs::create_dir_all(parent)?;
+    }
+    std::fs::write(path, clamped.to_string())?;
+    Ok(clamped)
+}
+
+// ---------------------------------------------------------------------------
+// Model status / download
+// ---------------------------------------------------------------------------
+
+pub fn tagger_model_status(app_data_dir: &Path) -> TaggerModelStatus {
+    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
+        .iter()
+        .filter(|file| {
+            let path = if file.starts_with("onnxruntime/") {
+                caption_model_dir.join(file)
+            } else {
+                local_dir.join(file)
+            };
+            !path.exists()
+        })
+        .map(|file| (*file).to_string())
+        .collect::<Vec<_>>();
+
+    TaggerModelStatus {
+        model_id: WD_TAGGER_MODEL_ID,
+        model_name: WD_TAGGER_MODEL_NAME,
+        local_dir: local_dir.to_string_lossy().to_string(),
+        ready: missing_files.is_empty(),
+        missing_files,
+    }
+}
+
+pub fn prepare_tagger_model_with_progress(
+    app_data_dir: &Path,
+    emit_progress: impl Fn(TaggerModelProgress),
+) -> Result<TaggerModelStatus> {
+    let local_dir = model_dir(app_data_dir);
+    std::fs::create_dir_all(&local_dir)?;
+
+    // Only download the two tagger-specific files; ONNX runtime DLLs are
+    // already handled by the captioner download flow.
+    const DOWNLOAD_FILES: &[&str] = &["model.onnx", "selected_tags.csv"];
+
+    let api = Api::new()?;
+    let repo = api.repo(Repo::new(WD_TAGGER_MODEL_ID.to_string(), RepoType::Model));
+
+    let mut completed_files = DOWNLOAD_FILES
+        .iter()
+        .filter(|file| local_dir.join(file).exists())
+        .count();
+
+    emit_progress(TaggerModelProgress {
+        total_files: DOWNLOAD_FILES.len(),
+        completed_files,
+        current_file: None,
+        done: completed_files == DOWNLOAD_FILES.len(),
+    });
+
+    for file in DOWNLOAD_FILES {
+        let destination = local_dir.join(file);
+        if destination.exists() {
+            continue;
+        }
+        emit_progress(TaggerModelProgress {
+            total_files: DOWNLOAD_FILES.len(),
+            completed_files,
+            current_file: Some((*file).to_string()),
+            done: false,
+        });
+        let cached = repo.get(file)?;
+        std::fs::copy(cached, destination)?;
+        completed_files += 1;
+        emit_progress(TaggerModelProgress {
+            total_files: DOWNLOAD_FILES.len(),
+            completed_files,
+            current_file: Some((*file).to_string()),
+            done: completed_files == DOWNLOAD_FILES.len(),
+        });
+    }
+
+    emit_progress(TaggerModelProgress {
+        total_files: DOWNLOAD_FILES.len(),
+        completed_files,
+        current_file: None,
+        done: true,
+    });
+
+    Ok(tagger_model_status(app_data_dir))
+}
+
+pub fn delete_tagger_model(app_data_dir: &Path) -> Result<TaggerModelStatus> {
+    let local_dir = model_dir(app_data_dir);
+    if local_dir.exists() {
+        std::fs::remove_dir_all(&local_dir)?;
+    }
+    Ok(tagger_model_status(app_data_dir))
+}
+
+pub fn probe_tagger_runtime(app_data_dir: &Path) -> Result<TaggerRuntimeProbe> {
+    let status = tagger_model_status(app_data_dir);
+    if !status.ready {
+        anyhow::bail!(
+            "WD Tagger model is missing {} required file(s): {}",
+            status.missing_files.len(),
+            status.missing_files.join(", ")
+        );
+    }
+
+    let local_dir = model_dir(app_data_dir);
+    let caption_model_dir = app_data_dir.join("models").join("florence-2-base-ft");
+    crate::captioner::ensure_onnx_runtime(&caption_model_dir)?;
+
+    let acceleration = tagger_acceleration(app_data_dir);
+    let model_path = local_dir.join("model.onnx");
+
+    // Verify that the model file exists and has non-zero size before trying
+    // to create an ORT session (better error message on corruption).
+    let metadata = std::fs::metadata(&model_path)?;
+    if metadata.len() == 0 {
+        anyhow::bail!("model.onnx is empty");
+    }
+
+    // Actually create a session to verify the EP loads correctly.
+    let loaded_acceleration = match acceleration {
+        TaggerAcceleration::Cpu => {
+            create_tagger_session(&model_path, TaggerAcceleration::Cpu)?;
+            TaggerAcceleration::Cpu
+        }
+        TaggerAcceleration::Auto => {
+            // Try DirectML explicitly; if it fails the real session would have
+            // fallen back to CPU silently — report what would actually run.
+            let directml_ok =
+                create_tagger_session(&model_path, TaggerAcceleration::Directml).is_ok();
+            if directml_ok {
+                TaggerAcceleration::Directml
+            } else {
+                create_tagger_session(&model_path, TaggerAcceleration::Cpu)?;
+                TaggerAcceleration::Cpu
+            }
+        }
+        TaggerAcceleration::Directml => {
+            create_tagger_session(&model_path, TaggerAcceleration::Directml)?;
+            TaggerAcceleration::Directml
+        }
+    };
+
+    Ok(TaggerRuntimeProbe {
+        ready: true,
+        acceleration: loaded_acceleration,
+        session: TaggerRuntimeSessionProbe {
+            file: "model.onnx",
+            inputs: vec!["pixel_values".to_string()],
+            outputs: vec![format!("output [EP: {:?}]", loaded_acceleration)],
+        },
+    })
+}
+
+// ---------------------------------------------------------------------------
+// Top-level inference entry point
+// ---------------------------------------------------------------------------
+
+// ---------------------------------------------------------------------------
+// Tagger implementation
+// ---------------------------------------------------------------------------
+
+pub struct WdTagger {
+    session: Session,
+    labels: Vec<TagEntry>,
+    threshold: f32,
+    input_size: usize,
+}
+
+impl WdTagger {
+    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!(
+                "WD tagger model is missing {} required file(s): {}",
+                status.missing_files.len(),
+                status.missing_files.join(", ")
+            );
+        }
+
+        let local_dir = model_dir(app_data_dir);
+
+        // The ONNX runtime DLLs are shared with the captioner; use the
+        // captioner's shared ORT init lock to avoid double-initialisation.
+        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 = tagger_threshold(app_data_dir);
+        let model_path = local_dir.join("model.onnx");
+        let labels_path = local_dir.join("selected_tags.csv");
+
+        let session = create_tagger_session(&model_path, acceleration)?;
+
+        // Determine the input spatial size from the ONNX model graph.
+        // WD v3 models use (1, H, W, 3) where H == W, typically 448.
+        let input_size = {
+            let inputs = session.inputs();
+            let dim = inputs
+                .first()
+                .and_then(|inp| {
+                    if let ort::value::ValueType::Tensor { shape, .. } = inp.dtype() {
+                        shape
+                            .get(1)
+                            .and_then(|&d| if d > 0 { Some(d as usize) } else { None })
+                    } else {
+                        None
+                    }
+                })
+                .unwrap_or(448);
+            dim
+        };
+
+        let labels = load_labels(&labels_path)?;
+
+        println!(
+            "WD tagger loaded in {:?} ({} labels, input {}x{}, {:?} acceleration)",
+            started_at.elapsed(),
+            labels.len(),
+            input_size,
+            input_size,
+            acceleration,
+        );
+
+        Ok(Self {
+            session,
+            labels,
+            threshold,
+            input_size,
+        })
+    }
+
+    pub fn run(&mut self, image_path: &Path, max_tags: usize) -> Result<TaggerOutput> {
+        let started_at = Instant::now();
+
+        let image_array = preprocess_image(image_path, self.input_size)?;
+        let batch_size = 1usize;
+        let input = Tensor::from_array((
+            [batch_size, self.input_size, self.input_size, 3usize],
+            image_array.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 (_, probabilities) = outputs[0]
+            .try_extract_tensor::<f32>()
+            .map_err(|error| anyhow::anyhow!("{error}"))?;
+
+        let probs: &[f32] = &probabilities;
+
+        if probs.len() != self.labels.len() {
+            anyhow::bail!(
+                "Model output length {} does not match label count {}",
+                probs.len(),
+                self.labels.len()
+            );
+        }
+
+        // Collect rating scores (category 9) - pick the argmax as the rating.
+        let rating = self
+            .labels
+            .iter()
+            .zip(probs.iter())
+            .filter(|(entry, _)| entry.category == RATING_CATEGORY)
+            .max_by(|(_, a), (_, b)| a.total_cmp(b))
+            .map(|(entry, _)| entry.name.clone())
+            .unwrap_or_else(|| "general".to_string());
+
+        // Collect general + character tags above threshold, sorted by confidence.
+        let mut tags: Vec<TagResult> = self
+            .labels
+            .iter()
+            .zip(probs.iter())
+            .filter(|(entry, prob)| {
+                (entry.category == GENERAL_CATEGORY || entry.category == CHARACTER_CATEGORY)
+                    && **prob >= self.threshold
+            })
+            .map(|(entry, prob)| TagResult {
+                tag: entry.name.clone(),
+                confidence: *prob,
+            })
+            .collect();
+
+        tags.sort_by(|a, b| b.confidence.total_cmp(&a.confidence));
+        tags.truncate(max_tags);
+
+        println!(
+            "WD tagger: {} tags (threshold={}, rating={}) in {:?} for {}",
+            tags.len(),
+            self.threshold,
+            rating,
+            started_at.elapsed(),
+            image_path.display(),
+        );
+
+        Ok(TaggerOutput { tags, rating })
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Session creation – mirrors captioner's `create_session`
+// ---------------------------------------------------------------------------
+
+fn create_tagger_session(path: &Path, acceleration: TaggerAcceleration) -> Result<Session> {
+    let builder = Session::builder().map_err(|error| anyhow::anyhow!("{error}"))?;
+    let builder = builder
+        .with_optimization_level(GraphOptimizationLevel::Level3)
+        .map_err(|error| anyhow::anyhow!("{error}"))?;
+
+    let use_directml = matches!(
+        acceleration,
+        TaggerAcceleration::Auto | TaggerAcceleration::Directml
+    );
+
+    let builder = builder
+        .with_memory_pattern(!use_directml)
+        .map_err(|error| anyhow::anyhow!("{error}"))?;
+    let builder = builder
+        .with_parallel_execution(false)
+        .map_err(|error| anyhow::anyhow!("{error}"))?;
+    let builder = builder
+        .with_intra_threads(1)
+        .map_err(|error| anyhow::anyhow!("{error}"))?;
+
+    let mut builder = match acceleration {
+        TaggerAcceleration::Cpu => {
+            println!("WD tagger: using CPU execution provider");
+            builder
+        }
+        TaggerAcceleration::Auto => builder
+            .with_execution_providers([ep::DirectML::default().build().fail_silently()])
+            .unwrap_or_else(|error| {
+                println!("WD tagger: DirectML unavailable, falling back to CPU");
+                error.recover()
+            }),
+        TaggerAcceleration::Directml => builder
+            .with_execution_providers([ep::DirectML::default().build().error_on_failure()])
+            .map_err(|error| anyhow::anyhow!("{error}"))?,
+    };
+
+    let session = builder
+        .commit_from_file(path)
+        .map_err(|error| anyhow::anyhow!("{error}"))?;
+
+    Ok(session)
+}
+// ---------------------------------------------------------------------------
+// Label loading
+// ---------------------------------------------------------------------------
+
+fn load_labels(path: &Path) -> Result<Vec<TagEntry>> {
+    let mut reader = csv::Reader::from_path(path)?;
+    let mut entries = Vec::new();
+    for result in reader.records() {
+        let record = result?;
+        // CSV columns: tag_id, name, category, count
+        let name = record
+            .get(1)
+            .ok_or_else(|| anyhow::anyhow!("Missing name column in selected_tags.csv"))?
+            .replace('_', " ");
+        let category: u8 = record
+            .get(2)
+            .ok_or_else(|| anyhow::anyhow!("Missing category column in selected_tags.csv"))?
+            .parse()
+            .unwrap_or(0);
+        entries.push(TagEntry { name, category });
+    }
+    if entries.is_empty() {
+        anyhow::bail!("selected_tags.csv is empty or could not be parsed");
+    }
+    Ok(entries)
+}
+
+// ---------------------------------------------------------------------------
+// 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).
+// ---------------------------------------------------------------------------
+
+fn preprocess_image(image_path: &Path, target_size: usize) -> Result<Vec<f32>> {
+    let image = ImageReader::open(image_path)?.decode()?;
+
+    // Composite any alpha channel onto a white background.
+    let image_rgba = image.to_rgba8();
+    let (width, height) = image_rgba.dimensions();
+    let mut canvas_rgba =
+        image::RgbaImage::from_pixel(width, height, image::Rgba([255, 255, 255, 255]));
+    image::imageops::overlay(&mut canvas_rgba, &image_rgba, 0, 0);
+    let image_rgb = DynamicImage::ImageRgba8(canvas_rgba).to_rgb8();
+
+    // Pad to 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(
+        &square,
+        target_size as u32,
+        target_size as u32,
+        FilterType::CatmullRom,
+    );
+
+    // Flatten to (H, W, 3) float32 in BGR order, values in [0, 255].
+    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;
+        // BGR order
+        pixel_values[base] = f32::from(pixel[2]); // B
+        pixel_values[base + 1] = f32::from(pixel[1]); // G
+        pixel_values[base + 2] = f32::from(pixel[0]); // R
+    }
+
+    Ok(pixel_values)
+}

src-tauri/src/vector.rs

@@ -80,7 +80,12 @@ pub fn upsert_caption_embedding(conn: &Connection, image_id: i64, embedding: &[f
     Ok(())
 }
 
-pub fn find_similar_image_ids(conn: &Connection, image_id: i64, limit: usize) -> Result<Vec<i64>> {
+pub fn find_similar_image_ids(
+    conn: &Connection,
+    image_id: i64,
+    limit: usize,
+    folder_id: Option<i64>,
+) -> Result<Vec<i64>> {
     let embedding: Vec<u8> = match conn.query_row(
         "SELECT embedding FROM image_vec WHERE image_id = ?1",
         [image_id],
@@ -91,19 +96,37 @@ pub fn find_similar_image_ids(conn: &Connection, image_id: i64, limit: usize) ->
         Err(error) => return Err(error.into()),
     };
 
+    let allowed_folder_ids = match folder_id {
+        Some(folder_id) => Some(image_ids_for_folder(conn, folder_id)?),
+        None => None,
+    };
+    let search_limit = if allowed_folder_ids.is_some() {
+        count_image_vectors(conn)?.max(1) as usize
+    } else {
+        limit + 1
+    };
     let mut stmt = conn.prepare(
         "SELECT image_id
          FROM image_vec
          WHERE embedding MATCH vec_f32(?1)
            AND k = ?2",
     )?;
-    let rows = stmt.query_map((&embedding, (limit as i64) + 1), |row| row.get::<_, i64>(0))?;
+    let rows = stmt
+        .query_map((&embedding, search_limit as i64), |row| {
+            row.get::<_, i64>(0)
+        })?
+        .collect::<rusqlite::Result<Vec<_>>>()?;
 
     let mut ids = Vec::new();
     for row in rows {
-        let candidate_id = row?;
-        if candidate_id != image_id {
-            ids.push(candidate_id);
+        if row != image_id {
+            if allowed_folder_ids
+                .as_ref()
+                .is_some_and(|folder_ids| !folder_ids.contains(&row))
+            {
+                continue;
+            }
+            ids.push(row);
         }
         if ids.len() >= limit {
             break;
@@ -112,6 +135,15 @@ pub fn find_similar_image_ids(conn: &Connection, image_id: i64, limit: usize) ->
     Ok(ids)
 }
 
+fn image_ids_for_folder(
+    conn: &Connection,
+    folder_id: i64,
+) -> Result<std::collections::HashSet<i64>> {
+    let mut stmt = conn.prepare("SELECT id FROM images WHERE folder_id = ?1")?;
+    let rows = stmt.query_map([folder_id], |row| row.get::<_, i64>(0))?;
+    Ok(rows.collect::<rusqlite::Result<std::collections::HashSet<_>>>()?)
+}
+
 /// Returns all stored image embeddings with their image IDs, optionally filtered to one folder.
 /// Each entry is `(image_id, normalized_f32_embedding)`.
 pub fn get_all_image_embeddings_with_ids(