| 1 | import argparse |
| 2 | from pathlib import Path |
| 3 | from typing import NamedTuple |
| 4 | |
| 5 | from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps |
| 6 | from ltx_core.quantization import QuantizationPolicy |
| 7 | from ltx_pipelines.utils.constants import ( |
| 8 | DEFAULT_IMAGE_CRF, |
| 9 | DEFAULT_LORA_STRENGTH, |
| 10 | DEFAULT_NEGATIVE_PROMPT, |
| 11 | LTX_2_3_HQ_PARAMS, |
| 12 | LTX_2_3_PARAMS, |
| 13 | PipelineParams, |
| 14 | ) |
| 15 | |
| 16 | |
| 17 | class ImageConditioningInput(NamedTuple): |
| 18 | path: str |
| 19 | frame_idx: int |
| 20 | strength: float |
| 21 | crf: int = DEFAULT_IMAGE_CRF |
| 22 | |
| 23 | |
| 24 | class VideoConditioningAction(argparse.Action): |
| 25 | def __call__( |
| 26 | self, |
| 27 | parser: argparse.ArgumentParser, # noqa: ARG002 |
| 28 | namespace: argparse.Namespace, |
| 29 | values: list[str], |
| 30 | option_string: str | None = None, # noqa: ARG002 |
| 31 | ) -> None: |
| 32 | path, strength_str = values |
| 33 | resolved_path = resolve_path(path) |
| 34 | strength = float(strength_str) |
| 35 | current = getattr(namespace, self.dest) or [] |
| 36 | current.append((resolved_path, strength)) |
| 37 | setattr(namespace, self.dest, current) |
| 38 | |
| 39 | |
| 40 | class VideoMaskConditioningAction(argparse.Action): |
| 41 | """Parse ``--conditioning-attention-mask PATH STRENGTH``. |
| 42 | Stores a ``(mask_path, strength)`` tuple on the namespace. The mask video |
| 43 | should be grayscale with pixel values in [0, 1] controlling per-region |
| 44 | conditioning attention strength. The scalar *STRENGTH* is multiplied with |
| 45 | the spatial mask before it is applied. |
| 46 | """ |
| 47 | |
| 48 | def __call__( |
| 49 | self, |
| 50 | parser: argparse.ArgumentParser, # noqa: ARG002 |
| 51 | namespace: argparse.Namespace, |
| 52 | values: list[str], |
| 53 | option_string: str | None = None, |
| 54 | ) -> None: |
| 55 | if len(values) != 2: |
| 56 | msg = f"{option_string} requires exactly 2 arguments (MASK_PATH STRENGTH), got {len(values)}" |
| 57 | raise argparse.ArgumentError(self, msg) |
| 58 | |
| 59 | mask_path = resolve_path(values[0]) |
| 60 | strength = float(values[1]) |
| 61 | setattr(namespace, self.dest, (mask_path, strength)) |
| 62 | |
| 63 | |
| 64 | class ImageAction(argparse.Action): |
| 65 | def __call__( |
| 66 | self, |
| 67 | parser: argparse.ArgumentParser, # noqa: ARG002 |
| 68 | namespace: argparse.Namespace, |
| 69 | values: list[str], |
| 70 | option_string: str | None = None, |
| 71 | ) -> None: |
| 72 | if len(values) not in (3, 4): |
| 73 | msg = f"{option_string} requires 3 or 4 arguments (PATH FRAME_IDX STRENGTH [CRF]), got {len(values)}" |
| 74 | raise argparse.ArgumentError(self, msg) |
| 75 | |
| 76 | conditioning = ImageConditioningInput( |
| 77 | path=resolve_path(values[0]), |
| 78 | frame_idx=int(values[1]), |
| 79 | strength=float(values[2]), |
| 80 | crf=int(values[3]) if len(values) > 3 else DEFAULT_IMAGE_CRF, |
| 81 | ) |
| 82 | current = getattr(namespace, self.dest) or [] |
| 83 | current.append(conditioning) |
| 84 | setattr(namespace, self.dest, current) |
| 85 | |
| 86 | |
| 87 | class LoraAction(argparse.Action): |
| 88 | def __call__( |
| 89 | self, |
| 90 | parser: argparse.ArgumentParser, # noqa: ARG002 |
| 91 | namespace: argparse.Namespace, |
| 92 | values: list[str], |
| 93 | option_string: str | None = None, |
| 94 | ) -> None: |
| 95 | if len(values) > 2: |
| 96 | msg = f"{option_string} accepts at most 2 arguments (PATH and optional STRENGTH), got {len(values)} values" |
| 97 | raise argparse.ArgumentError(self, msg) |
| 98 | |
| 99 | path = values[0] |
| 100 | strength_str = values[1] if len(values) > 1 else str(DEFAULT_LORA_STRENGTH) |
| 101 | |
| 102 | resolved_path = resolve_path(path) |
| 103 | strength = float(strength_str) |
| 104 | |
| 105 | current = getattr(namespace, self.dest) or [] |
| 106 | current.append(LoraPathStrengthAndSDOps(resolved_path, strength, LTXV_LORA_COMFY_RENAMING_MAP)) |
| 107 | setattr(namespace, self.dest, current) |
| 108 | |
| 109 | |
| 110 | def resolve_path(path: str) -> str: |
| 111 | return str(Path(path).expanduser().resolve().as_posix()) |
| 112 | |
| 113 | |
| 114 | QUANTIZATION_POLICIES = ("fp8-cast", "fp8-scaled-mm") |
| 115 | |
| 116 | |
| 117 | class QuantizationAction(argparse.Action): |
| 118 | def __call__( |
| 119 | self, |
| 120 | parser: argparse.ArgumentParser, # noqa: ARG002 |
| 121 | namespace: argparse.Namespace, |
| 122 | values: list[str], |
| 123 | option_string: str | None = None, |
| 124 | ) -> None: |
| 125 | if len(values) > 2: |
| 126 | msg = ( |
| 127 | f"{option_string} accepts at most 2 arguments (POLICY and optional AMAX_PATH), got {len(values)} values" |
| 128 | ) |
| 129 | raise argparse.ArgumentError(self, msg) |
| 130 | |
| 131 | policy_name = values[0] |
| 132 | if policy_name not in QUANTIZATION_POLICIES: |
| 133 | msg = f"Unknown quantization policy '{policy_name}'. Choose from: {', '.join(QUANTIZATION_POLICIES)}" |
| 134 | raise argparse.ArgumentError(self, msg) |
| 135 | |
| 136 | if policy_name == "fp8-cast": |
| 137 | if len(values) > 1: |
| 138 | msg = f"{option_string} fp8-cast does not accept additional arguments" |
| 139 | raise argparse.ArgumentError(self, msg) |
| 140 | policy = QuantizationPolicy.fp8_cast() |
| 141 | elif policy_name == "fp8-scaled-mm": |
| 142 | amax_path = resolve_path(values[1]) if len(values) > 1 else None |
| 143 | policy = QuantizationPolicy.fp8_scaled_mm(amax_path) |
| 144 | |
| 145 | setattr(namespace, self.dest, policy) |
| 146 | |
| 147 | |
| 148 | def detect_checkpoint_path(distilled: bool = False) -> str: |
| 149 | """Pre-parse argv to extract the checkpoint path before building the full parser.""" |
| 150 | pre = argparse.ArgumentParser(add_help=False) |
| 151 | flag = "--distilled-checkpoint-path" if distilled else "--checkpoint-path" |
| 152 | pre.add_argument(flag, type=resolve_path, required=True) |
| 153 | known, _ = pre.parse_known_args() |
| 154 | return known.distilled_checkpoint_path if distilled else known.checkpoint_path |
| 155 | |
| 156 | |
| 157 | def basic_arg_parser( |
| 158 | params: PipelineParams = LTX_2_3_PARAMS, |
| 159 | distilled: bool = False, |
| 160 | ) -> argparse.ArgumentParser: |
| 161 | parser = argparse.ArgumentParser() |
| 162 | if distilled: |
| 163 | parser.add_argument( |
| 164 | "--distilled-checkpoint-path", |
| 165 | type=resolve_path, |
| 166 | required=True, |
| 167 | help="Path to LTX-2 distilled model checkpoint (.safetensors file).", |
| 168 | ) |
| 169 | else: |
| 170 | parser.add_argument( |
| 171 | "--checkpoint-path", |
| 172 | type=resolve_path, |
| 173 | required=True, |
| 174 | help="Path to LTX-2 model checkpoint (.safetensors file).", |
| 175 | ) |
| 176 | parser.add_argument( |
| 177 | "--gemma-root", |
| 178 | type=resolve_path, |
| 179 | required=True, |
| 180 | help="Path to the root directory containing the Gemma text encoder model files.", |
| 181 | ) |
| 182 | parser.add_argument( |
| 183 | "--prompt", |
| 184 | type=str, |
| 185 | required=True, |
| 186 | help="Text prompt describing the desired video content to be generated by the model.", |
| 187 | ) |
| 188 | parser.add_argument( |
| 189 | "--output-path", |
| 190 | type=resolve_path, |
| 191 | required=True, |
| 192 | help="Path to the output video file (MP4 format).", |
| 193 | ) |
| 194 | parser.add_argument( |
| 195 | "--seed", |
| 196 | type=int, |
| 197 | default=params.seed, |
| 198 | help=f"Random seed for reproducible generation (default: {params.seed}).", |
| 199 | ) |
| 200 | parser.add_argument( |
| 201 | "--height", |
| 202 | type=int, |
| 203 | default=params.stage_1_height, |
| 204 | help=f"Video height in pixels, divisible by 32 (default: {params.stage_1_height}).", |
| 205 | ) |
| 206 | parser.add_argument( |
| 207 | "--width", |
| 208 | type=int, |
| 209 | default=params.stage_1_width, |
| 210 | help=f"Width of the generated video in pixels, should be divisible by 32 (default: {params.stage_1_width}).", |
| 211 | ) |
| 212 | parser.add_argument( |
| 213 | "--num-frames", |
| 214 | type=int, |
| 215 | default=params.num_frames, |
| 216 | help=f"Number of frames to generate in the output video sequence, num-frames = (8 x K) + 1, " |
| 217 | f"where k is a non-negative integer (default: {params.num_frames}).", |
| 218 | ) |
| 219 | parser.add_argument( |
| 220 | "--frame-rate", |
| 221 | type=float, |
| 222 | default=params.frame_rate, |
| 223 | help=f"Frame rate of the generated video (fps) (default: {params.frame_rate}).", |
| 224 | ) |
| 225 | parser.add_argument( |
| 226 | "--num-inference-steps", |
| 227 | type=int, |
| 228 | default=params.num_inference_steps, |
| 229 | help=( |
| 230 | f"Number of denoising steps in the diffusion sampling process. " |
| 231 | f"Higher values improve quality but increase generation time (default: {params.num_inference_steps})." |
| 232 | ), |
| 233 | ) |
| 234 | parser.add_argument( |
| 235 | "--image", |
| 236 | dest="images", |
| 237 | action=ImageAction, |
| 238 | nargs="+", |
| 239 | metavar="ARG", |
| 240 | default=[], |
| 241 | help=( |
| 242 | "Image conditioning input: PATH FRAME_IDX STRENGTH [CRF]. " |
| 243 | "PATH is the image file, FRAME_IDX is the target frame index, " |
| 244 | "STRENGTH is the conditioning strength (all three required). " |
| 245 | f"CRF is the optional H.264 compression quality (0=lossless, default: {DEFAULT_IMAGE_CRF}). " |
| 246 | "Can be specified multiple times. Example: --image path/to/image1.jpg 0 0.8 " |
| 247 | "--image path/to/image2.jpg 160 0.9 0" |
| 248 | ), |
| 249 | ) |
| 250 | parser.add_argument( |
| 251 | "--lora", |
| 252 | dest="lora", |
| 253 | action=LoraAction, |
| 254 | nargs="+", # Accept 1-2 arguments per use (path and optional strength); validation is handled in LoraAction |
| 255 | metavar=("PATH", "STRENGTH"), |
| 256 | default=[], |
| 257 | help=( |
| 258 | "LoRA (Low-Rank Adaptation) model: path to model file and optional strength " |
| 259 | f"(default strength: {DEFAULT_LORA_STRENGTH}). Can be specified multiple times. " |
| 260 | "Example: --lora path/to/lora1.safetensors 0.8 --lora path/to/lora2.safetensors" |
| 261 | ), |
| 262 | ) |
| 263 | |
| 264 | parser.add_argument("--enhance-prompt", action="store_true") |
| 265 | parser.add_argument( |
| 266 | "--quantization", |
| 267 | dest="quantization", |
| 268 | action=QuantizationAction, |
| 269 | nargs="+", |
| 270 | metavar=("POLICY", "AMAX_PATH"), |
| 271 | default=None, |
| 272 | help=( |
| 273 | f"Quantization policy: {', '.join(QUANTIZATION_POLICIES)}. " |
| 274 | "fp8-cast uses FP8 casting with upcasting during inference. " |
| 275 | "fp8-scaled-mm uses FP8 scaled matrix multiplication (optionally provide amax calibration file path). " |
| 276 | "Example: --quantization fp8-cast or --quantization fp8-scaled-mm /path/to/amax.json" |
| 277 | ), |
| 278 | ) |
| 279 | return parser |
| 280 | |
| 281 | |
| 282 | def default_1_stage_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argparse.ArgumentParser: |
| 283 | video_guider = params.video_guider_params |
| 284 | audio_guider = params.audio_guider_params |
| 285 | parser = basic_arg_parser(params=params) |
| 286 | parser.add_argument( |
| 287 | "--negative-prompt", |
| 288 | type=str, |
| 289 | default=DEFAULT_NEGATIVE_PROMPT, |
| 290 | help=( |
| 291 | "Negative prompt describing what should not appear in the generated video, " |
| 292 | "used to guide the diffusion process away from unwanted content. " |
| 293 | "Default: a comprehensive negative prompt covering common artifacts and quality issues." |
| 294 | ), |
| 295 | ) |
| 296 | parser.add_argument( |
| 297 | "--video-cfg-guidance-scale", |
| 298 | type=float, |
| 299 | default=video_guider.cfg_scale, |
| 300 | help=( |
| 301 | f"Classifier-free guidance (CFG) scale controlling how strongly " |
| 302 | f"the model adheres to the video prompt. Higher values increase prompt " |
| 303 | f"adherence but may reduce diversity. 1.0 means no effect " |
| 304 | f"(default: {video_guider.cfg_scale})." |
| 305 | ), |
| 306 | ) |
| 307 | parser.add_argument( |
| 308 | "--video-stg-guidance-scale", |
| 309 | type=float, |
| 310 | default=video_guider.stg_scale, |
| 311 | help=( |
| 312 | f"STG (Spatio-Temporal Guidance) scale controlling how strongly " |
| 313 | f"the model reacts to the perturbation of the video modality. Higher values increase " |
| 314 | f"the effect but may reduce quality. 0.0 means no effect " |
| 315 | f"(default: {video_guider.stg_scale})." |
| 316 | ), |
| 317 | ) |
| 318 | parser.add_argument( |
| 319 | "--video-rescale-scale", |
| 320 | type=float, |
| 321 | default=video_guider.rescale_scale, |
| 322 | help=( |
| 323 | f"Rescale scale controlling how strongly " |
| 324 | f"the model rescales the video modality after applying other guidance. Higher values tend to decrease " |
| 325 | f"oversaturation effects. 0.0 means no effect (default: {video_guider.rescale_scale})." |
| 326 | ), |
| 327 | ) |
| 328 | parser.add_argument( |
| 329 | "--video-stg-blocks", |
| 330 | type=int, |
| 331 | nargs="*", |
| 332 | default=video_guider.stg_blocks, |
| 333 | help=(f"Which transformer blocks to perturb for STG. Default: {video_guider.stg_blocks}."), |
| 334 | ) |
| 335 | parser.add_argument( |
| 336 | "--a2v-guidance-scale", |
| 337 | type=float, |
| 338 | default=video_guider.modality_scale, |
| 339 | help=( |
| 340 | f"A2V (Audio-to-Video) guidance scale controlling how strongly " |
| 341 | f"the model reacts to the perturbation of the audio-to-video cross-attention. Higher values may increase " |
| 342 | f"lipsync quality. 1.0 means no effect (default: {video_guider.modality_scale})." |
| 343 | ), |
| 344 | ) |
| 345 | parser.add_argument( |
| 346 | "--video-skip-step", |
| 347 | type=int, |
| 348 | default=video_guider.skip_step, |
| 349 | help=( |
| 350 | "Video skip step N controls periodic skipping during the video diffusion process: " |
| 351 | "only steps where step_index % (N + 1) == 0 are processed, all others are skipped " |
| 352 | f"(e.g., 0 = no skipping; 1 = skip every other step; 2 = skip 2 of every 3 steps; " |
| 353 | f"default: {video_guider.skip_step})." |
| 354 | ), |
| 355 | ) |
| 356 | parser.add_argument( |
| 357 | "--audio-cfg-guidance-scale", |
| 358 | type=float, |
| 359 | default=audio_guider.cfg_scale, |
| 360 | help=( |
| 361 | f"Audio CFG (Classifier-free guidance) scale controlling how strongly " |
| 362 | f"the model adheres to the audio prompt. Higher values increase prompt " |
| 363 | f"adherence but may reduce diversity. 1.0 means no effect " |
| 364 | f"(default: {audio_guider.cfg_scale})." |
| 365 | ), |
| 366 | ) |
| 367 | parser.add_argument( |
| 368 | "--audio-stg-guidance-scale", |
| 369 | type=float, |
| 370 | default=audio_guider.stg_scale, |
| 371 | help=( |
| 372 | f"Audio STG (Spatio-Temporal Guidance) scale controlling how strongly " |
| 373 | f"the model reacts to the perturbation of the audio modality. Higher values increase " |
| 374 | f"the effect but may reduce quality. 0.0 means no effect " |
| 375 | f"(default: {audio_guider.stg_scale})." |
| 376 | ), |
| 377 | ) |
| 378 | parser.add_argument( |
| 379 | "--audio-rescale-scale", |
| 380 | type=float, |
| 381 | default=audio_guider.rescale_scale, |
| 382 | help=( |
| 383 | f"Audio rescale scale controlling how strongly " |
| 384 | f"the model rescales the audio modality after applying other guidance. " |
| 385 | f"Experimental. 0.0 means no effect (default: {audio_guider.rescale_scale})." |
| 386 | ), |
| 387 | ) |
| 388 | parser.add_argument( |
| 389 | "--audio-stg-blocks", |
| 390 | type=int, |
| 391 | nargs="*", |
| 392 | default=audio_guider.stg_blocks, |
| 393 | help=(f"Which transformer blocks to perturb for Audio STG. Default: {audio_guider.stg_blocks}."), |
| 394 | ) |
| 395 | parser.add_argument( |
| 396 | "--v2a-guidance-scale", |
| 397 | type=float, |
| 398 | default=audio_guider.modality_scale, |
| 399 | help=( |
| 400 | f"V2A (Video-to-Audio) guidance scale controlling how strongly " |
| 401 | f"the model reacts to the perturbation of the video-to-audio cross-attention. Higher values may increase " |
| 402 | f"lipsync quality. 1.0 means no effect (default: {audio_guider.modality_scale})." |
| 403 | ), |
| 404 | ) |
| 405 | parser.add_argument( |
| 406 | "--audio-skip-step", |
| 407 | type=int, |
| 408 | default=audio_guider.skip_step, |
| 409 | help=( |
| 410 | "Audio skip step N controls periodic skipping during the audio diffusion process: " |
| 411 | "only steps where step_index % (N + 1) == 0 are processed, all others are skipped " |
| 412 | f"(e.g., 0 = no skipping; 1 = skip every other step; 2 = skip 2 of every 3 steps; " |
| 413 | f"default: {audio_guider.skip_step})." |
| 414 | ), |
| 415 | ) |
| 416 | return parser |
| 417 | |
| 418 | |
| 419 | def default_2_stage_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argparse.ArgumentParser: |
| 420 | parser = default_1_stage_arg_parser(params=params) |
| 421 | parser.set_defaults(height=params.stage_2_height, width=params.stage_2_width) |
| 422 | # Update help text to reflect 2-stage defaults |
| 423 | for action in parser._actions: |
| 424 | if "--height" in action.option_strings: |
| 425 | action.help = ( |
| 426 | f"Height of the generated video in pixels, should be divisible by 64 " |
| 427 | f"(default: {params.stage_2_height})." |
| 428 | ) |
| 429 | if "--width" in action.option_strings: |
| 430 | action.help = ( |
| 431 | f"Width of the generated video in pixels, should be divisible by 64 (default: {params.stage_2_width})." |
| 432 | ) |
| 433 | parser.add_argument( |
| 434 | "--distilled-lora", |
| 435 | dest="distilled_lora", |
| 436 | action=LoraAction, |
| 437 | nargs="+", # Accept 1-2 arguments per use (path and optional strength); validation is handled in LoraAction |
| 438 | metavar=("PATH", "STRENGTH"), |
| 439 | required=True, |
| 440 | help=( |
| 441 | "Distilled LoRA (Low-Rank Adaptation) model used in the second stage (upscaling and refinement): " |
| 442 | f"path to model file and optional strength (default strength: {DEFAULT_LORA_STRENGTH}). " |
| 443 | "The second stage upsamples the video by 2x resolution and refines it using a distilled " |
| 444 | "denoising schedule (fewer steps, no CFG). The distilled LoRA is specifically trained " |
| 445 | "for this refinement process to improve quality at higher resolutions. " |
| 446 | "Example: --distilled-lora path/to/distilled_lora.safetensors 0.8" |
| 447 | ), |
| 448 | ) |
| 449 | parser.add_argument( |
| 450 | "--spatial-upsampler-path", |
| 451 | type=resolve_path, |
| 452 | required=True, |
| 453 | help=( |
| 454 | "Path to the spatial upsampler model used to increase the resolution " |
| 455 | "of the generated video in the latent space." |
| 456 | ), |
| 457 | ) |
| 458 | return parser |
| 459 | |
| 460 | |
| 461 | def hq_2_stage_arg_parser(params: PipelineParams = LTX_2_3_HQ_PARAMS) -> argparse.ArgumentParser: |
| 462 | parser = default_2_stage_arg_parser(params=params) |
| 463 | parser.add_argument( |
| 464 | "--distilled-lora-strength-stage-1", |
| 465 | type=float, |
| 466 | default=0.25, |
| 467 | help=(f"Strength of the distilled LoRA used in the first stage (default: {0.25})."), |
| 468 | ) |
| 469 | parser.add_argument( |
| 470 | "--distilled-lora-strength-stage-2", |
| 471 | type=float, |
| 472 | default=0.5, |
| 473 | help=(f"Strength of the distilled LoRA used in the second stage (default: {0.5})."), |
| 474 | ) |
| 475 | return parser |
| 476 | |
| 477 | |
| 478 | def default_2_stage_distilled_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argparse.ArgumentParser: |
| 479 | parser = basic_arg_parser(params=params, distilled=True) |
| 480 | parser.set_defaults(height=params.stage_2_height, width=params.stage_2_width) |
| 481 | # Update help text to reflect 2-stage defaults |
| 482 | for action in parser._actions: |
| 483 | if "--height" in action.option_strings: |
| 484 | action.help = ( |
| 485 | f"Height of the generated video in pixels, should be divisible by 64 " |
| 486 | f"(default: {params.stage_2_height})." |
| 487 | ) |
| 488 | if "--width" in action.option_strings: |
| 489 | action.help = ( |
| 490 | f"Width of the generated video in pixels, should be divisible by 64 (default: {params.stage_2_width})." |
| 491 | ) |
| 492 | parser.add_argument( |
| 493 | "--spatial-upsampler-path", |
| 494 | type=resolve_path, |
| 495 | required=True, |
| 496 | help=( |
| 497 | "Path to the spatial upsampler model used to increase the resolution " |
| 498 | "of the generated video in the latent space." |
| 499 | ), |
| 500 | ) |
| 501 | return parser |
| 502 |