import os
from huggingface_hub import hf_hub_download

REPO_ID = "tencent/HunyuanVideo-Avatar"
BASE_PATH = "ckpts"
LOCAL_BASE = os.path.join(os.getcwd(), "weights", "ckpts")

# List of essential files/folders to download (you can expand this if needed)
ESSENTIAL_PATHS = [
    # Transformers checkpoints
    "hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states_fp8.pt",
    "hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt",
    "hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states_fp8_map.pt",

    # VAE
    "hunyuan-video-t2v-720p/vae/config.json",
    "hunyuan-video-t2v-720p/vae/pytorch_model.pt",

    # llava_llama_image shard files (adjust count if needed)
    "llava_llama_image/model-00001-of-00004.safetensors",
    "llava_llama_image/model-00002-of-00004.safetensors",
    "llava_llama_image/model-00003-of-00004.safetensors",
    "llava_llama_image/model-00004-of-00004.safetensors",
    "llava_llama_image/config.json",

    # text_encoder_2
    "text_encoder_2/config.json",
    "text_encoder_2/pytorch_model.bin",

    # whisper-tiny
    "whisper-tiny/config.json",
    "whisper-tiny/pytorch_model.bin",
    "whisper-tiny/tokenizer.json",
    "whisper-tiny/tokenizer_config.json",
    "whisper-tiny/vocab.json",

    # det_align
    "det_align/config.json",
    "det_align/pytorch_model.bin",
]

def download_files():
    for relative_path in ESSENTIAL_PATHS:
        source_path = f"{BASE_PATH}/{relative_path}"
        local_dir = os.path.join(LOCAL_BASE, os.path.dirname(relative_path))
        os.makedirs(local_dir, exist_ok=True)

        print(f"⬇️ Downloading {source_path} ...")
        try:
            hf_hub_download(
                repo_id=REPO_ID,
                filename=source_path,
                repo_type="model",
                local_dir=local_dir,
                local_dir_use_symlinks=False
            )
        except Exception as e:
            print(f"❌ Failed to download {source_path}: {e}")

if __name__ == "__main__":
    download_files()
    print("\n✅ All selected model weights downloaded to:")
    print(f"{os.path.abspath(LOCAL_BASE)}")