From b84f167d5a187fd1c0c4bd948ea0014b07aa73f0 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 16 Mar 2026 05:53:34 +0000 Subject: [PATCH] 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 {