#!/usr/bin/env python3
"""Pull all our trained checkpoints from HF in parallel."""
import os, shutil, time
from huggingface_hub import hf_hub_download
import concurrent.futures as cf

TOKEN = "REDACTED_SET_HF_TOKEN_ENV"
DM = "/workspace/runpod-slim/ComfyUI/models/diffusion_models"
os.makedirs(DM, exist_ok=True)

# (repo, filename, local_name)
PULLS = [
    ("advokat/anima-finish-checkpoints", "anima_finish_v1_step4739.safetensors", "anima_finish_v1_step4739.safetensors"),
    ("advokat/anima-finish-checkpoints", "anima_finish_v1_step5059.safetensors", "anima_finish_v1_step5059.safetensors"),
    ("advokat/anima-finish-checkpoints", "anima_ft_plus_v1_magaware.safetensors", "anima_ft_plus_v1_magaware.safetensors"),
    ("advokat/anima-polish-checkpoints", "anima_polish_v1_step524.safetensors", "anima_polish_v1_step524.safetensors"),
    ("advokat/anima-polish-checkpoints", "anima_polish_v1_step846.safetensors", "anima_polish_v1_step846.safetensors"),
]

def fetch(item):
    repo, fn, local = item
    t0 = time.time()
    p = hf_hub_download(repo_id=repo, filename=fn, token=TOKEN)
    dst = os.path.join(DM, local)
    if not os.path.exists(dst) or os.path.getsize(dst) != os.path.getsize(p):
        shutil.copy(p, dst)
    return f"{local}: {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")