models/dreamgen/examples/debug_dreamgen_predictor.py

#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Offline predictor diagnostics for DreamGen GR00T LVG.

Runs ONE dense LVG rollout with PredictorProbe hooked on pipe.dit, scoring -- per
block, per net call -- the relative reuse error of each candidate RFE predictor
WITHOUT changing the real output (read-only forward hooks). This tells us *why*
the online linear predictor underperforms `copy` on DreamGen:

  copy        : O_a
  linear      : O_a + s (I - I_a)
  band_linear : O_a + sum_b s_b band_b(I - I_a)
  snr_shrink  : O_a + g s (I - I_a)
  clamp{c}    : O_a + clip(s, 1/c, c) (I - I_a)        (two-sided clamp sweep)
  step_extrap : O_{t-1} + (O_{t-1} - O_{t-2})          (WorldCache-style temporal)

CFG note: DreamGen calls dit twice per step (cond, uncond). We route the two
branches to independent probe anchors via set_branch(call_parity % 2).

Outputs (under --out-dir):
  records.json   : raw per-(block, call) errors
  summary.json   : mean error per predictor, per-block, per-step-bin, ratio_s stats
  *.png          : error charts

Usage:
  python -m examples.debug_dreamgen_predictor \
     --model_size 14B --gr00t_variant droid \
     --batch_input_json .../batch_input.json --num_chunks 2 \
     --out-dir .../outputs/dreamgen_predictor_debug
"""
import argparse
import json
import os
import sys
import tempfile
from pathlib import Path

os.environ["TOKENIZERS_PARALLELISM"] = "false"

import numpy as np
import torch

from imaginaire.utils import log
from imaginaire.utils.io import save_image_or_video
from examples.video2world import _DEFAULT_NEGATIVE_PROMPT, validate_input_file
from examples.video2world_gr00t_lvg import setup_pipeline

_AM_DIT_SRC = os.environ.get(
    "AM_DIT_SRC", "/pfss/mlde/workspaces/mlde_wsp_IAS_SAMMerge/VLA/doanh/AM_DiT/src")
if _AM_DIT_SRC not in sys.path:
    sys.path.insert(0, _AM_DIT_SRC)
# Probe lives under src/snr/experiments
_EXP = os.path.join(_AM_DIT_SRC, "snr", "experiments")
if _EXP not in sys.path:
    sys.path.insert(0, _EXP)


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--model_size", choices=["2B", "14B"], default="14B")
    p.add_argument("--dit_path", type=str, default="")
    p.add_argument("--load_ema", action="store_true")
    p.add_argument("--gr00t_variant", type=str, required=True, choices=["gr1", "droid"])
    p.add_argument("--batch_input_json", type=str, required=True)
    p.add_argument("--num_chunks", type=int, default=2)
    p.add_argument("--num_conditional_frames", type=int, default=1, choices=[1, 5])
    p.add_argument("--aspect_ratio", default="16:9", type=str)
    p.add_argument("--guidance", type=float, default=7)
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--num_gpus", type=int, default=1)
    p.add_argument("--disable_guardrail", action="store_true")
    p.add_argument("--prompt_prefix", type=str, default="")
    p.add_argument("--negative_prompt", type=str, default=_DEFAULT_NEGATIVE_PROMPT)
    p.add_argument("--enable_prompt_refiner_per_chunk", action="store_true")
    p.add_argument("--num-episodes", type=int, default=1)
    p.add_argument("--warmup", type=int, default=3)
    p.add_argument("--snr-shrink-scale", type=float, default=5.0)
    p.add_argument("--downsample", type=int, default=2,
                   help="Spatial avg-pool factor for stored anchors (memory). "
                        "2 = 4x smaller; relative-error ordering is preserved.")
    p.add_argument("--no-bands", action="store_true",
                   help="Disable the band-linear candidate (saves the most memory).")
    p.add_argument("--out-dir", type=str, required=True)
    # Sparse-backend args are referenced by setup_pipeline via getattr; supply none.
    from examples.video2world_gr00t_lvg import _import_sparse_backends
    (add_pisa, add_svg, add_radial, add_sito, add_itm, add_cache, _b) = _import_sparse_backends()
    add_pisa(p); add_svg(p); add_radial(p); add_sito(p); add_itm(p); add_cache(p)
    return p.parse_args()


def _get_dit(pipe):
    return pipe.dit


def attach_probe(pipe, args):
    from _predictor_probe import PredictorProbe

    dit = _get_dit(pipe)
    probe = PredictorProbe(dit, warmup=args.warmup, snr_shrink_scale=args.snr_shrink_scale,
                           s_clamps=(1.0, 1.5, 2.0, 3.0),
                           downsample=args.downsample, use_bands=not args.no_bands)
    original_forward = dit.forward
    parity = {"n": 0}

    def wrapped(*fa, **fk):
        branch = parity["n"] % 2
        probe.set_branch(branch)
        # advance the global call counter once per full step (after uncond branch)
        if branch == 0:
            probe.call_idx += 1
        parity["n"] += 1
        return original_forward(*fa, **fk)

    dit.forward = wrapped

    def restore():
        dit.forward = original_forward
        probe.cleanup()

    return probe, restore, parity


def run_one(pipe, input_path, prompt, args, probe, parity):
    if not validate_input_file(input_path, args.num_conditional_frames):
        return False
    full_prompt = args.prompt_prefix + prompt
    current = input_path
    with tempfile.TemporaryDirectory() as tmp:
        for chunk_id in range(args.num_chunks):
            parity["n"] = 0  # reset cond/uncond parity each chunk (fresh denoise)
            video, _ = pipe(
                prompt=full_prompt, negative_prompt=args.negative_prompt,
                aspect_ratio=args.aspect_ratio, input_path=current,
                num_conditional_frames=args.num_conditional_frames,
                guidance=args.guidance, seed=args.seed + chunk_id, return_prompt=True,
            )
            last = video[:, :, -args.num_conditional_frames:, :, :]
            ext = "png" if args.num_conditional_frames == 1 else "mp4"
            p = os.path.join(tmp, f"c{chunk_id}.{ext}")
            save_image_or_video(last, p, fps=16)
            current = p
    return True


def summarize_and_plot(records, out_dir: Path):
    out_dir.mkdir(parents=True, exist_ok=True)
    (out_dir / "records.json").write_text(json.dumps(records[:5000], indent=2))

    keys = [k for k in records[0].keys() if k.startswith("err_")] if records else []

    def col(name):
        return np.array([r[name] for r in records if r.get(name) is not None], dtype=float)

    means = {k: (float(np.mean(col(k))) if len(col(k)) else None) for k in keys}
    med = {k: (float(np.median(col(k))) if len(col(k)) else None) for k in keys}
    s_vals = np.array([r["ratio_s"] for r in records if r.get("ratio_s") is not None], dtype=float)

    summary = {
        "n_records": len(records),
        "mean_err": means,
        "median_err": med,
        "ratio_s": {
            "mean": float(np.mean(s_vals)) if len(s_vals) else None,
            "median": float(np.median(s_vals)) if len(s_vals) else None,
            "p90": float(np.percentile(s_vals, 90)) if len(s_vals) else None,
            "max": float(np.max(s_vals)) if len(s_vals) else None,
            "frac_gt_1": float(np.mean(s_vals > 1.0)) if len(s_vals) else None,
        },
    }

    # Per-block mean error for copy vs linear vs best.
    blocks = sorted({r["block"] for r in records})
    per_block = {}
    for b in blocks:
        rb = [r for r in records if r["block"] == b]
        per_block[b] = {k: (float(np.mean([r[k] for r in rb if r.get(k) is not None]))
                            if any(r.get(k) is not None for r in rb) else None)
                        for k in keys}
    summary["per_block"] = per_block
    (out_dir / "summary.json").write_text(json.dumps(summary, indent=2))

    # ---- Charts ----
    try:
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
    except ImportError:
        log.warning("matplotlib not available; skipping charts")
        return summary

    # 1) mean error per predictor (bar)
    fig, ax = plt.subplots(figsize=(8, 4))
    names = [k.replace("err_", "") for k in keys]
    vals = [means[k] if means[k] is not None else 0 for k in keys]
    order = np.argsort(vals)
    ax.bar([names[i] for i in order], [vals[i] for i in order], color="steelblue")
    ax.set_ylabel("mean relative reuse error")
    ax.set_title("Predictor reuse error (lower=better) — DreamGen")
    ax.tick_params(axis="x", rotation=45)
    fig.tight_layout(); fig.savefig(out_dir / "err_by_predictor.png", dpi=120); plt.close(fig)

    # 2) per-block: copy vs linear vs step_extrap
    fig, ax = plt.subplots(figsize=(9, 4))
    for k, c in [("err_copy", "gray"), ("err_linear", "crimson"),
                 ("err_step_extrap", "green"), ("err_clamp1.5", "orange")]:
        if k in keys:
            ax.plot(blocks, [per_block[b][k] for b in blocks], label=k.replace("err_", ""), color=c, marker=".")
    ax.set_xlabel("DiT block index"); ax.set_ylabel("mean reuse error")
    ax.set_title("Per-block reuse error"); ax.legend()
    fig.tight_layout(); fig.savefig(out_dir / "err_by_block.png", dpi=120); plt.close(fig)

    # 3) error vs denoise step (call_idx bins)
    calls = sorted({r["call_idx"] for r in records})
    fig, ax = plt.subplots(figsize=(9, 4))
    for k, c in [("err_copy", "gray"), ("err_linear", "crimson"), ("err_step_extrap", "green")]:
        if k in keys:
            ys = []
            for ci in calls:
                rc = [r[k] for r in records if r["call_idx"] == ci and r.get(k) is not None]
                ys.append(np.mean(rc) if rc else np.nan)
            ax.plot(calls, ys, label=k.replace("err_", ""), color=c, marker=".")
    ax.set_xlabel("denoise step (call_idx)"); ax.set_ylabel("mean reuse error")
    ax.set_title("Reuse error vs denoise step"); ax.legend()
    fig.tight_layout(); fig.savefig(out_dir / "err_by_step.png", dpi=120); plt.close(fig)

    # 4) ratio_s histogram
    if len(s_vals):
        fig, ax = plt.subplots(figsize=(7, 4))
        ax.hist(np.clip(s_vals, 0, 10), bins=60, color="purple", alpha=0.8)
        ax.axvline(1.0, color="k", ls="--", label="s=1")
        ax.set_xlabel("fitted ratio_s = ||dO||/||dI||"); ax.set_ylabel("count")
        ax.set_title(f"ratio_s distribution (median={summary['ratio_s']['median']:.2f}, "
                     f"p90={summary['ratio_s']['p90']:.2f})")
        ax.legend(); fig.tight_layout(); fig.savefig(out_dir / "ratio_s_hist.png", dpi=120); plt.close(fig)

    return summary


def main():
    args = parse_args()
    out_dir = Path(args.out_dir)
    pipe = setup_pipeline(args)
    probe, restore, parity = attach_probe(pipe, args)

    with open(args.batch_input_json) as f:
        batch = json.load(f)[: args.num_episodes]

    all_records = []
    try:
        for idx, item in enumerate(batch):
            iv, pr = item.get("input_video", ""), item.get("prompt", "")
            if not iv or not pr:
                continue
            probe.reset_episode()
            log.info(f"[debug] episode {idx}: {iv}")
            run_one(pipe, iv, pr, args, probe, parity)
            for r in probe.records:
                r["episode"] = idx
            all_records.extend(probe.records)
    finally:
        restore()

    if not all_records:
        log.error("No probe records collected.")
        return
    summary = summarize_and_plot(all_records, out_dir)

    print("\n=== Predictor reuse error (mean, lower=better) ===")
    for k, v in sorted(summary["mean_err"].items(), key=lambda kv: (kv[1] is None, kv[1] or 0)):
        print(f"  {k.replace('err_',''):<14}: {v:.4f}" if v is not None else f"  {k}: n/a")
    rs = summary["ratio_s"]
    print(f"\nratio_s: median={rs['median']:.3f}  p90={rs['p90']:.3f}  "
          f"max={rs['max']:.3f}  frac>1={rs['frac_gt_1']:.2f}")
    print(f"\nCharts + JSON in: {out_dir}")


if __name__ == "__main__":
    main()