| 1 | """Model adapter used by causal autoregressive inference.""" |
| 2 | |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | from dataclasses import replace |
| 6 | |
| 7 | import torch |
| 8 | |
| 9 | from ltx_core.guidance.perturbations import BatchedPerturbationConfig |
| 10 | from ltx_core.model.transformer.modality import Modality |
| 11 | from ltx_core.utils import to_denoised |
| 12 | |
| 13 | from .cache import CausalCacheConfig |
| 14 | |
| 15 | |
| 16 | class CausalModelWrapper(torch.nn.Module): |
| 17 | """Adapt the release LTX velocity model to causal x0 prediction.""" |
| 18 | |
| 19 | def __init__(self, model: torch.nn.Module, patches_per_frame: int, cache: CausalCacheConfig): |
| 20 | super().__init__() |
| 21 | cache.validate() |
| 22 | self.model = model |
| 23 | self.patches_per_frame = patches_per_frame |
| 24 | self.cache = cache |
| 25 | |
| 26 | def forward( |
| 27 | self, |
| 28 | video: Modality, |
| 29 | audio: Modality | None, |
| 30 | action_cond: dict[str, torch.Tensor] | None, |
| 31 | kv_caches: list[dict], |
| 32 | video_start_frame: int, |
| 33 | audio_start_frame: int, |
| 34 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 35 | video = replace(video, sigma=torch.ones_like(video.sigma)) |
| 36 | if audio is not None: |
| 37 | audio = replace(audio, sigma=torch.ones_like(audio.sigma)) |
| 38 | velocity_video, velocity_audio = self.model( |
| 39 | video=video, |
| 40 | audio=audio, |
| 41 | perturbations=BatchedPerturbationConfig.empty(video.latent.shape[0]), |
| 42 | action_cond=action_cond, |
| 43 | kv_caches=kv_caches, |
| 44 | current_video_token_start=video_start_frame * self.patches_per_frame, |
| 45 | current_audio_token_start=audio_start_frame, |
| 46 | ) |
| 47 | video_sigma = video.timesteps.unsqueeze(-1) if video.timesteps.ndim == 2 else video.timesteps |
| 48 | video_x0 = to_denoised(video.latent, velocity_video, video_sigma) |
| 49 | audio_x0 = None |
| 50 | if audio is not None and velocity_audio is not None: |
| 51 | audio_sigma = audio.timesteps.unsqueeze(-1) if audio.timesteps.ndim == 2 else audio.timesteps |
| 52 | audio_x0 = to_denoised(audio.latent, velocity_audio, audio_sigma) |
| 53 | return video_x0, audio_x0 |
| 54 | |
| 55 | def init_caches( |
| 56 | self, |
| 57 | *, |
| 58 | batch_size: int, |
| 59 | video_frames: int, |
| 60 | audio_frames: int, |
| 61 | text_seq_len: int, |
| 62 | device: torch.device, |
| 63 | dtype: torch.dtype, |
| 64 | ) -> list[dict]: |
| 65 | patches_per_frame = self.patches_per_frame |
| 66 | return self.model.init_av_kv_caches( |
| 67 | batch_size=batch_size, |
| 68 | max_video_tokens=video_frames * patches_per_frame, |
| 69 | max_audio_tokens=audio_frames, |
| 70 | text_seq_len=text_seq_len, |
| 71 | device=device, |
| 72 | dtype=dtype, |
| 73 | video_local_attn_tokens=self.cache.video_local_attn_size * patches_per_frame, |
| 74 | video_sink_tokens=self.cache.video_sink_size * patches_per_frame, |
| 75 | video_ucpe_local_attn_tokens=self.cache.video_local_attn_size * patches_per_frame, |
| 76 | video_ucpe_sink_tokens=self.cache.video_sink_size * patches_per_frame, |
| 77 | audio_local_attn_tokens=self.cache.audio_local_attn_size, |
| 78 | audio_sink_tokens=self.cache.audio_sink_size, |
| 79 | ) |
| 80 |