Automated PR - 2026-01-05

2026-01-05 20:10:38 +00:00
parent fc3b319d34
commit 9ce438b353
153 changed files with 28100 additions and 0 deletions
@@ -0,0 +1,277 @@
+import argparse
+from pathlib import Path
+
+from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
+from ltx_pipelines.utils.constants import (
+    DEFAULT_1_STAGE_HEIGHT,
+    DEFAULT_1_STAGE_WIDTH,
+    DEFAULT_2_STAGE_HEIGHT,
+    DEFAULT_2_STAGE_WIDTH,
+    DEFAULT_CFG_GUIDANCE_SCALE,
+    DEFAULT_FRAME_RATE,
+    DEFAULT_LORA_STRENGTH,
+    DEFAULT_NEGATIVE_PROMPT,
+    DEFAULT_NUM_FRAMES,
+    DEFAULT_NUM_INFERENCE_STEPS,
+    DEFAULT_SEED,
+)
+
+
+class VideoConditioningAction(argparse.Action):
+    def __call__(
+        self,
+        parser: argparse.ArgumentParser,  # noqa: ARG002
+        namespace: argparse.Namespace,
+        values: list[str],
+        option_string: str | None = None,  # noqa: ARG002
+    ) -> None:
+        path, strength_str = values
+        resolved_path = resolve_path(path)
+        strength = float(strength_str)
+        current = getattr(namespace, self.dest) or []
+        current.append((resolved_path, strength))
+        setattr(namespace, self.dest, current)
+
+
+class ImageAction(argparse.Action):
+    def __call__(
+        self,
+        parser: argparse.ArgumentParser,  # noqa: ARG002
+        namespace: argparse.Namespace,
+        values: list[str],
+        option_string: str | None = None,  # noqa: ARG002
+    ) -> None:
+        path, frame_idx, strength_str = values
+        resolved_path = resolve_path(path)
+        frame_idx = int(frame_idx)
+        strength = float(strength_str)
+        current = getattr(namespace, self.dest) or []
+        current.append((resolved_path, frame_idx, strength))
+        setattr(namespace, self.dest, current)
+
+
+class LoraAction(argparse.Action):
+    def __call__(
+        self,
+        parser: argparse.ArgumentParser,  # noqa: ARG002
+        namespace: argparse.Namespace,
+        values: list[str],
+        option_string: str | None = None,
+    ) -> None:
+        if len(values) > 2:
+            msg = f"{option_string} accepts at most 2 arguments (PATH and optional STRENGTH), got {len(values)} values"
+            raise argparse.ArgumentError(self, msg)
+
+        path = values[0]
+        strength_str = values[1] if len(values) > 1 else str(DEFAULT_LORA_STRENGTH)
+
+        resolved_path = resolve_path(path)
+        strength = float(strength_str)
+
+        current = getattr(namespace, self.dest) or []
+        current.append(LoraPathStrengthAndSDOps(resolved_path, strength, LTXV_LORA_COMFY_RENAMING_MAP))
+        setattr(namespace, self.dest, current)
+
+
+def resolve_path(path: str) -> str:
+    return str(Path(path).expanduser().resolve().as_posix())
+
+
+def basic_arg_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--checkpoint-path",
+        type=resolve_path,
+        required=True,
+        help="Path to LTX-2 model checkpoint (.safetensors file).",
+    )
+    parser.add_argument(
+        "--gemma-root",
+        type=resolve_path,
+        required=True,
+        help="Path to the root directory containing the Gemma text encoder model files.",
+    )
+    parser.add_argument(
+        "--prompt",
+        type=str,
+        required=True,
+        help="Text prompt describing the desired video content to be generated by the model.",
+    )
+    parser.add_argument(
+        "--output-path",
+        type=resolve_path,
+        required=True,
+        help="Path to the output video file (MP4 format).",
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=DEFAULT_SEED,
+        help=(
+            f"Random seed value used to initialize the noise tensor for "
+            f"reproducible generation (default: {DEFAULT_SEED})."
+        ),
+    )
+    parser.add_argument(
+        "--height",
+        type=int,
+        default=DEFAULT_1_STAGE_HEIGHT,
+        help=f"Height of the generated video in pixels, should be divisible by 32 (default: {DEFAULT_1_STAGE_HEIGHT}).",
+    )
+    parser.add_argument(
+        "--width",
+        type=int,
+        default=DEFAULT_1_STAGE_WIDTH,
+        help=f"Width of the generated video in pixels, should be divisible by 32 (default: {DEFAULT_1_STAGE_WIDTH}).",
+    )
+    parser.add_argument(
+        "--num-frames",
+        type=int,
+        default=DEFAULT_NUM_FRAMES,
+        help=f"Number of frames to generate in the output video sequence, num-frames = (8 x K) + 1, "
+        f"where k is a non-negative integer (default: {DEFAULT_NUM_FRAMES}).",
+    )
+    parser.add_argument(
+        "--frame-rate",
+        type=float,
+        default=DEFAULT_FRAME_RATE,
+        help=f"Frame rate of the generated video (fps) (default: {DEFAULT_FRAME_RATE}).",
+    )
+    parser.add_argument(
+        "--num-inference-steps",
+        type=int,
+        default=DEFAULT_NUM_INFERENCE_STEPS,
+        help=(
+            f"Number of denoising steps in the diffusion sampling process. "
+            f"Higher values improve quality but increase generation time (default: {DEFAULT_NUM_INFERENCE_STEPS})."
+        ),
+    )
+    parser.add_argument(
+        "--image",
+        dest="images",
+        action=ImageAction,
+        nargs=3,
+        metavar=("PATH", "FRAME_IDX", "STRENGTH"),
+        default=[],
+        help=(
+            "Image conditioning input: path to image file, target frame index, "
+            "and conditioning strength (all three required). Default: empty list [] (no image conditioning). "
+            "Can be specified multiple times. Example: --image path/to/image1.jpg 0 0.8 "
+            "--image path/to/image2.jpg 160 0.9"
+        ),
+    )
+    parser.add_argument(
+        "--lora",
+        dest="lora",
+        action=LoraAction,
+        nargs="+",  # Accept 1-2 arguments per use (path and optional strength); validation is handled in LoraAction
+        metavar=("PATH", "STRENGTH"),
+        default=[],
+        help=(
+            "LoRA (Low-Rank Adaptation) model: path to model file and optional strength "
+            f"(default strength: {DEFAULT_LORA_STRENGTH}). Can be specified multiple times. "
+            "Example: --lora path/to/lora1.safetensors 0.8 --lora path/to/lora2.safetensors"
+        ),
+    )
+    parser.add_argument(
+        "--enable-fp8",
+        action="store_true",
+        help="Enable FP8 mode to reduce memory footprint by keeping model in lower precision. "
+        "Note that calculations are still performed in bfloat16 precision.",
+    )
+    parser.add_argument("--enhance-prompt", action="store_true")
+    return parser
+
+
+def default_1_stage_arg_parser() -> argparse.ArgumentParser:
+    parser = basic_arg_parser()
+    parser.add_argument(
+        "--cfg-guidance-scale",
+        type=float,
+        default=DEFAULT_CFG_GUIDANCE_SCALE,
+        help=(
+            f"Classifier-free guidance (CFG) scale controlling how strongly "
+            f"the model adheres to the prompt. Higher values increase prompt "
+            f"adherence but may reduce diversity (default: {DEFAULT_CFG_GUIDANCE_SCALE})."
+        ),
+    )
+    parser.add_argument(
+        "--negative-prompt",
+        type=str,
+        default=DEFAULT_NEGATIVE_PROMPT,
+        help=(
+            "Negative prompt describing what should not appear in the generated video, "
+            "used to guide the diffusion process away from unwanted content. "
+            "Default: a comprehensive negative prompt covering common artifacts and quality issues."
+        ),
+    )
+
+    return parser
+
+
+def default_2_stage_arg_parser() -> argparse.ArgumentParser:
+    parser = default_1_stage_arg_parser()
+    parser.set_defaults(height=DEFAULT_2_STAGE_HEIGHT, width=DEFAULT_2_STAGE_WIDTH)
+    # Update help text to reflect 2-stage defaults
+    for action in parser._actions:
+        if "--height" in action.option_strings:
+            action.help = (
+                f"Height of the generated video in pixels, should be divisible by 64 "
+                f"(default: {DEFAULT_2_STAGE_HEIGHT})."
+            )
+        if "--width" in action.option_strings:
+            action.help = (
+                f"Width of the generated video in pixels, should be divisible by 64 (default: {DEFAULT_2_STAGE_WIDTH})."
+            )
+    parser.add_argument(
+        "--distilled-lora",
+        dest="distilled_lora",
+        action=LoraAction,
+        nargs="+",  # Accept 1-2 arguments per use (path and optional strength); validation is handled in LoraAction
+        metavar=("PATH", "STRENGTH"),
+        required=True,
+        help=(
+            "Distilled LoRA (Low-Rank Adaptation) model used in the second stage (upscaling and refinement): "
+            f"path to model file and optional strength (default strength: {DEFAULT_LORA_STRENGTH}). "
+            "The second stage upsamples the video by 2x resolution and refines it using a distilled "
+            "denoising schedule (fewer steps, no CFG). The distilled LoRA is specifically trained "
+            "for this refinement process to improve quality at higher resolutions. "
+            "Example: --distilled-lora path/to/distilled_lora.safetensors 0.8"
+        ),
+    )
+    parser.add_argument(
+        "--spatial-upsampler-path",
+        type=resolve_path,
+        required=True,
+        help=(
+            "Path to the spatial upsampler model used to increase the resolution "
+            "of the generated video in the latent space."
+        ),
+    )
+    return parser
+
+
+def default_2_stage_distilled_arg_parser() -> argparse.ArgumentParser:
+    parser = basic_arg_parser()
+    parser.set_defaults(height=DEFAULT_2_STAGE_HEIGHT, width=DEFAULT_2_STAGE_WIDTH)
+    # Update help text to reflect 2-stage defaults
+    for action in parser._actions:
+        if "--height" in action.option_strings:
+            action.help = (
+                f"Height of the generated video in pixels, should be divisible by 64 "
+                f"(default: {DEFAULT_2_STAGE_HEIGHT})."
+            )
+        if "--width" in action.option_strings:
+            action.help = (
+                f"Width of the generated video in pixels, should be divisible by 64 (default: {DEFAULT_2_STAGE_WIDTH})."
+            )
+    parser.add_argument(
+        "--spatial-upsampler-path",
+        type=resolve_path,
+        required=True,
+        help=(
+            "Path to the spatial upsampler model used to increase the resolution "
+            "of the generated video in the latent space."
+        ),
+    )
+    return parser