From 1dfba866be136828e4ad3bd639b8ce86a240bdc8 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 15 Mar 2026 07:57:04 +0000
Subject: [PATCH] studio: fix download progress -- track per-variant, include
 incomplete blobs

1. Progress endpoint now takes a variant parameter and only counts
   .gguf files matching that variant (not all files in the repo cache,
   which would include previously downloaded variants)

2. Tracks .incomplete files in HF blobs dir for in-progress single-shard
   downloads, capping at 99% until the file is fully committed

3. Fixed loading text: "Loading model..." for cached, "Downloading
   model..." for new downloads, with appropriate descriptions

4. Wording: "Downloading and loading model. Large models can take a
   while." instead of "This may include downloading."
---
 studio/backend/routes/models.py               | 31 +++++++++++++------
 .../src/features/chat/api/chat-api.ts         |  2 ++
 .../chat/hooks/use-chat-model-runtime.ts      |  4 +--
 3 files changed, 25 insertions(+), 12 deletions(-)

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);