# AOTI compiler Space for the CLIP-guided diffusion museum demo.
# Compiles one full guided DDPM timestep (UNet + guidance grad + update; tag diffstep3-dyn:
# respaced schedule + guidance scales are graph INPUTS, so one package per backbone serves
# any steps/scales) on ZeroGPU for each backbone missing from the dataset.
# Requires an HF_TOKEN Space secret with write access to the dataset repo.

import spaces  # must precede any torch/CUDA import

import json
import os
import sys
import tempfile
import time
import traceback
import importlib.metadata as importlib_metadata
from datetime import datetime, timezone

sys.argv = [sys.argv[0]]  # diffusion_fast/generate_fast parse CLI args at import
import torch
import gradio as gr
from huggingface_hub import HfApi, hf_hub_download

import diffusion_fast as D
import backbones as bb
from guided_diffusion.script_util import create_model_and_diffusion, model_and_diffusion_defaults

IMAGE_SIZE = 512
CUTN = 16
CUT_POW = 1.0
NP = 4  # fixed prompt slots; museum app zero-pads (weight 0 => exact zero contribution)
UNET_REPO = "lowlevelware/512x512_diffusion_unconditional_ImageNet"
UNET_FILE = "512x512_diffusion_uncond_finetune_008100.pt"
DATASET_REPO = os.environ.get("AOTI_DATASET_REPO", "multimodalart/clip-guided-diffusion-aoti-zerogpu")
CLIP_MODELS = ["ViT-B/16", "ViT-B/32", "ViT-L/14", "google/siglip2-large-patch16-256"]
CFG_DIRS = {"ViT-B/16": "diffstep3dyn_vitb16_512_cutn16_bf16_np4",
            "ViT-B/32": "diffstep3dyn_vitb32_512_cutn16_bf16_np4",
            "ViT-L/14": "diffstep3dyn_vitl14_512_cutn16_bf16_np4",
            "google/siglip2-large-patch16-256": "diffstep3dyn_siglip2l16_256_512_cutn16_bf16_np4"}


def make_cfg(clip_model):
    return {"tag": "diffstep3-dyn", "clip_model": clip_model, "unet": UNET_FILE,
            "image_size": IMAGE_SIZE, "cutn": CUTN, "dtype": "torch.bfloat16", "n_prompts": NP}


print("Downloading UNet checkpoint...")
CKPT = hf_hub_download(UNET_REPO, UNET_FILE)

print("Loading UNet + backbones (CPU)...")
_cfg = model_and_diffusion_defaults()
_cfg.update(D.MODEL_CFG)
_cfg.update({"image_size": IMAGE_SIZE, "timestep_respacing": "250"})
unet, _ = create_model_and_diffusion(**_cfg)
unet.load_state_dict(torch.load(CKPT, map_location="cpu"))
unet.requires_grad_(False).eval()

backbones, step_mods = {}, {}
for _name in CLIP_MODELS:
    backbones[_name] = bb.load_backbone(_name)
    step_mods[_name] = bb.build_fusedstep(unet, backbones[_name], CUTN, IMAGE_SIZE)


@spaces.GPU(duration=900)
def gpu_compile(clip_model, out_dir, max_autotune):
    dev = torch.device("cuda")
    t0 = time.time()
    cap = torch.cuda.get_device_capability()
    arch = f"sm{cap[0]}{cap[1]}"
    yield f"[{clip_model}] GPU: {torch.cuda.get_device_name(0)} ({arch}), torch {torch.__version__}"

    step_mod = step_mods[clip_model].to(dev)
    cut_size = backbones[clip_model]["cut_size"]
    shape = (1, 3, IMAGE_SIZE, IMAGE_SIZE)
    gen = torch.Generator(device=dev).manual_seed(0)
    ex = D.sample_step_rands(gen, CUTN, cut_size, IMAGE_SIZE, CUT_POW, shape)
    e0 = torch.zeros(NP, backbones[clip_model]["output_dim"], device=dev)
    w0 = torch.full((NP,), 0.25, device=dev)

    # every example input must be a DISTINCT tensor object: make_fx binds repeated
    # tensors to one placeholder, silently miswiring the runtime inputs
    def s(v):
        return torch.full((), float(v), device=dev)

    ex_sched = (torch.full((1,), 999., device=dev),
                s(1.1), s(1.2), s(0.9), s(0.5), s(0.6), s(-9.), s(-8.), s(1.))
    ex_args = (torch.randn(shape, device=dev), *ex_sched, e0, w0,
               s(1000.), s(150.), s(50.), *ex)

    yield f"[{clip_model}] [{time.time()-t0:.0f}s] tracing (make_fx, real mode)..."
    from torch.fx.experimental.proxy_tensor import make_fx
    with torch.device(dev):
        gm = make_fx(lambda *a: step_mod(*a), tracing_mode="real")(*ex_args)

    yield f"[{clip_model}] [{time.time()-t0:.0f}s] torch.export..."
    ep = torch.export.export(gm, ex_args, strict=False)

    yield f"[{clip_model}] [{time.time()-t0:.0f}s] AOTI compile (max_autotune={bool(max_autotune)})..."
    spaces.aoti_compile_and_save(package_dir=out_dir, exported_program=ep,
                                 inductor_configs={"max_autotune": bool(max_autotune)})
    weights = {**ep.state_dict, **dict(ep.constants)}
    torch.save(weights, os.path.join(out_dir, "weights.pt"))

    yield f"[{clip_model}] [{time.time()-t0:.0f}s] verifying package vs eager..."
    from spaces.zero.torch.aoti import LazyAOTIModel
    w = {k: v.to(dev) for k, v in weights.items()}
    w.update({f"{k}_cuda0": v for k, v in list(w.items())})  # inductor device aliases
    fn = LazyAOTIModel(os.path.join(out_dir, "root", "package.pt2")).with_weights(w)
    gv = torch.Generator(device=dev).manual_seed(7)
    xv = torch.randn(shape, generator=gv, device=dev)
    ev = torch.nn.functional.normalize(
        torch.randn((NP, e0.shape[1]), generator=gv, device=dev), dim=1)
    vr = D.sample_step_rands(torch.Generator(device=dev).manual_seed(11),
                             CUTN, cut_size, IMAGE_SIZE, CUT_POW, shape)
    v_sched = (torch.full((1,), 501., device=dev),
               s(1.05), s(0.4), s(0.7), s(0.55), s(0.45), s(-7.), s(-6.5), s(1.))
    v_args = (xv, *v_sched, ev, w0, s(1000.), s(150.), s(50.), *vr)
    s1, p1 = fn(*v_args)
    s0_, p0 = step_mod(*v_args)
    sd = (s1 - s0_).abs().max().item()
    pd = (p1 - p0).abs().max().item()
    yield (f"[{clip_model}] [{time.time()-t0:.0f}s] sample maxdiff={sd:.4g}, "
           f"pred_xstart maxdiff={pd:.4g} (bf16 kernel-order noise; expect <0.1)")

    meta = {"config": make_cfg(clip_model), "torch": torch.__version__, "arch": arch,
            "device": torch.cuda.get_device_name(0), "cuda": torch.version.cuda,
            "spaces": importlib_metadata.version("spaces"),
            "created": datetime.now(timezone.utc).isoformat()}
    with open(os.path.join(out_dir, "metadata.json"), "w") as f:
        json.dump(meta, f, indent=1)
    yield f"[{clip_model}] [{time.time()-t0:.0f}s] package ready ({arch}, torch {meta['torch']})"


def dataset_has_config(clip_model):
    try:
        p = hf_hub_download(DATASET_REPO, f"{CFG_DIRS[clip_model]}/metadata.json",
                            repo_type="dataset", force_download=True)
        with open(p) as f:
            meta = json.load(f)
        return meta.get("config") == make_cfg(clip_model) and meta.get("torch") == torch.__version__
    except Exception:
        return False


def compile_and_publish(max_autotune):
    logs = []
    def out(s):
        logs.append(s)
        return "\n".join(logs)

    token = os.environ.get("HF_TOKEN")
    if not token:
        yield out("ERROR: HF_TOKEN Space secret not set (needs write access to "
                  f"{DATASET_REPO}). Set it in Settings > Variables and secrets.")
        return
    api = HfApi(token=token)
    for clip_model in CLIP_MODELS:
        if dataset_has_config(clip_model):
            yield out(f"[{clip_model}] up-to-date package already in {DATASET_REPO} — skipping")
            continue
        out_dir = tempfile.mkdtemp(prefix="diff_aoti_")
        try:
            for line in gpu_compile(clip_model, out_dir, max_autotune):
                yield out(line)
        except Exception:
            yield out(f"[{clip_model}] COMPILE ERROR:\n" + traceback.format_exc())
            continue
        if not os.path.isfile(os.path.join(out_dir, "metadata.json")):
            yield out(f"[{clip_model}] ERROR: no package produced")
            continue
        yield out(f"[{clip_model}] uploading to {DATASET_REPO}/{CFG_DIRS[clip_model]} ...")
        try:
            api.create_repo(DATASET_REPO, repo_type="dataset", exist_ok=True)
            api.upload_folder(repo_id=DATASET_REPO, repo_type="dataset",
                              folder_path=out_dir, path_in_repo=CFG_DIRS[clip_model],
                              commit_message=f"AOTI package {CFG_DIRS[clip_model]}")
            yield out(f"[{clip_model}] done: https://huggingface.co/datasets/{DATASET_REPO}"
                      f"/tree/main/{CFG_DIRS[clip_model]}")
        except Exception:
            yield out(f"[{clip_model}] UPLOAD ERROR:\n" + traceback.format_exc())
    yield out("All configs processed.")


with gr.Blocks(title="CLIP-Guided Diffusion AOTI compiler") as app:
    gr.Markdown("# CLIP-Guided Diffusion — ZeroGPU AOTI compiler")
    gr.Markdown(f"Compiles one full guided DDPM timestep ({IMAGE_SIZE}px, cutn {CUTN}, bf16, "
                f"{NP} prompt slots; schedule + guidance scales are graph inputs, so one "
                f"package serves any steps/scales) for each backbone missing from "
                f"`{DATASET_REPO}` ({', '.join(CLIP_MODELS)}). Configs already matching this "
                f"torch version are skipped. Needs the `HF_TOKEN` secret (write).")
    max_autotune = gr.Checkbox(label="max_autotune (slower compile, faster steps)", value=True)
    run_btn = gr.Button("Compile & publish missing configs", variant="primary", size="lg")
    log = gr.Textbox(label="Log", lines=22, interactive=False)
    run_btn.click(compile_and_publish, inputs=[max_autotune], outputs=[log],
                  concurrency_limit=1)

if __name__ == "__main__":
    app.launch()  # Spaces sets GRADIO_SERVER_PORT; hardcoding a port breaks the healthcheck