aoti_dit.py

"""Lazy, in-worker AOTInductor compilation of the LTX-2 DiT blocks (via aokit).

This is the ahead-of-time replacement for torch.compile on the transformer
(the DiT). DreamVerse JIT-compiles the blocks; we AOTI-compile ONE block
(weight-less) and run all 48 through that single optimized graph, each bound to
its own weights — aokit's regional design.

Wiring constraint: FastVideo runs the blocks inside a spawned worker, and AOTI
needs a real-shape example, so we hook `LTXModel._process_transformer_blocks`
and compile lazily on the first real forward (captures the true video+audio
shapes automatically). Installed at import time so the patch is live in the
worker (spawn re-imports the app module). Any failure → eager fallback, so the
app never breaks.

Set DREAMVERSE_AOTI=0 to disable.
"""
from __future__ import annotations
import os
import tempfile
import time
import traceback
from dataclasses import replace

import torch

ENABLED = os.getenv("DREAMVERSE_AOTI", "1") == "1"

# Order of TransformerArgs fields we marshal (tuples expand to 2 tensors).
_TUPLE_FIELDS = {"positional_embeddings", "cross_positional_embeddings"}
_FIELDS = ["x", "context", "context_mask", "timesteps", "embedded_timestep",
           "positional_embeddings", "cross_positional_embeddings",
           "cross_scale_shift_timestep", "cross_gate_timestep"]


def _flatten(ta):
    """TransformerArgs -> (list[Tensor], spec). spec records presence/shape."""
    tensors, spec = [], []
    for f in _FIELDS:
        v = getattr(ta, f)
        if v is None:
            spec.append((f, "none"))
        elif f in _TUPLE_FIELDS:
            tensors.extend(v)
            spec.append((f, "tuple2"))
        else:
            tensors.append(v)
            spec.append((f, "tensor"))
    return tensors, spec


def _unflatten(tensors, spec):
    """(iterator of Tensors, spec) -> TransformerArgs."""
    from fastvideo.models.dits.ltx2 import TransformerArgs
    it = iter(tensors)
    kw = {}
    for f, kind in spec:
        if kind == "none":
            kw[f] = None
        elif kind == "tuple2":
            kw[f] = (next(it), next(it))
        else:
            kw[f] = next(it)
    kw["enabled"] = True
    return TransformerArgs(**kw)


class _BlockWrapper(torch.nn.Module):
    """Flat-tensor view of ONE real AV block, for export."""
    def __init__(self, block, vspec, aspec, n_video, scalars):
        super().__init__()
        self.block = block
        self.vspec, self.aspec = vspec, aspec
        self.n_video = n_video
        self.scalars = scalars  # video_original_seq_len, audio_original_seq_len

    def forward(self, *flat):
        video = _unflatten(flat[:self.n_video], self.vspec)
        audio = _unflatten(flat[self.n_video:], self.aspec)
        vout, aout = self.block(
            video=video, audio=audio,
            video_original_seq_len=self.scalars[0],
            audio_original_seq_len=self.scalars[1],
            skip_cross_modal_attn=False,
            skip_video_self_attn=False,
            skip_audio_self_attn=False,
        )
        return vout.x, aout.x


def _make_block_forward(orig_forward, lazy_with_weights, vspec, aspec, baked_vx_shape, baked_ax_shape):
    """Replacement block.forward: AOTI when shapes/flags match, else eager."""
    def forward(video=None, audio=None, video_original_seq_len=None, audio_original_seq_len=None,
                skip_cross_modal_attn=False, skip_video_self_attn=False, skip_audio_self_attn=False):
        ok = (video is not None and audio is not None
              and not skip_cross_modal_attn and not skip_video_self_attn and not skip_audio_self_attn
              and tuple(video.x.shape) == baked_vx_shape and tuple(audio.x.shape) == baked_ax_shape)
        if not ok:
            return orig_forward(video=video, audio=audio,
                                video_original_seq_len=video_original_seq_len,
                                audio_original_seq_len=audio_original_seq_len,
                                skip_cross_modal_attn=skip_cross_modal_attn,
                                skip_video_self_attn=skip_video_self_attn,
                                skip_audio_self_attn=skip_audio_self_attn)
        vt, _ = _flatten(video)
        at, _ = _flatten(audio)
        vx, ax = lazy_with_weights(*vt, *at)
        return replace(video, x=vx), replace(audio, x=ax)
    return forward


AOTI_REPO = os.getenv("DREAMVERSE_AOTI_REPO", "multimodalart/dreamverse-flashinfer-cache")
AOTI_LOCAL = os.path.expanduser("~/.cache/dreamverse_aoti")
_INDUCTOR_CFG = {"max_autotune": True, "max_autotune_gemm": True,
                 "coordinate_descent_tuning": True, "triton.cudagraphs": False}


def _shape_key(video, audio):
    return ("vx" + "x".join(map(str, video.x.shape)) +
            "_ax" + "x".join(map(str, audio.x.shape)))


def _resolve_pt2(key, wrapper, flat):
    """Get the block .pt2 for this shape: local cache -> HF bucket -> compile.

    The .pt2 is aokit's weight-less, RELOCATABLE AOTInductor artifact, so a
    prebaked one loads with NO runtime compile (the ZeroGPU-clean path). We only
    compile if neither a local nor a bucket copy exists.
    """
    import aokit
    import shutil
    local = os.path.join(AOTI_LOCAL, key)
    pt2 = os.path.join(local, "submodules", "b", "package.pt2")
    if os.path.exists(pt2):
        print(f"[AOTI] local prebaked .pt2 ({key})", flush=True)
        return pt2
    try:
        from huggingface_hub import snapshot_download
        p = snapshot_download(repo_id=AOTI_REPO, repo_type="dataset",
                              allow_patterns=f"aoti/{key}/*")
        src = os.path.join(p, "aoti", key, "submodules", "b", "package.pt2")
        if os.path.exists(src):
            os.makedirs(os.path.dirname(pt2), exist_ok=True)
            shutil.copy(src, pt2)
            print(f"[AOTI] pulled prebaked .pt2 from {AOTI_REPO} ({key}) — no compile", flush=True)
            return pt2
    except Exception as e:
        print(f"[AOTI] .pt2 pull failed ({e}); compiling", flush=True)
    # Fallback: compile (and keep locally so this boot reuses it).
    print(f"[AOTI] no prebaked .pt2 for {key}; exporting + AOTInductor compiling ...", flush=True)
    with torch.no_grad():
        exported = torch.export.export(wrapper, flat, {})
    t0 = time.perf_counter()
    aokit.compile_and_save(local, exported, inductor_configs=_INDUCTOR_CFG, submodule="b")
    print(f"[AOTI] compiled in {time.perf_counter()-t0:.0f}s -> {os.path.getsize(pt2)//1024} KB", flush=True)
    return pt2


def _compile_blocks(model, video, audio, vsl, asl):
    blocks = model.transformer_blocks
    block0 = blocks[0]
    vt, vspec = _flatten(video)
    at, aspec = _flatten(audio)
    n_video = len(vt)
    wrapper = _BlockWrapper(block0, vspec, aspec, n_video, (vsl, asl)).eval()
    flat = (*vt, *at)
    key = _shape_key(video, audio)
    print(f"[AOTI] AV block: video.x={tuple(video.x.shape)} audio.x={tuple(audio.x.shape)} key={key}", flush=True)

    pt2 = _resolve_pt2(key, wrapper, flat)
    try:
        from aokit.aokit import LazyAOTIModel
    except Exception:
        from aokit import LazyAOTIModel  # type: ignore
    lazy = LazyAOTIModel(pt2)

    baked_vx = tuple(video.x.shape)
    baked_ax = tuple(audio.x.shape)
    n = 0
    for blk in blocks:
        # Weights are external; the graph's constants are "block.<fqn>".
        weights = {f"block.{k}": v for k, v in blk.state_dict().items()}
        lww = lazy.with_weights(weights)
        blk.forward = _make_block_forward(blk.forward, lww, vspec, aspec, baked_vx, baked_ax)
        n += 1
    print(f"[AOTI] installed AOTI forward on {n} blocks (shape {baked_vx})", flush=True)


def install():
    if not ENABLED:
        return
    if os.getenv("DREAMVERSE_NVFP4") == "1":
        # NVFP4 path runs FP4 (flashinfer) eager + prebaked kernels. AOTInductor
        # can't C++-compile the FP4 op (torch.export traces it, but inductor
        # codegen fails), so stacking AOTI on NVFP4 just wastes a ~50s compile
        # then falls back to eager. Skip the hook entirely in NVFP4 mode.
        print("[AOTI] NVFP4 mode -> skipping AOTI hook (FP4 runs eager + prebaked kernels)", flush=True)
        return
    try:
        from fastvideo.models.dits.ltx2 import LTXModel
    except Exception as e:
        print(f"[AOTI] fastvideo not importable here ({e}); skipping", flush=True)
        return
    if getattr(LTXModel, "_aoti_patched", False):
        return
    LTXModel._aoti_patched = True
    orig = LTXModel._process_transformer_blocks

    def patched(self, video, audio, video_original_seq_len=None, audio_original_seq_len=None,
                skip_cross_modal_attn=False, skip_video_self_attn_blocks=None,
                skip_audio_self_attn_blocks=None):
        if (video is not None and audio is not None
                and not getattr(self, "_aoti_ready", False) and not getattr(self, "_aoti_failed", False)):
            try:
                _compile_blocks(self, video, audio, video_original_seq_len, audio_original_seq_len)
                self._aoti_ready = True
            except Exception as e:
                traceback.print_exc()
                print(f"[AOTI] compile failed -> eager fallback: {e}", flush=True)
                self._aoti_failed = True
        return orig(self, video, audio, video_original_seq_len, audio_original_seq_len,
                    skip_cross_modal_attn, skip_video_self_attn_blocks, skip_audio_self_attn_blocks)

    LTXModel._process_transformer_blocks = patched
    print("[AOTI] installed lazy in-worker DiT compiler hook", flush=True)