studio: context-aware loading text + download progress bar

1. Loading text: shows "Loading model..." for cached models,
   "Downloading model..." for new downloads. Toast description
   adapts accordingly.

2. Download progress: polls /api/models/gguf-download-progress every
   2s during downloads, updating the toast with percentage and GB
   downloaded. Progress is estimated by checking the HF cache folder
   size against the expected total bytes.

3. Passes isDownloaded and expectedBytes through the full chain from
   variant click to selectModel for accurate UI state.
This commit is contained in:
Daniel Han 2026-03-15 07:47:43 +00:00
commit 475ba417dc
6 changed files with 145 additions and 20 deletions

View file

@ -604,6 +604,46 @@ async def get_gguf_variants(
)
@router.get("/gguf-download-progress")
async def get_gguf_download_progress(
repo_id: str = Query(..., description = "HuggingFace repo ID"),
expected_bytes: int = Query(0, description = "Expected total download size in bytes"),
current_subject: str = Depends(get_current_subject),
):
"""Return download progress by checking current size of cached GGUF files."""
import re as _re
try:
if not _re.fullmatch(r"[A-Za-z0-9._-]+/[A-Za-z0-9._-]+", repo_id):
return {"downloaded_bytes": 0, "expected_bytes": expected_bytes, "progress": 0}
from huggingface_hub import constants as hf_constants
cache_dir = Path(hf_constants.HF_HUB_CACHE)
target = f"models--{repo_id.replace('/', '--')}".lower()
downloaded_bytes = 0
for entry in cache_dir.iterdir():
if entry.name.lower() == target:
# Sum .gguf files in snapshots + incomplete downloads in blobs
for f in entry.rglob("*.gguf"):
downloaded_bytes += f.stat().st_size
# Also check incomplete downloads (blobs without extension)
blobs_dir = entry / "blobs"
if blobs_dir.is_dir():
for f in blobs_dir.iterdir():
if f.is_file():
downloaded_bytes += f.stat().st_size
break
progress = min(downloaded_bytes / expected_bytes, 1.0) if expected_bytes > 0 else 0
return {
"downloaded_bytes": downloaded_bytes,
"expected_bytes": expected_bytes,
"progress": round(progress, 3),
}
except Exception:
return {"downloaded_bytes": 0, "expected_bytes": expected_bytes, "progress": 0}
@router.get("/cached-gguf")
async def list_cached_gguf(
current_subject: str = Depends(get_current_subject),

View file

@ -184,11 +184,13 @@ function GgufVariantExpander({
}, [repoId]);
const handleVariantClick = useCallback(
(quant: string) => {
(quant: string, downloaded?: boolean, sizeBytes?: number) => {
onSelect(repoId, {
source: "hub",
isLora: false,
ggufVariant: quant,
isDownloaded: downloaded,
expectedBytes: sizeBytes,
});
},
[repoId, onSelect],
@ -294,7 +296,7 @@ function GgufVariantExpander({
<button
key={v.filename}
type="button"
onClick={() => handleVariantClick(v.quant)}
onClick={() => handleVariantClick(v.quant, v.downloaded, v.size_bytes)}
className={cn(
"flex w-full items-center justify-between gap-2 rounded-md px-2.5 py-1 text-left text-sm transition-colors hover:bg-accent",
)}

View file

@ -21,5 +21,7 @@ export interface ModelSelectorChangeMeta {
source: "hub" | "lora" | "exported";
isLora: boolean;
ggufVariant?: string;
isDownloaded?: boolean;
expectedBytes?: number;
}

View file

@ -103,6 +103,18 @@ export interface CachedGgufRepo {
cache_path: string;
}
export async function getGgufDownloadProgress(
repoId: string,
expectedBytes: number,
): Promise<{ downloaded_bytes: number; expected_bytes: number; progress: number }> {
const params = new URLSearchParams({
repo_id: repoId,
expected_bytes: String(expectedBytes),
});
const response = await authFetch(`/api/models/gguf-download-progress?${params}`);
return parseJsonOrThrow(response);
}
export async function listCachedGguf(): Promise<CachedGgufRepo[]> {
const response = await authFetch("/api/models/cached-gguf");
const data = await parseJsonOrThrow<{ cached: CachedGgufRepo[] }>(response);

View file

@ -337,7 +337,7 @@ export function ChatPage(): ReactElement {
}, [inferenceParams.checkpoint, lorasFromStore]);
const handleCheckpointChange = useCallback(
(value: string, meta?: { isLora: boolean; ggufVariant?: string }) => {
(value: string, meta?: { isLora: boolean; ggufVariant?: string; isDownloaded?: boolean; expectedBytes?: number }) => {
const store = useChatRuntimeStore.getState();
const currentCheckpoint = store.params.checkpoint;
const currentVariant = store.activeGgufVariant;
@ -374,6 +374,8 @@ export function ChatPage(): ReactElement {
id: value,
isLora: meta?.isLora,
ggufVariant: meta?.ggufVariant,
isDownloaded: meta?.isDownloaded,
expectedBytes: meta?.expectedBytes,
});
})();
},
@ -607,11 +609,13 @@ export function ChatPage(): ReactElement {
{loadingModel ? (
<div
className="flex items-center gap-1.5 text-muted-foreground"
title={`Loading ${loadingModel.displayName}. This may include downloading.`}
title={loadingModel.isDownloaded
? `Loading ${loadingModel.displayName} from cache.`
: `Loading ${loadingModel.displayName}. This may include downloading.`}
>
<Spinner className="size-3.5 shrink-0" />
<span className="text-xs">
Downloading model
{loadingModel.isDownloaded ? "Loading model…" : "Downloading model…"}
</span>
<button
type="button"

View file

@ -4,6 +4,7 @@
import { useCallback, useState } from "react";
import { toast } from "sonner";
import {
getGgufDownloadProgress,
getInferenceStatus,
listLoras,
listModels,
@ -24,6 +25,8 @@ type SelectedModelInput = {
isLora?: boolean;
ggufVariant?: string;
loadingDescription?: string;
isDownloaded?: boolean;
expectedBytes?: number;
};
const LORA_SUFFIX_RE = /_(\d{9,})$/;
@ -146,6 +149,7 @@ export function useChatModelRuntime() {
const [loadingModel, setLoadingModel] = useState<{
id: string;
displayName: string;
isDownloaded?: boolean;
} | null>(null);
const [loadAbortController, setLoadAbortController] =
useState<AbortController | null>(null);
@ -189,6 +193,8 @@ export function useChatModelRuntime() {
typeof selection === "string" ? undefined : selection.isLora;
const extraLoadingDescription =
typeof selection === "string" ? undefined : selection.loadingDescription;
const isDownloaded =
typeof selection === "string" ? false : selection.isDownloaded ?? false;
const model = models.find((entry) => entry.id === modelId);
const lora = loras.find((entry) => entry.id === modelId);
const isLora =
@ -210,13 +216,15 @@ export function useChatModelRuntime() {
const loadingDescription = [
currentCheckpoint ? "Unloading previous model first." : null,
extraLoadingDescription ?? null,
"This may include downloading. Large models can take a while.",
isDownloaded
? "Loading cached model into memory."
: "This may include downloading. Large models can take a while.",
]
.filter(Boolean)
.join(" ");
setModelsError(null);
setLoadingModel({ id: modelId, displayName });
setLoadingModel({ id: modelId, displayName, isDownloaded });
const abortCtrl = new AbortController();
setLoadAbortController(abortCtrl);
try {
@ -280,21 +288,77 @@ export function useChatModelRuntime() {
}
}
const toastId = toast.loading("Loading model…", {
description: loadingDescription,
action: {
label: "Cancel",
onClick: () => {
abortCtrl.abort();
setLoadingModel(null);
setLoadAbortController(null);
unloadModel({ model_path: modelId }).catch(() => {});
clearCheckpoint();
toast.dismiss(toastId);
toast.info("Model loading cancelled");
const toastId = toast.loading(
isDownloaded ? "Loading model…" : "Downloading model…",
{
description: loadingDescription,
duration: Infinity,
action: {
label: "Cancel",
onClick: () => {
abortCtrl.abort();
setLoadingModel(null);
setLoadAbortController(null);
unloadModel({ model_path: modelId }).catch(() => {});
clearCheckpoint();
toast.dismiss(toastId);
toast.info("Model loading cancelled");
},
},
},
});
);
// Poll download progress for non-cached models
let progressInterval: ReturnType<typeof setInterval> | null = null;
if (!isDownloaded && ggufVariant) {
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, 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: Infinity,
action: {
label: "Cancel",
onClick: () => {
abortCtrl.abort();
setLoadingModel(null);
setLoadAbortController(null);
unloadModel({ model_path: modelId }).catch(() => {});
clearCheckpoint();
toast.dismiss(toastId);
toast.info("Model loading cancelled");
},
},
},
);
} else if (prog.progress >= 1) {
toast.loading("Loading model…", {
id: toastId,
description: "Download complete. Starting inference server…",
duration: Infinity,
});
if (progressInterval) clearInterval(progressInterval);
}
} catch {
// Ignore polling errors
}
}, 2000);
}
}
try {
await performLoad();
@ -308,6 +372,7 @@ export function useChatModelRuntime() {
}
throw err;
} finally {
if (progressInterval) clearInterval(progressInterval);
setLoadingModel(null);
setLoadAbortController(null);
}