cache method for model downloads

app.py

@@ -1,31 +1,34 @@
+import torch
+import os
+import shutil
+import subprocess
+import gradio as gr 
+import json
+import tempfile
 from huggingface_hub import snapshot_download
 
 # Download All Required Models using `snapshot_download`
 
-# Download Wan2.1-I2V-14B-480P model
-wan_model_path = snapshot_download(
-    repo_id="Wan-AI/Wan2.1-I2V-14B-480P",
-    local_dir="./weights/Wan2.1-I2V-14B-480P",
-    #local_dir_use_symlinks=False
-)
+def download_and_extract(repo_id, target_dir):
+    """
+    Downloads a model repo (cached) and copies its contents to a local target directory.
+    If the target_dir exists, it will be updated (not re-downloaded if cache is present).
+    """
+    print(f"Downloading {repo_id} into cache...")
+    snapshot_path = snapshot_download(repo_id)
+    
+    print(f"Copying files to {target_dir}...")
+    os.makedirs(target_dir, exist_ok=True)
+    shutil.copytree(snapshot_path, target_dir, dirs_exist_ok=True)
 
-# Download Chinese wav2vec2 model
-wav2vec_path = snapshot_download(
-    repo_id="TencentGameMate/chinese-wav2vec2-base",
-    local_dir="./weights/chinese-wav2vec2-base",
-    #local_dir_use_symlinks=False
-)
+    print(f"Done: {repo_id} extracted to {target_dir}")
+    return target_dir
 
-# Download MeiGen MultiTalk weights
-multitalk_path = snapshot_download(
-    repo_id="MeiGen-AI/MeiGen-MultiTalk",
-    local_dir="./weights/MeiGen-MultiTalk",
-    #local_dir_use_symlinks=False
-)
 
+wan_model_path = download_and_extract("Wan-AI/Wan2.1-I2V-14B-480P", "./weights/Wan2.1-I2V-14B-480P")
+wav2vec_path   = download_and_extract("TencentGameMate/chinese-wav2vec2-base", "./weights/chinese-wav2vec2-base")
+multitalk_path = download_and_extract("MeiGen-AI/MeiGen-MultiTalk", "./weights/MeiGen-MultiTalk")
 
-import os
-import shutil
 
 # Define paths
 base_model_dir = "./weights/Wan2.1-I2V-14B-480P"
@@ -55,7 +58,6 @@ shutil.copy2(
 print("Copied MultiTalk files into base model directory.")
 
 
-import torch
 
 # Check if CUDA-compatible GPU is available
 if torch.cuda.is_available():
@@ -81,11 +83,7 @@ GPU_TO_VRAM_PARAMS = {
 USED_VRAM_PARAMS = GPU_TO_VRAM_PARAMS[gpu_name]
 print("Using", USED_VRAM_PARAMS, "for num_persistent_param_in_dit")
 
-import subprocess
 
-import json
-import tempfile
-#import os
 
 def create_temp_input_json(prompt: str, cond_image_path: str, cond_audio_path: str) -> str:
     """
@@ -136,7 +134,6 @@ def infer(prompt, cond_image_path, cond_audio_path):
 
     return "multi_long_mediumvra_exp.mp4"
 
-import gradio as gr
 
 with gr.Blocks(title="MultiTalk Inference") as demo:
     gr.Markdown("## 🎤 MultiTalk Inference Demo")