| 1 | #!/usr/bin/env python3 |
| 2 | """Run a checked-in WM case with causal 4-step inference.""" |
| 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_causal_cases" |
| 15 | DEFAULT_CHECKPOINT = ROOT / "checkpoints" / "echo-wm-flash.safetensors" |
| 16 | |
| 17 | |
| 18 | def _case_dirs() -> list[Path]: |
| 19 | return sorted(path for path in CASES.iterdir() if path.is_dir() and (path / "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 | missing = [key for key in ("prompt", "action") if key not in data] |
| 25 | if missing: |
| 26 | raise ValueError(f"{case_dir}/case.json missing fields: {', '.join(missing)}") |
| 27 | data["image"] = data.get("image", "input.jpg") |
| 28 | image = case_dir / data["image"] |
| 29 | if not image.is_file(): |
| 30 | raise FileNotFoundError(f"Case image not found: {image}") |
| 31 | return data |
| 32 | |
| 33 | |
| 34 | def _num_frames_from_action(action: str) -> int: |
| 35 | """Infer output length from Action DSL durations, including the first frame.""" |
| 36 | segments = "".join(action.replace("，", ",").split()).split(",") |
| 37 | try: |
| 38 | durations = [int(segment.rsplit("-", 1)[1]) for segment in segments] |
| 39 | except (IndexError, ValueError) as error: |
| 40 | raise ValueError(f"Cannot infer frame count from action {action!r}") from error |
| 41 | if not durations or any(duration <= 0 for duration in durations): |
| 42 | raise ValueError(f"Action durations must be positive: {action!r}") |
| 43 | return 1 + sum(durations) |
| 44 | |
| 45 | |
| 46 | def build_command(args: argparse.Namespace) -> list[str]: |
| 47 | case_dir = (args.case if args.case.is_absolute() else ROOT / args.case).resolve() |
| 48 | data = _load_case(case_dir) |
| 49 | image = case_dir / data["image"] |
| 50 | num_frames = args.num_frames if args.num_frames is not None else _num_frames_from_action(data["action"]) |
| 51 | command = [ |
| 52 | sys.executable, str(ROOT / "inference_wm_causal.py"), |
| 53 | "--image", str(image), "--prompt", data["prompt"], |
| 54 | "--action-str", data["action"], |
| 55 | "--num-frames", str(num_frames), |
| 56 | "--fov-deg", str(data.get("fov_deg", 70.0)), |
| 57 | "--output", str(args.output_dir / case_dir.name / "result.mp4"), |
| 58 | "--checkpoint", str(args.checkpoint), |
| 59 | ] |
| 60 | if args.gemma_path: |
| 61 | command += ["--gemma-path", str(args.gemma_path)] |
| 62 | for flag, value in ( |
| 63 | ("--width", args.width), |
| 64 | ("--height", args.height), |
| 65 | ("--video-local-attn-size", args.video_local_attn_size), |
| 66 | ("--video-sink-size", args.video_sink_size), |
| 67 | ("--video-chunk-size", args.video_chunk_size), |
| 68 | ): |
| 69 | if value is not None: |
| 70 | command += [flag, str(value)] |
| 71 | command += ["--action-overlay" if args.action_overlay else "--no-action-overlay"] |
| 72 | if "seed" in data: |
| 73 | command += ["--seed", str(data["seed"])] |
| 74 | return command |
| 75 | |
| 76 | |
| 77 | def parse_args(argv: list[str] | None = None) -> argparse.Namespace: |
| 78 | parser = argparse.ArgumentParser(description=__doc__) |
| 79 | parser.add_argument("--case", type=Path) |
| 80 | parser.add_argument("--checkpoint", type=Path, default=DEFAULT_CHECKPOINT) |
| 81 | parser.add_argument("--gemma-path", type=Path) |
| 82 | parser.add_argument("--output-dir", type=Path, default=ROOT / "outputs" / "wm_cases_causal") |
| 83 | parser.add_argument("--dry-run", action="store_true") |
| 84 | parser.add_argument("--list", action="store_true") |
| 85 | parser.add_argument("--num-frames", type=int) |
| 86 | parser.add_argument("--width", type=int) |
| 87 | parser.add_argument("--height", type=int) |
| 88 | parser.add_argument( |
| 89 | "--video-local-attn-size", "--video_local_attn_size", |
| 90 | dest="video_local_attn_size", type=int, |
| 91 | ) |
| 92 | parser.add_argument( |
| 93 | "--video-sink-size", "--video_sink_size", dest="video_sink_size", type=int, |
| 94 | ) |
| 95 | parser.add_argument( |
| 96 | "--video-chunk-size", "--video_chunk_size", dest="video_chunk_size", type=int, |
| 97 | ) |
| 98 | parser.add_argument( |
| 99 | "--action-overlay", action=argparse.BooleanOptionalAction, default=True, |
| 100 | help="Also write the HUD copy (default: enabled).", |
| 101 | ) |
| 102 | return parser.parse_args(argv) |
| 103 | |
| 104 | |
| 105 | def main() -> None: |
| 106 | args = parse_args() |
| 107 | if args.list: |
| 108 | for case_dir in _case_dirs(): |
| 109 | print(f"{case_dir.name}\t{_load_case(case_dir).get('description', 'causal I2V case')}") |
| 110 | return |
| 111 | if args.case is None: |
| 112 | raise SystemExit("--case is required unless --list is used") |
| 113 | command = build_command(args) |
| 114 | print(" ".join(shlex.quote(part) for part in command)) |
| 115 | if not args.dry_run: |
| 116 | subprocess.run(command, check=True, cwd=ROOT) |
| 117 | |
| 118 | |
| 119 | if __name__ == "__main__": |
| 120 | main() |
| 121 |