| 1 | """Autoregressive causal rollout for joint video and audio generation.""" |
| 2 | |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | from dataclasses import dataclass |
| 6 | |
| 7 | import torch |
| 8 | |
| 9 | from ltx_core.model.transformer.modality import Modality |
| 10 | |
| 11 | from .cache import configure_bounded_caches |
| 12 | from .causal_wrapper import CausalModelWrapper |
| 13 | from .scheduling import ( |
| 14 | DEFAULT_CAUSAL_TIMESTEPS, |
| 15 | causal_audio_blocks, |
| 16 | causal_audio_frames, |
| 17 | causal_video_blocks, |
| 18 | resolve_causal_sigmas, |
| 19 | ) |
| 20 | |
| 21 | Block = tuple[int, int] |
| 22 | |
| 23 | |
| 24 | def _modality( |
| 25 | latent: torch.Tensor, |
| 26 | positions: torch.Tensor, |
| 27 | context: torch.Tensor, |
| 28 | sigma: float, |
| 29 | *, |
| 30 | context_mask: torch.Tensor | None, |
| 31 | ) -> Modality: |
| 32 | return Modality( |
| 33 | latent=latent, |
| 34 | sigma=torch.ones(latent.shape[0], device=latent.device, dtype=latent.dtype), |
| 35 | timesteps=torch.full(latent.shape[:2], sigma, device=latent.device, dtype=latent.dtype), |
| 36 | positions=positions, |
| 37 | context=context, |
| 38 | context_mask=context_mask, |
| 39 | ) |
| 40 | |
| 41 | |
| 42 | def _random_like(reference: torch.Tensor, generator: torch.Generator) -> torch.Tensor: |
| 43 | return torch.randn( |
| 44 | reference.shape, |
| 45 | device=reference.device, |
| 46 | dtype=reference.dtype, |
| 47 | generator=generator, |
| 48 | ) |
| 49 | |
| 50 | |
| 51 | def _advance_sample( |
| 52 | denoised: torch.Tensor, |
| 53 | next_sigma: float, |
| 54 | generator: torch.Generator, |
| 55 | ) -> torch.Tensor: |
| 56 | return (1 - next_sigma) * denoised + next_sigma * _random_like(denoised, generator) |
| 57 | |
| 58 | |
| 59 | @dataclass(frozen=True) |
| 60 | class _BlockForward: |
| 61 | """Bind rollout-wide context and expose a compact per-block model call.""" |
| 62 | |
| 63 | wrapper: CausalModelWrapper |
| 64 | caches: list[dict] |
| 65 | video_positions: torch.Tensor |
| 66 | audio_positions: torch.Tensor |
| 67 | video_context: torch.Tensor |
| 68 | audio_context: torch.Tensor |
| 69 | context_mask: torch.Tensor | None |
| 70 | action_cond: dict[str, torch.Tensor] |
| 71 | |
| 72 | @classmethod |
| 73 | def create( |
| 74 | cls, |
| 75 | *, |
| 76 | wrapper: CausalModelWrapper, |
| 77 | clean_video: torch.Tensor, |
| 78 | clean_audio: torch.Tensor, |
| 79 | video_positions: torch.Tensor, |
| 80 | audio_positions: torch.Tensor, |
| 81 | video_context: torch.Tensor, |
| 82 | audio_context: torch.Tensor, |
| 83 | context_mask: torch.Tensor | None, |
| 84 | action_cond: dict[str, torch.Tensor], |
| 85 | ) -> _BlockForward: |
| 86 | caches = wrapper.init_caches( |
| 87 | batch_size=clean_video.shape[0], |
| 88 | video_frames=video_positions.shape[2] // wrapper.patches_per_frame, |
| 89 | audio_frames=clean_audio.shape[1], |
| 90 | text_seq_len=video_context.shape[1], |
| 91 | device=clean_video.device, |
| 92 | dtype=clean_video.dtype, |
| 93 | ) |
| 94 | configure_bounded_caches( |
| 95 | wrapper, |
| 96 | caches, |
| 97 | video_positions, |
| 98 | audio_positions, |
| 99 | action_cond, |
| 100 | clean_video.dtype, |
| 101 | ) |
| 102 | return cls( |
| 103 | wrapper=wrapper, |
| 104 | caches=caches, |
| 105 | video_positions=video_positions, |
| 106 | audio_positions=audio_positions, |
| 107 | video_context=video_context, |
| 108 | audio_context=audio_context, |
| 109 | context_mask=context_mask, |
| 110 | action_cond=action_cond, |
| 111 | ) |
| 112 | |
| 113 | def __call__( |
| 114 | self, |
| 115 | video_latent: torch.Tensor, |
| 116 | video_block: Block, |
| 117 | video_sigma: float, |
| 118 | audio_latent: torch.Tensor | None = None, |
| 119 | audio_block: Block = (0, 0), |
| 120 | audio_sigma: float | None = None, |
| 121 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 122 | video_start, video_end = video_block |
| 123 | audio_start, audio_end = audio_block |
| 124 | patches_per_frame = self.wrapper.patches_per_frame |
| 125 | sliced_action = { |
| 126 | key: value[:, video_start:video_end] |
| 127 | if key in {"ucpe_viewmats", "ucpe_Ks"} |
| 128 | else value |
| 129 | for key, value in self.action_cond.items() |
| 130 | } |
| 131 | video = _modality( |
| 132 | video_latent, |
| 133 | self.video_positions[ |
| 134 | :, :, video_start * patches_per_frame : video_end * patches_per_frame |
| 135 | ], |
| 136 | self.video_context, |
| 137 | video_sigma, |
| 138 | context_mask=self.context_mask, |
| 139 | ) |
| 140 | audio = None |
| 141 | if audio_latent is not None: |
| 142 | audio = _modality( |
| 143 | audio_latent, |
| 144 | self.audio_positions[:, :, audio_start:audio_end], |
| 145 | self.audio_context, |
| 146 | video_sigma if audio_sigma is None else audio_sigma, |
| 147 | context_mask=self.context_mask, |
| 148 | ) |
| 149 | return self.wrapper( |
| 150 | video, |
| 151 | audio, |
| 152 | sliced_action, |
| 153 | self.caches, |
| 154 | video_start, |
| 155 | audio_start, |
| 156 | ) |
| 157 | |
| 158 | |
| 159 | @dataclass |
| 160 | class _RolloutBuffers: |
| 161 | """Noise sources and output tensors shared by all rollout blocks.""" |
| 162 | |
| 163 | initial_video: torch.Tensor |
| 164 | initial_audio: torch.Tensor |
| 165 | video_output: torch.Tensor |
| 166 | audio_output: torch.Tensor |
| 167 | clean_image: torch.Tensor |
| 168 | |
| 169 | @classmethod |
| 170 | def create( |
| 171 | cls, |
| 172 | clean_video: torch.Tensor, |
| 173 | clean_audio: torch.Tensor, |
| 174 | patches_per_frame: int, |
| 175 | generator: torch.Generator, |
| 176 | ) -> _RolloutBuffers: |
| 177 | clean_image = clean_video[:, :patches_per_frame] |
| 178 | video_output = torch.zeros_like(clean_video) |
| 179 | video_output[:, :patches_per_frame] = clean_image |
| 180 | return cls( |
| 181 | initial_video=_random_like(clean_video, generator), |
| 182 | initial_audio=_random_like(clean_audio, generator), |
| 183 | video_output=video_output, |
| 184 | audio_output=torch.zeros_like(clean_audio), |
| 185 | clean_image=clean_image, |
| 186 | ) |
| 187 | |
| 188 | |
| 189 | def _denoise_audio_prefix( |
| 190 | forward: _BlockForward, |
| 191 | clean_image: torch.Tensor, |
| 192 | initial_audio: torch.Tensor, |
| 193 | video_block: Block, |
| 194 | audio_block: Block, |
| 195 | sigmas: list[float], |
| 196 | generator: torch.Generator, |
| 197 | ) -> torch.Tensor: |
| 198 | audio_start, audio_end = audio_block |
| 199 | audio_sample = initial_audio[:, audio_start:audio_end] |
| 200 | for step, sigma in enumerate(sigmas): |
| 201 | _, denoised_audio = forward( |
| 202 | clean_image, |
| 203 | video_block, |
| 204 | 0.0, |
| 205 | audio_sample, |
| 206 | audio_block, |
| 207 | sigma, |
| 208 | ) |
| 209 | if denoised_audio is None: |
| 210 | raise RuntimeError("causal AV model returned no audio") |
| 211 | audio_sample = ( |
| 212 | denoised_audio |
| 213 | if step == len(sigmas) - 1 |
| 214 | else _advance_sample(denoised_audio, sigmas[step + 1], generator) |
| 215 | ) |
| 216 | return audio_sample |
| 217 | |
| 218 | |
| 219 | def _denoise_av_block( |
| 220 | forward: _BlockForward, |
| 221 | initial_video: torch.Tensor, |
| 222 | initial_audio: torch.Tensor, |
| 223 | video_block: Block, |
| 224 | audio_block: Block, |
| 225 | sigmas: list[float], |
| 226 | generator: torch.Generator, |
| 227 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 228 | video_start, video_end = video_block |
| 229 | audio_start, audio_end = audio_block |
| 230 | patches_per_frame = forward.wrapper.patches_per_frame |
| 231 | video_sample = initial_video[ |
| 232 | :, video_start * patches_per_frame : video_end * patches_per_frame |
| 233 | ] |
| 234 | audio_sample = initial_audio[:, audio_start:audio_end] |
| 235 | |
| 236 | for step, sigma in enumerate(sigmas): |
| 237 | denoised_video, denoised_audio = forward( |
| 238 | video_sample, |
| 239 | video_block, |
| 240 | sigma, |
| 241 | audio_sample, |
| 242 | audio_block, |
| 243 | sigma, |
| 244 | ) |
| 245 | if denoised_audio is None: |
| 246 | raise RuntimeError("causal AV model returned no audio") |
| 247 | if step == len(sigmas) - 1: |
| 248 | video_sample, audio_sample = denoised_video, denoised_audio |
| 249 | else: |
| 250 | next_sigma = sigmas[step + 1] |
| 251 | video_sample = _advance_sample(denoised_video, next_sigma, generator) |
| 252 | audio_sample = _advance_sample(denoised_audio, next_sigma, generator) |
| 253 | return video_sample, audio_sample |
| 254 | |
| 255 | |
| 256 | def _generate_audio_prefix( |
| 257 | forward: _BlockForward, |
| 258 | buffers: _RolloutBuffers, |
| 259 | video_block: Block, |
| 260 | audio_block: Block, |
| 261 | sigmas: list[float], |
| 262 | generator: torch.Generator, |
| 263 | ) -> None: |
| 264 | """Commit the image sink, generate its audio prefix, then refresh caches.""" |
| 265 | forward(buffers.clean_image, video_block, 0.0) |
| 266 | audio_prefix = _denoise_audio_prefix( |
| 267 | forward, |
| 268 | buffers.clean_image, |
| 269 | buffers.initial_audio, |
| 270 | video_block, |
| 271 | audio_block, |
| 272 | sigmas, |
| 273 | generator, |
| 274 | ) |
| 275 | audio_start, audio_end = audio_block |
| 276 | buffers.audio_output[:, audio_start:audio_end] = audio_prefix |
| 277 | forward(buffers.clean_image, video_block, 0.0, audio_prefix, audio_block, 0.0) |
| 278 | |
| 279 | |
| 280 | def _generate_av_blocks( |
| 281 | forward: _BlockForward, |
| 282 | buffers: _RolloutBuffers, |
| 283 | video_blocks: list[Block], |
| 284 | audio_blocks: list[Block], |
| 285 | sigmas: list[float], |
| 286 | generator: torch.Generator, |
| 287 | ) -> None: |
| 288 | """Generate, store, and cache all blocks after the image sink.""" |
| 289 | patches_per_frame = forward.wrapper.patches_per_frame |
| 290 | for video_block, audio_block in zip(video_blocks[1:], audio_blocks[1:], strict=True): |
| 291 | video_sample, audio_sample = _denoise_av_block( |
| 292 | forward, |
| 293 | buffers.initial_video, |
| 294 | buffers.initial_audio, |
| 295 | video_block, |
| 296 | audio_block, |
| 297 | sigmas, |
| 298 | generator, |
| 299 | ) |
| 300 | video_start, video_end = video_block |
| 301 | audio_start, audio_end = audio_block |
| 302 | buffers.video_output[ |
| 303 | :, video_start * patches_per_frame : video_end * patches_per_frame |
| 304 | ] = video_sample |
| 305 | buffers.audio_output[:, audio_start:audio_end] = audio_sample |
| 306 | forward(video_sample, video_block, 0.0, audio_sample, audio_block, 0.0) |
| 307 | |
| 308 | |
| 309 | @torch.no_grad() |
| 310 | def causal_rollout( # noqa: PLR0913 |
| 311 | *, |
| 312 | wrapper: CausalModelWrapper, |
| 313 | clean_video: torch.Tensor, |
| 314 | clean_audio: torch.Tensor, |
| 315 | video_positions: torch.Tensor, |
| 316 | audio_positions: torch.Tensor, |
| 317 | video_context: torch.Tensor, |
| 318 | audio_context: torch.Tensor, |
| 319 | context_mask: torch.Tensor | None, |
| 320 | action_cond: dict[str, torch.Tensor], |
| 321 | seed: int, |
| 322 | timesteps: tuple[int, ...] | list[int] = DEFAULT_CAUSAL_TIMESTEPS, |
| 323 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 324 | """Generate all audio-video blocks and refresh caches with clean outputs.""" |
| 325 | patches_per_frame = wrapper.patches_per_frame |
| 326 | video_frames = video_positions.shape[2] // patches_per_frame |
| 327 | chunk_size = wrapper.cache.video_chunk_size |
| 328 | video_blocks = causal_video_blocks(video_frames, chunk_size) |
| 329 | audio_blocks = causal_audio_blocks(video_frames, chunk_size) |
| 330 | if clean_audio.shape[1] != causal_audio_frames(video_frames, chunk_size): |
| 331 | raise ValueError("audio latent length does not match the causal AV block layout") |
| 332 | |
| 333 | sigmas = resolve_causal_sigmas(timesteps) |
| 334 | generator = torch.Generator(device=clean_video.device).manual_seed(seed) |
| 335 | buffers = _RolloutBuffers.create( |
| 336 | clean_video, |
| 337 | clean_audio, |
| 338 | patches_per_frame, |
| 339 | generator, |
| 340 | ) |
| 341 | forward = _BlockForward.create( |
| 342 | wrapper=wrapper, |
| 343 | clean_video=clean_video, |
| 344 | clean_audio=clean_audio, |
| 345 | video_positions=video_positions, |
| 346 | audio_positions=audio_positions, |
| 347 | video_context=video_context, |
| 348 | audio_context=audio_context, |
| 349 | context_mask=context_mask, |
| 350 | action_cond=action_cond, |
| 351 | ) |
| 352 | |
| 353 | _generate_audio_prefix( |
| 354 | forward, |
| 355 | buffers, |
| 356 | video_blocks[0], |
| 357 | audio_blocks[0], |
| 358 | sigmas, |
| 359 | generator, |
| 360 | ) |
| 361 | _generate_av_blocks(forward, buffers, video_blocks, audio_blocks, sigmas, generator) |
| 362 | return buffers.video_output, buffers.audio_output |
| 363 |