fix: use fine-grained xet-core callback for smoother tqdm progress

Switch xet_get() from a 1-arg to 2-arg callback signature so xet-core
provides network-level progress updates (~200KB) instead of coarse
disk-write updates (~8MB). Network transfer progress is scaled to file
size so the bar always tracks 0-100% correctly.

Fixes #4058

Author: tobocop2

src/huggingface_hub/file_download.py

@@ -9,7 +9,7 @@
 import warnings
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, BinaryIO, Literal, NoReturn, overload
+from typing import TYPE_CHECKING, Any, BinaryIO, Literal, NoReturn, overload
 from urllib.parse import quote, urlparse
 
 import httpx
@@ -56,6 +56,10 @@
 from .utils.tqdm import _get_progress_bar_context
 
 
+if TYPE_CHECKING:
+    from hf_xet import PyItemProgressUpdate, PyTotalProgressUpdate
+
+
 logger = logging.get_logger(__name__)
 
 # Return value when trying to load a file from cache but the file does not exist in the distant repo.
@@ -551,9 +555,29 @@ def token_refresher() -> tuple[str, int]:
     xet_headers.pop("authorization", None)
 
     with progress_cm as progress:
-
-        def progress_updater(progress_bytes: float):
-            progress.update(progress_bytes)
+        # Use 2-arg callback signature so xet-core provides fine-grained
+        # network-level progress (~200KB) instead of coarse disk-write progress (~8MB).
+        # Scale network transfer fraction to expected file size so the bar
+        # always tracks 0-100% correctly even when transfer bytes differ from
+        # file size (e.g. due to xet deduplication).
+        _cumulative_transfer = 0
+
+        def progress_updater(
+            total_update: "PyTotalProgressUpdate",
+            item_updates: "list[PyItemProgressUpdate]",
+        ) -> None:
+            nonlocal _cumulative_transfer
+            increment = total_update.total_transfer_bytes_completion_increment
+            if increment > 0:
+                _cumulative_transfer += increment
+                transfer_total = total_update.total_transfer_bytes
+                if transfer_total and transfer_total > 0 and expected_size:
+                    new_n = min(int(_cumulative_transfer / transfer_total * expected_size), expected_size)
+                    advance = new_n - progress.n
+                    if advance > 0:
+                        progress.update(advance)
+                else:
+                    progress.update(int(increment))
 
         download_files(
             xet_download_info,

tests/test_xet_download.py

@@ -354,3 +354,143 @@ def token_refresher() -> tuple[str, int]:
         )
 
         assert os.path.exists(file_path)
+
+
+@requires("hf_xet")
+class TestXetProgressGranularity:
+    """Test that xet_get uses the fine-grained 2-arg callback for tqdm progress."""
+
+    _XET_FILE_DATA = XetFileData(file_hash="mock_hash", refresh_route="mock/route")
+    _CONNECTION_INFO = XetConnectionInfo(
+        endpoint="mock_endpoint", access_token="mock_token", expiration_unix_epoch=9999999999
+    )
+
+    @staticmethod
+    def _make_mock_total_update(transfer_increment: float, transfer_total: float):
+        """Create a mock PyTotalProgressUpdate with the given transfer fields."""
+        update = Mock()
+        update.total_transfer_bytes_completion_increment = transfer_increment
+        update.total_transfer_bytes = transfer_total
+        return update
+
+    def _call_xet_get_and_capture(self, tmp_path, mock_download, expected_size=1000, mock_progress_cm=None):
+        """Call xet_get and return the captured progress callback."""
+        incomplete_path = tmp_path / "test_file.bin"
+        incomplete_path.touch()
+
+        captured = {}
+
+        def capture(*args, **kwargs):
+            captured["callback"] = kwargs["progress_updater"][0]
+
+        mock_download.side_effect = capture
+
+        if mock_progress_cm is not None:
+            mock_bar = Mock()
+            mock_bar.n = 0
+            mock_progress_cm.return_value.__enter__ = Mock(return_value=mock_bar)
+            mock_progress_cm.return_value.__exit__ = Mock(return_value=False)
+        else:
+            mock_bar = None
+
+        xet_get(
+            incomplete_path=incomplete_path,
+            xet_file_data=self._XET_FILE_DATA,
+            headers={"authorization": "Bearer token"},
+            expected_size=expected_size,
+        )
+
+        return captured["callback"], mock_bar
+
+    @patch("huggingface_hub.file_download._get_progress_bar_context")
+    @patch("hf_xet.download_files")
+    @patch(
+        "huggingface_hub.file_download.refresh_xet_connection_info",
+        return_value=XetConnectionInfo(
+            endpoint="mock_endpoint", access_token="mock_token", expiration_unix_epoch=9999999999
+        ),
+    )
+    def test_callback_uses_two_arg_signature(self, _mock_conn, mock_download, mock_progress_cm, tmp_path):
+        """Verify xet_get passes a 2-arg callback to download_files, triggering
+        xet-core's fine-grained network-level progress dispatch."""
+        callback, _ = self._call_xet_get_and_capture(tmp_path, mock_download, mock_progress_cm=mock_progress_cm)
+
+        # Call with 2 args (total_update, item_updates) to confirm it accepts them.
+        # A 1-arg callback would raise TypeError here.
+        total_update = self._make_mock_total_update(transfer_increment=200, transfer_total=1000)
+        callback(total_update, [])  # should not raise
+
+    @patch("huggingface_hub.file_download._get_progress_bar_context")
+    @patch("hf_xet.download_files")
+    @patch(
+        "huggingface_hub.file_download.refresh_xet_connection_info",
+        return_value=XetConnectionInfo(
+            endpoint="mock_endpoint", access_token="mock_token", expiration_unix_epoch=9999999999
+        ),
+    )
+    def test_progress_bar_scales_network_to_file_size(self, _mock_conn, mock_download, mock_progress_cm, tmp_path):
+        """When transfer bytes differ from file size, the progress bar should
+        scale to expected_size so it always reaches 100%."""
+        expected_size = 10_000
+        transfer_total = 5_000  # fewer bytes due to deduplication
+
+        callback, mock_bar = self._call_xet_get_and_capture(
+            tmp_path, mock_download, expected_size=expected_size, mock_progress_cm=mock_progress_cm
+        )
+
+        def update_side_effect(n):
+            mock_bar.n += n
+
+        mock_bar.update = Mock(side_effect=update_side_effect)
+
+        # Simulate 5 updates of 1000 transfer bytes each (total: 5000 transfer bytes)
+        for _ in range(5):
+            total_update = self._make_mock_total_update(transfer_increment=1000, transfer_total=transfer_total)
+            callback(total_update, [])
+
+        # After transferring 5000/5000 bytes, bar should be at expected_size (10000)
+        assert mock_bar.n == expected_size
+
+    @patch("huggingface_hub.file_download._get_progress_bar_context")
+    @patch("hf_xet.download_files")
+    @patch(
+        "huggingface_hub.file_download.refresh_xet_connection_info",
+        return_value=XetConnectionInfo(
+            endpoint="mock_endpoint", access_token="mock_token", expiration_unix_epoch=9999999999
+        ),
+    )
+    def test_progress_bar_capped_at_expected_size(self, _mock_conn, mock_download, mock_progress_cm, tmp_path):
+        """Progress bar should never exceed expected_size."""
+        expected_size = 1000
+
+        callback, mock_bar = self._call_xet_get_and_capture(
+            tmp_path, mock_download, expected_size=expected_size, mock_progress_cm=mock_progress_cm
+        )
+
+        def update_side_effect(n):
+            mock_bar.n += n
+
+        mock_bar.update = Mock(side_effect=update_side_effect)
+
+        # Send more transfer bytes than total (edge case)
+        total_update = self._make_mock_total_update(transfer_increment=1200, transfer_total=1000)
+        callback(total_update, [])
+
+        assert mock_bar.n <= expected_size
+
+    @patch("huggingface_hub.file_download._get_progress_bar_context")
+    @patch("hf_xet.download_files")
+    @patch(
+        "huggingface_hub.file_download.refresh_xet_connection_info",
+        return_value=XetConnectionInfo(
+            endpoint="mock_endpoint", access_token="mock_token", expiration_unix_epoch=9999999999
+        ),
+    )
+    def test_zero_increment_skipped(self, _mock_conn, mock_download, mock_progress_cm, tmp_path):
+        """Zero-increment updates should not call progress.update."""
+        callback, mock_bar = self._call_xet_get_and_capture(tmp_path, mock_download, mock_progress_cm=mock_progress_cm)
+
+        total_update = self._make_mock_total_update(transfer_increment=0, transfer_total=1000)
+        callback(total_update, [])
+
+        mock_bar.update.assert_not_called()