| 1 | #!/usr/bin/env python3 |
| 2 | """Run a checked-in single-image WM case through inference_wm.py.""" |
| 3 | |
| 4 | from __future__ import annotations |
| 5 | |
| 6 | import argparse |
| 7 | import json |
| 8 | import shlex |
| 9 | import subprocess |
| 10 | import sys |
| 11 | from pathlib import Path |
| 12 | |
| 13 | ROOT = Path(__file__).resolve().parents[1] |
| 14 | CASES = ROOT / "examples" / "wm_cases" |
| 15 | DEFAULT_CHECKPOINT = ROOT / "checkpoints" / "echo-wm-base.safetensors" |
| 16 | |
| 17 | |
| 18 | def _case_dirs() -> list[Path]: |
| 19 | return sorted(p for p in CASES.iterdir() if p.is_dir() and (p / "case.json").is_file()) |
| 20 | |
| 21 | |
| 22 | def _load_case(case_dir: Path) -> dict: |
| 23 | data = json.loads((case_dir / "case.json").read_text(encoding="utf-8")) |
| 24 | required = ("prompt", "action") |
| 25 | missing = [key for key in required if key not in data] |
| 26 | if missing: |
| 27 | raise ValueError(f"{case_dir}/case.json missing fields: {', '.join(missing)}") |
| 28 | data["name"] = case_dir.name |
| 29 | data["image"] = data.get("image", "input.png") |
| 30 | data["action_str"] = data["action"] |
| 31 | data["fov_deg"] = data.get("fov_deg", 70.0) |
| 32 | image = case_dir / data["image"] |
| 33 | if not image.is_file(): |
| 34 | raise FileNotFoundError(f"Case image not found: {image}") |
| 35 | return data |
| 36 | |
| 37 | |
| 38 | def main() -> None: |
| 39 | parser = argparse.ArgumentParser(description=__doc__) |
| 40 | parser.add_argument("--case", type=Path) |
| 41 | parser.add_argument("--checkpoint", type=Path, default=DEFAULT_CHECKPOINT) |
| 42 | parser.add_argument("--gemma-path", type=Path, default=None) |
| 43 | parser.add_argument("--output-dir", type=Path, default=ROOT / "outputs" / "wm_cases") |
| 44 | parser.add_argument("--list", action="store_true") |
| 45 | parser.add_argument("--dry-run", action="store_true") |
| 46 | parser.add_argument("--num-frames", type=int) |
| 47 | parser.add_argument("--width", type=int) |
| 48 | parser.add_argument("--height", type=int) |
| 49 | parser.add_argument("--steps", type=int) |
| 50 | parser.add_argument( |
| 51 | "--action-overlay", action=argparse.BooleanOptionalAction, default=True, |
| 52 | help="Also write the HUD copy (default: enabled).", |
| 53 | ) |
| 54 | args = parser.parse_args() |
| 55 | if args.list: |
| 56 | for case_dir in _case_dirs(): |
| 57 | data = _load_case(case_dir) |
| 58 | print(f"{case_dir.name}\t{data.get('description', 'single-image I2V case')}") |
| 59 | return |
| 60 | if args.case is None: |
| 61 | parser.error("--case is required unless --list is used") |
| 62 | case_dir = args.case if args.case.is_absolute() else ROOT / args.case |
| 63 | case_dir = case_dir.resolve() |
| 64 | data = _load_case(case_dir) |
| 65 | output = args.output_dir / data["name"] / "result.mp4" |
| 66 | command = [ |
| 67 | sys.executable, str(ROOT / "inference_wm.py"), |
| 68 | "--image", str(case_dir / data["image"]), |
| 69 | "--prompt", data["prompt"], "--action-str", data["action_str"], |
| 70 | "--fov-deg", str(data["fov_deg"]), "--output", str(output), |
| 71 | ] |
| 72 | if args.checkpoint: |
| 73 | command += ["--checkpoint", str(args.checkpoint)] |
| 74 | if args.gemma_path: |
| 75 | command += ["--gemma-path", str(args.gemma_path)] |
| 76 | command += ["--action-overlay" if args.action_overlay else "--no-action-overlay"] |
| 77 | if "seed" in data: |
| 78 | command += ["--seed", str(data["seed"])] |
| 79 | for flag, value in (("--num-frames", args.num_frames), ("--width", args.width), |
| 80 | ("--height", args.height), ("--steps", args.steps)): |
| 81 | if value is not None: |
| 82 | command += [flag, str(value)] |
| 83 | print(" ".join(shlex.quote(part) for part in command)) |
| 84 | if not args.dry_run: |
| 85 | subprocess.run(command, check=True, cwd=ROOT) |
| 86 | |
| 87 | |
| 88 | if __name__ == "__main__": |
| 89 | main() |
| 90 |