From b84f167d5a187fd1c0c4bd948ea0014b07aa73f0 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 16 Mar 2026 05:53:34 +0000 Subject: [PATCH 01/19] Add download progress bar for non-GGUF models in Chat Previously only GGUF models showed download progress in Chat. Non-GGUF models (safetensors, bnb quantized, etc.) showed a static message with no progress indication. This adds progress tracking for all model types and fixes several related issues. Backend: - Add /api/models/download-progress endpoint that checks the HF cache blobs directory for completed and .incomplete files. Uses model_info() (cached per repo) to determine expected total size for percentage. - Add /api/models/cached-models endpoint that lists non-GGUF model repos from the HF cache via scan_cache_dir(). - Fix progress stuck at 0.99: when no .incomplete files remain, report 1.0 immediately (blob deduplication can make byte totals mismatch). Frontend: - Remove the ggufVariant gate so download progress polling works for all non-cached models, not just GGUFs. - Use GGUF-specific endpoint when variant + expectedBytes available, otherwise use the general download-progress endpoint. - Fix toast stuck after load: check loadingModelRef.current before and after the async poll to prevent overwriting the success toast. - First poll at 500ms instead of waiting for the 2s interval. - Show downloaded non-GGUF models in the Hub model picker "Downloaded" section alongside GGUFs. --- studio/backend/routes/models.py | 121 ++++++++++++++++++ .../assistant-ui/model-selector/pickers.tsx | 20 ++- .../src/features/chat/api/chat-api.ts | 19 +++ .../chat/hooks/use-chat-model-runtime.ts | 118 ++++++++++------- 4 files changed, 228 insertions(+), 50 deletions(-) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 586b18718e..71d4a5e104 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -686,6 +686,90 @@ async def get_gguf_download_progress( return {"downloaded_bytes": 0, "expected_bytes": expected_bytes, "progress": 0} +@router.get("/download-progress") +async def get_download_progress( + repo_id: str = Query(..., description = "HuggingFace repo ID"), + current_subject: str = Depends(get_current_subject), +): + """Return download progress for any HuggingFace model repo. + + Checks the local HF cache for completed blobs and in-progress + (.incomplete) downloads. Uses the HF API to determine the expected + total size on the first call, then caches it for subsequent polls. + """ + _empty = {"downloaded_bytes": 0, "expected_bytes": 0, "progress": 0} + try: + if not _is_valid_repo_id(repo_id): + return _empty + + from huggingface_hub import constants as hf_constants + + cache_dir = Path(hf_constants.HF_HUB_CACHE) + target = f"models--{repo_id.replace('/', '--')}".lower() + completed_bytes = 0 + in_progress_bytes = 0 + + for entry in cache_dir.iterdir(): + if entry.name.lower() != target: + continue + blobs_dir = entry / "blobs" + if not blobs_dir.is_dir(): + break + for f in blobs_dir.iterdir(): + if not f.is_file(): + continue + if f.name.endswith(".incomplete"): + in_progress_bytes += f.stat().st_size + else: + completed_bytes += f.stat().st_size + break + + downloaded_bytes = completed_bytes + in_progress_bytes + if downloaded_bytes == 0: + return _empty + + # Get expected size from HF API (cached per repo_id) + expected_bytes = _get_repo_size_cached(repo_id) + if expected_bytes <= 0: + # Cannot determine total; report bytes only, no percentage + return { + "downloaded_bytes": downloaded_bytes, + "expected_bytes": 0, + "progress": 0, + } + + # No .incomplete files means download is done regardless of byte totals + # (blob deduplication can make completed_bytes differ from expected_bytes) + if in_progress_bytes == 0 and completed_bytes > 0: + progress = 1.0 + else: + progress = min(downloaded_bytes / expected_bytes, 0.99) + return { + "downloaded_bytes": downloaded_bytes, + "expected_bytes": expected_bytes, + "progress": round(progress, 3), + } + except Exception: + return _empty + + +_repo_size_cache: dict[str, int] = {} + + +def _get_repo_size_cached(repo_id: str) -> int: + if repo_id in _repo_size_cache: + return _repo_size_cache[repo_id] + try: + from huggingface_hub import model_info as hf_model_info + + info = hf_model_info(repo_id, token = None, files_metadata = True) + total = sum(s.size for s in info.siblings if s.size) + _repo_size_cache[repo_id] = total + return total + except Exception: + return 0 + + @router.get("/cached-gguf") async def list_cached_gguf( current_subject: str = Depends(get_current_subject), @@ -733,6 +817,43 @@ async def list_cached_gguf( return {"cached": []} +@router.get("/cached-models") +async def list_cached_models( + current_subject: str = Depends(get_current_subject), +): + """List non-GGUF model repos that have been downloaded to the HF cache.""" + try: + from huggingface_hub import scan_cache_dir + + hf_cache = scan_cache_dir() + seen_lower: dict[str, dict] = {} + for repo_info in hf_cache.repos: + if repo_info.repo_type != "model": + continue + repo_id = repo_info.repo_id + if repo_id.upper().endswith("-GGUF"): + continue + total_size = sum( + f.size_on_disk + for rev in repo_info.revisions + for f in rev.files + ) + if total_size == 0: + continue + key = repo_id.lower() + existing = seen_lower.get(key) + if existing is None or total_size > existing["size_bytes"]: + seen_lower[key] = { + "repo_id": repo_id, + "size_bytes": total_size, + } + cached = sorted(seen_lower.values(), key = lambda c: c["repo_id"]) + return {"cached": cached} + except Exception as e: + logger.error(f"Error listing cached models: {e}", exc_info = True) + return {"cached": []} + + @router.get("/checkpoints", response_model = CheckpointListResponse) async def list_checkpoints( outputs_dir: str = Query( diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx index fe4226c4af..9ad33936c3 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -8,8 +8,8 @@ import { TooltipContent, TooltipTrigger, } from "@/components/ui/tooltip"; -import { listCachedGguf, listGgufVariants } from "@/features/chat/api/chat-api"; -import type { CachedGgufRepo } from "@/features/chat/api/chat-api"; +import { listCachedGguf, listCachedModels, listGgufVariants } from "@/features/chat/api/chat-api"; +import type { CachedGgufRepo, CachedModelRepo } from "@/features/chat/api/chat-api"; import type { GgufVariantDetail } from "@/features/chat/types/api"; import { usePlatformStore } from "@/config/env"; import { @@ -361,10 +361,12 @@ export function HubModelPicker({ // Track which GGUF repo is expanded for variant selection const [expandedGguf, setExpandedGguf] = useState(null); - // Cached (already downloaded) GGUF repos + // Cached (already downloaded) repos const [cachedGguf, setCachedGguf] = useState([]); + const [cachedModels, setCachedModels] = useState([]); useEffect(() => { listCachedGguf().then(setCachedGguf).catch(() => {}); + listCachedModels().then(setCachedModels).catch(() => {}); }, []); const recommendedIds = useMemo( @@ -472,7 +474,7 @@ export function HubModelPicker({
- {!showHfSection && cachedGguf.length > 0 ? ( + {!showHfSection && (cachedGguf.length > 0 || cachedModels.length > 0) ? ( <> Downloaded {cachedGguf.map((c) => ( @@ -489,6 +491,16 @@ export function HubModelPicker({ )}
))} + {cachedModels.map((c) => ( + onSelect(c.repo_id, { source: "hub", isLora: false, isDownloaded: true })} + vramStatus={null} + /> + ))} ) : null} diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts index be105e86be..01bc762d86 100644 --- a/studio/frontend/src/features/chat/api/chat-api.ts +++ b/studio/frontend/src/features/chat/api/chat-api.ts @@ -117,12 +117,31 @@ export async function getGgufDownloadProgress( return parseJsonOrThrow(response); } +export async function getDownloadProgress( + repoId: string, +): Promise<{ downloaded_bytes: number; expected_bytes: number; progress: number }> { + const params = new URLSearchParams({ repo_id: repoId }); + const response = await authFetch(`/api/models/download-progress?${params}`); + return parseJsonOrThrow(response); +} + export async function listCachedGguf(): Promise { const response = await authFetch("/api/models/cached-gguf"); const data = await parseJsonOrThrow<{ cached: CachedGgufRepo[] }>(response); return data.cached; } +export interface CachedModelRepo { + repo_id: string; + size_bytes: number; +} + +export async function listCachedModels(): Promise { + const response = await authFetch("/api/models/cached-models"); + const data = await parseJsonOrThrow<{ cached: CachedModelRepo[] }>(response); + return data.cached; +} + export async function listGgufVariants( repoId: string, hfToken?: string, diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index ebcbe8e966..ecdd901983 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -4,6 +4,7 @@ import { useCallback, useRef, useState } from "react"; import { toast } from "sonner"; import { + getDownloadProgress, getGgufDownloadProgress, getInferenceStatus, listLoras, @@ -326,57 +327,82 @@ export function useChatModelRuntime() { // Poll download progress for non-cached models let progressInterval: ReturnType | null = null; - if (!isDownloaded && ggufVariant) { + if (!isDownloaded) { const expectedBytes = typeof selection !== "string" ? selection.expectedBytes ?? 0 : 0; - if (expectedBytes > 0) { - progressInterval = setInterval(async () => { - if (abortCtrl.signal.aborted) { - if (progressInterval) clearInterval(progressInterval); - return; - } - try { - const prog = await getGgufDownloadProgress(modelId, ggufVariant ?? "", expectedBytes); - if (prog.progress > 0 && prog.progress < 1) { - const dlGb = prog.downloaded_bytes / (1024 ** 3); - const totalGb = prog.expected_bytes / (1024 ** 3); - const pct = Math.round(prog.progress * 100); - toast.loading( - `Downloading model… ${pct}%`, - { - id: toastId, - description: `${dlGb.toFixed(1)} / ${totalGb.toFixed(1)} GB`, - duration: 10000, - action: { - label: "Cancel", - onClick: () => { - abortCtrl.abort(); - setLoadingModel(null); - setLoadAbortController(null); - loadingModelRef.current = null; - loadAbortRef.current = null; - loadToastIdRef.current = null; - unloadModel({ model_path: modelId }).catch(() => {}); - clearCheckpoint(); - toast.dismiss(toastId); - toast.info("Model loading cancelled"); - }, - }, - }, - ); - } else if (prog.progress >= 1) { - toast.loading("Loading model…", { + + const cancelAction = { + label: "Cancel", + onClick: () => { + abortCtrl.abort(); + setLoadingModel(null); + setLoadAbortController(null); + loadingModelRef.current = null; + loadAbortRef.current = null; + loadToastIdRef.current = null; + unloadModel({ model_path: modelId }).catch(() => {}); + clearCheckpoint(); + toast.dismiss(toastId); + toast.info("Model loading cancelled"); + }, + }; + + const pollProgress = async () => { + // Stop if cancelled or if loading already finished + if (abortCtrl.signal.aborted || !loadingModelRef.current) { + if (progressInterval) clearInterval(progressInterval); + return; + } + try { + const prog = ggufVariant && expectedBytes > 0 + ? await getGgufDownloadProgress(modelId, ggufVariant, expectedBytes) + : await getDownloadProgress(modelId); + + // Re-check after await -- load may have finished while polling + if (!loadingModelRef.current) return; + + if (prog.downloaded_bytes > 0 && prog.progress > 0 && prog.progress < 1) { + const dlGb = prog.downloaded_bytes / (1024 ** 3); + const totalGb = prog.expected_bytes / (1024 ** 3); + const pct = Math.round(prog.progress * 100); + toast.loading( + `Downloading model... ${pct}%`, + { id: toastId, - description: "Download complete. Starting inference server…", + description: totalGb > 0 + ? `${dlGb.toFixed(1)} / ${totalGb.toFixed(1)} GB` + : `${dlGb.toFixed(1)} GB downloaded`, duration: 10000, - }); - if (progressInterval) clearInterval(progressInterval); - } - } catch { - // Ignore polling errors + action: cancelAction, + }, + ); + } else if (prog.downloaded_bytes > 0 && prog.expected_bytes === 0) { + const dlGb = prog.downloaded_bytes / (1024 ** 3); + toast.loading( + "Downloading model...", + { + id: toastId, + description: `${dlGb.toFixed(1)} GB downloaded`, + duration: 10000, + action: cancelAction, + }, + ); + } else if (prog.progress >= 1) { + toast.loading("Loading model...", { + id: toastId, + description: "Download complete. Loading into memory...", + duration: 10000, + }); + if (progressInterval) clearInterval(progressInterval); } - }, 2000); - } + } catch { + // Ignore polling errors + } + }; + + // First poll after 500ms, then every 2s + setTimeout(pollProgress, 500); + progressInterval = setInterval(pollProgress, 2000); } try { From e03a809994283fa6c7da85f9c6939c25b76c45b3 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 16 Mar 2026 05:55:40 +0000 Subject: [PATCH 02/19] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/routes/models.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 71d4a5e104..77bb1cf329 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -834,9 +834,7 @@ async def list_cached_models( if repo_id.upper().endswith("-GGUF"): continue total_size = sum( - f.size_on_disk - for rev in repo_info.revisions - for f in rev.files + f.size_on_disk for rev in repo_info.revisions for f in rev.files ) if total_size == 0: continue From 1471c63b9653d5910da04d82e9c1ab5502eca3cb Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 16 Mar 2026 06:04:32 +0000 Subject: [PATCH 03/19] Fix download progress bugs: false completion, stale UI, dedup Three fixes on top of the download progress feature: 1. Backend: Replace broken "no .incomplete = done" completion check with a 95% byte threshold. HF downloads files sequentially, so between files there are briefly no .incomplete files even though the download is far from done (e.g. Kimi-K2.5 reported "done" after downloading 22KB of config files out of 595GB). 2. Frontend: Track hasShownProgress flag. Only show "Download complete. Loading into memory..." if we actually displayed download progress before. For already-cached models where the first poll returns progress=1.0, this avoids the misleading "Download complete" message. 3. Frontend: Deduplicate recommended vs downloaded -- filter out models already in the "Downloaded" section. Cache the fetched lists at module level so re-mounting the popover does not flash an empty "Downloaded" section. --- studio/backend/routes/models.py | 9 ++++-- .../assistant-ui/model-selector/pickers.tsx | 28 ++++++++++++++----- .../chat/hooks/use-chat-model-runtime.ts | 12 ++++++-- 3 files changed, 36 insertions(+), 13 deletions(-) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 77bb1cf329..02a810997a 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -738,9 +738,12 @@ async def get_download_progress( "progress": 0, } - # No .incomplete files means download is done regardless of byte totals - # (blob deduplication can make completed_bytes differ from expected_bytes) - if in_progress_bytes == 0 and completed_bytes > 0: + # Use 95% threshold for completion (blob deduplication can make + # completed_bytes differ slightly from expected_bytes). + # Do NOT use "no .incomplete files" as a completion signal -- + # HF downloads files sequentially, so between files there are + # no .incomplete files even though the download is far from done. + if completed_bytes >= expected_bytes * 0.95: progress = 1.0 else: progress = min(downloaded_bytes / expected_bytes, 0.99) diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx index 9ad33936c3..b8b464c56b 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -340,6 +340,10 @@ function isGgufRepo(id: string): boolean { return id.toUpperCase().includes("-GGUF"); } +// Module-level caches so re-mounting the popover shows results instantly +let _cachedGgufCache: CachedGgufRepo[] = []; +let _cachedModelsCache: CachedModelRepo[] = []; + // ── Hub Model Picker ────────────────────────────────────────── export function HubModelPicker({ @@ -361,17 +365,27 @@ export function HubModelPicker({ // Track which GGUF repo is expanded for variant selection const [expandedGguf, setExpandedGguf] = useState(null); - // Cached (already downloaded) repos - const [cachedGguf, setCachedGguf] = useState([]); - const [cachedModels, setCachedModels] = useState([]); + // Cached (already downloaded) repos -- use module-level cache so + // re-mounting the popover does not flash an empty "Downloaded" section. + const [cachedGguf, setCachedGguf] = useState(_cachedGgufCache); + const [cachedModels, setCachedModels] = useState(_cachedModelsCache); useEffect(() => { - listCachedGguf().then(setCachedGguf).catch(() => {}); - listCachedModels().then(setCachedModels).catch(() => {}); + listCachedGguf().then((v) => { _cachedGgufCache = v; setCachedGguf(v); }).catch(() => {}); + listCachedModels().then((v) => { _cachedModelsCache = v; setCachedModels(v); }).catch(() => {}); }, []); + // Deduplicate: don't show downloaded models in the recommended list + const downloadedSet = useMemo(() => { + const s = new Set(); + for (const c of cachedGguf) s.add(c.repo_id); + for (const c of cachedModels) s.add(c.repo_id); + return s; + }, [cachedGguf, cachedModels]); + const recommendedIds = useMemo( - () => dedupe([...models.map((model) => model.id), value ?? ""]), - [models, value], + () => dedupe([...models.map((model) => model.id), value ?? ""]) + .filter((id) => !downloadedSet.has(id)), + [models, value, downloadedSet], ); const { paramCountById: recommendedParamCountById } = diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index ecdd901983..f1d0e16d0c 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -347,6 +347,8 @@ export function useChatModelRuntime() { }, }; + let hasShownProgress = false; + const pollProgress = async () => { // Stop if cancelled or if loading already finished if (abortCtrl.signal.aborted || !loadingModelRef.current) { @@ -361,7 +363,8 @@ export function useChatModelRuntime() { // Re-check after await -- load may have finished while polling if (!loadingModelRef.current) return; - if (prog.downloaded_bytes > 0 && prog.progress > 0 && prog.progress < 1) { + if (prog.progress > 0 && prog.progress < 1) { + hasShownProgress = true; const dlGb = prog.downloaded_bytes / (1024 ** 3); const totalGb = prog.expected_bytes / (1024 ** 3); const pct = Math.round(prog.progress * 100); @@ -376,7 +379,9 @@ export function useChatModelRuntime() { action: cancelAction, }, ); - } else if (prog.downloaded_bytes > 0 && prog.expected_bytes === 0) { + } else if (prog.downloaded_bytes > 0 && prog.expected_bytes === 0 && prog.progress === 0) { + // Have bytes but no total size -- show bytes only + hasShownProgress = true; const dlGb = prog.downloaded_bytes / (1024 ** 3); toast.loading( "Downloading model...", @@ -387,7 +392,8 @@ export function useChatModelRuntime() { action: cancelAction, }, ); - } else if (prog.progress >= 1) { + } else if (prog.progress >= 1 && hasShownProgress) { + // Only show "download complete" if we actually showed progress toast.loading("Loading model...", { id: toastId, description: "Download complete. Loading into memory...", From f20c7ca54ddd04b64caabfef938d70fbabd8a993 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 16 Mar 2026 06:09:38 +0000 Subject: [PATCH 04/19] Friendlier unsupported model errors, show estimated download size 1. Backend: When a model fails with "No config file found" or similar unsupported-model errors, wrap the message with "This model is not supported yet. Try a different model." instead of showing the raw Unsloth exception. 2. Frontend: Compute estimated download size from the HF search API's safetensors.parameters dtype breakdown (BF16=2B/param, I32=4B/param, F32=4B/param, etc.) and show it in the model picker instead of just the param count. For example, Kimi-K2.5 now shows "~554 GB" instead of "171B" (which was misleading since 171B params != 171GB download). --- studio/backend/routes/inference.py | 12 +++++++++- .../assistant-ui/model-selector/pickers.tsx | 9 +++++-- .../frontend/src/hooks/use-hf-model-search.ts | 24 ++++++++++++++++++- 3 files changed, 41 insertions(+), 4 deletions(-) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index c7fa9ed720..13062f8b29 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -288,7 +288,17 @@ async def load_model( raise except Exception as e: logger.error(f"Error loading model: {e}", exc_info = True) - raise HTTPException(status_code = 500, detail = f"Failed to load model: {str(e)}") + msg = str(e) + # Surface a friendlier message for models that Unsloth cannot load + not_supported_hints = [ + "No config file found", + "not yet supported", + "is not supported", + "does not support", + ] + if any(h.lower() in msg.lower() for h in not_supported_hints): + msg = f"This model is not supported yet. Try a different model. (Original error: {msg})" + raise HTTPException(status_code = 500, detail = f"Failed to load model: {msg}") @router.post("/validate", response_model = ValidateModelResponse) diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx index b8b464c56b..22a1cb7e72 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -408,8 +408,13 @@ export function HubModelPicker({ () => new Map( results - .filter((result) => result.totalParams) - .map((result) => [result.id, formatCompact(result.totalParams!)]), + .filter((result) => result.totalParams || result.estimatedSizeBytes) + .map((result) => [ + result.id, + result.estimatedSizeBytes + ? `~${formatBytes(result.estimatedSizeBytes)}` + : formatCompact(result.totalParams!), + ]), ), [results], ); diff --git a/studio/frontend/src/hooks/use-hf-model-search.ts b/studio/frontend/src/hooks/use-hf-model-search.ts index a092a895a8..f7b06ab65a 100644 --- a/studio/frontend/src/hooks/use-hf-model-search.ts +++ b/studio/frontend/src/hooks/use-hf-model-search.ts @@ -11,6 +11,7 @@ export interface HfModelResult { downloads: number; likes: number; totalParams?: number; + estimatedSizeBytes?: number; } const EXCLUDED_TAGS = new Set([ @@ -54,13 +55,33 @@ function withPopularitySort( return fetch(url, init); } +/** Bytes per parameter for each dtype. */ +const DTYPE_BYTES: Record = { + F64: 8, F32: 4, F16: 2, BF16: 2, + I64: 8, I32: 4, I16: 2, I8: 1, U8: 1, + // Quantized types (4-bit) + NF4: 0.5, FP4: 0.5, INT4: 0.5, GPTQ: 0.5, +}; + +function estimateSizeFromDtypes( + params: Record | undefined, +): number | undefined { + if (!params) return undefined; + let total = 0; + for (const [dtype, count] of Object.entries(params)) { + const bpp = DTYPE_BYTES[dtype.toUpperCase()] ?? 2; // default BF16 + total += count * bpp; + } + return total > 0 ? total : undefined; +} + function makeMapModel(excludeGguf: boolean) { return (raw: unknown): HfModelResult | null => { const m = raw as { name: string; downloads: number; likes: number; - safetensors?: { total: number }; + safetensors?: { total: number; parameters?: Record }; tags?: string[]; }; const isEmbedding = m.tags?.some((t) => EMBEDDING_TAGS.has(t)); @@ -75,6 +96,7 @@ function makeMapModel(excludeGguf: boolean) { downloads: m.downloads, likes: m.likes, totalParams: m.safetensors?.total, + estimatedSizeBytes: estimateSizeFromDtypes(m.safetensors?.parameters), }; }; } From 39854f44292d124bc30b9962d73123da6b607a04 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 16 Mar 2026 06:16:25 +0000 Subject: [PATCH 05/19] Auto-download mmproj for vision-capable GGUF models GGUF repos with mmproj files (e.g. Qwen3.5-0.8B-GGUF) are already detected as vision-capable by list_gguf_variants(), and is_vision is set correctly in ModelConfig. However, the HF download path only downloaded the main GGUF file without the mmproj projection file, so llama-server started without --mmproj and rejected image uploads with "text-only model" errors. Add _download_mmproj() to LlamaCppBackend that: - Lists repo files for mmproj*.gguf matches - Prefers mmproj-F16.gguf (best quality), falls back to any mmproj - Downloads via hf_hub_download (uses the same HF cache) In load_model(), when is_vision=True and no explicit mmproj_path was provided (HF mode), auto-download the mmproj after the main GGUF. The downloaded path is passed to llama-server via --mmproj. --- studio/backend/core/inference/llama_cpp.py | 48 ++++++++++++++++++++++ 1 file changed, 48 insertions(+) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 3ed4589dc0..f20a4fbe10 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -535,6 +535,48 @@ class LlamaCppBackend: logger.info(f"GGUF downloaded to: {local_path}") return local_path + def _download_mmproj( + self, + *, + hf_repo: str, + hf_token: Optional[str] = None, + ) -> Optional[str]: + """Download the mmproj (vision projection) file from a GGUF repo. + + Prefers mmproj-F16.gguf, falls back to any mmproj*.gguf file. + Returns the local path, or None if no mmproj file exists. + """ + try: + from huggingface_hub import hf_hub_download, list_repo_files + + files = list_repo_files(hf_repo, token = hf_token) + mmproj_files = sorted( + f for f in files + if f.endswith(".gguf") and "mmproj" in f.lower() + ) + if not mmproj_files: + return None + + # Prefer F16 variant + target = None + for f in mmproj_files: + if "f16" in f.lower(): + target = f + break + if target is None: + target = mmproj_files[0] + + logger.info(f"Downloading mmproj: {hf_repo}/{target}") + local_path = hf_hub_download( + repo_id = hf_repo, + filename = target, + token = hf_token, + ) + return local_path + except Exception as e: + logger.warning(f"Could not download mmproj: {e}") + return None + # ── Lifecycle ───────────────────────────────────────────────── def load_model( @@ -588,6 +630,12 @@ class LlamaCppBackend: hf_variant = hf_variant, hf_token = hf_token, ) + # Auto-download mmproj for vision models + if is_vision and not mmproj_path: + mmproj_path = self._download_mmproj( + hf_repo = hf_repo, + hf_token = hf_token, + ) elif gguf_path: if not Path(gguf_path).is_file(): raise FileNotFoundError(f"GGUF file not found: {gguf_path}") From c842e019d8bd24030f6e14cec10be853f1a60f5e Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 16 Mar 2026 06:16:35 +0000 Subject: [PATCH 06/19] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/inference/llama_cpp.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index f20a4fbe10..c0208d654d 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -551,8 +551,7 @@ class LlamaCppBackend: files = list_repo_files(hf_repo, token = hf_token) mmproj_files = sorted( - f for f in files - if f.endswith(".gguf") and "mmproj" in f.lower() + f for f in files if f.endswith(".gguf") and "mmproj" in f.lower() ) if not mmproj_files: return None From d41740708794707bd23f994ed4534fb4584b4fc4 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 16 Mar 2026 06:20:30 +0000 Subject: [PATCH 07/19] Convert images to PNG before sending to llama-server llama-server uses stb_image internally which does not support WebP, TIFF, AVIF, and other formats that browsers accept for upload. Uploading a WebP image to a vision GGUF model caused a 400 error: "Failed to load image or audio file" / "failed to decode image bytes". Convert all uploaded images to PNG via PIL before base64-encoding and forwarding to llama-server. This handles WebP, TIFF, BMP, GIF, AVIF, and any other format PIL supports. RGBA images are converted to RGB first since PNG with alpha can cause issues in some vision pipelines. --- studio/backend/routes/inference.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 13062f8b29..6740abbc58 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -860,6 +860,23 @@ async def openai_chat_completions( detail = "Image provided but current GGUF model does not support vision.", ) + # Convert image to PNG for llama-server (stb_image has limited format support) + if image_b64: + try: + import base64 as _b64 + from io import BytesIO as _BytesIO + from PIL import Image as _Image + + raw = _b64.b64decode(image_b64) + img = _Image.open(_BytesIO(raw)) + if img.mode == "RGBA": + img = img.convert("RGB") + buf = _BytesIO() + img.save(buf, format = "PNG") + image_b64 = _b64.b64encode(buf.getvalue()).decode("ascii") + except Exception as e: + raise HTTPException(status_code = 400, detail = f"Failed to process image: {e}") + # Build message list with system prompt prepended gguf_messages = [] if system_prompt: From a45babc6203f624759bddc5c664961742f31f56e Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 16 Mar 2026 06:20:48 +0000 Subject: [PATCH 08/19] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/routes/inference.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 6740abbc58..1b2ee91785 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -875,7 +875,9 @@ async def openai_chat_completions( img.save(buf, format = "PNG") image_b64 = _b64.b64encode(buf.getvalue()).decode("ascii") except Exception as e: - raise HTTPException(status_code = 400, detail = f"Failed to process image: {e}") + raise HTTPException( + status_code = 400, detail = f"Failed to process image: {e}" + ) # Build message list with system prompt prepended gguf_messages = [] From 2642f6d21d702e7aa92b098dfbc2c7f9728d384a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 16 Mar 2026 06:28:39 +0000 Subject: [PATCH 09/19] Add sloth emoji to section labels, friendlier network error - Add sloth emoji prefix to "Downloaded" and "Recommended" section labels in the Hub model picker so they are visually distinct. - Replace browser network errors ("NetworkError when attempting to fetch resource" / "Failed to fetch") with a clearer message: "Studio isn't running -- please relaunch it." --- .../components/assistant-ui/model-selector/pickers.tsx | 4 ++-- studio/frontend/src/features/auth/api.ts | 10 +++++++++- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx index 22a1cb7e72..6baeb9b55b 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -495,7 +495,7 @@ export function HubModelPicker({
{!showHfSection && (cachedGguf.length > 0 || cachedModels.length > 0) ? ( <> - Downloaded + {"\uD83E\uDDA5"} Downloaded {cachedGguf.map((c) => (
- Recommended + {"\uD83E\uDDA5"} Recommended {recommendedIds.length === 0 ? (
No default models. diff --git a/studio/frontend/src/features/auth/api.ts b/studio/frontend/src/features/auth/api.ts index e2e4932936..3bb2c9139c 100644 --- a/studio/frontend/src/features/auth/api.ts +++ b/studio/frontend/src/features/auth/api.ts @@ -84,7 +84,15 @@ export async function authFetch( headers.set("Authorization", `Bearer ${accessToken}`); } - const response = await fetch(input, { ...init, headers }); + let response: Response; + try { + response = await fetch(input, { ...init, headers }); + } catch (err) { + if (err instanceof TypeError) { + throw new Error("Studio isn't running -- please relaunch it."); + } + throw err; + } if (await isPasswordChangeRequiredResponse(response)) { void redirectToAuth(); return response; From 3a5d751f19bee95e0c9ed6ef8883480d70515a7c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 16 Mar 2026 06:31:32 +0000 Subject: [PATCH 10/19] Add logging to download-progress exception handler --- studio/backend/routes/models.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 02a810997a..bd64ae7196 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -752,7 +752,8 @@ async def get_download_progress( "expected_bytes": expected_bytes, "progress": round(progress, 3), } - except Exception: + except Exception as e: + logger.warning(f"Error checking download progress for {repo_id}: {e}") return _empty From f4d54a8de79a21f8d371a467b7677cd9a312fe8a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 16 Mar 2026 06:40:34 +0000 Subject: [PATCH 11/19] Fix vision detection subprocess using undefined logger The _VISION_CHECK_SCRIPT subprocess used logger.info() but logger was never defined in the subprocess context. This caused a NameError on every vision check, making all transformers 5.x models (Qwen3.5, GLM, etc.) fall back to text-only mode even when they support vision. Replace logger.info() with print() since the parent process reads the subprocess stdout via result.stdout. --- studio/backend/utils/models/model_config.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 8df53a28c9..aaf15994be 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -473,10 +473,10 @@ try: model_type = getattr(config, "model_type", "unknown") archs = getattr(config, "architectures", []) - logger.info(json.dumps({"is_vision": is_vlm, "model_type": model_type, + print(json.dumps({"is_vision": is_vlm, "model_type": model_type, "architectures": archs})) except Exception as exc: - logger.info(json.dumps({"error": str(exc)})) + print(json.dumps({"error": str(exc)})) sys.exit(1) """ From 042598d9f1cf7688e52152f6604faab00c8f892d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 16 Mar 2026 06:46:16 +0000 Subject: [PATCH 12/19] Suppress model-switch warning on empty chat threads Don't show "Model changed for this chat" toast when the thread has no messages. On a fresh page load with a stale thread from a previous session, this warning is confusing. The warning is only useful mid-conversation to alert about image compatibility with the new model. When messages.length === 0, silently update the thread's modelId and proceed with loading. --- studio/frontend/src/features/chat/chat-page.tsx | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index 0e06d89b28..bd506f6b42 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -352,6 +352,18 @@ export function ChatPage(): ReactElement { .where("threadId") .equals(activeThreadId) .toArray(); + if (messages.length === 0) { + // No history -- just switch silently + await db.threads.update(activeThreadId, { modelId: value }); + await selectModel({ + id: value, + isLora: meta?.isLora, + ggufVariant: meta?.ggufVariant, + isDownloaded: meta?.isDownloaded, + expectedBytes: meta?.expectedBytes, + }); + return; + } const hasImage = messages.some(messageHasImage); const targetModel = modelsFromStore.find((model) => model.id === value); const nonVisionWithImages = hasImage && targetModel?.isVision === false; From 9cbeecc16aaf75b28c17b499bcd5b9606e58f2f8 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 16 Mar 2026 06:51:52 +0000 Subject: [PATCH 13/19] Incorporate PR #4304 toast UX improvements Merge the toast UX refactor from PR #4304 (by @Shine1i): - Toast duration 5s default with close button (X) for manual dismiss - Inline progress bar component (ModelLoadInlineStatus) shown in the header after toast is dismissed - Model switch warning only for image compatibility (not generic) - activeThreadId tracked in store via ActiveThreadSync - Loading state cleanup via resetLoadingUi helper - Toast uses Infinity duration during loading with onDismiss handler Re-applied non-GGUF download progress additions on top: - getDownloadProgress for all models (not just GGUF) - hasShownProgress flag, loadingModelRef race condition checks - First poll at 500ms, bytes-only fallback when expected size unknown --- studio/frontend/src/components/ui/sonner.tsx | 10 +- .../frontend/src/features/chat/chat-page.tsx | 103 +++---- .../chat/components/model-load-status.tsx | 106 +++++++ .../chat/hooks/use-chat-model-runtime.ts | 286 ++++++++++++------ .../src/features/chat/runtime-provider.tsx | 17 ++ .../chat/stores/chat-runtime-store.ts | 4 + 6 files changed, 359 insertions(+), 167 deletions(-) create mode 100644 studio/frontend/src/features/chat/components/model-load-status.tsx diff --git a/studio/frontend/src/components/ui/sonner.tsx b/studio/frontend/src/components/ui/sonner.tsx index a20f106d42..f211c4acd0 100644 --- a/studio/frontend/src/components/ui/sonner.tsx +++ b/studio/frontend/src/components/ui/sonner.tsx @@ -16,11 +16,11 @@ const Toaster = ({ ...props }: ToasterProps) => { const { theme = "system" } = useTheme(); return ( - { - if (view.mode !== "single") { - return undefined; - } - if (view.threadId) { - return view.threadId; - } - - // New-thread flow keeps threadId undefined in local view state. - // Fall back to most recent regular base thread. - const candidates = await db.threads.where("modelType").equals("base").toArray(); - const latest = candidates - .filter((thread) => !thread.archived && !thread.pairId) - .sort((a, b) => b.createdAt - a.createdAt)[0]; - return latest?.id; -} - const SingleContent = memo(function SingleContent({ threadId, newThreadNonce, @@ -321,7 +304,16 @@ export function ChatPage(): ReactElement { const modelsFromStore = useChatRuntimeStore((state) => state.models); const lorasFromStore = useChatRuntimeStore((state) => state.loras); const modelsError = useChatRuntimeStore((state) => state.modelsError); - const { refresh, selectModel, ejectModel, cancelLoading, loadingModel } = + const activeThreadId = useChatRuntimeStore((state) => state.activeThreadId); + const { + refresh, + selectModel, + ejectModel, + cancelLoading, + loadingModel, + loadProgress, + loadToastDismissed, + } = useChatModelRuntime(); const refreshRef = useRef(refresh); const selectModelRef = useRef(selectModel); @@ -343,42 +335,27 @@ export function ChatPage(): ReactElement { const currentVariant = store.activeGgufVariant; if (!value || (value === currentCheckpoint && (meta?.ggufVariant ?? null) === (currentVariant ?? null))) return; void (async () => { - let switchNote: string | undefined; - const activeThreadId = await resolveActiveSingleThreadId(view); - if (activeThreadId) { + let showImageCompatibilityWarning = false; + if (view.mode === "single" && activeThreadId) { const thread = await db.threads.get(activeThreadId); if (thread?.modelId && thread.modelId !== value) { const messages = await db.messages .where("threadId") .equals(activeThreadId) .toArray(); - if (messages.length === 0) { - // No history -- just switch silently - await db.threads.update(activeThreadId, { modelId: value }); - await selectModel({ - id: value, - isLora: meta?.isLora, - ggufVariant: meta?.ggufVariant, - isDownloaded: meta?.isDownloaded, - expectedBytes: meta?.expectedBytes, - }); - return; + if (messages.length > 0) { + const hasImage = messages.some(messageHasImage); + const targetModel = modelsFromStore.find((model) => model.id === value); + showImageCompatibilityWarning = + hasImage && targetModel?.isVision === false; } - const hasImage = messages.some(messageHasImage); - const targetModel = modelsFromStore.find((model) => model.id === value); - const nonVisionWithImages = hasImage && targetModel?.isVision === false; - - switchNote = nonVisionWithImages - ? "Full chat history will be sent to the new model. This chat has images; text-only models may fail." - : hasImage - ? "Full chat history will be sent to the new model. This chat includes images." - : "Full chat history will be sent to the new model."; } } - if (switchNote) { - toast.warning("Model changed for this chat", { - description: switchNote, + if (showImageCompatibilityWarning) { + toast.warning("Selected model may not handle earlier images", { + description: + "This chat already includes images. Text-only models can ignore them or fail on follow-up replies.", duration: 6000, }); } @@ -391,13 +368,16 @@ export function ChatPage(): ReactElement { }); })(); }, - [modelsFromStore, selectModel, view], + [activeThreadId, modelsFromStore, selectModel, view], ); const handleEject = useCallback(() => { void ejectModel(); }, [ejectModel]); const handleNewThread = useCallback( - () => setView({ mode: "single", newThreadNonce: crypto.randomUUID() }), + () => { + useChatRuntimeStore.getState().setActiveThreadId(null); + setView({ mode: "single", newThreadNonce: crypto.randomUUID() }); + }, [], ); const handleNewCompare = useCallback( @@ -618,25 +598,22 @@ export function ChatPage(): ReactElement { contentDataTour="chat-model-selector-popover" className="max-w-[62vw] sm:max-w-none" /> - {loadingModel ? ( -
- - - {loadingModel.isDownloaded ? "Loading model…" : "Downloading model…"} - - -
+ progressPercent={loadProgress?.percent} + progressLabel={loadProgress?.label} + onStop={cancelLoading} + /> ) : null}
{modelsError && ( diff --git a/studio/frontend/src/features/chat/components/model-load-status.tsx b/studio/frontend/src/features/chat/components/model-load-status.tsx new file mode 100644 index 0000000000..9686011aca --- /dev/null +++ b/studio/frontend/src/features/chat/components/model-load-status.tsx @@ -0,0 +1,106 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { Progress } from "@/components/ui/progress"; +import { Spinner } from "@/components/ui/spinner"; +import { Button } from "@/components/ui/button"; + +type ModelLoadDescriptionProps = { + message?: string | null; + progressPercent?: number | null; + progressLabel?: string | null; + onStop?: () => void; + stopLabel?: string; +}; + +function clampProgress(value: number): number { + return Math.max(0, Math.min(100, value)); +} + +export function ModelLoadDescription({ + message, + progressPercent, + progressLabel, + onStop, + stopLabel = "Stop loading", +}: ModelLoadDescriptionProps) { + const hasProgress = typeof progressPercent === "number"; + + return ( +
+
+ {hasProgress ? ( +
+
+ {progressLabel} + {Math.round(clampProgress(progressPercent))}% +
+ +
+ ) : message ? ( +

{message}

+ ) : null} +
+ {onStop ? ( + + ) : null} +
+ ); +} + +type ModelLoadInlineStatusProps = { + label: string; + title: string; + progressPercent?: number | null; + progressLabel?: string | null; + onStop?: () => void; +}; + +export function ModelLoadInlineStatus({ + label, + title, + progressPercent, + progressLabel, + onStop, +}: ModelLoadInlineStatusProps) { + const hasProgress = typeof progressPercent === "number"; + + return ( +
+
+ + {label} +
+ {hasProgress ? ( +
+
+ +
+
+ {progressLabel} + {Math.round(clampProgress(progressPercent))}% +
+
+ ) : null} + {onStop ? ( + + ) : null} +
+ ); +} diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index f1d0e16d0c..e5c58c7be7 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -1,8 +1,10 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -import { useCallback, useRef, useState } from "react"; +import { createElement, useCallback, useRef, useState } from "react"; import { toast } from "sonner"; +import { Spinner } from "@/components/ui/spinner"; +import { ModelLoadDescription } from "../components/model-load-status"; import { getDownloadProgress, getGgufDownloadProgress, @@ -30,6 +32,15 @@ type SelectedModelInput = { expectedBytes?: number; }; +const MODEL_LOAD_TOAST_CLASSNAMES = { + toast: "items-start gap-2.5 pr-8", + content: "gap-0.5", + title: "leading-5", + description: "mt-0", + closeButton: + "!left-auto !right-1 !top-2 !translate-x-0 !translate-y-0 !border-transparent !bg-transparent !shadow-none hover:!bg-transparent hover:opacity-70", +} as const; + const LORA_SUFFIX_RE = /_(\d{9,})$/; function parseTrailingEpoch(input: string): number | undefined { @@ -152,11 +163,46 @@ export function useChatModelRuntime() { displayName: string; isDownloaded?: boolean; } | null>(null); - const [_loadAbortController, setLoadAbortController] = - useState(null); + const [loadToastDismissed, setLoadToastDismissed] = useState(false); + const [loadProgress, setLoadProgress] = useState<{ + percent: number | null; + label: string | null; + phase: "downloading" | "starting"; + } | null>(null); const loadAbortRef = useRef(null); const loadingModelRef = useRef(null); const loadToastIdRef = useRef(null); + const loadToastDismissedRef = useRef(false); + + const setLoadToastDismissedState = useCallback((dismissed: boolean) => { + loadToastDismissedRef.current = dismissed; + setLoadToastDismissed(dismissed); + }, []); + + const resetLoadingUi = useCallback(() => { + setLoadingModel(null); + setLoadProgress(null); + loadingModelRef.current = null; + loadAbortRef.current = null; + loadToastIdRef.current = null; + setLoadToastDismissedState(false); + }, [setLoadToastDismissedState]); + + const renderLoadDescription = useCallback( + ( + message: string, + progressPercent?: number | null, + progressLabel?: string | null, + onStop?: () => void, + ) => + createElement(ModelLoadDescription, { + message, + progressPercent, + progressLabel, + onStop, + }), + [], + ); const refresh = useCallback(async () => { setModelsError(null); @@ -183,6 +229,26 @@ export function useChatModelRuntime() { } }, [setCheckpoint, setLoras, setModels, setModelsError]); + const cancelLoading = useCallback(() => { + const model = loadingModelRef.current; + if (!model) return; + loadAbortRef.current?.abort(); + loadAbortRef.current = null; + loadingModelRef.current = null; + const tid = loadToastIdRef.current; + loadToastIdRef.current = null; + setLoadingModel(null); + setLoadProgress(null); + setLoadToastDismissedState(false); + clearCheckpoint(); + if (tid != null) toast.dismiss(tid); + toast.info("Stopped loading model", { + description: "The current download may still finish in the background.", + }); + // Fire-and-forget: tell backend to stop, don't block UI + unloadModel({ model_path: model.id }).catch(() => {}); + }, [clearCheckpoint, setLoadToastDismissedState]); + const selectModel = useCallback( async (selection: string | SelectedModelInput) => { const modelId = typeof selection === "string" ? selection : selection.id; @@ -218,21 +284,23 @@ export function useChatModelRuntime() { const previousIsLora = previousModel?.isLora ?? (previousLora ? true : false); const loadingDescription = [ - currentCheckpoint ? "Unloading previous model first." : null, + currentCheckpoint ? "Switching models." : null, extraLoadingDescription ?? null, - isDownloaded - ? "Loading cached model into memory." - : "Downloading and loading model. Large models can take a while.", + isDownloaded ? "Loading cached model into memory." : null, ] .filter(Boolean) .join(" "); - setModelsError(null); + setLoadToastDismissedState(false); const loadInfo = { id: modelId, displayName, isDownloaded }; setLoadingModel(loadInfo); + setLoadProgress( + isDownloaded + ? { percent: null, label: null, phase: "starting" } + : { percent: 0, label: "Preparing download", phase: "downloading" }, + ); loadingModelRef.current = loadInfo; const abortCtrl = new AbortController(); - setLoadAbortController(abortCtrl); loadAbortRef.current = abortCtrl; try { async function performLoad(): Promise { @@ -301,56 +369,37 @@ export function useChatModelRuntime() { } } - const toastId = toast.loading( - isDownloaded ? "Loading model…" : "Downloading model…", + const toastId = toast( + isDownloaded ? "Starting model…" : "Downloading model…", { - description: loadingDescription, - duration: 10000, - action: { - label: "Cancel", - onClick: () => { - abortCtrl.abort(); - setLoadingModel(null); - setLoadAbortController(null); - loadingModelRef.current = null; - loadAbortRef.current = null; - loadToastIdRef.current = null; - unloadModel({ model_path: modelId }).catch(() => {}); - clearCheckpoint(); - toast.dismiss(toastId); - toast.info("Model loading cancelled"); - }, + icon: createElement(Spinner, { className: "size-4" }), + description: renderLoadDescription( + loadingDescription, + isDownloaded ? null : 0, + isDownloaded ? null : "Preparing download", + cancelLoading, + ), + duration: Infinity, + closeButton: true, + classNames: MODEL_LOAD_TOAST_CLASSNAMES, + onDismiss: (dismissedToast) => { + if (loadToastIdRef.current !== dismissedToast.id) { + return; + } + setLoadToastDismissedState(true); }, }, ); loadToastIdRef.current = toastId; - // Poll download progress for non-cached models + // Poll download progress for non-cached models (GGUF and non-GGUF) let progressInterval: ReturnType | null = null; if (!isDownloaded) { const expectedBytes = typeof selection !== "string" ? selection.expectedBytes ?? 0 : 0; - - const cancelAction = { - label: "Cancel", - onClick: () => { - abortCtrl.abort(); - setLoadingModel(null); - setLoadAbortController(null); - loadingModelRef.current = null; - loadAbortRef.current = null; - loadToastIdRef.current = null; - unloadModel({ model_path: modelId }).catch(() => {}); - clearCheckpoint(); - toast.dismiss(toastId); - toast.info("Model loading cancelled"); - }, - }; - let hasShownProgress = false; const pollProgress = async () => { - // Stop if cancelled or if loading already finished if (abortCtrl.signal.aborted || !loadingModelRef.current) { if (progressInterval) clearInterval(progressInterval); return; @@ -360,7 +409,6 @@ export function useChatModelRuntime() { ? await getGgufDownloadProgress(modelId, ggufVariant, expectedBytes) : await getDownloadProgress(modelId); - // Re-check after await -- load may have finished while polling if (!loadingModelRef.current) return; if (prog.progress > 0 && prog.progress < 1) { @@ -368,36 +416,69 @@ export function useChatModelRuntime() { const dlGb = prog.downloaded_bytes / (1024 ** 3); const totalGb = prog.expected_bytes / (1024 ** 3); const pct = Math.round(prog.progress * 100); - toast.loading( - `Downloading model... ${pct}%`, + const progressLabel = totalGb > 0 + ? `${dlGb.toFixed(1)} of ${totalGb.toFixed(1)} GB` + : `${dlGb.toFixed(1)} GB downloaded`; + setLoadProgress({ + percent: pct, + label: progressLabel, + phase: "downloading", + }); + if (loadToastDismissedRef.current) return; + toast( + "Downloading model…", { id: toastId, - description: totalGb > 0 - ? `${dlGb.toFixed(1)} / ${totalGb.toFixed(1)} GB` - : `${dlGb.toFixed(1)} GB downloaded`, - duration: 10000, - action: cancelAction, + icon: createElement(Spinner, { className: "size-4" }), + description: renderLoadDescription( + loadingDescription, + pct, + progressLabel, + cancelLoading, + ), + duration: Infinity, + closeButton: true, + classNames: MODEL_LOAD_TOAST_CLASSNAMES, + onDismiss: (dismissedToast) => { + if (loadToastIdRef.current !== dismissedToast.id) return; + setLoadToastDismissedState(true); + }, }, ); } else if (prog.downloaded_bytes > 0 && prog.expected_bytes === 0 && prog.progress === 0) { - // Have bytes but no total size -- show bytes only hasShownProgress = true; const dlGb = prog.downloaded_bytes / (1024 ** 3); - toast.loading( - "Downloading model...", - { - id: toastId, - description: `${dlGb.toFixed(1)} GB downloaded`, - duration: 10000, - action: cancelAction, - }, - ); + setLoadProgress({ + percent: null, + label: `${dlGb.toFixed(1)} GB downloaded`, + phase: "downloading", + }); } else if (prog.progress >= 1 && hasShownProgress) { - // Only show "download complete" if we actually showed progress - toast.loading("Loading model...", { + setLoadProgress({ + percent: 100, + label: "Download complete", + phase: "starting", + }); + if (loadToastDismissedRef.current) { + if (progressInterval) clearInterval(progressInterval); + return; + } + toast("Starting model…", { id: toastId, - description: "Download complete. Loading into memory...", - duration: 10000, + icon: createElement(Spinner, { className: "size-4" }), + description: renderLoadDescription( + "Download complete. Loading the model into memory.", + 100, + "Download complete", + cancelLoading, + ), + duration: Infinity, + closeButton: true, + classNames: MODEL_LOAD_TOAST_CLASSNAMES, + onDismiss: (dismissedToast) => { + if (loadToastIdRef.current !== dismissedToast.id) return; + setLoadToastDismissedState(true); + }, }); if (progressInterval) clearInterval(progressInterval); } @@ -406,40 +487,62 @@ export function useChatModelRuntime() { } }; - // First poll after 500ms, then every 2s setTimeout(pollProgress, 500); progressInterval = setInterval(pollProgress, 2000); } try { await performLoad(); - toast.success(`${displayName} loaded`, { id: toastId }); + if (loadToastDismissedRef.current) { + toast.success(`${displayName} loaded`); + } else { + toast.success(`${displayName} loaded`, { + id: toastId, + description: undefined, + closeButton: false, + duration: 2000, + }); + } } catch (err) { if (!abortCtrl.signal.aborted) { - toast.error( - err instanceof Error ? err.message : "Failed to load model", - { id: toastId }, - ); + const message = + err instanceof Error ? err.message : "Failed to load model"; + if (loadToastDismissedRef.current) { + toast.error(message); + } else { + toast.error(message, { + id: toastId, + description: undefined, + closeButton: false, + duration: 5000, + }); + } } throw err; } finally { if (progressInterval) clearInterval(progressInterval); - setLoadingModel(null); - setLoadAbortController(null); - loadingModelRef.current = null; - loadAbortRef.current = null; - loadToastIdRef.current = null; + resetLoadingUi(); } } catch (error) { if (abortCtrl.signal.aborted) return; // User cancelled, nothing to report - setLoadingModel(null); - loadingModelRef.current = null; + resetLoadingUi(); const message = error instanceof Error ? error.message : "Failed to load model"; setModelsError(message); } }, - [loras, models, params.checkpoint, refresh, setModelsError, setParams], + [ + cancelLoading, + loras, + models, + params.checkpoint, + refresh, + renderLoadDescription, + resetLoadingUi, + setLoadToastDismissedState, + setModelsError, + setParams, + ], ); const ejectModel = useCallback(async () => { @@ -468,28 +571,13 @@ export function useChatModelRuntime() { } }, [clearCheckpoint, params.checkpoint, refresh, setModelsError]); - const cancelLoading = useCallback(() => { - const model = loadingModelRef.current; - if (!model) return; - loadAbortRef.current?.abort(); - loadAbortRef.current = null; - loadingModelRef.current = null; - const tid = loadToastIdRef.current; - loadToastIdRef.current = null; - setLoadingModel(null); - setLoadAbortController(null); - clearCheckpoint(); - if (tid != null) toast.dismiss(tid); - toast.info("Model loading cancelled"); - // Fire-and-forget: tell backend to stop, don't block UI - unloadModel({ model_path: model.id }).catch(() => {}); - }, [clearCheckpoint]); - return { refresh, selectModel, ejectModel, cancelLoading, loadingModel, + loadProgress, + loadToastDismissed, }; } diff --git a/studio/frontend/src/features/chat/runtime-provider.tsx b/studio/frontend/src/features/chat/runtime-provider.tsx index dbc631a1d5..0ad6005509 100644 --- a/studio/frontend/src/features/chat/runtime-provider.tsx +++ b/studio/frontend/src/features/chat/runtime-provider.tsx @@ -559,6 +559,22 @@ function ThreadNewChatSwitch({ return null; } +function ActiveThreadSync({ + enabled, +}: { enabled: boolean }): ReactElement | null { + const mainThreadId = useAuiState(({ threads }) => threads.mainThreadId); + const setActiveThreadId = useChatRuntimeStore((state) => state.setActiveThreadId); + + useEffect(() => { + if (!enabled) { + return; + } + setActiveThreadId(mainThreadId ?? null); + }, [enabled, mainThreadId, setActiveThreadId]); + + return null; +} + export function ChatRuntimeProvider({ children, modelType = "base", @@ -586,6 +602,7 @@ export function ChatRuntimeProvider({ return ( + {initialThreadId && } {!initialThreadId && newThreadNonce && ( diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index 87e162ed56..07d3aecc36 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -43,6 +43,7 @@ type ChatRuntimeStore = { autoTitle: boolean; modelsError: string | null; activeGgufVariant: string | null; + activeThreadId: string | null; pendingAudioBase64: string | null; pendingAudioName: string | null; setParams: (params: InferenceParams) => void; @@ -52,6 +53,7 @@ type ChatRuntimeStore = { setAutoTitle: (enabled: boolean) => void; setModelsError: (error: string | null) => void; setCheckpoint: (modelId: string, ggufVariant?: string | null) => void; + setActiveThreadId: (threadId: string | null) => void; clearCheckpoint: () => void; setPendingAudio: (base64: string, name: string) => void; clearPendingAudio: () => void; @@ -65,6 +67,7 @@ export const useChatRuntimeStore = create((set) => ({ autoTitle: loadBool(AUTO_TITLE_KEY, false), modelsError: null, activeGgufVariant: null, + activeThreadId: null, pendingAudioBase64: null, pendingAudioName: null, setParams: (params) => set({ params }), @@ -94,6 +97,7 @@ export const useChatRuntimeStore = create((set) => ({ }, activeGgufVariant: ggufVariant ?? null, })), + setActiveThreadId: (activeThreadId) => set({ activeThreadId }), clearCheckpoint: () => set((state) => ({ params: { From 8ffd86012fdaef3e4a3445fdfbc32e385649cbaf Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 16 Mar 2026 06:58:46 +0000 Subject: [PATCH 14/19] Change "Stop loading" to outlined "Stop" button --- .../features/chat/components/model-load-status.tsx | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/studio/frontend/src/features/chat/components/model-load-status.tsx b/studio/frontend/src/features/chat/components/model-load-status.tsx index 9686011aca..19eff054f7 100644 --- a/studio/frontend/src/features/chat/components/model-load-status.tsx +++ b/studio/frontend/src/features/chat/components/model-load-status.tsx @@ -10,7 +10,6 @@ type ModelLoadDescriptionProps = { progressPercent?: number | null; progressLabel?: string | null; onStop?: () => void; - stopLabel?: string; }; function clampProgress(value: number): number { @@ -22,7 +21,6 @@ export function ModelLoadDescription({ progressPercent, progressLabel, onStop, - stopLabel = "Stop loading", }: ModelLoadDescriptionProps) { const hasProgress = typeof progressPercent === "number"; @@ -45,11 +43,11 @@ export function ModelLoadDescription({ ) : null}
@@ -94,8 +92,8 @@ export function ModelLoadInlineStatus({