| 1 | """ |
| 2 | LTX-2 Diffusion Model Wrapper for DMD distillation. |
| 3 | |
| 4 | This wrapper adapts LTX-2's audio-video joint generation model for use in |
| 5 | DMD (Distribution Matching Distillation) training. |
| 6 | |
| 7 | Model Architecture: |
| 8 | - patch_size = (1, 1, 1): No spatial/temporal grouping |
| 9 | - Patchification: Simple reshape [B, C, F, H, W] → [B, F*H*W, C] |
| 10 | - Each token: 128-dimensional latent vector (one per spatial-temporal position) |
| 11 | - Model input projection: Linear(128, 4096) |
| 12 | """ |
| 13 | |
| 14 | from typing import Any, Dict, Optional, Tuple |
| 15 | |
| 16 | import torch |
| 17 | import torch.nn as nn |
| 18 | |
| 19 | from ltx_core.components.patchifiers import ( |
| 20 | AudioPatchifier, |
| 21 | VideoLatentPatchifier, |
| 22 | get_pixel_coords, |
| 23 | ) |
| 24 | from ltx_core.guidance.perturbations import ( |
| 25 | BatchedPerturbationConfig, |
| 26 | Perturbation, |
| 27 | PerturbationConfig, |
| 28 | PerturbationType, |
| 29 | ) |
| 30 | from ltx_core.loader import LoraPathStrengthAndSDOps |
| 31 | from ltx_core.loader.registry import Registry |
| 32 | from ltx_core.model.transformer import LTXModel |
| 33 | from ltx_core.quantization import QuantizationPolicy |
| 34 | from ltx_core.model.transformer.modality import Modality |
| 35 | from ltx_core.types import ( |
| 36 | SpatioTemporalScaleFactors, |
| 37 | VideoLatentShape, |
| 38 | ) |
| 39 | |
| 40 | |
| 41 | class LTX2DiffusionWrapper(nn.Module): |
| 42 | """ |
| 43 | Wrapper for LTX-2 model to provide DMD-compatible interface. |
| 44 | |
| 45 | Handles: |
| 46 | - Input format conversion: [B, F, C, H, W] -> Modality |
| 47 | - Timestep handling: sigma values for all tokens |
| 48 | - Position computation for video (3D) and audio (1D) |
| 49 | - Output format: x0 predictions for both video and audio |
| 50 | |
| 51 | Uses official LTX-2 patchifiers (patch_size=1) to ensure consistency |
| 52 | with the pretrained model weights. |
| 53 | """ |
| 54 | |
| 55 | # Time alignment constants |
| 56 | VIDEO_LATENT_FPS = 3.0 # 24fps / 8 (VAE compression) |
| 57 | AUDIO_LATENT_FPS = 25.0 # 16kHz / 160 / 4 (mel hop / VAE compression) |
| 58 | ALIGNMENT_RATIO = AUDIO_LATENT_FPS / VIDEO_LATENT_FPS # ~8.33 |
| 59 | |
| 60 | # Video FPS for position computation |
| 61 | VIDEO_FPS = 24.0 |
| 62 | |
| 63 | # VAE scale factors (temporal=8, height=32, width=32) |
| 64 | DEFAULT_SCALE_FACTORS = SpatioTemporalScaleFactors.default() |
| 65 | |
| 66 | def __init__( |
| 67 | self, |
| 68 | model: LTXModel, |
| 69 | video_height: int = 512, |
| 70 | video_width: int = 768, |
| 71 | vae_spatial_compression: int = 32, |
| 72 | ): |
| 73 | """ |
| 74 | Args: |
| 75 | model: X0Model instance (wraps velocity model, returns x0 predictions) |
| 76 | video_height: Video height in pixels |
| 77 | video_width: Video width in pixels |
| 78 | vae_spatial_compression: VAE spatial compression factor |
| 79 | """ |
| 80 | super().__init__() |
| 81 | self.model = model |
| 82 | self.video_height = video_height |
| 83 | self.video_width = video_width |
| 84 | self.vae_spatial_compression = vae_spatial_compression |
| 85 | |
| 86 | # Compute latent dimensions |
| 87 | self.latent_height = video_height // vae_spatial_compression # 16 |
| 88 | self.latent_width = video_width // vae_spatial_compression # 24 |
| 89 | |
| 90 | # Official patchifiers with patch_size=1 (no spatial grouping) |
| 91 | self.video_patchifier = VideoLatentPatchifier(patch_size=1) |
| 92 | self.audio_patchifier = AudioPatchifier(patch_size=1) |
| 93 | |
| 94 | # Frame sequence length: with patch_size=1, each spatial position is one token |
| 95 | # For 512x768: H'*W' = 16*24 = 384 tokens per frame |
| 96 | self.video_frame_seqlen = self.latent_height * self.latent_width # 384 |
| 97 | |
| 98 | def set_module_grad(self, module_grad: Dict[str, bool]) -> None: |
| 99 | """ |
| 100 | Set gradient requirements for model components. |
| 101 | |
| 102 | Args: |
| 103 | module_grad: Dict mapping component names to requires_grad flags |
| 104 | """ |
| 105 | if module_grad.get("model", True): |
| 106 | self.model.requires_grad_(True) |
| 107 | else: |
| 108 | self.model.requires_grad_(False) |
| 109 | self.model.eval() |
| 110 | |
| 111 | def enable_gradient_checkpointing(self) -> None: |
| 112 | """Enable gradient checkpointing for memory efficiency.""" |
| 113 | if hasattr(self.model, "velocity_model"): |
| 114 | self.model.velocity_model.set_gradient_checkpointing(True) |
| 115 | elif hasattr(self.model, "set_gradient_checkpointing"): |
| 116 | self.model.set_gradient_checkpointing(True) |
| 117 | |
| 118 | def _flatten_video_latent( |
| 119 | self, |
| 120 | video_latent: torch.Tensor, |
| 121 | ) -> torch.Tensor: |
| 122 | """ |
| 123 | Flatten video latent from [B, F, C, H, W] to [B, T, C] using patch_size=1. |
| 124 | |
| 125 | With patch_size=1, this is a simple reshape — no spatial grouping. |
| 126 | The official VideoLatentPatchifier(patch_size=1) does: |
| 127 | "b c (f 1) (h 1) (w 1) -> b (f h w) (c 1 1 1)" = "b c f h w -> b (f h w) c" |
| 128 | |
| 129 | Args: |
| 130 | video_latent: Shape [B, F, C, H, W] where |
| 131 | - F: number of latent frames |
| 132 | - C: latent channels (128) |
| 133 | - H, W: latent spatial dimensions (16, 24) |
| 134 | |
| 135 | Returns: |
| 136 | Flattened tensor [B, T, C] where: |
| 137 | - T = F * H * W (e.g., 16 * 16 * 24 = 6144) |
| 138 | - C = 128 (unchanged, since patch_size=1) |
| 139 | """ |
| 140 | B, F, C, H, W = video_latent.shape |
| 141 | assert C == 128, ( |
| 142 | f"Expected video latent C=128 at dim 2, got shape {video_latent.shape}. " |
| 143 | f"Input should be [B, F, C, H, W] with C=128." |
| 144 | ) |
| 145 | |
| 146 | # Convert from [B, F, C, H, W] to [B, C, F, H, W] (official format) |
| 147 | video_latent = video_latent.permute(0, 2, 1, 3, 4) |
| 148 | |
| 149 | # Use official patchifier: [B, C, F, H, W] -> [B, F*H*W, C] |
| 150 | # With patch_size=1 this is equivalent to: |
| 151 | # einops.rearrange(x, "b c f h w -> b (f h w) c") |
| 152 | video_latent = self.video_patchifier.patchify(video_latent) |
| 153 | |
| 154 | return video_latent |
| 155 | |
| 156 | def _unflatten_video_latent( |
| 157 | self, |
| 158 | flat_latent: torch.Tensor, |
| 159 | num_frames: int, |
| 160 | ) -> torch.Tensor: |
| 161 | """ |
| 162 | Unflatten video latent from [B, T, C] back to [B, F, C, H, W]. |
| 163 | |
| 164 | Args: |
| 165 | flat_latent: Shape [B, T, C] where C = 128 (patch_size=1) |
| 166 | num_frames: Number of latent frames F |
| 167 | |
| 168 | Returns: |
| 169 | Video latent [B, F, C, H, W] |
| 170 | """ |
| 171 | B, T, C = flat_latent.shape |
| 172 | H = self.latent_height |
| 173 | W = self.latent_width |
| 174 | F = num_frames |
| 175 | |
| 176 | # Use official unpatchifier: [B, T, C] -> [B, C, F, H, W] |
| 177 | output_shape = VideoLatentShape( |
| 178 | batch=B, channels=C, frames=F, height=H, width=W |
| 179 | ) |
| 180 | video_latent = self.video_patchifier.unpatchify(flat_latent, output_shape) |
| 181 | |
| 182 | # Convert from [B, C, F, H, W] to [B, F, C, H, W] (DMD format) |
| 183 | video_latent = video_latent.permute(0, 2, 1, 3, 4) |
| 184 | |
| 185 | return video_latent |
| 186 | |
| 187 | def _compute_video_positions( |
| 188 | self, |
| 189 | video_latent: torch.Tensor, |
| 190 | downscale_factor: int = 1, |
| 191 | start_frame: int = 0, |
| 192 | ) -> torch.Tensor: |
| 193 | """ |
| 194 | Compute 3D position indices for video tokens with [start, end) bounds. |
| 195 | |
| 196 | Uses the official VideoLatentPatchifier.get_patch_grid_bounds() and |
| 197 | get_pixel_coords() to ensure consistency with the pretrained model. |
| 198 | |
| 199 | The RoPE computation expects positions in the format [B, 3, T, 2] where: |
| 200 | - dim 1 (size 3): temporal, height, width dimensions |
| 201 | - dim 3 (size 2): [start, end) bounds for each patch |
| 202 | |
| 203 | Returns: |
| 204 | Position tensor of shape [B, 3, T, 2] with patch bounds in pixel space |
| 205 | """ |
| 206 | B, F, C, H, W = video_latent.shape |
| 207 | device = video_latent.device |
| 208 | |
| 209 | # Build VideoLatentShape for the patchifier |
| 210 | video_shape = VideoLatentShape( |
| 211 | batch=B, channels=C, frames=F, height=H, width=W |
| 212 | ) |
| 213 | |
| 214 | # Get patch grid bounds in latent coordinates: [B, 3, T, 2] |
| 215 | # With patch_size=1, each token covers [i, i+1) in each dimension |
| 216 | latent_coords = self.video_patchifier.get_patch_grid_bounds( |
| 217 | output_shape=video_shape, |
| 218 | device=device, |
| 219 | ) |
| 220 | if start_frame != 0: |
| 221 | latent_coords = latent_coords.clone() |
| 222 | latent_coords[:, 0, :, :] += int(start_frame) |
| 223 | |
| 224 | # Convert to pixel coordinates using official helper |
| 225 | # Applies scale_factors (temporal=8, height=32, width=32) |
| 226 | # and causal_fix (first frame temporal offset) |
| 227 | pixel_coords = get_pixel_coords( |
| 228 | latent_coords=latent_coords, |
| 229 | scale_factors=self.DEFAULT_SCALE_FACTORS, |
| 230 | causal_fix=True, |
| 231 | ).float() |
| 232 | |
| 233 | # Convert temporal dimension from frames to seconds (divide by fps=24) |
| 234 | # This matches VideoLatentTools.create_initial_state |
| 235 | pixel_coords[:, 0, ...] = pixel_coords[:, 0, ...] / self.VIDEO_FPS |
| 236 | |
| 237 | if downscale_factor != 1: |
| 238 | pixel_coords = pixel_coords.clone() |
| 239 | pixel_coords[:, 1, ...] *= downscale_factor |
| 240 | pixel_coords[:, 2, ...] *= downscale_factor |
| 241 | |
| 242 | return pixel_coords |
| 243 | |
| 244 | def _compute_slot_zero_video_positions( |
| 245 | self, |
| 246 | video_latent: torch.Tensor, |
| 247 | downscale_factor: int = 1, |
| 248 | ) -> torch.Tensor: |
| 249 | """Place every memory slot at the same local temporal position.""" |
| 250 | _, num_frames, _, _, _ = video_latent.shape |
| 251 | if num_frames <= 0: |
| 252 | return self._compute_video_positions( |
| 253 | video_latent, downscale_factor=downscale_factor |
| 254 | ) |
| 255 | single_frame_positions = self._compute_video_positions( |
| 256 | video_latent[:, :1], downscale_factor=downscale_factor |
| 257 | ) |
| 258 | return single_frame_positions.repeat(1, 1, num_frames, 1) |
| 259 | |
| 260 | def _compute_slot_center_video_positions( |
| 261 | self, |
| 262 | video_latent: torch.Tensor, |
| 263 | *, |
| 264 | downscale_factor: int = 1, |
| 265 | position_offset: float = 0.0, |
| 266 | slot_stride: float = 50.0, |
| 267 | negative: bool = False, |
| 268 | ) -> torch.Tensor: |
| 269 | """Center each memory slot on a fixed virtual timestamp.""" |
| 270 | batch_size, num_frames, _, height, width = video_latent.shape |
| 271 | if num_frames <= 0: |
| 272 | return self._compute_video_positions( |
| 273 | video_latent, downscale_factor=downscale_factor |
| 274 | ) |
| 275 | |
| 276 | single_frame_positions = self._compute_video_positions( |
| 277 | video_latent[:, :1], downscale_factor=downscale_factor |
| 278 | ) |
| 279 | positions = single_frame_positions.repeat(1, 1, num_frames, 1) |
| 280 | tokens_per_frame = int(height) * int(width) |
| 281 | slot_indices = torch.arange( |
| 282 | num_frames, device=video_latent.device, dtype=positions.dtype |
| 283 | ) |
| 284 | if negative: |
| 285 | centers = float(position_offset) - ( |
| 286 | float(num_frames - 1) - slot_indices |
| 287 | ) * float(slot_stride) |
| 288 | else: |
| 289 | centers = float(position_offset) + slot_indices * float(slot_stride) |
| 290 | midpoint = ( |
| 291 | single_frame_positions[:, 0, :1, 0] |
| 292 | + single_frame_positions[:, 0, :1, 1] |
| 293 | ) * 0.5 |
| 294 | shifts = centers.repeat_interleave(tokens_per_frame).view(1, 1, -1, 1) |
| 295 | shifts = shifts - midpoint.view(batch_size, 1, 1, 1) |
| 296 | positions = positions.clone() |
| 297 | positions[:, 0, ...] = positions[:, 0, ...] + shifts[:, 0, ...] |
| 298 | return positions |
| 299 | |
| 300 | # Audio timing constants (from AudioPatchifier defaults) |
| 301 | AUDIO_SAMPLE_RATE = 16000 |
| 302 | AUDIO_HOP_LENGTH = 160 |
| 303 | AUDIO_LATENT_DOWNSAMPLE_FACTOR = 4 |
| 304 | AUDIO_IS_CAUSAL = True |
| 305 | |
| 306 | def _get_audio_latent_time_in_sec( |
| 307 | self, |
| 308 | start_latent: int, |
| 309 | end_latent: int, |
| 310 | dtype: torch.dtype, |
| 311 | device: torch.device, |
| 312 | ) -> torch.Tensor: |
| 313 | """ |
| 314 | Converts latent indices into real-time seconds while honoring causal |
| 315 | offsets and the configured hop length. |
| 316 | |
| 317 | Matches AudioPatchifier._get_audio_latent_time_in_sec exactly. |
| 318 | """ |
| 319 | audio_latent_frame = torch.arange(start_latent, end_latent, dtype=dtype, device=device) |
| 320 | audio_mel_frame = audio_latent_frame * self.AUDIO_LATENT_DOWNSAMPLE_FACTOR |
| 321 | |
| 322 | if self.AUDIO_IS_CAUSAL: |
| 323 | # Frame offset for causal alignment. |
| 324 | causal_offset = 1 |
| 325 | audio_mel_frame = (audio_mel_frame + causal_offset - self.AUDIO_LATENT_DOWNSAMPLE_FACTOR).clip(min=0) |
| 326 | |
| 327 | return audio_mel_frame * self.AUDIO_HOP_LENGTH / self.AUDIO_SAMPLE_RATE |
| 328 | |
| 329 | def _compute_audio_positions( |
| 330 | self, |
| 331 | audio_latent: torch.Tensor, |
| 332 | start_frame: int = 0, |
| 333 | ) -> torch.Tensor: |
| 334 | """ |
| 335 | Compute 1D temporal positions for audio tokens with [start, end) bounds. |
| 336 | |
| 337 | The RoPE computation expects positions in the format [B, 1, T, 2] where: |
| 338 | - dim 1 (size 1): temporal dimension only (audio is 1D) |
| 339 | - dim 3 (size 2): [start, end) bounds in seconds |
| 340 | |
| 341 | Returns: |
| 342 | Position tensor of shape [B, 1, T, 2] with temporal bounds in seconds |
| 343 | """ |
| 344 | B, T, C = audio_latent.shape |
| 345 | device = audio_latent.device |
| 346 | |
| 347 | # Compute start timings for each audio frame |
| 348 | start_timings = self._get_audio_latent_time_in_sec( |
| 349 | int(start_frame), int(start_frame) + T, torch.float32, device |
| 350 | ) |
| 351 | start_timings = start_timings.unsqueeze(0).expand(B, -1).unsqueeze(1) # [B, 1, T] |
| 352 | |
| 353 | # Compute end timings for each audio frame (shifted by 1) |
| 354 | end_timings = self._get_audio_latent_time_in_sec( |
| 355 | int(start_frame) + 1, int(start_frame) + T + 1, torch.float32, device |
| 356 | ) |
| 357 | end_timings = end_timings.unsqueeze(0).expand(B, -1).unsqueeze(1) # [B, 1, T] |
| 358 | |
| 359 | # Stack to create [B, 1, T, 2] with [start, end) bounds |
| 360 | positions = torch.stack([start_timings, end_timings], dim=-1) |
| 361 | |
| 362 | return positions |
| 363 | |
| 364 | @staticmethod |
| 365 | def _normalize_memory_segment_lengths( |
| 366 | segment_lengths: tuple[tuple[int, ...], ...] | tuple[int, ...] | None, |
| 367 | total_seq_len: int, |
| 368 | ) -> list[int]: |
| 369 | if total_seq_len <= 0: |
| 370 | return [] |
| 371 | lengths_obj: Any = segment_lengths |
| 372 | if ( |
| 373 | lengths_obj |
| 374 | and len(lengths_obj) > 0 |
| 375 | and isinstance(lengths_obj[0], (tuple, list)) |
| 376 | ): |
| 377 | lengths_obj = lengths_obj[0] |
| 378 | if lengths_obj: |
| 379 | lengths = [max(0, int(length)) for length in lengths_obj] |
| 380 | if sum(lengths) == total_seq_len and all(length > 0 for length in lengths): |
| 381 | return lengths |
| 382 | return [int(total_seq_len)] |
| 383 | |
| 384 | def _compute_slot_center_audio_positions( |
| 385 | self, |
| 386 | audio_latent: torch.Tensor, |
| 387 | *, |
| 388 | segment_lengths: tuple[tuple[int, ...], ...] | tuple[int, ...] | None = None, |
| 389 | position_offset: float = 0.0, |
| 390 | slot_stride: float = 50.0, |
| 391 | negative: bool = False, |
| 392 | ) -> torch.Tensor: |
| 393 | batch_size, total_seq_len, _ = audio_latent.shape |
| 394 | if total_seq_len <= 0: |
| 395 | return self._compute_audio_positions(audio_latent) |
| 396 | |
| 397 | lengths = self._normalize_memory_segment_lengths( |
| 398 | segment_lengths, total_seq_len |
| 399 | ) |
| 400 | pieces: list[torch.Tensor] = [] |
| 401 | start = 0 |
| 402 | for slot_idx, length in enumerate(lengths): |
| 403 | segment = audio_latent[:, start : start + length] |
| 404 | piece = self._compute_audio_positions(segment) |
| 405 | if negative: |
| 406 | center = float(position_offset) - float( |
| 407 | len(lengths) - 1 - slot_idx |
| 408 | ) * float(slot_stride) |
| 409 | else: |
| 410 | center = float(position_offset) + float(slot_idx) * float(slot_stride) |
| 411 | midpoint = (piece[:, 0, :1, 0] + piece[:, 0, -1:, 1]) * 0.5 |
| 412 | piece = piece.clone() |
| 413 | piece[:, 0, ...] = piece[:, 0, ...] + ( |
| 414 | center - midpoint |
| 415 | ).view(batch_size, 1, 1) |
| 416 | pieces.append(piece) |
| 417 | start += length |
| 418 | return torch.cat(pieces, dim=2) |
| 419 | |
| 420 | def _compute_timesteps_for_tokens( |
| 421 | self, |
| 422 | sigma: torch.Tensor, |
| 423 | num_tokens: int, |
| 424 | tokens_per_frame: int, |
| 425 | ) -> torch.Tensor: |
| 426 | """ |
| 427 | Expand sigma to per-token timesteps. |
| 428 | |
| 429 | In the official pipeline, timesteps = denoise_mask * sigma, producing |
| 430 | shape [B, T, 1]. Here we replicate sigma to each token belonging to |
| 431 | the same frame and add a trailing dimension for broadcasting with |
| 432 | the latent channels. |
| 433 | |
| 434 | Args: |
| 435 | sigma: Shape [B] or [B, F] - sigma values per frame |
| 436 | num_tokens: Total number of tokens |
| 437 | tokens_per_frame: Number of tokens per frame |
| 438 | |
| 439 | Returns: |
| 440 | Timesteps tensor [B, T, 1] for correct broadcasting with [B, T, C] |
| 441 | """ |
| 442 | B = sigma.shape[0] |
| 443 | |
| 444 | if sigma.dim() == 1: |
| 445 | # Single sigma per sample -> expand to all tokens |
| 446 | return sigma.view(B, 1, 1).expand(B, num_tokens, 1) |
| 447 | else: |
| 448 | # Per-frame sigma [B, F] -> expand to per-token [B, T, 1] |
| 449 | F = sigma.shape[1] |
| 450 | expanded = sigma.unsqueeze(2).expand(B, F, tokens_per_frame).reshape(B, -1) |
| 451 | return expanded.unsqueeze(-1) # [B, T, 1] |
| 452 | |
| 453 | def forward( |
| 454 | self, |
| 455 | noisy_image_or_video: torch.Tensor, |
| 456 | conditional_dict: Dict[str, Any], |
| 457 | timestep: torch.Tensor, |
| 458 | noisy_audio: Optional[torch.Tensor] = None, |
| 459 | audio_timestep: Optional[torch.Tensor] = None, |
| 460 | memory_video: Optional[torch.Tensor] = None, |
| 461 | memory_audio: Optional[torch.Tensor] = None, |
| 462 | memory_audio_timestep: Optional[torch.Tensor] = None, |
| 463 | memory_audio_segment_lengths: tuple[tuple[int, ...], ...] | None = None, |
| 464 | memory_position_mode: str = "reference", |
| 465 | memory_position_offset: float = 0.0, |
| 466 | memory_position_slot_stride: float = 50.0, |
| 467 | memory_downscale_factor: int = 1, |
| 468 | skip_a2v_cross_attn: bool = False, |
| 469 | skip_v2a_cross_attn: bool = False, |
| 470 | skip_video_self_attn: bool = False, |
| 471 | skip_audio_self_attn: bool = False, |
| 472 | use_causal_timestep: bool = False, # ignored, for API compatibility |
| 473 | **kwargs, |
| 474 | ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: |
| 475 | """ |
| 476 | Forward pass for DMD distillation. |
| 477 | |
| 478 | Args: |
| 479 | noisy_image_or_video: Noisy video latent [B, F, C, H, W] |
| 480 | conditional_dict: Dictionary containing: |
| 481 | - video_context: [B, seq_len, dim] |
| 482 | - audio_context: [B, seq_len, dim] |
| 483 | - attention_mask: [B, seq_len] |
| 484 | timestep: Sigma values [B] or [B, F] |
| 485 | noisy_audio: Noisy audio latent [B, F_a, C_audio] (optional) |
| 486 | where C_audio = 128 (= 8 channels * 16 mel_bins, post-patchify) |
| 487 | audio_timestep: Audio sigma values [B] or [B, F_a] (optional) |
| 488 | memory_audio: Optional clean memory-audio prefix [B, F_mem_a, C_audio] |
| 489 | memory_audio_timestep: Optional memory-audio sigma values [B] or [B, F_mem_a] |
| 490 | |
| 491 | Returns: |
| 492 | Tuple of (video_x0_pred, audio_x0_pred) |
| 493 | - video_x0_pred: [B, F, C, H, W] |
| 494 | - audio_x0_pred: [B, F_a, C_audio] or None |
| 495 | """ |
| 496 | B = noisy_image_or_video.shape[0] |
| 497 | num_video_frames = noisy_image_or_video.shape[1] |
| 498 | device = noisy_image_or_video.device |
| 499 | memory_position_mode = str(memory_position_mode).lower() |
| 500 | if memory_position_mode == "reference": |
| 501 | memory_position_mode = "legacy" |
| 502 | allowed_memory_position_modes = { |
| 503 | "legacy", |
| 504 | "prefix_continuous", |
| 505 | "slot_zero", |
| 506 | "slot_center", |
| 507 | "negative_slot_center", |
| 508 | "reference_offset", |
| 509 | } |
| 510 | if memory_position_mode not in allowed_memory_position_modes: |
| 511 | raise ValueError( |
| 512 | "memory_position_mode must be one of {'reference', 'legacy', " |
| 513 | "'prefix_continuous', 'slot_zero', 'slot_center', " |
| 514 | "'negative_slot_center', 'reference_offset'}, " |
| 515 | f"got {memory_position_mode}" |
| 516 | ) |
| 517 | if memory_video is not None and int(memory_video.shape[1]) == 0: |
| 518 | memory_video = None |
| 519 | |
| 520 | # Flatten target video latent: [B, F, C, H, W] -> [B, T, C] |
| 521 | # With patch_size=1: T = F*H*W, C = 128 |
| 522 | target_video_flat = self._flatten_video_latent(noisy_image_or_video) |
| 523 | num_target_video_tokens = target_video_flat.shape[1] |
| 524 | |
| 525 | # Compute target video positions / timesteps |
| 526 | target_video_position_start = ( |
| 527 | int(memory_video.shape[1]) |
| 528 | if memory_position_mode == "prefix_continuous" and memory_video is not None |
| 529 | else 0 |
| 530 | ) |
| 531 | target_video_positions = self._compute_video_positions( |
| 532 | noisy_image_or_video, |
| 533 | start_frame=target_video_position_start, |
| 534 | ) |
| 535 | target_video_timesteps = self._compute_timesteps_for_tokens( |
| 536 | timestep, num_target_video_tokens, self.video_frame_seqlen |
| 537 | ) |
| 538 | |
| 539 | memory_seq_len = 0 |
| 540 | if memory_video is not None: |
| 541 | memory_video_flat = self._flatten_video_latent(memory_video) |
| 542 | if memory_position_mode == "slot_zero": |
| 543 | memory_video_positions = self._compute_slot_zero_video_positions( |
| 544 | memory_video, |
| 545 | downscale_factor=memory_downscale_factor, |
| 546 | ) |
| 547 | elif memory_position_mode in {"slot_center", "negative_slot_center"}: |
| 548 | memory_video_positions = self._compute_slot_center_video_positions( |
| 549 | memory_video, |
| 550 | downscale_factor=memory_downscale_factor, |
| 551 | position_offset=float(memory_position_offset), |
| 552 | slot_stride=float(memory_position_slot_stride), |
| 553 | negative=memory_position_mode == "negative_slot_center", |
| 554 | ) |
| 555 | else: |
| 556 | memory_video_positions = self._compute_video_positions( |
| 557 | memory_video, downscale_factor=memory_downscale_factor |
| 558 | ) |
| 559 | if memory_position_mode == "reference_offset": |
| 560 | memory_video_positions = memory_video_positions.clone() |
| 561 | memory_video_positions[:, 0, ...] += float( |
| 562 | memory_position_offset |
| 563 | ) |
| 564 | memory_video_timesteps = torch.zeros( |
| 565 | B, |
| 566 | memory_video_flat.shape[1], |
| 567 | 1, |
| 568 | device=device, |
| 569 | dtype=target_video_timesteps.dtype, |
| 570 | ) |
| 571 | video_flat = torch.cat([memory_video_flat, target_video_flat], dim=1) |
| 572 | video_positions = torch.cat([memory_video_positions, target_video_positions], dim=2) |
| 573 | video_timesteps = torch.cat([memory_video_timesteps, target_video_timesteps], dim=1) |
| 574 | memory_seq_len = memory_video_flat.shape[1] |
| 575 | else: |
| 576 | video_flat = target_video_flat |
| 577 | video_positions = target_video_positions |
| 578 | video_timesteps = target_video_timesteps |
| 579 | |
| 580 | # Build video modality |
| 581 | video_sigma = timestep if timestep.dim() == 1 else timestep[:, 0] |
| 582 | video_modality = Modality( |
| 583 | latent=video_flat, |
| 584 | sigma=video_sigma, |
| 585 | timesteps=video_timesteps, |
| 586 | positions=video_positions, |
| 587 | context=conditional_dict["video_context"], |
| 588 | context_mask=conditional_dict.get("attention_mask"), |
| 589 | enabled=True, |
| 590 | ) |
| 591 | |
| 592 | # Build audio modality if provided |
| 593 | audio_modality = None |
| 594 | memory_audio_seq_len = 0 |
| 595 | if noisy_audio is None and (memory_audio is not None or memory_audio_timestep is not None): |
| 596 | raise ValueError("memory_audio requires noisy_audio") |
| 597 | if noisy_audio is not None: |
| 598 | target_audio = noisy_audio |
| 599 | target_audio_frames = target_audio.shape[1] |
| 600 | |
| 601 | # Use provided audio timestep or derive from video timestep |
| 602 | if audio_timestep is None: |
| 603 | # In bidirectional mode, audio uses same sigma as video. |
| 604 | # video timestep could be [B] or [B, F_v]. For audio we need [B] |
| 605 | # or [B, F_a]. If timestep is [B, F_v] (per-frame video), take the |
| 606 | # first frame's sigma since bidirectional uses uniform sigma anyway. |
| 607 | if timestep.dim() == 1: |
| 608 | audio_timestep = timestep # [B] |
| 609 | else: |
| 610 | # All video frames have same sigma in bidirectional mode, |
| 611 | # take the first frame's value and broadcast to audio frames |
| 612 | audio_timestep = timestep[:, 0] # [B] |
| 613 | |
| 614 | if audio_timestep.dim() == 1: |
| 615 | target_audio_timestep = audio_timestep[:, None].expand(B, target_audio_frames) |
| 616 | elif audio_timestep.shape == (B, target_audio_frames): |
| 617 | target_audio_timestep = audio_timestep |
| 618 | else: |
| 619 | raise ValueError( |
| 620 | "audio_timestep must have shape [B] or [B, F_a], " |
| 621 | f"got {tuple(audio_timestep.shape)} vs {(B, target_audio_frames)}" |
| 622 | ) |
| 623 | |
| 624 | if memory_audio_timestep is not None and memory_audio is None: |
| 625 | raise ValueError("memory_audio_timestep requires memory_audio") |
| 626 | |
| 627 | if memory_audio is not None: |
| 628 | memory_audio = memory_audio.to(device=device, dtype=target_audio.dtype) |
| 629 | memory_audio_seq_len = memory_audio.shape[1] |
| 630 | if memory_audio_timestep is None: |
| 631 | prefix_audio_timestep = torch.zeros( |
| 632 | B, |
| 633 | memory_audio_seq_len, |
| 634 | device=device, |
| 635 | dtype=target_audio_timestep.dtype, |
| 636 | ) |
| 637 | elif memory_audio_timestep.dim() == 1: |
| 638 | prefix_audio_timestep = memory_audio_timestep[:, None].expand(B, memory_audio_seq_len) |
| 639 | elif memory_audio_timestep.shape == (B, memory_audio_seq_len): |
| 640 | prefix_audio_timestep = memory_audio_timestep |
| 641 | else: |
| 642 | raise ValueError( |
| 643 | "memory_audio_timestep must have shape [B] or [B, F_mem_a], " |
| 644 | f"got {tuple(memory_audio_timestep.shape)} vs {(B, memory_audio_seq_len)}" |
| 645 | ) |
| 646 | |
| 647 | noisy_audio = torch.cat([memory_audio, target_audio], dim=1) |
| 648 | combined_audio_timestep = torch.cat([prefix_audio_timestep, target_audio_timestep], dim=1) |
| 649 | else: |
| 650 | noisy_audio = target_audio |
| 651 | combined_audio_timestep = target_audio_timestep |
| 652 | |
| 653 | num_audio_tokens = noisy_audio.shape[1] |
| 654 | audio_timesteps = self._compute_timesteps_for_tokens(combined_audio_timestep, num_audio_tokens, 1) |
| 655 | if memory_audio_seq_len > 0: |
| 656 | if memory_position_mode in {"slot_center", "negative_slot_center"}: |
| 657 | memory_audio_positions = self._compute_slot_center_audio_positions( |
| 658 | memory_audio, |
| 659 | segment_lengths=memory_audio_segment_lengths, |
| 660 | position_offset=float(memory_position_offset), |
| 661 | slot_stride=float(memory_position_slot_stride), |
| 662 | negative=memory_position_mode == "negative_slot_center", |
| 663 | ) |
| 664 | else: |
| 665 | memory_audio_positions = self._compute_audio_positions(memory_audio) |
| 666 | if memory_position_mode == "reference_offset": |
| 667 | memory_audio_positions = memory_audio_positions.clone() |
| 668 | memory_audio_positions[:, 0, ...] += float( |
| 669 | memory_position_offset |
| 670 | ) |
| 671 | target_audio_position_start = memory_audio_seq_len if memory_position_mode == "prefix_continuous" else 0 |
| 672 | target_audio_positions = self._compute_audio_positions( |
| 673 | target_audio, |
| 674 | start_frame=target_audio_position_start, |
| 675 | ) |
| 676 | audio_positions = torch.cat([memory_audio_positions, target_audio_positions], dim=2) |
| 677 | else: |
| 678 | audio_positions = self._compute_audio_positions(noisy_audio) |
| 679 | audio_sigma = target_audio_timestep[:, 0] |
| 680 | audio_modality = Modality( |
| 681 | latent=noisy_audio, |
| 682 | sigma=audio_sigma, |
| 683 | timesteps=audio_timesteps, |
| 684 | positions=audio_positions, |
| 685 | context=conditional_dict.get("audio_context", conditional_dict["video_context"]), |
| 686 | context_mask=conditional_dict.get("attention_mask"), |
| 687 | enabled=True, |
| 688 | ) |
| 689 | |
| 690 | # Forward through model. The optional perturbation flags let inference |
| 691 | # freeze one direction of cross-modal interaction without modifying the |
| 692 | # shared core transformer implementation. |
| 693 | perturbation_items: list[Perturbation] = [] |
| 694 | if skip_a2v_cross_attn: |
| 695 | perturbation_items.append( |
| 696 | Perturbation( |
| 697 | type=PerturbationType.SKIP_A2V_CROSS_ATTN, |
| 698 | blocks=None, |
| 699 | ) |
| 700 | ) |
| 701 | if skip_v2a_cross_attn: |
| 702 | perturbation_items.append( |
| 703 | Perturbation( |
| 704 | type=PerturbationType.SKIP_V2A_CROSS_ATTN, |
| 705 | blocks=None, |
| 706 | ) |
| 707 | ) |
| 708 | if skip_video_self_attn: |
| 709 | perturbation_items.append( |
| 710 | Perturbation( |
| 711 | type=PerturbationType.SKIP_VIDEO_SELF_ATTN, |
| 712 | blocks=None, |
| 713 | ) |
| 714 | ) |
| 715 | if skip_audio_self_attn: |
| 716 | perturbation_items.append( |
| 717 | Perturbation( |
| 718 | type=PerturbationType.SKIP_AUDIO_SELF_ATTN, |
| 719 | blocks=None, |
| 720 | ) |
| 721 | ) |
| 722 | |
| 723 | if perturbation_items: |
| 724 | perturbation_config = PerturbationConfig(perturbations=perturbation_items) |
| 725 | perturbations = BatchedPerturbationConfig( |
| 726 | [perturbation_config for _ in range(B)] |
| 727 | ) |
| 728 | else: |
| 729 | perturbations = BatchedPerturbationConfig.empty(batch_size=B) |
| 730 | |
| 731 | # The model returns x0 predictions (X0Model wraps velocity model) |
| 732 | video_x0, audio_x0 = self.model( |
| 733 | video=video_modality, |
| 734 | audio=audio_modality, |
| 735 | perturbations=perturbations, |
| 736 | ) |
| 737 | |
| 738 | # Unflatten video output: [B, T, C] -> [B, F, C, H, W] |
| 739 | if video_x0 is not None: |
| 740 | if memory_seq_len > 0: |
| 741 | video_x0 = video_x0[:, memory_seq_len:, :] |
| 742 | video_x0 = self._unflatten_video_latent(video_x0, num_video_frames) |
| 743 | if audio_x0 is not None and memory_audio_seq_len > 0: |
| 744 | audio_x0 = audio_x0[:, memory_audio_seq_len:, :] |
| 745 | |
| 746 | return video_x0, audio_x0 |
| 747 | |
| 748 | def load_state_dict(self, state_dict: Dict[str, Any], strict: bool = True): |
| 749 | """Load state dict, handling potential key mismatches.""" |
| 750 | # Remove 'model.' prefix if present |
| 751 | new_state_dict = {} |
| 752 | for k, v in state_dict.items(): |
| 753 | if k.startswith("model."): |
| 754 | new_state_dict[k] = v |
| 755 | else: |
| 756 | new_state_dict[f"model.{k}"] = v |
| 757 | |
| 758 | return super().load_state_dict(new_state_dict, strict=strict) |
| 759 | |
| 760 | |
| 761 | def create_ltx2_wrapper( |
| 762 | checkpoint_path: str, |
| 763 | gemma_path: str, |
| 764 | device: torch.device, |
| 765 | dtype: torch.dtype = torch.bfloat16, |
| 766 | video_height: int = 512, |
| 767 | video_width: int = 768, |
| 768 | loras: tuple[LoraPathStrengthAndSDOps, ...] = (), |
| 769 | registry: Registry | None = None, |
| 770 | quantization: QuantizationPolicy | None = None, |
| 771 | ) -> LTX2DiffusionWrapper: |
| 772 | """ |
| 773 | Factory function to create LTX2DiffusionWrapper from checkpoint. |
| 774 | |
| 775 | Args: |
| 776 | checkpoint_path: Path to LTX-2 checkpoint |
| 777 | gemma_path: Path to Gemma text encoder |
| 778 | device: Target device |
| 779 | dtype: Model dtype |
| 780 | video_height: Video height |
| 781 | video_width: Video width |
| 782 | |
| 783 | Returns: |
| 784 | Configured LTX2DiffusionWrapper |
| 785 | """ |
| 786 | from ltx_pipelines.utils.model_ledger import ModelLedger |
| 787 | |
| 788 | # IMPORTANT: Load to CPU first, then move to target device |
| 789 | # safetensors doesn't support device indices like "cuda:4" |
| 790 | # It only accepts "cuda" or "cpu" |
| 791 | ledger = ModelLedger( |
| 792 | dtype=dtype, |
| 793 | device=torch.device("cpu"), # Load to CPU first |
| 794 | checkpoint_path=checkpoint_path, |
| 795 | gemma_root_path=gemma_path, |
| 796 | loras=loras, |
| 797 | registry=registry, |
| 798 | quantization=quantization, |
| 799 | ) |
| 800 | |
| 801 | # Get X0Model (wraps velocity model) |
| 802 | x0_model = ledger.transformer() |
| 803 | |
| 804 | # Move to target device |
| 805 | if quantization is None: |
| 806 | x0_model = x0_model.to(device=device, dtype=dtype) |
| 807 | else: |
| 808 | # Quantized modules deliberately mix FP8 weights and FP32 scales. |
| 809 | x0_model = x0_model.to(device=device) |
| 810 | |
| 811 | wrapper = LTX2DiffusionWrapper( |
| 812 | model=x0_model, |
| 813 | video_height=video_height, |
| 814 | video_width=video_width, |
| 815 | ) |
| 816 | |
| 817 | return wrapper |
| 818 |