Update app.py

app.py

@@ -1,78 +1,57 @@
+
 import os
 import sys
-from huggingface_hub import hf_hub_download, list_repo_files
 import subprocess
+import time
+from huggingface_hub import snapshot_download
 
 MODEL_REPO = "tencent/HunyuanVideo-Avatar"
 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",
+    "hunyuan-video-t2v-720p",
+    "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...")
+    os.makedirs(WEIGHTS_DIR, exist_ok=True)
+
+    snapshot_download(
+        repo_id=MODEL_REPO,
+        local_dir=WEIGHTS_DIR,
+        local_dir_use_symlinks=False
+    )
+
+    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)
 
-# Only download this specific file from transformers
-ESSENTIAL_PATHS = [
-    "hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states_fp8.pt",
-]
-
-# Download everything from these folders
-FULL_DIRS = [
-    "hunyuan-video-t2v-720p/vae",
-    #"llava_llama_image",
-    "text_encoder_2",
-    "whisper-tiny",
-    "det_align",
-]
-
-
-def list_ckpt_files():
-    """Return a list of files to download: specific files + all from some folders."""
-    try:
-        all_files = list_repo_files(MODEL_REPO)
-    except Exception as e:
-        print(f"❌ Failed to list files from {MODEL_REPO}: {e}")
-        return []
-
-    files_to_download = ESSENTIAL_PATHS.copy()
-    for path in all_files:
-        if any(path.startswith(folder + "/") for folder in FULL_DIRS):
-            files_to_download.append(path)
-
-    return files_to_download
-
-
-def download_ckpts(files):
-    """Download selected files to local directory, preserving structure."""
-    for file_path in files:
-        local_path = os.path.join(WEIGHTS_DIR, file_path)
-        if os.path.exists(local_path):
-            print(f"✅ Already exists: {local_path}")
-            continue
-
-        print(f"⬇️  Downloading: {file_path}")
-        try:
-            hf_hub_download(
-                repo_id=MODEL_REPO,
-                filename=file_path,
-                local_dir=WEIGHTS_DIR,
-                local_dir_use_symlinks=False,
-                resume_download=True,
-            )
-        except EntryNotFoundError:
-            print(f"❌ Entry not found: {file_path}")
-        except Exception as e:
-            print(f"❌ Failed to download {file_path}: {e}")
+    print("✅ Model downloaded successfully.")
 
 def run_sample_gpu_poor():
-    ckpt_fp8 = os.path.join(WEIGHTS_DIR, "ckpts", "hunyuan-video-t2v-720p", "transformers", "mp_rank_00_model_states_fp8.pt")
-
-    if not os.path.isfile(ckpt_fp8):
-        print(f"❌ Missing checkpoint: {ckpt_fp8}")
-        return
-
+    print("🎬 Running sample_gpu_poor.py...")
     cmd = [
         "python3", "hymm_sp/sample_gpu_poor.py",
         "--input", "assets/test.csv",
-        "--ckpt", ckpt_fp8,
+        "--ckpt", CHECKPOINT_FP8_FILE,
         "--sample-n-frames", "129",
         "--seed", "128",
         "--image-size", "704",
@@ -92,27 +71,49 @@ def run_sample_gpu_poor():
     env["CPU_OFFLOAD"] = "1"
     env["CUDA_VISIBLE_DEVICES"] = "0"
 
-    print("🚀 Running sample_gpu_poor.py...")
-    result = subprocess.run(cmd, env=env, capture_output=True, text=True)
+    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.")
 
-    if result.returncode != 0:
-        print(f"❌ sample_gpu_poor.py failed:\n{result.stderr}")
-    else:
-        print("✅ sample_gpu_poor.py ran successfully!")
+def run_flask_audio():
+    print("🚀 Starting flask_audio.py...")
+    cmd = [
+        "torchrun",
+        "--nnodes=1",
+        "--nproc_per_node=8",
+        "--master_port=29605",
+        "hymm_gradio/flask_audio.py",
+        "--input", "assets/test.csv",
+        "--ckpt", CHECKPOINT_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"
+    ]
+    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():
-    files = list_ckpt_files()
-    if not files:
-        print("❌ No checkpoint files found in repo under ckpts folder.")
-        sys.exit(1)
+    if os.path.isfile(CHECKPOINT_FILE) and os.path.isfile(CHECKPOINT_FP8_FILE):
+        print("✅ Model checkpoint already exists. Skipping download.")
+    else:
+        download_model()
 
-    download_ckpts(files)
     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()