Update app.py

app.py

@@ -5,9 +5,11 @@ import time
 from huggingface_hub import snapshot_download
 
 MODEL_REPO = "tencent/HunyuanVideo-Avatar"
-BASE_DIR = os.getcwd()  # current working directory absolute path
+BASE_DIR = os.getcwd()
 WEIGHTS_DIR = os.path.join(BASE_DIR, "weights")
+OUTPUT_BASEPATH = os.path.join(BASE_DIR, "results-poor")
 
+# Specific checkpoint to use in the poor sampling run
 CHECKPOINT_FILE = os.path.join(
     WEIGHTS_DIR,
     "ckpts",
@@ -15,9 +17,16 @@ CHECKPOINT_FILE = os.path.join(
     "transformers",
     "mp_rank_00_model_states.pt"
 )
+CHECKPOINT_FP8_FILE = os.path.join(
+    WEIGHTS_DIR,
+    "ckpts",
+    "hunyuan-video-t2v-720p",
+    "transformers",
+    "mp_rank_00_model_states_fp8.pt"
+)
 
 def download_model():
-    print("⬇️ Model not found. Downloading with snapshot_download into weights directory...")
+    print("⬇️  Model not found. Downloading with snapshot_download into weights directory...")
     os.makedirs(WEIGHTS_DIR, exist_ok=True)
 
     snapshot_download(
@@ -29,9 +38,44 @@ def download_model():
     if not os.path.isfile(CHECKPOINT_FILE):
         print(f"❌ Checkpoint file not found at {CHECKPOINT_FILE} after download.")
         sys.exit(1)
+    
+    if not os.path.isfile(CHECKPOINT_FP8_FILE):
+        print(f"❌ FP8 checkpoint file not found at {CHECKPOINT_FP8_FILE}. Cannot proceed with sample_gpu_poor.py.")
+        sys.exit(1)
 
     print("✅ Model downloaded successfully.")
 
+def run_sample_gpu_poor():
+    print("🎬 Running sample_gpu_poor.py...")
+    cmd = [
+        "python3", "hymm_sp/sample_gpu_poor.py",
+        "--input", "assets/test.csv",
+        "--ckpt", CHECKPOINT_FP8_FILE,
+        "--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",
+        "--save-path", OUTPUT_BASEPATH,
+        "--use-fp8",
+        "--cpu-offload",
+        "--infer-min"
+    ]
+
+    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.")
+        sys.exit(1)
+    print("✅ sample_gpu_poor.py completed successfully.")
+
 def run_flask_audio():
     print("🚀 Starting flask_audio.py...")
     cmd = [
@@ -50,24 +94,25 @@ def run_flask_audio():
         "--use-deepcache", "1",
         "--flow-shift-eval-video", "5.0"
     ]
-    proc = subprocess.Popen(cmd)
-    return proc
+    subprocess.Popen(cmd)
 
 def run_gradio_ui():
     print("🟢 Starting gradio_audio.py UI...")
     cmd = ["python3", "hymm_gradio/gradio_audio.py"]
-    proc = subprocess.Popen(cmd)
-    return proc
+    subprocess.Popen(cmd)
 
 def main():
-    if os.path.isfile(CHECKPOINT_FILE):
+    if os.path.isfile(CHECKPOINT_FILE) and os.path.isfile(CHECKPOINT_FP8_FILE):
         print("✅ Model checkpoint already exists. Skipping download.")
     else:
         download_model()
 
-    flask_proc = run_flask_audio()
-    time.sleep(5)  # Wait a bit for flask_audio.py to start
-    gradio_proc = run_gradio_ui()
+    run_sample_gpu_poor()
+
+    # Optional: Start Flask and Gradio UIs after poor sample run
+    run_flask_audio()
+    time.sleep(5)
+    run_gradio_ui()
 
 if __name__ == "__main__":
-    main() 
+    main()