| 1 | """Public 4-step autoregressive pure-UCPE image-to-video inference entrypoint.""" |
| 2 | |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | import argparse |
| 6 | import json |
| 7 | import subprocess |
| 8 | import sys |
| 9 | from pathlib import Path |
| 10 | |
| 11 | import torch |
| 12 | import yaml |
| 13 | |
| 14 | ROOT = Path(__file__).resolve().parent |
| 15 | for package in ("ltx-core/src", "ltx-causal/src", "ltx-pipelines/src"): |
| 16 | sys.path.insert(0, str(ROOT / package)) |
| 17 | |
| 18 | DEFAULT_CONFIG = ROOT / "configs" / "inference_wm_causal.yaml" |
| 19 | |
| 20 | |
| 21 | def _load_config(path: Path) -> dict: |
| 22 | return yaml.safe_load(path.read_text(encoding="utf-8")) or {} |
| 23 | |
| 24 | |
| 25 | def _override(value, default): |
| 26 | return default if value is None else value |
| 27 | |
| 28 | |
| 29 | def _auto_fov(image: Path, model: str, python_bin: str, width: int, height: int) -> float: |
| 30 | raw = subprocess.run( |
| 31 | [python_bin, str(ROOT / "helpers" / "moge_fov.py"), "--image", str(image), "--model", model, |
| 32 | "--target-width", str(width), "--target-height", str(height)], |
| 33 | check=True, capture_output=True, text=True, |
| 34 | ).stdout.strip() |
| 35 | return float(json.loads(raw)["fov_x_deg"]) |
| 36 | |
| 37 | |
| 38 | def parse_args(argv: list[str] | None = None) -> argparse.Namespace: |
| 39 | parser = argparse.ArgumentParser(description=__doc__) |
| 40 | parser.add_argument("--config", type=Path, default=DEFAULT_CONFIG) |
| 41 | parser.add_argument("--image", type=Path, required=True) |
| 42 | parser.add_argument("--prompt", required=True) |
| 43 | parser.add_argument("--action-str", required=True) |
| 44 | parser.add_argument("--checkpoint", type=Path) |
| 45 | parser.add_argument("--gemma-path", type=Path) |
| 46 | parser.add_argument("--output", type=Path, default=Path("outputs/echo_wm_causal.mp4")) |
| 47 | parser.add_argument("--auto-fov", action="store_true") |
| 48 | parser.add_argument("--moge-model", default="Ruicheng/moge-2-vitl-normal") |
| 49 | parser.add_argument("--moge-python", default=sys.executable) |
| 50 | parser.add_argument("--fov-deg", type=float) |
| 51 | parser.add_argument("--translation-speed", type=float) |
| 52 | parser.add_argument("--rotation-speed-deg", type=float) |
| 53 | parser.add_argument("--pitch-limit-deg", type=float) |
| 54 | parser.add_argument("--width", type=int) |
| 55 | parser.add_argument("--height", type=int) |
| 56 | parser.add_argument("--num-frames", type=int) |
| 57 | parser.add_argument("--fps", type=float) |
| 58 | parser.add_argument("--timesteps", type=int, nargs="+") |
| 59 | parser.add_argument( |
| 60 | "--video-local-attn-size", "--video_local_attn_size", |
| 61 | dest="video_local_attn_size", type=int, |
| 62 | ) |
| 63 | parser.add_argument( |
| 64 | "--video-sink-size", "--video_sink_size", dest="video_sink_size", type=int, |
| 65 | ) |
| 66 | parser.add_argument( |
| 67 | "--video-chunk-size", "--video_chunk_size", dest="video_chunk_size", type=int, |
| 68 | ) |
| 69 | parser.add_argument("--seed", type=int) |
| 70 | parser.add_argument("--no-audio", action="store_true") |
| 71 | parser.add_argument( |
| 72 | "--action-overlay", action=argparse.BooleanOptionalAction, default=True, |
| 73 | help="Write a second MP4 with a Genie-style action HUD overlay " |
| 74 | "(default: enabled; disable with --no-action-overlay).", |
| 75 | ) |
| 76 | return parser.parse_args(argv) |
| 77 | |
| 78 | |
| 79 | @torch.inference_mode() |
| 80 | def main() -> None: |
| 81 | args = parse_args() |
| 82 | # Keep --help usable before optional media/CUDA dependencies are installed. |
| 83 | from ltx_causal import CausalCacheConfig |
| 84 | from ltx_core.model.video_vae.tiling import TilingConfig |
| 85 | from ltx_core.model.video_vae.video_vae import get_video_chunks_number |
| 86 | from ltx_pipelines.causal_ti2vid import CausalTI2VidPipeline |
| 87 | from ltx_pipelines.utils.args import ImageConditioningInput |
| 88 | from ltx_pipelines.utils.media_io import encode_video |
| 89 | |
| 90 | from helpers.action_camera import ( |
| 91 | DEFAULT_PITCH_LIMIT_DEG, |
| 92 | DEFAULT_ROTATION_SPEED_DEG, |
| 93 | DEFAULT_TRANSLATION_SPEED, |
| 94 | ) |
| 95 | from helpers.action_condition import action_config, build_action_trajectory, build_causal_action_condition |
| 96 | from helpers.action_overlay import overlay_genie_on_video |
| 97 | |
| 98 | if not args.image.is_file(): |
| 99 | raise FileNotFoundError(f"I2V first-frame image not found: {args.image}") |
| 100 | cfg = _load_config(args.config) |
| 101 | model_cfg, video_cfg, causal_cfg, action_cfg = ( |
| 102 | cfg.get("model", {}), cfg.get("video", {}), cfg.get("causal", {}), cfg.get("action", {}) |
| 103 | ) |
| 104 | checkpoint = args.checkpoint or ROOT / model_cfg.get("checkpoint", "checkpoints/echo-wm-flash.safetensors") |
| 105 | gemma_path = args.gemma_path or ROOT / model_cfg["gemma_path"] |
| 106 | width = _override(args.width, video_cfg.get("width", 1280)) |
| 107 | height = _override(args.height, video_cfg.get("height", 704)) |
| 108 | num_frames = _override(args.num_frames, video_cfg.get("num_frames", 241)) |
| 109 | fps = _override(args.fps, video_cfg.get("fps", 24.0)) |
| 110 | seed = _override(args.seed, video_cfg.get("seed", 42)) |
| 111 | fov = _override(args.fov_deg, action_cfg.get("fov_deg", 70.0)) |
| 112 | if args.auto_fov: |
| 113 | fov = _auto_fov(args.image, args.moge_model, args.moge_python, width, height) |
| 114 | cache = CausalCacheConfig( |
| 115 | video_local_attn_size=_override( |
| 116 | args.video_local_attn_size, causal_cfg.get("video_local_attn_size", 19) |
| 117 | ), |
| 118 | video_sink_size=_override( |
| 119 | args.video_sink_size, causal_cfg.get("video_sink_size", 7) |
| 120 | ), |
| 121 | video_chunk_size=_override( |
| 122 | args.video_chunk_size, causal_cfg.get("video_chunk_size", 3) |
| 123 | ), |
| 124 | ) |
| 125 | cache.validate() |
| 126 | timesteps = tuple(_override(args.timesteps, causal_cfg.get("timesteps", [1000, 750, 500, 250]))) |
| 127 | action_kwargs = dict( |
| 128 | action=args.action_str, num_frames=num_frames, width=width, height=height, |
| 129 | translation_speed=_override(args.translation_speed, action_cfg.get("translation_speed", DEFAULT_TRANSLATION_SPEED)), |
| 130 | rotation_speed_deg=_override(args.rotation_speed_deg, action_cfg.get("rotation_speed_deg", DEFAULT_ROTATION_SPEED_DEG)), |
| 131 | pitch_limit_deg=_override(args.pitch_limit_deg, action_cfg.get("pitch_limit_deg", DEFAULT_PITCH_LIMIT_DEG)), |
| 132 | fov_deg=fov, fps=fps, |
| 133 | ) |
| 134 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| 135 | action = build_causal_action_condition(device=device, **action_kwargs) |
| 136 | trajectory = build_action_trajectory( |
| 137 | args.action_str, num_frames=num_frames, |
| 138 | translation_speed=action_kwargs["translation_speed"], |
| 139 | rotation_speed_deg=action_kwargs["rotation_speed_deg"], |
| 140 | pitch_limit_deg=action_kwargs["pitch_limit_deg"], fps=fps, |
| 141 | ) if args.action_overlay else None |
| 142 | |
| 143 | pipeline = CausalTI2VidPipeline( |
| 144 | checkpoint_path=str(checkpoint), gemma_root=str(gemma_path), device=device, |
| 145 | action_config=action_config(width, height), cache_config=cache, |
| 146 | ) |
| 147 | video, audio = pipeline( |
| 148 | prompt=args.prompt, seed=seed, height=height, width=width, |
| 149 | num_frames=num_frames, frame_rate=fps, |
| 150 | images=[ImageConditioningInput(str(args.image), 0, 1.0)], |
| 151 | action_cond=action, timesteps=timesteps, |
| 152 | video_tiling_config=TilingConfig.default(), |
| 153 | ) |
| 154 | args.output.parent.mkdir(parents=True, exist_ok=True) |
| 155 | encode_video( |
| 156 | video=video, fps=int(fps), audio=None if args.no_audio else audio, |
| 157 | output_path=str(args.output), |
| 158 | video_chunks_number=get_video_chunks_number(num_frames, TilingConfig.default()), |
| 159 | ) |
| 160 | overlay_output = None |
| 161 | if trajectory is not None: |
| 162 | overlay_output = args.output.with_name(f"{args.output.stem}_action{args.output.suffix}") |
| 163 | overlay_genie_on_video(args.output, trajectory, output_path=overlay_output) |
| 164 | metadata = { |
| 165 | "mode": "causal_4_step", "prompt": args.prompt, "action": args.action_str, |
| 166 | "checkpoint": str(checkpoint), "timesteps": list(timesteps), "seed": seed, |
| 167 | "width": width, "height": height, "num_frames": num_frames, "fps": fps, |
| 168 | "video_local_attn_size": cache.video_local_attn_size, |
| 169 | "video_sink_size": cache.video_sink_size, |
| 170 | "video_chunk_size": cache.video_chunk_size, |
| 171 | "audio_local_attn_size": cache.audio_local_attn_size, |
| 172 | "audio_sink_size": cache.audio_sink_size, "cfg": False, |
| 173 | "cache_policy": "bounded sink-plus-FIFO", "camera_policy": "bounded anchor translation", |
| 174 | "action_overlay": bool(args.action_overlay), |
| 175 | "overlay_output": overlay_output.name if overlay_output else None, |
| 176 | } |
| 177 | args.output.with_suffix(".json").write_text(json.dumps(metadata, indent=2), encoding="utf-8") |
| 178 | print(f"Saved {args.output}") |
| 179 | if overlay_output: |
| 180 | print(f"Saved {overlay_output}") |
| 181 | |
| 182 | |
| 183 | if __name__ == "__main__": |
| 184 | main() |
| 185 |