import os
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"

def download_checkpoints():
    print("⬇️ Downloading specific checkpoint files from Tencent repo...")

    fp8_ckpt = hf_hub_download(
        repo_id=MODEL_REPO,
        filename="ckpts/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states_fp8.pt"
    )

    fp32_ckpt = hf_hub_download(
        repo_id=MODEL_REPO,
        filename="ckpts/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt"
    )

    print("✅ Checkpoints downloaded.")
    return fp8_ckpt, fp32_ckpt

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", fp8_ckpt_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",
        "--save-path", OUTPUT_DIR,
        "--use-fp8",
        "--cpu-offload",
        "--infer-min"
    ]

    env = os.environ.copy()
    env["PYTHONPATH"] = "./"
    env["CUDA_VISIBLE_DEVICES"] = "0"

    proc = subprocess.run(cmd, env=env)
    if proc.returncode != 0:
        print("❌ sample_gpu_poor.py failed.")
        exit(1)
    print("✅ sample_gpu_poor.py completed successfully.")

def run_gradio_ui():
    print("🟢 Launching Gradio interface...")
    subprocess.Popen(["python3", "hymm_gradio/gradio_audio.py"])

def main():
    fp8_ckpt, _ = download_checkpoints()
    run_sample_gpu_poor(fp8_ckpt)
    time.sleep(5)
    run_gradio_ui()

if __name__ == "__main__":
    main()