| |
| """ |
| 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 = [(args.input_video, args.structure_strength)] |
| images = [] |
|
|
| if args.pos_mode == "reference": |
| |
| |
| |
| |
| video_conditioning.append((args.ref_image, args.ref_strength)) |
| else: |
| |
| |
| 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() |
|
|