Update app.py

app.py

@@ -1,50 +1,64 @@
 import os
 import sys
 import subprocess
-import time
-from huggingface_hub import hf_hub_download
-
-BASE_DIR = os.getcwd()
-WEIGHTS_DIR = os.path.join(BASE_DIR, "weights")
-OUTPUT_BASEPATH = os.path.join(BASE_DIR, "results-poor")
-
-# Download specific files from Hugging Face repo
-def download_checkpoints():
-    os.makedirs(WEIGHTS_DIR, exist_ok=True)
-    print("⬇️  Downloading necessary checkpoint files...")
-
-    try:
-        # Download FP8 checkpoint
-        checkpoint_fp8 = hf_hub_download(
-            repo_id="tencent/HunyuanVideo-Avatar",
-            filename="ckpts/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states_fp8.pt",
-            cache_dir=WEIGHTS_DIR,
-            local_dir=WEIGHTS_DIR,
-            local_dir_use_symlinks=False
-        )
-
-        # Download normal checkpoint for Flask/Gradio UI
-        checkpoint = hf_hub_download(
-            repo_id="tencent/HunyuanVideo-Avatar",
-            filename="ckpts/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt",
-            cache_dir=WEIGHTS_DIR,
-            local_dir=WEIGHTS_DIR,
-            local_dir_use_symlinks=False
-        )
-
-        return checkpoint, checkpoint_fp8
-
-    except Exception as e:
-        print(f"❌ Error during checkpoint download: {e}")
+from huggingface_hub import hf_hub_download, hf_hub_list
+from pathlib import Path
+
+MODEL_REPO = "tencent/HunyuanVideo-Avatar"
+BASE_DIR = Path.cwd()
+WEIGHTS_DIR = BASE_DIR / "weights"
+CKPTS_DIR = WEIGHTS_DIR / "ckpts"
+
+# The root folder in the repo you want to download files from
+REPO_ROOT_FOLDER = "ckpts"
+
+# List all files under the folder 'ckpts' in the repo (recursively)
+def list_ckpt_files():
+    print(f"🔍 Listing files under '{REPO_ROOT_FOLDER}' in repo {MODEL_REPO}...")
+    files = hf_hub_list(repo_id=MODEL_REPO, repo_type="model", revision="main", allow_regex=f"^{REPO_ROOT_FOLDER}/.*")
+    return files
+
+# Download all files in the list into weights/ckpts preserving folder structure
+def download_ckpts(files):
+    print(f"⬇️ Downloading {len(files)} files under '{REPO_ROOT_FOLDER}' ...")
+    for file in files:
+        # Target local path preserving folder structure after 'ckpts/'
+        relative_path = Path(file).relative_to(REPO_ROOT_FOLDER)
+        local_path = CKPTS_DIR / relative_path
+        
+        # Make sure directory exists
+        local_path.parent.mkdir(parents=True, exist_ok=True)
+
+        # Download the file if not present
+        if not local_path.is_file():
+            print(f"Downloading {file} -> {local_path}")
+            hf_hub_download(
+                repo_id=MODEL_REPO,
+                filename=file,
+                repo_type="model",
+                cache_dir=str(WEIGHTS_DIR),
+                local_dir=str(WEIGHTS_DIR),
+                local_dir_use_symlinks=False,
+                force_filename=str(local_path.name),
+                revision="main",
+            )
+        else:
+            print(f"File {local_path} already exists, skipping download.")
+    print("✅ All checkpoint files downloaded.")
+
+def run_sample_gpu_poor():
+    checkpoint_fp8 = CKPTS_DIR / "hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states_fp8.pt"
+    if not checkpoint_fp8.is_file():
+        print(f"❌ Checkpoint file {checkpoint_fp8} not found!")
         sys.exit(1)
 
-def run_sample_gpu_poor(checkpoint_fp8_path):
-    print("🎬 Running sample_gpu_poor.py...")
+    OUTPUT_BASEPATH = BASE_DIR / "results-poor"
+    OUTPUT_BASEPATH.mkdir(parents=True, exist_ok=True)
 
     cmd = [
         "python3", "hymm_sp/sample_gpu_poor.py",
         "--input", "assets/test.csv",
-        "--ckpt", checkpoint_fp8_path,
+        "--ckpt", str(checkpoint_fp8),
         "--sample-n-frames", "129",
         "--seed", "128",
         "--image-size", "704",
@@ -52,7 +66,7 @@ def run_sample_gpu_poor(checkpoint_fp8_path):
         "--infer-steps", "50",
         "--use-deepcache", "1",
         "--flow-shift-eval-video", "5.0",
-        "--save-path", OUTPUT_BASEPATH,
+        "--save-path", str(OUTPUT_BASEPATH),
         "--use-fp8",
         "--cpu-offload",
         "--infer-min"
@@ -60,53 +74,25 @@ def run_sample_gpu_poor(checkpoint_fp8_path):
 
     env = os.environ.copy()
     env["PYTHONPATH"] = "./"
-    env["MODEL_BASE"] = WEIGHTS_DIR
+    env["MODEL_BASE"] = str(WEIGHTS_DIR)
     env["CPU_OFFLOAD"] = "1"
     env["CUDA_VISIBLE_DEVICES"] = "0"
 
+    print("🎬 Running sample_gpu_poor.py...")
     proc = subprocess.run(cmd, env=env)
     if proc.returncode != 0:
         print("❌ sample_gpu_poor.py failed.")
         sys.exit(1)
-
     print("✅ sample_gpu_poor.py completed successfully.")
 
-def run_flask_audio(checkpoint_path):
-    print("🚀 Starting flask_audio.py...")
-    cmd = [
-        "torchrun",
-        "--nnodes=1",
-        "--nproc_per_node=1",
-        "--master_port=29605",
-        "hymm_gradio/flask_audio.py",
-        "--input", "assets/test.csv",
-        "--ckpt", checkpoint_path,
-        "--sample-n-frames", "129",
-        "--seed", "128",
-        "--image-size", "704",
-        "--cfg-scale", "7.5",
-        "--infer-steps", "50",
-        "--use-deepcache", "1",
-        "--flow-shift-eval-video", "5.0"
-    ]
-    subprocess.Popen(cmd)
-
-def run_gradio_ui():
-    print("🟢 Starting gradio_audio.py UI...")
-    cmd = ["python3", "hymm_gradio/gradio_audio.py"]
-    subprocess.Popen(cmd)
-
 def main():
-    # Step 1: Download only needed files from Hugging Face repo
-    checkpoint, checkpoint_fp8 = download_checkpoints()
-
-    # Step 2: Run poor sample video generation
-    run_sample_gpu_poor(checkpoint_fp8)
+    files = list_ckpt_files()
+    if not files:
+        print("❌ No checkpoint files found in repo under ckpts folder.")
+        sys.exit(1)
 
-    # Step 3: Launch Flask + Gradio UIs
-    run_flask_audio(checkpoint)
-    time.sleep(5)
-    run_gradio_ui()
+    download_ckpts(files)
+    run_sample_gpu_poor()
 
 if __name__ == "__main__":
     main()