#!/usr/bin/env python3
"""Pull anima base models, text encoder, VAE from circlestone-labs/anima."""
import os, shutil, time
from huggingface_hub import hf_hub_download
import concurrent.futures as cf

TOKEN = "REDACTED_SET_HF_TOKEN_ENV"
ROOT = "/workspace/runpod-slim/ComfyUI/models"
REPO = "circlestone-labs/anima"

PULLS = [
    ("split_files/diffusion_models/anima-base-v1.0.safetensors",      f"{ROOT}/diffusion_models/anima-base-v1.0.safetensors"),
    ("split_files/diffusion_models/anima-preview3-base.safetensors",  f"{ROOT}/diffusion_models/anima-preview3-base.safetensors"),
    ("split_files/text_encoders/qwen_3_06b_base.safetensors",         f"{ROOT}/text_encoders/qwen_3_06b_base.safetensors"),
    ("split_files/vae/qwen_image_vae.safetensors",                    f"{ROOT}/vae/qwen_image_vae.safetensors"),
]

def fetch(item):
    fn, dst = item
    os.makedirs(os.path.dirname(dst), exist_ok=True)
    t0 = time.time()
    p = hf_hub_download(repo_id=REPO, filename=fn, token=TOKEN)
    if not os.path.exists(dst) or os.path.getsize(dst) != os.path.getsize(p):
        shutil.copy(p, dst)
    return f"{os.path.basename(dst)}: {os.path.getsize(dst)/1e9:.2f} GB in {time.time()-t0:.0f}s"

with cf.ThreadPoolExecutor(max_workers=4) as ex:
    for r in ex.map(fetch, PULLS):
        print(r, flush=True)
print("done")