diff --git a/v2/models/ltx2/loop.py b/v2/models/ltx2/loop.py index beb4a3ae5..63f48d781 100644 --- a/v2/models/ltx2/loop.py +++ b/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, diff --git a/v2/platform/backends/torch_adapters.py b/v2/platform/backends/torch_adapters.py index e1457a6da..a59e27166 100644 --- a/v2/platform/backends/torch_adapters.py +++ b/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 (/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) diff --git a/v2/platform/backends/torch_ltx2.py b/v2/platform/backends/torch_ltx2.py new file mode 100644 index 000000000..97494eada --- /dev/null +++ b/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))