feat(tagger): real byte progress for model download instead of a spinner

The 1.3 GB model.onnx downloaded via hf-hub's repo.get() and the ONNX
runtime DLLs via a read_to_end — neither reported bytes, so the user saw
only an indeterminate spinner on a multi-minute download.

- tagger model files now download via repo.download_with_progress with a
  HubProgress adapter over hf-hub's Progress trait (throttled 200ms events
  carrying downloaded/total bytes)
- captioner's nuget DLL download streams in 64 KB chunks, parsing
  content-length, and reports per-chunk byte progress
- TaggerModelProgress carries downloaded_bytes/total_bytes; onboarding and
  Settings show a determinate bar plus 'X MB / Y MB' for the current file
This commit is contained in:
2026-06-13 08:56:49 +01:00
parent 9c135179a3
commit f8e981c6f6
5 changed files with 213 additions and 63 deletions
+58 -12
View File
@@ -669,20 +669,45 @@ pub fn ensure_onnx_runtime(local_dir: &Path) -> Result<()> {
Ok(())
}
/// Download any ONNX Runtime DLLs that are missing from `local_dir`. Unlike
/// 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.
pub fn provision_onnx_runtime(local_dir: &Path) -> Result<()> {
/// 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;
}
download_nuget_file(source_url, archive_path, &destination)?;
// 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!(
@@ -692,22 +717,43 @@ fn download_onnx_runtime_files(local_dir: &Path) -> Result<()> {
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)?;
download_nuget_file(source_url, archive_path, &destination, |_, _| {})?;
}
Ok(())
}
fn download_nuget_file(source_url: &str, archive_path: &str, destination: &Path) -> Result<()> {
fn download_nuget_file(
source_url: &str,
archive_path: &str,
destination: &Path,
mut on_progress: impl FnMut(u64, Option<u64>),
) -> 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 total_bytes = response
.headers()
.get("content-length")
.and_then(|value| value.to_str().ok())
.and_then(|value| value.parse::<u64>().ok());
// Stream into memory (the zip reader needs Seek) while reporting bytes.
let mut bytes = Vec::with_capacity(total_bytes.unwrap_or(0) as usize);
let mut reader = response.body_mut().as_reader();
let mut buffer = [0u8; 64 * 1024];
let mut downloaded = 0u64;
loop {
let read = reader
.read(&mut buffer)
.map_err(|error| anyhow::anyhow!("{error}"))?;
if read == 0 {
break;
}
bytes.extend_from_slice(&buffer[..read]);
downloaded += read as u64;
on_progress(downloaded, total_bytes);
}
let mut archive = zip::ZipArchive::new(Cursor::new(bytes))?;
let mut dll = archive.by_name(archive_path)?;
+107 -29
View File
@@ -87,9 +87,63 @@ pub struct TaggerModelProgress {
pub total_files: usize,
pub completed_files: usize,
pub current_file: Option<String>,
// Byte progress for the file currently downloading. None for files whose
// size isn't known up front (or between files).
pub downloaded_bytes: Option<u64>,
pub total_bytes: Option<u64>,
pub done: bool,
}
/// Adapts hf-hub's `Progress` trait to throttled `tagger-model-progress`
/// events so the 1.3 GB model download reports real bytes, not a spinner.
struct HubProgress<'a, F: Fn(TaggerModelProgress)> {
emit: &'a F,
total_files: usize,
completed_files: usize,
current_file: String,
downloaded: u64,
total: u64,
last_emit: Instant,
}
impl<F: Fn(TaggerModelProgress)> hf_hub::api::Progress for HubProgress<'_, F> {
fn init(&mut self, size: usize, _filename: &str) {
self.total = size as u64;
self.downloaded = 0;
self.last_emit = Instant::now();
(self.emit)(self.snapshot(false));
}
fn update(&mut self, size: usize) {
self.downloaded = self.downloaded.saturating_add(size as u64);
if self.last_emit.elapsed() >= std::time::Duration::from_millis(200) {
self.last_emit = Instant::now();
(self.emit)(self.snapshot(false));
}
}
fn finish(&mut self) {
(self.emit)(self.snapshot(false));
}
}
impl<F: Fn(TaggerModelProgress)> HubProgress<'_, F> {
fn snapshot(&self, done: bool) -> TaggerModelProgress {
TaggerModelProgress {
total_files: self.total_files,
completed_files: self.completed_files,
current_file: Some(self.current_file.clone()),
downloaded_bytes: Some(self.downloaded),
total_bytes: if self.total > 0 {
Some(self.total)
} else {
None
},
done,
}
}
}
// ---------------------------------------------------------------------------
// Runtime probe types exposed to the frontend
// ---------------------------------------------------------------------------
@@ -250,59 +304,83 @@ pub fn prepare_tagger_model_with_progress(
let local_dir = model_dir(app_data_dir);
std::fs::create_dir_all(&local_dir)?;
// Ensure the shared ONNX runtime DLLs are present. The tagger shares
// 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 (ensure_onnx_runtime only initializes, never downloads).
// 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).
const DOWNLOAD_FILES: &[&str] = &["model.onnx", "selected_tags.csv"];
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.
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
// 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 model_pending = DOWNLOAD_FILES
.iter()
.filter(|file| local_dir.join(file).exists())
.filter(|file| !local_dir.join(file).exists())
.count();
let total_files = dll_count + model_pending;
let mut completed_files = 0usize;
emit_progress(TaggerModelProgress {
total_files: DOWNLOAD_FILES.len(),
total_files,
completed_files,
current_file: None,
done: completed_files == DOWNLOAD_FILES.len(),
downloaded_bytes: None,
total_bytes: None,
done: total_files == 0,
});
// ── 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,
|label, downloaded, total| {
if last_emit.elapsed() >= std::time::Duration::from_millis(200) {
last_emit = Instant::now();
emit_progress(TaggerModelProgress {
total_files,
completed_files,
current_file: Some(format!("ONNX Runtime: {label}")),
downloaded_bytes: Some(downloaded),
total_bytes: total,
done: false,
});
}
},
)?;
completed_files += dll_count;
}
crate::captioner::ensure_onnx_runtime(&caption_model_dir)?;
// ── Tagger model files (model.onnx is ~1.3 GB) ──
let api = Api::new()?;
let repo = api.repo(Repo::new(WD_TAGGER_MODEL_ID.to_string(), RepoType::Model));
for file in DOWNLOAD_FILES {
let destination = local_dir.join(file);
if destination.exists() {
continue;
}
emit_progress(TaggerModelProgress {
total_files: DOWNLOAD_FILES.len(),
let progress = HubProgress {
emit: &emit_progress,
total_files,
completed_files,
current_file: Some((*file).to_string()),
done: false,
});
let cached = repo.get(file)?;
current_file: (*file).to_string(),
downloaded: 0,
total: 0,
last_emit: Instant::now(),
};
let cached = repo.download_with_progress(file, progress)?;
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(),
total_files,
completed_files,
current_file: None,
downloaded_bytes: None,
total_bytes: None,
done: true,
});