refactor(backend): extract download and onnx_runtime modules from captioner
github/actions/ci GitHub Actions CI finished: success

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.
This commit is contained in:
2026-07-05 20:21:47 +01:00
parent a791f112f5
commit b23212ea1c
6 changed files with 408 additions and 374 deletions
+8 -339
View File
@@ -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<bool> rather than OnceLock<Result>: 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<bool> = 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<CaptionRuntimeProbe>
}
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<Ca
}
let local_dir = model_dir(app_data_dir);
ensure_onnx_runtime(&local_dir)?;
onnx_runtime::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()))
@@ -455,7 +402,7 @@ impl FlorenceCaptioner {
}
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)?;
let caption_detail = caption_detail(app_data_dir);
@@ -664,284 +611,6 @@ fn probe_vision_session(
})
}
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<u64>),
) -> 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<u64>),
) -> 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/<total>` 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<u64> {
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::<u64>() {
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<u64>,
on_progress: &mut impl FnMut(u64, Option<u64>),
) -> 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<u64>),
) -> 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<TensorData> {
let pixels = preprocess_image(image_path)?;
let input = Tensor::from_array(([1usize, 3, 768, 768], pixels.into_boxed_slice()))