Use temporary HF cache and fix tar extraction

Use a temporary directory for HuggingFace cache (tempfile.TemporaryDirectory) so HF cache folders are not created under the target_dir. Improve tar.gz handling in _handle_downloaded_file by extracting members with tar.extractfile and shutil.copyfileobj (and fall back to shutil.copy2 for non-tar files), and remove the previous symlink-resolution and explicit HF-folder removal logic. Update hf_hub_download call to use repo_id/filename named args and thread per-file rename mappings through to the extractor. Adjusted tests to reflect that the HF cache is no longer created inside the target directory and to ensure the final artifact still exists.

Author: C-Achard

dlclibrary/dlcmodelzoo/modelzoo_download.py

@@ -13,6 +13,8 @@
 import json
 import os
 import tarfile
+import shutil
+import tempfile
 from pathlib import Path
 
 from huggingface_hub import hf_hub_download
@@ -111,32 +113,27 @@ def get_available_models(dataset: str) -> list[str]:
     return list(_load_pytorch_dataset_models(dataset)["pose_models"].keys())
 
 
+
 def _handle_downloaded_file(
     file_path: str, target_dir: str, rename_mapping: dict | None = None
 ):
-    """Handle the downloaded file from HuggingFace"""
+    """Handle the downloaded file from HuggingFace cache and place the final artifact in target_dir."""
     file_name = os.path.basename(file_path)
+
     try:
         with tarfile.open(file_path, mode="r:gz") as tar:
             for member in tar:
                 if not member.isdir():
                     fname = Path(member.name).name
-                    tar.makefile(member, os.path.join(target_dir, fname))
-    except tarfile.ReadError:  # The model is a .pt file
+                    extracted_path = os.path.join(target_dir, fname)
+                    with tar.extractfile(member) as src, open(extracted_path, "wb") as dst:
+                        shutil.copyfileobj(src, dst)
+    except tarfile.ReadError:
         if rename_mapping is not None:
             file_name = rename_mapping.get(file_name, file_name)
-        if os.path.islink(file_path):
-            file_path_ = os.readlink(file_path)
-            if not os.path.isabs(file_path_):
-                file_path_ = os.path.abspath(
-                    os.path.join(os.path.dirname(file_path), file_path_)
-                )
-            file_path = file_path_
-        import shutil
         shutil.copy2(file_path, os.path.join(target_dir, file_name))
 
 
-
 def download_huggingface_model(
     model_name: str,
     target_dir: str = ".",
@@ -190,27 +187,22 @@ def download_huggingface_model(
 
     if not os.path.isabs(target_dir):
         target_dir = os.path.abspath(target_dir)
-
     os.makedirs(target_dir, exist_ok=True)
 
-    last_hf_folder = None
-
-    for url in urls:
-        url = url.split("/")
-        repo_id, targzfn = url[0] + "/" + url[1], str(url[-1])
-
-        downloaded = hf_hub_download(repo_id, targzfn, cache_dir=str(target_dir))
-
-        if isinstance(rename_mapping, str):
-            mapping = {targzfn: rename_mapping}
-        else:
-            mapping = rename_mapping
-
-        _handle_downloaded_file(downloaded, target_dir, mapping)
+    with tempfile.TemporaryDirectory(prefix="dlc_hf_") as hf_cache_dir:
+        for url in urls:
+            url = url.split("/")
+            repo_id, targzfn = url[0] + "/" + url[1], str(url[-1])
 
-        last_hf_folder = f"models--{url[0]}--{url[1]}"
+            downloaded = hf_hub_download(
+                repo_id=repo_id,
+                filename=targzfn,
+                cache_dir=hf_cache_dir,
+            )
 
-    if remove_hf_folder and last_hf_folder is not None:
-        import shutil
-        shutil.rmtree(os.path.join(target_dir, last_hf_folder), ignore_errors=True)
+            if isinstance(rename_mapping, str):
+                mapping = {targzfn: rename_mapping}
+            else:
+                mapping = rename_mapping
 
+            _handle_downloaded_file(downloaded, target_dir, mapping)

tests/test_modeldownload.py

@@ -171,10 +171,7 @@ def test_download_with_rename_mapping_for_pt(tmp_path, mock_modelzoo):
     assert not any(name.startswith("models--") for name in files)
 
 
-def test_keep_hf_folder_when_requested(tmp_path, mock_modelzoo):
-    """
-    If remove_hf_folder=False, the cache structure should still exist.
-    """
+def test_remove_hf_folder_flag_no_longer_affects_target_dir(tmp_path, mock_modelzoo):
     folder = tmp_path / "keep_cache"
     folder.mkdir()
 
@@ -185,4 +182,9 @@ def test_keep_hf_folder_when_requested(tmp_path, mock_modelzoo):
     )
 
     files = {p.name for p in folder.iterdir()}
-    assert any(name.startswith("models--") for name in files)
+
+    # Final artifact should still be there
+    assert "full_cat.pt" in files
+
+    # HF cache should no longer be created inside target_dir
+    assert not any(name.startswith("models--") for name in files)