"""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.". 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)