"""Find the AdaLN linear precision boundary that matches ComfyUI."""

import torch
import torch.nn.functional as functional


capture_dir = "/artifacts/capture"
inputs = torch.load(f"{capture_dir}/input.pt", map_location="cuda", weights_only=False)
expected = torch.load(f"{capture_dir}/block0_norm1_adaln.pt", map_location="cuda", weights_only=False)
embedding = inputs["t_emb"]
weight = expected["weight"]
bias = expected["bias"]
hidden = expected["shift"].shape[-1]

variants = {
    "fp32": (embedding.float(), weight.float(), bias.float()),
    "bf16_weight_fp32": (embedding.float(), weight.bfloat16().float(), bias.float()),
    "bf16_all": (embedding.bfloat16(), weight.bfloat16(), bias.bfloat16()),
    "bf16_input_fp32": (embedding.bfloat16().float(), weight.float(), bias.float()),
}

for name, (x, w, b) in variants.items():
    values = functional.linear(x, w, b).float().view(x.shape[0] * 3, 6 * hidden)
    shift, scale, *_ = values.chunk(6, dim=-1)
    shift_delta = (shift - expected["shift"].float()).abs()
    scale_delta = (scale - expected["scale"].float()).abs()
    print(f"{name} shift_mean={shift_delta.mean().item():.6g} shift_max={shift_delta.max().item():.6g} scale_mean={scale_delta.mean().item():.6g} scale_max={scale_delta.max().item():.6g}")