scripts/infer_v2v.py

#!/usr/bin/env python
"""
Video-to-video generation with a first-frame reference image.

Input : a source video (structure/motion signal, see data/README.md) + a
        reference image (target identity/appearance) + a text prompt.
Output: an MP4 video that follows the source video's structure/motion while
        adopting the reference image's appearance.

Two switchable ways to position-encode the reference image relative to the
denoised video (--pos-mode):

  first-frame (default) - the reference image REUSES the denoised video's own
      frame-0 RoPE position. Implemented via ICLoraPipeline's native `images`
      parameter -> VideoConditionByLatentIndex/KeyframeIndex. No LoRA required
      for this path; it's a base-model capability.

  reference - the reference image gets its OWN, disjoint RoPE position range
      (shifted to sit just before the earliest position already used in the
      sequence, so it never aliases with the target's frame 0). Implemented via
      `video_conditioning` -> VideoConditionByReferenceLatent. This is the
      mechanism the custom-trained IC-LoRAs in this repo
      (configs/ref_image_ic_lora.yaml, configs/v2v_reference_ic_lora.yaml) were
      actually trained against, and generally gives better results once you
      have a matching LoRA. See README.md "Two ways to position-encode the
      reference image" for the full explanation with code references.

The source video always goes through `video_conditioning` (there is no
first-frame-aligned way to fold a whole video into the `images` mechanism,
which only conditions a single frame at a fixed index).
"""

import argparse
import logging
import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import torch

REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
WEIGHTS_DIR = os.path.join(REPO_ROOT, "weights")

DEFAULT_CKPT = os.path.join(WEIGHTS_DIR, "ltx-2.3", "ltx-2.3-22b-dev.safetensors")
DEFAULT_SPATIAL_UPSCALER = os.path.join(WEIGHTS_DIR, "ltx-2.3", "ltx-2.3-spatial-upscaler-x2-1.1.safetensors")
DEFAULT_GEMMA_ROOT = os.path.join(WEIGHTS_DIR, "gemma-3-12b-it-qat-q4_0-unquantized")


def build_dev_sigmas(steps: int) -> tuple[torch.Tensor, torch.Tensor]:
    """Non-distilled sigma schedules. The custom LoRAs here are trained on the
    DEV base, so ICLoraPipeline's default distilled 8-step schedule (meant for
    the distilled base) would produce garbage - build a proper multi-step
    schedule instead, mirroring the trainer's own ValidationRunner.
    """
    from ltx_core.components.schedulers import LTX2Scheduler

    stage_1_sigmas = LTX2Scheduler().execute(steps=steps).float()
    full = LTX2Scheduler().execute(steps=steps).float()
    stage_2_sigmas = full[full <= 0.5]
    if stage_2_sigmas.numel() < 2 or stage_2_sigmas[-1].item() != 0.0:
        stage_2_sigmas = torch.tensor([0.5, 0.35, 0.22, 0.1, 0.0], dtype=torch.float32)
    return stage_1_sigmas, stage_2_sigmas


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--input-video", required=True, help="Source/structure video path.")
    ap.add_argument("--ref-image", required=True, help="Reference image (identity/appearance) path.")
    ap.add_argument("--prompt", required=True, help="Text prompt.")
    ap.add_argument("--output", required=True, help="Output .mp4 path.")

    ap.add_argument(
        "--pos-mode",
        choices=["first-frame", "reference"],
        default="first-frame",
        help="How the reference image is position-encoded relative to the denoised "
        "video. 'first-frame' (default): reuse frame-0's position. 'reference': give "
        "it its own disjoint position (requires a LoRA trained for this, see "
        "configs/ref_image_ic_lora.yaml or configs/v2v_reference_ic_lora.yaml).",
    )

    ap.add_argument("--structure-lora", required=True, help="LoRA checkpoint (.safetensors) trained to interpret "
                     "--input-video via video_conditioning (e.g. your trained "
                     "configs/v2v_reference_ic_lora.yaml checkpoint).")
    ap.add_argument("--structure-lora-strength", type=float, default=1.0)
    ap.add_argument("--structure-strength", type=float, default=1.0,
                     help="video_conditioning strength for --input-video (0=ignore, 1=full).")
    ap.add_argument("--ref-strength", type=float, default=1.0,
                     help="Conditioning strength for --ref-image (0=ignore, 1=full).")

    ap.add_argument("--checkpoint", default=DEFAULT_CKPT, help="Dev base checkpoint (.safetensors).")
    ap.add_argument("--spatial-upscaler", default=DEFAULT_SPATIAL_UPSCALER)
    ap.add_argument("--gemma-root", default=DEFAULT_GEMMA_ROOT)

    ap.add_argument("--height", type=int, default=1024)
    ap.add_argument("--width", type=int, default=1920)
    ap.add_argument("--num-frames", type=int, default=241)
    ap.add_argument("--frame-rate", type=float, default=24.0)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--steps", type=int, default=30)
    ap.add_argument("--conditioning-attention-strength", type=float, default=1.0)
    ap.add_argument("--skip-stage-2", action="store_true", help="Half-res output, faster, lower VRAM.")
    ap.add_argument("--no-offload", action="store_true", help="Disable CPU offload (faster, more VRAM).")
    ap.add_argument("--tile", action="store_true", help="Force tiled VAE decode (lower peak VRAM, slower).")
    args = ap.parse_args()

    logging.basicConfig(level=logging.INFO)

    from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
    from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
    from ltx_pipelines.ic_lora import ICLoraPipeline
    from ltx_pipelines.utils.media_io import encode_video
    from ltx_pipelines.utils.types import OffloadMode

    lora = LoraPathStrengthAndSDOps(args.structure_lora, args.structure_lora_strength, LTXV_LORA_COMFY_RENAMING_MAP)

    pipeline = ICLoraPipeline(
        distilled_checkpoint_path=args.checkpoint,
        spatial_upsampler_path=args.spatial_upscaler,
        gemma_root=args.gemma_root,
        loras=[lora],
        offload_mode=OffloadMode.NONE if args.no_offload else OffloadMode.CPU,
    )

    stage_1_sigmas, stage_2_sigmas = build_dev_sigmas(args.steps)
    tiling_config = TilingConfig.default() if args.tile else None

    # video_conditioning: always carries the input (structure) video.
    video_conditioning = [(args.input_video, args.structure_strength)]
    images = []

    if args.pos_mode == "reference":
        # Reference image rides the SAME video_conditioning channel as the
        # structure video, one entry after another -> gets appended as its own
        # disjoint RoPE time range (VideoConditionByReferenceLatent), matching
        # how configs/v2v_reference_ic_lora.yaml / ref_image_ic_lora.yaml train it.
        video_conditioning.append((args.ref_image, args.ref_strength))
    else:
        # first-frame: reference image reuses frame-0's position via the base
        # model's native image-conditioning path. frame_idx=0 -> first frame.
        images.append((args.ref_image, 0, args.ref_strength))

    logging.info(
        "[infer_v2v] pos_mode=%s | video_conditioning entries=%d | images entries=%d",
        args.pos_mode,
        len(video_conditioning),
        len(images),
    )

    with torch.no_grad():
        video, audio = pipeline(
            prompt=args.prompt,
            seed=args.seed,
            height=args.height,
            width=args.width,
            num_frames=args.num_frames,
            frame_rate=args.frame_rate,
            images=images,
            video_conditioning=video_conditioning,
            tiling_config=tiling_config,
            conditioning_attention_strength=args.conditioning_attention_strength,
            skip_stage_2=args.skip_stage_2,
            stage_1_sigmas=stage_1_sigmas,
            stage_2_sigmas=stage_2_sigmas,
        )

    video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)
    encode_video(
        video=video,
        fps=args.frame_rate,
        audio=audio,
        output_path=args.output,
        video_chunks_number=video_chunks_number,
    )
    peak = torch.cuda.max_memory_allocated() / 1024**3
    print(f"Done. Peak VRAM: {peak:.1f} GB -> {args.output}")


if __name__ == "__main__":
    main()