"""Download the image models used for teaser first frames into HF_HOME (NVMe)."""
import sys
from concurrent.futures import ThreadPoolExecutor
from huggingface_hub import snapshot_download

JOBS = [
    ("Qwen/Qwen-Image-2.1-Turbo", None),
    ("Qwen/Qwen-Image-Edit-2511", None),
    ("lightx2v/Qwen-Image-Edit-2511-Lightning", ["Qwen-Image-Edit-2511-Lightning-8steps-V1.0-bf16.safetensors",
                                                 "Qwen-Image-Edit-2511-Lightning-4steps-V1.0-bf16.safetensors", "README.md"]),
    ("Tongyi-MAI/Z-Image-Turbo", None),
    ("alibaba-pai/Z-Image-Turbo-Fun-Controlnet-Union", None),
]


def one(job):
    repo, pats = job
    p = snapshot_download(repo, allow_patterns=pats, max_workers=8)
    print("DONE", repo, p, flush=True)


with ThreadPoolExecutor(len(JOBS)) as ex:
    list(ex.map(one, JOBS))
print("ALL DONE", flush=True)