@sigit / siti / commits / d08ce59

Add support for UQFF models

paydii committed Aug 10, 2026 at 22:01 UTC d08ce593ed736287dfb9a91d5ef226d9659c0ce1
6 files changed +301 -13
src-tauri/src/chat/command_download_progress.rs new
+41
@@ -0,0 +1,41 @@
1 +//! Tauri command `chat_download_progress`: bytes fetched so far for a model.
2 +//!
3 +//! Weight downloads run inside mistral.rs, which reports progress only to an
4 +//! `indicatif` bar destined for the log — nothing the WebView can observe. The
5 +//! HF cache directory is therefore the one shared surface between the transfer
6 +//! and the UI, so Settings polls this while a model is loading to turn an
7 +//! indeterminate spinner into a real percentage.
8 +//!
9 +//! A 646 MB UQFF model on a phone hotspot takes ten minutes or more; without
10 +//! this the panel is indistinguishable from a hang for that whole time.
11 +
12 +/// Bytes of `model_id`'s weights present in the local HF cache, counting the
13 +/// partially written file of a download still in flight.
14 +///
15 +/// The caller already knows the model's expected total (`ModelInfo.size_bytes`)
16 +/// and computes the percentage from it, so this stays a single number rather
17 +/// than duplicating the size table on both sides of the IPC boundary.
18 +///
19 +/// Returns `0` for an unknown id or an untouched cache, which the frontend
20 +/// renders as an indeterminate state — the same thing it showed before.
21 +#[cfg(any(
22 + target_os = "macos",
23 + target_os = "ios",
24 + target_os = "android",
25 + target_os = "windows"
26 +))]
27 +#[tauri::command]
28 +pub async fn chat_download_progress(model_id: String) -> u64 {
29 + super::model_downloaded_bytes(&model_id)
30 +}
31 +
32 +#[cfg(not(any(
33 + target_os = "macos",
34 + target_os = "ios",
35 + target_os = "android",
36 + target_os = "windows"
37 +)))]
38 +#[tauri::command]
39 +pub async fn chat_download_progress(_model_id: String) -> u64 {
40 + 0
41 +}
src-tauri/src/chat/mod.rs
+135 -8
@@ -13,6 +13,7 @@
13 // ── Command submodules (one Tauri command per file) ──────────────────────────
14
15 pub mod command_clear_history;
16 +pub mod command_download_progress;
17 pub mod command_get_history;
18 pub mod command_get_status;
19 pub mod command_list_models;
@@ -25,6 +26,7 @@ pub mod command_unload_model;
26 // ── Re-exports so lib.rs can pull in commands with a flat path ───────────────
27
28 pub use command_clear_history::chat_clear_history;
29 +pub use command_download_progress::chat_download_progress;
30 pub use command_get_history::chat_get_history;
31 pub use command_get_status::chat_get_status;
32 pub use command_list_models::chat_list_models;
@@ -229,8 +231,12 @@ pub(crate) fn config_for_model_id(id: &str) -> Option<GgufModelConfig> {
231 // UQFF (Universal Quantized File Format) stores pre-quantised weights and loads
232 // directly through mistral.rs's `UqffTextModelBuilder`, avoiding the ISQ path's
233 // full-precision download + in-memory quantisation spike. onde exposes it via
232 -// `ChatEngine::load_uqff_model`, which is available on every inference platform
233 -// (macOS/iOS/Android), so — unlike Gemma ISQ — these are offered on all three.
234 +// `ChatEngine::load_uqff_model` on every platform its engine builds for — the
235 +// `q4k` shards are a GGML-family type that dequantises on CPU as well as Metal
236 +// — so, unlike Gemma ISQ, these are offered everywhere Siti does inference
237 +// (macOS/iOS/Android/Windows). Keep these `cfg`s in step with the rest of the
238 +// module: a narrower gate here breaks the Windows build, because
239 +// `resolve_model_id` and `chat_list_models` reference the table unconditionally.
240 //
241 // We ship the `mistralrs-community` Qwen 3 UQFF repos: each is self-contained
242 // (base `config.json` + tokenizer + `residual.safetensors` + the `q4k` shard),
@@ -240,7 +246,12 @@ pub(crate) fn config_for_model_id(id: &str) -> Option<GgufModelConfig> {
246 // `expected_size_bytes` counts just those, matching the on-disk footprint.
247
248 /// Static metadata for a Qwen 3 UQFF model offered in Siti's model list.
243 -#[cfg(any(target_os = "macos", target_os = "ios", target_os = "android"))]
249 +#[cfg(any(
250 + target_os = "macos",
251 + target_os = "ios",
252 + target_os = "android",
253 + target_os = "windows"
254 +))]
255 pub(crate) struct UqffModelEntry {
256 /// HuggingFace repo id; also the value passed back to `chat_set_model`.
257 pub id: &'static str,
@@ -259,7 +270,12 @@ pub(crate) struct UqffModelEntry {
270 /// The Qwen 3 UQFF models Siti can load. Sizes are the measured
271 /// `q4k-0.uqff` + `residual.safetensors` totals of each `mistralrs-community`
272 /// repo (small config/tokenizer files add a negligible remainder).
262 -#[cfg(any(target_os = "macos", target_os = "ios", target_os = "android"))]
273 +#[cfg(any(
274 + target_os = "macos",
275 + target_os = "ios",
276 + target_os = "android",
277 + target_os = "windows"
278 +))]
279 pub(crate) const UQFF_MODELS: &[UqffModelEntry] = &[
280 UqffModelEntry {
281 id: "mistralrs-community/Qwen3-0.6B-UQFF",
@@ -332,7 +348,12 @@ pub(crate) const UQFF_MODELS: &[UqffModelEntry] = &[
348 /// Every `mistralrs-community` Qwen 3 UQFF repo names its 4-bit shard
349 /// `q4k-0.uqff`; passing that single shard is enough for mistral.rs to resolve
350 /// the base config, tokenizer, and residual weights.
335 -#[cfg(any(target_os = "macos", target_os = "ios", target_os = "android"))]
351 +#[cfg(any(
352 + target_os = "macos",
353 + target_os = "ios",
354 + target_os = "android",
355 + target_os = "windows"
356 +))]
357 pub(crate) fn uqff_config_for_model_id(id: &str) -> Option<UqffModelConfig> {
358 let entry = UQFF_MODELS.iter().find(|m| m.id == id)?;
359 Some(UqffModelConfig {
@@ -476,19 +497,40 @@ const DOWNLOAD_COMPLETE_THRESHOLD: f64 = 0.99;
497 /// Only used against the HF cache's `blobs/` directory, which holds the actual
498 /// downloaded files (not the `snapshots/` symlink/hard-link views), so nothing
499 /// is double-counted.
500 +///
501 +/// `include_incomplete` decides how an in-flight download is treated. hf-hub
502 +/// streams each blob into a sibling `<sha>.part` file and only renames it to the
503 +/// final blob name once the transfer completes, holding a `<sha>.lock` alongside
504 +/// it meanwhile. Those two answer different questions:
505 +///
506 +/// - **Completeness** (`false`) must ignore them, or a download interrupted far
507 +/// enough into the last file reads as finished — for a two-file UQFF model,
508 +/// one finished shard plus a nearly-full `.part` clears the 99% threshold
509 +/// while the model is still unloadable.
510 +/// - **Progress** (`true`) wants exactly the opposite: the `.part` file *is* the
511 +/// live byte count, and excluding it would peg the UI at 0% for the entire
512 +/// download.
513 #[cfg(any(
514 target_os = "macos",
515 target_os = "ios",
516 target_os = "android",
517 target_os = "windows"
518 ))]
485 -fn blobs_dir_size(path: &std::path::Path) -> u64 {
519 +fn blobs_dir_size(path: &std::path::Path, include_incomplete: bool) -> u64 {
520 let mut total = 0;
521 if let Ok(entries) = std::fs::read_dir(path) {
522 for entry in entries.flatten() {
523 let p = entry.path();
524 + if !include_incomplete
525 + && matches!(
526 + p.extension().and_then(|e| e.to_str()),
527 + Some("part") | Some("lock")
528 + )
529 + {
530 + continue;
531 + }
532 match entry.metadata() {
491 - Ok(md) if md.is_dir() => total += blobs_dir_size(&p),
533 + Ok(md) if md.is_dir() => total += blobs_dir_size(&p, include_incomplete),
534 Ok(md) => total += md.len(),
535 Err(_) => {}
536 }
@@ -515,13 +557,34 @@ pub(crate) fn is_model_downloaded(id: &str, expected_size_bytes: u64) -> bool {
557 match onde::hf_cache::model_cache_path(id) {
558 Some(root) => {
559 let blobs = root.join("blobs");
518 - blobs_dir_size(&blobs) as f64
560 + // Completed blobs only — a `.part` still being streamed must not
561 + // count towards "downloaded".
562 + blobs_dir_size(&blobs, false) as f64
563 >= expected_size_bytes as f64 * DOWNLOAD_COMPLETE_THRESHOLD
564 }
565 None => false,
566 }
567 }
568
569 +/// Bytes of `id`'s weights currently on disk, **including** the partially
570 +/// written `.part` file of an in-flight download.
571 +///
572 +/// Drives the download progress indicator in Settings. mistral.rs owns the
573 +/// actual transfer and only reports progress to an `indicatif` bar that goes to
574 +/// the log, so the cache directory is the one place the UI can observe it.
575 +#[cfg(any(
576 + target_os = "macos",
577 + target_os = "ios",
578 + target_os = "android",
579 + target_os = "windows"
580 +))]
581 +pub(crate) fn model_downloaded_bytes(id: &str) -> u64 {
582 + match onde::hf_cache::model_cache_path(id) {
583 + Some(root) => blobs_dir_size(&root.join("blobs"), true),
584 + None => 0,
585 + }
586 +}
587 +
588 /// Return the sampling config for the currently selected model.
589 ///
590 /// iOS and Android use conservative settings to bound KV cache memory and
@@ -621,3 +684,67 @@ pub(crate) fn emit_chat_status(
684 error!("Failed to emit chat status event: {:?}", e);
685 }
686 }
687 +
688 +// ── Tests ────────────────────────────────────────────────────────────────────
689 +
690 +#[cfg(all(
691 + test,
692 + any(
693 + target_os = "macos",
694 + target_os = "ios",
695 + target_os = "android",
696 + target_os = "windows"
697 + )
698 +))]
699 +mod tests {
700 + use super::*;
701 +
702 + /// A two-file model whose download died partway through the second file:
703 + /// one finished blob, plus the `.part`/`.lock` pair hf-hub leaves behind
704 + /// while streaming the next. Sized so the in-flight bytes land *above* the
705 + /// completeness threshold and the finished bytes alone land below it, which
706 + /// is precisely the case that used to report an unloadable model as ready.
707 + const FIXTURE_EXPECTED_TOTAL: u64 = 2000;
708 +
709 + fn blobs_fixture(name: &str) -> std::path::PathBuf {
710 + let dir = std::env::temp_dir().join(format!("siti-blobs-{name}"));
711 + let _ = std::fs::remove_dir_all(&dir);
712 + std::fs::create_dir_all(&dir).expect("create fixture dir");
713 + std::fs::write(dir.join("aaaa"), vec![0u8; 1000]).expect("write blob");
714 + std::fs::write(dir.join("bbbb.part"), vec![0u8; 990]).expect("write part");
715 + std::fs::write(dir.join("bbbb.lock"), vec![0u8; 10]).expect("write lock");
716 + dir
717 + }
718 +
719 + #[test]
720 + fn blobs_dir_size_excludes_incomplete_downloads() {
721 + let dir = blobs_fixture("complete");
722 + assert_eq!(blobs_dir_size(&dir, false), 1000);
723 + let _ = std::fs::remove_dir_all(&dir);
724 + }
725 +
726 + #[test]
727 + fn blobs_dir_size_counts_partial_for_progress() {
728 + let dir = blobs_fixture("progress");
729 + assert_eq!(blobs_dir_size(&dir, true), 2000);
730 + let _ = std::fs::remove_dir_all(&dir);
731 + }
732 +
733 + /// The regression this guards: one finished blob plus a nearly-complete
734 + /// `.part` clears the 99% threshold, so counting in-flight bytes towards
735 + /// completeness reports an unloadable model as downloaded.
736 + #[test]
737 + fn partial_download_does_not_read_as_complete() {
738 + let dir = blobs_fixture("threshold");
739 + let cutoff = FIXTURE_EXPECTED_TOTAL as f64 * DOWNLOAD_COMPLETE_THRESHOLD;
740 + assert!(
741 + blobs_dir_size(&dir, true) as f64 >= cutoff,
742 + "fixture must cross the threshold when in-flight bytes are counted"
743 + );
744 + assert!(
745 + (blobs_dir_size(&dir, false) as f64) < cutoff,
746 + "completed bytes alone must stay below the threshold"
747 + );
748 + let _ = std::fs::remove_dir_all(&dir);
749 + }
750 +}
src-tauri/src/lib.rs
+4 -2
@@ -11,8 +11,9 @@ pub mod inference {
11
12 use app_info::app_build_version;
13 use chat::{
14 - chat_clear_history, chat_get_history, chat_get_status, chat_list_models, chat_load_model,
15 - chat_remove_model, chat_send_message, chat_set_model, chat_unload_model,
14 + chat_clear_history, chat_download_progress, chat_get_history, chat_get_status,
15 + chat_list_models, chat_load_model, chat_remove_model, chat_send_message, chat_set_model,
16 + chat_unload_model,
17 };
18
19 #[cfg_attr(mobile, tauri::mobile_entry_point)]
@@ -51,6 +52,7 @@ pub fn run() {
52 chat_list_models,
53 chat_set_model,
54 chat_remove_model,
55 + chat_download_progress,
56 // App metadata
57 app_build_version,
58 ])
src/App.css
+25
@@ -632,6 +632,31 @@ body {
632 background: #e5484d;
633 }
634
635 +/* Weight-download progress. Determinate only: an indeterminate bar would say
636 + no more than the pulsing status dot already does. */
637 +.model-progress {
638 + margin-top: 10px;
639 + height: 4px;
640 + border-radius: 2px;
641 + /* `--border` rather than a dimmed `--text-secondary`: opacity here would
642 + cascade into the fill and mute it too. Tracks both colour schemes. */
643 + background: var(--border);
644 + overflow: hidden;
645 +}
646 +
647 +.model-progress-fill {
648 + height: 100%;
649 + border-radius: inherit;
650 + background: #ff9800;
651 + transition: width 0.4s ease-out;
652 +}
653 +
654 +@media (prefers-reduced-motion: reduce) {
655 + .model-progress-fill {
656 + transition: none;
657 + }
658 +}
659 +
660 /* ── Supported Models list ─────────────────────────────────────────────── */
661
662 .model-list-row {
src/Settings.tsx
+88 -3
@@ -1,6 +1,11 @@
1 import { useEffect, useMemo, useState } from "react";
2 import type { ChatStatus, ModelInfo } from "./api";
3 -import { formatSize, getAppVersion, getBuildVersion } from "./api";
3 +import {
4 + formatSize,
5 + getAppVersion,
6 + getBuildVersion,
7 + getDownloadProgress,
8 +} from "./api";
9 import { openUrl } from "@tauri-apps/plugin-opener";
10
11 interface SettingsProps {
@@ -62,9 +67,40 @@ function DownloadedIcon() {
67 );
68 }
69
65 -function statusLine(status: ChatStatus, error: string | null) {
70 +/** Live weight-download progress, or `null` when there is nothing to report. */
71 +interface DownloadProgress {
72 + /** 0–1, clamped. */
73 + fraction: number;
74 + downloaded: number;
75 + total: number;
76 +}
77 +
78 +/**
79 + * Treat the download as finished slightly before the byte counts match. The
80 + * small JSON assets (config, tokenizer) aren't in the model's expected size, and
81 + * the last blob is renamed off `.part` only once complete, so the tail of the
82 + * transfer is jumpy. Past this point the wait is the model loading into memory,
83 + * which is what the user should be told.
84 + */
85 +const DOWNLOAD_DONE_FRACTION = 0.999;
86 +
87 +function statusLine(
88 + status: ChatStatus,
89 + error: string | null,
90 + progress: DownloadProgress | null
91 +) {
92 switch (status) {
93 case "loading":
94 + if (progress && progress.fraction < DOWNLOAD_DONE_FRACTION) {
95 + const pct = Math.floor(progress.fraction * 100);
96 + // Byte detail is noise before the first chunk lands.
97 + const detail =
98 + progress.downloaded > 0
99 + ? ` · ${formatSize(progress.downloaded)} of ${formatSize(progress.total)}`
100 + : "";
101 + return { text: `Downloading… ${pct}%${detail}`, cls: "loading" };
102 + }
103 + if (progress) return { text: "Loading into memory…", cls: "loading" };
104 return { text: "Downloading / loading…", cls: "loading" };
105 case "ready":
106 return { text: "Ready", cls: "ready" };
@@ -91,7 +127,41 @@ export default function Settings({
127 [models]
128 );
129 const busy = status === "loading";
94 - const line = statusLine(status, statusError);
130 +
131 + // Poll the cache while a model loads so the indeterminate spinner becomes a
132 + // real percentage. mistral.rs owns the transfer and reports progress only to
133 + // the log, so bytes-on-disk is the only signal available to the UI.
134 + const [downloadedBytes, setDownloadedBytes] = useState<number | null>(null);
135 + const selectedId = selected?.id;
136 + useEffect(() => {
137 + if (!open || status !== "loading" || !selectedId) {
138 + setDownloadedBytes(null);
139 + return;
140 + }
141 + let active = true;
142 + const tick = async () => {
143 + const bytes = await getDownloadProgress(selectedId).catch(() => null);
144 + if (active && bytes !== null) setDownloadedBytes(bytes);
145 + };
146 + void tick();
147 + const timer = window.setInterval(tick, 1000);
148 + return () => {
149 + active = false;
150 + window.clearInterval(timer);
151 + };
152 + }, [open, status, selectedId]);
153 +
154 + const progress = useMemo<DownloadProgress | null>(() => {
155 + const total = selected?.size_bytes;
156 + if (status !== "loading" || !total || downloadedBytes === null) return null;
157 + return {
158 + fraction: Math.min(downloadedBytes / total, 1),
159 + downloaded: downloadedBytes,
160 + total,
161 + };
162 + }, [status, selected?.size_bytes, downloadedBytes]);
163 +
164 + const line = statusLine(status, statusError, progress);
165
166 const handleOpenUrl = (url: string) => (e: React.MouseEvent) => {
167 e.preventDefault();
@@ -168,6 +238,21 @@ export default function Settings({
238 <span className="model-status-dot" />
239 {line.text}
240 </div>
241 + {progress && progress.fraction < DOWNLOAD_DONE_FRACTION && (
242 + <div
243 + className="model-progress"
244 + role="progressbar"
245 + aria-valuenow={Math.floor(progress.fraction * 100)}
246 + aria-valuemin={0}
247 + aria-valuemax={100}
248 + aria-label="Model download progress"
249 + >
250 + <div
251 + className="model-progress-fill"
252 + style={{ width: `${progress.fraction * 100}%` }}
253 + />
254 + </div>
255 + )}
256 </div>
257 )}
258 </div>
src/api.ts
+8
@@ -68,6 +68,14 @@ export const getHistory = () =>
68
69 export const clearHistory = () => invoke<void>("chat_clear_history");
70
71 +/**
72 + * Bytes of a model's weights on disk so far, including a download still in
73 + * flight. Poll while the model is loading and compare against the model's
74 + * `size_bytes` to render real progress.
75 + */
76 +export const getDownloadProgress = (modelId: string) =>
77 + invoke<number>("chat_download_progress", { modelId });
78 +
79 // ── App metadata ──────────────────────────────────────────────────────────
80
81 /** Marketing version (CFBundleShortVersionString) from tauri.conf.json. */