"""LTX2DenoiseLoop — a custom-step driven loop (design_v3 §5.1 escape hatch; §15 LTX-2 2-stage). LTX-2's guidance is braided (multi-pass: positive / negative / STG-perturbed, combined per the LTX formula), so the step body is hand-written using samplers + the DiT as a library — exactly the legitimate custom-step pattern. It is still a Loop (scheduled, admitted, observed, parity- gated). One ``LTX2DenoiseLoop`` is instantiated per stage (base / refine) with that stage's distilled sigma schedule; both stages bind the SAME ``transformer`` component (shared by reference — the MoT 'one instance, many loops' property, demonstrated even here). Real LTX-2 distilled schedules (from the repo): base : [1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875, 0.0] (8 steps) refine: [0.909375, 0.725, 0.421875, 0.0] (3 steps) The real DiT predicts ``denoised`` (x0); the CPU toy predicts velocity, integrated by the same flow-match Euler step — the 2-stage structure, distilled schedules, and multi-pass guidance are what this models faithfully. """ from __future__ import annotations import numpy as np from ..._enums import ExecutionProfile, WorkUnitKind from ...loop.contracts import ( Done, LoopResult, LoopState, ResourceRequest, ShapeSignature, StepResult, WorkPlan, ) from ...platform import FLOW_MATCH_STEP from ..backend import LATENT_CHANNELS BASE_SIGMAS = [1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875, 0.0] REFINE_SIGMAS = [0.909375, 0.725, 0.421875, 0.0] def ltx_base_latent_shape(req) -> tuple[int, int, int, int]: d = req.diffusion t = max(1, d.num_frames // 40) # refine halves requested resolution for stage-1; stage-2 upsamples 2× (design §15) h = max(2, d.height // 240) w = max(2, d.width // 240) return (LATENT_CHANNELS, t, h, w) class LTX2DenoiseLoop: def __init__(self, *, loop_id, stage, sigmas, cfg_scale, stg_scale, cost, input_slot=None, seed_offset=0, audio_input_slot=None): self.loop_id = loop_id self.stage = stage # "base" | "refine" self.sigmas = list(sigmas) self.cfg_scale = cfg_scale self.stg_scale = stg_scale self.cost = cost self.input_slot = input_slot # None => fresh noise (base); slot name => read latents (refine) self.audio_input_slot = audio_input_slot # joint A/V: where to read the threaded audio latent self.seed_offset = seed_offset def init(self, req, model, ctx) -> LoopState: seed = (req.diffusion.seed if req.diffusion.seed is not None else 0) + self.seed_offset rng = np.random.default_rng(seed) sig = self.sigmas if self.input_slot is None: x = (rng.standard_normal(ltx_base_latent_shape(req)) * float(sig[0])).astype("float32") else: x = np.asarray(ctx.slots[self.input_slot], dtype="float32") st = LoopState(loop_id=self.loop_id, instance_id=model.card.model_id, request_id=req.request_id, profile=ctx.profile, rng=rng, seed=seed, latents={"video": x}, sigmas=[float(s) for s in sig], timesteps=[float(s) * 1000.0 for s in sig]) st.cond["prompt_embeds"] = ctx.slots.get("text_embeds") st.cond["negative_prompt_embeds"] = ctx.slots.get("neg_text_embeds") # joint A/V (T2VS): denoise an audio latent alongside the video, with per-modality guidance. # Gated on the request asking for audio ⇒ the T2V path below is byte-for-byte unchanged. want_audio = "audio" in getattr(req.outputs, "modalities", frozenset()) gpm = req.diffusion.guidance_per_modality or {} st.scratch["want_audio"] = want_audio st.scratch["v_guidance"] = float(gpm.get("video", self.cfg_scale)) # == cfg_scale for T2V st.scratch["a_guidance"] = float(gpm.get("audio", self.cfg_scale)) if want_audio: if self.audio_input_slot is None: a_len = max(2, req.diffusion.num_frames // 20) au = (rng.standard_normal((LATENT_CHANNELS, a_len, 1, 1)) * float(sig[0])).astype("float32") else: au = np.asarray(ctx.slots[self.audio_input_slot], dtype="float32") st.latents["audio"] = au return st def next(self, st: LoopState): i = st.step_idx if i >= len(st.sigmas) - 1: return Done() sigma_t, sigma_next = st.sigmas[i], st.sigmas[i + 1] x = st.latents["video"] pe, ne = st.cond["prompt_embeds"], st.cond["negative_prompt_embeds"] v_guidance, stg_scale = st.scratch["v_guidance"], self.stg_scale # per-modality video CFG want_audio, a_guidance = st.scratch.get("want_audio", False), st.scratch.get("a_guidance", 0.0) au = st.latents.get("audio") def run(model, override=None): fm = model.platform.kernels.get(FLOW_MATCH_STEP) # solver dispatched per (device, arch) if override is not None and "noise_pred" in override: velocity = np.asarray(override["noise_pred"], dtype="float32") else: dit = model.component("transformer") pos = dit(x, pe, sigma_t) # full conditioned pass neg = dit(x, ne, sigma_t) # negative (CFG) pass # STG-perturbed pass (here: drop text conditioning) — the multi-pass braid ptb = dit(x, np.zeros_like(pe) if pe is not None else None, sigma_t) velocity = (pos + (v_guidance - 1.0) * (pos - neg) + stg_scale * (pos - ptb)) x_next = fm(x, np.asarray(velocity, dtype="float32"), sigma_t, sigma_next) out = {"noise_pred": np.asarray(velocity, dtype="float32"), "latents": x_next.astype("float32")} if want_audio and au is not None: # joint audio denoise, conditioned on video dit = model.component("transformer") a_pos = dit(au, pe, sigma_t, context=x) # context=x ⇒ audio synced to the video a_neg = dit(au, ne, sigma_t, context=x) a_vel = a_pos + (a_guidance - 1.0) * (a_pos - a_neg) # audio's OWN guidance scale out["audio_latents"] = fm(au, a_vel, sigma_t, sigma_next).astype("float32") return StepResult(output=out) res = ResourceRequest( compute_seconds=self.cost.predict(int(np.prod(x.shape)), 3.0), # 3 passes resident_bytes=int(x.nbytes), peak_activation_bytes=int(x.nbytes * 3)) return WorkPlan( loop_id=self.loop_id, instance_id=st.instance_id, kind=WorkUnitKind.DIFFUSION_STEP, shape_sig=ShapeSignature(WorkUnitKind.DIFFUSION_STEP, dims=tuple(x.shape), extra=(("stage", self.stage),)), resources=res, payload={"branch": "combined", "step": i, "stage": self.stage}, run=run, label=f"ltx2.{self.stage}.{i}") def advance(self, st: LoopState, result: StepResult) -> LoopState: st.latents["video"] = result.output["latents"] if "audio_latents" in result.output: st.latents["audio"] = result.output["audio_latents"] if st.profile == ExecutionProfile.ROLLOUT: st.trajectory.append({"step": st.step_idx, "stage": self.stage, "latents": np.asarray(st.latents["video"]).copy()}) st.step_idx += 1 return st def finalize(self, st: LoopState) -> LoopResult: outs = {"latents": st.latents["video"]} if st.scratch.get("want_audio") and "audio" in st.latents: outs["audio_latents"] = st.latents["audio"] # threaded to the next stage / decoded return LoopResult(outputs=outs, metrics={f"{self.stage}_steps": float(st.step_idx)}, behavior=st.trajectory or None)