Override xet refresh route's base URL with HF Endpoint (#3180)

* override the xet refresh route base url with hf endpoint

* simpler

* Update src/huggingface_hub/utils/_xet.py

Co-authored-by: Lucain <lucain@huggingface.co>

* add a check before replacing the base url

---------

Co-authored-by: Lucain <lucain@huggingface.co>

Author: hanouticelina

src/huggingface_hub/file_download.py

@@ -1407,6 +1407,7 @@ def get_hf_file_metadata(
     library_version: Optional[str] = None,
     user_agent: Union[Dict, str, None] = None,
     headers: Optional[Dict[str, str]] = None,
+    endpoint: Optional[str] = None,
 ) -> HfFileMetadata:
     """Fetch metadata of a file versioned on the Hub for a given url.
 
@@ -1432,6 +1433,8 @@ def get_hf_file_metadata(
             The user-agent info in the form of a dictionary or a string.
         headers (`dict`, *optional*):
             Additional headers to be sent with the request.
+        endpoint (`str`, *optional*):
+            Endpoint of the Hub. Defaults to <https://huggingface.co>.
 
     Returns:
         A [`HfFileMetadata`] object containing metadata such as location, etag, size and
@@ -1471,7 +1474,7 @@ def get_hf_file_metadata(
         size=_int_or_none(
             r.headers.get(constants.HUGGINGFACE_HEADER_X_LINKED_SIZE) or r.headers.get("Content-Length")
         ),
-        xet_file_data=parse_xet_file_data_from_response(r),  # type: ignore
+        xet_file_data=parse_xet_file_data_from_response(r, endpoint=endpoint),  # type: ignore
     )
 
 
@@ -1531,7 +1534,7 @@ def _get_metadata_or_catch_error(
         try:
             try:
                 metadata = get_hf_file_metadata(
-                    url=url, proxies=proxies, timeout=etag_timeout, headers=headers, token=token
+                    url=url, proxies=proxies, timeout=etag_timeout, headers=headers, token=token, endpoint=endpoint
                 )
             except EntryNotFoundError as http_error:
                 if storage_folder is not None and relative_filename is not None:

src/huggingface_hub/hf_api.py

@@ -5392,6 +5392,7 @@ def get_hf_file_metadata(
             library_name=self.library_name,
             library_version=self.library_version,
             user_agent=self.user_agent,
+            endpoint=self.endpoint,
         )
 
     @validate_hf_hub_args

src/huggingface_hub/utils/_xet.py

@@ -26,7 +26,9 @@ class XetConnectionInfo:
     endpoint: str
 
 
-def parse_xet_file_data_from_response(response: requests.Response) -> Optional[XetFileData]:
+def parse_xet_file_data_from_response(
+    response: requests.Response, endpoint: Optional[str] = None
+) -> Optional[XetFileData]:
     """
     Parse XET file metadata from an HTTP response.
 
@@ -52,7 +54,9 @@ def parse_xet_file_data_from_response(response: requests.Response) -> Optional[X
             refresh_route = response.headers[constants.HUGGINGFACE_HEADER_X_XET_REFRESH_ROUTE]
     except KeyError:
         return None
-
+    endpoint = endpoint if endpoint is not None else constants.ENDPOINT
+    if refresh_route.startswith(constants.HUGGINGFACE_CO_URL_HOME):
+        refresh_route = refresh_route.replace(constants.HUGGINGFACE_CO_URL_HOME.rstrip("/"), endpoint.rstrip("/"))
     return XetFileData(
         file_hash=file_hash,
         refresh_route=refresh_route,

tests/test_xet_utils.py

@@ -51,6 +51,34 @@ def test_parse_invalid_headers_file_info() -> None:
     assert parse_xet_file_data_from_response(mock_response) is None
 
 
+@pytest.mark.parametrize(
+    "refresh_route, expected_refresh_route",
+    [
+        (
+            "/api/refresh",
+            "/api/refresh",
+        ),
+        (
+            "https://huggingface.co/api/refresh",
+            "https://xet.example.com/api/refresh",
+        ),
+    ],
+)
+def test_parse_header_file_info_with_endpoint(refresh_route: str, expected_refresh_route: str) -> None:
+    mock_response = MagicMock()
+    mock_response.headers = {
+        "X-Xet-Hash": "sha256:abcdef",
+        "X-Xet-Refresh-Route": refresh_route,
+    }
+    mock_response.links = {}
+
+    file_data = parse_xet_file_data_from_response(mock_response, endpoint="https://xet.example.com")
+
+    assert file_data is not None
+    assert file_data.refresh_route == expected_refresh_route
+    assert file_data.file_hash == "sha256:abcdef"
+
+
 def test_parse_valid_headers_connection_info() -> None:
     headers = {
         "X-Xet-Cas-Url": "https://xet.example.com",