fix(tagger): provision ONNX runtime DLLs during tagger model prep

On a clean install the WD tagger download failed instantly: only the
(UI-removed) caption model prep ever downloaded the shared ONNX Runtime
DLLs, while ensure_onnx_runtime — despite its comment — only initializes
and errors when the DLL is absent. Worse, the OnceLock cached that error
for the whole session, and the onboarding step swallowed the message.

- new captioner::provision_onnx_runtime downloads missing DLLs; tagger
  prep calls it before init
- ensure_onnx_runtime no longer caches failure (Mutex<bool>, success-only
  latch) so a later download can recover in-session
- onboarding AI step now displays taggerModelError
This commit is contained in:
2026-06-13 00:00:35 +01:00
parent 3b7190d353
commit b72f140737
3 changed files with 43 additions and 20 deletions
+34 -18
View File
@@ -11,7 +11,7 @@ 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::sync::Mutex;
use std::time::Instant;
use tokenizers::Tokenizer;
@@ -62,7 +62,10 @@ const REQUIRED_FILES: &[&str] = &[
"onnx/embed_tokens_fp16.onnx",
];
static ORT_RUNTIME_INIT: OnceLock<std::result::Result<(), String>> = OnceLock::new();
// 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.
@@ -648,23 +651,36 @@ 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);
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)
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 that are missing from `local_dir`. Unlike
/// `ensure_onnx_runtime` (init only), this actually provisions the files —
/// callers that can run on a clean install must call this first.
pub fn provision_onnx_runtime(local_dir: &Path) -> Result<()> {
for (destination_file, source_url, archive_path) in ONNX_RUNTIME_FILES {
let destination = local_dir.join(destination_file);
if destination.exists() {
continue;
}
download_nuget_file(source_url, archive_path, &destination)?;
}
Ok(())
}
fn download_onnx_runtime_files(local_dir: &Path) -> Result<()> {
+3 -2
View File
@@ -251,11 +251,12 @@ pub fn prepare_tagger_model_with_progress(
std::fs::create_dir_all(&local_dir)?;
// Ensure the shared ONNX runtime DLLs are present. The tagger shares
// them with the captioner; install them here so the tagger is fully
// them with the captioner; download them here so the tagger is fully
// functional even on a clean install where the caption model has never
// been downloaded.
// been downloaded (ensure_onnx_runtime only initializes, never downloads).
let caption_model_dir = crate::captioner::model_dir(app_data_dir);
std::fs::create_dir_all(&caption_model_dir)?;
crate::captioner::provision_onnx_runtime(&caption_model_dir)?;
crate::captioner::ensure_onnx_runtime(&caption_model_dir)?;
// Download the two tagger-specific files.