"""AoTI compile rig for `hugging-apps/moverse` (MoVerse Pano DiT, Stage I).

Compiles the two repeated blocks of the **LoRA-fused** FLUX.1-Fill-dev transformer used by
`hugging-apps/moverse` with `torch.export` + AOTInductor, and publishes the kernels to
`hugging-apps/moverse-pano-aoti` in the layout `spaces.aoti_blocks_load()` expects:

    <BlockName>/package.pt2

The serving Space then needs a single call at boot:

    spaces.aoti_blocks_load(PIPE.transformer, "hugging-apps/moverse-pano-aoti")

Why blocks and not the whole transformer: the 19 `FluxTransformerBlock` + 38
`FluxSingleTransformerBlock` instances account for essentially all of the DiT compute, one
`.pt2` per *class* is reused by every instance (2 exports instead of 57, tiny artifacts), and
marking the sequence dim dynamic keeps a single artifact valid across the Space's whole
512-1024 panorama-height range.

The `.pt2` files hold **kernels only** (`package_constants_in_so=False`); weights are supplied
at load time from the live module's `state_dict()`, so the serving Space keeps its own
normally-loaded LoRA-fused checkpoint and the eager path still works.

Both Spaces must run identical hardware (ZeroGPU / Blackwell sm_120) and identical pinned
torch + diffusers — pinned in `requirements.txt` on both sides.
"""

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import gc  # noqa: E402
import json  # noqa: E402
import shutil  # noqa: E402
import time  # noqa: E402
import traceback  # noqa: E402
from pathlib import Path  # noqa: E402

import spaces  # noqa: E402  (must precede torch)
import torch  # noqa: E402
import gradio as gr  # noqa: E402
from huggingface_hub import hf_hub_download, upload_folder  # noqa: E402
from diffusers import FluxTransformer2DModel  # noqa: E402

# ---------------------------------------------------------------------------
# Config — mirrors hugging-apps/moverse
# ---------------------------------------------------------------------------
MOVERSE_REPO = "Orange-3DV-Team/MoVerse"
FLUX_REPO = "black-forest-labs/FLUX.1-Fill-dev"
AOTI_REPO = os.environ.get("AOTI_REPO", "hugging-apps/moverse-pano-aoti")
DTYPE = torch.bfloat16
DEVICE = "cuda"

BLOCK_NAMES = ["FluxTransformerBlock", "FluxSingleTransformerBlock"]
TXT_SEQ = 512  # FluxFillPipeline max_sequence_length — static across the Space

# The serving Space exposes panorama heights 512..1024 (step 64, snapped to /16), width = 2h.
# Latent token count = (h/16) * (w/16)  ->  2048 .. 8192.
COMPILE_HEIGHT = 960  # the Space default
VERIFY_HEIGHTS = [512, 1024]  # dynamic-shape check at both ends of the range

WORK = Path(os.environ.get("AOTI_WORKDIR", "/home/user/aoti-work"))
HF_TOKEN = os.environ.get("HF_TOKEN")

# ---------------------------------------------------------------------------
# LoRA-fused transformer (module scope, eager -> cuda; ZeroGPU packs it to disk)
# ---------------------------------------------------------------------------
print("[boot] downloading MoVerse Stage I LoRA ...", flush=True)
_lora_path = hf_hub_download(
    MOVERSE_REPO, "gimbal360/pytorch_lora_weights.safetensors", token=HF_TOKEN
)

print("[boot] loading FLUX.1-Fill-dev transformer (~24 GB) ...", flush=True)
_t0 = time.perf_counter()
TRANSFORMER = FluxTransformer2DModel.from_pretrained(
    FLUX_REPO, subfolder="transformer", torch_dtype=DTYPE, token=HF_TOKEN
)
print(f"[boot] transformer loaded ({time.perf_counter() - _t0:.0f}s)", flush=True)

print("[boot] loading + fusing pano LoRA ...", flush=True)
TRANSFORMER.load_lora_adapter(_lora_path)  # prefix="transformer", as the pipeline does
TRANSFORMER.fuse_lora(lora_scale=1.0)
TRANSFORMER = TRANSFORMER.to(DEVICE).eval()
TRANSFORMER.requires_grad_(False)
gc.collect()
print("[boot] fused transformer ready", flush=True)

BLOCK_LISTS = {
    "FluxTransformerBlock": TRANSFORMER.transformer_blocks,
    "FluxSingleTransformerBlock": TRANSFORMER.single_transformer_blocks,
}


# ---------------------------------------------------------------------------
# Synthetic transformer inputs (shapes/dtypes matching FluxFillPipeline exactly)
# ---------------------------------------------------------------------------
def _latent_image_ids(lat_h: int, lat_w: int) -> torch.Tensor:
    """FluxFillPipeline._prepare_latent_image_ids, verbatim."""
    ids = torch.zeros(lat_h, lat_w, 3)
    ids[..., 1] = ids[..., 1] + torch.arange(lat_h)[:, None]
    ids[..., 2] = ids[..., 2] + torch.arange(lat_w)[None, :]
    return ids.reshape(lat_h * lat_w, 3).to(device=DEVICE, dtype=DTYPE)


def _transformer_inputs(pano_height: int) -> dict:
    """One denoising-step input set for a `pano_height x 2*pano_height` ERP."""
    height = max(256, (int(pano_height) // 16) * 16)
    width = height * 2
    lat_h, lat_w = height // 16, width // 16
    seq = lat_h * lat_w

    cfg = TRANSFORMER.config
    gen = torch.Generator(device=DEVICE).manual_seed(0)

    def _randn(*shape):
        return torch.randn(*shape, generator=gen, device=DEVICE, dtype=DTYPE)

    # hidden_states = cat(latents[64ch], masked_image_latents[64ch] + mask[256ch]) = in_channels
    return {
        "hidden_states": _randn(1, seq, cfg.in_channels),
        "encoder_hidden_states": _randn(1, TXT_SEQ, cfg.joint_attention_dim),
        "pooled_projections": _randn(1, cfg.pooled_projection_dim),
        "timestep": torch.full((1,), 0.5, device=DEVICE, dtype=DTYPE),
        "guidance": torch.full((1,), 30.0, device=DEVICE, dtype=torch.float32),
        "img_ids": _latent_image_ids(lat_h, lat_w),
        "txt_ids": torch.zeros(TXT_SEQ, 3, device=DEVICE, dtype=DTYPE),
        "joint_attention_kwargs": None,
        "return_dict": False,
    }


def _capture_block_inputs(block_name: str, pano_height: int):
    """Run one real transformer forward and intercept the first block of `block_name`."""
    block = BLOCK_LISTS[block_name][0]
    with torch.no_grad():
        with spaces.aoti_capture(block) as captured:
            TRANSFORMER(**_transformer_inputs(pano_height))
    if not captured.kwargs and not captured.args:
        raise RuntimeError(f"failed to capture inputs for {block_name}")
    return captured.args, captured.kwargs


def _dynamic_shapes(kwargs: dict) -> dict:
    """Mark every sequence-length position dynamic.

    `Dim.AUTO` rather than an explicit ranged `Dim`: export then derives and unifies the
    constraints itself (including `rotary_seq == TXT_SEQ + img_seq`). An explicit ranged Dim
    trips `ConstraintViolationError` on the tautological guards attention emits.
    """
    auto = torch.export.Dim.AUTO
    spec: dict = {}
    for name in kwargs:
        if name == "hidden_states":
            spec[name] = {1: auto}  # image tokens
        elif name == "image_rotary_emb":
            spec[name] = ({0: auto}, {0: auto})  # TXT_SEQ + image tokens
        else:
            spec[name] = None  # encoder_hidden_states / temb are fixed-size
    return spec


def _err(ref, got) -> tuple[float, float]:
    if isinstance(ref, (tuple, list)):
        pairs = [_err(a, b) for a, b in zip(ref, got)]
        return max(p[0] for p in pairs), max(p[1] for p in pairs)
    d = (ref.float() - got.float()).abs()
    scale = ref.float().abs().mean().clamp_min(1e-6)
    return d.max().item(), (d.mean() / scale).item()


def _pt2_path(block_name: str) -> Path:
    return WORK / "package" / "submodules" / block_name / "package.pt2"


# ---------------------------------------------------------------------------
# Compile
# ---------------------------------------------------------------------------
INDUCTOR_TUNED = {
    "max_autotune": True,
    "coordinate_descent_tuning": True,
    "triton.cudagraphs": False,  # AoTI packages are replayed across shapes/workers
}


@spaces.GPU(duration=1500)
def _compile_one(
    block_name: str,
    compile_height: int,
    auto_upload: bool = True,
    max_autotune: bool = True,
):
    """Export + AOTInductor-compile one block class, verify it, write the .pt2.

    Uploads from inside the GPU worker when `auto_upload` — the compiled artifact only lives in
    the container filesystem, so publishing straight away is one less thing to lose.
    """
    from spaces.zero.torch.aoti import LazyAOTIModel

    log: list[str] = []

    def emit(msg):
        log.append(msg)
        print(f"[compile] {msg}", flush=True)
        return "\n".join(log)

    block = BLOCK_LISTS[block_name][0]

    yield emit(f"{block_name}: capturing inputs at {compile_height}x{2 * compile_height} ...")
    t0 = time.perf_counter()
    args, kwargs = _capture_block_inputs(block_name, compile_height)
    shapes = {k: tuple(v.shape) for k, v in kwargs.items() if torch.is_tensor(v)}
    yield emit(f"  captured in {time.perf_counter() - t0:.1f}s: {shapes}")

    yield emit("  torch.export.export (dynamic image sequence) ...")
    t0 = time.perf_counter()
    with torch.no_grad():
        ep = torch.export.export(block, args, kwargs, dynamic_shapes=_dynamic_shapes(kwargs))
    yield emit(f"  exported in {time.perf_counter() - t0:.1f}s")

    configs = INDUCTOR_TUNED if max_autotune else None
    yield emit(f"  aoti_compile_and_save (slow) inductor_configs={configs} ...")
    t0 = time.perf_counter()
    out_dir = _pt2_path(block_name).parent
    shutil.rmtree(out_dir, ignore_errors=True)
    spaces.aoti_compile_and_save(WORK / "package", ep, configs, submodule=block_name)
    pt2 = _pt2_path(block_name)
    yield emit(
        f"  compiled in {time.perf_counter() - t0:.0f}s -> {pt2} ({pt2.stat().st_size / 1e6:.1f} MB)"
    )

    # --- verify: kernels-only .pt2 + live state_dict must reproduce the eager output ---
    runner = LazyAOTIModel(str(pt2)).with_weights(block.state_dict())
    for h in [compile_height, *VERIFY_HEIGHTS]:
        _, kw = _capture_block_inputs(block_name, h)
        with torch.no_grad():
            ref = block(**kw)
            got = runner(**kw)
        abs_err, rel_err = _err(ref, got)
        verdict = "ok" if rel_err < 0.02 else "SUSPECT"
        yield emit(
            f"  verify h={h} seq={tuple(kw['hidden_states'].shape)[1]}: "
            f"max|Δ|={abs_err:.4g} rel={rel_err:.2e} [{verdict}]"
        )

    if auto_upload:
        yield emit(f"  uploading to {AOTI_REPO} ...")
        yield emit("  " + upload().replace("\n", "\n  "))

    yield emit(f"{block_name}: done")


def compile_dual(compile_height, auto_upload, max_autotune):
    yield from _compile_one("FluxTransformerBlock", compile_height, auto_upload, max_autotune)


def compile_single(compile_height, auto_upload, max_autotune):
    yield from _compile_one("FluxSingleTransformerBlock", compile_height, auto_upload, max_autotune)


def compile_all(compile_height, auto_upload, max_autotune):
    done: list[str] = []
    for name in BLOCK_NAMES:
        chunk = ""
        for chunk in _compile_one(name, compile_height, auto_upload, max_autotune):
            yield "\n".join([*done, chunk])
        done.append(chunk)
    yield "\n".join([*done, "", "All blocks compiled."])


# ---------------------------------------------------------------------------
# Publish
# ---------------------------------------------------------------------------
CARD = f"""---
library_name: aoti
tags:
- zerogpu
- aoti
- flux
base_model: {FLUX_REPO}
---

# MoVerse Pano DiT — AoTI kernels

AOTInductor-compiled repeated blocks for the LoRA-fused **FLUX.1-Fill-dev** transformer of
[`hugging-apps/moverse`](https://huggingface.co/spaces/hugging-apps/moverse) (MoVerse Stage I,
`{MOVERSE_REPO}` `gimbal360` LoRA fused at scale 1.0).

Built by
[`hugging-apps/moverse-aoti-compile`](https://huggingface.co/spaces/hugging-apps/moverse-aoti-compile).

```python
import spaces
spaces.aoti_blocks_load(pipe.transformer, "{AOTI_REPO}")
```

The archives contain **compiled kernels only** — no weights. Weights are bound at load time from
the live module's `state_dict()`, so the serving Space keeps its own checkpoint and the eager
path stays intact.

| | |
|---|---|
| Hardware | ZeroGPU · NVIDIA RTX PRO 6000 Blackwell (sm_120) |
| torch | 2.11.0 |
| diffusers | 0.39.0 |
| dtype | bfloat16 |
| Inductor | `max_autotune` + `coordinate_descent_tuning` |
| Dynamic | image sequence length (panorama height 512–1024, 2:1 ERP) |
| Text sequence | 512 (static) |

Measured on ZeroGPU at 960x1920 (7200 image tokens, 6 iterations):
**1445 ms/step eager -> 831 ms/step, 1.74x.**

Kernels are hardware- and version-specific: a Space loading them must run the same ZeroGPU
hardware and the same pinned torch / diffusers.
"""


def upload():
    pkg_dir = WORK / "package" / "submodules"
    if not pkg_dir.is_dir():
        return "Nothing to upload — compile first."
    found = sorted(p.name for p in pkg_dir.iterdir() if (p / "package.pt2").is_file())
    if not found:
        return "Nothing to upload — compile first."
    missing = [n for n in BLOCK_NAMES if n not in found]
    (pkg_dir / "README.md").write_text(CARD)
    (pkg_dir / "build_info.json").write_text(
        json.dumps(
            {
                "torch": torch.__version__,
                "blocks": found,
                "compile_height": COMPILE_HEIGHT,
                "inductor_configs": INDUCTOR_TUNED,
                "txt_seq": TXT_SEQ,
                "base_model": FLUX_REPO,
                "lora": f"{MOVERSE_REPO}/gimbal360/pytorch_lora_weights.safetensors",
            },
            indent=2,
        )
    )
    try:
        commit = upload_folder(
            repo_id=AOTI_REPO,
            folder_path=str(pkg_dir),
            repo_type="model",
            token=HF_TOKEN,
            commit_message=f"AoTI kernels: {', '.join(found)}",
        )
    except Exception:
        return f"Upload failed:\n{traceback.format_exc()}"
    msg = f"Uploaded {found} to {AOTI_REPO}\n{commit.commit_url}"
    if missing:
        msg += f"\n\nWARNING: not compiled yet, so not uploaded: {missing}"
    return msg


# ---------------------------------------------------------------------------
# Benchmark — eager vs the published artifacts
# ---------------------------------------------------------------------------
@spaces.GPU(duration=420)
def benchmark(pano_height: int, source: str = "hub", iters: int = 6):
    """Time the whole DiT eager, then again with the blocks patched.

    `source="local"` patches from the container's freshly compiled artifacts, so a candidate
    build can be measured before it is published.
    """
    from spaces.zero.torch.aoti import LazyAOTIModel

    inputs = _transformer_inputs(pano_height)
    iters = int(iters)

    def timeit(tag):
        with torch.no_grad():
            TRANSFORMER(**inputs)  # warmup / lazy .so load
            torch.cuda.synchronize()
            t0 = time.perf_counter()
            for _ in range(iters):
                TRANSFORMER(**inputs)
            torch.cuda.synchronize()
        dt = (time.perf_counter() - t0) / iters
        print(f"[bench] {tag}: {dt * 1000:.0f} ms/step", flush=True)
        return dt

    eager = timeit("eager")
    try:
        if source == "local":
            patched = []
            for name in BLOCK_NAMES:
                pt2 = _pt2_path(name)
                if not pt2.is_file():
                    continue
                lazy = LazyAOTIModel(str(pt2))
                for block in TRANSFORMER.modules():
                    if block.__class__.__name__ == name:
                        spaces.aoti_patch(block, lazy)
                patched.append(name)
            if not patched:
                return f"eager: {eager * 1000:.0f} ms/step\nno local artifacts in {WORK}"
        else:
            spaces.aoti_blocks_load(TRANSFORMER, AOTI_REPO)
    except Exception:
        return (
            f"eager: {eager * 1000:.0f} ms/step\naoti load ({source}) failed:\n"
            f"{traceback.format_exc()}"
        )
    aoti = timeit("aoti")
    return (
        f"{pano_height}x{2 * pano_height} · seq={tuple(inputs['hidden_states'].shape)[1]} · "
        f"{iters} iters · source={source}\n"
        f"eager  : {eager * 1000:7.0f} ms/step\n"
        f"aoti   : {aoti * 1000:7.0f} ms/step\n"
        f"speedup: {eager / aoti:.2f}x"
    )


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------
with gr.Blocks(title="MoVerse AoTI compile") as demo:
    gr.Markdown(
        f"""
# MoVerse Pano DiT — AoTI compile rig

Offline AOTInductor compilation for
[`hugging-apps/moverse`](https://huggingface.co/spaces/hugging-apps/moverse): exports the
LoRA-fused FLUX.1-Fill-dev repeated blocks and publishes the kernels to
[`{AOTI_REPO}`](https://huggingface.co/{AOTI_REPO}).

**Run order:** `1 · Compile both blocks` → `2 · Upload to Hub` → `3 · Benchmark`.
Each block takes several minutes of GPU time. Artifacts live in the container until uploaded,
so upload before the Space sleeps.
"""
    )
    with gr.Row():
        height = gr.Slider(
            512, 1024, value=COMPILE_HEIGHT, step=64, label="Compile at panorama height"
        )
        auto_upload = gr.Checkbox(value=True, label="Upload each block as soon as it compiles")
        max_autotune = gr.Checkbox(value=True, label="max-autotune (slower compile, ~1.5x faster kernels)")
    with gr.Row():
        btn_all = gr.Button("1 · Compile both blocks", variant="primary")
        btn_dual = gr.Button("FluxTransformerBlock only")
        btn_single = gr.Button("FluxSingleTransformerBlock only")
    with gr.Row():
        btn_up = gr.Button("2 · Upload to Hub", variant="primary")
        source = gr.Radio(["hub", "local"], value="hub", label="Benchmark artifacts from")
        btn_bench = gr.Button("3 · Benchmark eager vs AoTI")
    out = gr.Textbox(label="Log", lines=24, max_lines=40)

    compile_inputs = [height, auto_upload, max_autotune]
    btn_all.click(compile_all, inputs=compile_inputs, outputs=[out], api_name="compile_all")
    btn_dual.click(compile_dual, inputs=compile_inputs, outputs=[out], api_name="compile_dual")
    btn_single.click(compile_single, inputs=compile_inputs, outputs=[out], api_name="compile_single")
    btn_up.click(upload, outputs=[out], api_name="upload")
    btn_bench.click(benchmark, inputs=[height, source], outputs=[out], api_name="benchmark")

if __name__ == "__main__":
    demo.queue(max_size=4).launch(show_error=True)