| |
| """分辨率决策的两个快测：(1) 角色区域的 VAE 重建上限：832×480 / 960×544 / 1280×704（都在源分辨率上和 GT 比）； |
| (2) DiT 前向耗时：21 cell 双向前向在三种分辨率下的相对开销（实时可行性的粗估）。""" |
| import json |
| import os |
| import sys |
| import time |
|
|
| import cv2 |
| import numpy as np |
| import torch |
|
|
| import geometry as GEO |
| import infer as INF |
| import model as GM |
| import sample as S |
|
|
| M = "/opt/dlami/nvme/zhiyangdeng/maliang" |
| D = f"{M}/outputs/char_diag" |
| RES = {"832x480": (832, 480), "960x544": (960, 544), "1024x576": (1024, 576), "1152x640": (1152, 640), "1280x704": (1280, 704)} |
| NF = 41 |
|
|
|
|
| def enc(vae, frames): |
| """infer.encode_video 的任意分辨率版（fp32 编码，非 tiled）。""" |
| x = torch.from_numpy(frames).permute(3, 0, 1, 2).float().div_(127.5).sub_(1.0) |
| with torch.no_grad(), INF._vae_fp32(vae, True): |
| z = vae.vae.encode([x], device=vae.device, tiled=False) |
| return z.to(dtype=torch.bfloat16, device=vae.device) |
|
|
|
|
| def dec(vae, z): |
| with torch.no_grad(): |
| v = vae.vae.decode(z.to(vae.torch_dtype), device=vae.device, tiled=False) |
| return v[0].float().add_(1.0).mul_(127.5).round_().clamp_(0, 255).to(torch.uint8).permute(1, 2, 3, 0).cpu().numpy() |
|
|
|
|
| def to_res(src, wh): |
| """源 1280×704 → 目标分辨率：832/960 先按 gc25d 约定裁成目标宽高比再 INTER_AREA；1280×704 原样。""" |
| W, H = wh |
| if (W, H) == (1280, 704): |
| return src |
| cw = int(round(704 * W / H)) |
| x0 = (1280 - cw) // 2 |
| return np.stack([cv2.resize(f[:, x0:x0 + cw], (W, H), interpolation=cv2.INTER_AREA) for f in src]), x0, cw |
|
|
|
|
| def main(): |
| clips = json.load(open(f"{D}/clips.json")) |
| vae = INF.load_vae("cuda", S.MODEL_DIR) |
| out = {k: [] for k in RES} |
| for cid in clips: |
| ref = np.load(f"{D}/_ref/{cid}.npz") |
| src, ids = ref["gt_src"][:NF], ref["ids"] |
| for name, wh in RES.items(): |
| if name == "1280x704": |
| x, x0, cw = src, 0, 1280 |
| else: |
| x, x0, cw = to_res(src, wh) |
| rec = dec(vae, enc(vae, np.ascontiguousarray(x))) |
| for t in range(8, NF, 8): |
| m = np.isin(ids[t], GM.PLAYER_PART_IDS) |
| if m.sum() < 600: |
| continue |
| ys, xs = np.nonzero(m) |
| p = GEO.dst_to_src(np.array([[xs.min() - 10, ys.min() - 10], [xs.max() + 10, ys.max() + 10]], float)) |
| sx0, sy0 = int(max(0, p[0, 0])), int(max(0, p[0, 1])) |
| sx1, sy1 = int(min(1280, p[1, 0])), int(min(704, p[1, 1])) |
| gt_c = src[t, sy0:sy1, sx0:sx1] |
| |
| if name == "1280x704": |
| rc = rec[t, sy0:sy1, sx0:sx1] |
| else: |
| up = cv2.resize(rec[t], (cw, 704), interpolation=cv2.INTER_CUBIC) |
| full = np.zeros_like(src[t]); full[:, x0:x0 + cw] = up |
| rc = full[sy0:sy1, sx0:sx1] |
| mse = np.mean((rc.astype(np.float32) - gt_c.astype(np.float32)) ** 2) |
| g = lambda a: float(np.abs(cv2.Laplacian(cv2.cvtColor(a, cv2.COLOR_RGB2GRAY).astype(np.float32), cv2.CV_32F)).mean()) |
| out[name].append((10 * np.log10(255 ** 2 / max(mse, 1e-6)), g(rc) / max(g(gt_c), 1e-6))) |
| print(cid, {k: len(v) for k, v in out.items()}, flush=True) |
| res = {k: {"n": len(v), "player_crop_psnr_vs_native_gt": round(float(np.mean([a for a, _ in v])), 2), |
| "detail_ratio_vs_native_gt": round(float(np.mean([b for _, b in v])), 3)} for k, v in out.items()} |
| print(json.dumps(res, indent=1), flush=True) |
| del vae |
| torch.cuda.empty_cache() |
|
|
| |
| dit = GM.load_dit(S.MODEL_DIR, "cuda") |
| GM.install_sem_cond(dit, "parts") |
| dit = dit.cuda().eval().requires_grad_(False) |
| dit.freqs = tuple(t.cuda() for t in dit.freqs) |
| tt = S.TextTable(os.path.join(S.DATA_ROOT, "text_table_gc25d.pt")) |
| ctx, cids = tt.context([sorted(tt.table)[0]] * 21, "cuda") |
| timing = {} |
| for name, (W, H) in RES.items(): |
| lat = torch.randn(1, 48, 21, H // 16, W // 16, dtype=torch.bfloat16, device="cuda") |
| sem = torch.ones(1, 81, H, W, dtype=torch.uint8, device="cuda") |
| ts = [] |
| for i in range(4): |
| torch.cuda.synchronize(); t0 = time.time() |
| with torch.no_grad(): |
| GM.gc_forward(dit, lat, torch.tensor([500.0], dtype=torch.bfloat16, device="cuda"), ctx, sem, |
| first_frame_cond=True, context_ids=cids) |
| torch.cuda.synchronize(); ts.append(time.time() - t0) |
| timing[name] = {"tokens_per_cell": (H // 32) * (W // 32), "sec_per_forward": round(float(np.median(ts[1:])), 3)} |
| print(name, timing[name], flush=True) |
| base = timing["832x480"]["sec_per_forward"] |
| for k in timing: |
| timing[k]["ratio_vs_480"] = round(timing[k]["sec_per_forward"] / base, 2) |
| json.dump({"vae_ceiling": res, "dit_timing": timing}, open(f"{D}/res_test.json", "w"), indent=1) |
| print(json.dumps(timing, indent=1)) |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
