app.py

import os
import sys
import subprocess
import time
import shutil
import psutil
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")
CUSTOM_CACHE_DIR = "/data/hf_cache"  # ✅ change if different
HF_HOME_ENV = {"HF_HOME": CUSTOM_CACHE_DIR}

# Paths to required checkpoints
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 check_disk_space(min_free_gb=10):
    total, used, free = shutil.disk_usage("/")
    free_gb = free // (2**30)
    print(f"💾 Disk space: {free_gb}GB free")
    if free_gb < min_free_gb:
        print(f"❌ Not enough disk space. {free_gb}GB available, {min_free_gb}GB required.")
        sys.exit(1)

def clear_hf_cache():
    hf_cache = os.path.expanduser("~/.cache/huggingface")
    if os.path.exists(hf_cache):
        print("🧹 Cleaning Hugging Face cache...")
        shutil.rmtree(hf_cache)

def download_model():
    print("⬇️  Downloading model to weights directory...")
    os.makedirs(WEIGHTS_DIR, exist_ok=True)

    snapshot_download(
        repo_id=MODEL_REPO,
        local_dir=WEIGHTS_DIR,
        local_dir_use_symlinks=False,
        **HF_HOME_ENV
    )

    if not os.path.isfile(CHECKPOINT_FILE):
        print(f"❌ Checkpoint missing at {CHECKPOINT_FILE}")
        sys.exit(1)
    if not os.path.isfile(CHECKPOINT_FP8_FILE):
        print(f"❌ FP8 Checkpoint missing at {CHECKPOINT_FP8_FILE}")
        sys.exit(1)

    print("✅ Model downloaded suc