File size: 7,502 Bytes
eafbe80
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
#!/usr/bin/env python3
"""Batch multiview revisit from a first-frame list."""
from __future__ import annotations

import argparse
import csv
import json
import os
import subprocess
import sys
from typing import Dict, List

COMBO_CHUNK_FILES = (
    "chunk0_rotate_left_45.json",
    "chunk1_translate_forward.json",
    "chunk2_rotate_right_45.json",
    "chunk3_translate_backward.json",
)


def _validate_prereqs(args: argparse.Namespace, runner: str) -> None:
    if not os.path.isfile(args.ckpt):
        raise FileNotFoundError(f"[run_multiview_revisit] CKPT 不是文件: {args.ckpt}")
    if not os.path.isfile(args.firstframe_list):
        raise FileNotFoundError(f"[run_multiview_revisit] firstframe_list 不存在: {args.firstframe_list}")
    if not os.path.isdir(args.action_combo_dir):
        raise FileNotFoundError(f"[run_multiview_revisit] action_combo_dir 不是目录: {args.action_combo_dir}")
    for fn in COMBO_CHUNK_FILES:
        p = os.path.join(args.action_combo_dir, fn)
        if not os.path.isfile(p):
            raise FileNotFoundError(f"[run_multiview_revisit] 缺少 combo 动作文件(请先 build_action_combo): {p}")
    if not os.path.isfile(runner):
        raise FileNotFoundError(f"[run_multiview_revisit] runner 不存在: {runner}")


def _load_items(path: str) -> List[Dict[str, str]]:
    ext = os.path.splitext(path)[1].lower()
    items: List[Dict[str, str]] = []
    if ext == ".jsonl":
        with open(path, "r", encoding="utf-8") as f:
            for ln in f:
                ln = ln.strip()
                if not ln:
                    continue
                d = json.loads(ln)
                items.append(
                    {
                        "view_id": str(d.get("view_id") or len(items)),
                        "first_frame_image": str(d.get("first_frame_image") or d.get("image") or ""),
                        "prompt": str(d.get("prompt") or "A scene."),
                    }
                )
        return items
    if ext == ".csv":
        with open(path, "r", encoding="utf-8") as f:
            for row in csv.DictReader(f):
                items.append(
                    {
                        "view_id": str(row.get("view_id") or len(items)),
                        "first_frame_image": str(row.get("first_frame_image") or row.get("image") or ""),
                        "prompt": str(row.get("prompt") or "A scene."),
                    }
                )
        return items
    # txt: each line -> image_path[tab prompt]
    with open(path, "r", encoding="utf-8") as f:
        for i, ln in enumerate(f):
            ln = ln.strip()
            if not ln or ln.startswith("#"):
                continue
            parts = ln.split("\t", 1)
            items.append(
                {
                    "view_id": str(i),
                    "first_frame_image": parts[0],
                    "prompt": parts[1] if len(parts) > 1 else "A scene.",
                }
            )
    return items


def main() -> int:
    ap = argparse.ArgumentParser(description="Run combo revisit for a list of edited first frames")
    ap.add_argument("--ckpt", required=True)
    ap.add_argument("--firstframe_list", required=True, help="txt/csv/jsonl")
    ap.add_argument("--action_combo_dir", required=True)
    ap.add_argument("--output_root", required=True)
    ap.add_argument("--runner", default=None, help="default: eval_v2/static/run_combo_revisit_fixed_first.py")
    ap.add_argument("--chunk_frames", type=int, default=81)
    ap.add_argument("--context_frames", type=int, default=1)
    ap.add_argument("--sigma_shift", type=float, default=5.0)
    ap.add_argument("--num_inference_steps", type=int, default=50)
    ap.add_argument("--cfg_scale", type=float, default=5.0)
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument(
        "--camera_inject_mode",
        type=str,
        default=None,
        help="与 evals_ep0 一致;默认不传则由子进程读环境 CAMERA_INJECT_MODE",
    )
    # REMAINDER: 子进程参数若以 - 开头,nargs='*' 会被 argparse 误当作本脚本的选项而报错
    ap.add_argument(
        "--extra_args",
        nargs=argparse.REMAINDER,
        default=[],
        help="传给 run_combo_revisit_fixed_first.py 的额外参数;须放在命令行最后(如 MEM_ARGS)",
    )
    args = ap.parse_args()
    # 允许用户写「占位」-- 与 shell 的 -- 一致
    extra = list(args.extra_args or [])
    if extra and extra[0] == "--":
        extra = extra[1:]
    args.extra_args = extra

    out_root = os.path.abspath(args.output_root)
    if args.runner:
        runner = os.path.abspath(args.runner)
    else:
        here = os.path.dirname(os.path.abspath(__file__))
        runner = os.path.join(here, "run_combo_revisit_fixed_first.py")

    _validate_prereqs(args, runner)

    os.makedirs(out_root, exist_ok=True)
    items = _load_items(os.path.abspath(args.firstframe_list))
    if not items:
        print("[run_multiview_revisit] no entries found")
        return 0

    missing_ff = [
        str(it.get("first_frame_image") or "")
        for it in items
        if not os.path.isfile(str(it.get("first_frame_image") or ""))
    ]
    if missing_ff:
        raise FileNotFoundError(
            f"[run_multiview_revisit] 首帧图不存在: {missing_ff[:5]}{'...' if len(missing_ff) > 5 else ''}"
        )

    summary: List[Dict[str, object]] = []
    for i, it in enumerate(items):
        view_id = str(it["view_id"])
        ff = str(it["first_frame_image"])
        prompt = str(it["prompt"])
        out_dir = os.path.join(out_root, f"view_{view_id}")
        os.makedirs(out_dir, exist_ok=True)
        cmd = [
            sys.executable,
            runner,
            "--ckpt",
            args.ckpt,
        ]
        if (args.camera_inject_mode or "").strip():
            cmd.extend(["--camera_inject_mode", str(args.camera_inject_mode).strip()])
        cmd.extend(
            [
            "--first_frame_image",
            ff,
            "--output_dir",
            out_dir,
            "--prompt",
            prompt,
            "--action_combo_dir",
            args.action_combo_dir,
            "--context_frames",
            str(args.context_frames),
            "--chunk_frames",
            str(args.chunk_frames),
            "--sigma_shift",
            str(args.sigma_shift),
            "--num_inference_steps",
            str(args.num_inference_steps),
            "--cfg_scale",
            str(args.cfg_scale),
            "--seed",
            str(args.seed + i),
            "--no_camera_encoder_separate_t_r",
            ]
        )
        cmd.extend(list(args.extra_args))
        rc = subprocess.run(cmd, check=False).returncode
        ok = rc == 0
        summary.append(
            {
                "view_id": view_id,
                "first_frame_image": ff,
                "prompt": prompt,
                "ok": ok,
                "output_dir": out_dir,
            }
        )

    out = {
        "ckpt": args.ckpt,
        "firstframe_list": os.path.abspath(args.firstframe_list),
        "num_items": len(items),
        "summary": summary,
    }
    with open(os.path.join(out_root, "multiview_revisit_summary.json"), "w", encoding="utf-8") as f:
        json.dump(out, f, indent=2)
    print(f"[run_multiview_revisit] wrote {out_root}/multiview_revisit_summary.json")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())