| 1 | """Autoregressive text/image-to-video rollout for Echo-WM Flash.""" |
| 2 | |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | from collections.abc import Iterator |
| 6 | |
| 7 | import torch |
| 8 | |
| 9 | from ltx_core.components.noisers import GaussianNoiser |
| 10 | from ltx_core.loader import LoraPathStrengthAndSDOps |
| 11 | from ltx_core.model.audio_vae import decode_audio as vae_decode_audio |
| 12 | from ltx_core.model.video_vae import decode_video as vae_decode_video |
| 13 | from ltx_core.model.video_vae.tiling import TilingConfig |
| 14 | from ltx_core.quantization import QuantizationPolicy |
| 15 | from ltx_core.tools import AudioLatentTools |
| 16 | from ltx_core.types import Audio, AudioLatentShape, VideoPixelShape |
| 17 | from ltx_causal import ( |
| 18 | DEFAULT_CAUSAL_TIMESTEPS, |
| 19 | CausalCacheConfig, |
| 20 | CausalModelWrapper, |
| 21 | causal_audio_frames, |
| 22 | causal_rollout, |
| 23 | causal_video_blocks, |
| 24 | ) |
| 25 | from ltx_pipelines.utils import ModelLedger, assert_resolution, cleanup_memory, combined_image_conditionings, encode_prompts, get_device |
| 26 | from ltx_pipelines.utils.args import ImageConditioningInput |
| 27 | from ltx_pipelines.utils.helpers import create_noised_state, noise_video_state |
| 28 | from ltx_pipelines.utils.types import PipelineComponents |
| 29 | |
| 30 | device = get_device() |
| 31 | |
| 32 | |
| 33 | class CausalTI2VidPipeline: |
| 34 | """Inference-only 4-step autoregressive I2V pipeline.""" |
| 35 | |
| 36 | def __init__( |
| 37 | self, |
| 38 | checkpoint_path: str, |
| 39 | gemma_root: str, |
| 40 | loras: tuple[LoraPathStrengthAndSDOps, ...] = (), |
| 41 | device: torch.device = device, |
| 42 | quantization: QuantizationPolicy | None = None, |
| 43 | action_config=None, |
| 44 | cache_config: CausalCacheConfig = CausalCacheConfig(), |
| 45 | ) -> None: |
| 46 | self.dtype = torch.bfloat16 |
| 47 | self.device = device |
| 48 | self.action_config = action_config |
| 49 | cache_config.validate() |
| 50 | self.cache_config = cache_config |
| 51 | self.model_ledger = ModelLedger( |
| 52 | dtype=self.dtype, |
| 53 | device=device, |
| 54 | checkpoint_path=checkpoint_path, |
| 55 | gemma_root_path=gemma_root, |
| 56 | loras=loras, |
| 57 | quantization=quantization, |
| 58 | ) |
| 59 | self.pipeline_components = PipelineComponents(dtype=self.dtype, device=device) |
| 60 | |
| 61 | @torch.inference_mode() |
| 62 | def __call__( # noqa: PLR0913 |
| 63 | self, |
| 64 | *, |
| 65 | prompt: str, |
| 66 | seed: int, |
| 67 | height: int, |
| 68 | width: int, |
| 69 | num_frames: int, |
| 70 | frame_rate: float, |
| 71 | images: list[ImageConditioningInput], |
| 72 | action_cond: dict[str, torch.Tensor], |
| 73 | timesteps: tuple[int, ...] | list[int] = DEFAULT_CAUSAL_TIMESTEPS, |
| 74 | video_tiling_config: TilingConfig | None = None, |
| 75 | ) -> tuple[Iterator[torch.Tensor], Audio]: |
| 76 | assert_resolution(height=height, width=width, is_two_stage=False) |
| 77 | latent_frames = (num_frames - 1) // 8 + 1 |
| 78 | if num_frames != (latent_frames - 1) * 8 + 1: |
| 79 | raise ValueError("causal --num-frames must be 1 + 8*n output frames") |
| 80 | causal_video_blocks(latent_frames, self.cache_config.video_chunk_size) |
| 81 | # The causal student is trained on one positive conditioning branch. |
| 82 | encoded_prompt, = encode_prompts([prompt], self.model_ledger) |
| 83 | if encoded_prompt.audio_encoding is None: |
| 84 | raise ValueError("the causal AV checkpoint must provide audio text embeddings") |
| 85 | |
| 86 | output_shape = VideoPixelShape(1, num_frames, height, width, frame_rate) |
| 87 | video_encoder = self.model_ledger.video_encoder() |
| 88 | conditionings = combined_image_conditionings( |
| 89 | images, height, width, video_encoder, self.dtype, self.device |
| 90 | ) |
| 91 | if self.device.type == "cuda": |
| 92 | torch.cuda.synchronize() |
| 93 | del video_encoder |
| 94 | cleanup_memory() |
| 95 | |
| 96 | generator = torch.Generator(device=self.device).manual_seed(seed) |
| 97 | noiser = GaussianNoiser(generator) |
| 98 | video_state, video_tools = noise_video_state( |
| 99 | output_shape, noiser, conditionings, self.pipeline_components, |
| 100 | self.dtype, self.device, |
| 101 | ) |
| 102 | audio_frames = causal_audio_frames(latent_frames, self.cache_config.video_chunk_size) |
| 103 | audio_shape = AudioLatentShape(batch=1, channels=8, frames=audio_frames, mel_bins=16) |
| 104 | audio_tools = AudioLatentTools(self.pipeline_components.audio_patchifier, audio_shape) |
| 105 | audio_state = create_noised_state( |
| 106 | audio_tools, [], noiser, self.dtype, self.device |
| 107 | ) |
| 108 | |
| 109 | x0_model = self.model_ledger.transformer(action_config=self.action_config) |
| 110 | wrapper = CausalModelWrapper( |
| 111 | x0_model.velocity_model, |
| 112 | patches_per_frame=(height // 32) * (width // 32), |
| 113 | cache=self.cache_config, |
| 114 | ) |
| 115 | generated_video, generated_audio = causal_rollout( |
| 116 | wrapper=wrapper, |
| 117 | clean_video=video_state.clean_latent, |
| 118 | clean_audio=audio_state.clean_latent, |
| 119 | video_positions=video_state.positions, |
| 120 | audio_positions=audio_state.positions, |
| 121 | video_context=encoded_prompt.video_encoding, |
| 122 | audio_context=encoded_prompt.audio_encoding, |
| 123 | context_mask=encoded_prompt.attention_mask, |
| 124 | action_cond=action_cond, |
| 125 | seed=seed, |
| 126 | timesteps=timesteps, |
| 127 | ) |
| 128 | del wrapper, x0_model |
| 129 | cleanup_memory() |
| 130 | |
| 131 | video_state = video_tools.unpatchify(video_tools.clear_conditioning( |
| 132 | video_state.__class__( |
| 133 | latent=generated_video, |
| 134 | denoise_mask=video_state.denoise_mask, |
| 135 | positions=video_state.positions, |
| 136 | clean_latent=video_state.clean_latent, |
| 137 | attention_mask=None, |
| 138 | ) |
| 139 | )) |
| 140 | audio_state = audio_tools.unpatchify(audio_tools.clear_conditioning( |
| 141 | audio_state.__class__( |
| 142 | latent=generated_audio, |
| 143 | denoise_mask=audio_state.denoise_mask, |
| 144 | positions=audio_state.positions, |
| 145 | clean_latent=audio_state.clean_latent, |
| 146 | attention_mask=None, |
| 147 | ) |
| 148 | )) |
| 149 | decoded_video = vae_decode_video( |
| 150 | video_state.latent, |
| 151 | self.model_ledger.video_decoder(), |
| 152 | tiling_config=video_tiling_config, |
| 153 | generator=generator, |
| 154 | ) |
| 155 | decoded_audio = vae_decode_audio( |
| 156 | audio_state.latent, |
| 157 | self.model_ledger.audio_decoder(), |
| 158 | self.model_ledger.vocoder(), |
| 159 | ) |
| 160 | return decoded_video, decoded_audio |
| 161 |