@ondeinference / onde
Code Commits Issues Pulls Cloud Agent CI/CD
copilot/assess-inference-engine-complexity
copilot/assess-inference-engine-complexity copilot/fix-onde-npm-sdk-v1-1-2-release development feature/diffusion-vision main
rs 1,221 lines 44.1 KB
Raw
1 //! HuggingFace hub cache inspection, repair, and model download.
2 //!
3 //! This module owns all on-disk state management for locally cached
4 //! HuggingFace models and exposes a progress-callback-based download API
5 //! that is completely decoupled from any application framework.
6 //!
7 //! 1. Call [`list_local_hf_models`] / [`list_supported_hf_models`] /
8 //! [`delete_local_hf_model`] directly from your command handlers.
9 //! 2. Call [`download_model`] with an `Arc<dyn Fn(ModelDownloadProgress)>`
10 //! that forwards progress to whatever event system your app uses.
11
12 use {
13 crate::inference::models::{SUPPORTED_MODELS, SUPPORTED_MODEL_INFO},
14 log::{debug, info, warn},
15 serde::{Deserialize, Serialize},
16 std::{
17 fs,
18 path::{Path, PathBuf},
19 },
20 tsync::tsync,
21 };
22
23 // ---------------------------------------------------------------------------
24 // Public types
25 // ---------------------------------------------------------------------------
26
27 /// Information about a locally cached HuggingFace model.
28 #[tsync]
29 #[derive(Debug, Clone, Serialize, Deserialize)]
30 pub struct LocalHfModel {
31 /// The full model identifier, e.g. "bartowski/Qwen2.5-1.5B-Instruct-GGUF".
32 pub model_id: String,
33 /// The organisation or user that published the model.
34 pub org: String,
35 /// The model name without the org prefix.
36 pub name: String,
37 /// Total size on disk in bytes.
38 pub size_bytes: u64,
39 /// Human-readable size string, e.g. "4.2 GB".
40 pub size_display: String,
41 /// Absolute path to the model cache directory.
42 pub path: String,
43 /// List of snapshot revisions that are locally available.
44 pub revisions: Vec<String>,
45 }
46
47 /// Response returned by [`list_local_hf_models`].
48 #[tsync]
49 #[derive(Debug, Clone, Serialize, Deserialize)]
50 pub struct LocalHfModelsResponse {
51 /// The models found in the local HuggingFace cache.
52 pub models: Vec<LocalHfModel>,
53 /// The path that was scanned.
54 pub cache_path: String,
55 /// Total size of all cached models in bytes.
56 pub total_size_bytes: u64,
57 /// Human-readable total size string.
58 pub total_size_display: String,
59 }
60
61 /// A supported model that may or may not be downloaded locally.
62 #[tsync]
63 #[derive(Debug, Clone, Serialize, Deserialize)]
64 pub struct SupportedHfModel {
65 /// The full model identifier, e.g. "bartowski/Qwen2.5-1.5B-Instruct-GGUF".
66 pub model_id: String,
67 /// Human-friendly display name.
68 pub name: String,
69 /// Organisation or publisher display name.
70 pub org: String,
71 /// Short description of the model.
72 pub description: String,
73 /// Whether this model is fully downloaded locally.
74 pub is_downloaded: bool,
75 /// Whether a partial (incomplete) download exists on disk.
76 /// True when a cache directory exists but its size is less than 99 %
77 /// of `expected_size_bytes`.
78 pub is_incomplete: bool,
79 /// Bytes currently on disk for this model (0 when not present at all).
80 pub local_size_bytes: u64,
81 /// Human-readable local size string, e.g. "4.2 GB".
82 pub local_size_display: String,
83 /// Approximate total size in bytes when fully downloaded.
84 pub expected_size_bytes: u64,
85 /// Human-readable expected size string, e.g. "23.80 GB".
86 pub expected_size_display: String,
87 }
88
89 /// Response returned by [`list_supported_hf_models`].
90 #[tsync]
91 #[derive(Debug, Clone, Serialize, Deserialize)]
92 pub struct SupportedHfModelsResponse {
93 /// All models the app supports, with download status.
94 pub models: Vec<SupportedHfModel>,
95 }
96
97 /// Progress payload passed to the [`download_model`] callback on every poll
98 /// tick and on completion.
99 ///
100 /// Forward this value to your app's event system so the frontend can render
101 /// a progress bar.
102 #[tsync]
103 #[derive(Debug, Clone, Serialize, Deserialize)]
104 pub struct ModelDownloadProgress {
105 /// The model being downloaded.
106 pub model_id: String,
107 /// Bytes downloaded so far (observed from cache directory size).
108 pub downloaded_bytes: u64,
109 /// Human-readable downloaded size, e.g. "1.2 GB".
110 pub downloaded_display: String,
111 /// Expected total size in bytes (approximate).
112 pub total_bytes: u64,
113 /// Human-readable total size, e.g. "23.80 GB".
114 pub total_display: String,
115 /// Download progress as a value between 0.0 and 1.0.
116 pub progress: f64,
117 /// Whether the download has completed.
118 pub done: bool,
119 }
120
121 // ---------------------------------------------------------------------------
122 // Internal constants
123 // ---------------------------------------------------------------------------
124
125 /// Fraction of `expected_size_bytes` that must be present on disk for a
126 /// model to be considered fully downloaded (99 %).
127 const DOWNLOAD_COMPLETE_THRESHOLD: f64 = 0.99;
128
129 // ---------------------------------------------------------------------------
130 // Private helpers
131 // ---------------------------------------------------------------------------
132
133 /// Resolve the HuggingFace hub cache directory.
134 ///
135 /// Respects the `HF_HOME` / `HUGGINGFACE_HUB_CACHE` environment variables.
136 /// Falls back to `~/.cache/huggingface/hub`.
137 ///
138 /// On iOS the `HF_HOME` variable is set during app setup (see
139 /// `setup_application_filesystem`) to a path inside the app's writable
140 /// `Library/Application Support/` container. We create the directory if it
141 /// doesn't exist yet so that first-run and freshly-installed apps work
142 /// correctly.
143 fn hf_cache_dir() -> Option<PathBuf> {
144 // 1. Explicit hub cache override
145 if let Ok(cache) = std::env::var("HUGGINGFACE_HUB_CACHE") {
146 let p = PathBuf::from(cache);
147 let _ = fs::create_dir_all(&p);
148 if p.is_dir() {
149 return Some(p);
150 }
151 }
152
153 // 2. HF_HOME override (hub is a subdirectory)
154 if let Ok(hf_home) = std::env::var("HF_HOME") {
155 let p = PathBuf::from(hf_home).join("hub");
156 let _ = fs::create_dir_all(&p);
157 if p.is_dir() {
158 return Some(p);
159 }
160 }
161
162 // 3. Default: ~/.cache/huggingface/hub
163 if let Some(home) = home::home_dir() {
164 let p = home.join(".cache").join("huggingface").join("hub");
165 let _ = fs::create_dir_all(&p);
166 if p.is_dir() {
167 return Some(p);
168 }
169 }
170
171 None
172 }
173
174 /// Recursively compute the total size of a directory in bytes.
175 ///
176 /// Follows symlinks to count the real file size and also counts in-progress
177 /// `.incomplete` staging files so that progress polling reflects partial
178 /// downloads before they are renamed to their final blob name.
179 fn dir_size(path: &PathBuf) -> u64 {
180 let mut total: u64 = 0;
181 if let Ok(entries) = fs::read_dir(path) {
182 for entry in entries.flatten() {
183 let entry_path = entry.path();
184 if entry_path.is_dir() {
185 total += dir_size(&entry_path);
186 } else {
187 // Count both completed blobs AND in-progress staging files
188 // (hf-hub downloads to `<sha256>.incomplete` and renames on
189 // completion — without this, progress stays at 0% for the
190 // entire download and jumps to 100% at the end).
191 if let Ok(meta) = entry_path.symlink_metadata() {
192 if meta.is_symlink() {
193 // Follow the symlink to get the real file size.
194 if let Ok(real_meta) = entry_path.metadata() {
195 total += real_meta.len();
196 }
197 } else {
198 total += meta.len();
199 }
200 }
201 }
202 }
203 }
204 total
205 }
206
207 /// Format a byte count into a human-readable string.
208 fn format_size(bytes: u64) -> String {
209 const KB: u64 = 1024;
210 const MB: u64 = KB * 1024;
211 const GB: u64 = MB * 1024;
212 const TB: u64 = GB * 1024;
213
214 if bytes >= TB {
215 format!("{:.2} TB", bytes as f64 / TB as f64)
216 } else if bytes >= GB {
217 format!("{:.2} GB", bytes as f64 / GB as f64)
218 } else if bytes >= MB {
219 format!("{:.1} MB", bytes as f64 / MB as f64)
220 } else if bytes >= KB {
221 format!("{:.0} KB", bytes as f64 / KB as f64)
222 } else {
223 format!("{} B", bytes)
224 }
225 }
226
227 /// List the snapshot revisions available for a cached model directory.
228 fn list_revisions(model_dir: &Path) -> Vec<String> {
229 let snapshots_dir = model_dir.join("snapshots");
230 if !snapshots_dir.is_dir() {
231 return Vec::new();
232 }
233
234 let mut revisions = Vec::new();
235 if let Ok(entries) = fs::read_dir(&snapshots_dir) {
236 for entry in entries.flatten() {
237 if entry.path().is_dir() {
238 if let Some(name) = entry.file_name().to_str() {
239 revisions.push(name.to_string());
240 }
241 }
242 }
243 }
244 revisions.sort();
245 revisions
246 }
247
248 /// Build a [`ModelDownloadProgress`] payload from raw byte counts.
249 #[cfg(any(
250 target_os = "macos",
251 target_os = "ios",
252 target_os = "tvos",
253 target_os = "visionos",
254 target_os = "watchos",
255 target_os = "android"
256 ))]
257 fn make_progress(model_id: &str, downloaded: u64, total: u64, done: bool) -> ModelDownloadProgress {
258 let progress = if total > 0 {
259 (downloaded as f64 / total as f64).min(1.0)
260 } else {
261 0.0
262 };
263 ModelDownloadProgress {
264 model_id: model_id.to_string(),
265 downloaded_bytes: downloaded,
266 downloaded_display: format_size(downloaded),
267 total_bytes: total,
268 total_display: format_size(total),
269 progress,
270 done,
271 }
272 }
273
274 // ---------------------------------------------------------------------------
275 // Public cache utilities
276 // ---------------------------------------------------------------------------
277
278 /// Return the cache sub-directory path for a given model ID.
279 ///
280 /// E.g. for `"bartowski/Qwen2.5-1.5B-Instruct-GGUF"` the cache directory is
281 /// `<hf_cache>/models--bartowski--Qwen2.5-1.5B-Instruct-GGUF`.
282 pub fn model_cache_path(model_id: &str) -> Option<PathBuf> {
283 let cache_dir = hf_cache_dir()?;
284 let dir_name = format!("models--{}", model_id.replace('/', "--"));
285 Some(cache_dir.join(dir_name))
286 }
287
288 /// Pre-create the HuggingFace cache directory tree for `model_id` so that the
289 /// progress-monitor thread can find a valid path immediately, even before
290 /// hf-hub creates it during the first download.
291 ///
292 /// hf-hub won't be confused by a pre-existing (empty) directory — it will
293 /// simply populate it as normal.
294 pub fn ensure_model_cache_dir(model_id: &str) {
295 if let Some(model_path) = model_cache_path(model_id) {
296 let blobs_dir = model_path.join("blobs");
297 if let Err(e) = fs::create_dir_all(&blobs_dir) {
298 warn!(
299 "Could not pre-create model cache dir {}: {}",
300 blobs_dir.display(),
301 e
302 );
303 } else {
304 info!("Pre-created model cache dir: {}", model_path.display());
305 }
306 }
307 }
308
309 /// Remove stale `.lock` and orphaned `.part` files from the HF cache
310 /// `blobs/` directory for a model.
311 ///
312 /// The HF hub client creates `.lock` files during download. If a download is
313 /// interrupted these locks persist and prevent any subsequent access (download
314 /// **or** model load). Calling this before loading or resuming a download
315 /// allows the hub client to re-acquire the lock and proceed normally.
316 pub fn clean_stale_lock_files(model_id: &str) {
317 if let Some(cache_path) = model_cache_path(model_id) {
318 let blobs_dir = cache_path.join("blobs");
319 if blobs_dir.is_dir() {
320 // Collect the names of complete blobs (no extension) so we can
321 // decide whether a `.part` file is stale.
322 let complete_blobs: Vec<String> = fs::read_dir(&blobs_dir)
323 .into_iter()
324 .flatten()
325 .flatten()
326 .filter_map(|e| {
327 let p = e.path();
328 if p.is_file() && p.extension().is_none() {
329 p.file_name().map(|n| n.to_string_lossy().to_string())
330 } else {
331 None
332 }
333 })
334 .collect();
335
336 if let Ok(entries) = fs::read_dir(&blobs_dir) {
337 for entry in entries.flatten() {
338 let path = entry.path();
339 let ext = path.extension().and_then(|e| e.to_str());
340
341 match ext {
342 Some("lock") => {
343 info!("Removing stale lock file: {}", path.display());
344 if let Err(e) = fs::remove_file(&path) {
345 warn!("Failed to remove lock file {}: {}", path.display(), e);
346 }
347 }
348 Some("part") => {
349 // A `.part` file is a partial download left by hf-hub.
350 // If the corresponding complete blob exists (same name
351 // without `.part`), the partial file is stale and
352 // causes hf-hub to issue a spurious Range-resume
353 // download instead of using the complete blob.
354 let stem = path.file_stem().and_then(|s| s.to_str()).unwrap_or("");
355 if complete_blobs.iter().any(|b| b == stem) {
356 let part_size = fs::metadata(&path).map(|m| m.len()).unwrap_or(0);
357 info!(
358 "Removing stale .part file ({} bytes, complete blob exists): {}",
359 part_size,
360 path.display()
361 );
362 if let Err(e) = fs::remove_file(&path) {
363 warn!("Failed to remove .part file {}: {}", path.display(), e);
364 }
365 } else {
366 let part_size = fs::metadata(&path).map(|m| m.len()).unwrap_or(0);
367 debug!(
368 "Keeping .part file ({} bytes, no complete blob yet): {}",
369 part_size,
370 path.display()
371 );
372 }
373 }
374 _ => {}
375 }
376 }
377 }
378 }
379 }
380 }
381
382 /// Log the complete HF cache structure for a model to aid debugging.
383 ///
384 /// This is called before model load so the developer can see exactly what's
385 /// present in the cache, whether `refs/main` resolves, whether symlinks are
386 /// intact, and whether any `.part` files are left over from interrupted
387 /// downloads.
388 pub fn diagnose_hf_cache(model_id: &str) {
389 let cache_path = match model_cache_path(model_id) {
390 Some(p) => p,
391 None => {
392 info!("diagnose_hf_cache({}): cache path not found", model_id);
393 return;
394 }
395 };
396
397 if !cache_path.is_dir() {
398 info!(
399 "diagnose_hf_cache({}): directory does not exist: {}",
400 model_id,
401 cache_path.display()
402 );
403 return;
404 }
405
406 info!(
407 "diagnose_hf_cache({}): scanning {}",
408 model_id,
409 cache_path.display()
410 );
411
412 // ── refs/ ────────────────────────────────────────────────────────────
413 let refs_dir = cache_path.join("refs");
414 if refs_dir.is_dir() {
415 if let Ok(entries) = fs::read_dir(&refs_dir) {
416 for entry in entries.flatten() {
417 let name = entry.file_name().to_string_lossy().to_string();
418 match fs::read_to_string(entry.path()) {
419 Ok(content) => {
420 let hash = content.trim();
421 info!(
422 " refs/{} → \"{}\" (len={})",
423 name,
424 if hash.len() > 12 {
425 format!("{}…", &hash[..12])
426 } else {
427 hash.to_string()
428 },
429 hash.len()
430 );
431 }
432 Err(e) => {
433 warn!(" refs/{} → <unreadable: {}>", name, e);
434 }
435 }
436 }
437 }
438 } else {
439 warn!(" refs/ directory missing!");
440 }
441
442 // ── blobs/ ───────────────────────────────────────────────────────────
443 let blobs_dir = cache_path.join("blobs");
444 if blobs_dir.is_dir() {
445 if let Ok(entries) = fs::read_dir(&blobs_dir) {
446 for entry in entries.flatten() {
447 let path = entry.path();
448 let name = entry.file_name().to_string_lossy().to_string();
449 let size = fs::metadata(&path).map(|m| m.len()).unwrap_or(0);
450 let kind = if name.ends_with(".lock") {
451 "LOCK"
452 } else if name.ends_with(".part") {
453 "PART (incomplete)"
454 } else {
455 "blob"
456 };
457 info!(
458 " blobs/{} — {} ({} bytes / {:.1} MB)",
459 name,
460 kind,
461 size,
462 size as f64 / 1_048_576.0
463 );
464 }
465 }
466 } else {
467 warn!(" blobs/ directory missing!");
468 }
469
470 // ── snapshots/ ───────────────────────────────────────────────────────
471 let snapshots_dir = cache_path.join("snapshots");
472 if snapshots_dir.is_dir() {
473 if let Ok(revs) = fs::read_dir(&snapshots_dir) {
474 for rev_entry in revs.flatten() {
475 let rev_name = rev_entry.file_name().to_string_lossy().to_string();
476 let rev_display = if rev_name.len() > 12 {
477 format!("{}…", &rev_name[..12])
478 } else {
479 rev_name.clone()
480 };
481 info!(" snapshots/{}/", rev_display);
482
483 if let Ok(files) = fs::read_dir(rev_entry.path()) {
484 for file_entry in files.flatten() {
485 let file_name = file_entry.file_name().to_string_lossy().to_string();
486 let file_path = file_entry.path();
487
488 if file_path.is_symlink() {
489 let target = fs::read_link(&file_path)
490 .map(|t| t.to_string_lossy().to_string())
491 .unwrap_or_else(|_| "<unreadable>".to_string());
492 let accessible = fs::metadata(&file_path).is_ok();
493 let status = if accessible { "OK" } else { "BROKEN" };
494 info!(" {} → {} [symlink: {}]", file_name, target, status);
495 } else if file_path.is_file() {
496 let size = fs::metadata(&file_path).map(|m| m.len()).unwrap_or(0);
497 info!(
498 " {} [{} bytes / {:.1} MB, real file]",
499 file_name,
500 size,
501 size as f64 / 1_048_576.0
502 );
503 } else {
504 warn!(" {} [not a file or symlink!]", file_name);
505 }
506 }
507 }
508 }
509 }
510 } else {
511 warn!(" snapshots/ directory missing!");
512 }
513
514 info!("diagnose_hf_cache({}): done", model_id);
515 }
516
517 /// Repair broken symlinks in the HF hub cache for a given model.
518 ///
519 /// The `hf-hub` crate stores downloaded files as blobs keyed by SHA-256 and
520 /// creates **symlinks** from `snapshots/<rev>/<filename>` → `../../blobs/<hash>`.
521 /// When the cache directory is copied into an iOS app container (e.g. via
522 /// Xcode "Replace Container" or `ideviceinstaller`), symlinks are often
523 /// silently dropped or broken. The `hf-hub` client then thinks the file is
524 /// missing and re-downloads it — even though the 940 MB blob is right there.
525 ///
526 /// This function walks every snapshot entry for `model_id`. If a path is a
527 /// broken symlink **or** a missing file, and a matching blob exists, it
528 /// replaces the symlink with a hard copy of the blob so that `hf-hub`'s
529 /// `Cache::get()` finds a real file at the expected path.
530 ///
531 /// It also ensures a `refs/main` file exists pointing at the snapshot
532 /// revision, since that file is needed for `hf-hub` to resolve the "main"
533 /// ref to a concrete snapshot directory.
534 ///
535 /// This is idempotent and safe to call on every app launch.
536 pub fn repair_hf_cache_symlinks(model_id: &str) {
537 let cache_path = match model_cache_path(model_id) {
538 Some(p) if p.is_dir() => p,
539 _ => return,
540 };
541
542 let blobs_dir = cache_path.join("blobs");
543 let snapshots_dir = cache_path.join("snapshots");
544 let refs_dir = cache_path.join("refs");
545
546 if !blobs_dir.is_dir() || !snapshots_dir.is_dir() {
547 return;
548 }
549
550 // Collect blob filenames (excluding .lock files) for lookup.
551 let blob_names: Vec<String> = match fs::read_dir(&blobs_dir) {
552 Ok(entries) => entries
553 .flatten()
554 .filter_map(|e| {
555 let name = e.file_name().to_string_lossy().to_string();
556 if name.ends_with(".lock") {
557 None
558 } else if e.path().is_file() {
559 Some(name)
560 } else {
561 None
562 }
563 })
564 .collect(),
565 Err(_) => return,
566 };
567
568 if blob_names.is_empty() {
569 return;
570 }
571
572 // Walk each snapshot revision directory.
573 let snapshot_revs: Vec<(String, PathBuf)> = match fs::read_dir(&snapshots_dir) {
574 Ok(entries) => entries
575 .flatten()
576 .filter(|e| e.path().is_dir())
577 .map(|e| {
578 let rev = e.file_name().to_string_lossy().to_string();
579 (rev, e.path())
580 })
581 .collect(),
582 Err(_) => return,
583 };
584
585 for (rev, snap_dir) in &snapshot_revs {
586 let entries = match fs::read_dir(snap_dir) {
587 Ok(e) => e,
588 Err(_) => continue,
589 };
590
591 for entry in entries.flatten() {
592 let snap_file = entry.path();
593 let file_name = entry.file_name().to_string_lossy().to_string();
594
595 // Check if the file is accessible (broken symlink → metadata fails).
596 let needs_repair = if snap_file.is_symlink() {
597 // Symlink exists but may be broken — check if target is readable.
598 fs::metadata(&snap_file).is_err()
599 } else if snap_file.exists() {
600 // Real file, no repair needed.
601 false
602 } else {
603 // Path doesn't exist at all (dangling entry in readdir).
604 true
605 };
606
607 if !needs_repair {
608 continue;
609 }
610
611 // Try to find a matching blob. For GGUF repos there's usually
612 // just one blob; for multi-file repos we match by the symlink
613 // target name if we can read it, otherwise fall back to the
614 // single-blob heuristic.
615 let blob_name = if let Ok(target) = fs::read_link(&snap_file) {
616 target
617 .file_name()
618 .map(|n| n.to_string_lossy().to_string())
619 .filter(|n| blob_names.contains(n))
620 } else {
621 None
622 };
623
624 let resolved_blob = if let Some(ref name) = blob_name {
625 Some(blobs_dir.join(name))
626 } else if blob_names.len() == 1 {
627 // Single-blob repo (typical for GGUF) — safe to assume it
628 // corresponds to this snapshot entry.
629 Some(blobs_dir.join(&blob_names[0]))
630 } else {
631 None
632 };
633
634 if let Some(blob_path) = resolved_blob {
635 if !blob_path.is_file() {
636 continue;
637 }
638
639 // Remove the broken symlink before copying.
640 if snap_file.is_symlink() || snap_file.exists() {
641 let _ = fs::remove_file(&snap_file);
642 }
643
644 info!(
645 "repair_hf_cache_symlinks: replacing broken symlink with copy: {} → {}",
646 blob_path.display(),
647 snap_file.display()
648 );
649
650 // Prefer hard link (instant, no extra disk space) over copy.
651 if fs::hard_link(&blob_path, &snap_file).is_err() {
652 if let Err(e) = fs::copy(&blob_path, &snap_file) {
653 warn!(
654 "repair_hf_cache_symlinks: failed to copy blob to snapshot: {}",
655 e
656 );
657 }
658 }
659 } else {
660 warn!(
661 "repair_hf_cache_symlinks: broken symlink for '{}' but no matching blob found",
662 file_name
663 );
664 }
665 }
666
667 // Ensure refs/main exists and points to this revision so that
668 // `hf-hub` can resolve the snapshot.
669 if !refs_dir.join("main").is_file() {
670 let _ = fs::create_dir_all(&refs_dir);
671 if let Err(e) = fs::write(refs_dir.join("main"), rev.as_bytes()) {
672 warn!("repair_hf_cache_symlinks: failed to write refs/main: {}", e);
673 } else {
674 info!("repair_hf_cache_symlinks: created refs/main → {}", rev);
675 }
676 }
677 }
678 }
679
680 // ---------------------------------------------------------------------------
681 // Query / mutation functions
682 // ---------------------------------------------------------------------------
683
684 /// Scan the local HuggingFace hub cache and return all downloaded models that
685 /// the inference engine supports.
686 ///
687 /// The cache directory is typically `~/.cache/huggingface/hub/` and contains
688 /// sub-directories named `models--{org}--{model}` for each cached model.
689 pub fn list_local_hf_models() -> LocalHfModelsResponse {
690 let cache_dir = match hf_cache_dir() {
691 Some(dir) => dir,
692 None => {
693 debug!("HuggingFace cache directory not found.");
694 return LocalHfModelsResponse {
695 models: Vec::new(),
696 cache_path: String::new(),
697 total_size_bytes: 0,
698 total_size_display: "0 B".to_string(),
699 };
700 }
701 };
702
703 let cache_path_str = cache_dir.to_string_lossy().to_string();
704 debug!("Scanning HuggingFace cache at: {}", cache_path_str);
705
706 let mut models: Vec<LocalHfModel> = Vec::new();
707
708 let entries = match fs::read_dir(&cache_dir) {
709 Ok(entries) => entries,
710 Err(e) => {
711 warn!("Failed to read HuggingFace cache directory: {:?}", e);
712 return LocalHfModelsResponse {
713 models: Vec::new(),
714 cache_path: cache_path_str,
715 total_size_bytes: 0,
716 total_size_display: "0 B".to_string(),
717 };
718 }
719 };
720
721 for entry in entries.flatten() {
722 let entry_path = entry.path();
723 if !entry_path.is_dir() {
724 continue;
725 }
726
727 let dir_name = match entry.file_name().to_str() {
728 Some(name) => name.to_string(),
729 None => continue,
730 };
731
732 // HuggingFace hub caches model repos as "models--{org}--{name}"
733 if !dir_name.starts_with("models--") {
734 continue;
735 }
736
737 // Strip the "models--" prefix and split by "--" to get org/name
738 let remainder = &dir_name["models--".len()..];
739 let parts: Vec<&str> = remainder.splitn(2, "--").collect();
740 if parts.len() != 2 {
741 // Some models have no org (single segment), handle gracefully
742 let model_id = remainder.replace("--", "/");
743 let size = dir_size(&entry_path);
744 let revisions = list_revisions(&entry_path);
745
746 models.push(LocalHfModel {
747 model_id: model_id.clone(),
748 org: String::new(),
749 name: model_id,
750 size_bytes: size,
751 size_display: format_size(size),
752 path: entry_path.to_string_lossy().to_string(),
753 revisions,
754 });
755 continue;
756 }
757
758 let org = parts[0].to_string();
759 let name = parts[1].to_string();
760 let model_id = format!("{}/{}", org, name);
761 let size = dir_size(&entry_path);
762 let revisions = list_revisions(&entry_path);
763
764 models.push(LocalHfModel {
765 model_id,
766 org,
767 name,
768 size_bytes: size,
769 size_display: format_size(size),
770 path: entry_path.to_string_lossy().to_string(),
771 revisions,
772 });
773 }
774
775 // Only keep models that the inference engine actually supports.
776 models.retain(|m| SUPPORTED_MODELS.iter().any(|&s| s == m.model_id));
777
778 // Sort by model_id for a stable, predictable order.
779 models.sort_by_key(|a| a.model_id.to_lowercase());
780
781 let total_size_bytes: u64 = models.iter().map(|m| m.size_bytes).sum();
782
783 debug!(
784 "Found {} local HuggingFace model(s), total size: {}",
785 models.len(),
786 format_size(total_size_bytes)
787 );
788
789 LocalHfModelsResponse {
790 models,
791 cache_path: cache_path_str,
792 total_size_bytes,
793 total_size_display: format_size(total_size_bytes),
794 }
795 }
796
797 /// Delete a locally cached HuggingFace model by removing its directory from
798 /// the hub cache.
799 ///
800 /// `model_id` should be the full identifier, e.g. "bartowski/Qwen2.5-1.5B-Instruct-GGUF".
801 /// Returns `Ok(())` on success or an error string describing what went wrong.
802 pub fn delete_local_hf_model(model_id: String) -> Result<(), String> {
803 let cache_dir = hf_cache_dir().ok_or("HuggingFace cache directory not found.".to_string())?;
804
805 // Convert model_id (e.g. "org/name") to the cache directory name ("models--org--name")
806 let dir_name = format!("models--{}", model_id.replace('/', "--"));
807 let model_path = cache_dir.join(&dir_name);
808
809 if !model_path.exists() {
810 return Err(format!(
811 "Model cache directory not found: {}",
812 model_path.display()
813 ));
814 }
815
816 if !model_path.is_dir() {
817 return Err(format!(
818 "Expected a directory but found a file: {}",
819 model_path.display()
820 ));
821 }
822
823 debug!(
824 "Deleting local HuggingFace model: {} at {:?}",
825 model_id, model_path
826 );
827
828 fs::remove_dir_all(&model_path).map_err(|e| {
829 let msg = format!("Failed to delete model {}: {}", model_id, e);
830 warn!("{}", msg);
831 msg
832 })?;
833
834 debug!("Successfully deleted local model: {}", model_id);
835 Ok(())
836 }
837
838 /// Return the full list of models the inference engine supports, together with
839 /// flags indicating whether each one is fully downloaded, partially downloaded,
840 /// or not present at all.
841 pub fn list_supported_hf_models() -> SupportedHfModelsResponse {
842 let models = SUPPORTED_MODEL_INFO
843 .iter()
844 .map(|info| {
845 let cache_path = model_cache_path(info.id);
846 let local_size = cache_path
847 .as_ref()
848 .filter(|p| p.exists())
849 .map(dir_size)
850 .unwrap_or(0);
851
852 let has_cache = local_size > 0;
853 let is_complete = info.expected_size_bytes > 0
854 && local_size as f64
855 >= info.expected_size_bytes as f64 * DOWNLOAD_COMPLETE_THRESHOLD;
856
857 SupportedHfModel {
858 model_id: info.id.to_string(),
859 name: info.name.to_string(),
860 org: info.org.to_string(),
861 description: info.description.to_string(),
862 is_downloaded: has_cache && is_complete,
863 is_incomplete: has_cache && !is_complete,
864 local_size_bytes: local_size,
865 local_size_display: format_size(local_size),
866 expected_size_bytes: info.expected_size_bytes,
867 expected_size_display: format_size(info.expected_size_bytes),
868 }
869 })
870 .collect();
871
872 SupportedHfModelsResponse { models }
873 }
874
875 // ---------------------------------------------------------------------------
876 // Model download — platform-specific implementations
877 // ---------------------------------------------------------------------------
878 //
879 // The download function accepts a generic progress callback so that callers
880 // can forward progress to any event system without coupling to this crate.
881 //
882 // Callback signature: `Fn(ModelDownloadProgress) + Send + Sync + Clone + 'static`
883 // - Called every ~2 s from a native monitor thread.
884 // - Called once with `done: true` after the build completes.
885
886 // ── Apple (macOS + iOS): Metal-accelerated, all model types ─────────────────
887
888 #[cfg(any(
889 target_os = "macos",
890 target_os = "ios",
891 target_os = "tvos",
892 target_os = "visionos",
893 target_os = "watchos"
894 ))]
895 pub async fn download_model<F>(
896 model_id: String,
897 on_progress: F,
898 app_data_dir: Option<PathBuf>,
899 ) -> Result<(), String>
900 where
901 F: Fn(ModelDownloadProgress) + Send + Sync + Clone + 'static,
902 {
903 use crate::inference::models::{
904 BARTOWSKI_QWEN25_1_5B_INSTRUCT_GGUF, BARTOWSKI_QWEN25_3B_INSTRUCT_GGUF,
905 QWEN25_1_5B_GGUF_FILE, QWEN25_3B_GGUF_FILE,
906 };
907 use crate::inference::token::hf_token_source;
908 use mistralrs::GgufModelBuilder;
909 use std::sync::{
910 atomic::{AtomicBool, Ordering},
911 Arc,
912 };
913 use std::time::Duration;
914
915 // When app_data_dir is provided, seed the HF env vars so that hf-hub,
916 // mistral.rs, and hf_cache_dir() all write inside the app's writable
917 // container. This is required on every sandboxed platform:
918 // - iOS: ~/.cache is outside the container (os error 1).
919 // - macOS App Store: the app is sandboxed, ~/.cache is inaccessible.
920 // It is harmless on non-sandboxed desktop builds (the cache simply
921 // lives under the app data dir instead of ~/.cache).
922 if let Some(ref data_dir) = app_data_dir {
923 let hf_home = data_dir.join("models");
924 let hf_hub_cache = hf_home.join("hub");
925 fs::create_dir_all(&hf_hub_cache)
926 .map_err(|e| format!("Cannot create HF cache dir: {e}"))?;
927 std::env::set_var("HF_HUB_CACHE", &hf_hub_cache);
928 std::env::set_var("HF_HOME", &hf_home);
929 info!(
930 "HF hub cache resolved to app data path: {}",
931 hf_hub_cache.display()
932 );
933 }
934
935 if !SUPPORTED_MODELS.contains(&model_id.as_str()) {
936 return Err(format!("Model {} is not supported.", model_id));
937 }
938
939 let expected_size = SUPPORTED_MODEL_INFO
940 .iter()
941 .find(|m| m.id == model_id)
942 .map(|m| m.expected_size_bytes)
943 .unwrap_or(0);
944
945 info!(
946 "Starting download for model: {} (expected ~{})",
947 model_id,
948 format_size(expected_size)
949 );
950
951 clean_stale_lock_files(&model_id);
952 repair_hf_cache_symlinks(&model_id);
953 ensure_model_cache_dir(&model_id);
954
955 // Emit an initial progress event based on the actual cache size so that
956 // resumed downloads don't visually restart from 0%.
957 let initial_bytes = model_cache_path(&model_id)
958 .filter(|p| p.exists())
959 .map(|p| dir_size(&p))
960 .unwrap_or(0);
961 on_progress(make_progress(
962 &model_id,
963 initial_bytes,
964 expected_size,
965 false,
966 ));
967
968 // Shared flag so the monitor task knows when to stop.
969 let finished = Arc::new(AtomicBool::new(false));
970 let finished_clone = Arc::clone(&finished);
971 let monitor_model_id = model_id.clone();
972 let monitor_cb = on_progress.clone();
973
974 // Spawn a native OS thread to poll the cache directory size every 2 seconds.
975 // We use std::thread instead of tokio::spawn because the model builder's
976 // .build().await internally calls the synchronous hf-hub API (api.get()),
977 // which blocks the tokio worker thread. This starves any tokio-spawned
978 // monitor tasks, preventing progress updates. A native thread runs
979 // independently of the tokio runtime and is immune to this starvation.
980 let monitor_handle = std::thread::spawn(move || loop {
981 std::thread::sleep(Duration::from_secs(2));
982
983 if finished_clone.load(Ordering::Relaxed) {
984 break;
985 }
986
987 let current_bytes = model_cache_path(&monitor_model_id)
988 .filter(|p| p.exists())
989 .map(|p| dir_size(&p))
990 .unwrap_or(0);
991
992 monitor_cb(make_progress(
993 &monitor_model_id,
994 current_bytes,
995 expected_size,
996 false,
997 ));
998 });
999
1000 // Build (and immediately drop) the model using the GGUF chat builder.
1001 // Building triggers the HuggingFace hub download.
1002 let result: Result<mistralrs::Model, anyhow::Error> = match model_id.as_str() {
1003 BARTOWSKI_QWEN25_1_5B_INSTRUCT_GGUF => {
1004 GgufModelBuilder::new(&model_id, vec![QWEN25_1_5B_GGUF_FILE])
1005 .with_token_source(hf_token_source())
1006 .build()
1007 .await
1008 }
1009 BARTOWSKI_QWEN25_3B_INSTRUCT_GGUF => {
1010 GgufModelBuilder::new(&model_id, vec![QWEN25_3B_GGUF_FILE])
1011 .with_token_source(hf_token_source())
1012 .build()
1013 .await
1014 }
1015 _ => Err(anyhow::anyhow!(
1016 "Model '{}' is not part of the chat-only release.",
1017 model_id
1018 )),
1019 };
1020
1021 finished.store(true, Ordering::Relaxed);
1022 let _ = monitor_handle.join();
1023
1024 match result {
1025 Ok(_model) => {
1026 let final_bytes = model_cache_path(&model_id)
1027 .filter(|p| p.exists())
1028 .map(|p| dir_size(&p))
1029 .unwrap_or(expected_size);
1030 on_progress(make_progress(&model_id, final_bytes, final_bytes, true));
1031 info!("Successfully downloaded model: {}", model_id);
1032 Ok(())
1033 }
1034 Err(e) => {
1035 let msg = format!("Failed to download model {}: {}", model_id, e);
1036 warn!("{}", msg);
1037 Err(msg)
1038 }
1039 }
1040 }
1041
1042 // ── Android: GGUF-only download ───────────────────────────────────────────
1043 //
1044 // On Android there is no Metal GPU, so only pre-quantized GGUF models
1045 // (CPU inference via candle) are supported. We mirror the Apple
1046 // implementation but restrict to GGUF presets and skip all Metal / ISQ
1047 // model builders.
1048 //
1049 // `app_data_dir` is **required** on Android: `dirs::home_dir()` returns
1050 // `None` inside the Android sandbox, which causes `hf_hub::Cache::default()`
1051 // to panic. The caller resolves the app data directory and passes it here.
1052
1053 #[cfg(target_os = "android")]
1054 pub async fn download_model<F>(
1055 model_id: String,
1056 on_progress: F,
1057 app_data_dir: Option<PathBuf>,
1058 ) -> Result<(), String>
1059 where
1060 F: Fn(ModelDownloadProgress) + Send + Sync + Clone + 'static,
1061 {
1062 use crate::inference::models::{
1063 BARTOWSKI_QWEN25_1_5B_INSTRUCT_GGUF, BARTOWSKI_QWEN25_3B_INSTRUCT_GGUF,
1064 QWEN25_1_5B_GGUF_FILE, QWEN25_1_5B_TOK_MODEL_ID, QWEN25_3B_GGUF_FILE,
1065 QWEN25_3B_TOK_MODEL_ID,
1066 };
1067 use crate::inference::token::hf_token_source;
1068 use hf_hub::Cache;
1069 use mistralrs::GgufModelBuilder;
1070 use mistralrs_core::GLOBAL_HF_CACHE;
1071 use std::sync::{
1072 atomic::{AtomicBool, Ordering},
1073 Arc,
1074 };
1075 use std::time::Duration;
1076
1077 let (gguf_file, tok_model_id) = match model_id.as_str() {
1078 id if id == BARTOWSKI_QWEN25_1_5B_INSTRUCT_GGUF => {
1079 (QWEN25_1_5B_GGUF_FILE, QWEN25_1_5B_TOK_MODEL_ID)
1080 }
1081 id if id == BARTOWSKI_QWEN25_3B_INSTRUCT_GGUF => {
1082 (QWEN25_3B_GGUF_FILE, QWEN25_3B_TOK_MODEL_ID)
1083 }
1084 _ => {
1085 return Err(format!(
1086 "Model '{}' is not supported on Android. \
1087 Only Qwen 2.5 1.5B or 3B (GGUF) are available.",
1088 model_id
1089 ));
1090 }
1091 };
1092
1093 // ── Resolve the HF cache path ─────────────────────────────────────────
1094 //
1095 // `app_data_dir` is required on Android — `dirs::home_dir()` returns None
1096 // in the Android sandbox and causes Cache::default() to panic.
1097 let resolved_app_data = app_data_dir
1098 .ok_or_else(|| "app_data_dir is required on Android for HF cache resolution".to_string())?;
1099
1100 let hf_home = resolved_app_data.join("models");
1101 let hf_hub_cache = hf_home.join("hub");
1102
1103 fs::create_dir_all(&hf_hub_cache)
1104 .map_err(|e| format!("Failed to create HF hub cache dir: {e}"))?;
1105
1106 // Seed GLOBAL_HF_CACHE so get_paths_gguf! never falls back to Cache::default().
1107 GLOBAL_HF_CACHE.get_or_init(|| Cache::new(hf_hub_cache.clone()));
1108
1109 // Keep env vars in sync so hf_hub_cache_dir() and hf_cache_dir() agree.
1110 std::env::set_var("HF_HUB_CACHE", &hf_hub_cache);
1111 std::env::set_var("HF_HOME", &hf_home);
1112
1113 info!(
1114 "Android HF hub cache resolved to: {}",
1115 hf_hub_cache.display()
1116 );
1117
1118 let expected_size = SUPPORTED_MODEL_INFO
1119 .iter()
1120 .find(|m| m.id == model_id)
1121 .map(|m| m.expected_size_bytes)
1122 .unwrap_or(0);
1123
1124 info!(
1125 "Starting Android download for model: {} (expected ~{})",
1126 model_id,
1127 format_size(expected_size)
1128 );
1129
1130 clean_stale_lock_files(&model_id);
1131 ensure_model_cache_dir(&model_id);
1132
1133 let initial_bytes = model_cache_path(&model_id)
1134 .filter(|p| p.exists())
1135 .map(|p| dir_size(&p))
1136 .unwrap_or(0);
1137 on_progress(make_progress(
1138 &model_id,
1139 initial_bytes,
1140 expected_size,
1141 false,
1142 ));
1143
1144 // Native monitor thread — same reasoning as the Apple implementation.
1145 let finished = Arc::new(AtomicBool::new(false));
1146 let finished_clone = Arc::clone(&finished);
1147 let monitor_model_id = model_id.clone();
1148 let monitor_cb = on_progress.clone();
1149
1150 let monitor_handle = std::thread::spawn(move || loop {
1151 std::thread::sleep(Duration::from_secs(2));
1152 if finished_clone.load(Ordering::Relaxed) {
1153 break;
1154 }
1155 let current_bytes = model_cache_path(&monitor_model_id)
1156 .filter(|p| p.exists())
1157 .map(|p| dir_size(&p))
1158 .unwrap_or(0);
1159 monitor_cb(make_progress(
1160 &monitor_model_id,
1161 current_bytes,
1162 expected_size,
1163 false,
1164 ));
1165 });
1166
1167 let result = GgufModelBuilder::new(&model_id, vec![gguf_file])
1168 .with_tok_model_id(tok_model_id)
1169 .with_token_source(hf_token_source())
1170 .with_logging()
1171 .build()
1172 .await;
1173
1174 finished.store(true, Ordering::Relaxed);
1175 let _ = monitor_handle.join();
1176
1177 match result {
1178 Ok(_model) => {
1179 let final_bytes = model_cache_path(&model_id)
1180 .filter(|p| p.exists())
1181 .map(|p| dir_size(&p))
1182 .unwrap_or(expected_size);
1183 on_progress(make_progress(&model_id, final_bytes, final_bytes, true));
1184 info!("Successfully downloaded model on Android: {}", model_id);
1185 Ok(())
1186 }
1187 Err(e) => {
1188 let msg = format!("Failed to download model {}: {}", model_id, e);
1189 warn!("{}", msg);
1190 Err(msg)
1191 }
1192 }
1193 }
1194
1195 // ── Other platforms: no-op ────────────────────────────────────────────────
1196
1197 #[cfg(not(any(
1198 target_os = "macos",
1199 target_os = "ios",
1200 target_os = "tvos",
1201 target_os = "visionos",
1202 target_os = "watchos",
1203 target_os = "android"
1204 )))]
1205 pub async fn download_model<F>(
1206 model_id: String,
1207 _on_progress: F,
1208 _app_data_dir: Option<PathBuf>,
1209 ) -> Result<(), String>
1210 where
1211 F: Fn(ModelDownloadProgress) + Send + Sync + Clone + 'static,
1212 {
1213 debug!(
1214 "download_model: not supported on this platform, ignoring request for {}.",
1215 model_id
1216 );
1217 Err(format!(
1218 "Model downloads are not supported on this platform (requested: {}).",
1219 model_id
1220 ))
1221 }