diff --git a/src-tauri/src/captioner.rs b/src-tauri/src/captioner.rs index 7d27f5a..37c6f79 100644 --- a/src-tauri/src/captioner.rs +++ b/src-tauri/src/captioner.rs @@ -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), +) -> 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), +) -> 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::().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)?; diff --git a/src-tauri/src/tagger.rs b/src-tauri/src/tagger.rs index cc06abc..5bc1a1f 100644 --- a/src-tauri/src/tagger.rs +++ b/src-tauri/src/tagger.rs @@ -87,9 +87,63 @@ pub struct TaggerModelProgress { pub total_files: usize, pub completed_files: usize, pub current_file: Option, + // Byte progress for the file currently downloading. None for files whose + // size isn't known up front (or between files). + pub downloaded_bytes: Option, + pub total_bytes: Option, 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 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 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, }); diff --git a/src/components/SettingsModal.tsx b/src/components/SettingsModal.tsx index c11f0b6..b2779b9 100644 --- a/src/components/SettingsModal.tsx +++ b/src/components/SettingsModal.tsx @@ -9,6 +9,12 @@ const SECTIONS: { id: SettingsSection; label: string; detail: string }[] = [ { id: "general", label: "General", detail: "App data and diagnostics" }, ]; +function formatBytesShort(bytes: number): string { + if (bytes >= 1024 * 1024 * 1024) return `${(bytes / (1024 * 1024 * 1024)).toFixed(1)} GB`; + if (bytes >= 1024 * 1024) return `${(bytes / (1024 * 1024)).toFixed(0)} MB`; + return `${(bytes / 1024).toFixed(0)} KB`; +} + function StatusPill({ children, tone }: { children: React.ReactNode; tone: "ready" | "muted" | "busy" }) { const className = tone === "ready" @@ -235,16 +241,24 @@ export function SettingsModal() { : "no folders selected"; const thresholdDisplay = taggerThresholdDraft ?? String(taggerThreshold); const batchSizeDisplay = taggerBatchSizeDraft ?? String(taggerBatchSize); - const taggerDownloadLabel = taggerModelProgress - ? `Downloading ${taggerModelProgress.completed_files}/${taggerModelProgress.total_files}` - : taggerModelPreparing - ? "Preparing WD Tagger..." - : taggerReady - ? "Installed" - : "Install model"; - const taggerDownloadPercent = taggerModelProgress - ? Math.round((taggerModelProgress.completed_files / Math.max(taggerModelProgress.total_files, 1)) * 100) - : 0; + const taggerBytesKnown = + taggerModelProgress?.downloaded_bytes != null && + taggerModelProgress.total_bytes != null && + taggerModelProgress.total_bytes > 0; + const taggerDownloadLabel = taggerBytesKnown + ? `Downloading ${formatBytesShort(taggerModelProgress!.downloaded_bytes!)} / ${formatBytesShort(taggerModelProgress!.total_bytes!)}` + : taggerModelProgress + ? `Downloading ${taggerModelProgress.completed_files}/${taggerModelProgress.total_files}` + : taggerModelPreparing + ? "Preparing WD Tagger..." + : taggerReady + ? "Installed" + : "Install model"; + const taggerDownloadPercent = taggerBytesKnown + ? Math.round((taggerModelProgress!.downloaded_bytes! / taggerModelProgress!.total_bytes!) * 100) + : taggerModelProgress + ? Math.round((taggerModelProgress.completed_files / Math.max(taggerModelProgress.total_files, 1)) * 100) + : 0; const runQueueAction = (action: "queue" | "clear") => { const selectedIds = taggingQueueFolderIds; diff --git a/src/components/onboarding/StepAiFeatures.tsx b/src/components/onboarding/StepAiFeatures.tsx index c24fd6a..ecff547 100644 --- a/src/components/onboarding/StepAiFeatures.tsx +++ b/src/components/onboarding/StepAiFeatures.tsx @@ -1,9 +1,20 @@ import { useEffect } from "react"; -import { useGalleryStore } from "../../store"; -import { FakeProgressBar } from "./fakes"; +import { TaggerModelProgress, useGalleryStore } from "../../store"; +import { FakeProgressBar, formatBytes } from "./fakes"; const FAKE_TAGS = ["landscape", "sunset", "outdoors", "no_humans", "ocean", "cloudy_sky"]; +// Prefer the current file's byte fraction (the 1.3 GB model dominates); fall +// back to the coarse step count; null renders an indeterminate bar. +function taggerDownloadFraction(progress: TaggerModelProgress | null): number | null { + if (!progress) return null; + if (progress.downloaded_bytes != null && progress.total_bytes != null && progress.total_bytes > 0) { + return progress.downloaded_bytes / progress.total_bytes; + } + if (progress.total_files > 0) return progress.completed_files / progress.total_files; + return null; +} + export function StepAiFeatures() { const taggerModelStatus = useGalleryStore((s) => s.taggerModelStatus); const taggerModelPreparing = useGalleryStore((s) => s.taggerModelPreparing); @@ -61,16 +72,15 @@ export function StepAiFeatures() { {taggerModelPreparing ? (
- 0 - ? taggerModelProgress.completed_files / taggerModelProgress.total_files - : null - } - /> - {taggerModelProgress?.current_file ? ( -

{taggerModelProgress.current_file}

- ) : null} + +
+ {taggerModelProgress?.current_file ?? "Preparing..."} + {taggerModelProgress?.downloaded_bytes != null && taggerModelProgress.total_bytes != null ? ( + + {formatBytes(taggerModelProgress.downloaded_bytes)} / {formatBytes(taggerModelProgress.total_bytes)} + + ) : null} +
) : null} {!taggerModelPreparing && taggerModelError ? ( diff --git a/src/store.ts b/src/store.ts index a150bb5..8a58414 100644 --- a/src/store.ts +++ b/src/store.ts @@ -112,6 +112,8 @@ export interface TaggerModelProgress { total_files: number; completed_files: number; current_file: string | null; + downloaded_bytes: number | null; + total_bytes: number | null; done: boolean; }