From b23212ea1cad3d2adef745c9f7e44b6604bbdea0 Mon Sep 17 00:00:00 2001 From: LyAhn Date: Sun, 5 Jul 2026 20:21:47 +0100 Subject: [PATCH] refactor(backend): extract download and onnx_runtime modules from captioner captioner.rs had grown into two things: the (currently disabled) Florence-2 captioner plus generic infrastructure the live tagger depends on. Split the neutral parts out: - download.rs: resilient curl downloader (resume, stall detection) and NuGet package extraction - onnx_runtime.rs: shared ONNX Runtime/DirectML DLL manifest, provisioning, and ort init, with runtime_dir() as the single definition of the DLL location (kept inside the caption model dir so existing installs do not re-download) Also dedupes the tagger-side copies of the DLL list and four hardcoded caption-model path literals, and rewords tagger runtime errors that wrongly told users to download the caption model. --- CLAUDE.md | 2 + src-tauri/src/captioner.rs | 347 +--------------------------------- src-tauri/src/download.rs | 238 +++++++++++++++++++++++ src-tauri/src/lib.rs | 2 + src-tauri/src/onnx_runtime.rs | 137 ++++++++++++++ src-tauri/src/tagger.rs | 56 ++---- 6 files changed, 408 insertions(+), 374 deletions(-) create mode 100644 src-tauri/src/download.rs create mode 100644 src-tauri/src/onnx_runtime.rs diff --git a/CLAUDE.md b/CLAUDE.md index dcd4f8c..763c895 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -68,6 +68,8 @@ Key modules: | `hnsw_index.rs` | In-memory HNSW index wrapper (hnsw_rs) | | `tagger.rs` | WD tagger: ONNX model download, inference, CSV tag loading | | `captioner.rs` | AI captioning (ONNX, disabled in workers but code intact) | +| `download.rs` | Resilient file downloads via the system `curl` (resume, stall detection) | +| `onnx_runtime.rs` | Shared ONNX Runtime DLL provisioning + `ort` init (used by tagger and captioner; DLLs live in the caption model dir for legacy reasons) | | `thumbnail.rs` | Thumbnail generation (image crate + fast_image_resize, FFmpeg for video) | | `media.rs` | FFmpeg sidecar provisioning and probing | | `storage.rs` | `StorageProfile` for tuning worker counts | diff --git a/src-tauri/src/captioner.rs b/src-tauri/src/captioner.rs index 91ae2bf..ea47e1b 100644 --- a/src-tauri/src/captioner.rs +++ b/src-tauri/src/captioner.rs @@ -1,3 +1,6 @@ +use crate::onnx_runtime::{ + self, DIRECTML_DLL_FILE, ONNX_RUNTIME_DLL_FILE, ONNX_RUNTIME_PROVIDERS_DLL_FILE, +}; use anyhow::Result; use hf_hub::{api::sync::Api, Repo, RepoType}; use image::{imageops::FilterType, ImageReader}; @@ -8,67 +11,16 @@ use ort::session::{builder::GraphOptimizationLevel, Session}; use ort::value::{Shape, Tensor}; use serde::{Deserialize, Serialize}; use std::borrow::Cow; -use std::io::Read; use std::path::{Path, PathBuf}; -use std::process::Command; use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::Mutex; use std::time::Instant; use tokenizers::Tokenizer; -// Suppress the console window when spawning curl from the GUI app. -#[cfg(target_os = "windows")] -use std::os::windows::process::CommandExt; -#[cfg(target_os = "windows")] -const CREATE_NO_WINDOW: u32 = 0x08000000; - -/// Discard target for curl output we only need the headers of. -#[cfg(target_os = "windows")] -const NULL_DEVICE: &str = "NUL"; -#[cfg(not(target_os = "windows"))] -const NULL_DEVICE: &str = "/dev/null"; - -/// Build a `curl` command with platform quirks applied. Windows resolves the -/// bare name to `curl.exe` (bundled since Windows 10 1803) and needs the -/// no-window flag; macOS ships curl; Linux is expected to have it installed. -fn curl_command() -> Command { - #[allow(unused_mut)] - let mut command = Command::new("curl"); - #[cfg(target_os = "windows")] - command.creation_flags(CREATE_NO_WINDOW); - command -} - 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, @@ -86,11 +38,6 @@ const REQUIRED_FILES: &[&str] = &[ "onnx/embed_tokens_fp16.onnx", ]; -// Mutex rather than OnceLock: a failed attempt (DLL not yet -// downloaded) must NOT be cached, or a later successful download could never -// recover within the same app session. -static ORT_RUNTIME_INIT: Mutex = Mutex::new(false); - /// 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); @@ -311,11 +258,11 @@ pub fn prepare_caption_model_with_progress( if let Some(parent) = destination.parent() { std::fs::create_dir_all(parent)?; } - if ONNX_RUNTIME_FILES + if onnx_runtime::ONNX_RUNTIME_FILES .iter() .any(|(runtime_file, _, _)| runtime_file == file) { - download_onnx_runtime_files(&local_dir)?; + onnx_runtime::download_onnx_runtime_files(&local_dir)?; completed_files = REQUIRED_FILES .iter() .filter(|file| local_dir.join(file).exists()) @@ -361,7 +308,7 @@ pub fn probe_caption_runtime(app_data_dir: &Path) -> Result } let local_dir = model_dir(app_data_dir); - ensure_onnx_runtime(&local_dir)?; + onnx_runtime::ensure_onnx_runtime(&local_dir)?; let tokenizer = Tokenizer::from_file(local_dir.join("tokenizer.json")).map_err(anyhow::Error::msg)?; @@ -410,7 +357,7 @@ pub fn probe_caption_vision(app_data_dir: &Path, image_path: &Path) -> Result Result<()> { - let mut initialized = ORT_RUNTIME_INIT - .lock() - .map_err(|_| anyhow::anyhow!("ONNX runtime init lock poisoned"))?; - if *initialized { - return Ok(()); - } - let dll_path = local_dir.join(ONNX_RUNTIME_DLL_FILE); - if !dll_path.exists() { - anyhow::bail!("ONNX Runtime DLL is missing: {}", dll_path.display()); - } - ort::environment::init_from(&dll_path) - .map_err(|error| anyhow::anyhow!(error.to_string()))? - .with_name("phokus-florence") - .commit(); - *initialized = true; - Ok(()) -} - -/// Download any ONNX Runtime DLLs missing from `local_dir`, reporting per-file -/// byte progress as `(short_label, downloaded_bytes, total_bytes)`. -/// `total_bytes` is `None` when the server omits Content-Length. Unlike -/// `ensure_onnx_runtime` (init only), this actually provisions the files — -/// callers that can run on a clean install must call this first. The callback -/// fires per chunk; callers should throttle. -pub fn provision_onnx_runtime_with_progress( - local_dir: &Path, - mut on_progress: impl FnMut(&str, u64, Option), -) -> Result<()> { - for (destination_file, source_url, archive_path) in ONNX_RUNTIME_FILES { - let destination = local_dir.join(destination_file); - if destination.exists() { - continue; - } - // Strip the "onnxruntime/" prefix for a clean label. - let label = destination_file - .rsplit('/') - .next() - .unwrap_or(destination_file); - download_nuget_file( - source_url, - archive_path, - &destination, - |downloaded, total| on_progress(label, downloaded, total), - )?; - } - Ok(()) -} - -/// Number of ONNX Runtime DLLs still missing from `local_dir` (for progress -/// step counts before downloading). -pub fn missing_onnx_runtime_count(local_dir: &Path) -> usize { - ONNX_RUNTIME_FILES - .iter() - .filter(|(destination_file, _, _)| !local_dir.join(destination_file).exists()) - .count() -} - -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(()) -} - -// Give up only after this many *consecutive* curl runs that download nothing; -// a run that makes any progress resets the counter, so a large file completes -// across however many resumes it takes. Kept low so a hard stall (e.g. a -// broken VM NIC) fails in a couple of minutes — surfacing a retryable error — -// rather than locking the UI on "preparing" for many minutes. -const MAX_STALL_RETRIES: usize = 3; - -/// Resiliently download `url` to `destination` using the system `curl`. -/// -/// ureq's read timeout does not fire on a stalled large transfer on Windows -/// (schannel doesn't honor the socket read timeout), so a stall there hangs -/// forever. curl detects an inactivity stall (`--speed-time`), resumes from -/// the partial file (`-C -`), and retries internally — the same behavior a -/// browser gets. We monitor the `.part` file's size for the progress bar and -/// wrap curl in an outer progress-aware retry as a backstop. The partial file -/// survives an app restart, so a later retry continues from disk. -pub fn download_file_resilient( - url: &str, - destination: &Path, - mut on_progress: impl FnMut(u64, Option), -) -> Result<()> { - if let Some(parent) = destination.parent() { - std::fs::create_dir_all(parent)?; - } - let part = match destination.extension() { - Some(ext) => destination.with_extension(format!("{}.part", ext.to_string_lossy())), - None => destination.with_extension("part"), - }; - let name = destination - .file_name() - .map(|n| n.to_string_lossy().into_owned()) - .unwrap_or_else(|| url.to_string()); - - log::info!("{name}: resolving download size"); - let total = remote_content_length(url); - - // Reconcile any existing `.part` against the real size: exactly complete → - // finish; oversized (stale/corrupt) → discard so curl restarts cleanly - // (otherwise `curl -C -` would 416 forever). - if let Some(total) = total { - let size = std::fs::metadata(&part).map(|m| m.len()).unwrap_or(0); - if size == total { - std::fs::rename(&part, destination)?; - return Ok(()); - } - if size > total { - let _ = std::fs::remove_file(&part); - } - } - log::info!( - "{name}: downloading via curl ({} bytes)", - total - .map(|t| t.to_string()) - .unwrap_or_else(|| "unknown size".into()) - ); - - let mut stalls = 0usize; - loop { - let before = std::fs::metadata(&part).map(|m| m.len()).unwrap_or(0); - match run_curl_download(url, &part, total, &mut on_progress) { - Ok(()) => break, - Err(error) => { - let after = std::fs::metadata(&part).map(|m| m.len()).unwrap_or(0); - if after > before { - log::warn!("{name}: curl interrupted at {after} bytes, resuming: {error}"); - stalls = 0; - } else { - stalls += 1; - log::warn!( - "{name}: curl made no progress ({stalls}/{MAX_STALL_RETRIES}): {error}" - ); - if stalls >= MAX_STALL_RETRIES { - // Discard the partial so a future attempt restarts clean - // rather than getting stuck re-resuming a bad file. - let _ = std::fs::remove_file(&part); - return Err(error); - } - } - std::thread::sleep(std::time::Duration::from_secs(2)); - } - } - } - - if let Some(total) = total { - let got = std::fs::metadata(&part).map(|m| m.len()).unwrap_or(0); - if got < total { - anyhow::bail!("{name}: incomplete after curl ({got}/{total} bytes)"); - } - } - std::fs::rename(&part, destination)?; - Ok(()) -} - -/// Size probe via `curl -r 0-0` (a 1-byte Range request), parsing the total -/// from the `Content-Range: bytes 0-0/` header. Uses curl rather than -/// ureq so no part of the download path depends on ureq (which hangs on this -/// VM's TLS stack). Returns None if the server doesn't report a size. -fn remote_content_length(url: &str) -> Option { - let mut command = curl_command(); - command.args([ - "-sL", - "-r", - "0-0", - "-D", - "-", - "-o", - NULL_DEVICE, - "--connect-timeout", - "30", - "--max-time", - "30", - url, - ]); - let output = command.output().ok()?; - let headers = String::from_utf8_lossy(&output.stdout); - for line in headers.lines() { - if let Some(rest) = line.to_ascii_lowercase().strip_prefix("content-range:") { - if let Some(total) = rest.rsplit('/').next().map(str::trim) { - if let Ok(n) = total.parse::() { - return Some(n); - } - } - } - } - None -} - -/// Run one `curl` download to `dest`, resuming from any partial file, while -/// reporting progress from the growing file size. Returns an error (leaving the -/// partial in place) if curl exits non-zero. -fn run_curl_download( - url: &str, - dest: &Path, - total: Option, - on_progress: &mut impl FnMut(u64, Option), -) -> Result<()> { - let mut command = curl_command(); - command - .arg("-fSL") // fail on HTTP errors, follow redirects, show errors - .args(["-C", "-"]) // resume from the existing output file - .args(["--retry", "3", "--retry-delay", "1", "--retry-connrefused"]) - .args(["--connect-timeout", "30"]) - // Abort (then --retry resumes) if under 1 KB/s for 30s — a real - // inactivity timeout, which is what ureq couldn't deliver here. - .args(["--speed-limit", "1024", "--speed-time", "30"]) - .arg("-s") // no progress meter (we watch the file instead) - .arg("-o") - .arg(dest) - .arg(url) - .stdout(std::process::Stdio::null()) - .stderr(std::process::Stdio::piped()); - - let mut child = command - .spawn() - .map_err(|e| anyhow::anyhow!("failed to launch curl (required for downloads): {e}"))?; - - loop { - if let Some(status) = child.try_wait()? { - if status.success() { - return Ok(()); - } - let mut stderr = String::new(); - if let Some(mut pipe) = child.stderr.take() { - let _ = pipe.read_to_string(&mut stderr); - } - anyhow::bail!( - "curl exited with {}: {}", - status - .code() - .map(|c| c.to_string()) - .unwrap_or_else(|| "signal".into()), - stderr.trim() - ); - } - let downloaded = std::fs::metadata(dest).map(|m| m.len()).unwrap_or(0); - on_progress(downloaded, total); - std::thread::sleep(std::time::Duration::from_millis(300)); - } -} - -fn download_nuget_file( - source_url: &str, - archive_path: &str, - destination: &Path, - on_progress: impl FnMut(u64, Option), -) -> Result<()> { - // Download the .nupkg (a zip) resiliently, then extract the one DLL. - let package = destination.with_extension("nupkg"); - download_file_resilient(source_url, &package, on_progress)?; - - log::info!("extracting {archive_path} from package"); - let file = std::fs::File::open(&package)?; - let mut archive = zip::ZipArchive::new(file)?; - let mut dll = archive.by_name(archive_path)?; - let temp_destination = destination.with_extension("tmp"); - { - let mut out = std::fs::File::create(&temp_destination)?; - std::io::copy(&mut dll, &mut out)?; - } - std::fs::rename(&temp_destination, destination)?; - let _ = std::fs::remove_file(&package); - log::info!("extracted {archive_path}"); - Ok(()) -} - fn run_vision_encoder(session: &mut Session, image_path: &Path) -> Result { let pixels = preprocess_image(image_path)?; let input = Tensor::from_array(([1usize, 3, 768, 768], pixels.into_boxed_slice())) diff --git a/src-tauri/src/download.rs b/src-tauri/src/download.rs new file mode 100644 index 0000000..71d2c32 --- /dev/null +++ b/src-tauri/src/download.rs @@ -0,0 +1,238 @@ +//! Resilient file downloads via the system `curl` binary. +//! +//! Used for all large model/runtime downloads (tagger models, ONNX Runtime +//! DLLs, caption model). ureq's read timeout does not fire on a stalled large +//! transfer on Windows (schannel doesn't honor the socket read timeout), so a +//! stall there hangs forever. curl detects an inactivity stall +//! (`--speed-time`), resumes from the partial file (`-C -`), and retries +//! internally — the same behavior a browser gets. + +use anyhow::Result; +use std::io::Read; +use std::path::Path; +use std::process::Command; + +// Suppress the console window when spawning curl from the GUI app. +#[cfg(target_os = "windows")] +use std::os::windows::process::CommandExt; +#[cfg(target_os = "windows")] +const CREATE_NO_WINDOW: u32 = 0x08000000; + +/// Discard target for curl output we only need the headers of. +#[cfg(target_os = "windows")] +const NULL_DEVICE: &str = "NUL"; +#[cfg(not(target_os = "windows"))] +const NULL_DEVICE: &str = "/dev/null"; + +/// Build a `curl` command with platform quirks applied. Windows resolves the +/// bare name to `curl.exe` (bundled since Windows 10 1803) and needs the +/// no-window flag; macOS ships curl; Linux is expected to have it installed. +fn curl_command() -> Command { + #[allow(unused_mut)] + let mut command = Command::new("curl"); + #[cfg(target_os = "windows")] + command.creation_flags(CREATE_NO_WINDOW); + command +} + +// Give up only after this many *consecutive* curl runs that download nothing; +// a run that makes any progress resets the counter, so a large file completes +// across however many resumes it takes. Kept low so a hard stall (e.g. a +// broken VM NIC) fails in a couple of minutes — surfacing a retryable error — +// rather than locking the UI on "preparing" for many minutes. +const MAX_STALL_RETRIES: usize = 3; + +/// Resiliently download `url` to `destination` using the system `curl`. +/// +/// We monitor the `.part` file's size for the progress bar and wrap curl in an +/// outer progress-aware retry as a backstop. The partial file survives an app +/// restart, so a later retry continues from disk. +pub fn download_file_resilient( + url: &str, + destination: &Path, + mut on_progress: impl FnMut(u64, Option), +) -> Result<()> { + if let Some(parent) = destination.parent() { + std::fs::create_dir_all(parent)?; + } + let part = match destination.extension() { + Some(ext) => destination.with_extension(format!("{}.part", ext.to_string_lossy())), + None => destination.with_extension("part"), + }; + let name = destination + .file_name() + .map(|n| n.to_string_lossy().into_owned()) + .unwrap_or_else(|| url.to_string()); + + log::info!("{name}: resolving download size"); + let total = remote_content_length(url); + + // Reconcile any existing `.part` against the real size: exactly complete → + // finish; oversized (stale/corrupt) → discard so curl restarts cleanly + // (otherwise `curl -C -` would 416 forever). + if let Some(total) = total { + let size = std::fs::metadata(&part).map(|m| m.len()).unwrap_or(0); + if size == total { + std::fs::rename(&part, destination)?; + return Ok(()); + } + if size > total { + let _ = std::fs::remove_file(&part); + } + } + log::info!( + "{name}: downloading via curl ({} bytes)", + total + .map(|t| t.to_string()) + .unwrap_or_else(|| "unknown size".into()) + ); + + let mut stalls = 0usize; + loop { + let before = std::fs::metadata(&part).map(|m| m.len()).unwrap_or(0); + match run_curl_download(url, &part, total, &mut on_progress) { + Ok(()) => break, + Err(error) => { + let after = std::fs::metadata(&part).map(|m| m.len()).unwrap_or(0); + if after > before { + log::warn!("{name}: curl interrupted at {after} bytes, resuming: {error}"); + stalls = 0; + } else { + stalls += 1; + log::warn!( + "{name}: curl made no progress ({stalls}/{MAX_STALL_RETRIES}): {error}" + ); + if stalls >= MAX_STALL_RETRIES { + // Discard the partial so a future attempt restarts clean + // rather than getting stuck re-resuming a bad file. + let _ = std::fs::remove_file(&part); + return Err(error); + } + } + std::thread::sleep(std::time::Duration::from_secs(2)); + } + } + } + + if let Some(total) = total { + let got = std::fs::metadata(&part).map(|m| m.len()).unwrap_or(0); + if got < total { + anyhow::bail!("{name}: incomplete after curl ({got}/{total} bytes)"); + } + } + std::fs::rename(&part, destination)?; + Ok(()) +} + +/// Size probe via `curl -r 0-0` (a 1-byte Range request), parsing the total +/// from the `Content-Range: bytes 0-0/` header. Uses curl rather than +/// ureq so no part of the download path depends on ureq (which hangs on this +/// VM's TLS stack). Returns None if the server doesn't report a size. +fn remote_content_length(url: &str) -> Option { + let mut command = curl_command(); + command.args([ + "-sL", + "-r", + "0-0", + "-D", + "-", + "-o", + NULL_DEVICE, + "--connect-timeout", + "30", + "--max-time", + "30", + url, + ]); + let output = command.output().ok()?; + let headers = String::from_utf8_lossy(&output.stdout); + for line in headers.lines() { + if let Some(rest) = line.to_ascii_lowercase().strip_prefix("content-range:") { + if let Some(total) = rest.rsplit('/').next().map(str::trim) { + if let Ok(n) = total.parse::() { + return Some(n); + } + } + } + } + None +} + +/// Run one `curl` download to `dest`, resuming from any partial file, while +/// reporting progress from the growing file size. Returns an error (leaving the +/// partial in place) if curl exits non-zero. +fn run_curl_download( + url: &str, + dest: &Path, + total: Option, + on_progress: &mut impl FnMut(u64, Option), +) -> Result<()> { + let mut command = curl_command(); + command + .arg("-fSL") // fail on HTTP errors, follow redirects, show errors + .args(["-C", "-"]) // resume from the existing output file + .args(["--retry", "3", "--retry-delay", "1", "--retry-connrefused"]) + .args(["--connect-timeout", "30"]) + // Abort (then --retry resumes) if under 1 KB/s for 30s — a real + // inactivity timeout, which is what ureq couldn't deliver here. + .args(["--speed-limit", "1024", "--speed-time", "30"]) + .arg("-s") // no progress meter (we watch the file instead) + .arg("-o") + .arg(dest) + .arg(url) + .stdout(std::process::Stdio::null()) + .stderr(std::process::Stdio::piped()); + + let mut child = command + .spawn() + .map_err(|e| anyhow::anyhow!("failed to launch curl (required for downloads): {e}"))?; + + loop { + if let Some(status) = child.try_wait()? { + if status.success() { + return Ok(()); + } + let mut stderr = String::new(); + if let Some(mut pipe) = child.stderr.take() { + let _ = pipe.read_to_string(&mut stderr); + } + anyhow::bail!( + "curl exited with {}: {}", + status + .code() + .map(|c| c.to_string()) + .unwrap_or_else(|| "signal".into()), + stderr.trim() + ); + } + let downloaded = std::fs::metadata(dest).map(|m| m.len()).unwrap_or(0); + on_progress(downloaded, total); + std::thread::sleep(std::time::Duration::from_millis(300)); + } +} + +/// Download a NuGet package (a zip) resiliently, then extract the single file +/// at `archive_path` into `destination`. +pub fn download_nuget_file( + source_url: &str, + archive_path: &str, + destination: &Path, + on_progress: impl FnMut(u64, Option), +) -> Result<()> { + let package = destination.with_extension("nupkg"); + download_file_resilient(source_url, &package, on_progress)?; + + log::info!("extracting {archive_path} from package"); + let file = std::fs::File::open(&package)?; + let mut archive = zip::ZipArchive::new(file)?; + let mut dll = archive.by_name(archive_path)?; + let temp_destination = destination.with_extension("tmp"); + { + let mut out = std::fs::File::create(&temp_destination)?; + std::io::copy(&mut dll, &mut out)?; + } + std::fs::rename(&temp_destination, destination)?; + let _ = std::fs::remove_file(&package); + log::info!("extracted {archive_path}"); + Ok(()) +} diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 706012a..7807e1f 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -3,10 +3,12 @@ mod captioner; mod color; mod commands; mod db; +mod download; mod embedder; mod hnsw_index; mod indexer; mod media; +mod onnx_runtime; mod storage; mod tagger; mod thumbnail; diff --git a/src-tauri/src/onnx_runtime.rs b/src-tauri/src/onnx_runtime.rs new file mode 100644 index 0000000..6610ec1 --- /dev/null +++ b/src-tauri/src/onnx_runtime.rs @@ -0,0 +1,137 @@ +//! Shared ONNX Runtime provisioning: downloading the runtime + DirectML DLLs +//! and initializing `ort` from them. +//! +//! Both ONNX consumers go through here — the tagger (live) and the Florence-2 +//! captioner (backend intact, UI disabled). The DLLs are Windows/DirectML +//! specific; a future cross-platform build needs a per-OS runtime strategy. + +use crate::download; +use anyhow::Result; +use std::path::{Path, PathBuf}; +use std::sync::Mutex; + +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"; + +pub const ONNX_RUNTIME_DLL_FILE: &str = "onnxruntime/onnxruntime.dll"; +pub const ONNX_RUNTIME_PROVIDERS_DLL_FILE: &str = "onnxruntime/onnxruntime_providers_shared.dll"; +pub const DIRECTML_DLL_FILE: &str = "onnxruntime/DirectML.dll"; + +/// The shared runtime DLLs, as paths relative to [`runtime_dir`]. +pub const RUNTIME_DLLS: &[&str] = &[ + ONNX_RUNTIME_DLL_FILE, + ONNX_RUNTIME_PROVIDERS_DLL_FILE, + DIRECTML_DLL_FILE, +]; + +/// `(destination_file, nuget_package_url, path_inside_package)` for each DLL. +pub 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", + ), +]; + +// Mutex rather than OnceLock: a failed attempt (DLL not yet +// downloaded) must NOT be cached, or a later successful download could never +// recover within the same app session. +static ORT_RUNTIME_INIT: Mutex = Mutex::new(false); + +/// Directory the shared runtime DLLs are provisioned into. +/// +/// Historically the DLLs were downloaded as part of the Florence-2 caption +/// model, so they live inside that model's directory +/// (`models/florence-2-base-ft/onnxruntime/`). The location is kept even +/// though the tagger is now the main consumer, so existing installs don't +/// have to re-download the runtime. +pub fn runtime_dir(app_data_dir: &Path) -> PathBuf { + crate::captioner::model_dir(app_data_dir) +} + +/// Initialize `ort` from the already-downloaded runtime DLL in `local_dir`. +/// Fails if the DLL is missing — use [`provision_onnx_runtime_with_progress`] +/// to download it first on a clean install. +pub fn ensure_onnx_runtime(local_dir: &Path) -> Result<()> { + let mut initialized = ORT_RUNTIME_INIT + .lock() + .map_err(|_| anyhow::anyhow!("ONNX runtime init lock poisoned"))?; + if *initialized { + return Ok(()); + } + let dll_path = local_dir.join(ONNX_RUNTIME_DLL_FILE); + if !dll_path.exists() { + anyhow::bail!("ONNX Runtime DLL is missing: {}", dll_path.display()); + } + ort::environment::init_from(&dll_path) + .map_err(|error| anyhow::anyhow!(error.to_string()))? + .with_name("phokus-florence") + .commit(); + *initialized = true; + Ok(()) +} + +/// Download any ONNX Runtime DLLs missing from `local_dir`, reporting per-file +/// byte progress as `(short_label, downloaded_bytes, total_bytes)`. +/// `total_bytes` is `None` when the server omits Content-Length. Unlike +/// `ensure_onnx_runtime` (init only), this actually provisions the files — +/// callers that can run on a clean install must call this first. The callback +/// fires per chunk; callers should throttle. +pub fn provision_onnx_runtime_with_progress( + local_dir: &Path, + mut on_progress: impl FnMut(&str, u64, Option), +) -> Result<()> { + for (destination_file, source_url, archive_path) in ONNX_RUNTIME_FILES { + let destination = local_dir.join(destination_file); + if destination.exists() { + continue; + } + // Strip the "onnxruntime/" prefix for a clean label. + let label = destination_file + .rsplit('/') + .next() + .unwrap_or(destination_file); + download::download_nuget_file( + source_url, + archive_path, + &destination, + |downloaded, total| on_progress(label, downloaded, total), + )?; + } + Ok(()) +} + +/// Number of ONNX Runtime DLLs still missing from `local_dir` (for progress +/// step counts before downloading). +pub fn missing_onnx_runtime_count(local_dir: &Path) -> usize { + ONNX_RUNTIME_FILES + .iter() + .filter(|(destination_file, _, _)| !local_dir.join(destination_file).exists()) + .count() +} + +/// Download all missing runtime DLLs without progress reporting. +pub fn download_onnx_runtime_files(local_dir: &Path) -> Result<()> { + if !cfg!(target_os = "windows") { + anyhow::bail!("ONNX Runtime DLL 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::download_nuget_file(source_url, archive_path, &destination, |_, _| {})?; + } + + Ok(()) +} diff --git a/src-tauri/src/tagger.rs b/src-tauri/src/tagger.rs index d2a7f03..417afc0 100644 --- a/src-tauri/src/tagger.rs +++ b/src-tauri/src/tagger.rs @@ -31,15 +31,6 @@ const JOYTAG_STD: [f32; 3] = [0.268_629_54, 0.261_302_6, 0.275_777_1]; // JoyTag's recommended detection threshold. const JOYTAG_DEFAULT_THRESHOLD: f32 = 0.4; -// Shared ONNX Runtime DLLs (live in the caption model's `onnxruntime/` dir and -// are reused by the tagger). The per-model weight + label files are listed by -// `TaggerModel::download_files`. -const TAGGER_RUNTIME_DLLS: &[&str] = &[ - "onnxruntime/onnxruntime.dll", - "onnxruntime/onnxruntime_providers_shared.dll", - "onnxruntime/DirectML.dll", -]; - // 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 @@ -354,12 +345,11 @@ pub fn set_tagger_batch_size(app_data_dir: &Path, batch_size: usize) -> Result TaggerModelStatus { let model = tagger_model(app_data_dir); let local_dir = model_dir(app_data_dir); - // The ONNX runtime DLLs live in the caption model dir; reuse them. - let caption_model_dir = app_data_dir.join("models").join("florence-2-base-ft"); + let runtime_dir = crate::onnx_runtime::runtime_dir(app_data_dir); - let mut missing_files: Vec = TAGGER_RUNTIME_DLLS + let mut missing_files: Vec = crate::onnx_runtime::RUNTIME_DLLS .iter() - .filter(|file| !caption_model_dir.join(file).exists()) + .filter(|file| !runtime_dir.join(file).exists()) .map(|file| (*file).to_string()) .collect(); missing_files.extend( @@ -387,15 +377,14 @@ pub fn prepare_tagger_model_with_progress( let local_dir = model_dir(app_data_dir); std::fs::create_dir_all(&local_dir)?; - // The tagger shares the ONNX runtime DLLs with the captioner; download - // them here so the tagger works even on a clean install where the caption - // model has never been fetched (ensure_onnx_runtime only initializes). + // Download the shared ONNX Runtime DLLs here so the tagger works even on + // a clean install (ensure_onnx_runtime only initializes). let download_files = model.download_files(); - let caption_model_dir = crate::captioner::model_dir(app_data_dir); - std::fs::create_dir_all(&caption_model_dir)?; + let runtime_dir = crate::onnx_runtime::runtime_dir(app_data_dir); + std::fs::create_dir_all(&runtime_dir)?; // Unified step count across DLLs and tagger files so the bar is coherent. - let dll_count = crate::captioner::missing_onnx_runtime_count(&caption_model_dir); + let dll_count = crate::onnx_runtime::missing_onnx_runtime_count(&runtime_dir); let model_pending = download_files .iter() .filter(|file| !local_dir.join(file).exists()) @@ -415,8 +404,8 @@ pub fn prepare_tagger_model_with_progress( // ── ONNX runtime DLLs (small, but non-trivial on a slow link) ── { let mut last_emit = Instant::now() - std::time::Duration::from_secs(1); - crate::captioner::provision_onnx_runtime_with_progress( - &caption_model_dir, + crate::onnx_runtime::provision_onnx_runtime_with_progress( + &runtime_dir, |label, downloaded, total| { if last_emit.elapsed() >= std::time::Duration::from_millis(200) { last_emit = Instant::now(); @@ -434,7 +423,7 @@ pub fn prepare_tagger_model_with_progress( completed_files += dll_count; } log::info!("Tagger: ONNX runtime DLLs ready; initializing runtime"); - crate::captioner::ensure_onnx_runtime(&caption_model_dir)?; + crate::onnx_runtime::ensure_onnx_runtime(&runtime_dir)?; log::info!("Tagger: runtime initialized; downloading model files"); // ── Tagger model files (model.onnx is ~446 MB) ── @@ -452,7 +441,7 @@ pub fn prepare_tagger_model_with_progress( let url = repo.url(file); let label = (*file).to_string(); let mut last_emit = Instant::now() - std::time::Duration::from_secs(1); - crate::captioner::download_file_resilient(&url, &destination, |downloaded, total| { + crate::download::download_file_resilient(&url, &destination, |downloaded, total| { if last_emit.elapsed() >= std::time::Duration::from_millis(200) { last_emit = Instant::now(); emit_progress(TaggerModelProgress { @@ -500,8 +489,7 @@ pub fn probe_tagger_runtime(app_data_dir: &Path) -> Result { } 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)?; + crate::onnx_runtime::ensure_onnx_runtime(&crate::onnx_runtime::runtime_dir(app_data_dir))?; let acceleration = tagger_acceleration(app_data_dir); let model_path = local_dir.join("model.onnx"); @@ -663,13 +651,11 @@ impl WdTagger { 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| { + let runtime_dir = crate::onnx_runtime::runtime_dir(app_data_dir); + crate::onnx_runtime::ensure_onnx_runtime(&runtime_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}" + "ONNX Runtime not initialised — the shared runtime DLLs are missing; \ + re-run the tagger model download. Original error: {e}" ) })?; @@ -856,11 +842,11 @@ impl JoyTagger { let local_dir = model_dir(app_data_dir); // Shared ONNX runtime DLLs (see WdTagger::new). - let caption_model_dir = app_data_dir.join("models").join("florence-2-base-ft"); - crate::captioner::ensure_onnx_runtime(&caption_model_dir).map_err(|e| { + let runtime_dir = crate::onnx_runtime::runtime_dir(app_data_dir); + crate::onnx_runtime::ensure_onnx_runtime(&runtime_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}" + "ONNX Runtime not initialised — the shared runtime DLLs are missing; \ + re-run the tagger model download. Original error: {e}" ) })?;