| 1 | from __future__ import annotations |
| 2 | |
| 3 | import argparse |
| 4 | import logging |
| 5 | from collections.abc import Iterator |
| 6 | from dataclasses import dataclass |
| 7 | |
| 8 | import torch |
| 9 | |
| 10 | from ltx_core.components.diffusion_steps import EulerDiffusionStep |
| 11 | from ltx_core.components.guiders import MultiModalGuider, MultiModalGuiderParams |
| 12 | from ltx_core.components.noisers import GaussianNoiser |
| 13 | from ltx_core.components.patchifiers import get_pixel_coords |
| 14 | from ltx_core.components.protocols import DiffusionStepProtocol |
| 15 | from ltx_core.components.schedulers import LTX2Scheduler |
| 16 | from ltx_core.conditioning import ConditioningItem |
| 17 | from ltx_core.loader import LoraPathStrengthAndSDOps |
| 18 | from ltx_core.model.audio_vae import decode_audio as vae_decode_audio |
| 19 | from ltx_core.model.audio_vae import encode_audio as vae_encode_audio |
| 20 | from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number |
| 21 | from ltx_core.model.video_vae import decode_video as vae_decode_video |
| 22 | from ltx_core.quantization import QuantizationPolicy |
| 23 | from ltx_core.tools import LatentTools |
| 24 | from ltx_core.types import ( |
| 25 | Audio, |
| 26 | AudioLatentShape, |
| 27 | LatentState, |
| 28 | SpatioTemporalScaleFactors, |
| 29 | VideoPixelShape, |
| 30 | ) |
| 31 | from ltx_pipelines.utils import ModelLedger |
| 32 | from ltx_pipelines.utils.args import QuantizationAction |
| 33 | from ltx_pipelines.utils.constants import DISTILLED_SIGMA_VALUES, detect_params |
| 34 | from ltx_pipelines.utils.helpers import ( |
| 35 | cleanup_memory, |
| 36 | encode_prompts, |
| 37 | get_device, |
| 38 | multi_modal_guider_denoising_func, |
| 39 | noise_audio_state, |
| 40 | noise_video_state, |
| 41 | simple_denoising_func, |
| 42 | ) |
| 43 | from ltx_pipelines.utils.media_io import ( |
| 44 | decode_audio_from_file, |
| 45 | encode_video, |
| 46 | get_videostream_metadata, |
| 47 | load_video_conditioning, |
| 48 | ) |
| 49 | from ltx_pipelines.utils.samplers import euler_denoising_loop |
| 50 | from ltx_pipelines.utils.types import PipelineComponents |
| 51 | |
| 52 | device = get_device() |
| 53 | |
| 54 | |
| 55 | def _encode_video_for_retake( |
| 56 | video_encoder: torch.nn.Module, |
| 57 | video_path: str, |
| 58 | output_shape: VideoPixelShape, |
| 59 | dtype: torch.dtype, |
| 60 | device: torch.device, |
| 61 | ) -> torch.Tensor: |
| 62 | """Load video and encode to latents.""" |
| 63 | pixel_video = load_video_conditioning( |
| 64 | video_path=video_path, |
| 65 | height=output_shape.height, |
| 66 | width=output_shape.width, |
| 67 | frame_cap=output_shape.frames, |
| 68 | dtype=dtype, |
| 69 | device=device, |
| 70 | ) # (1, C, F, H, W) |
| 71 | return video_encoder(pixel_video) |
| 72 | |
| 73 | |
| 74 | def _encode_audio_for_retake( |
| 75 | audio_encoder: torch.nn.Module, |
| 76 | waveform: torch.Tensor, |
| 77 | waveform_sr: int, |
| 78 | output_shape: VideoPixelShape, |
| 79 | dtype: torch.dtype, |
| 80 | ) -> torch.Tensor: |
| 81 | """Encode audio to latents and trim/pad to match output_shape.""" |
| 82 | waveform_batch = waveform.unsqueeze(0) if waveform.dim() == 2 else waveform |
| 83 | initial_audio_latent = vae_encode_audio( |
| 84 | Audio(waveform=waveform_batch.to(dtype), sampling_rate=waveform_sr), audio_encoder, None |
| 85 | ) |
| 86 | expected_audio_shape = AudioLatentShape.from_video_pixel_shape(output_shape) |
| 87 | expected_frames = expected_audio_shape.frames |
| 88 | actual_frames = initial_audio_latent.shape[2] |
| 89 | if actual_frames > expected_frames: |
| 90 | initial_audio_latent = initial_audio_latent[:, :, :expected_frames, :] |
| 91 | elif actual_frames < expected_frames: |
| 92 | pad = torch.zeros( |
| 93 | initial_audio_latent.shape[0], |
| 94 | initial_audio_latent.shape[1], |
| 95 | expected_frames - actual_frames, |
| 96 | initial_audio_latent.shape[3], |
| 97 | device=initial_audio_latent.device, |
| 98 | dtype=initial_audio_latent.dtype, |
| 99 | ) |
| 100 | initial_audio_latent = torch.cat([initial_audio_latent, pad], dim=2) |
| 101 | return initial_audio_latent |
| 102 | |
| 103 | |
| 104 | # --------------------------------------------------------------------------- |
| 105 | # Custom conditioning item: temporal region mask |
| 106 | # --------------------------------------------------------------------------- |
| 107 | |
| 108 | |
| 109 | @dataclass(frozen=True) |
| 110 | class TemporalRegionMask: |
| 111 | """Conditioning item that sets ``denoise_mask = 0`` outside a time range |
| 112 | and ``1`` inside, so only the specified temporal region is regenerated. |
| 113 | Uses ``start_time`` and ``end_time`` in seconds. Works in *patchified* |
| 114 | (token) space using the patchifier's ``get_patch_grid_bounds``: for video |
| 115 | coords are latent frame indices (converted from seconds via ``fps``), for |
| 116 | audio coords are already in seconds. |
| 117 | """ |
| 118 | |
| 119 | start_time: float # seconds, inclusive |
| 120 | end_time: float # seconds, exclusive |
| 121 | fps: float |
| 122 | |
| 123 | def apply_to(self, latent_state: LatentState, latent_tools: LatentTools) -> LatentState: |
| 124 | coords = latent_tools.patchifier.get_patch_grid_bounds( |
| 125 | latent_tools.target_shape, device=latent_state.denoise_mask.device |
| 126 | ) |
| 127 | # coords: [B, 3, N, 2] (video) or [B, 1, N, 2] (audio); temporal dim is index 0 |
| 128 | if coords.shape[1] == 1: |
| 129 | # Audio: patchifier returns seconds |
| 130 | t_start = coords[:, 0, :, 0] # [B, N] |
| 131 | t_end = coords[:, 0, :, 1] # [B, N] |
| 132 | in_region = (t_end > self.start_time) & (t_start < self.end_time) |
| 133 | else: |
| 134 | # Video: get pixel bounds per patch, find patches for start/end frame, read latent from coords. |
| 135 | scale_factors = getattr(latent_tools, "scale_factors", SpatioTemporalScaleFactors.default()) |
| 136 | pixel_bounds = get_pixel_coords(coords, scale_factors, causal_fix=getattr(latent_tools, "causal_fix", True)) |
| 137 | timestamp_bounds = pixel_bounds[0, 0] / self.fps |
| 138 | t_start, t_end = timestamp_bounds.unbind(dim=-1) |
| 139 | in_region = (t_end > self.start_time) & (t_start < self.end_time) |
| 140 | state = latent_state.clone() |
| 141 | mask_val = in_region.to(state.denoise_mask.dtype) |
| 142 | if state.denoise_mask.dim() == 3: |
| 143 | mask_val = mask_val.unsqueeze(-1) |
| 144 | state.denoise_mask.copy_(mask_val) |
| 145 | return state |
| 146 | |
| 147 | |
| 148 | # --------------------------------------------------------------------------- |
| 149 | # Pipeline |
| 150 | # --------------------------------------------------------------------------- |
| 151 | |
| 152 | |
| 153 | class RetakePipeline: |
| 154 | """Regenerate a time region (retake) of an existing video. |
| 155 | Given a source video file and a time window ``[start_time, end_time]`` |
| 156 | (in seconds), this pipeline keeps the video/audio outside that window |
| 157 | unchanged and *regenerates* the content inside the window from a text |
| 158 | prompt using the LTX-2 diffusion model. |
| 159 | Parameters |
| 160 | ---------- |
| 161 | checkpoint_path : str |
| 162 | Path to the LTX-2 model checkpoint. |
| 163 | gemma_root : str |
| 164 | Root directory containing Gemma text-encoder weights. |
| 165 | loras : list[LoraPathStrengthAndSDOps] |
| 166 | Optional LoRA configs applied to the transformer. |
| 167 | device : torch.device |
| 168 | Target device (default: CUDA if available). |
| 169 | quantization : QuantizationPolicy | None |
| 170 | Optional quantization policy for the transformer. |
| 171 | """ |
| 172 | |
| 173 | def __init__( |
| 174 | self, |
| 175 | checkpoint_path: str, |
| 176 | gemma_root: str, |
| 177 | loras: list[LoraPathStrengthAndSDOps], |
| 178 | device: torch.device = device, |
| 179 | quantization: QuantizationPolicy | None = None, |
| 180 | ): |
| 181 | self.device = device |
| 182 | self.dtype = torch.bfloat16 |
| 183 | self.model_ledger = ModelLedger( |
| 184 | dtype=self.dtype, |
| 185 | device=device, |
| 186 | checkpoint_path=checkpoint_path, |
| 187 | gemma_root_path=gemma_root, |
| 188 | loras=loras, |
| 189 | quantization=quantization, |
| 190 | ) |
| 191 | self.pipeline_components = PipelineComponents( |
| 192 | dtype=self.dtype, |
| 193 | device=device, |
| 194 | ) |
| 195 | |
| 196 | # --------------------------------------------------------------------- # |
| 197 | # Public entry point # |
| 198 | # --------------------------------------------------------------------- # |
| 199 | |
| 200 | def __call__( # noqa: PLR0913, PLR0915 |
| 201 | self, |
| 202 | video_path: str, |
| 203 | prompt: str, |
| 204 | start_time: float, |
| 205 | end_time: float, |
| 206 | seed: int, |
| 207 | *, |
| 208 | negative_prompt: str = "", |
| 209 | num_inference_steps: int = 40, |
| 210 | video_guider_params: MultiModalGuiderParams | None = None, |
| 211 | audio_guider_params: MultiModalGuiderParams | None = None, |
| 212 | regenerate_video: bool = True, |
| 213 | regenerate_audio: bool = True, |
| 214 | enhance_prompt: bool = False, |
| 215 | distilled: bool = False, |
| 216 | tiling_config: TilingConfig | None = None, |
| 217 | ) -> tuple[Iterator[torch.Tensor], torch.Tensor]: |
| 218 | """Regenerate ``[start_time, end_time]`` of the source video (retake). |
| 219 | Parameters |
| 220 | ---------- |
| 221 | video_path : str |
| 222 | Path to the source video file (must contain video; audio is optional). |
| 223 | prompt : str |
| 224 | Text prompt describing the *regenerated* section. |
| 225 | start_time, end_time : float |
| 226 | Time window (in seconds) of the section to regenerate. |
| 227 | seed : int |
| 228 | Random seed for reproducibility. |
| 229 | negative_prompt : str |
| 230 | Negative prompt for CFG guidance (ignored in distilled mode). |
| 231 | num_inference_steps : int |
| 232 | Number of Euler denoising steps (ignored in distilled mode which |
| 233 | uses a fixed 8-step schedule). |
| 234 | video_guider_params, audio_guider_params : MultiModalGuiderParams | None |
| 235 | Guidance parameters for video and audio modalities. Ignored in |
| 236 | distilled mode. |
| 237 | regenerate_video : bool |
| 238 | If ``True`` (default), preserve video outside ``[start_time, end_time]`` |
| 239 | and only regenerate the masked region. If ``False``, fully regenerate |
| 240 | all video frames (the encoded video is still used as the initial latent |
| 241 | but with ``denoise_mask = 1`` everywhere). |
| 242 | regenerate_audio : bool |
| 243 | If True, regenerate audio in the [start_time, end_time] window; if False, |
| 244 | audio is preserved as-is (no regeneration). |
| 245 | enhance_prompt : bool |
| 246 | Whether to enhance the prompt via the text encoder. |
| 247 | distilled : bool |
| 248 | If ``True``, use the distilled sigma schedule |
| 249 | (``DISTILLED_SIGMA_VALUES``) and a simple (non-guided) denoising |
| 250 | function. The model checkpoint must be the distilled variant. |
| 251 | Returns |
| 252 | ------- |
| 253 | tuple[Iterator[torch.Tensor], torch.Tensor] |
| 254 | ``(video_frames_iterator, audio_waveform)`` |
| 255 | """ |
| 256 | if start_time >= end_time: |
| 257 | raise ValueError(f"start_time ({start_time}) must be less than end_time ({end_time})") |
| 258 | |
| 259 | effective_seed = torch.randint(0, 2**31, (1,), device=self.device).item() if seed < 0 else seed |
| 260 | generator = torch.Generator(device=self.device).manual_seed(effective_seed) |
| 261 | noiser = GaussianNoiser(generator=generator) |
| 262 | stepper = EulerDiffusionStep() |
| 263 | dtype = self.dtype |
| 264 | |
| 265 | video_encoder = self.model_ledger.video_encoder() |
| 266 | |
| 267 | # Use av to get metadata |
| 268 | fps, num_pixel_frames, src_width, src_height = get_videostream_metadata(video_path) |
| 269 | |
| 270 | output_shape = VideoPixelShape( |
| 271 | batch=1, |
| 272 | frames=num_pixel_frames, |
| 273 | width=src_width, |
| 274 | height=src_height, |
| 275 | fps=fps, |
| 276 | ) |
| 277 | initial_video_latent = _encode_video_for_retake( |
| 278 | video_encoder=video_encoder, |
| 279 | video_path=video_path, |
| 280 | output_shape=output_shape, |
| 281 | dtype=dtype, |
| 282 | device=self.device, |
| 283 | ) |
| 284 | video_conditionings: list[ConditioningItem] = [ |
| 285 | TemporalRegionMask( |
| 286 | start_time=start_time if regenerate_video else 0.0, |
| 287 | end_time=end_time if regenerate_video else 0.0, |
| 288 | fps=fps, |
| 289 | ) |
| 290 | ] |
| 291 | del video_encoder |
| 292 | cleanup_memory() |
| 293 | |
| 294 | initial_audio_latent: torch.Tensor | None = None |
| 295 | audio_conditionings: list[ConditioningItem] = [] |
| 296 | |
| 297 | audio_in = decode_audio_from_file(video_path, self.device) |
| 298 | audio_encoder = self.model_ledger.audio_encoder() |
| 299 | |
| 300 | if audio_in is not None: |
| 301 | waveform = audio_in.waveform.squeeze(0) |
| 302 | waveform_sr = audio_in.sampling_rate |
| 303 | else: |
| 304 | waveform, waveform_sr = None, None |
| 305 | if waveform is not None: |
| 306 | initial_audio_latent = _encode_audio_for_retake( |
| 307 | audio_encoder=audio_encoder, |
| 308 | waveform=waveform, |
| 309 | waveform_sr=waveform_sr, |
| 310 | output_shape=output_shape, |
| 311 | dtype=dtype, |
| 312 | ) |
| 313 | audio_conditionings = [ |
| 314 | TemporalRegionMask( |
| 315 | start_time=start_time if regenerate_audio else 0.0, |
| 316 | end_time=end_time if regenerate_audio else 0.0, |
| 317 | fps=fps, |
| 318 | ) |
| 319 | ] |
| 320 | |
| 321 | del audio_encoder |
| 322 | cleanup_memory() |
| 323 | |
| 324 | prompts_to_encode = [prompt] if distilled else [prompt, negative_prompt] |
| 325 | contexts = encode_prompts( |
| 326 | prompts_to_encode, |
| 327 | self.model_ledger, |
| 328 | enhance_first_prompt=enhance_prompt, |
| 329 | enhance_prompt_seed=effective_seed, |
| 330 | ) |
| 331 | |
| 332 | v_context_p, a_context_p = contexts[0].video_encoding, contexts[0].audio_encoding |
| 333 | if not distilled: |
| 334 | v_context_n, a_context_n = contexts[1].video_encoding, contexts[1].audio_encoding |
| 335 | |
| 336 | transformer = self.model_ledger.transformer() |
| 337 | |
| 338 | sigmas = ( |
| 339 | torch.tensor(DISTILLED_SIGMA_VALUES) if distilled else LTX2Scheduler().execute(steps=num_inference_steps) |
| 340 | ).to(dtype=torch.float32, device=self.device) |
| 341 | if distilled: |
| 342 | denoise_fn = simple_denoising_func( |
| 343 | video_context=v_context_p, |
| 344 | audio_context=a_context_p, |
| 345 | transformer=transformer, |
| 346 | ) |
| 347 | else: |
| 348 | video_guider = MultiModalGuider( |
| 349 | params=video_guider_params, |
| 350 | negative_context=v_context_n, |
| 351 | ) |
| 352 | audio_guider = MultiModalGuider( |
| 353 | params=audio_guider_params, |
| 354 | negative_context=a_context_n, |
| 355 | ) |
| 356 | denoise_fn = multi_modal_guider_denoising_func( |
| 357 | video_guider, |
| 358 | audio_guider, |
| 359 | v_context=v_context_p, |
| 360 | a_context=a_context_p, |
| 361 | transformer=transformer, |
| 362 | ) |
| 363 | |
| 364 | def denoising_loop( |
| 365 | sigmas: torch.Tensor, |
| 366 | video_state: LatentState, |
| 367 | audio_state: LatentState, |
| 368 | stepper: DiffusionStepProtocol, |
| 369 | ) -> tuple[LatentState, LatentState]: |
| 370 | return euler_denoising_loop( |
| 371 | sigmas=sigmas, |
| 372 | video_state=video_state, |
| 373 | audio_state=audio_state, |
| 374 | stepper=stepper, |
| 375 | denoise_fn=denoise_fn, |
| 376 | ) |
| 377 | |
| 378 | # Build noised states with the encoded latents as initial values and |
| 379 | # the temporal masks applied via conditionings. |
| 380 | video_state, video_tools = noise_video_state( |
| 381 | output_shape=output_shape, |
| 382 | noiser=noiser, |
| 383 | conditionings=video_conditionings, |
| 384 | components=self.pipeline_components, |
| 385 | dtype=dtype, |
| 386 | device=self.device, |
| 387 | initial_latent=initial_video_latent, |
| 388 | ) |
| 389 | audio_state, audio_tools = noise_audio_state( |
| 390 | output_shape=output_shape, |
| 391 | noiser=noiser, |
| 392 | conditionings=audio_conditionings, |
| 393 | components=self.pipeline_components, |
| 394 | dtype=dtype, |
| 395 | device=self.device, |
| 396 | initial_latent=initial_audio_latent, |
| 397 | ) |
| 398 | |
| 399 | video_state, audio_state = denoising_loop(sigmas, video_state, audio_state, stepper) |
| 400 | |
| 401 | video_state = video_tools.clear_conditioning(video_state) |
| 402 | video_state = video_tools.unpatchify(video_state) |
| 403 | audio_state = audio_tools.clear_conditioning(audio_state) |
| 404 | audio_state = audio_tools.unpatchify(audio_state) |
| 405 | |
| 406 | torch.cuda.synchronize() |
| 407 | del transformer |
| 408 | cleanup_memory() |
| 409 | |
| 410 | decoded_video = vae_decode_video( |
| 411 | video_state.latent, self.model_ledger.video_decoder(), tiling_config, generator |
| 412 | ) |
| 413 | decoded_audio = vae_decode_audio( |
| 414 | audio_state.latent, self.model_ledger.audio_decoder(), self.model_ledger.vocoder() |
| 415 | ) |
| 416 | |
| 417 | return decoded_video, decoded_audio |
| 418 | |
| 419 | |
| 420 | @torch.inference_mode() |
| 421 | def main() -> None: |
| 422 | """CLI entry point for retake (regenerate a time region).""" |
| 423 | logging.getLogger().setLevel(logging.INFO) |
| 424 | parser = argparse.ArgumentParser(description="Retake: regenerate a time region of a video with LTX-2.") |
| 425 | parser.add_argument("--video-path", type=str, required=True, help="Path to the source video.") |
| 426 | parser.add_argument("--prompt", type=str, required=True, help="Text prompt for the regenerated region.") |
| 427 | parser.add_argument("--start-time", type=float, required=True, help="Start time of the region to regenerate (s).") |
| 428 | parser.add_argument("--end-time", type=float, required=True, help="End time of the region to regenerate (s).") |
| 429 | parser.add_argument("--output-path", type=str, required=True, help="Path for the output video.") |
| 430 | parser.add_argument("--checkpoint-path", type=str, required=True, help="Path to the LTX-2 checkpoint.") |
| 431 | parser.add_argument("--gemma-root", type=str, required=True, help="Path to Gemma text encoder weights.") |
| 432 | parser.add_argument("--seed", type=int, default=42, help="Random seed. Use -1 for a random seed.") |
| 433 | parser.add_argument("--loras", nargs="*", default=[], help="LoRA paths (optional).") |
| 434 | parser.add_argument( |
| 435 | "--quantization", |
| 436 | dest="quantization", |
| 437 | action=QuantizationAction, |
| 438 | nargs="+", |
| 439 | metavar=("POLICY", "AMAX_PATH"), |
| 440 | default=None, |
| 441 | help="Quantization policy: fp8-cast or fp8-scaled-mm [AMAX_PATH].", |
| 442 | ) |
| 443 | args = parser.parse_args() |
| 444 | |
| 445 | if args.start_time >= args.end_time: |
| 446 | raise ValueError("start_time must be less than end_time") |
| 447 | |
| 448 | # Validate frame count (8k+1) and resolution (multiples of 32) at CLI stage |
| 449 | video_scale = SpatioTemporalScaleFactors.default() |
| 450 | fps, num_frames, width, height = get_videostream_metadata(args.video_path) |
| 451 | if (num_frames - 1) % video_scale.time != 0: |
| 452 | snapped = ((num_frames - 1) // video_scale.time) * video_scale.time + 1 |
| 453 | raise ValueError( |
| 454 | f"Video frame count must satisfy 8k+1 (e.g. 97, 193). Got {num_frames}; use a video with {snapped} frames." |
| 455 | ) |
| 456 | if width % 32 != 0 or height % 32 != 0: |
| 457 | raise ValueError(f"Video width and height must be multiples of 32. Got {width}x{height}.") |
| 458 | |
| 459 | pipeline = RetakePipeline( |
| 460 | checkpoint_path=args.checkpoint_path, |
| 461 | gemma_root=args.gemma_root, |
| 462 | loras=tuple(args.loras) if args.loras else (), |
| 463 | quantization=args.quantization, |
| 464 | ) |
| 465 | params = detect_params(args.checkpoint_path) |
| 466 | tiling_config = TilingConfig.default() |
| 467 | video_iter, audio = pipeline( |
| 468 | video_path=args.video_path, |
| 469 | prompt=args.prompt, |
| 470 | start_time=args.start_time, |
| 471 | end_time=args.end_time, |
| 472 | seed=args.seed, |
| 473 | video_guider_params=params.video_guider_params, |
| 474 | audio_guider_params=params.audio_guider_params, |
| 475 | tiling_config=tiling_config, |
| 476 | ) |
| 477 | video_chunks_number = get_video_chunks_number(num_frames, tiling_config) |
| 478 | encode_video( |
| 479 | video=video_iter, |
| 480 | fps=int(fps), |
| 481 | audio=audio, |
| 482 | output_path=args.output_path, |
| 483 | video_chunks_number=video_chunks_number, |
| 484 | ) |
| 485 | |
| 486 | |
| 487 | if __name__ == "__main__": |
| 488 | main() |
| 489 |