import spaces  # MUST be the first import (before torch/diffusers/peft)
import os

import inspect
import json
import shutil
import tempfile
import time
import traceback
from pathlib import Path

import torch
import gradio as gr

BASE_MODEL = "Qwen/Qwen-Image-Layered"
LORA_REPO = "StabilityLabs/Stable-Layers"
LORA_SUBFOLDER = "model"
ARTIFACT_REPO = "multimodalart/stable-layers-aoti"
BLOCK_NAME = "QwenImageTransformerBlock"

INDUCTOR_CONFIGS = {
    "max_autotune": True,
    "coordinate_descent_tuning": True,
    "triton.cudagraphs": False,
}

# Verification shapes: (height, width, num_layers, txt_len)
CAPTURE_SHAPE = (640, 640, 4, 64)
VERIFY_SHAPES = [(640, 640, 4, 64), (640, 384, 2, 32), (448, 640, 6, 77)]

# ---------------------------------------------------------------------------
# Module scope: load transformer only, fuse LoRA on CPU, then .to("cuda")
# (no forward passes here — that would poison the GPU worker forks)
# ---------------------------------------------------------------------------

print(f"Loading transformer from {BASE_MODEL} ...")
from diffusers import QwenImageTransformer2DModel

transformer = QwenImageTransformer2DModel.from_pretrained(
    BASE_MODEL, subfolder="transformer", torch_dtype=torch.bfloat16
)

print(f"Applying LoRA {LORA_REPO}/{LORA_SUBFOLDER} and fusing on CPU ...")
from peft import PeftModel

_peft = PeftModel.from_pretrained(
    transformer, LORA_REPO, subfolder=LORA_SUBFOLDER, torch_device="cpu"
)
transformer = _peft.merge_and_unload()
del _peft

transformer.eval().requires_grad_(False)
transformer.to("cuda")
print("Transformer fused and placed on cuda (hijacked).")

IN_CHANNELS = transformer.config.in_channels            # 64
JOINT_DIM = transformer.config.joint_attention_dim      # 3584


# ---------------------------------------------------------------------------
# Helpers (only ever called inside the @spaces.GPU function)
# ---------------------------------------------------------------------------

def make_transformer_inputs(height, width, num_layers, txt_len):
    """Synthetic full-transformer inputs with realistic shapes/dtypes."""
    ph = (height // 8) // 2
    pw = (width // 8) // 2
    seq = (num_layers + 2) * ph * pw  # (L+1) generated frames + 1 condition frame
    return dict(
        hidden_states=torch.randn(1, seq, IN_CHANNELS, dtype=torch.bfloat16, device="cuda"),
        timestep=torch.tensor([0.5], dtype=torch.bfloat16, device="cuda"),
        encoder_hidden_states=torch.randn(1, txt_len, JOINT_DIM, dtype=torch.bfloat16, device="cuda"),
        encoder_hidden_states_mask=torch.ones(1, txt_len, dtype=torch.long, device="cuda"),
        img_shapes=[[(1, ph, pw)] * (num_layers + 1) + [(1, ph, pw)]],
        guidance=None,
        additional_t_cond=torch.zeros(1, dtype=torch.long, device="cuda"),
        return_dict=False,
    )


def _clone(x):
    if torch.is_tensor(x):
        return x.detach().clone()
    if isinstance(x, tuple):
        return tuple(_clone(v) for v in x)
    if isinstance(x, list):
        return [_clone(v) for v in x]
    if isinstance(x, dict):
        return {k: _clone(v) for k, v in x.items()}
    return x


def capture_block_inputs(height, width, num_layers, txt_len, use_mask=True):
    """Capture REAL inputs to transformer_blocks[0] by aborting a forward pass.

    use_mask=False calls the transformer with encoder_hidden_states_mask=None,
    matching a demo that passes no text mask (all-ones mask is equivalent)."""
    block = transformer.transformer_blocks[0]
    inputs = make_transformer_inputs(height, width, num_layers, txt_len)
    if not use_mask:
        inputs["encoder_hidden_states_mask"] = None
    with spaces.aoti_capture(block) as call:
        with torch.no_grad():
            transformer(**inputs)
    if not call.args and not call.kwargs:
        raise RuntimeError("aoti_capture captured nothing — block forward was never reached")
    return _clone(call.args), _clone(call.kwargs)


def describe(x):
    if torch.is_tensor(x):
        return f"Tensor{tuple(x.shape)} {str(x.dtype).replace('torch.', '')} {x.device.type}"
    if isinstance(x, tuple):
        return "(" + ", ".join(describe(v) for v in x) + ")"
    if isinstance(x, (list, int, float)):
        return f"NON-TENSOR<{type(x).__name__}>: {x!r}"
    return repr(x)


def flat_tensor_check(kwargs):
    """Return names of kwargs that are neither tensors, tuples of tensors, nor None."""
    bad = []
    for k, v in kwargs.items():
        if v is None or torch.is_tensor(v):
            continue
        if isinstance(v, tuple) and all(torch.is_tensor(t) for t in v):
            continue
        bad.append(f"{k}={describe(v)}")
    return bad


def block_eager_call(block, args, kwargs):
    with torch.no_grad():
        return block(*args, **kwargs)


def bench(fn, n=20, warmup=3):
    for _ in range(warmup):
        fn()
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    for _ in range(n):
        fn()
    torch.cuda.synchronize()
    return (time.perf_counter() - t0) / n * 1000.0  # ms/call


def compare_outputs(eager, compiled, atol=1e-2, rtol=1e-2):
    """bf16 comparison. This block emits huge activations (absmax ~1e8; the source
    even fp16-clips at 65504), so elementwise allclose(atol=1e-2) is vacuous at
    that scale. Pass criterion is relative L2 error + cosine similarity; the raw
    allclose result is still reported."""
    res = []
    for i, (e, c) in enumerate(zip(eager, compiled)):
        e32, c32 = e.float(), c.float()
        diff = e32 - c32
        rel_l2 = float(diff.norm() / e32.norm().clamp_min(1e-12))
        cos = float(torch.nn.functional.cosine_similarity(
            e32.flatten(), c32.flatten(), dim=0
        ))
        res.append({
            "output_index": i,
            "shape": list(e.shape),
            "pass (rel_l2<2e-2 and cos>0.999)": bool(rel_l2 < 2e-2 and cos > 0.999),
            "rel_l2_error": rel_l2,
            "cosine_similarity": cos,
            "allclose(atol=1e-2,rtol=1e-2)": bool(torch.allclose(e32, c32, atol=atol, rtol=rtol)),
            "max_abs_diff": float(diff.abs().max()),
            "eager_absmax": float(e32.abs().max()),
        })
    return res


# ---------------------------------------------------------------------------
# The GPU task
# ---------------------------------------------------------------------------

@spaces.GPU(duration=1200, size="xlarge")
def compile_and_publish():
    report = {"torch": torch.__version__}
    try:
        import diffusers
        report["diffusers"] = diffusers.__version__

        # 1. Sanity: repeated blocks + signature
        repeated = getattr(transformer, "_repeated_blocks", None)
        report["_repeated_blocks"] = repeated
        if not repeated or BLOCK_NAME not in repeated:
            report["error"] = f"_repeated_blocks does not contain {BLOCK_NAME}"
            return report
        block = transformer.transformer_blocks[0]
        report["block_forward_signature"] = str(inspect.signature(block.forward))

        # 2. Capture real block inputs at the reference shape
        t0 = time.perf_counter()
        args, kwargs = capture_block_inputs(*CAPTURE_SHAPE)
        report["capture_shape (H,W,L,T)"] = list(CAPTURE_SHAPE)
        report["captured_args"] = [describe(a) for a in args]
        report["captured_kwargs"] = {k: describe(v) for k, v in kwargs.items()}
        bad = flat_tensor_check(kwargs) + [describe(a) for a in args if not torch.is_tensor(a)]
        if bad:
            report["error"] = (
                "Non-tensor varying inputs reach the block forward — refusing to compile "
                "(baked-in constants would silently corrupt other shapes): " + "; ".join(bad)
            )
            return report

        # 3. Export with Dim.AUTO on every sequence-varying dim
        DA = torch.export.Dim.AUTO
        dyn = {}
        for k, v in kwargs.items():
            if k in ("hidden_states", "encoder_hidden_states", "encoder_hidden_states_mask"):
                dyn[k] = {1: DA}
            elif k == "image_rotary_emb" and isinstance(v, tuple):
                dyn[k] = tuple({0: DA} for _ in v)
            else:
                dyn[k] = None
        report["dynamic_shapes"] = {
            k: (repr(v) if not isinstance(v, tuple) else repr(tuple(v))) for k, v in dyn.items()
        }
        ep = torch.export.export(block, args, kwargs, dynamic_shapes=dyn, strict=False)
        report["export_seconds"] = round(time.perf_counter() - t0, 1)

        # 4. Anonymous tensor constants check
        anon = [n for n in ep.constants if n.startswith("_tensor_constant")]
        report["anonymous_tensor_constants"] = len(anon)
        if len(anon) > 10:
            report["error"] = (
                f"Exported program has {len(anon)} anonymous _tensor_constant entries — "
                "these cannot be supplied by the weight loader and cause illegal memory "
                "access at serve time. Not compiling."
            )
            return report

        # 5. Compile
        t0 = time.perf_counter()
        package_dir = tempfile.mkdtemp(prefix="aoti_pkg_")
        spaces.aoti_compile_and_save(package_dir, ep, INDUCTOR_CONFIGS)
        pt2_path = Path(package_dir) / "root" / "package.pt2"
        report["compile_seconds"] = round(time.perf_counter() - t0, 1)
        report["pt2_bytes"] = pt2_path.stat().st_size

        # 6. Verify vs eager at three shapes (+ benchmark) BEFORE publishing
        cases = []
        for shape in VERIFY_SHAPES:
            a, k = capture_block_inputs(*shape)
            eager_out = block_eager_call(block, a, k)
            cases.append({"shape": shape, "args": a, "kwargs": k, "eager": eager_out})
        big = cases[0]
        eager_ms = bench(lambda: block_eager_call(block, big["args"], big["kwargs"]))

        from spaces.zero.torch.aoti import LazyAOTIModel, aoti_patch
        aoti_patch(block, LazyAOTIModel(str(pt2_path)))

        verification = {}
        all_ok = True
        for case in cases:
            with torch.no_grad():
                compiled_out = block(*case["args"], **case["kwargs"])
            res = compare_outputs(case["eager"], compiled_out)
            verification[f"H{case['shape'][0]}xW{case['shape'][1]}_L{case['shape'][2]}_T{case['shape'][3]}"] = res
            all_ok = all_ok and all(r["pass (rel_l2<2e-2 and cos>0.999)"] for r in res)
        report["verification"] = verification

        def compiled_call():
            with torch.no_grad():
                return block(*big["args"], **big["kwargs"])
        compiled_ms = bench(compiled_call)
        report["benchmark_640x640_L4"] = {
            "eager_ms_per_block_call": round(eager_ms, 2),
            "compiled_ms_per_block_call": round(compiled_ms, 2),
            "speedup": round(eager_ms / compiled_ms, 3),
        }

        if not all_ok:
            report["error"] = "Compiled outputs mismatch eager at one or more shapes — NOT uploading."
            return report

        # 7. Upload with layout <BLOCK_NAME>/package.pt2 at repo root
        from huggingface_hub import HfApi
        token = os.environ.get("HF_TOKEN")
        if not token:
            report["error"] = "HF_TOKEN secret is not set — cannot upload."
            return report
        updir = Path(tempfile.mkdtemp(prefix="aoti_upload_"))
        (updir / BLOCK_NAME).mkdir(parents=True)
        shutil.copy2(pt2_path, updir / BLOCK_NAME / "package.pt2")
        api = HfApi(token=token)
        api.create_repo(ARTIFACT_REPO, repo_type="model", exist_ok=True)
        commit = api.upload_folder(
            repo_id=ARTIFACT_REPO,
            repo_type="model",
            folder_path=str(updir),
            commit_message=f"AoTI {BLOCK_NAME} kernels (Stable-Layers fused, torch {torch.__version__})",
        )
        report["uploaded"] = f"{BLOCK_NAME}/package.pt2"
        report["artifact_repo"] = f"https://huggingface.co/{ARTIFACT_REPO}"
        report["commit_url"] = str(getattr(commit, "commit_url", commit))
        report["status"] = "SUCCESS"
        return report

    except Exception:
        report["error"] = traceback.format_exc()
        return report


# ---------------------------------------------------------------------------
# Phase 1: eager A/B benchmark (mask handling x SDPA backend)
# ---------------------------------------------------------------------------

import contextlib
from torch.nn.attention import SDPBackend, sdpa_kernel

SDPA_BACKENDS = {
    "default": None,
    "cudnn": [SDPBackend.CUDNN_ATTENTION],
    "flash": [SDPBackend.FLASH_ATTENTION],
    "efficient": [SDPBackend.EFFICIENT_ATTENTION],
}


def _timed_block(block, args, kwargs, backends=None, n=20, warmup=3):
    cm = sdpa_kernel(backends) if backends else contextlib.nullcontext()
    with torch.no_grad(), cm:
        for _ in range(warmup):
            block(*args, **kwargs)
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        for _ in range(n):
            block(*args, **kwargs)
        torch.cuda.synchronize()
    return (time.perf_counter() - t0) / n * 1000.0


def _rel_metrics(e, c):
    e32, c32 = e.float(), c.float()
    return {
        "rel_l2_error": float((e32 - c32).norm() / e32.norm().clamp_min(1e-12)),
        "cosine_similarity": float(torch.nn.functional.cosine_similarity(
            e32.flatten(), c32.flatten(), dim=0)),
    }


@spaces.GPU(duration=600, size="xlarge")
def phase1_bench():
    import sys
    report = {"torch": torch.__version__, "python": sys.version.split()[0]}
    try:
        block = transformer.transformer_blocks[0]

        # Runtime attention-dispatch facts
        report["transformer_has_set_attention_backend"] = bool(
            hasattr(transformer, "set_attention_backend")
        )
        proc = block.attn.processor
        report["processor_class"] = type(proc).__name__
        report["processor._attention_backend"] = repr(getattr(proc, "_attention_backend", "<absent>"))
        try:
            from diffusers.models.attention_dispatch import _AttentionBackendRegistry
            report["diffusers_active_backend"] = repr(_AttentionBackendRegistry._active_backend)
        except Exception as e:
            report["diffusers_active_backend"] = f"lookup failed: {e}"

        # Captures
        args, kwargs = capture_block_inputs(*CAPTURE_SHAPE, use_mask=True)
        _, kwargs_nomask_cap = capture_block_inputs(*CAPTURE_SHAPE, use_mask=False)
        report["captured_kwargs_with_mask"] = {k: describe(v) for k, v in kwargs.items()}
        report["captured_kwargs_transformer_mask_none"] = {
            k: describe(v) for k, v in kwargs_nomask_cap.items()
        }

        # Equality: same inputs, mask all-ones vs mask=None (default dispatch)
        kwargs_nomask = dict(kwargs)
        kwargs_nomask["encoder_hidden_states_mask"] = None
        with torch.no_grad():
            out_mask = block(*args, **kwargs)
            out_none = block(*args, **kwargs_nomask)
        report["equality_mask_vs_none"] = [
            _rel_metrics(m, n_) for m, n_ in zip(out_mask, out_none)
        ]

        # Matrix
        matrix = []
        for mask_mode, kw in [("bool_mask", kwargs), ("mask_none", kwargs_nomask)]:
            for backend_name, backends in SDPA_BACKENDS.items():
                cell = {"mask": mask_mode, "sdpa_backend": backend_name}
                try:
                    ms = _timed_block(block, args, kw, backends)
                    cell["ms_per_block_call"] = round(ms, 3)
                    cell["s_per_eval_60_blocks"] = round(ms * 60 / 1000, 3)
                except Exception as e:
                    cell["error"] = f"{type(e).__name__}: {str(e)[:300]}"
                matrix.append(cell)
        report["matrix"] = matrix

        # Bonus rows
        try:
            import flash_attn  # noqa: F401
            report["flash_attn_package"] = flash_attn.__version__
        except Exception as e:
            report["flash_attn_package"] = f"not installed ({type(e).__name__})"
        try:
            import sageattention  # noqa: F401
            report["sageattention_package"] = "installed"
        except Exception as e:
            report["sageattention_package"] = (
                "not installed — no sm_120 wheel for torch 2.11/cp312 in "
                "multimodalart/zerogpu-blackwell-wheels; row skipped (would be lossy anyway)"
            )
        return report
    except Exception:
        report["error"] = traceback.format_exc()
        return report


# ---------------------------------------------------------------------------
# Phase 2: recompile with winning config, publish as variant v2
# ---------------------------------------------------------------------------

@spaces.GPU(duration=1200, size="xlarge")
def phase2_compile_publish(use_mask: bool = False, backend: str = "default"):
    report = {
        "torch": torch.__version__,
        "config": {"use_mask": bool(use_mask), "sdpa_backend": backend},
    }
    try:
        block = transformer.transformer_blocks[0]
        backends = SDPA_BACKENDS.get(backend)
        ctx = (lambda: sdpa_kernel(backends)) if backends else contextlib.nullcontext

        # Capture at reference shape with the chosen mask mode
        args, kwargs = capture_block_inputs(*CAPTURE_SHAPE, use_mask=use_mask)
        report["captured_kwargs"] = {k: describe(v) for k, v in kwargs.items()}
        bad = flat_tensor_check(kwargs) + [describe(a) for a in args if not torch.is_tensor(a)]
        if bad:
            report["error"] = "Non-tensor varying inputs reach the block: " + "; ".join(bad)
            return report

        DA = torch.export.Dim.AUTO
        dyn = {}
        for k, v in kwargs.items():
            if k in ("hidden_states", "encoder_hidden_states"):
                dyn[k] = {1: DA}
            elif k == "encoder_hidden_states_mask":
                dyn[k] = {1: DA} if torch.is_tensor(v) else None
            elif k == "image_rotary_emb" and isinstance(v, tuple):
                dyn[k] = tuple({0: DA} for _ in v)
            else:
                dyn[k] = None
        report["dynamic_shapes"] = {k: repr(v) for k, v in dyn.items()}

        t0 = time.perf_counter()
        with ctx():
            ep = torch.export.export(block, args, kwargs, dynamic_shapes=dyn, strict=False)
        report["export_seconds"] = round(time.perf_counter() - t0, 1)

        anon = [n for n in ep.constants if n.startswith("_tensor_constant")]
        report["anonymous_tensor_constants"] = len(anon)
        if len(anon) > 10:
            report["error"] = f"{len(anon)} anonymous _tensor_constant entries — not compiling."
            return report

        t0 = time.perf_counter()
        package_dir = tempfile.mkdtemp(prefix="aoti_pkg_v2_")
        with ctx():
            spaces.aoti_compile_and_save(package_dir, ep, INDUCTOR_CONFIGS)
        pt2_path = Path(package_dir) / "root" / "package.pt2"
        report["compile_seconds"] = round(time.perf_counter() - t0, 1)
        report["pt2_bytes"] = pt2_path.stat().st_size

        # Verify at three shapes (eager reference uses the SAME mask mode, default dispatch)
        cases = []
        for shape in VERIFY_SHAPES:
            a, k = capture_block_inputs(*shape, use_mask=use_mask)
            with torch.no_grad():
                eager_out = block(*a, **k)
            cases.append({"shape": shape, "args": a, "kwargs": k, "eager": eager_out})
        big = cases[0]

        # Benchmarks: baseline (bool mask, default), best-eager (chosen config), compiled
        base_args, base_kwargs = capture_block_inputs(*CAPTURE_SHAPE, use_mask=True)
        baseline_ms = _timed_block(block, base_args, base_kwargs, None)
        best_eager_ms = _timed_block(block, big["args"], big["kwargs"], backends)

        from spaces.zero.torch.aoti import LazyAOTIModel, aoti_patch
        aoti_patch(block, LazyAOTIModel(str(pt2_path)))

        verification = {}
        all_ok = True
        for case in cases:
            with torch.no_grad():
                compiled_out = block(*case["args"], **case["kwargs"])
            res = compare_outputs(case["eager"], compiled_out)
            key = f"H{case['shape'][0]}xW{case['shape'][1]}_L{case['shape'][2]}_T{case['shape'][3]}"
            verification[key] = res
            all_ok = all_ok and all(r["pass (rel_l2<2e-2 and cos>0.999)"] for r in res)
        report["verification"] = verification

        compiled_ms = _timed_block(block, big["args"], big["kwargs"], None)
        report["benchmark_640x640_L4"] = {
            "baseline_eager_boolmask_default_ms": round(baseline_ms, 3),
            "best_eager_ms": round(best_eager_ms, 3),
            "compiled_ms": round(compiled_ms, 3),
            "compiled_vs_baseline_speedup": round(baseline_ms / compiled_ms, 3),
            "compiled_vs_best_eager_speedup": round(best_eager_ms / compiled_ms, 3),
        }

        if not all_ok:
            report["error"] = "Verification failed — NOT uploading."
            return report
        if compiled_ms > best_eager_ms:
            report["error"] = (
                f"Compiled ({compiled_ms:.2f} ms) is slower than best eager "
                f"({best_eager_ms:.2f} ms) — NOT uploading per gate."
            )
            return report

        from huggingface_hub import HfApi
        token = os.environ.get("HF_TOKEN")
        if not token:
            report["error"] = "HF_TOKEN secret is not set — cannot upload."
            return report
        updir = Path(tempfile.mkdtemp(prefix="aoti_upload_v2_"))
        (updir / f"{BLOCK_NAME}.v2").mkdir(parents=True)
        shutil.copy2(pt2_path, updir / f"{BLOCK_NAME}.v2" / "package.pt2")
        api = HfApi(token=token)
        commit = api.upload_folder(
            repo_id=ARTIFACT_REPO,
            repo_type="model",
            folder_path=str(updir),
            commit_message=(
                f"AoTI {BLOCK_NAME} v2 kernels (mask={'bool' if use_mask else 'none'}, "
                f"sdpa={backend}, torch {torch.__version__})"
            ),
        )
        report["uploaded"] = f"{BLOCK_NAME}.v2/package.pt2"
        report["commit_url"] = str(getattr(commit, "commit_url", commit))
        report["status"] = "SUCCESS"
        return report
    except Exception:
        report["error"] = traceback.format_exc()
        return report


# ---------------------------------------------------------------------------
# Diagnostics
# ---------------------------------------------------------------------------

class _RotMicro(torch.nn.Module):
    """Mirrors apply_rotary_emb_qwen(use_real=False) math."""
    def forward(self, x, freqs):
        xr = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
        out = torch.view_as_real(xr * freqs.unsqueeze(1)).flatten(3)
        return out.type_as(x)


def _stats(t):
    t32 = t.float()
    return {"absmax": float(t32.abs().max()), "mean": float(t32.mean()), "std": float(t32.std())}


@spaces.GPU(duration=900, size="xlarge")
def diagnose():
    report = {"torch": torch.__version__}
    try:
        DA = torch.export.Dim.AUTO
        block = transformer.transformer_blocks[0]

        # ---- A. block export + compile, then inspect constant FQNs vs state_dict
        args, kwargs = capture_block_inputs(*CAPTURE_SHAPE)
        with torch.no_grad():
            eager_out = block(*args, **kwargs)
        dyn = {}
        for k, v in kwargs.items():
            if k in ("hidden_states", "encoder_hidden_states", "encoder_hidden_states_mask"):
                dyn[k] = {1: DA}
            elif k == "image_rotary_emb" and isinstance(v, tuple):
                dyn[k] = tuple({0: DA} for _ in v)
            else:
                dyn[k] = None
        ep = torch.export.export(block, args, kwargs, dynamic_shapes=dyn, strict=False)
        package_dir = tempfile.mkdtemp(prefix="diag_pkg_")
        spaces.aoti_compile_and_save(package_dir, ep, INDUCTOR_CONFIGS)
        pt2 = Path(package_dir) / "root" / "package.pt2"

        compiled = torch._inductor.aoti_load_package(str(pt2))
        fqns = list(compiled.get_constant_fqns())
        sd = block.state_dict()
        report["n_constant_fqns"] = len(fqns)
        report["n_state_dict_keys"] = len(sd)
        report["fqn_sample"] = sorted(fqns)[:8]
        report["sd_key_sample"] = sorted(sd.keys())[:8]
        missing = sorted(set(fqns) - set(sd.keys()))
        extra = sorted(set(sd.keys()) - set(fqns))
        report["fqns_missing_from_state_dict"] = missing[:20]
        report["state_dict_keys_not_in_fqns"] = extra[:20]

        constant_map = {n: sd[n] for n in fqns if n in sd}
        report["n_loaded_constants"] = len(constant_map)
        try:
            compiled.load_constants(constant_map, check_full_update=True, user_managed=True)
            report["load_constants_full_update"] = "ok"
        except Exception as e:
            report["load_constants_full_update_error"] = str(e)[:500]
            compiled.load_constants(constant_map, check_full_update=False, user_managed=True)

        with torch.no_grad():
            c1 = compiled(*args, **kwargs)
            c2 = compiled(*args, **kwargs)
        report["block_eager_stats"] = [_stats(t) for t in eager_out]
        report["block_compiled_stats"] = [_stats(t) for t in c1]
        report["block_allclose"] = [
            bool(torch.allclose(e.float(), c.float(), atol=1e-2, rtol=1e-2))
            for e, c in zip(eager_out, c1)
        ]
        report["block_max_abs_diff"] = [
            float((e.float() - c.float()).abs().max()) for e, c in zip(eager_out, c1)
        ]
        report["compiled_deterministic"] = [
            bool(torch.equal(a, b)) for a, b in zip(c1, c2)
        ]

        # ---- B. same block WITHOUT max_autotune
        package_dir2 = tempfile.mkdtemp(prefix="diag_pkg_plain_")
        spaces.aoti_compile_and_save(package_dir2, ep, {"triton.cudagraphs": False})
        pt2b = Path(package_dir2) / "root" / "package.pt2"
        compiled_b = torch._inductor.aoti_load_package(str(pt2b))
        cmap_b = {n: sd[n] for n in compiled_b.get_constant_fqns() if n in sd}
        compiled_b.load_constants(cmap_b, check_full_update=False, user_managed=True)
        with torch.no_grad():
            cb = compiled_b(*args, **kwargs)
        report["plain_block_allclose"] = [
            bool(torch.allclose(e.float(), c.float(), atol=1e-2, rtol=1e-2))
            for e, c in zip(eager_out, cb)
        ]
        report["plain_block_max_abs_diff"] = [
            float((e.float() - c.float()).abs().max()) for e, c in zip(eager_out, cb)
        ]

        # ---- C. micro complex-rotary module (no weights) with both config sets
        x = torch.randn(1, 1024, 24, 128, dtype=torch.bfloat16, device="cuda")
        fr = torch.polar(
            torch.ones(1024, 64, device="cuda"), torch.randn(1024, 64, device="cuda")
        )
        micro = _RotMicro().to("cuda").eval()
        with torch.no_grad():
            micro_eager = micro(x, fr)
        ep_m = torch.export.export(
            micro, (x, fr), dynamic_shapes=(({1: DA}, {0: DA})), strict=False
        )
        for label, cfg in [("autotune", INDUCTOR_CONFIGS), ("plain", {"triton.cudagraphs": False})]:
            d = tempfile.mkdtemp(prefix=f"diag_micro_{label}_")
            spaces.aoti_compile_and_save(d, ep_m, cfg)
            cm = torch._inductor.aoti_load_package(str(Path(d) / "root" / "package.pt2"))
            mf = cm.get_constant_fqns()
            if mf:
                cm.load_constants({}, check_full_update=False, user_managed=True)
            with torch.no_grad():
                mo = cm(x, fr)
            if isinstance(mo, (list, tuple)):
                mo = mo[0]
            report[f"micro_{label}_allclose"] = bool(
                torch.allclose(micro_eager.float(), mo.float(), atol=1e-2, rtol=1e-2)
            )
            report[f"micro_{label}_max_abs_diff"] = float(
                (micro_eager.float() - mo.float()).abs().max()
            )
        return report
    except Exception:
        report["error"] = traceback.format_exc()
        return report


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------

with gr.Blocks(title="Stable Layers AoTI compile") as demo:
    gr.Markdown(
        "# Stable Layers — AoTI block compiler\n"
        f"Exports + AOT-compiles `{BLOCK_NAME}` from `{BASE_MODEL}` (transformer only, "
        f"`{LORA_REPO}` LoRA fused), verifies vs eager at 3 shapes, then uploads "
        f"`{BLOCK_NAME}/package.pt2` to [{ARTIFACT_REPO}](https://huggingface.co/{ARTIFACT_REPO})."
    )
    btn = gr.Button("Compile, verify & publish", variant="primary")
    diag_btn = gr.Button("Diagnose (no publish)")
    with gr.Row():
        p1_btn = gr.Button("Phase 1: eager mask/backend matrix")
        use_mask_in = gr.Checkbox(label="v2: keep bool mask", value=False)
        backend_in = gr.Dropdown(
            label="v2: sdpa backend", choices=list(SDPA_BACKENDS.keys()), value="default"
        )
        p2_btn = gr.Button("Phase 2: compile v2 & publish")
    out = gr.JSON(label="Report")
    btn.click(fn=compile_and_publish, inputs=None, outputs=out)
    diag_btn.click(fn=diagnose, inputs=None, outputs=out)
    p1_btn.click(fn=phase1_bench, inputs=None, outputs=out)
    p2_btn.click(fn=phase2_compile_publish, inputs=[use_mask_in, backend_in], outputs=out)

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