Update hf_model.py

hf_model.py

@@ -131,32 +131,32 @@ class UniBioTransferModel(LatentDiffusion, PyTorchModelHubMixin):
         arcface_path = cache_dir / "Other_dependencies" / "arcface" / "model_ir_se50.pth"
         face_parsing_path = cache_dir / "Other_dependencies" / "face_parsing" / "79999_iter.pth"
         
-        def _download_file(repo, filename, local_path):
-            local_path = Path(local_path)
-            local_path.parent.mkdir(parents=True, exist_ok=True)
+        def _download_file(repo, filename, target_root):
+            target_root = Path(target_root)
+            target_root.mkdir(parents=True, exist_ok=True)
             print(f"Downloading {filename} from {repo}...")
             token = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN")
             hf_hub_download(
                 repo_id=repo,
                 filename=filename,
-                local_dir=str(local_path.parent),
+                local_dir=str(target_root),
                 local_dir_use_symlinks=False,
                 token=token,
             )
-        
+
         if not ckpt_path.exists():
-            _download_file(repo_id, "checkpoints/pretrained.ckpt", ckpt_path)
+            _download_file(repo_id, "checkpoints/pretrained.ckpt", cache_dir)
         if not json_path.exists():
-            _download_file(repo_id, "checkpoints/pretrained.json", json_path)
-        
+            _download_file(repo_id, "checkpoints/pretrained.json", cache_dir)
+
         if download_sd14 and not sd14_path.exists():
-            _download_file(SD14_REPO, SD14_FILENAME, sd14_path)
+            _download_file(SD14_REPO, SD14_FILENAME, cache_dir)
         
         if download_deps:
             if not arcface_path.exists():
-                _download_file(repo_id, "Other_dependencies/arcface/model_ir_se50.pth", arcface_path)
+                _download_file(repo_id, "Other_dependencies/arcface/model_ir_se50.pth", cache_dir)
             if not face_parsing_path.exists():
-                _download_file(repo_id, "Other_dependencies/face_parsing/79999_iter.pth", face_parsing_path)
+                _download_file(repo_id, "Other_dependencies/face_parsing/79999_iter.pth", cache_dir)
         
         seed_everything(42)