Update app.py

app.py

@@ -1,58 +1,34 @@
 import os
-import sys
 import subprocess
 import time
-from pathlib import Path
-from huggingface_hub import snapshot_download
+from huggingface_hub import hf_hub_download
 
-# --------------------
-# CONFIGURATION
-# --------------------
 MODEL_REPO = "tencent/HunyuanVideo-Avatar"
-HF_CACHE_DIR = Path("/home/user/.cache/huggingface/hf_cache/hunyuan_avatar")
-HF_CACHE_DIR.mkdir(parents=True, exist_ok=True)
-
-CHECKPOINT_FILE = HF_CACHE_DIR / "ckpts/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt"
-CHECKPOINT_FP8_FILE = HF_CACHE_DIR / "ckpts/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states_fp8.pt"
-
+OUTPUT_DIR = "results-poor"
 ASSETS_CSV = "assets/test.csv"
-OUTPUT_DIR = Path("results-poor")
-OUTPUT_DIR.mkdir(exist_ok=True)
 
-# --------------------
-# Download the model (if needed)
-# --------------------
-def download_model():
-    if CHECKPOINT_FILE.exists() and CHECKPOINT_FP8_FILE.exists():
-        print("✅ Model checkpoint already exists. Skipping download.")
-        return
+def download_checkpoints():
+    print("⬇️ Downloading specific checkpoint files from Tencent repo...")
 
-    print("⬇️ Downloading model into HF Space cache...")
-    snapshot_download(
+    fp8_ckpt = hf_hub_download(
         repo_id=MODEL_REPO,
-        local_dir=HF_CACHE_DIR,
-        local_dir_use_symlinks=False
+        filename="ckpts/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states_fp8.pt"
     )
 
-    if not CHECKPOINT_FILE.exists():
-        print(f"❌ Missing checkpoint: {CHECKPOINT_FILE}")
-        sys.exit(1)
-
-    if not CHECKPOINT_FP8_FILE.exists():
-        print(f"❌ Missing FP8 checkpoint: {CHECKPOINT_FP8_FILE}")
-        sys.exit(1)
+    fp32_ckpt = hf_hub_download(
+        repo_id=MODEL_REPO,
+        filename="ckpts/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt"
+    )
 
-    print("✅ Model download complete.")
+    print("✅ Checkpoints downloaded.")
+    return fp8_ckpt, fp32_ckpt
 
-# --------------------
-# Run sample_gpu_poor.py
-# --------------------
-def run_sample_gpu_poor():
+def run_sample_gpu_poor(fp8_ckpt_path):
     print("🎬 Running sample_gpu_poor.py...")
     cmd = [
         "python3", "hymm_sp/sample_gpu_poor.py",
         "--input", ASSETS_CSV,
-        "--ckpt", str(CHECKPOINT_FP8_FILE),
+        "--ckpt", fp8_ckpt_path,
         "--sample-n-frames", "129",
         "--seed", "128",
         "--image-size", "704",
@@ -60,7 +36,7 @@ def run_sample_gpu_poor():
         "--infer-steps", "50",
         "--use-deepcache", "1",
         "--flow-shift-eval-video", "5.0",
-        "--save-path", str(OUTPUT_DIR),
+        "--save-path", OUTPUT_DIR,
         "--use-fp8",
         "--cpu-offload",
         "--infer-min"
@@ -68,30 +44,21 @@ def run_sample_gpu_poor():
 
     env = os.environ.copy()
     env["PYTHONPATH"] = "./"
-    env["MODEL_BASE"] = str(HF_CACHE_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.")
-        sys.exit(1)
+        exit(1)
     print("✅ sample_gpu_poor.py completed successfully.")
 
-# --------------------
-# Optional: Start UI
-# --------------------
 def run_gradio_ui():
     print("🟢 Launching Gradio interface...")
-    cmd = ["python3", "hymm_gradio/gradio_audio.py"]
-    subprocess.Popen(cmd)
+    subprocess.Popen(["python3", "hymm_gradio/gradio_audio.py"])
 
-# --------------------
-# Entry point
-# --------------------
 def main():
-    download_model()
-    run_sample_gpu_poor()
+    fp8_ckpt, _ = download_checkpoints()
+    run_sample_gpu_poor(fp8_ckpt)
     time.sleep(5)
     run_gradio_ui()