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,16 +555,36 @@ 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 _xet_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,
             endpoint=connection_info.endpoint,
             token_info=(connection_info.access_token, connection_info.expiration_unix_epoch),
             token_refresher=token_refresher,
-            progress_updater=[progress_updater],
+            progress_updater=[_xet_progress_updater],
             request_headers=xet_headers,
         )
 

tests/test_xet_download.py

@@ -354,3 +354,185 @@ 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."""
+
+    def _make_mock_total_update(self, 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 test_callback_uses_two_arg_signature(self, tmp_path):
+        """Verify xet_get passes a 2-arg callback to download_files, triggering
+        xet-core's fine-grained network-level progress dispatch."""
+        incomplete_path = tmp_path / "test_file.bin"
+        incomplete_path.touch()
+
+        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
+        )
+
+        with patch("huggingface_hub.file_download.refresh_xet_connection_info", return_value=connection_info):
+            with patch("hf_xet.download_files") as mock_download:
+                xet_get(
+                    incomplete_path=incomplete_path,
+                    xet_file_data=xet_file_data,
+                    headers={"authorization": "Bearer token"},
+                    expected_size=1000,
+                )
+
+                mock_download.assert_called_once()
+                callbacks = mock_download.call_args.kwargs["progress_updater"]
+                assert len(callbacks) == 1
+                callback = callbacks[0]
+
+                # 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
+
+    def test_progress_bar_scales_network_to_file_size(self, tmp_path):
+        """When transfer bytes differ from file size, the progress bar should
+        scale to expected_size so it always reaches 100%."""
+        incomplete_path = tmp_path / "test_file.bin"
+        incomplete_path.touch()
+
+        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
+        )
+
+        expected_size = 10_000
+        # Simulate xet transferring fewer bytes than file size (deduplication)
+        transfer_total = 5_000
+
+        captured_callback = None
+
+        def capture_callback(*args, **kwargs):
+            nonlocal captured_callback
+            captured_callback = kwargs["progress_updater"][0]
+
+        with patch("huggingface_hub.file_download.refresh_xet_connection_info", return_value=connection_info):
+            with patch("hf_xet.download_files", side_effect=capture_callback):
+                with patch("huggingface_hub.file_download._get_progress_bar_context") as mock_progress_cm:
+                    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)
+
+                    xet_get(
+                        incomplete_path=incomplete_path,
+                        xet_file_data=xet_file_data,
+                        headers={"authorization": "Bearer token"},
+                        expected_size=expected_size,
+                    )
+
+        assert captured_callback is not None
+
+        # Simulate 5 updates of 1000 transfer bytes each (total: 5000 transfer bytes)
+        for i in range(5):
+            total_update = self._make_mock_total_update(
+                transfer_increment=1000,
+                transfer_total=transfer_total,
+            )
+
+            # Simulate tqdm advancing bar.n after each update call
+            def update_side_effect(n):
+                mock_bar.n += n
+
+            mock_bar.update = Mock(side_effect=update_side_effect)
+            captured_callback(total_update, [])
+
+        # After transferring 5000/5000 bytes, bar should be at expected_size (10000)
+        assert mock_bar.n == expected_size
+
+    def test_progress_bar_capped_at_expected_size(self, tmp_path):
+        """Progress bar should never exceed expected_size."""
+        incomplete_path = tmp_path / "test_file.bin"
+        incomplete_path.touch()
+
+        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
+        )
+
+        expected_size = 1000
+        transfer_total = 1000
+
+        captured_callback = None
+
+        def capture_callback(*args, **kwargs):
+            nonlocal captured_callback
+            captured_callback = kwargs["progress_updater"][0]
+
+        with patch("huggingface_hub.file_download.refresh_xet_connection_info", return_value=connection_info):
+            with patch("hf_xet.download_files", side_effect=capture_callback):
+                with patch("huggingface_hub.file_download._get_progress_bar_context") as mock_progress_cm:
+                    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)
+
+                    xet_get(
+                        incomplete_path=incomplete_path,
+                        xet_file_data=xet_file_data,
+                        headers={"authorization": "Bearer token"},
+                        expected_size=expected_size,
+                    )
+
+        assert captured_callback is not None
+
+        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=transfer_total)
+        captured_callback(total_update, [])
+
+        assert mock_bar.n <= expected_size
+
+    def test_zero_increment_skipped(self, tmp_path):
+        """Zero-increment updates should not call progress.update."""
+        incomplete_path = tmp_path / "test_file.bin"
+        incomplete_path.touch()
+
+        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
+        )
+
+        captured_callback = None
+
+        def capture_callback(*args, **kwargs):
+            nonlocal captured_callback
+            captured_callback = kwargs["progress_updater"][0]
+
+        with patch("huggingface_hub.file_download.refresh_xet_connection_info", return_value=connection_info):
+            with patch("hf_xet.download_files", side_effect=capture_callback):
+                with patch("huggingface_hub.file_download._get_progress_bar_context") as mock_progress_cm:
+                    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)
+
+                    xet_get(
+                        incomplete_path=incomplete_path,
+                        xet_file_data=xet_file_data,
+                        headers={"authorization": "Bearer token"},
+                        expected_size=1000,
+                    )
+
+        assert captured_callback is not None
+
+        total_update = self._make_mock_total_update(transfer_increment=0, transfer_total=1000)
+        captured_callback(total_update, [])
+
+        mock_bar.update.assert_not_called()