| 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.guiders import ( |
| 8 | MultiModalGuiderFactory, |
| 9 | MultiModalGuiderParams, |
| 10 | create_multimodal_guider_factory, |
| 11 | ) |
| 12 | from ltx_core.components.noisers import GaussianNoiser |
| 13 | from ltx_core.components.protocols import DiffusionStepProtocol |
| 14 | from ltx_core.components.schedulers import LTX2Scheduler |
| 15 | from ltx_core.loader import LoraPathStrengthAndSDOps |
| 16 | from ltx_core.model.audio_vae import decode_audio as vae_decode_audio |
| 17 | from ltx_core.model.upsampler import upsample_video |
| 18 | from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number |
| 19 | from ltx_core.model.video_vae import decode_video as vae_decode_video |
| 20 | from ltx_core.quantization import QuantizationPolicy |
| 21 | from ltx_core.types import Audio, LatentState, VideoPixelShape |
| 22 | from ltx_pipelines.utils import ModelLedger |
| 23 | from ltx_pipelines.utils.args import ImageConditioningInput, default_2_stage_arg_parser, detect_checkpoint_path |
| 24 | from ltx_pipelines.utils.constants import STAGE_2_DISTILLED_SIGMA_VALUES, detect_params |
| 25 | from ltx_pipelines.utils.helpers import ( |
| 26 | assert_resolution, |
| 27 | cleanup_memory, |
| 28 | denoise_audio_video, |
| 29 | encode_prompts, |
| 30 | get_device, |
| 31 | image_conditionings_by_adding_guiding_latent, |
| 32 | multi_modal_guider_factory_denoising_func, |
| 33 | simple_denoising_func, |
| 34 | ) |
| 35 | from ltx_pipelines.utils.media_io import encode_video |
| 36 | from ltx_pipelines.utils.samplers import euler_denoising_loop |
| 37 | from ltx_pipelines.utils.types import PipelineComponents |
| 38 | |
| 39 | device = get_device() |
| 40 | |
| 41 | |
| 42 | class KeyframeInterpolationPipeline: |
| 43 | """ |
| 44 | Keyframe-based Two-stage video interpolation pipeline. |
| 45 | Interpolates between keyframes to generate a video with smoother transitions. |
| 46 | Stage 1 generates video at half of the target resolution, then Stage 2 upsamples |
| 47 | by 2x and refines with additional denoising steps for higher quality output. |
| 48 | Stage 1 uses full model while Stage 2 uses distilled LORA for efficiency, |
| 49 | as the upsampled video already has good quality and just needs refinement. |
| 50 | """ |
| 51 | |
| 52 | def __init__( |
| 53 | self, |
| 54 | checkpoint_path: str, |
| 55 | distilled_lora: list[LoraPathStrengthAndSDOps], |
| 56 | spatial_upsampler_path: str, |
| 57 | gemma_root: str, |
| 58 | loras: list[LoraPathStrengthAndSDOps], |
| 59 | device: torch.device = device, |
| 60 | quantization: QuantizationPolicy | None = None, |
| 61 | ): |
| 62 | self.device = device |
| 63 | self.dtype = torch.bfloat16 |
| 64 | self.stage_1_model_ledger = ModelLedger( |
| 65 | dtype=self.dtype, |
| 66 | device=device, |
| 67 | checkpoint_path=checkpoint_path, |
| 68 | spatial_upsampler_path=spatial_upsampler_path, |
| 69 | gemma_root_path=gemma_root, |
| 70 | loras=loras, |
| 71 | quantization=quantization, |
| 72 | ) |
| 73 | self.stage_2_model_ledger = self.stage_1_model_ledger.with_additional_loras( |
| 74 | loras=distilled_lora, |
| 75 | ) |
| 76 | self.pipeline_components = PipelineComponents( |
| 77 | dtype=self.dtype, |
| 78 | device=device, |
| 79 | ) |
| 80 | |
| 81 | def __call__( # noqa: PLR0913 |
| 82 | self, |
| 83 | prompt: str, |
| 84 | negative_prompt: str, |
| 85 | seed: int, |
| 86 | height: int, |
| 87 | width: int, |
| 88 | num_frames: int, |
| 89 | frame_rate: float, |
| 90 | num_inference_steps: int, |
| 91 | video_guider_params: MultiModalGuiderParams | MultiModalGuiderFactory, |
| 92 | audio_guider_params: MultiModalGuiderParams | MultiModalGuiderFactory, |
| 93 | images: list[ImageConditioningInput], |
| 94 | tiling_config: TilingConfig | None = None, |
| 95 | enhance_prompt: bool = False, |
| 96 | ) -> tuple[Iterator[torch.Tensor], Audio]: |
| 97 | assert_resolution(height=height, width=width, is_two_stage=True) |
| 98 | |
| 99 | generator = torch.Generator(device=self.device).manual_seed(seed) |
| 100 | noiser = GaussianNoiser(generator=generator) |
| 101 | stepper = EulerDiffusionStep() |
| 102 | dtype = torch.bfloat16 |
| 103 | |
| 104 | ctx_p, ctx_n = encode_prompts( |
| 105 | [prompt, negative_prompt], |
| 106 | self.stage_1_model_ledger, |
| 107 | enhance_first_prompt=enhance_prompt, |
| 108 | enhance_prompt_image=images[0][0] if len(images) > 0 else None, |
| 109 | enhance_prompt_seed=seed, |
| 110 | ) |
| 111 | v_context_p, a_context_p = ctx_p.video_encoding, ctx_p.audio_encoding |
| 112 | v_context_n, a_context_n = ctx_n.video_encoding, ctx_n.audio_encoding |
| 113 | |
| 114 | # Stage 1: Initial low resolution video generation. |
| 115 | video_encoder = self.stage_1_model_ledger.video_encoder() |
| 116 | transformer = self.stage_1_model_ledger.transformer() |
| 117 | sigmas = LTX2Scheduler().execute(steps=num_inference_steps).to(dtype=torch.float32, device=self.device) |
| 118 | |
| 119 | def first_stage_denoising_loop( |
| 120 | sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol |
| 121 | ) -> tuple[LatentState, LatentState]: |
| 122 | return euler_denoising_loop( |
| 123 | sigmas=sigmas, |
| 124 | video_state=video_state, |
| 125 | audio_state=audio_state, |
| 126 | stepper=stepper, |
| 127 | denoise_fn=multi_modal_guider_factory_denoising_func( |
| 128 | video_guider_factory=create_multimodal_guider_factory( |
| 129 | params=video_guider_params, |
| 130 | negative_context=v_context_n, |
| 131 | ), |
| 132 | audio_guider_factory=create_multimodal_guider_factory( |
| 133 | params=audio_guider_params, |
| 134 | negative_context=a_context_n, |
| 135 | ), |
| 136 | v_context=v_context_p, |
| 137 | a_context=a_context_p, |
| 138 | transformer=transformer, # noqa: F821 |
| 139 | ), |
| 140 | ) |
| 141 | |
| 142 | stage_1_output_shape = VideoPixelShape( |
| 143 | batch=1, |
| 144 | frames=num_frames, |
| 145 | width=width // 2, |
| 146 | height=height // 2, |
| 147 | fps=frame_rate, |
| 148 | ) |
| 149 | stage_1_conditionings = image_conditionings_by_adding_guiding_latent( |
| 150 | images=images, |
| 151 | height=stage_1_output_shape.height, |
| 152 | width=stage_1_output_shape.width, |
| 153 | video_encoder=video_encoder, |
| 154 | dtype=dtype, |
| 155 | device=self.device, |
| 156 | ) |
| 157 | video_state, audio_state = denoise_audio_video( |
| 158 | output_shape=stage_1_output_shape, |
| 159 | conditionings=stage_1_conditionings, |
| 160 | noiser=noiser, |
| 161 | sigmas=sigmas, |
| 162 | stepper=stepper, |
| 163 | denoising_loop_fn=first_stage_denoising_loop, |
| 164 | components=self.pipeline_components, |
| 165 | dtype=dtype, |
| 166 | device=self.device, |
| 167 | ) |
| 168 | |
| 169 | torch.cuda.synchronize() |
| 170 | del transformer |
| 171 | cleanup_memory() |
| 172 | |
| 173 | # Stage 2: Upsample and refine the video at higher resolution with distilled LORA. |
| 174 | upscaled_video_latent = upsample_video( |
| 175 | latent=video_state.latent[:1], |
| 176 | video_encoder=video_encoder, |
| 177 | upsampler=self.stage_2_model_ledger.spatial_upsampler(), |
| 178 | ) |
| 179 | |
| 180 | torch.cuda.synchronize() |
| 181 | cleanup_memory() |
| 182 | |
| 183 | transformer = self.stage_2_model_ledger.transformer() |
| 184 | distilled_sigmas = torch.Tensor(STAGE_2_DISTILLED_SIGMA_VALUES).to(self.device) |
| 185 | |
| 186 | def second_stage_denoising_loop( |
| 187 | sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol |
| 188 | ) -> tuple[LatentState, LatentState]: |
| 189 | return euler_denoising_loop( |
| 190 | sigmas=sigmas, |
| 191 | video_state=video_state, |
| 192 | audio_state=audio_state, |
| 193 | stepper=stepper, |
| 194 | denoise_fn=simple_denoising_func( |
| 195 | video_context=v_context_p, |
| 196 | audio_context=a_context_p, |
| 197 | transformer=transformer, # noqa: F821 |
| 198 | ), |
| 199 | ) |
| 200 | |
| 201 | stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate) |
| 202 | stage_2_conditionings = image_conditionings_by_adding_guiding_latent( |
| 203 | images=images, |
| 204 | height=stage_2_output_shape.height, |
| 205 | width=stage_2_output_shape.width, |
| 206 | video_encoder=video_encoder, |
| 207 | dtype=dtype, |
| 208 | device=self.device, |
| 209 | ) |
| 210 | video_state, audio_state = denoise_audio_video( |
| 211 | output_shape=stage_2_output_shape, |
| 212 | conditionings=stage_2_conditionings, |
| 213 | noiser=noiser, |
| 214 | sigmas=distilled_sigmas, |
| 215 | stepper=stepper, |
| 216 | denoising_loop_fn=second_stage_denoising_loop, |
| 217 | components=self.pipeline_components, |
| 218 | dtype=dtype, |
| 219 | device=self.device, |
| 220 | noise_scale=distilled_sigmas[0], |
| 221 | initial_video_latent=upscaled_video_latent, |
| 222 | initial_audio_latent=audio_state.latent, |
| 223 | ) |
| 224 | |
| 225 | torch.cuda.synchronize() |
| 226 | del transformer |
| 227 | del video_encoder |
| 228 | cleanup_memory() |
| 229 | |
| 230 | decoded_video = vae_decode_video( |
| 231 | video_state.latent, self.stage_2_model_ledger.video_decoder(), tiling_config, generator |
| 232 | ) |
| 233 | decoded_audio = vae_decode_audio( |
| 234 | audio_state.latent, self.stage_2_model_ledger.audio_decoder(), self.stage_2_model_ledger.vocoder() |
| 235 | ) |
| 236 | return decoded_video, decoded_audio |
| 237 | |
| 238 | |
| 239 | @torch.inference_mode() |
| 240 | def main() -> None: |
| 241 | logging.getLogger().setLevel(logging.INFO) |
| 242 | checkpoint_path = detect_checkpoint_path() |
| 243 | params = detect_params(checkpoint_path) |
| 244 | parser = default_2_stage_arg_parser(params=params) |
| 245 | args = parser.parse_args() |
| 246 | pipeline = KeyframeInterpolationPipeline( |
| 247 | checkpoint_path=args.checkpoint_path, |
| 248 | distilled_lora=args.distilled_lora, |
| 249 | spatial_upsampler_path=args.spatial_upsampler_path, |
| 250 | gemma_root=args.gemma_root, |
| 251 | loras=tuple(args.lora) if args.lora else (), |
| 252 | quantization=args.quantization, |
| 253 | ) |
| 254 | tiling_config = TilingConfig.default() |
| 255 | video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config) |
| 256 | video, audio = pipeline( |
| 257 | prompt=args.prompt, |
| 258 | negative_prompt=args.negative_prompt, |
| 259 | seed=args.seed, |
| 260 | height=args.height, |
| 261 | width=args.width, |
| 262 | num_frames=args.num_frames, |
| 263 | frame_rate=args.frame_rate, |
| 264 | num_inference_steps=args.num_inference_steps, |
| 265 | video_guider_params=MultiModalGuiderParams( |
| 266 | cfg_scale=args.video_cfg_guidance_scale, |
| 267 | stg_scale=args.video_stg_guidance_scale, |
| 268 | rescale_scale=args.video_rescale_scale, |
| 269 | modality_scale=args.a2v_guidance_scale, |
| 270 | skip_step=args.video_skip_step, |
| 271 | stg_blocks=args.video_stg_blocks, |
| 272 | ), |
| 273 | audio_guider_params=MultiModalGuiderParams( |
| 274 | cfg_scale=args.audio_cfg_guidance_scale, |
| 275 | stg_scale=args.audio_stg_guidance_scale, |
| 276 | rescale_scale=args.audio_rescale_scale, |
| 277 | modality_scale=args.v2a_guidance_scale, |
| 278 | skip_step=args.audio_skip_step, |
| 279 | stg_blocks=args.audio_stg_blocks, |
| 280 | ), |
| 281 | images=args.images, |
| 282 | tiling_config=tiling_config, |
| 283 | ) |
| 284 | |
| 285 | encode_video( |
| 286 | video=video, |
| 287 | fps=args.frame_rate, |
| 288 | audio=audio, |
| 289 | output_path=args.output_path, |
| 290 | video_chunks_number=video_chunks_number, |
| 291 | ) |
| 292 | |
| 293 | |
| 294 | if __name__ == "__main__": |
| 295 | main() |
| 296 |