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:
parent
061de08f86
commit
475ba417dc
6 changed files with 145 additions and 20 deletions
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
)}
|
||||
|
|
|
|||
|
|
@ -21,5 +21,7 @@ export interface ModelSelectorChangeMeta {
|
|||
source: "hub" | "lora" | "exported";
|
||||
isLora: boolean;
|
||||
ggufVariant?: string;
|
||||
isDownloaded?: boolean;
|
||||
expectedBytes?: number;
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue