#!/usr/bin/env python3
"""Export the Stable Audio 3 medium DiT (1.45B) to ONNX.

Part of the SA3 post-processing refiner port. Exports the single-forward core
(_forward): the model is 8-step distilled at cfg_scale=1.0, so there is no CFG
dual-pass — the sampler loop lives outside the graph (C++/numpy).

Tensor specs (fp32):
  "x"               [1, 256, T]   noised latents (T = latent frames, dynamic)
  "t"               [1]           current timestep in [0,1] (rf convention)
  "cross_attn_cond" [1, S, 768]   prompt tokens + seconds_total embed (S dynamic)
  "cross_attn_mask" [1, S]        bool
  "global_embed"    [1, 768]      seconds_total embed
  "local_add_cond"  [1, 257, T]   inpaint_mask (1ch) + inpaint_masked_input (256ch)
  "padding_mask"    [1, T]        bool, True = valid
  -> "v"            [1, 256, T]   rf_denoiser output

Also verifies whether the traced graph generalizes over T (dynamic_axes) by
running ORT at a different length; if that fails, the C++ side uses bucketed
static graphs + padding_mask instead.

Runs in the StableAudio3 uv venv:
  cd d:/Ace-Step-Latest/StableAudio3
  uv run --with onnx --with onnxruntime python \
      d:/Ace-Step-Latest/hot-step-cpp/tools/onnx-export/export_sa3_dit.py
"""

import argparse
import os
import sys
import time

import numpy as np
import torch
import torch.nn as nn

sys.path.insert(0, r"d:/Ace-Step-Latest/StableAudio3")

import stable_audio_3.models.transformer as sat
sat.flash_attn_func = None
sat.flash_attn_kvpacked_func = None
sat.flex_attention_available = False
sat.flex_attention_compiled = None

# aten::rms_norm has no ONNX symbolic in the TS exporter — decompose to
# primitive ops (identical math; pow/mean/rsqrt export cleanly and TRT likes them).
import torch.nn.functional as F_patch

def _rms_norm_decomposed(input, normalized_shape, weight=None, eps=None):
    if eps is None:
        eps = torch.finfo(input.dtype).eps
    dims = tuple(range(-len(normalized_shape), 0))
    out = input * torch.rsqrt(input.pow(2).mean(dim=dims, keepdim=True) + eps)
    if weight is not None:
        out = out * weight
    return out

F_patch.rms_norm = _rms_norm_decomposed

from stable_audio_3.model import StableAudioModel

T_TRACE = 1024   # latent frames used for tracing (~97s of audio)
T_ALT = 640      # different length to probe dynamic-shape generalization
S_TRACE = 257    # 256 prompt tokens + 1 seconds_total token


class DiTCore(nn.Module):
    """Flattens the conditioning dict interface to plain tensors around _forward."""

    def __init__(self, dit):
        super().__init__()
        self.dit = dit  # DiffusionTransformer

    def forward(self, x, t, cross_attn_cond, cross_attn_mask, global_embed,
                local_add_cond, padding_mask):
        return self.dit._forward(
            x, t,
            cross_attn_cond=cross_attn_cond,
            cross_attn_cond_mask=cross_attn_mask,
            global_embed=global_embed,
            local_add_cond=local_add_cond,
            padding_mask=padding_mask,
        )


def make_inputs(T, device, seed=0):
    g = torch.Generator(device="cpu").manual_seed(seed)
    x = torch.randn(1, 256, T, generator=g).to(device)
    t = torch.tensor([0.3], dtype=torch.float32, device=device)
    cross = torch.randn(1, S_TRACE, 768, generator=g).to(device)
    cross_mask = torch.ones(1, S_TRACE, dtype=torch.bool, device=device)
    glob = torch.randn(1, 768, generator=g).to(device)
    local = torch.zeros(1, 257, T, device=device)
    pad = torch.ones(1, T, dtype=torch.bool, device=device)
    return (x, t, cross, cross_mask, glob, local, pad)


_ORT_DTYPES = {"tensor(float)": np.float32, "tensor(bool)": np.bool_, "tensor(int64)": np.int64}
INPUT_NAMES = ["x", "t", "cross_attn_cond", "cross_attn_mask",
               "global_embed", "local_add_cond", "padding_mask"]


def run_ort(sess, inputs):
    # Feed by NAME: the exporter prunes graph-unused inputs (e.g. cross_attn_mask —
    # the model never forwards it; learned padding replaces masking), so positional
    # zipping misaligns.
    named = dict(zip(INPUT_NAMES, inputs))
    feed = {}
    for meta in sess.get_inputs():
        arr = named[meta.name].cpu().numpy()
        want = _ORT_DTYPES.get(meta.type)
        if want is not None and arr.dtype != want:
            arr = arr.astype(want)
        feed[meta.name] = arr
    return sess.run(None, feed)[0]


def compare(tag, out, ref):
    ref_np = ref.float().cpu().numpy()
    max_abs = np.abs(out - ref_np).max()
    denom = np.linalg.norm(out.ravel()) * np.linalg.norm(ref_np.ravel())
    cos = float(np.dot(out.ravel(), ref_np.ravel()) / denom) if denom > 0 else 0.0
    print(f"  [{tag}] max_abs_diff={max_abs:.3e}  cosine={cos:.6f}")
    return cos


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--output-dir", default=r"d:/Ace-Step-Latest/hot-step-cpp/models/onnx/sa3")
    ap.add_argument("--parity-only", action="store_true",
                    help="Skip export; run parity against an existing sa3-dit.onnx")
    args = ap.parse_args()
    os.makedirs(args.output_dir, exist_ok=True)

    print("Loading stable-audio-3-medium (fp32)...")
    model = StableAudioModel.from_pretrained("medium", model_half=False)
    dit = model.model.model.model  # StableAudioModel -> CondWrapper -> DiTWrapper -> DiffusionTransformer
    dit.eval().requires_grad_(False)
    device = next(dit.parameters()).device
    core = DiTCore(dit)

    inputs = make_inputs(T_TRACE, device)
    with torch.no_grad():
        ref = core(*inputs)
    print(f"Reference out: {tuple(ref.shape)}")

    dit_path = os.path.join(args.output_dir, "sa3-dit.onnx")
    if args.parity_only:
        assert os.path.exists(dit_path), f"{dit_path} not found"
        print("(--parity-only: skipping export)")
        import onnxruntime as ort
        sess = ort.InferenceSession(dit_path, providers=["CPUExecutionProvider"])
        print("Parity at traced length:")
        out = run_ort(sess, inputs)
        c1 = compare(f"T={T_TRACE}", out, ref)
        print("Dynamic-shape probe at different length:")
        alt_inputs = make_inputs(T_ALT, device, seed=1)
        with torch.no_grad():
            alt_ref = core(*alt_inputs)
        try:
            alt_out = run_ort(sess, alt_inputs)
            c2 = compare(f"T={T_ALT}", alt_out, alt_ref)
            dynamic_ok = c2 > 0.999
        except Exception as e:
            print(f"  [T={T_ALT}] FAILED to run: {type(e).__name__}: {str(e)[:300]}")
            dynamic_ok = False
        print(f"PARITY {'OK' if c1 > 0.999 else 'FAILED'}; DYNAMIC-T {'OK' if dynamic_ok else 'NOT SUPPORTED -> use bucketed static graphs'}")
        return 0 if c1 > 0.999 else 1

    t0 = time.time()
    torch.onnx.export(
        core, inputs, dit_path,
        input_names=["x", "t", "cross_attn_cond", "cross_attn_mask",
                     "global_embed", "local_add_cond", "padding_mask"],
        output_names=["v"],
        dynamic_axes={
            "x": {2: "T"}, "local_add_cond": {2: "T"}, "padding_mask": {1: "T"},
            "cross_attn_cond": {1: "S"}, "cross_attn_mask": {1: "S"},
            "v": {2: "T"},
        },
        opset_version=18, dynamo=False,
    )
    total = sum(os.path.getsize(os.path.join(args.output_dir, f))
                for f in os.listdir(args.output_dir)
                if f.startswith("sa3-dit"))
    print(f"Exported {dit_path} ({total/1e9:.2f} GB incl. external data, {time.time()-t0:.0f}s)")

    import onnxruntime as ort
    sess = ort.InferenceSession(dit_path, providers=["CPUExecutionProvider"])

    print("Parity at traced length:")
    out = run_ort(sess, inputs)
    c1 = compare(f"T={T_TRACE}", out, ref)

    print("Dynamic-shape probe at different length:")
    alt_inputs = make_inputs(T_ALT, device, seed=1)
    with torch.no_grad():
        alt_ref = core(*alt_inputs)
    try:
        alt_out = run_ort(sess, alt_inputs)
        c2 = compare(f"T={T_ALT}", alt_out, alt_ref)
        dynamic_ok = c2 > 0.999
    except Exception as e:
        print(f"  [T={T_ALT}] FAILED to run: {type(e).__name__}: {str(e)[:300]}")
        dynamic_ok = False

    print(f"PARITY {'OK' if c1 > 0.999 else 'FAILED'}; DYNAMIC-T {'OK' if dynamic_ok else 'NOT SUPPORTED -> use bucketed static graphs'}")
    return 0 if c1 > 0.999 else 1


if __name__ == "__main__":
    sys.exit(main())