| |
| """gc25d 离线评测（SFT 双向模型 / 基座）：val clip 的首帧 latent + parts id 视频 + 逐 cell 文本 ⇒ 81 帧，对原始 GT 算 PSNR。 |
| |
| 改自 infer.py（演示版推理）的 clip 评测路径，适配 maliang 的数据布局（SPEC §5）： |
| * 采样直接 import infer.sample（TI2V：FlowMatchScheduler("Wan")、帧 0 每步写回 z0、组合式 CFG），前向 = model.gc_forward； |
| * z0 = latents/<clip>.pt 的 cell 0（因果 VAE ⇒ 等于单独编码帧 0）； |
| * ids = cond/<clip>.npz 的 parts raw id；模型是 ids 词表 ⇒ 先 model.collapse_parts；--sem_mode unknown ⇒ 全 255； |
| * GT = 原始 tar（<raw_root>/<subset>/<scene>/<run>/frames-*.tar）里帧 [start, start+81) 的 jpg → geometry.resample_rgb |
| （不是 VAE 重建；infer.read_raw_frames 读，按 index/clips.jsonl 的 subset / scene / run / start 定位）； |
| 另把缓存 latent 解码一遍当"VAE 重建上限"（recon_ceiling，--no_recon_ceiling 关掉）； |
| * 文本 = <data_root>/text_table_gc25d.pt 查表，context_ids = 排序下标（与 dataset.key_to_id / infer._KEY_TO_ID 一致）； |
| * 模型 = --ckpt（model.load_gc_weights，词表按 ckpt 自动定）或 --base（基座 + 零初始化 gc_sem_cond ⇒ 等于基座 TI2V）。 |
| infer.load_models 不能直接用：它先按默认 "ids" 词表 install_sem_cond，parts ckpt 会报词表不符。 |
| |
| 分辨率 = geometry（env GC_RES，默认 832x480）：--data_root 的 index/prep_args.json "res" 与 --ckpt metadata 的 "res" |
| （缺失 ⇒ 832x480）都必须等于 GC_RES，否则加载前报错。 |
| |
| 输出（--out_dir）：<clip>.mp4（GT | generated | cond 色图：静态类沿用 geometry.PALETTE，玩家 / Boss 的 6 个部位各一色， |
| 255 灰色斜纹）、<clip>.json（逐帧 PSNR / SSIM、设置、耗时）、metrics.json（汇总：帧 1..80 的 PSNR 均值按 clip 平均）。 |
| |
| CUDA_VISIBLE_DEVICES=2 .venv/bin/python code/gc25d/sample.py --ckpt outputs/sft_v10/step-20000.safetensors \ |
| --val_n 8 --out_dir outputs/sft_v10/eval_step-20000 |
| CUDA_VISIBLE_DEVICES=2 .venv/bin/python code/gc25d/sample.py --base --clip_ids <id> --steps 8 --out_dir outputs/... |
| 用户模式（自带首帧 / ids / 文本，无 GT）见 sample_user.py：与这里共用 model_ids / generate（同输入同种子逐位相同； |
| 进程里先碰 VAE 再采样的调用方须先 warmup，见该函数）。 |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import os |
| import sys |
| import time |
|
|
| HERE = os.path.dirname(os.path.abspath(__file__)) |
| ROOT = os.path.dirname(os.path.dirname(HERE)) |
| os.environ.setdefault("DIFFSYNTH_SKIP_DOWNLOAD", "True") |
| os.environ.setdefault("AROPE_ROOT", os.path.join(ROOT, "code", "arope")) |
| os.environ.setdefault("TMPDIR", os.path.join(ROOT, ".tmp")) |
| for _p in (HERE, os.environ["AROPE_ROOT"]): |
| if _p not in sys.path: |
| sys.path.append(_p) |
|
|
| import numpy as np |
| import torch |
|
|
| import geometry as GEO |
| import infer as INF |
| import metrics as MET |
| import model as GM |
| import text as TXT |
| import viz as VIZ |
|
|
| MODEL_DIR = os.path.join(ROOT, "models", "Wan2.2-TI2V-5B") |
| DATA_ROOT = os.path.join(ROOT, "data", "gc25d_v10") |
| RAW_ROOT = os.path.join(ROOT, "data", "GameCrafter-Data-v10", "data") |
| F_CLIP, N_CELLS = 81, 21 |
|
|
| |
| PART_COLORS = { |
| 12: (255, 150, 255), 13: (150, 60, 255), 14: (255, 70, 160), 15: (110, 40, 190), 16: (0, 255, 255), |
| 17: (255, 150, 110), 18: (255, 120, 0), 19: (200, 0, 70), 20: (130, 20, 20), 21: (0, 90, 255), |
| } |
| PARTS_PALETTE = GEO.PALETTE.copy() |
| for _raw, _c in PART_COLORS.items(): |
| PARTS_PALETTE[_raw] = _c |
|
|
|
|
| def colorize_parts(ids: np.ndarray) -> np.ndarray: |
| """uint8 [..., H, W] → RGB：geometry.colorize_ids（静态类 + 255 斜纹）再把 12..21 的部位换成 PART_COLORS。""" |
| out = GEO.colorize_ids(ids) |
| m = (ids >= 12) & (ids <= 21) |
| if m.any(): |
| out[m] = PARTS_PALETTE[ids[m]] |
| return out |
|
|
|
|
| class PartsSideBySide(VIZ.SideBySide): |
| """viz.SideBySide，只把 ids 面板换成 colorize_parts（保留面板分隔线）。""" |
|
|
| def frame(self, i: int) -> np.ndarray: |
| out = super().frame(i) |
| x0 = (self.gen_index + 1) * VIZ.PANEL_W |
| out[VIZ.HEADER_H:, x0 + 1:x0 + VIZ.PANEL_W] = colorize_parts(self.ids[i])[:, 1:] |
| return out |
|
|
|
|
| |
| def clip_index(data_root: str) -> dict[str, dict]: |
| return {r["clip_id"]: r for r in INF.read_jsonl(os.path.join(data_root, "index", "clips.jsonl"))} |
|
|
|
|
| def select_val_clips(index: dict[str, dict], n: int) -> list[str]: |
| """val 且 drop=false 的 clip 里确定性地挑 n 条：按 run 轮流分配名额，run 内等距取。""" |
| by_run: dict[str, list[str]] = {} |
| for cid, r in index.items(): |
| if r["split"] == "val" and not r.get("drop"): |
| by_run.setdefault(r["run"], []).append(cid) |
| if not by_run: |
| raise RuntimeError("index/clips.jsonl 里没有可用的 val clip") |
| runs = sorted(by_run) |
| for r in runs: |
| by_run[r].sort() |
| n = min(n, sum(len(v) for v in by_run.values())) |
| quota, left, k = {r: 0 for r in runs}, n, 0 |
| while left > 0: |
| r = runs[k % len(runs)] |
| if quota[r] < len(by_run[r]): |
| quota[r] += 1 |
| left -= 1 |
| k += 1 |
| out = [] |
| for r in runs: |
| L, m = by_run[r], quota[r] |
| out += [L[int((j + 0.5) * len(L) / m)] for j in range(m)] |
| return out |
|
|
|
|
| def load_clip(data_root: str, cid: str) -> dict: |
| d = torch.load(os.path.join(data_root, "latents", cid + ".pt"), map_location="cpu", weights_only=False) |
| lat = d["latent"] |
| assert lat.dtype == torch.bfloat16 and tuple(lat.shape) == (48, N_CELLS, GEO.LAT_H, GEO.LAT_W), (lat.dtype, tuple(lat.shape)) |
| prompts = list(d["prompts"]) |
| assert len(prompts) == N_CELLS, len(prompts) |
| blocks = d["actions_block"] |
| blocks = np.asarray(blocks.numpy() if torch.is_tensor(blocks) else blocks).astype(np.uint8) |
| with np.load(os.path.join(data_root, "cond", cid + ".npz"), allow_pickle=False) as z: |
| ids = z["ids"] |
| assert ids.dtype == np.uint8 and ids.shape == (F_CLIP, GEO.DST_H, GEO.DST_W), (ids.dtype, ids.shape) |
| bad = sorted(set(np.unique(ids).tolist()) - GM.valid_raw_ids("parts")) |
| assert not bad, f"{cid}: ids 含 parts 词表外的值 {bad}" |
| return {"latent": lat, "prompts": prompts, "blocks": blocks, "ids": ids} |
|
|
|
|
| def gt_frames(raw_root: str, row: dict) -> np.ndarray: |
| raw = INF.read_raw_frames(row["scene"], row["run"], int(row["start"]), F_CLIP, "jpg", |
| raw_root=os.path.join(raw_root, row["subset"])) |
| return np.stack([GEO.resample_rgb(f) for f in raw], 0) |
|
|
|
|
| class TextTable: |
| """text_table_gc25d.pt 查表；context_ids = 排序下标（与训练一致）。""" |
|
|
| def __init__(self, path: str): |
| self.table = torch.load(path, map_location="cpu", weights_only=False) |
| self.key_to_id = {k: i for i, k in enumerate(sorted(self.table))} |
|
|
| def context(self, prompts: list[str], device): |
| missing = sorted({p for p in prompts if p not in self.table}) |
| if missing: |
| raise KeyError(f"文本表里没有 {missing[:3]}（gc25d 的 19 条串应该都在）") |
| ctx = torch.stack([self.table[p] for p in prompts], 0).to(torch.bfloat16) |
| ids = torch.tensor([self.key_to_id[p] for p in prompts], dtype=torch.int64) |
| return ctx.unsqueeze(0).to(device), ids.unsqueeze(0).to(device) |
|
|
|
|
| |
| def load_dit(ckpt: str | None, base: bool, vocab: str, model_dir: str, device: str): |
| if bool(ckpt) == bool(base): |
| raise ValueError("必须二选一：--ckpt <safetensors> 或 --base") |
| if ckpt: |
| GEO.check_ckpt_res(ckpt, "--ckpt") |
| t0 = time.time() |
| dit = GM.load_dit(model_dir, device=device) |
| if ckpt: |
| info = GM.load_gc_weights(dit, ckpt) |
| info = {k: v for k, v in info.items() if k != "missing"} | {"n_missing": len(info["missing"])} |
| else: |
| GM.install_sem_cond(dit, vocab) |
| info = {"vocab": vocab, "has_sem_cond": False} |
| dit = dit.to(device).eval().requires_grad_(False) |
| dit.freqs = tuple(t.to(device) for t in dit.freqs) |
| info.update(ckpt=ckpt, base=bool(base), model_dir=model_dir, res=GEO.RES_TAG, dit_load_sec=time.time() - t0, |
| sem_cond_zero=getattr(dit, GM.SEM_COND_NAME).is_zero_output()) |
| return dit, info |
|
|
|
|
| |
| def model_ids(ids: np.ndarray, vocab: str, sem_mode: str) -> np.ndarray: |
| """parts raw ids → 模型实际收到的条件：ids 词表的模型先 collapse_parts；sem_mode=unknown ⇒ 全 255。""" |
| out = GM.collapse_parts(ids) if vocab == "ids" else ids |
| return np.full_like(out, GM.UNKNOWN_ID) if sem_mode == "unknown" else out |
|
|
|
|
| def generate(args, dit, tt: TextTable, z0: torch.Tensor, ids_model: np.ndarray, prompts: list[str]): |
| """z0 [1,48,1,H/16,W/16] + 条件 + 21 条 prompt → (latents [1,48,21,H/16,W/16], 采样秒数)。两种模式走同一段代码 ⇒ 同输入同种子逐位相同。 |
| 用 args.device / steps / shift / cfg / sem_cfg / seed。""" |
| dev = args.device |
| context, context_ids = tt.context(prompts, dev) |
| neg_context = neg_ids = None |
| if args.cfg != 1.0: |
| neg_context, neg_ids = tt.context([TXT.NULL_PROMPT] * N_CELLS, dev) |
| sem = torch.from_numpy(np.ascontiguousarray(ids_model)).unsqueeze(0).to(dev) |
|
|
| torch.cuda.reset_peak_memory_stats() |
| t0 = time.time() |
| latents = INF.sample(dit, z0, context, context_ids, sem, steps=args.steps, shift=args.shift, cfg=args.cfg, |
| sem_cfg=args.sem_cfg, seed=args.seed, neg_context=neg_context, neg_context_ids=neg_ids, |
| forward=GM.gc_forward) |
| torch.cuda.synchronize() |
| return latents, time.time() - t0 |
|
|
|
|
| @torch.no_grad() |
| def warmup(args, dit, tt: TextTable, ids_model: np.ndarray, prompts: list[str]) -> float: |
| """在本进程跑任何 VAE 运算之前，用真实形状空跑一次 DiT 前向，返回秒数（clip 模式不调用）。 |
| |
| 默认 SDPA 后端（H200 上 = cuDNN attention）首次调用时建图并缓存，建出的图与进程里之前跑过的 cuDNN 运算有关： |
| 实测先做一次 fp32 VAE 编码再采样，同输入下第 0 个 block 的自注意力输出就不同（flash / mem-efficient 后端不受影响）， |
| 50 步后 latents 最大差 0.035。clip 模式第一次 DiT 前向之前没跑过 VAE；要先编码首帧的调用方（sample_user.py） |
| 先调这个，之后的 generate 才与 clip 模式逐位相同。交叉注意力的形状恒为 [21, 512, ·]（逐 cell gather），一次就够。""" |
| t0 = time.time() |
| context, context_ids = tt.context(prompts, args.device) |
| sem = torch.from_numpy(np.ascontiguousarray(ids_model)).unsqueeze(0).to(args.device) |
| lat = torch.zeros((1, 48, N_CELLS, GEO.LAT_H, GEO.LAT_W), dtype=torch.bfloat16, device=args.device) |
| GM.gc_forward(dit, lat, torch.tensor([1000.0], dtype=torch.bfloat16, device=args.device), context, sem, |
| first_frame_cond=True, context_ids=context_ids) |
| torch.cuda.synchronize() |
| return time.time() - t0 |
|
|
|
|
| def panel_names(args) -> tuple[str, str]: |
| tag = os.path.basename(args.ckpt) if args.ckpt else "base" |
| gen_name = f"generated [{tag} sem={args.sem_mode} cfg={args.cfg:g} sem_cfg={args.sem_cfg:g} {args.steps} steps seed={args.seed}]" |
| ids_name = "cond = model input (parts coloured)" if args.sem_mode == "cond" else "cond = all 255 (sem_mode=unknown)" |
| return gen_name, ids_name |
|
|
|
|
| |
| def run_clip(args, dit, vae_pipe, tt: TextTable, model_info: dict, cid: str, row: dict) -> dict: |
| dev = args.device |
| c = load_clip(args.data_root, cid) |
| gt = gt_frames(args.raw_root, row) |
| vocab = model_info["vocab"] |
| ids_model = model_ids(c["ids"], vocab, args.sem_mode) |
| z0 = c["latent"][:, 0:1].unsqueeze(0).to(dev) |
| latents, sample_sec = generate(args, dit, tt, z0, ids_model, c["prompts"]) |
| if getattr(args, "save_latents", False): |
| torch.save(latents.cpu(), os.path.join(args.out_dir, cid + "_latents.pt")) |
| frames = INF.decode_video(vae_pipe, latents) |
| info = {"clip_id": cid, "split": row.get("split"), "subset": row.get("subset"), "scene": row.get("scene"), |
| "run": row.get("run"), "start": row.get("start"), "vocab": vocab, "sem_mode": args.sem_mode, |
| "steps": args.steps, "shift": args.shift, "cfg": args.cfg, "sem_cfg": args.sem_cfg, "seed": args.seed, |
| "prompt_runs": INF.prompt_runs(c["prompts"]), "sample_sec": sample_sec, |
| "peak_mem_gb": torch.cuda.max_memory_allocated() / 1e9, "model": model_info} |
| info["metrics"] = MET.clip_metrics(frames, gt, with_ssim=not args.no_ssim) |
| if args.recon_ceiling: |
| recon = INF.decode_video(vae_pipe, c["latent"].unsqueeze(0).to(dev)) |
| rc = MET.clip_metrics(recon, gt, with_ssim=not args.no_ssim) |
| info["recon_ceiling"] = {k: v for k, v in rc.items() if k not in ("psnr", "ssim")} |
| m = info["metrics"] |
| print(f" {cid}: PSNR 帧0 {m['psnr_frame0']:.2f} · 帧1–80 均值 {m['psnr_mean']:.2f} dB" |
| + (f" · SSIM {m['ssim_mean']:.4f}" if "ssim_mean" in m else "") |
| + (f"（VAE 上限 {info['recon_ceiling']['psnr_mean']:.2f}）" if "recon_ceiling" in info else "") |
| + f" · 采样 {sample_sec:.1f}s", flush=True) |
|
|
| out_mp4 = os.path.join(args.out_dir, cid + ".mp4") |
| gen_name, ids_name = panel_names(args) |
| |
| sbs = PartsSideBySide(frames, ids_model, gt=gt, prompts=c["prompts"], blocks=c["blocks"], title=cid, |
| gen_name=gen_name, ids_name=ids_name) |
| VIZ.write_mp4((sbs.frame(i) for i in range(sbs.n_frames)), out_mp4) |
| info["out"] = out_mp4 |
| with open(os.path.join(args.out_dir, cid + ".json"), "w", encoding="utf-8") as fh: |
| json.dump(info, fh, ensure_ascii=False, indent=1, default=INF._json_default) |
| return info |
|
|
|
|
| |
| def build_parser() -> argparse.ArgumentParser: |
| ap = argparse.ArgumentParser(description="gc25d 离线评测（SFT / 基座，val clip vs 原始 GT）") |
| ap.add_argument("--ckpt", default=None, help="train.py 导出的 step-N.safetensors（831 键，词表自动识别）") |
| ap.add_argument("--base", action="store_true", help="基座 Wan2.2-TI2V-5B + 零初始化 gc_sem_cond") |
| ap.add_argument("--vocab", choices=tuple(GM.VOCABS), default="parts", help="只对 --base 有意义（--ckpt 按 ckpt 定）") |
| ap.add_argument("--model_dir", default=MODEL_DIR) |
| ap.add_argument("--data_root", default=DATA_ROOT) |
| ap.add_argument("--raw_root", default=RAW_ROOT, help="GameCrafter-Data-v10/data（下面是 <subset>/<scene>/<run>/frames-*.tar）") |
| ap.add_argument("--text_table", default=None, help="默认 <data_root>/text_table_gc25d.pt") |
| ap.add_argument("--clip_ids", default=None, help="逗号分隔，或 @文件（每行一个）") |
| ap.add_argument("--val_n", type=int, default=0, help="确定性地挑 N 条 val clip（drop=false，按 run 均分）") |
| ap.add_argument("--steps", type=int, default=50) |
| ap.add_argument("--shift", type=float, default=5.0) |
| ap.add_argument("--cfg", type=float, default=1.0, help="文本 CFG（负支 NULL_PROMPT × 21）") |
| ap.add_argument("--sem_cfg", type=float, default=1.0, help="语义 CFG（负支 ids 全 255）") |
| ap.add_argument("--sem_mode", choices=("cond", "unknown"), default="cond") |
| ap.add_argument("--seed", type=int, default=0) |
| ap.add_argument("--out_dir", required=True) |
| ap.add_argument("--no_ssim", action="store_true") |
| ap.add_argument("--no_recon_ceiling", dest="recon_ceiling", action="store_false") |
| ap.add_argument("--save_latents", action="store_true", help="另存 <clip>_latents.pt（采样结果 bf16 [1,48,21,H/16,W/16]，复现核对用）") |
| ap.add_argument("--device", default="cuda") |
| return ap |
|
|
|
|
| def main(argv=None) -> dict: |
| ap = build_parser() |
| args = ap.parse_args(argv) |
| if bool(args.ckpt) == bool(args.base): |
| ap.error("必须二选一：--ckpt 或 --base") |
| if bool(args.clip_ids) == bool(args.val_n): |
| ap.error("必须二选一：--clip_ids 或 --val_n") |
| if args.sem_mode != "cond" and args.sem_cfg != 1.0: |
| ap.error("--sem_cfg 只在 --sem_mode cond 下有意义") |
| t_all = time.time() |
| GEO.check_root_res(args.data_root, "--data_root") |
| index = clip_index(args.data_root) |
| if args.clip_ids: |
| if args.clip_ids.startswith("@"): |
| with open(args.clip_ids[1:], encoding="utf-8") as fh: |
| clips = [ln.strip() for ln in fh if ln.strip()] |
| else: |
| clips = [c.strip() for c in args.clip_ids.split(",") if c.strip()] |
| else: |
| clips = select_val_clips(index, args.val_n) |
| unknown = [c for c in clips if c not in index] |
| if unknown: |
| raise KeyError(f"{len(unknown)} 条 clip 不在 {args.data_root}/index/clips.jsonl：{unknown[:3]}") |
| |
| for cid in clips: |
| r = index[cid] |
| for p in (os.path.join(args.data_root, "latents", cid + ".pt"), os.path.join(args.data_root, "cond", cid + ".npz"), |
| os.path.join(args.raw_root, r["subset"], r["scene"], r["run"])): |
| if not os.path.exists(p): |
| raise FileNotFoundError(f"{cid}: 缺 {p}") |
| not_val = [c for c in clips if index[c]["split"] != "val"] |
| if not_val: |
| print(f" 注意：{len(not_val)} / {len(clips)} 条不是 val clip：{not_val[:3]}") |
| os.makedirs(args.out_dir, exist_ok=True) |
| tt = TextTable(args.text_table or os.path.join(args.data_root, "text_table_gc25d.pt")) |
|
|
| dit, model_info = load_dit(args.ckpt, args.base, args.vocab, args.model_dir, args.device) |
| print(f"[model] {'ckpt ' + args.ckpt if args.ckpt else 'base'}：vocab={model_info['vocab']} " |
| f"gc_sem_cond 零输出={model_info['sem_cond_zero']}（{model_info['dit_load_sec']:.0f}s）", flush=True) |
| vae_pipe = INF.load_vae(args.device, args.model_dir) |
|
|
| per_clip = [] |
| for n, cid in enumerate(clips): |
| print(f"[{n + 1}/{len(clips)}] {cid}", flush=True) |
| info = run_clip(args, dit, vae_pipe, tt, model_info, cid, index[cid]) |
| per_clip.append({"clip_id": cid, "split": info["split"], "run": info["run"], "scene": info["scene"], |
| **info["metrics"], "sample_sec": info["sample_sec"], |
| "recon_psnr_mean": info.get("recon_ceiling", {}).get("psnr_mean")}) |
| summary = { |
| "config": {k: getattr(args, k) for k in ("ckpt", "base", "steps", "shift", "cfg", "sem_cfg", "sem_mode", "seed", |
| "data_root", "raw_root", "val_n")} | {"res": GEO.RES_TAG}, |
| "model": {k: v for k, v in model_info.items() if k != "metadata"}, |
| "ckpt_metadata": model_info.get("metadata"), |
| "clip_ids": clips, |
| "aggregate": MET.aggregate(per_clip), |
| "recon_psnr_mean": (float(np.mean([c["recon_psnr_mean"] for c in per_clip])) |
| if all(c["recon_psnr_mean"] is not None for c in per_clip) else None), |
| "per_clip": per_clip, |
| "total_sec": time.time() - t_all, |
| } |
| with open(os.path.join(args.out_dir, "metrics.json"), "w", encoding="utf-8") as fh: |
| json.dump(summary, fh, ensure_ascii=False, indent=1, default=INF._json_default) |
| agg = summary["aggregate"] |
| print(f"DONE {len(clips)} 条：PSNR(帧1–80) {agg['psnr_mean']:.2f} dB" |
| + (f" · SSIM {agg['ssim_mean']:.4f}" if "ssim_mean" in agg else "") |
| + (f"（VAE 上限 {summary['recon_psnr_mean']:.2f}）" if summary["recon_psnr_mean"] is not None else "") |
| + f" → {args.out_dir}/metrics.json · {summary['total_sec']:.0f}s", flush=True) |
| return summary |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
