#!/usr/bin/env python3
"""Download all flux-redux workflow models into ComfyUI/models/."""
import os, shutil, sys
os.environ.pop("HF_HUB_ENABLE_HF_TRANSFER", None)
os.environ.pop("HF_XET_HIGH_PERFORMANCE", None)
from huggingface_hub import hf_hub_download

TOKEN = os.environ["HUGGING_FACE_ACCESS_TOKEN"]
M = "/workspace/ComfyUI/models"

JOBS = [
    # (repo, filename_in_repo, local_dir, local_name)
    ("black-forest-labs/FLUX.1-Redux-dev", "flux1-redux-dev.safetensors", "style_models", None),
    ("Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro-2.0", "diffusion_pytorch_model.safetensors",
     "controlnet", "FLUX.1-dev-ControlNet-Union-Pro-2.0.safetensors"),
    ("black-forest-labs/FLUX.1-Depth-dev-lora", "flux1-depth-dev-lora.safetensors", "loras", None),
    ("aleph65/ComfyUI", "models/text_encoders/t5xxl_fp16.safetensors", "text_encoders", "t5xxl_fp16.safetensors"),
    ("aleph65/ComfyUI", "models/text_encoders/clip_l.safetensors", "text_encoders", "clip_l.safetensors"),
    ("aleph65/ComfyUI", "models/vae/ae.safetensors", "vae", "ae.safetensors"),
    ("aleph65/ComfyUI", "models/clip_vision/sigclip_vision_patch14_384.safetensors",
     "clip_vision", "sigclip_vision_patch14_384.safetensors"),
    ("aleph65/ComfyUI", "models/loras/flux1-turbo-alpha.safetensors", "loras", "flux1-turbo-alpha.safetensors"),
    ("black-forest-labs/FLUX.1-schnell", "flux1-schnell.safetensors", "diffusion_models", None),
    ("black-forest-labs/FLUX.1-dev", "flux1-dev.safetensors", "diffusion_models", None),
]

failures = []
for repo, fname, subdir, rename in JOBS:
    target = os.path.join(M, subdir, rename or os.path.basename(fname))
    if os.path.exists(target) and os.path.getsize(target) > 1e6:
        print(f"SKIP (exists): {target}", flush=True)
        continue
    print(f"DOWNLOADING {repo}/{fname} -> {target}", flush=True)
    try:
        p = hf_hub_download(repo_id=repo, filename=fname, token=TOKEN)
        os.makedirs(os.path.dirname(target), exist_ok=True)
        shutil.copy(p, target)
        print(f"DONE: {target} ({os.path.getsize(target)/1e9:.2f} GB)", flush=True)
    except Exception as e:
        print(f"FAILED: {repo}/{fname}: {e}", flush=True)
        failures.append((repo, fname))

if failures:
    print("FAILURES:", failures, flush=True)
    sys.exit(1)
print("ALL DOWNLOADS COMPLETE", flush=True)