lb_aggregate.py

#!/usr/bin/env python
"""Summarise the LingBot-World-V2 evaluation into eval_out_lb/summaries/final_summary.{json,csv}.

Latency: ratio of means against the FFFF latency measured in the same process for
the same job (matched FFFF).  Pixel metrics: PSNR from the mean per-video MSE, SSIM
and LPIPS averaged.  VBench: the eight validation25 dimensions, normalised with
VBench's official ranges; the headline "VBench-I2V-100" is the 251-protocol
formula with temporal_flickering in the quality group and overall_consistency as
the semantic group (there is no scene dimension for I2V):
    quality  = (sc + bc + tf + ms + 0.5*dd + aq + iq) / 6.5
    semantic = oc
    selected = (4*quality + semantic) / 5
"core5" (mean of raw sc, bc, ms, aq, iq) is kept for comparison with the DEV
project's tables.
"""
import csv, glob, json, math, os, sys
ROOT = os.path.dirname(os.path.abspath(__file__))
OUT = os.path.join(ROOT, "eval_out_lb")
# Two headline numbers over the same five dimensions, both reported:
#
#   core5     the plain mean of the raw scores -- the HY-WorldPlay DEV protocol's
#             "VBench Core5" for validation25, so the tables line up with it;
#   score_5d  the mean after VBench's *official* min-max normalisation, which is
#             what the 251-prompt tables use for their quality group.
#
# VBench reports imaging_quality *per video* on 0..100 but already divides its
# own aggregate by 100, and the aggregate is what is collected here -- so no
# further scaling is applied (dividing again would cost ~14 points of Core5).
# aesthetic_quality and imaging_quality live on [0, 1], so their official
# normalisation is the identity.
DIMS = ["subject_consistency", "background_consistency", "motion_smoothness",
        "aesthetic_quality", "imaging_quality"]
NORM_RANGE = {"subject_consistency": (0.1462, 1.0),
              "background_consistency": (0.2615, 1.0),
              "motion_smoothness": (0.7060, 0.9975),
              "aesthetic_quality": (0.0, 1.0),
              "imaging_quality": (0.0, 1.0)}


def normalize_score(dim, value):
    return value


def official_normalize(dim, value):
    lo, hi = NORM_RANGE[dim]
    return (value - lo) / (hi - lo)
only = sys.argv[1] if len(sys.argv) > 1 else None

rows, problems = [], []
for d in sorted(glob.glob(os.path.join(OUT, "per_prompt", "lb_*"))):
    s = os.path.basename(d)
    recs = [json.load(open(p)) for p in sorted(glob.glob(os.path.join(d, "*.json")))]
    recs = [r for r in recs if r.get("status") == "complete"]
    if len(recs) != 100:
        problems.append(f"{s}: {len(recs)} records"); continue
    sc = os.path.join(OUT, "vbench", "scores", s + ".json")
    vb = json.load(open(sc)) if os.path.exists(sc) else None
    if vb is None:
        problems.append(f"{s}: no VBench scores")
    lat = sum(r["policy_latency_ms"] for r in recs) / 100
    matched = [r.get("matched_ffff_policy_latency_ms") for r in recs]
    if all(m is not None for m in matched):
        ref = sum(matched) / 100
        latency_source = "paired with FFFF in the same process"
    else:
        # FFFF was not regenerated for these jobs: compare against lb_ffff's own records.
        ff = [json.load(open(p)) for p in sorted(glob.glob(os.path.join(OUT, "per_prompt", "lb_ffff", "*.json")))]
        ref = sum(r["policy_latency_ms"] for r in ff if r.get("status") == "complete") / 100
        latency_source = "generation run (FFFF from its own records)"
        if all(str(r.get("latency_source", "")).startswith("DEV generator") for r in recs):
            latency_source = recs[0]["latency_source"]
    ce = sum(r["cache_diagnostics"]["compute_equivalent_forwards"] for r in recs) / 100
    fw = recs[0]["cache_diagnostics"]["denoise_forwards"]
    mse = sum(r["pixel_metrics_vs_ffff"]["mean_mse"] for r in recs) / 100
    row = {"strategy": s, "base_model": "lingbot_world_v2_1.3b_causal_fast", "method": recs[0]["method"],
           "schedule": recs[0]["schedule"], "target_speedup": recs[0]["target_speedup"],
           "param": recs[0]["param"], "param_value": recs[0]["param_value"], "num_videos": 100,
           "policy_latency_ms": lat, "matched_ffff_policy_latency_ms": ref,
           "latency_ratio_vs_ffff": lat / ref, "speedup_percent_vs_ffff": 100 * (1 - lat / ref),
           "latency_source": latency_source,
           "first_chunk_schedule": recs[0].get("first_chunk_schedule"),
           "mp4_reference_records": sum(1 for r in recs if r.get("reference_source") == "ffff_mp4"),
           "excluded_context_kv_latency_ms": sum(r["excluded_context_kv_latency_ms"] for r in recs) / 100,
           "denoise_forwards": fw, "mean_compute_equivalent_forwards": ce,
           "mean_full_forwards": (sum(r["cache_diagnostics"]["full_forwards"] for r in recs) / 100
                                  if all("full_forwards" in r["cache_diagnostics"] for r in recs) else None),
           "mean_predictor_forwards": (sum(r["cache_diagnostics"]["predictor_forwards"] for r in recs) / 100
                                       if all("predictor_forwards" in r["cache_diagnostics"] for r in recs) else None),
           "psnr": -10 * math.log10(max(mse, 1e-12)),
           "ssim": sum(r["pixel_metrics_vs_ffff"]["ssim"] for r in recs) / 100,
           "lpips": sum(r["pixel_metrics_vs_ffff"]["lpips"] for r in recs) / 100,
           "peak_mem_gib": max(r.get("peak_mem_gib", 0) for r in recs)}
    if vb:
        scores, normed = {}, {}
        for k in DIMS:
            row["raw_" + k] = vb["raw"][k]
            scores[k] = normalize_score(k, vb["raw"][k])
            normed[k] = official_normalize(k, scores[k])
            row["score_" + k] = scores[k]
            row["normalized_" + k] = normed[k]
        core5 = sum(scores.values()) / len(DIMS)
        score_5d = sum(normed.values()) / len(DIMS)
        row.update(core5=core5, core5_percent=100 * core5,
                   score_5d=score_5d, score_5d_percent=100 * score_5d,
                   selected_vbench_score=score_5d, selected_vbench_percent=100 * score_5d)
    rows.append(row)
rows.sort(key=lambda r: (["none", "teacache", "flowcache", "taylorseer", "motioncache"].index(r["method"])
                         if r["method"] in ("none", "teacache", "flowcache", "taylorseer", "motioncache") else 9,
                         r["schedule"] != "FxxF", r["target_speedup"]))
os.makedirs(os.path.join(OUT, "summaries"), exist_ok=True)
json.dump({"protocol": "LingBot-World-V2 VBench-I2V-100 (validation25 x 4 actions, Core5)",
           "num_strategies": len(rows), "rows": rows, "problems": problems, "dimensions": DIMS,
           "aggregation": {
               "core5": "mean of raw(sc, bc, ms, aq, iq/100)",
               "score_5d": "mean of ((sc-0.1462)/0.8538, (bc-0.2615)/0.7385, "
                           "(ms-0.7060)/0.2915, aq, iq/100)"},
           "normalize_range": NORM_RANGE},
          open(os.path.join(OUT, "summaries", "final_summary.json"), "w"), indent=2)
keys = sorted({k for r in rows for k in r}, key=lambda k: (k not in ("strategy", "method"), k))
with open(os.path.join(OUT, "summaries", "final_summary.csv"), "w", newline="") as f:
    w = csv.DictWriter(f, fieldnames=keys); w.writeheader(); [w.writerow(r) for r in rows]
print(f"{'strategy':26s} {'lat_ms':>8} {'ffff_ms':>8} {'x':>6} {'Full':>6} {'PSNR':>6} {'LPIPS':>6} {'raw%':>7} {'norm%':>7}")
for r in rows:
    print(f"{r['strategy']:26s} {r['policy_latency_ms']:8.0f} {r['matched_ffff_policy_latency_ms']:8.0f} "
          f"{1/r['latency_ratio_vs_ffff']:6.3f} {r['mean_compute_equivalent_forwards']:6.2f} {r['psnr']:6.2f} "
          f"{r['lpips']:6.3f} {r.get('core5_percent', float('nan')):7.2f} "
          f"{r.get('score_5d_percent', float('nan')):7.2f}")
if problems: print("problems:", problems)