Update app.py

app.py

@@ -1,34 +1,50 @@
 import os
+import sys
 import subprocess
 import time
 from huggingface_hub import hf_hub_download
 
-MODEL_REPO = "tencent/HunyuanVideo-Avatar"
-OUTPUT_DIR = "results-poor"
-ASSETS_CSV = "assets/test.csv"
+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():
-    print("⬇️ Downloading specific checkpoint files from Tencent repo...")
+    os.makedirs(WEIGHTS_DIR, exist_ok=True)
+    print("⬇️  Downloading necessary checkpoint files...")
 
-    fp8_ckpt = hf_hub_download(
-        repo_id=MODEL_REPO,
-        filename="ckpts/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states_fp8.pt"
-    )
+    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
+        )
 
-    fp32_ckpt = hf_hub_download(
-        repo_id=MODEL_REPO,
-        filename="ckpts/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt"
-    )
+        # 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
+        )
 
-    print("✅ Checkpoints downloaded.")
-    return fp8_ckpt, fp32_ckpt
+        return checkpoint, checkpoint_fp8
 
-def run_sample_gpu_poor(fp8_ckpt_path):
+    except Exception as e:
+        print(f"❌ Error during checkpoint download: {e}")
+        sys.exit(1)
+
+def run_sample_gpu_poor(checkpoint_fp8_path):
     print("🎬 Running sample_gpu_poor.py...")
+
     cmd = [
         "python3", "hymm_sp/sample_gpu_poor.py",
-        "--input", ASSETS_CSV,
-        "--ckpt", fp8_ckpt_path,
+        "--input", "assets/test.csv",
+        "--ckpt", checkpoint_fp8_path,
         "--sample-n-frames", "129",
         "--seed", "128",
         "--image-size", "704",
@@ -36,7 +52,7 @@ def run_sample_gpu_poor(fp8_ckpt_path):
         "--infer-steps", "50",
         "--use-deepcache", "1",
         "--flow-shift-eval-video", "5.0",
-        "--save-path", OUTPUT_DIR,
+        "--save-path", OUTPUT_BASEPATH,
         "--use-fp8",
         "--cpu-offload",
         "--infer-min"
@@ -44,21 +60,51 @@ def run_sample_gpu_poor(fp8_ckpt_path):
 
     env = os.environ.copy()
     env["PYTHONPATH"] = "./"
+    env["MODEL_BASE"] = WEIGHTS_DIR
+    env["CPU_OFFLOAD"] = "1"
     env["CUDA_VISIBLE_DEVICES"] = "0"
 
     proc = subprocess.run(cmd, env=env)
     if proc.returncode != 0:
         print("❌ sample_gpu_poor.py failed.")
-        exit(1)
+        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("🟢 Launching Gradio interface...")
-    subprocess.Popen(["python3", "hymm_gradio/gradio_audio.py"])
+    print("🟢 Starting gradio_audio.py UI...")
+    cmd = ["python3", "hymm_gradio/gradio_audio.py"]
+    subprocess.Popen(cmd)
 
 def main():
-    fp8_ckpt, _ = download_checkpoints()
-    run_sample_gpu_poor(fp8_ckpt)
+    # 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)
+
+    # Step 3: Launch Flask + Gradio UIs
+    run_flask_audio(checkpoint)
     time.sleep(5)
     run_gradio_ui()