Update app.py

app.py

@@ -1,5 +1,6 @@
 import os
 os.system("pip install pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py310_cu121_pyt221/download.html")
+import shutil
 
 from huggingface_hub import snapshot_download
 
@@ -10,29 +11,35 @@ snapshot_download(
     local_dir="pretrained_models/diffposetalk"
 )
 
-flame = snapshot_download(
-    repo_id="Skywork/SkyReels-A1",
-    local_dir="pretrained_models/FLAME",
-    allow_patterns="extra_models/FLAME/**"
-)
-print(os.listdir(flame))
+base_dir = "pretrained_models"
+os.makedirs(base_dir, exist_ok=True)
 
-snapshot_download(
-    repo_id="Skywork/SkyReels-A1",
-    local_dir="pretrained_models/mediapipe",
-    allow_patterns="extra_models/mediapipe/**"
-)
-
-snapshot_download(
-    repo_id="Skywork/SkyReels-A1",
-    local_dir="pretrained_models/smirk",
-    allow_patterns="extra_models/smirk/**"
-)
+# Download FLAME, mediapipe, and smirk
+for model in ["FLAME", "mediapipe", "smirk"]:
+    # Download to a temp folder first
+    temp_dir = f"{base_dir}/{model}_temp"
+    snapshot_download(
+        repo_id="Skywork/SkyReels-A1",
+        local_dir=temp_dir,
+        allow_patterns=f"extra_models/{model}/**"
+    )
+    
+    # Move files from nested extra_models/model to the proper location
+    src_dir = f"{temp_dir}/extra_models/{model}"
+    dst_dir = f"{base_dir}/{model}"
+    os.makedirs(dst_dir, exist_ok=True)
+    
+    # Move all contents
+    for item in os.listdir(src_dir):
+        shutil.move(f"{src_dir}/{item}", f"{dst_dir}/{item}")
+    
+    # Clean up temp directory
+    shutil.rmtree(temp_dir)
 
+# Download SkyReels-A1-5B
 snapshot_download(
     repo_id="Skywork/SkyReels-A1",
-    local_dir="pretrained_models/SkyReels-A1-5B",
-    allow_patterns="SkyReels-A1-5B/**"
+    local_dir=f"{base_dir}/SkyReels-A1-5B", 
 )
 
 import gradio as gr