diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 1e27b19b12..1331aa5fd3 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -607,10 +607,15 @@ async def get_gguf_variants( @router.get("/gguf-download-progress") async def get_gguf_download_progress( repo_id: str = Query(..., description = "HuggingFace repo ID"), + variant: str = Query("", description = "Quantization variant (e.g. UD-TQ1_0)"), 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.""" + """Return download progress by checking cached GGUF files for a specific variant. + + Tracks completed shard downloads in snapshots and in-progress downloads + in the blobs directory (incomplete files). + """ import re as _re try: @@ -625,25 +630,31 @@ async def get_gguf_download_progress( cache_dir = Path(hf_constants.HF_HUB_CACHE) target = f"models--{repo_id.replace('/', '--')}".lower() + variant_lower = variant.lower().replace("-", "").replace("_", "") downloaded_bytes = 0 + in_progress_bytes = 0 for entry in cache_dir.iterdir(): if entry.name.lower() == target: - # Sum .gguf files in snapshots + incomplete downloads in blobs + # Count completed .gguf files matching this variant in snapshots for f in entry.rglob("*.gguf"): - downloaded_bytes += f.stat().st_size - # Also check incomplete downloads (blobs without extension) + fname = f.name.lower().replace("-", "").replace("_", "") + if not variant_lower or variant_lower in fname: + downloaded_bytes += f.stat().st_size + # Check blobs for in-progress downloads (.incomplete files) 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 + if f.is_file() and f.name.endswith(".incomplete"): + in_progress_bytes += f.stat().st_size break - progress = ( - min(downloaded_bytes / expected_bytes, 1.0) if expected_bytes > 0 else 0 - ) + total_progress_bytes = downloaded_bytes + in_progress_bytes + progress = min(total_progress_bytes / expected_bytes, 0.99) if expected_bytes > 0 else 0 + # Only report 1.0 when all bytes are in completed files (not in-progress) + if expected_bytes > 0 and downloaded_bytes >= expected_bytes: + progress = 1.0 return { - "downloaded_bytes": downloaded_bytes, + "downloaded_bytes": total_progress_bytes, "expected_bytes": expected_bytes, "progress": round(progress, 3), } diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts index 5831642606..be105e86be 100644 --- a/studio/frontend/src/features/chat/api/chat-api.ts +++ b/studio/frontend/src/features/chat/api/chat-api.ts @@ -105,10 +105,12 @@ export interface CachedGgufRepo { export async function getGgufDownloadProgress( repoId: string, + variant: string, expectedBytes: number, ): Promise<{ downloaded_bytes: number; expected_bytes: number; progress: number }> { const params = new URLSearchParams({ repo_id: repoId, + variant, expected_bytes: String(expectedBytes), }); const response = await authFetch(`/api/models/gguf-download-progress?${params}`); 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 07a1208287..04633ae767 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 @@ -218,7 +218,7 @@ export function useChatModelRuntime() { extraLoadingDescription ?? null, isDownloaded ? "Loading cached model into memory." - : "This may include downloading. Large models can take a while.", + : "Downloading and loading model. Large models can take a while.", ] .filter(Boolean) .join(" "); @@ -320,7 +320,7 @@ export function useChatModelRuntime() { return; } try { - const prog = await getGgufDownloadProgress(modelId, expectedBytes); + 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);