[feat] v2 ltx2: two-stage distilled GPU bring-up (LTX2Transformer3DModel 18.88B)

Official FastVideo/LTX2-Distilled-Diffusers. New torch_ltx2.py adapters (build_torch_* dispatch on
class):
- TorchLTX2DiT: patchify-internal; per-token timestep ones(B,tok,1)*sigma (sigma direct) + per-sample
  video_sigma; DiT predicts x0 so the adapter returns velocity=(x_t-x0)/sigma for the v2 flow-match step.
- TorchLTX2VAE: CausalVideoAutoencoder decode (internal per-channel un_normalize).
- TorchGemma: LTX2GemmaTextEncoderModel (Gemma + feature-extractor + connectors) -> last_hidden_state.
- ltx2 loop: real 128-ch latent geometry on cuda (32x spatial / 8x temporal; half-res base, 2x upsample).
e2e two-stage (8+3 steps) -> coherent, high-quality video (surfers at sunset), (3,9,512,768). NOTE:
still uses the v2 program's np.repeat upsampler between stages (the refine regenerates from noise so
output is faithful-quality); real LTX2LatentUpsampler swap-in is a follow-up.

v2/models/ltx2/loop.py

@@ -35,8 +35,20 @@ BASE_SIGMAS = [1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875,
 REFINE_SIGMAS = [0.909375, 0.725, 0.421875, 0.0]
 
 
-def ltx_base_latent_shape(req) -> tuple[int, int, int, int]:
+# Real LTX-2 latent: 128 channels; CausalVideoAutoencoder ~32x spatial / 8x temporal. The base stage
+# runs at HALF the full latent spatial res (stage-2 upsamples 2x), so divide by 64 (=32*2).
+LTX2_LATENT_CHANNELS = 128
+LTX2_SPATIAL_RATIO = 32
+LTX2_TEMPORAL_RATIO = 8
+
+
+def ltx_base_latent_shape(req, model=None) -> tuple[int, int, int, int]:
     d = req.diffusion
+    if model is not None and getattr(getattr(model, "platform", None), "device", "cpu") == "cuda":
+        t = (max(1, d.num_frames) - 1) // LTX2_TEMPORAL_RATIO + 1
+        h = max(1, d.height // (LTX2_SPATIAL_RATIO * 2))      # half-res base; stage-2 upsamples 2x
+        w = max(1, d.width // (LTX2_SPATIAL_RATIO * 2))
+        return (LTX2_LATENT_CHANNELS, t, h, w)
     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)
@@ -62,7 +74,7 @@ class LTX2DenoiseLoop:
         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")
+            x = (rng.standard_normal(ltx_base_latent_shape(req, model)) * 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,

v2/platform/backends/torch_adapters.py

@@ -279,14 +279,22 @@ def build_torch_dit(spec, instance, platform):
     ckpt = _require_checkpoint(spec)          # fail fast on a mis-stamped card, before any dist init
     _ensure_fastvideo_runtime()
     module = _load_component("TransformerLoader", ckpt, _fastvideo_args(spec))
-    return TorchWanDiT(module, device=_device(platform), dtype=_native_dtype(module))
+    device, dtype = _device(platform), _native_dtype(module)
+    if "LTX2" in type(module).__name__:                          # LTX2Transformer3DModel
+        from .torch_ltx2 import TorchLTX2DiT
+        return TorchLTX2DiT(module, device=device, dtype=dtype)
+    return TorchWanDiT(module, device=device, dtype=dtype)       # WanTransformer3DModel / Causal*
 
 
 def build_torch_vae(spec, instance, platform):
     ckpt = _require_checkpoint(spec)
     _ensure_fastvideo_runtime()
     module = _load_component("VAELoader", ckpt, _fastvideo_args(spec))
-    return TorchWanVAE(module, device=_device(platform), dtype=_native_dtype(module))
+    device, dtype = _device(platform), _native_dtype(module)
+    if "LTX2" in type(module).__name__:                          # LTX2CausalVideoAutoencoder
+        from .torch_ltx2 import TorchLTX2VAE
+        return TorchLTX2VAE(module, device=device, dtype=dtype)
+    return TorchWanVAE(module, device=device, dtype=dtype)
 
 
 def build_torch_text_encoder(spec, instance, platform):
@@ -296,4 +304,8 @@ def build_torch_text_encoder(spec, instance, platform):
     module = _load_component("TextEncoderLoader", ckpt, args)        # text_encoder subfolder
     # The tokenizer is a SIBLING subfolder (<root>/tokenizer), not under text_encoder/. BRINGUP risk E.
     tokenizer = _load_component("TokenizerLoader", os.path.join(_model_root(spec), "tokenizer"), args)
-    return TorchT5Encoder(module, tokenizer, device=_device(platform), dtype=_native_dtype(module))
+    device, dtype = _device(platform), _native_dtype(module)
+    if "Gemma" in type(module).__name__:                          # LTX2GemmaTextEncoderModel
+        from .torch_ltx2 import TorchGemma
+        return TorchGemma(module, tokenizer, device=device, dtype=dtype)
+    return TorchT5Encoder(module, tokenizer, device=device, dtype=dtype)

v2/platform/backends/torch_ltx2.py

@@ -0,0 +1,109 @@
+"""Real torch adapters for LTX-2 (design_v3 §15 two-stage distilled). Imported lazily on a GPU box.
+
+LTX-2 differs from Wan in three ways the adapters absorb so the v2 loops keep their numpy
+``dit(latent,text,sigma)->velocity`` / ``vae.decode`` / ``text_encoder.encode`` surface:
+  * DiT (``LTX2Transformer3DModel``) patchifies internally and takes a PER-TOKEN timestep
+    ``ones(B, token_count, 1) * sigma`` (sigma direct, 0..1 — NOT sigma*1000) plus a per-sample
+    ``video_sigma``; it predicts ``denoised`` (x0), so the adapter returns the flow-match velocity
+    ``(x_t - x0)/sigma`` the loop integrates.
+  * VAE (``LTX2CausalVideoAutoencoder``) un-normalizes internally on decode (per-channel stats).
+  * Text encoder (``LTX2GemmaTextEncoderModel``) runs Gemma + feature-extractor + connectors
+    (the projection) inside its forward, returning the projected ``last_hidden_state``.
+All run at the model's native precision (bf16 per LTX2T2VConfig). BRINGUP: confirm on box.
+"""
+from __future__ import annotations
+
+import torch
+
+from .torch_adapters import _to_numpy, _to_torch
+
+
+class TorchLTX2DiT:
+    """dit(latent[C,T,H,W], text_embed[seq,dim], sigma) -> velocity[C,T,H,W] (numpy in/out)."""
+
+    def __init__(self, module, *, device, dtype):
+        self.module = module.to(device=device, dtype=dtype).eval()
+        self.device, self.dtype = device, dtype
+        from fastvideo.models.dits.ltx2 import VideoLatentShape
+        self._VideoLatentShape = VideoLatentShape
+
+    @torch.no_grad()
+    def __call__(self, latent, text_embed, sigma, context=None):
+        hs = _to_torch(latent, device=self.device, dtype=self.dtype).unsqueeze(0)   # [1,C,T,H,W]
+        ehs = _to_torch(text_embed, device=self.device, dtype=self.dtype)
+        if ehs is not None:
+            ehs = ehs.unsqueeze(0)
+        s = float(sigma)
+        vshape = self._VideoLatentShape.from_torch_shape(tuple(hs.shape))
+        token_count = self.module.patchifier.get_token_count(vshape)
+        # per-token timestep = template * sigma (sigma direct, 0..1); per-sample video_sigma.
+        timestep = torch.full((1, token_count, 1), s, device=self.device, dtype=torch.float32)
+        video_sigma = torch.tensor([s], device=self.device, dtype=torch.float32)
+        from fastvideo.forward_context import set_forward_context
+        with set_forward_context(current_timestep=s, attn_metadata=None):
+            denoised = self.module(hidden_states=hs, encoder_hidden_states=ehs, timestep=timestep,
+                                   video_sigma=video_sigma, encoder_attention_mask=None)
+        if isinstance(denoised, tuple):
+            denoised = denoised[0]                       # t2v: video_out only (audio path unused)
+        # LTX-2 DiT predicts x0; the v2 loop integrates velocity. velocity = (x_t - x0)/sigma.
+        velocity = (hs.float() - denoised.float()) / max(s, 1e-6)
+        return _to_numpy(velocity.squeeze(0))
+
+    def copy_from(self, other):
+        self.module.load_state_dict(other.module.state_dict())
+
+    def clone(self):
+        import copy
+        c = TorchLTX2DiT.__new__(TorchLTX2DiT)
+        c.module = copy.deepcopy(self.module)
+        c.device, c.dtype, c._VideoLatentShape = self.device, self.dtype, self._VideoLatentShape
+        return c
+
+    def mse_grad_step(self, *a, **k):
+        raise NotImplementedError("GPU training surface (mse_grad_step) is a separate workstream (Risk F)")
+
+
+class TorchLTX2VAE:
+    """vae.decode(latent[C,T,H,W]) -> video[3,T,H,W] in [-1,1]. The LTX2 VideoDecoder un-normalizes
+    internally (per-channel statistics), so the adapter just marshals numpy<->torch."""
+
+    def __init__(self, module, *, device, dtype):
+        self.module = module.to(device=device, dtype=dtype).eval()
+        self.device, self.dtype = device, dtype
+
+    @torch.no_grad()
+    def decode(self, latent):
+        z = _to_torch(latent, device=self.device, dtype=self.dtype).unsqueeze(0)
+        video = self.module.decode(z)
+        if hasattr(video, "sample"):
+            video = video.sample
+        return _to_numpy(video.squeeze(0))
+
+    @torch.no_grad()
+    def encode(self, video):
+        x = _to_torch(video, device=self.device, dtype=self.dtype).unsqueeze(0)
+        dist = self.module.encode(x)
+        z = dist.mode() if hasattr(dist, "mode") else (dist.sample() if hasattr(dist, "sample") else dist)
+        return _to_numpy(z.squeeze(0))
+
+
+class TorchGemma:
+    """text_encoder.encode(text) -> embedding[seq,dim] (numpy out). LTX2GemmaTextEncoderModel runs
+    Gemma + feature-extractor + connectors (projection) in its forward and returns last_hidden_state."""
+
+    def __init__(self, module, tokenizer, *, device, dtype, max_length: int = 256):
+        self.module = module.to(device=device, dtype=dtype).eval()
+        self.tokenizer = tokenizer
+        self.device, self.dtype, self.max_length = device, dtype, max_length
+
+    @torch.no_grad()
+    def encode(self, text):
+        toks = self.tokenizer(text or "", return_tensors="pt", max_length=self.max_length,
+                              truncation=True, padding="max_length")
+        ids = toks.input_ids.to(self.device)
+        mask = toks.attention_mask.to(self.device)
+        from fastvideo.forward_context import set_forward_context
+        with set_forward_context(current_timestep=0, attn_metadata=None):
+            out = self.module(input_ids=ids, attention_mask=mask)
+        hidden = out.last_hidden_state if hasattr(out, "last_hidden_state") else out[0]
+        return _to_numpy(hidden.squeeze(0))