Update app.py

app.py

@@ -11,28 +11,25 @@ import traceback
 # =========================================
 # 1. Define Hugging Face weights and paths
 # =========================================
+from huggingface_hub import hf_hub_download
 
-HF_DATASET_URL = "https://huggingface.co/datasets/roll-ai/FloVD-weights/resolve/main/ckpt"
-WEIGHT_FILES = {
-    "FVSM/FloVD_FVSM_Controlnet.pt": "FVSM/FloVD_FVSM_Controlnet.pt",
-    "OMSM/selected_blocks.safetensors": "OMSM/selected_blocks.safetensors",
-    "OMSM/pytorch_lora_weights.safetensors": "OMSM/pytorch_lora_weights.safetensors",
-    "others/depth_anything_v2_metric_hypersim_vitb.pth": "others/depth_anything_v2_metric_hypersim_vitb.pth"
-}
+HF_DATASET_REPO = "roll-ai/FloVD-weights"  # dataset repo ID
 
 def download_weights():
-    print("🔄 Downloading model weights...")
+    print("🔄 Downloading model weights via huggingface_hub...")
     for rel_path in WEIGHT_FILES.values():
-        save_path = Path("ckpt") / rel_path
-        if not save_path.exists():
-            save_path.parent.mkdir(parents=True, exist_ok=True)
-            url = f"{HF_DATASET_URL}/{rel_path}"
-            print(f"📥 Downloading {url} → {save_path}")
-            subprocess.run(["wget", "-q", "-O", str(save_path), url], check=True)
+        local_path = Path("ckpt") / rel_path
+        if not local_path.exists():
+            print(f"📥 Downloading {rel_path}")
+            hf_hub_download(
+                repo_id=HF_DATASET_REPO,
+                repo_type="dataset",
+                filename=rel_path,
+                local_dir="ckpt",
+                local_dir_use_symlinks=False,
+            )
         else:
-            print(f"✅ Already exists: {save_path}")
-
-download_weights()
+            print(f"✅ Already exists: {local_path}")
 
 from inference.flovd_demo import generate_video