fix progress of initial models download

Nowadays huggingface mostly uses xet protocol for downloads,
but we only monkey-patched the pure http-based flow.
Progress updates are driven by the patch, hence there were none
for xet based downloads.

In addition, increase granularity so we report downloaded bytes
before reaching the first Mb.

Author: alexger

backend/api_types.py

@@ -149,14 +149,14 @@ class ModelsStatusResponse(BaseModel):
 class DownloadProgressResponse(BaseModel):
     status: str
     current_downloading_file: ModelFileType | None
-    current_file_progress: int
-    total_progress: int
+    current_file_progress: float
+    total_progress: float
     total_downloaded_bytes: int
     expected_total_bytes: int
     completed_files: set[ModelFileType]
     all_files: set[ModelFileType]
     error: str | None
-    speed_mbps: int
+    speed_bytes_per_sec: float
 
 
 class SuggestGapPromptResponse(BaseModel):

backend/handlers/download_handler.py

@@ -72,7 +72,7 @@ def start_file(self, file_type: ModelFileType, target: str) -> None:
             file_type=file_type,
             target_path=target,
             downloaded_bytes=0,
-            speed_mbps=0.0,
+            speed_bytes_per_sec=0.0,
         )
 
     @with_state_lock
@@ -87,15 +87,15 @@ def finish_download(self) -> None:
         self.state.downloading_session = None
 
     @with_state_lock
-    def update_file_progress(self, file_type: ModelFileType, downloaded: int, speed_mbps: float) -> None:
+    def update_file_progress(self, file_type: ModelFileType, downloaded: int, speed_bytes_per_sec: float) -> None:
         session = self.state.downloading_session
         if session is None:
             return
         rf = session.current_running_file
         if rf is None or rf.file_type != file_type:
             return
         rf.downloaded_bytes = downloaded
-        rf.speed_mbps = speed_mbps
+        rf.speed_bytes_per_sec = speed_bytes_per_sec
 
     @with_state_lock
     def fail_download(self, error: str) -> None:
@@ -106,12 +106,25 @@ def fail_download(self, error: str) -> None:
             self.state.downloading_session = None
 
     def _make_progress_callback(self, file_type: ModelFileType) -> Callable[[int], None]:
-        start_time = time.monotonic()
+        last_sample_time = time.monotonic()
+        last_sample_bytes = 0
+        smoothed_speed = 0.0
 
         def on_progress(downloaded: int) -> None:
-            elapsed = time.monotonic() - start_time
-            speed_mbps = (downloaded / elapsed / (1024 * 1024)) if elapsed > 0 else 0.0
-            self.update_file_progress(file_type, downloaded, speed_mbps)
+            nonlocal last_sample_time, last_sample_bytes, smoothed_speed
+            now = time.monotonic()
+            elapsed = now - last_sample_time
+            if elapsed >= 1.0:
+                instant_speed = (downloaded - last_sample_bytes) / elapsed
+                # EWMA: weight new sample at 30%, keep 70% of previous.
+                # On first sample (smoothed_speed == 0) use instant value.
+                if smoothed_speed == 0.0:
+                    smoothed_speed = instant_speed
+                else:
+                    smoothed_speed = 0.3 * instant_speed + 0.7 * smoothed_speed
+                last_sample_time = now
+                last_sample_bytes = downloaded
+            self.update_file_progress(file_type, downloaded, smoothed_speed)
 
         return on_progress
 
@@ -130,15 +143,15 @@ def get_download_progress(self, session_id: str) -> DownloadProgressResponse:
                 self.config.spec_for(ft).expected_size_bytes for ft in session.files_to_download
             )
 
-            current_file_progress = 0
+            current_file_progress = 0.0
             if rf is not None:
                 spec = self.config.spec_for(rf.file_type)
                 if spec.expected_size_bytes > 0:
-                    current_file_progress = min(99, int(rf.downloaded_bytes / spec.expected_size_bytes * 100))
+                    current_file_progress = min(99.0, rf.downloaded_bytes / spec.expected_size_bytes * 100)
 
-            total_progress = 0
+            total_progress = 0.0
             if expected_total_bytes > 0:
-                total_progress = min(99, int(total_downloaded / expected_total_bytes * 100))
+                total_progress = min(99.0, total_downloaded / expected_total_bytes * 100)
 
             return DownloadProgressResponse(
                 status="downloading",
@@ -149,7 +162,7 @@ def get_download_progress(self, session_id: str) -> DownloadProgressResponse:
                 expected_total_bytes=expected_total_bytes,
                 completed_files=set(session.completed_files),
                 all_files=set(session.files_to_download),
-                speed_mbps=int(rf.speed_mbps) if rf else 0,
+                speed_bytes_per_sec=rf.speed_bytes_per_sec if rf else 0.0,
                 error=None,
             )
 
@@ -159,26 +172,26 @@ def get_download_progress(self, session_id: str) -> DownloadProgressResponse:
                 return DownloadProgressResponse(
                     status="complete",
                     current_downloading_file=None,
-                    current_file_progress=100,
-                    total_progress=100,
+                    current_file_progress=100.0,
+                    total_progress=100.0,
                     total_downloaded_bytes=0,
                     expected_total_bytes=0,
                     completed_files=set(),
                     all_files=set(),
-                    speed_mbps=0,
+                    speed_bytes_per_sec=0.0,
                     error=None,
                 )
             else:
                 return DownloadProgressResponse(
                     status="error",
                     current_downloading_file=None,
-                    current_file_progress=0,
-                    total_progress=0,
+                    current_file_progress=0.0,
+                    total_progress=0.0,
                     total_downloaded_bytes=0,
                     expected_total_bytes=0,
                     completed_files=set(),
                     all_files=set(),
-                    speed_mbps=0,
+                    speed_bytes_per_sec=0.0,
                     error=result,
                 )
 

backend/services/model_downloader/hugging_face_downloader.py

@@ -40,14 +40,15 @@ def update(self, n: float | int | None = 1) -> bool | None:  # type: ignore[repo
 
 
 @contextlib.contextmanager
-def _patch_http_get_progress(callback: Callable[[int], None]) -> Iterator[None]:
+def _patch_download_progress(callback: Callable[[int], None]) -> Iterator[None]:
     """Temporarily monkey-patch ``huggingface_hub.file_download.http_get``
-    to inject a custom tqdm bar that forwards progress to *callback*.
+    and ``xet_get`` to inject a custom tqdm bar that forwards progress to
+    *callback*.
 
     ``hf_hub_download`` does not expose a ``tqdm_class`` parameter (unlike
-    ``snapshot_download``), but its internal ``http_get`` accepts a private
-    ``_tqdm_bar`` kwarg.  We wrap ``http_get`` to inject our own bar when
-    the caller hasn't already provided one.
+    ``snapshot_download``), but its internal ``http_get`` and ``xet_get``
+    both accept a private ``_tqdm_bar`` kwarg.  We wrap them to inject our
+    own bar when the caller hasn't already provided one.
 
     See ``test_http_get_accepts_tqdm_bar`` — if that test breaks after a
     huggingface_hub upgrade, this patch needs to be revisited.
@@ -60,8 +61,20 @@ def _wrapped_http_get(*args: Any, **kwargs: Any) -> None:
             kwargs["_tqdm_bar"] = tqdm_cls(disable=True)
         return original_http_get(*args, **kwargs)
 
+    xet_get_fn: Callable[..., Any] | None = getattr(file_download, "xet_get", None)
+
+    def _wrapped_xet_get(*args: Any, **kwargs: Any) -> None:
+        if kwargs.get("_tqdm_bar") is None:
+            kwargs["_tqdm_bar"] = tqdm_cls(disable=True)
+        assert xet_get_fn is not None
+        return xet_get_fn(*args, **kwargs)
+
     with patch.object(file_download, "http_get", _wrapped_http_get):
-        yield
+        if xet_get_fn is not None:
+            with patch.object(file_download, "xet_get", _wrapped_xet_get):
+                yield
+        else:
+            yield
 
 
 class HuggingFaceDownloader:
@@ -74,7 +87,7 @@ def download_file(
         local_dir: str,
         on_progress: Callable[[int], None] | None = None,
     ) -> Path:
-        ctx = _patch_http_get_progress(on_progress) if on_progress is not None else contextlib.nullcontext()
+        ctx = _patch_download_progress(on_progress) if on_progress is not None else contextlib.nullcontext()
         with ctx:
             path: str = hf_hub_download(repo_id=repo_id, filename=filename, local_dir=local_dir)
         return Path(path)
@@ -85,7 +98,7 @@ def download_snapshot(
         local_dir: str,
         on_progress: Callable[[int], None] | None = None,
     ) -> Path:
-        ctx = _patch_http_get_progress(on_progress) if on_progress is not None else contextlib.nullcontext()
+        ctx = _patch_download_progress(on_progress) if on_progress is not None else contextlib.nullcontext()
         with ctx:
             path: str = snapshot_download(repo_id=repo_id, local_dir=local_dir)
         return Path(path)

backend/state/app_state_types.py

@@ -43,7 +43,7 @@ class FileDownloadRunning:
     file_type: ModelFileType
     target_path: str
     downloaded_bytes: int
-    speed_mbps: float
+    speed_bytes_per_sec: float
 
 
 @dataclass

backend/tests/test_models.py

@@ -105,7 +105,7 @@ def test_active(self, client, test_state):
                 file_type="checkpoint",
                 target_path="checkpoint",
                 downloaded_bytes=5_000_000_000,
-                speed_mbps=50,
+                speed_bytes_per_sec=50_000_000.0,
             ),
             files_to_download={"checkpoint"},
             completed_files=set(),
@@ -344,9 +344,10 @@ def test_failed_download_cleans_up_downloading_dir(self, test_state):
 class TestHuggingFaceInternals:
     """Guard tests for huggingface_hub internals we rely on.
 
-    We monkey-patch ``file_download.http_get`` to inject a custom tqdm bar
-    for progress tracking during ``hf_hub_download`` (which has no public
-    ``tqdm_class`` parameter, unlike ``snapshot_download``).
+    We monkey-patch ``file_download.http_get`` and ``file_download.xet_get``
+    to inject a custom tqdm bar for progress tracking during
+    ``hf_hub_download`` (which has no public ``tqdm_class`` parameter,
+    unlike ``snapshot_download``).
 
     If these tests break after a huggingface_hub upgrade, the internal API
     has changed.  Find an alternative approach and raise to a developer.
@@ -363,3 +364,12 @@ def test_http_get_accepts_tqdm_bar(self):
         assert "_tqdm_bar" in sig.parameters, (
             "file_download.http_get no longer accepts _tqdm_bar — progress patch for hf_hub_download is broken"
         )
+
+    def test_xet_get_accepts_tqdm_bar(self):
+        xet_get = getattr(file_download, "xet_get", None)
+        if xet_get is None:
+            return  # xet_get not present in this version; patch skips it gracefully
+        sig = inspect.signature(xet_get)
+        assert "_tqdm_bar" in sig.parameters, (
+            "file_download.xet_get no longer accepts _tqdm_bar — progress patch for xet downloads is broken"
+        )

backend/tests/test_response_models.py

@@ -39,7 +39,7 @@ def test_snake_case_keys(self, client, test_state):
                 file_type="checkpoint",
                 target_path="checkpoint",
                 downloaded_bytes=5_000_000_000,
-                speed_mbps=50,
+                speed_bytes_per_sec=50_000_000.0,
             ),
             files_to_download={"checkpoint"},
             completed_files=set(),
@@ -60,7 +60,7 @@ def test_snake_case_keys(self, client, test_state):
             "completed_files",
             "all_files",
             "error",
-            "speed_mbps",
+            "speed_bytes_per_sec",
         }
         assert set(data.keys()) == expected_keys
 

frontend/components/FirstRunSetup.tsx

@@ -22,7 +22,7 @@ interface DownloadProgress {
   completed_files: string[]
   all_files: string[]
   error: string | null
-  speed_mbps: number
+  speed_bytes_per_sec: number
 }
 
 interface RequiredModelsResponse {
@@ -82,11 +82,10 @@ export function LaunchGate({
 
   // Calculate ETA based on speed and remaining bytes
   const getTimeRemaining = (): string => {
-    if (!downloadProgress || downloadProgress.speed_mbps <= 0) return '--'
+    if (!downloadProgress || downloadProgress.speed_bytes_per_sec <= 0) return '--'
     const remainingBytes = downloadProgress.expected_total_bytes - downloadProgress.total_downloaded_bytes
     if (remainingBytes <= 0) return '--'
-    const speedBytesPerSec = downloadProgress.speed_mbps * 1024 * 1024
-    const secondsRemaining = remainingBytes / speedBytesPerSec
+    const secondsRemaining = remainingBytes / downloadProgress.speed_bytes_per_sec
     return formatTimeRemaining(secondsRemaining)
   }
 
@@ -701,7 +700,7 @@ export function LaunchGate({
                     {(downloadProgress?.total_progress || 0) > 85 ? 'Installing...' : 'Downloading...'}
                   </span>
                   <span style={{ fontSize: 13, color: '#A98BD9', fontWeight: 600 }}>
-                    {downloadProgress?.total_progress || 0}%
+                    {Math.round(downloadProgress?.total_progress || 0)}%
                   </span>
                 </div>
 
@@ -739,17 +738,17 @@ export function LaunchGate({
 
                   {/* Speed and ETA */}
                   <div style={{ display: 'flex', gap: 16, marginLeft: 16, flexShrink: 0 }}>
-                    {downloadProgress && downloadProgress.speed_mbps > 0 && (
+                    {downloadProgress && downloadProgress.speed_bytes_per_sec > 0 && (
                       <span style={{ color: '#6D28D9', fontWeight: 500 }}>
-                        {downloadProgress.speed_mbps.toFixed(1)} MB/s
+                        {(downloadProgress.speed_bytes_per_sec / (1024 * 1024)).toFixed(1)} MB/s
                       </span>
                     )}
                     {downloadProgress && downloadProgress.expected_total_bytes > 0 && (
                       <span>
                         {formatBytes(downloadProgress.total_downloaded_bytes)} / {formatBytes(downloadProgress.expected_total_bytes)}
                       </span>
                     )}
-                    {downloadProgress && downloadProgress.speed_mbps > 0 && (
+                    {downloadProgress && downloadProgress.speed_bytes_per_sec > 0 && (
                       <span>
                         ETA: {getTimeRemaining()}
                       </span>