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