| 1 | import logging |
| 2 | from collections.abc import Iterator |
| 3 | |
| 4 | import torch |
| 5 | |
| 6 | from ltx_core.components.diffusion_steps import EulerDiffusionStep |
| 7 | from ltx_core.components.noisers import GaussianNoiser |
| 8 | from ltx_core.components.protocols import DiffusionStepProtocol |
| 9 | from ltx_core.loader import LoraPathStrengthAndSDOps |
| 10 | from ltx_core.model.audio_vae import decode_audio as vae_decode_audio |
| 11 | from ltx_core.model.upsampler import upsample_video |
| 12 | from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number |
| 13 | from ltx_core.model.video_vae import decode_video as vae_decode_video |
| 14 | from ltx_core.quantization import QuantizationPolicy |
| 15 | from ltx_core.types import Audio, LatentState, VideoPixelShape |
| 16 | from ltx_pipelines.utils import ModelLedger, euler_denoising_loop |
| 17 | from ltx_pipelines.utils.args import ( |
| 18 | ImageConditioningInput, |
| 19 | default_2_stage_distilled_arg_parser, |
| 20 | detect_checkpoint_path, |
| 21 | ) |
| 22 | from ltx_pipelines.utils.constants import ( |
| 23 | DISTILLED_SIGMA_VALUES, |
| 24 | STAGE_2_DISTILLED_SIGMA_VALUES, |
| 25 | detect_params, |
| 26 | ) |
| 27 | from ltx_pipelines.utils.helpers import ( |
| 28 | assert_resolution, |
| 29 | cleanup_memory, |
| 30 | combined_image_conditionings, |
| 31 | denoise_audio_video, |
| 32 | encode_prompts, |
| 33 | get_device, |
| 34 | simple_denoising_func, |
| 35 | ) |
| 36 | from ltx_pipelines.utils.media_io import encode_video |
| 37 | from ltx_pipelines.utils.types import PipelineComponents |
| 38 | |
| 39 | device = get_device() |
| 40 | |
| 41 | |
| 42 | class DistilledPipeline: |
| 43 | """ |
| 44 | Two-stage distilled video generation pipeline. |
| 45 | Stage 1 generates video at half of the target resolution, then Stage 2 upsamples |
| 46 | by 2x and refines with additional denoising steps for higher quality output. |
| 47 | """ |
| 48 | |
| 49 | def __init__( |
| 50 | self, |
| 51 | distilled_checkpoint_path: str, |
| 52 | gemma_root: str, |
| 53 | spatial_upsampler_path: str, |
| 54 | loras: list[LoraPathStrengthAndSDOps], |
| 55 | device: torch.device = device, |
| 56 | quantization: QuantizationPolicy | None = None, |
| 57 | ): |
| 58 | self.device = device |
| 59 | self.dtype = torch.bfloat16 |
| 60 | |
| 61 | self.model_ledger = ModelLedger( |
| 62 | dtype=self.dtype, |
| 63 | device=device, |
| 64 | checkpoint_path=distilled_checkpoint_path, |
| 65 | spatial_upsampler_path=spatial_upsampler_path, |
| 66 | gemma_root_path=gemma_root, |
| 67 | loras=loras, |
| 68 | quantization=quantization, |
| 69 | ) |
| 70 | |
| 71 | self.pipeline_components = PipelineComponents( |
| 72 | dtype=self.dtype, |
| 73 | device=device, |
| 74 | ) |
| 75 | |
| 76 | def __call__( |
| 77 | self, |
| 78 | prompt: str, |
| 79 | seed: int, |
| 80 | height: int, |
| 81 | width: int, |
| 82 | num_frames: int, |
| 83 | frame_rate: float, |
| 84 | images: list[ImageConditioningInput], |
| 85 | tiling_config: TilingConfig | None = None, |
| 86 | enhance_prompt: bool = False, |
| 87 | ) -> tuple[Iterator[torch.Tensor], Audio]: |
| 88 | assert_resolution(height=height, width=width, is_two_stage=True) |
| 89 | |
| 90 | generator = torch.Generator(device=self.device).manual_seed(seed) |
| 91 | noiser = GaussianNoiser(generator=generator) |
| 92 | stepper = EulerDiffusionStep() |
| 93 | dtype = torch.bfloat16 |
| 94 | |
| 95 | (ctx_p,) = encode_prompts( |
| 96 | [prompt], |
| 97 | self.model_ledger, |
| 98 | enhance_first_prompt=enhance_prompt, |
| 99 | enhance_prompt_image=images[0][0] if len(images) > 0 else None, |
| 100 | ) |
| 101 | video_context, audio_context = ctx_p.video_encoding, ctx_p.audio_encoding |
| 102 | |
| 103 | # Stage 1: Initial low resolution video generation. |
| 104 | video_encoder = self.model_ledger.video_encoder() |
| 105 | transformer = self.model_ledger.transformer() |
| 106 | stage_1_sigmas = torch.Tensor(DISTILLED_SIGMA_VALUES).to(self.device) |
| 107 | |
| 108 | def denoising_loop( |
| 109 | sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol |
| 110 | ) -> tuple[LatentState, LatentState]: |
| 111 | return euler_denoising_loop( |
| 112 | sigmas=sigmas, |
| 113 | video_state=video_state, |
| 114 | audio_state=audio_state, |
| 115 | stepper=stepper, |
| 116 | denoise_fn=simple_denoising_func( |
| 117 | video_context=video_context, |
| 118 | audio_context=audio_context, |
| 119 | transformer=transformer, # noqa: F821 |
| 120 | ), |
| 121 | ) |
| 122 | |
| 123 | stage_1_output_shape = VideoPixelShape( |
| 124 | batch=1, |
| 125 | frames=num_frames, |
| 126 | width=width // 2, |
| 127 | height=height // 2, |
| 128 | fps=frame_rate, |
| 129 | ) |
| 130 | stage_1_conditionings = combined_image_conditionings( |
| 131 | images=images, |
| 132 | height=stage_1_output_shape.height, |
| 133 | width=stage_1_output_shape.width, |
| 134 | video_encoder=video_encoder, |
| 135 | dtype=dtype, |
| 136 | device=self.device, |
| 137 | ) |
| 138 | |
| 139 | video_state, audio_state = denoise_audio_video( |
| 140 | output_shape=stage_1_output_shape, |
| 141 | conditionings=stage_1_conditionings, |
| 142 | noiser=noiser, |
| 143 | sigmas=stage_1_sigmas, |
| 144 | stepper=stepper, |
| 145 | denoising_loop_fn=denoising_loop, |
| 146 | components=self.pipeline_components, |
| 147 | dtype=dtype, |
| 148 | device=self.device, |
| 149 | ) |
| 150 | |
| 151 | # Stage 2: Upsample and refine the video at higher resolution with distilled LORA. |
| 152 | upscaled_video_latent = upsample_video( |
| 153 | latent=video_state.latent[:1], video_encoder=video_encoder, upsampler=self.model_ledger.spatial_upsampler() |
| 154 | ) |
| 155 | |
| 156 | torch.cuda.synchronize() |
| 157 | cleanup_memory() |
| 158 | |
| 159 | stage_2_sigmas = torch.Tensor(STAGE_2_DISTILLED_SIGMA_VALUES).to(self.device) |
| 160 | stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate) |
| 161 | stage_2_conditionings = combined_image_conditionings( |
| 162 | images=images, |
| 163 | height=stage_2_output_shape.height, |
| 164 | width=stage_2_output_shape.width, |
| 165 | video_encoder=video_encoder, |
| 166 | dtype=dtype, |
| 167 | device=self.device, |
| 168 | ) |
| 169 | video_state, audio_state = denoise_audio_video( |
| 170 | output_shape=stage_2_output_shape, |
| 171 | conditionings=stage_2_conditionings, |
| 172 | noiser=noiser, |
| 173 | sigmas=stage_2_sigmas, |
| 174 | stepper=stepper, |
| 175 | denoising_loop_fn=denoising_loop, |
| 176 | components=self.pipeline_components, |
| 177 | dtype=dtype, |
| 178 | device=self.device, |
| 179 | noise_scale=stage_2_sigmas[0], |
| 180 | initial_video_latent=upscaled_video_latent, |
| 181 | initial_audio_latent=audio_state.latent, |
| 182 | ) |
| 183 | |
| 184 | torch.cuda.synchronize() |
| 185 | del transformer |
| 186 | del video_encoder |
| 187 | cleanup_memory() |
| 188 | |
| 189 | decoded_video = vae_decode_video( |
| 190 | video_state.latent, self.model_ledger.video_decoder(), tiling_config, generator |
| 191 | ) |
| 192 | decoded_audio = vae_decode_audio( |
| 193 | audio_state.latent, self.model_ledger.audio_decoder(), self.model_ledger.vocoder() |
| 194 | ) |
| 195 | return decoded_video, decoded_audio |
| 196 | |
| 197 | |
| 198 | @torch.inference_mode() |
| 199 | def main() -> None: |
| 200 | logging.getLogger().setLevel(logging.INFO) |
| 201 | checkpoint_path = detect_checkpoint_path(distilled=True) |
| 202 | params = detect_params(checkpoint_path) |
| 203 | parser = default_2_stage_distilled_arg_parser(params=params) |
| 204 | args = parser.parse_args() |
| 205 | pipeline = DistilledPipeline( |
| 206 | distilled_checkpoint_path=args.distilled_checkpoint_path, |
| 207 | spatial_upsampler_path=args.spatial_upsampler_path, |
| 208 | gemma_root=args.gemma_root, |
| 209 | loras=tuple(args.lora) if args.lora else (), |
| 210 | quantization=args.quantization, |
| 211 | ) |
| 212 | tiling_config = TilingConfig.default() |
| 213 | video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config) |
| 214 | video, audio = pipeline( |
| 215 | prompt=args.prompt, |
| 216 | seed=args.seed, |
| 217 | height=args.height, |
| 218 | width=args.width, |
| 219 | num_frames=args.num_frames, |
| 220 | frame_rate=args.frame_rate, |
| 221 | images=args.images, |
| 222 | tiling_config=tiling_config, |
| 223 | enhance_prompt=args.enhance_prompt, |
| 224 | ) |
| 225 | |
| 226 | encode_video( |
| 227 | video=video, |
| 228 | fps=args.frame_rate, |
| 229 | audio=audio, |
| 230 | output_path=args.output_path, |
| 231 | video_chunks_number=video_chunks_number, |
| 232 | ) |
| 233 | |
| 234 | |
| 235 | if __name__ == "__main__": |
| 236 | main() |
| 237 |