| |
| """gc25d 离线采样 · 用户模式：自带首帧 + parts id 视频 + 逐 cell 文本 ⇒ 81 帧 W×H（无 GT，不出 PSNR）。泛化测试（genbench）用。 |
| W×H = geometry.DST_W × DST_H（env GC_RES，默认 832×480）；--ckpt 的 metadata |
| "res"（缺失 ⇒ 832x480）必须等于 GC_RES，--export_clip 的 --data_root 同理。 |
| |
| 薄封装，复用 sample.py（load_dit / TextTable / model_ids / generate / PartsSideBySide）与 infer.py（load_image_rgb / |
| parse_actions / encode_first_frame / decode_video / prompt_runs）。采样核心 = sample.generate，与 clip 模式是同一段代码： |
| 同首帧 latent + 同 ids + 同 prompts + 同种子 ⇒ latents 逐位相同。首帧用 fp32 VAE 单帧编码再转 bf16（infer.encode_first_frame）， |
| 与 prep 存的 latent cell 0 逐位相同（verify_prep 抽查 cell0_bf16_exact_frac = 1.0），所以 clip 自己的帧 0 png 能逐位复现 clip 模式。 |
| 注意顺序：先 sample.warmup（真实形状空跑一次 DiT 前向）再编码首帧 —— cuDNN attention 的图在首次调用时建，建图结果受之前 |
| 跑过的 VAE 运算影响（不 warmup 时实测 latents 最大差 0.035），warmup 后与 clip 模式逐位相同（outputs/genbench/usermode_test）。 |
| |
| 输入 |
| --first_frame png/jpg：W×H 原样；1280×704（源分辨率）→ geometry.resample_rgb；其它尺寸居中裁成 W:H 再 INTER_AREA。 |
| --ids npz（'ids' 键或唯一数组）/ npy，uint8 [81,H,W] parts raw id，或源分辨率 [81,704,1280]（geometry.resample_ids）。 |
| 取值必须 ⊆ model.valid_raw_ids(ckpt 词表)（ids 词表的 ckpt 先 collapse_parts，与 clip 模式相同），否则报错。 |
| --prompts json：21 条串（cell 0..20，须都在文本表里）；或 |
| --actions "<spec>"（语法见 infer.py 模块文档，20 块，如 "D:10,A:10"）/ [20,5] 的 .npy → text.cell_prompts。 |
| --export_clip <clip_id>：只导出该 clip 的 帧 0（W×H + 源 1280×704 png）/ ids npz / prompts json / 动作块 npy 到 --out， |
| 不加载模型（L2 重绘首帧、复现测试的素材）。 |
| |
| 输出（--out 目录，前缀 --name）：<name>.mp4（generated | cond 色图，--thumb 在生成面板右上角嵌输入首帧缩略图）、 |
| <name>_gen.mp4（纯生成）、<name>_sheet.jpg（帧 0/20/40/60/80）、<name>.json（设置 / 输入 / 条件类别占比 / 耗时）、 |
| --save_latents ⇒ <name>_latents.pt（bf16 [1,48,21,H/16,W/16]）。 |
| |
| CUDA_VISIBLE_DEVICES=5 .venv/bin/python code/gc25d/sample_user.py --ckpt outputs/sft_v10/step-8000.safetensors \ |
| --first_frame f0.png --ids ids.npz --actions "D:10,A:10" --out outputs/genbench/x --name d10a10 --thumb |
| .venv/bin/python code/gc25d/sample_user.py --export_clip <clip_id> --out outputs/genbench/x |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import os |
| import struct |
| import time |
|
|
| import sample as S |
|
|
| import cv2 |
| import numpy as np |
| import torch |
| from PIL import Image |
|
|
| import geometry as GEO |
| import infer as INF |
| import model as GM |
| import text as TXT |
| import viz as VIZ |
|
|
| F_CLIP, N_CELLS, H, W = S.F_CLIP, S.N_CELLS, GEO.DST_H, GEO.DST_W |
| STATIC_NAMES = ("sky", "terrain", "paving", "architecture", "rock", "vegetation", "prop") |
| CLASS_NAMES = {**dict(enumerate(STATIC_NAMES)), |
| **{r: f"player_{n}" for r, n in zip(GM.PLAYER_PART_IDS, GM.PART_NAMES)}, |
| **{r: f"boss_{n}" for r, n in zip(GM.BOSS_PART_IDS, GM.PART_NAMES)}, GM.UNKNOWN_ID: "unknown"} |
| THUMB_W, THUMB_H = W // 4, H // 4 |
|
|
|
|
| |
| def ckpt_vocab(path: str) -> str: |
| """只读 safetensors 头（不碰 10 GB 数据）→ gc_sem_cond 的词表；加载模型前就能校验 ids。""" |
| with open(path, "rb") as fh: |
| hdr = json.loads(fh.read(struct.unpack("<Q", fh.read(8))[0])) |
| v = GM.vocab_from_state(hdr) |
| if v is None: |
| raise ValueError(f"{path} 里没有 gc_sem_cond（不是 gc25d 导出？）") |
| return v |
|
|
|
|
| def load_first_frame(path: str) -> tuple[np.ndarray, dict]: |
| with Image.open(path) as im: |
| w, h = im.size |
| mode = "as_is" if (h, w) == (H, W) else "resample_rgb" if (h, w) == (GEO.SRC_H, GEO.SRC_W) else "center_crop_area" |
| return INF.load_image_rgb(path), {"path": os.path.abspath(path), "src_size": [w, h], "resize": mode} |
|
|
|
|
| def load_ids(path: str, vocab: str) -> tuple[np.ndarray, dict]: |
| """→ (模型词表下的 raw ids uint8 [81,H,W], info)。取值严格 ⊆ valid_raw_ids(vocab)，不静默映射。""" |
| if path.endswith(".npz"): |
| with np.load(path, allow_pickle=False) as z: |
| keys = list(z.keys()) |
| if "ids" not in keys and len(keys) != 1: |
| raise ValueError(f"{path}：npz 里没有 'ids' 键且不止一个数组：{keys}") |
| arr = z["ids" if "ids" in keys else keys[0]] |
| elif path.endswith(".npy"): |
| arr = np.load(path, allow_pickle=False) |
| else: |
| raise ValueError(f"--ids 应为 .npz / .npy，得到 {path}") |
| if arr.dtype != np.uint8: |
| raise ValueError(f"{path}：ids 应为 uint8，得到 {arr.dtype}") |
| info = {"path": os.path.abspath(path), "src_shape": list(arr.shape), "resampled": False} |
| if arr.shape == (F_CLIP, GEO.SRC_H, GEO.SRC_W) and not GEO.IDENTITY: |
| arr, info["resampled"] = GEO.resample_ids(arr), True |
| if arr.shape != (F_CLIP, H, W): |
| raise ValueError(f"{path}：ids 应为 [81,{H},{W}]（或源分辨率 [81,704,1280]；{GEO.RES_ENV}={GEO.RES_TAG}），得到 {arr.shape}") |
| if vocab == "ids": |
| arr = GM.collapse_parts(arr) |
| cnt = np.bincount(arr.ravel(), minlength=256) |
| bad = {int(k): int(cnt[k]) for k in np.nonzero(cnt)[0] if int(k) not in GM.valid_raw_ids(vocab)} |
| if bad: |
| raise ValueError(f"{path}：ids 含 {vocab!r} 词表外的值 {{值: 像素数}} = {bad}" |
| f"（允许 {sorted(GM.valid_raw_ids(vocab))}；未知区域请写 255）") |
| return np.ascontiguousarray(arr), info |
|
|
|
|
| def load_prompts(args, tt: S.TextTable) -> tuple[list[str], np.ndarray | None, dict]: |
| """--prompts json（21 条）或 --actions → (prompts, cell 顺序动作块 [21,5] 或 None, info)。""" |
| if args.prompts: |
| with open(args.prompts, encoding="utf-8") as fh: |
| prompts = json.load(fh) |
| if not (isinstance(prompts, list) and len(prompts) == N_CELLS and all(isinstance(p, str) for p in prompts)): |
| raise ValueError(f"{args.prompts}：应为 {N_CELLS} 条字符串的 json 列表") |
| blocks, info = None, {"source": "prompts_json", "path": os.path.abspath(args.prompts)} |
| else: |
| blocks20 = INF.parse_actions(args.actions) |
| blocks, prompts = TXT.cell_blocks(blocks20), TXT.cell_prompts(blocks20) |
| info = {"source": "actions", "spec": args.actions} |
| missing = sorted({p for p in prompts if p not in tt.table}) |
| if missing: |
| raise KeyError(f"文本表里没有 {missing[:3]}（只能用 text.py 模板的 18 条 + NULL_PROMPT）") |
| return prompts, blocks, info |
|
|
|
|
| def cond_stats(ids: np.ndarray) -> dict: |
| """条件（模型收到的）类别占比：全部帧 / 帧 0；UNKNOWN 逐帧占比的 min / mean / max。""" |
| def frac(a): |
| cnt = np.bincount(a.ravel(), minlength=256) |
| return {CLASS_NAMES.get(int(k), str(int(k))): float(cnt[k]) / a.size for k in np.nonzero(cnt)[0]} |
| unk = (ids == GM.UNKNOWN_ID).mean(axis=(1, 2)) |
| return {"class_frac": frac(ids), "class_frac_frame0": frac(ids[0]), |
| "unknown_frac_per_frame": {"min": float(unk.min()), "mean": float(unk.mean()), "max": float(unk.max())}} |
|
|
|
|
| |
| def export_clip(args) -> dict: |
| GEO.check_root_res(args.data_root, "--data_root") |
| row = S.clip_index(args.data_root).get(args.export_clip) |
| if row is None: |
| raise KeyError(f"{args.export_clip} 不在 {args.data_root}/index/clips.jsonl") |
| c = S.load_clip(args.data_root, args.export_clip) |
| src0 = INF.read_raw_frames(row["scene"], row["run"], int(row["start"]), 1, "jpg", |
| raw_root=os.path.join(args.raw_root, row["subset"]))[0] |
| os.makedirs(args.out, exist_ok=True) |
| p = {k: os.path.join(args.out, f"{args.export_clip}_{k}") for k in |
| ("frame0.png", "frame0_src.png", "ids.npz", "prompts.json", "actions.npy")} |
| for key, img in (("frame0.png", GEO.resample_rgb(src0)), ("frame0_src.png", src0)): |
| assert cv2.imwrite(p[key], cv2.cvtColor(img, cv2.COLOR_RGB2BGR)), p[key] |
| np.savez_compressed(p["ids.npz"], ids=c["ids"]) |
| with open(p["prompts.json"], "w", encoding="utf-8") as fh: |
| json.dump(c["prompts"], fh, ensure_ascii=False, indent=1) |
| np.save(p["actions.npy"], c["blocks"][1:]) |
| print(f"导出 {args.export_clip}（split={row['split']}）→ " + ", ".join(os.path.basename(v) for v in p.values())) |
| return p |
|
|
|
|
| |
| class UserSideBySide(S.PartsSideBySide): |
| """generated | cond（parts 着色）；thumb ≠ None ⇒ 生成面板右上角嵌输入首帧缩略图。""" |
|
|
| def __init__(self, *a, thumb: np.ndarray | None = None, **kw): |
| super().__init__(*a, **kw) |
| self.thumb = thumb |
|
|
| def frame(self, i: int) -> np.ndarray: |
| out = super().frame(i) |
| if self.thumb is not None: |
| th, tw = self.thumb.shape[:2] |
| x0, y0 = (self.gen_index + 1) * VIZ.PANEL_W - tw - 6, VIZ.HEADER_H + 6 |
| out[y0:y0 + th, x0:x0 + tw] = self.thumb |
| return out |
|
|
|
|
| def make_thumb(img: np.ndarray) -> np.ndarray: |
| t = cv2.resize(img, (THUMB_W, THUMB_H), interpolation=cv2.INTER_AREA) |
| cv2.rectangle(t, (0, 0), (THUMB_W - 1, THUMB_H - 1), VIZ.ACCENT, 1) |
| VIZ._put(t, "input frame 0", 4, THUMB_H - 6, VIZ.ACCENT, 0.4) |
| return t |
|
|
|
|
| |
| def build_parser() -> argparse.ArgumentParser: |
| ap = argparse.ArgumentParser(description="gc25d 离线采样 · 用户模式（自带首帧 / ids / 文本）") |
| ap.add_argument("--ckpt", default=None) |
| ap.add_argument("--base", action="store_true") |
| ap.add_argument("--vocab", choices=tuple(GM.VOCABS), default="parts", help="只对 --base 有意义（--ckpt 按 ckpt 定）") |
| ap.add_argument("--model_dir", default=S.MODEL_DIR) |
| ap.add_argument("--data_root", default=S.DATA_ROOT) |
| ap.add_argument("--raw_root", default=S.RAW_ROOT, help="只有 --export_clip 用") |
| ap.add_argument("--text_table", default=None, help="默认 <data_root>/text_table_gc25d.pt") |
| ap.add_argument("--first_frame", default=None) |
| ap.add_argument("--ids", default=None) |
| ap.add_argument("--prompts", default=None, help="json：21 条 cell prompt") |
| ap.add_argument("--actions", default=None, help='动作脚本（如 "D:10,A:10"）或 [20,5] .npy') |
| ap.add_argument("--export_clip", default=None, help="只导出某 clip 的帧 0 / ids / prompts / 动作块，然后退出") |
| ap.add_argument("--out", required=True, help="输出目录") |
| ap.add_argument("--name", default=None, help="输出文件前缀（默认 = 首帧文件名）") |
| ap.add_argument("--title", default=None, help="视频标题栏右侧文字（默认 = name）") |
| 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("--thumb", action="store_true", help="生成面板右上角嵌输入首帧缩略图") |
| ap.add_argument("--save_latents", action="store_true") |
| ap.add_argument("--device", default="cuda") |
| return ap |
|
|
|
|
| def main(argv=None) -> dict: |
| ap = build_parser() |
| args = ap.parse_args(argv) |
| if args.export_clip: |
| return export_clip(args) |
| if bool(args.ckpt) == bool(args.base): |
| ap.error("必须二选一：--ckpt 或 --base") |
| if not (args.first_frame and args.ids): |
| ap.error("用户模式需要 --first_frame 和 --ids") |
| if bool(args.prompts) == bool(args.actions): |
| ap.error("必须二选一：--prompts <json> 或 --actions <spec>") |
| if args.sem_mode != "cond" and args.sem_cfg != 1.0: |
| ap.error("--sem_cfg 只在 --sem_mode cond 下有意义") |
| t_all = time.time() |
| name = args.name or os.path.splitext(os.path.basename(args.first_frame))[0] |
| stem = os.path.join(args.out, name) |
|
|
| |
| if args.ckpt: |
| GEO.check_ckpt_res(args.ckpt, "--ckpt") |
| vocab = ckpt_vocab(args.ckpt) if args.ckpt else args.vocab |
| tt = S.TextTable(args.text_table or os.path.join(args.data_root, "text_table_gc25d.pt")) |
| img, ff_info = load_first_frame(args.first_frame) |
| ids, ids_info = load_ids(args.ids, vocab) |
| prompts, blocks, p_info = load_prompts(args, tt) |
| ids_model = S.model_ids(ids, vocab, args.sem_mode) |
| os.makedirs(args.out, exist_ok=True) |
| print(f"[{name}] 首帧 {ff_info['src_size']} ({ff_info['resize']}) · ids {ids_info['src_shape']} · vocab={vocab} · " |
| f"文本（{p_info['source']}）：\n " + "\n ".join(INF.prompt_runs(prompts)), flush=True) |
|
|
| timing = {} |
| dit, model_info = S.load_dit(args.ckpt, args.base, vocab, args.model_dir, args.device) |
| assert model_info["vocab"] == vocab, (model_info["vocab"], vocab) |
| t0 = time.time() |
| vae_pipe = INF.load_vae(args.device, args.model_dir) |
| timing.update(dit_load_sec=model_info["dit_load_sec"], vae_load_sec=time.time() - t0) |
| print(f"[model] {args.ckpt or 'base'}：vocab={vocab}（DiT {timing['dit_load_sec']:.0f}s · VAE {timing['vae_load_sec']:.0f}s）", |
| flush=True) |
|
|
| timing["warmup_sec"] = S.warmup(args, dit, tt, ids_model, prompts) |
| t0 = time.time() |
| z0 = INF.encode_first_frame(vae_pipe, img) |
| torch.cuda.synchronize() |
| timing["encode_sec"] = time.time() - t0 |
| latents, timing["sample_sec"] = S.generate(args, dit, tt, z0, ids_model, prompts) |
| if args.save_latents: |
| torch.save(latents.cpu(), stem + "_latents.pt") |
| t0 = time.time() |
| frames = INF.decode_video(vae_pipe, latents) |
| timing["decode_sec"] = time.time() - t0 |
| peak = torch.cuda.max_memory_allocated() / 1e9 |
|
|
| t0 = time.time() |
| gen_name, ids_name = S.panel_names(args) |
| sbs = UserSideBySide(frames, ids_model, prompts=prompts, blocks=blocks, title=args.title or name, |
| gen_name=gen_name, ids_name=ids_name, thumb=make_thumb(img) if args.thumb else None) |
| VIZ.write_mp4((sbs.frame(i) for i in range(sbs.n_frames)), stem + ".mp4") |
| VIZ.write_video(stem + "_gen.mp4", frames) |
| VIZ.contact_sheet(stem + "_sheet.jpg", sbs) |
| timing["write_sec"] = time.time() - t0 |
| timing["total_sec"] = time.time() - t_all |
|
|
| info = {"name": name, "mode": "user", "res": GEO.RES_TAG, "vocab": vocab, "sem_mode": args.sem_mode, "steps": args.steps, "shift": args.shift, |
| "cfg": args.cfg, "sem_cfg": args.sem_cfg, "seed": args.seed, |
| "inputs": {"first_frame": ff_info, "ids": ids_info, "text": p_info}, |
| "prompt_runs": INF.prompt_runs(prompts), "prompts": prompts, |
| "actions_block": None if blocks is None else blocks.tolist(), |
| "cond": cond_stats(ids_model), "ids_supplied": cond_stats(ids) if args.sem_mode != "cond" else "= cond", |
| "timing": timing, "peak_mem_gb": peak, "model": model_info, |
| "out": {"mp4": stem + ".mp4", "gen_mp4": stem + "_gen.mp4", "sheet": stem + "_sheet.jpg", |
| "latents": stem + "_latents.pt" if args.save_latents else None}} |
| with open(stem + ".json", "w", encoding="utf-8") as fh: |
| json.dump(info, fh, ensure_ascii=False, indent=1, default=INF._json_default) |
| print(f"DONE {stem}.mp4 · warmup {timing['warmup_sec']:.1f}s · 编码 {timing['encode_sec']:.1f}s · 采样 {timing['sample_sec']:.1f}s · 解码 {timing['decode_sec']:.1f}s" |
| f" · 写 {timing['write_sec']:.1f}s · 总 {timing['total_sec']:.0f}s · 峰值显存 {peak:.1f} GB", flush=True) |
| return info |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
