Enable Hugging Face model URL support for remote checkpoint/config downloads

direct/utils/io.py

@@ -41,6 +41,20 @@
 USER_AGENT = "NKI-AI/direct"
 
 
+def _normalize_huggingface_url(url: str) -> str:
+    """Normalize Hugging Face page URLs to direct download URLs."""
+    parsed_url = urllib.parse.urlparse(url)
+    if parsed_url.netloc != "huggingface.co":
+        return url
+
+    path = parsed_url.path
+    if "/blob/" not in path:
+        return url
+
+    normalized_path = path.replace("/blob/", "/resolve/", 1)
+    return urllib.parse.urlunparse(parsed_url._replace(path=normalized_path))
+
+
 def read_json(fn: dict | str | pathlib.Path) -> dict:  # pragma: no cover
     """Read file and output dict, or take dict and output dict.
 
@@ -187,6 +201,7 @@ def check_integrity(fpath: str, md5: str | None = None) -> bool:  # pragma: no c
 
 def _get_redirect_url(url: str, max_hops: int = 3) -> str:  # pragma: no cover
     initial_url = url
+    url = _normalize_huggingface_url(url)
     headers = {"Method": "HEAD", "User-Agent": USER_AGENT}
 
     for _ in range(max_hops + 1):
@@ -444,6 +459,7 @@ def read_text_from_url(url, chunk_size: int = 1024):
     if not check_is_valid_url(url):
         raise ValueError(f"{url} is not a valid URL.")
 
+    url = _normalize_huggingface_url(url)
     scheme = urllib.parse.urlparse(url).scheme
     if scheme not in {"http", "https"}:
         raise ValueError(f"URL scheme not permitted: {scheme!r}")

tests/tests_utils/io_test.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 import pytest
 
-from direct.utils.io import check_is_valid_url
+from direct.utils.io import _normalize_huggingface_url, check_is_valid_url
 
 
 @pytest.mark.parametrize(
@@ -31,3 +31,24 @@ def test_check_valid_url(path, is_url):
         assert check_is_valid_url(path)
     else:
         assert not check_is_valid_url(path)
+
+
+@pytest.mark.parametrize(
+    ["url", "expected_url"],
+    [
+        (
+            "https://huggingface.co/NKI-AI/direct/blob/main/recurrentvarnet/model_148500.pt",
+            "https://huggingface.co/NKI-AI/direct/resolve/main/recurrentvarnet/model_148500.pt",
+        ),
+        (
+            "https://huggingface.co/datasets/NKI-AI/direct-dataset/blob/main/config.yaml?download=true",
+            "https://huggingface.co/datasets/NKI-AI/direct-dataset/resolve/main/config.yaml?download=true",
+        ),
+        (
+            "https://files.aiforoncology.nl/direct-project/recurrentvarnet.zip",
+            "https://files.aiforoncology.nl/direct-project/recurrentvarnet.zip",
+        ),
+    ],
+)
+def test_normalize_huggingface_url(url, expected_url):
+    assert _normalize_huggingface_url(url) == expected_url