#!/usr/bin/env python3
"""角色清晰度诊断·704p 生成端（char_gen.py 的 GC_RES 版）：同一批 char_diag clip（outputs/char_diag/clips.json），
704p ckpt 50 步采样，无损存解码帧；--recon 另存 704p VAE 重建（704p 的 VAE 上限）。
输出 <out>/<variant>/<clip>.npz（uint8 [81,H,W,3]）、<out>/_ref704/<clip>.npz（recon）。要 GC_RES 与 data_root 一致。
usage: GC_RES=1280x704 char_704.py --ckpt outputs/sft_v10_704/step-5000.safetensors --variant s5000_704 [--clips 0,1] [--recon]"""
import argparse
import json
import os
import sys
import types

import numpy as np
import torch

import infer as INF
import sample as S

M = "/opt/dlami/nvme/zhiyangdeng/maliang"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", required=True)
    ap.add_argument("--variant", required=True)
    ap.add_argument("--data_root", default=f"{M}/data/gc25d_v10_704")
    ap.add_argument("--clips", default="", help="clips.json 里的下标，逗号分隔（分卡用）；空 = 全部")
    ap.add_argument("--steps", type=int, default=50)
    ap.add_argument("--recon", action="store_true")
    ap.add_argument("--out", default=f"{M}/outputs/char_diag")
    a = ap.parse_args()
    clips = json.load(open(f"{a.out}/clips.json"))
    if a.clips:
        clips = [clips[int(i)] for i in a.clips.split(",")]
    os.makedirs(f"{a.out}/{a.variant}", exist_ok=True)
    args = types.SimpleNamespace(device="cuda", steps=a.steps, shift=5.0, cfg=1.0, sem_cfg=1.0, seed=0)
    dit, info = S.load_dit(a.ckpt, False, "parts", S.MODEL_DIR, "cuda")
    tt = S.TextTable(os.path.join(a.data_root, "text_table_gc25d.pt"))
    vae = INF.load_vae("cuda", S.MODEL_DIR)
    for cid in clips:
        dst = f"{a.out}/{a.variant}/{cid}.npz"
        c = S.load_clip(a.data_root, cid)
        if not os.path.exists(dst):
            z0 = c["latent"][:, 0:1].unsqueeze(0).to("cuda")
            lat, sec = S.generate(args, dit, tt, z0, S.model_ids(c["ids"], info["vocab"], "cond"), c["prompts"])
            np.savez(dst, frames=INF.decode_video(vae, lat))
            print(a.variant, cid, f"{sec:.1f}s", flush=True)
        if a.recon and not os.path.exists(f"{a.out}/_ref704/{cid}.npz"):
            os.makedirs(f"{a.out}/_ref704", exist_ok=True)
            np.savez(f"{a.out}/_ref704/{cid}.npz", recon=INF.decode_video(vae, c["latent"].unsqueeze(0).to("cuda")))
    print("DONE", a.variant, flush=True)


if __name__ == "__main__":
    sys.exit(main())