| 1 | #!/usr/bin/env python3 |
| 2 | """Gradio app for Echo-WM world model inference. |
| 3 | |
| 4 | Features: |
| 5 | - Image upload or selection from examples |
| 6 | - Six-field cinematic prompt from PROMPT_SKILL.md |
| 7 | - Action string input with presets (WASD + camera controls) |
| 8 | - Genie-style HUD overlay (on by default) |
| 9 | - Video + audio generation |
| 10 | |
| 11 | Run: |
| 12 | CUDA_VISIBLE_DEVICES=0 python gradio_echo_wm.py |
| 13 | # then open http://0.0.0.0:7860 |
| 14 | """ |
| 15 | from __future__ import annotations |
| 16 | |
| 17 | import argparse |
| 18 | import json |
| 19 | import os |
| 20 | import sys |
| 21 | import time |
| 22 | import traceback |
| 23 | from pathlib import Path |
| 24 | |
| 25 | import gradio as gr |
| 26 | import torch |
| 27 | import yaml |
| 28 | |
| 29 | # Setup paths |
| 30 | ROOT = Path(__file__).resolve().parent |
| 31 | REPO_ROOT = ROOT.parent |
| 32 | for package in ("ltx-core/src", "ltx-pipelines/src"): |
| 33 | sys.path.insert(0, str(ROOT / package)) |
| 34 | |
| 35 | from ltx_core.components.guiders import MultiModalGuiderParams # noqa: E402 |
| 36 | from ltx_pipelines.ti2vid_one_stage import TI2VidOneStagePipeline # noqa: E402 |
| 37 | from ltx_core.model.video_vae.tiling import TilingConfig # noqa: E402 |
| 38 | from ltx_core.model.video_vae.video_vae import get_video_chunks_number # noqa: E402 |
| 39 | from ltx_pipelines.utils.args import ImageConditioningInput # noqa: E402 |
| 40 | from ltx_pipelines.utils.media_io import encode_video # noqa: E402 |
| 41 | |
| 42 | from helpers.action_condition import ( # noqa: E402 |
| 43 | action_config, |
| 44 | build_action_condition, |
| 45 | build_action_trajectory, |
| 46 | ) |
| 47 | from helpers.action_camera import ( # noqa: E402 |
| 48 | DEFAULT_PITCH_LIMIT_DEG, |
| 49 | DEFAULT_ROTATION_SPEED_DEG, |
| 50 | DEFAULT_TRANSLATION_SPEED, |
| 51 | parse_action_string, |
| 52 | ) |
| 53 | from helpers.action_overlay import overlay_genie_on_video # noqa: E402 |
| 54 | |
| 55 | # Default paths |
| 56 | DEFAULT_CONFIG = ROOT / "configs" / "inference_wm.yaml" |
| 57 | DEFAULT_CHECKPOINT = ROOT / "checkpoints" / "echo-wm-base.safetensors" |
| 58 | DEFAULT_GEMMA = ROOT / "checkpoints" / "gemma-3" |
| 59 | EXAMPLES_DIR = ROOT / "examples" |
| 60 | OUTPUT_ROOT = ROOT / "outputs" / "gradio_app" |
| 61 | VLM_MODEL = os.environ.get("VLM_MODEL", "Qwen/Qwen3-VL-8B-Instruct") |
| 62 | |
| 63 | NEGATIVE_PROMPT = ( |
| 64 | "worst quality, inconsistent motion, blurry, jittery, distorted, " |
| 65 | "game UI, video game interface, HUD, heads-up display, menu, status bar, " |
| 66 | "health bar, score, minimap, crosshair, reticle, buttons, icons, subtitles, " |
| 67 | "captions, watermark, logo, text overlay, user interface" |
| 68 | ) |
| 69 | |
| 70 | # Action presets (~240 frames for 10s at 24fps) |
| 71 | ACTION_PRESETS = { |
| 72 | "forward (w-240)": "w-240", |
| 73 | "S-curve weave": "wj-60,wl-60,wj-60,wl-60", |
| 74 | "orbit left": "wj-40,dj-80,wj-40,dj-80", |
| 75 | "look around": "w-40,l-60,i-30,k-30,j-40,w-40", |
| 76 | "approach + tilt up": "w-70,wi-40,i-30,w-40,none-60", |
| 77 | "strafe scan": "w-40,d-50,a-100,d-50", |
| 78 | } |
| 79 | |
| 80 | ACTION_HELP = ( |
| 81 | "**Action DSL** — segments `<keys>-<frames>` joined by commas; keys held simultaneously.\n" |
| 82 | "`w`/`s` forward/back · `a`/`d` strafe left/right · `i`/`k` pitch up/down · `j`/`l` yaw (pan) left/right · " |
| 83 | "`none` holds still. Combine, e.g. `wj` = forward + pan-left. Frames total ≈ num_frames − 1 (240 for a 10s clip)." |
| 84 | ) |
| 85 | |
| 86 | # Case directories under examples/ that hold `<id>/input.png` + `<id>/case.json`. |
| 87 | # The causal cases target the chained multi-shot script, not this single-shot demo. |
| 88 | CASE_COLLECTIONS = ("wm_cases",) |
| 89 | |
| 90 | |
| 91 | def discover_cases() -> dict[str, dict]: |
| 92 | """Map "<collection>/<id>" -> case fields merged with its first-frame path. |
| 93 | |
| 94 | Cases live in examples/<collection>/<id>/ with a case.json carrying the public |
| 95 | semantic controls (prompt, action, optional fov_deg/seed) next to input.png. |
| 96 | """ |
| 97 | cases: dict[str, dict] = {} |
| 98 | for collection in CASE_COLLECTIONS: |
| 99 | root = EXAMPLES_DIR / collection |
| 100 | if not root.is_dir(): |
| 101 | continue |
| 102 | for case_dir in sorted(p for p in root.iterdir() if p.is_dir()): |
| 103 | meta_path = case_dir / "case.json" |
| 104 | image_path = case_dir / "input.png" |
| 105 | if not meta_path.is_file() or not image_path.is_file(): |
| 106 | continue |
| 107 | try: |
| 108 | meta = json.loads(meta_path.read_text()) |
| 109 | except json.JSONDecodeError as exc: |
| 110 | print(f"[cases] skipping {case_dir.name}: bad case.json ({exc})", flush=True) |
| 111 | continue |
| 112 | meta["image"] = str(image_path) |
| 113 | label = case_dir.name if len(CASE_COLLECTIONS) == 1 else f"{collection}/{case_dir.name}" |
| 114 | cases[label] = meta |
| 115 | return cases |
| 116 | |
| 117 | |
| 118 | CASES = discover_cases() |
| 119 | |
| 120 | |
| 121 | class EchoWMEngine: |
| 122 | """Loads the Echo-WM model once; generates on demand.""" |
| 123 | |
| 124 | def __init__( |
| 125 | self, |
| 126 | checkpoint: Path, |
| 127 | gemma_path: Path, |
| 128 | config_path: Path, |
| 129 | device: torch.device, |
| 130 | ): |
| 131 | self.device = device |
| 132 | self.checkpoint = checkpoint |
| 133 | self.gemma_path = gemma_path |
| 134 | self.config_path = config_path |
| 135 | |
| 136 | print(f"[engine] Loading config from {config_path}", flush=True) |
| 137 | self.cfg = yaml.safe_load(config_path.read_text()) or {} |
| 138 | self.video_cfg = self.cfg.get("video", {}) |
| 139 | self.action_cfg = self.cfg.get("action", {}) |
| 140 | |
| 141 | print(f"[engine] Loading Echo-WM model...", flush=True) |
| 142 | print(f" checkpoint: {checkpoint}", flush=True) |
| 143 | print(f" gemma: {gemma_path}", flush=True) |
| 144 | |
| 145 | # The pipeline only records paths here; the 47GB of weights are read lazily on |
| 146 | # the first generation. Probe the checkpoint header now so a wrong --checkpoint |
| 147 | # fails at startup instead of surfacing minutes later as a failed generation. |
| 148 | self._probe_checkpoint(checkpoint) |
| 149 | |
| 150 | self.pipeline = TI2VidOneStagePipeline( |
| 151 | checkpoint_path=str(checkpoint), |
| 152 | gemma_root=str(gemma_path), |
| 153 | loras=(), |
| 154 | device=device, |
| 155 | action_config=None, # Will be set per generation |
| 156 | ) |
| 157 | print("[engine] Ready (weights load on first generation).", flush=True) |
| 158 | |
| 159 | @staticmethod |
| 160 | def _probe_checkpoint(checkpoint: Path) -> None: |
| 161 | """Validate the checkpoint is a readable safetensors file before serving.""" |
| 162 | if not checkpoint.is_file(): |
| 163 | raise SystemExit(f"[engine] Checkpoint not found: {checkpoint}") |
| 164 | from safetensors import safe_open |
| 165 | |
| 166 | try: |
| 167 | with safe_open(str(checkpoint), framework="pt") as f: |
| 168 | n_tensors = len(f.keys()) |
| 169 | except Exception as exc: # noqa: BLE001 - surface any unreadable file the same way |
| 170 | raise SystemExit(f"[engine] Cannot read checkpoint {checkpoint}: {exc}") from exc |
| 171 | size_gb = checkpoint.stat().st_size / 1024**3 |
| 172 | print(f" verified: {n_tensors} tensors, {size_gb:.1f} GiB", flush=True) |
| 173 | |
| 174 | @torch.inference_mode() |
| 175 | def generate( |
| 176 | self, |
| 177 | image_path: str, |
| 178 | prompt: str, |
| 179 | action_str: str, |
| 180 | seed: int, |
| 181 | num_frames: int, |
| 182 | fps: float, |
| 183 | steps: int, |
| 184 | video_cfg: float, |
| 185 | audio_cfg: float, |
| 186 | width: int, |
| 187 | height: int, |
| 188 | fov_deg: float, |
| 189 | translation_speed: float, |
| 190 | rotation_speed_deg: float, |
| 191 | pitch_limit_deg: float, |
| 192 | generate_audio: bool, |
| 193 | overlay: bool, |
| 194 | out_dir: Path, |
| 195 | ) -> tuple[Path, Path | None, dict]: |
| 196 | """Returns (video_path, overlaid_path_or_None, timing).""" |
| 197 | timing: dict[str, float] = {} |
| 198 | |
| 199 | # Validate action string early |
| 200 | parse_action_string(action_str) |
| 201 | |
| 202 | # Build action condition |
| 203 | t0 = time.time() |
| 204 | action_cond = build_action_condition( |
| 205 | action_str, |
| 206 | num_frames=num_frames, |
| 207 | width=width, |
| 208 | height=height, |
| 209 | translation_speed=translation_speed, |
| 210 | rotation_speed_deg=rotation_speed_deg, |
| 211 | pitch_limit_deg=pitch_limit_deg, |
| 212 | fov_deg=fov_deg, |
| 213 | device=self.device, |
| 214 | fps=fps, |
| 215 | ) |
| 216 | timing["action_prep"] = time.time() - t0 |
| 217 | |
| 218 | # Generate |
| 219 | t0 = time.time() |
| 220 | # Update action config for this generation |
| 221 | self.pipeline.action_config = action_config(width, height) |
| 222 | |
| 223 | video, audio = self.pipeline( |
| 224 | prompt=prompt, |
| 225 | negative_prompt=self.cfg.get("negative_prompt", NEGATIVE_PROMPT), |
| 226 | seed=seed, |
| 227 | height=height, |
| 228 | width=width, |
| 229 | num_frames=num_frames, |
| 230 | frame_rate=fps, |
| 231 | num_inference_steps=steps, |
| 232 | video_guider_params=MultiModalGuiderParams( |
| 233 | cfg_scale=video_cfg, |
| 234 | stg_scale=self.video_cfg.get("stg_scale", 1.0), |
| 235 | stg_blocks=self.video_cfg.get("stg_blocks", [29]), |
| 236 | ), |
| 237 | audio_guider_params=MultiModalGuiderParams( |
| 238 | cfg_scale=audio_cfg, |
| 239 | stg_scale=self.video_cfg.get("stg_scale", 1.0), |
| 240 | stg_blocks=self.video_cfg.get("stg_blocks", [29]), |
| 241 | ), |
| 242 | images=[ImageConditioningInput(str(image_path), 0, 1.0)], |
| 243 | action_cond=action_cond, |
| 244 | video_tiling_config=TilingConfig.default(), |
| 245 | ) |
| 246 | timing["generate"] = time.time() - t0 |
| 247 | |
| 248 | # Save video |
| 249 | out_dir.mkdir(parents=True, exist_ok=True) |
| 250 | video_path = out_dir / "output.mp4" |
| 251 | |
| 252 | t0 = time.time() |
| 253 | encode_video( |
| 254 | video=video, |
| 255 | fps=int(fps), |
| 256 | audio=audio if generate_audio else None, |
| 257 | output_path=str(video_path), |
| 258 | video_chunks_number=get_video_chunks_number(num_frames, TilingConfig.default()), |
| 259 | ) |
| 260 | timing["encode"] = time.time() - t0 |
| 261 | |
| 262 | # Overlay if requested |
| 263 | overlaid_path = None |
| 264 | if overlay: |
| 265 | t0 = time.time() |
| 266 | trajectory = build_action_trajectory( |
| 267 | action_str, |
| 268 | num_frames=num_frames, |
| 269 | translation_speed=translation_speed, |
| 270 | rotation_speed_deg=rotation_speed_deg, |
| 271 | pitch_limit_deg=pitch_limit_deg, |
| 272 | fps=fps, |
| 273 | ) |
| 274 | overlaid_path = out_dir / "output_action.mp4" |
| 275 | overlay_genie_on_video(video_path, trajectory, output_path=overlaid_path) |
| 276 | timing["overlay"] = time.time() - t0 |
| 277 | |
| 278 | return video_path, overlaid_path, timing |
| 279 | |
| 280 | |
| 281 | |
| 282 | def build_ui(engine: EchoWMEngine) -> gr.Blocks: |
| 283 | """Build Gradio interface.""" |
| 284 | run_counter = {"n": 0} |
| 285 | |
| 286 | def on_preset(name: str): |
| 287 | return gr.update(value=ACTION_PRESETS.get(name, "w-240")) |
| 288 | |
| 289 | def on_case(name: str): |
| 290 | """Fill the first frame and the case's authored controls.""" |
| 291 | case = CASES.get(name) |
| 292 | if case is None: |
| 293 | return (gr.update(),) * 5 |
| 294 | return ( |
| 295 | gr.update(value=case["image"]), |
| 296 | gr.update(value=case.get("prompt", "")), |
| 297 | gr.update(value=case.get("action", "w-240")), |
| 298 | gr.update(value=case.get("fov_deg", 70.0)), |
| 299 | gr.update(value=case.get("seed", 42)), |
| 300 | ) |
| 301 | |
| 302 | def on_generate( |
| 303 | image_path, prompt, action_str, seed, num_frames, fps, steps, |
| 304 | video_cfg, audio_cfg, width, height, fov_deg, |
| 305 | translation_speed, rotation_speed, pitch_limit, |
| 306 | gen_audio, overlay, |
| 307 | ): |
| 308 | if not image_path: |
| 309 | yield "❌ Pick or upload an image first.", None, None |
| 310 | return |
| 311 | if not (prompt or "").strip(): |
| 312 | yield "❌ Prompt is empty.", None, None |
| 313 | return |
| 314 | try: |
| 315 | parse_action_string(action_str) |
| 316 | except Exception as e: |
| 317 | yield f"❌ Invalid action string: {e}", None, None |
| 318 | return |
| 319 | |
| 320 | run_counter["n"] += 1 |
| 321 | out_dir = OUTPUT_ROOT / f"run_{run_counter['n']:04d}" |
| 322 | |
| 323 | est_time = int(num_frames) * steps // 100 # rough estimate |
| 324 | yield ( |
| 325 | f"⏳ Generating {int(num_frames)}f @ {int(steps)} steps (~{est_time}s)…\n" |
| 326 | f"action=[{action_str}] seed={int(seed)}" |
| 327 | ), None, None |
| 328 | |
| 329 | t0 = time.time() |
| 330 | try: |
| 331 | video_path, overlaid_path, timing = engine.generate( |
| 332 | image_path=image_path, |
| 333 | prompt=prompt, |
| 334 | action_str=action_str, |
| 335 | seed=int(seed), |
| 336 | num_frames=int(num_frames), |
| 337 | fps=float(fps), |
| 338 | steps=int(steps), |
| 339 | video_cfg=float(video_cfg), |
| 340 | audio_cfg=float(audio_cfg), |
| 341 | width=int(width), |
| 342 | height=int(height), |
| 343 | fov_deg=float(fov_deg), |
| 344 | translation_speed=float(translation_speed), |
| 345 | rotation_speed_deg=float(rotation_speed), |
| 346 | pitch_limit_deg=float(pitch_limit), |
| 347 | generate_audio=bool(gen_audio), |
| 348 | overlay=bool(overlay), |
| 349 | out_dir=out_dir, |
| 350 | ) |
| 351 | except Exception as e: |
| 352 | yield f"❌ Generation failed: {e}\n{traceback.format_exc()[-800:]}", None, None |
| 353 | return |
| 354 | |
| 355 | shown = overlaid_path or video_path |
| 356 | parts = " ".join(f"{k}={v:.1f}s" for k, v in timing.items()) |
| 357 | msg = f"✅ Done in {time.time() - t0:.1f}s ({parts}).\n video: {video_path.name}" |
| 358 | if overlaid_path: |
| 359 | msg += f"\n overlay: {overlaid_path.name}" |
| 360 | yield msg, str(shown), str(video_path) |
| 361 | |
| 362 | with gr.Blocks(title="Echo-WM World Model") as demo: |
| 363 | gr.Markdown( |
| 364 | f"# Echo-WM: Action-Conditioned World Model\n" |
| 365 | f"Checkpoint: `{engine.checkpoint.name}` · Gemma: `{engine.gemma_path.name}`" |
| 366 | ) |
| 367 | |
| 368 | with gr.Row(): |
| 369 | with gr.Column(scale=1): |
| 370 | case_picker = gr.Dropdown( |
| 371 | list(CASES), |
| 372 | label="Example case (fills image, prompt, action, FOV, seed)", |
| 373 | value=None, |
| 374 | ) |
| 375 | image = gr.Image(label="First-frame image", type="filepath", height=300) |
| 376 | |
| 377 | with gr.Row(): |
| 378 | prompt = gr.Textbox( |
| 379 | label="Prompt", |
| 380 | lines=4, |
| 381 | placeholder="Describe the scene, style, perspective...", |
| 382 | ) |
| 383 | gr.Markdown(ACTION_HELP) |
| 384 | action = gr.Textbox(label="Action string", value="w-240") |
| 385 | preset = gr.Dropdown( |
| 386 | list(ACTION_PRESETS), |
| 387 | label="Action preset", |
| 388 | value="forward (w-240)", |
| 389 | ) |
| 390 | |
| 391 | with gr.Accordion("Video Settings", open=False): |
| 392 | with gr.Row(): |
| 393 | # Fixed: the model is trained at this resolution. |
| 394 | width = gr.Number(label="Width", value=1280, precision=0, interactive=False) |
| 395 | height = gr.Number(label="Height", value=704, precision=0, interactive=False) |
| 396 | with gr.Row(): |
| 397 | num_frames = gr.Number(label="Frames (241=10s)", value=241, precision=0) |
| 398 | fps = gr.Number(label="FPS", value=24, precision=1) |
| 399 | with gr.Row(): |
| 400 | steps = gr.Slider(10, 50, value=30, step=1, label="Inference steps") |
| 401 | seed = gr.Number(label="Seed", value=42, precision=0) |
| 402 | with gr.Row(): |
| 403 | video_cfg = gr.Slider(1.0, 8.0, value=4.0, step=0.5, label="Video CFG") |
| 404 | audio_cfg = gr.Slider(1.0, 8.0, value=2.0, step=0.5, label="Audio CFG") |
| 405 | |
| 406 | with gr.Accordion("Action Settings", open=False): |
| 407 | fov_deg = gr.Slider(30, 120, value=70, step=5, label="FOV (degrees)") |
| 408 | translation_speed = gr.Slider( |
| 409 | 0.005, 0.1, value=DEFAULT_TRANSLATION_SPEED, step=0.005, |
| 410 | label="Translation speed (w/s/a/d per frame)", |
| 411 | ) |
| 412 | rotation_speed = gr.Slider( |
| 413 | 0.1, 3.0, value=DEFAULT_ROTATION_SPEED_DEG, step=0.1, |
| 414 | label="Rotation speed (°/frame, i/k/j/l)", |
| 415 | ) |
| 416 | pitch_limit = gr.Slider( |
| 417 | 0, 90, value=DEFAULT_PITCH_LIMIT_DEG, step=5, |
| 418 | label="Pitch limit (degrees)", |
| 419 | ) |
| 420 | |
| 421 | with gr.Row(): |
| 422 | overlay = gr.Checkbox(label="Action HUD overlay", value=True) |
| 423 | gen_audio = gr.Checkbox(label="Generate audio", value=True) |
| 424 | |
| 425 | generate_btn = gr.Button("🚀 Generate", variant="primary", size="lg") |
| 426 | |
| 427 | with gr.Column(scale=1): |
| 428 | out_video = gr.Video(label="Result", height=400) |
| 429 | status = gr.Textbox(label="Status", lines=6, interactive=False) |
| 430 | raw_file = gr.File(label="Raw video (no overlay)", interactive=False) |
| 431 | |
| 432 | # Event handlers |
| 433 | case_picker.change( |
| 434 | on_case, inputs=case_picker, |
| 435 | outputs=[image, prompt, action, fov_deg, seed], |
| 436 | ) |
| 437 | preset.change(on_preset, inputs=preset, outputs=action) |
| 438 | generate_btn.click( |
| 439 | on_generate, |
| 440 | inputs=[ |
| 441 | image, prompt, action, seed, num_frames, fps, steps, |
| 442 | video_cfg, audio_cfg, width, height, fov_deg, |
| 443 | translation_speed, rotation_speed, pitch_limit, |
| 444 | gen_audio, overlay, |
| 445 | ], |
| 446 | outputs=[status, out_video, raw_file], |
| 447 | concurrency_limit=1, |
| 448 | ) |
| 449 | |
| 450 | return demo |
| 451 | |
| 452 | |
| 453 | def main() -> None: |
| 454 | parser = argparse.ArgumentParser(description=__doc__) |
| 455 | parser.add_argument("--checkpoint", type=Path, default=DEFAULT_CHECKPOINT) |
| 456 | parser.add_argument("--gemma-path", type=Path, default=DEFAULT_GEMMA) |
| 457 | parser.add_argument("--config", type=Path, default=DEFAULT_CONFIG) |
| 458 | parser.add_argument("--host", default="0.0.0.0") |
| 459 | parser.add_argument("--port", type=int, default=7860) |
| 460 | parser.add_argument("--share", action="store_true", help="Create public gradio share link") |
| 461 | args = parser.parse_args() |
| 462 | |
| 463 | if not torch.cuda.is_available(): |
| 464 | raise SystemExit("CUDA device required.") |
| 465 | |
| 466 | device = torch.device("cuda") |
| 467 | engine = EchoWMEngine( |
| 468 | checkpoint=args.checkpoint, |
| 469 | gemma_path=args.gemma_path, |
| 470 | config_path=args.config, |
| 471 | device=device, |
| 472 | ) |
| 473 | |
| 474 | demo = build_ui(engine) |
| 475 | OUTPUT_ROOT.mkdir(parents=True, exist_ok=True) |
| 476 | |
| 477 | print(f"[server] Serving on http://127.0.0.1:{args.port} " |
| 478 | f"(forward port {args.port} if you are on a remote host)", flush=True) |
| 479 | demo.queue().launch( |
| 480 | server_name=args.host, |
| 481 | server_port=args.port, |
| 482 | share=args.share, |
| 483 | allowed_paths=[str(OUTPUT_ROOT), str(EXAMPLES_DIR)], |
| 484 | show_error=True, |
| 485 | ) |
| 486 | |
| 487 | |
| 488 | if __name__ == "__main__": |
| 489 | main() |
| 490 |