#!/usr/bin/env python3
"""Laedt alle Krea-2-Komponenten parallel von HuggingFace in die ComfyUI-Ordner."""
import os, sys, time
from concurrent.futures import ThreadPoolExecutor, as_completed
from huggingface_hub import hf_hub_download

COMFY = os.environ.get("COMFY", "/workspace/ComfyUI")
LORA_REPO = os.environ.get("LORA_REPO", "Tirkaru/krea2-nvfp4-rig")
VARIANT = os.environ.get("VARIANT", "both")  # nvfp4 | fp8 | both

M = f"{COMFY}/models"

JOBS = [
    ("Comfy-Org/Krea-2", "text_encoders/qwen3vl_4b_fp8_scaled.safetensors", f"{M}/text_encoders"),
    ("Comfy-Org/Krea-2", "vae/qwen_image_vae.safetensors",                  f"{M}/vae"),
]
if VARIANT in ("nvfp4", "both"):
    JOBS.append(("Comfy-Org/Krea-2", "diffusion_models/krea2_turbo_nvfp4.safetensors", f"{M}/diffusion_models"))
if VARIANT in ("fp8", "both"):
    JOBS.append(("Comfy-Org/Krea-2", "diffusion_models/krea2_turbo_fp8_scaled.safetensors", f"{M}/diffusion_models"))

# LoRA kommt NICHT von HF - es wird direkt auf die Instanz geschoben,
# damit nie ein Account-Token auf der Mietmaschine liegt.
os.makedirs(f"{M}/loras", exist_ok=True)

def grab(repo, path, dest):
    os.makedirs(dest, exist_ok=True)
    t0 = time.time()
    p = hf_hub_download(repo, path, local_dir=dest + "/_tmp")
    final = os.path.join(dest, os.path.basename(path))
    if not os.path.exists(final):
        os.replace(p, final)
    sz = os.path.getsize(final) / 1e9
    return f"{os.path.basename(path):48s} {sz:6.2f} GB in {time.time()-t0:5.1f}s"

if __name__ == "__main__":
    t0 = time.time()
    with ThreadPoolExecutor(max_workers=4) as ex:
        futs = {ex.submit(grab, *j): j for j in JOBS}
        for f in as_completed(futs):
            try:
                print("  OK  " + f.result(), flush=True)
            except Exception as e:
                print(f"  FEHLER {futs[f][1]}: {e}", file=sys.stderr, flush=True)
    print(f"\nAlle Downloads fertig in {time.time()-t0:.1f}s")