File size: 5,323 Bytes
51ac7cc | 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 | #!/usr/bin/env python3
"""分辨率决策的两个快测：(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]
# 重建帧映回源分辨率坐标再裁同一块（832/960 的重建上采样回源尺寸，比较同一份 GT）
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()
# (2) DiT 前向耗时（bf16，no_grad，21 cell，双向；只看相对比例）
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())
|