| |
| """双向 SFT 的泛化测试 L2 / L3（L1 = sample.py --val_n 20，已跑）。 |
| L2 见过的地图 + 没见过的首帧风格：clip 的真 ids / prompts，首帧 = FLUX restyle 成训练里没有的画风。 |
| L3 没见过的地图 + 没见过的首帧：procmap 程序生成的新世界（真相机 + 真胶囊运动）+ FLUX map2frame 生成的首帧。 |
| usage: genbench_run.py [--ckpt ...] [--gpus 4,5] [--flux_gpu 3]""" |
| import argparse |
| import json |
| import os |
| import subprocess |
| import sys |
| from concurrent.futures import ThreadPoolExecutor |
|
|
| M = "/opt/dlami/nvme/zhiyangdeng/maliang" |
| G = f"{M}/outputs/genbench" |
| PY, FPY = f"{M}/.venv/bin/python", "/opt/dlami/nvme/zhiyangdeng/PCE/envs/pce_f3/bin/python" |
| ENV = dict(os.environ, DIFFSYNTH_SKIP_DOWNLOAD="True", PYTHONPATH=f"{M}/code/gc25d:{M}/code/arope", |
| HF_HOME="/opt/dlami/nvme/zhiyangdeng/.cache/huggingface", TMPDIR=f"{M}/.tmp", OMP_NUM_THREADS="8") |
| L2_CLIPS = { |
| "train_v10": "01_Sunlit_Mediterranean_Courtyard__R10_s1501__000672", |
| "train_v9": "06_Overgrown_Tropical_Water_Temple__P9C_golden_s6942__003648", |
| "val_v10": "01_Sunlit_Mediterranean_Courtyard__R10_s1512__001152", |
| "val_v9": "02_Autumn_Woodland_Shrine__P9_golden_s2909__001152", |
| } |
| STYLES = { |
| "inkwash": "traditional Chinese ink-wash painting (shuimo) on textured rice paper: black and grey ink strokes, soft ink " |
| "bleeding, sparse muted colour washes, empty paper tones", |
| "clay": "claymation stop-motion scene: everything sculpted from coloured plasticine and modelling clay, visible " |
| "fingerprints, rounded soft shapes, studio lighting", |
| "papercraft": "low-poly papercraft diorama: every object folded from coloured cardstock with flat polygon facets, " |
| "visible paper folds and cut edges, soft paper shadows", |
| } |
|
|
|
|
| def run(cmd, log, env=ENV): |
| os.makedirs(os.path.dirname(log), exist_ok=True) |
| with open(log, "w") as fh: |
| r = subprocess.run(cmd, stdout=fh, stderr=subprocess.STDOUT, env=env, cwd=f"{M}/.tmp") |
| if r.returncode: |
| raise RuntimeError(f"失败（{r.returncode}）：{' '.join(cmd[:4])}… 见 {log}") |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--ckpt", default=f"{M}/outputs/sft_v10/step-8000.safetensors") |
| ap.add_argument("--gpus", default="4,5") |
| ap.add_argument("--flux_gpu", default="3") |
| a = ap.parse_args() |
| gpus = a.gpus.split(",") |
|
|
| |
| for tag, cid in L2_CLIPS.items(): |
| d = f"{G}/L2/{tag}/src" |
| if not os.path.exists(f"{d}/done"): |
| run([PY, f"{M}/code/gc25d/sample_user.py", "--export_clip", cid, "--out", d], |
| f"{G}/logs/export_{tag}.log", dict(ENV, CUDA_VISIBLE_DEVICES=gpus[0])) |
| open(f"{d}/done", "w").write(cid) |
| print("export ok", flush=True) |
|
|
| |
| jobs = [] |
| for tag in L2_CLIPS: |
| src = f"{G}/L2/{tag}/src" |
| f0 = [f for f in os.listdir(src) if f.endswith("frame0.png")][0] |
| for s, text in STYLES.items(): |
| jobs.append({"mode": "restyle", "image": f"{src}/{f0}", "style": text, "out": f"{G}/L2/{tag}/{s}_f0.png"}) |
| cases = sorted(d for d in os.listdir(f"{G}/procmap") if os.path.isdir(f"{G}/procmap/{d}")) |
| for c in cases: |
| jobs.append({"mode": "map2frame", "map": f"{G}/procmap/{c}/ids.npz", "frame": 0, "out": f"{G}/L3/{c}/f0.png"}) |
| todo = [j for j in jobs if not os.path.exists(j["out"])] |
| if todo: |
| for j in todo: |
| os.makedirs(os.path.dirname(j["out"]), exist_ok=True) |
| json.dump(todo, open(f"{G}/flux_jobs.json", "w"), indent=1) |
| run([FPY, f"{M}/code/gc25d/firstframe_flux.py", "batch", "--jobs", f"{G}/flux_jobs.json"], f"{G}/logs/flux.log", |
| dict(ENV, CUDA_VISIBLE_DEVICES=a.flux_gpu, PYTHONPATH=f"{M}/code/gc25d")) |
| print("flux ok", flush=True) |
|
|
| |
| runs = [] |
| for tag in L2_CLIPS: |
| src = f"{G}/L2/{tag}/src" |
| ids = [f for f in os.listdir(src) if f.endswith("ids.npz")][0] |
| pr = [f for f in os.listdir(src) if f.endswith("prompts.json")][0] |
| f0 = [f for f in os.listdir(src) if f.endswith("frame0.png")][0] |
| firsts = {s: f"{G}/L2/{tag}/{s}_f0.png" for s in STYLES} |
| if tag.startswith("train"): |
| firsts["orig"] = f"{src}/{f0}" |
| for s, ff in firsts.items(): |
| runs.append(dict(name=f"{tag}_{s}", first=ff, ids=f"{src}/{ids}", prompts=f"{src}/{pr}", out=f"{G}/L2/{tag}")) |
| for c in cases: |
| p = f"{G}/procmap/{c}" |
| runs.append(dict(name=c, first=f"{G}/L3/{c}/f0.png", ids=f"{p}/ids.npz", prompts=f"{p}/prompts.json", out=f"{G}/L3/{c}")) |
|
|
| def gen(i_r): |
| i, r = i_r |
| if os.path.exists(f"{r['out']}/{r['name']}.mp4"): |
| return r["name"] |
| run([PY, f"{M}/code/gc25d/sample_user.py", "--ckpt", a.ckpt, "--first_frame", r["first"], "--ids", r["ids"], |
| "--prompts", r["prompts"], "--out", r["out"], "--name", r["name"], "--thumb"], |
| f"{G}/logs/gen_{r['name']}.log", dict(ENV, CUDA_VISIBLE_DEVICES=gpus[i % len(gpus)])) |
| return r["name"] |
|
|
| |
| queues = [[(i, r) for i, r in enumerate(runs) if i % len(gpus) == g] for g in range(len(gpus))] |
| with ThreadPoolExecutor(len(gpus)) as ex: |
| for done in ex.map(lambda q: [gen(x) for x in q], queues): |
| print("gen ok", done, flush=True) |
| print("ALL DONE", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
