scripts/evaluate_layer17_chunk_impact.py

#!/usr/bin/env python3
"""Measure local and downstream impact of isolated Layer-17 Predictor calls.

For every prompt, chunk 1..6 is independently evaluated with FPFF, FFPF, and
FPPF while every other chunk remains FFFF.  A shadow Full forward is executed
at each selected Predictor step to measure hidden/flow/x0 error on the exact
rollout state.  The shadow output is never used by the generated trajectory.
"""

from __future__ import annotations

import argparse
import csv
import json
import math
import os
import sys
import time
from pathlib import Path
from typing import Any


def _preparse_gpu() -> str:
    parser = argparse.ArgumentParser(add_help=False)
    parser.add_argument("--gpu", default="4")
    args, _ = parser.parse_known_args()
    os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu)
    return str(args.gpu)


PHYSICAL_GPU = _preparse_gpu()

import lpips
import torch
from omegaconf import OmegaConf

REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from scripts import evaluate_single_block_fppf as base
from utils.misc import set_seed
from utils.wan_wrapper import WanVAEWrapper


SCHEDULES = ("FPFF", "FFPF", "FPPF")
EPS = 1e-8


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--gpu", default=PHYSICAL_GPU)
    parser.add_argument(
        "--config_path", type=Path, default=Path("configs/self_forcing_sid.yaml")
    )
    parser.add_argument(
        "--checkpoint_path",
        type=Path,
        default=Path("checkpoints/self_forcing_dmd.pt"),
    )
    parser.add_argument(
        "--dataset_root",
        type=Path,
        default=Path("outputs/predictor_offline_100_all_blocks"),
    )
    parser.add_argument(
        "--sweep_dir",
        type=Path,
        default=Path("outputs/single_block_init_sweep"),
    )
    parser.add_argument(
        "--reference_root",
        type=Path,
        default=Path("outputs/single_block_fppf_eval"),
    )
    parser.add_argument(
        "--output_dir",
        type=Path,
        default=Path("outputs/layer17_chunk_impact_pilot"),
    )
    parser.add_argument(
        "--prompt_ids", type=int, nargs="*", default=list(range(80, 90))
    )
    parser.add_argument("--schedules", nargs="*", choices=SCHEDULES, default=list(SCHEDULES))
    parser.add_argument(
        "--chunks", type=int, nargs="*", default=list(range(1, base.NUM_CHUNKS))
    )
    parser.add_argument("--max_prompts", type=int, default=None)
    parser.add_argument("--metric_batch_size", type=int, default=4)
    parser.add_argument("--generation_seed", type=int, default=0)
    parser.add_argument(
        "--skip_lpips", action=argparse.BooleanOptionalAction, default=False
    )
    parser.add_argument("--overwrite", action="store_true")
    args = parser.parse_args()
    if not args.prompt_ids:
        parser.error("At least one prompt ID is required")
    if any(value < 0 or value >= 100 for value in args.prompt_ids):
        parser.error("Prompt IDs must be in [0, 99]")
    if args.metric_batch_size < 1:
        parser.error("--metric_batch_size must be positive")
    if not args.schedules:
        parser.error("At least one schedule is required")
    if not args.chunks or any(chunk < 1 or chunk >= base.NUM_CHUNKS for chunk in args.chunks):
        parser.error("--chunks must contain values in [1, 6]")
    return args


def resolve(path: Path) -> Path:
    return path.resolve() if path.is_absolute() else (REPO_ROOT / path).resolve()


def rms(value: torch.Tensor) -> torch.Tensor:
    return value.float().square().mean().sqrt()


def nrmse(prediction: torch.Tensor, target: torch.Tensor) -> float:
    return float(rms(prediction.float() - target.float()) / rms(target).clamp_min(EPS))


def chunk_frame_slice(chunk: int) -> slice:
    if chunk == 0:
        return slice(0, base.PIXEL_FRAMES_FIRST_CHUNK)
    start = base.PIXEL_FRAMES_FIRST_CHUNK + 12 * (chunk - 1)
    return slice(start, start + 12)


def summarize_frame_range(metrics: dict[str, Any], selected: slice) -> dict[str, float]:
    mse_values = metrics["mse_per_frame"][selected]
    ssim_values = metrics["ssim_per_frame"][selected]
    lpips_values = metrics["lpips_per_frame"][selected]
    mean_mse = sum(mse_values) / len(mse_values)
    return {
        "pixel_mse": mean_mse,
        "psnr": -10.0 * math.log10(max(mean_mse, 1e-12)),
        "ssim": sum(ssim_values) / len(ssim_values),
        "lpips": (
            sum(lpips_values) / len(lpips_values) if lpips_values else float("nan")
        ),
    }


@torch.inference_mode()
def generate_intervention(
    *,
    pipeline: Any,
    dataset_root: Path,
    prompt_id: int,
    generation_seed: int,
    device: torch.device,
    predictor: Any,
    source_layer: int,
    intervention_chunk: int,
    intervention_schedule: str,
) -> tuple[torch.Tensor, dict[str, Any]]:
    if intervention_schedule not in SCHEDULES:
        raise ValueError(intervention_schedule)
    if intervention_chunk < 1 or intervention_chunk >= base.NUM_CHUNKS:
        raise ValueError("Predictor intervention chunk must be 1..6")

    base.reset_kv_and_load_cross_cache(pipeline, dataset_root, prompt_id, device)
    set_seed(generation_seed)
    noise = torch.randn(
        1,
        base.NUM_CHUNKS * base.FRAMES_PER_CHUNK,
        base.LATENT_CHANNELS,
        base.LATENT_HEIGHT,
        base.LATENT_WIDTH,
        dtype=torch.bfloat16,
        device=device,
    )
    timesteps = pipeline.denoising_step_list.to(device=device)
    output_chunks: list[torch.Tensor] = []
    previous_chunk_hidden: list[torch.Tensor | None] | None = None
    teacher = pipeline.generator.model
    capture = base.FinalHiddenCapture(teacher)
    full_calls = 0
    predictor_calls = 0
    shadow_full_calls = 0
    local_errors: list[dict[str, float | int]] = []
    started = time.perf_counter()

    try:
        for chunk in range(base.NUM_CHUNKS):
            noisy_input = noise[
                :,
                chunk * base.FRAMES_PER_CHUNK : (chunk + 1) * base.FRAMES_PER_CHUNK,
            ]
            current_hidden: list[torch.Tensor | None] = [None] * base.NUM_DENOISING_STEPS
            denoised_pred: torch.Tensor | None = None
            timestep: torch.Tensor | None = None

            for step, current_timestep in enumerate(timesteps):
                timestep = torch.ones(
                    [1, base.FRAMES_PER_CHUNK], dtype=torch.int64, device=device
                ) * current_timestep
                selected = (
                    chunk == intervention_chunk
                    and intervention_schedule[step] == "P"
                )

                if selected:
                    anchor_hidden = current_hidden[step - 1]
                    if anchor_hidden is None or previous_chunk_hidden is None:
                        raise RuntimeError("Predictor inputs are unavailable")
                    previous_hidden = previous_chunk_hidden[step]
                    if previous_hidden is None:
                        raise RuntimeError("Previous-chunk hidden is unavailable")
                    history = pipeline.kv_cache1[source_layer]
                    cross = pipeline.crossattn_cache[source_layer]
                    pred_hidden, pred_flow, _ = base.predictor_step(
                        predictor=predictor,
                        teacher=teacher,
                        noisy_input=noisy_input,
                        timestep=timestep,
                        anchor_hidden=anchor_hidden,
                        previous_hidden=previous_hidden,
                        history_cache=history,
                        cross_cache=cross,
                        current_start=chunk * base.TOKENS_PER_CHUNK,
                    )
                    pred_x0 = pipeline.generator._convert_flow_pred_to_x0(
                        flow_pred=pred_flow.flatten(0, 1),
                        xt=noisy_input.flatten(0, 1),
                        timestep=timestep.flatten(0, 1),
                    ).unflatten(0, pred_flow.shape[:2])

                    # Shadow Full measures the exact counterfactual target on
                    # this rollout state. Its x0/hidden are never accepted.
                    capture.start()
                    full_flow, full_x0 = pipeline.generator(
                        noisy_image_or_video=noisy_input,
                        conditional_dict={
                            "prompt_embeds": torch.zeros(
                                1,
                                1,
                                int(teacher.text_embedding[0].in_features),
                                dtype=torch.bfloat16,
                                device=device,
                            )
                        },
                        timestep=timestep,
                        kv_cache=pipeline.kv_cache1,
                        crossattn_cache=pipeline.crossattn_cache,
                        current_start=chunk * base.TOKENS_PER_CHUNK,
                    )
                    full_hidden = capture.finish()
                    local_errors.append(
                        {
                            "step": step,
                            "timestep": float(current_timestep),
                            "hidden_nrmse": nrmse(pred_hidden, full_hidden),
                            "flow_nrmse": nrmse(pred_flow, full_flow),
                            "x0_nrmse": nrmse(pred_x0, full_x0),
                        }
                    )
                    current_hidden[step] = pred_hidden
                    denoised_pred = pred_x0
                    predictor_calls += 1
                    shadow_full_calls += 1
                    del full_flow, full_x0, full_hidden
                else:
                    capture.start()
                    _, denoised_pred = pipeline.generator(
                        noisy_image_or_video=noisy_input,
                        conditional_dict={
                            "prompt_embeds": torch.zeros(
                                1,
                                1,
                                int(teacher.text_embedding[0].in_features),
                                dtype=torch.bfloat16,
                                device=device,
                            )
                        },
                        timestep=timestep,
                        kv_cache=pipeline.kv_cache1,
                        crossattn_cache=pipeline.crossattn_cache,
                        current_start=chunk * base.TOKENS_PER_CHUNK,
                    )
                    current_hidden[step] = capture.finish()
                    full_calls += 1

                if step < base.NUM_DENOISING_STEPS - 1:
                    if denoised_pred is None:
                        raise RuntimeError("Denoising step produced no x0")
                    next_timestep = timesteps[step + 1]
                    flat = denoised_pred.flatten(0, 1)
                    noisy_input = pipeline.scheduler.add_noise(
                        flat,
                        torch.randn_like(flat),
                        next_timestep
                        * torch.ones(
                            [base.FRAMES_PER_CHUNK], dtype=torch.long, device=device
                        ),
                    ).unflatten(0, denoised_pred.shape[:2])

            if denoised_pred is None or timestep is None:
                raise RuntimeError("Chunk produced no clean latent")
            output_chunks.append(denoised_pred)
            context_timestep = torch.ones_like(timestep) * pipeline.args.context_noise
            pipeline.generator(
                noisy_image_or_video=denoised_pred,
                conditional_dict={
                    "prompt_embeds": torch.zeros(
                        1,
                        1,
                        int(teacher.text_embedding[0].in_features),
                        dtype=torch.bfloat16,
                        device=device,
                    )
                },
                timestep=context_timestep,
                kv_cache=pipeline.kv_cache1,
                crossattn_cache=pipeline.crossattn_cache,
                current_start=chunk * base.TOKENS_PER_CHUNK,
            )
            previous_chunk_hidden = current_hidden
    finally:
        capture.close()

    torch.cuda.synchronize()
    return torch.cat(output_chunks, dim=1), {
        "generation_time_s": time.perf_counter() - started,
        "full_calls": full_calls,
        "predictor_calls": predictor_calls,
        "shadow_full_calls": shadow_full_calls,
        "local_errors": local_errors,
    }


def average(values: list[float]) -> float:
    return sum(values) / len(values)


def finite_average(values: list[Any]) -> float:
    numeric = [
        float(value)
        for value in values
        if value is not None and math.isfinite(float(value))
    ]
    return average(numeric) if numeric else float("nan")


def rankdata(values: list[float]) -> list[float]:
    order = sorted(range(len(values)), key=values.__getitem__)
    ranks = [0.0] * len(values)
    start = 0
    while start < len(order):
        end = start + 1
        while end < len(order) and values[order[end]] == values[order[start]]:
            end += 1
        rank = 0.5 * (start + end - 1)
        for position in range(start, end):
            ranks[order[position]] = rank
        start = end
    return ranks


def pearson(left: list[float], right: list[float]) -> float:
    left_mean, right_mean = average(left), average(right)
    left_centered = [value - left_mean for value in left]
    right_centered = [value - right_mean for value in right]
    numerator = sum(a * b for a, b in zip(left_centered, right_centered))
    denominator = math.sqrt(
        sum(value * value for value in left_centered)
        * sum(value * value for value in right_centered)
    )
    return numerator / denominator if denominator > 0 else float("nan")


def spearman(left: list[float], right: list[float]) -> float:
    return pearson(rankdata(left), rankdata(right))


def write_csv(path: Path, rows: list[dict[str, Any]], fields: list[str]) -> None:
    temporary = path.with_suffix(path.suffix + ".tmp")
    with temporary.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=fields)
        writer.writeheader()
        writer.writerows(rows)
    os.replace(temporary, path)


def group_centered_values(
    records: list[dict[str, Any]], field: str
) -> list[float]:
    """Remove schedule-by-chunk means to isolate prompt/state variation."""
    groups: dict[tuple[str, int], list[float]] = {}
    for row in records:
        key = (str(row["schedule"]), int(row["chunk"]))
        groups.setdefault(key, []).append(float(row[field]))
    means = {key: average(values) for key, values in groups.items()}
    return [
        float(row[field]) - means[(str(row["schedule"]), int(row["chunk"]))]
        for row in records
    ]


def aggregate(records: list[dict[str, Any]], output_dir: Path) -> None:
    numeric_fields = [
        "hidden_nrmse_mean",
        "flow_nrmse_mean",
        "x0_nrmse_mean",
        "latent_all_nrmse",
        "latent_current_nrmse",
        "latent_tail_nrmse",
        "psnr",
        "ssim",
        "lpips",
        "current_psnr",
        "current_ssim",
        "current_lpips",
        "tail_psnr",
        "tail_ssim",
        "tail_lpips",
        "generation_time_s",
    ]
    summary_rows: list[dict[str, Any]] = []
    available_schedules = [
        schedule for schedule in SCHEDULES if any(row["schedule"] == schedule for row in records)
    ]
    available_chunks = sorted({int(row["chunk"]) for row in records})
    for schedule in available_schedules:
        for chunk in available_chunks:
            selected = [
                row
                for row in records
                if row["schedule"] == schedule and row["chunk"] == chunk
            ]
            row: dict[str, Any] = {
                "schedule": schedule,
                "chunk": chunk,
                "num_prompts": len(selected),
                "alpha_eligible_linear": (base.NUM_CHUNKS - 1 - chunk)
                / (base.NUM_CHUNKS - 2),
            }
            for field in numeric_fields:
                row[field] = finite_average([item[field] for item in selected])
            summary_rows.append(row)

    summary_fields = [
        "schedule",
        "chunk",
        "num_prompts",
        "alpha_eligible_linear",
        *numeric_fields,
    ]
    write_csv(output_dir / "summary_by_schedule_chunk.csv", summary_rows, summary_fields)

    correlation_rows: list[dict[str, Any]] = []
    for schedule in (*available_schedules, "ALL"):
        selected = (
            records if schedule == "ALL" else [r for r in records if r["schedule"] == schedule]
        )
        for local in ("hidden_nrmse_mean", "flow_nrmse_mean", "x0_nrmse_mean"):
            for downstream in ("tail_lpips", "latent_tail_nrmse", "tail_pixel_mse"):
                pairs = [
                    (float(row[local]), float(row[downstream]))
                    for row in selected
                    if row[local] is not None
                    and row[downstream] is not None
                    and math.isfinite(float(row[local]))
                    and math.isfinite(float(row[downstream]))
                ]
                correlation_rows.append(
                    {
                        "schedule": schedule,
                        "local_metric": local,
                        "downstream_metric": downstream,
                        "num_observations": len(pairs),
                        "spearman": (
                            spearman(
                                [pair[0] for pair in pairs],
                                [pair[1] for pair in pairs],
                            )
                            if len(pairs) >= 2
                            else float("nan")
                        ),
                    }
                )
    correlation_fields = [
        "schedule",
        "local_metric",
        "downstream_metric",
        "num_observations",
        "spearman",
    ]
    write_csv(output_dir / "local_downstream_correlations.csv", correlation_rows, correlation_fields)

    controlled_rows: list[dict[str, Any]] = []
    for local in ("hidden_nrmse_mean", "flow_nrmse_mean", "x0_nrmse_mean"):
        local_residual = group_centered_values(records, local)
        downstream_residual = group_centered_values(records, "tail_lpips")
        within_cell = []
        for schedule in available_schedules:
            for chunk in available_chunks:
                selected = [
                    row
                    for row in records
                    if row["schedule"] == schedule and row["chunk"] == chunk
                ]
                if len(selected) >= 2:
                    within_cell.append(
                        spearman(
                            [float(row[local]) for row in selected],
                            [float(row["tail_lpips"]) for row in selected],
                        )
                    )
        controlled_rows.append(
            {
                "local_metric": local,
                "downstream_metric": "tail_lpips",
                "controls": "schedule+chunk",
                "residual_spearman": spearman(local_residual, downstream_residual),
                "mean_within_cell_spearman": finite_average(within_cell),
                "positive_cells": sum(value > 0 for value in within_cell),
                "num_cells": len(within_cell),
            }
        )
    controlled_fields = [
        "local_metric",
        "downstream_metric",
        "controls",
        "residual_spearman",
        "mean_within_cell_spearman",
        "positive_cells",
        "num_cells",
    ]
    write_csv(
        output_dir / "controlled_local_downstream_correlations.csv",
        controlled_rows,
        controlled_fields,
    )

    report = [
        "# Layer-17 Predictor chunk-impact pilot",
        "",
        f"Prompts: {len(set(row['prompt_id'] for row in records))}; seed 0; "
        "all non-intervened chunks use FFFF.",
        "",
        "A shadow Full call measures local error at each Predictor decision, but the "
        "generated trajectory always consumes the Predictor output.",
        "",
    ]
    for schedule in available_schedules:
        report.extend(
            [
                f"## {schedule}",
                "",
                "| Chunk | x0 nRMSE | Tail PSNR | Tail LPIPS | Tail latent nRMSE |",
                "|---:|---:|---:|---:|---:|",
            ]
        )
        for row in summary_rows:
            if row["schedule"] != schedule:
                continue
            report.append(
                f"| {row['chunk']} | {row['x0_nrmse_mean']:.6f} | "
                f"{row['tail_psnr']:.4f} | {row['tail_lpips']:.6f} | "
                f"{row['latent_tail_nrmse']:.6f} |"
            )
        report.append("")
    report.extend(
        [
            "## Local-to-downstream Spearman correlations",
            "",
            "| Schedule | Local metric | Downstream metric | Spearman |",
            "|---|---|---|---:|",
        ]
    )
    for row in correlation_rows:
        if row["downstream_metric"] == "tail_lpips":
            report.append(
                f"| {row['schedule']} | {row['local_metric']} | tail LPIPS | "
                f"{row['spearman']:.4f} |"
            )
    report.extend(
        [
            "",
            "## Correlations after controlling schedule and chunk",
            "",
            "Residual correlations remove each schedule-by-chunk mean, so they test "
            "whether local error explains prompt/state risk beyond the position prior.",
            "",
            "| Local metric | Residual Spearman | Mean within-cell Spearman | Positive cells |",
            "|---|---:|---:|---:|",
        ]
    )
    for row in controlled_rows:
        report.append(
            f"| {row['local_metric']} | {row['residual_spearman']:.4f} | "
            f"{row['mean_within_cell_spearman']:.4f} | "
            f"{row['positive_cells']}/{row['num_cells']} |"
        )
    (output_dir / "REPORT.md").write_text("\n".join(report) + "\n", encoding="utf-8")


def main() -> None:
    args = parse_args()
    for name in (
        "config_path",
        "checkpoint_path",
        "dataset_root",
        "sweep_dir",
        "reference_root",
        "output_dir",
    ):
        setattr(args, name, resolve(getattr(args, name)))
    args.output_dir.mkdir(parents=True, exist_ok=True)
    prompt_ids = sorted(set(args.prompt_ids))
    if args.max_prompts is not None:
        prompt_ids = prompt_ids[: args.max_prompts]

    config = OmegaConf.merge(
        OmegaConf.load(REPO_ROOT / "configs/default_config.yaml"),
        OmegaConf.load(args.config_path),
    )
    device = torch.device("cuda")
    torch.set_grad_enabled(False)
    set_seed(args.generation_seed)
    manifest = {
        "status": "running",
        "gpu": str(args.gpu),
        "prompt_ids": prompt_ids,
        "generation_seed": args.generation_seed,
        "checkpoint_path": str(args.checkpoint_path),
        "predictor_weights": str(
            args.sweep_dir / "teacher_layer_17" / "predictor_final.safetensors"
        ),
        "dataset_root": str(args.dataset_root),
        "schedules": list(args.schedules),
        "chunks": list(args.chunks),
        "shadow_full_target": True,
    }
    base.atomic_json(args.output_dir / "manifest.json", manifest)

    print("[setup] loading VAE", flush=True)
    vae = WanVAEWrapper().to(device=device, dtype=torch.bfloat16).eval()
    missing_references = [
        prompt_id
        for prompt_id in prompt_ids
        if not (
            args.reference_root
            / "ffff_reference_frames"
            / f"prompt_{prompt_id:04d}.safetensors"
        ).exists()
    ]
    if missing_references:
        base.prepare_reference_frames(
            vae=vae,
            dataset_root=args.dataset_root,
            output_dir=args.reference_root,
            prompt_ids=missing_references,
            device=device,
            rebuild=False,
        )

    print("[setup] loading frozen generator and Layer-17 Predictor", flush=True)
    pipeline = base.build_pipeline(config, args.checkpoint_path, vae, device)
    experiment = base.discover_experiments(
        args.sweep_dir, ["teacher_layer_17"], None
    )[0]
    predictor = base.load_predictor(pipeline.generator.model, experiment, device)
    lpips_model = None
    if not args.skip_lpips:
        lpips_model = lpips.LPIPS(net="alex", verbose=False).to(device).eval()
        lpips_model.requires_grad_(False)

    records: list[dict[str, Any]] = []
    total = len(prompt_ids) * len(args.schedules) * len(args.chunks)
    completed = 0
    for prompt_id in prompt_ids:
        reference_latent = base.load_ffff_latent(args.dataset_root, prompt_id).to(
            device=device, dtype=torch.bfloat16
        )
        reference_u8 = base.load_reference_frames(args.reference_root, prompt_id)
        for schedule in args.schedules:
            for chunk in args.chunks:
                destination = (
                    args.output_dir
                    / "per_intervention"
                    / f"prompt_{prompt_id:04d}_{schedule}_chunk_{chunk:02d}.json"
                )
                if destination.exists() and not args.overwrite:
                    record = json.loads(destination.read_text(encoding="utf-8"))
                    records.append(record)
                    completed += 1
                    print(f"[cached] {completed}/{total} {destination.stem}", flush=True)
                    continue
                started = time.perf_counter()
                latent, diagnostics = generate_intervention(
                    pipeline=pipeline,
                    dataset_root=args.dataset_root,
                    prompt_id=prompt_id,
                    generation_seed=args.generation_seed,
                    device=device,
                    predictor=predictor,
                    source_layer=17,
                    intervention_chunk=chunk,
                    intervention_schedule=schedule,
                )
                with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
                    pixels = vae.decode_to_pixel(latent, use_cache=False)
                prediction_u8 = base.pixels_to_u8(pixels)
                frame = base.frame_metrics(
                    reference_u8=reference_u8,
                    prediction_u8=prediction_u8,
                    lpips_model=lpips_model,
                    batch_size=args.metric_batch_size,
                    device=device,
                )
                current_slice = chunk_frame_slice(chunk)
                current = summarize_frame_range(frame, current_slice)
                tail = summarize_frame_range(frame, slice(current_slice.start, None))
                latent_chunk_start = chunk * base.FRAMES_PER_CHUNK
                latent_chunk_end = latent_chunk_start + base.FRAMES_PER_CHUNK
                errors = diagnostics.pop("local_errors")
                record = {
                    "prompt_id": prompt_id,
                    "schedule": schedule,
                    "chunk": chunk,
                    "predictor_steps": [int(item["step"]) for item in errors],
                    "hidden_nrmse_mean": average(
                        [float(item["hidden_nrmse"]) for item in errors]
                    ),
                    "flow_nrmse_mean": average(
                        [float(item["flow_nrmse"]) for item in errors]
                    ),
                    "x0_nrmse_mean": average(
                        [float(item["x0_nrmse"]) for item in errors]
                    ),
                    "local_errors": errors,
                    "latent_all_nrmse": nrmse(latent, reference_latent),
                    "latent_current_nrmse": nrmse(
                        latent[:, latent_chunk_start:latent_chunk_end],
                        reference_latent[:, latent_chunk_start:latent_chunk_end],
                    ),
                    "latent_tail_nrmse": nrmse(
                        latent[:, latent_chunk_start:],
                        reference_latent[:, latent_chunk_start:],
                    ),
                    "psnr": frame["psnr"],
                    "ssim": frame["ssim"],
                    "lpips": frame["lpips"],
                    "current_pixel_mse": current["pixel_mse"],
                    "current_psnr": current["psnr"],
                    "current_ssim": current["ssim"],
                    "current_lpips": current["lpips"],
                    "tail_pixel_mse": tail["pixel_mse"],
                    "tail_psnr": tail["psnr"],
                    "tail_ssim": tail["ssim"],
                    "tail_lpips": tail["lpips"],
                    **diagnostics,
                    "total_time_s": time.perf_counter() - started,
                }
                base.atomic_json(destination, record)
                records.append(record)
                completed += 1
                print(
                    f"[run] {completed}/{total} p={prompt_id} {schedule} c={chunk} "
                    f"x0={record['x0_nrmse_mean']:.5f} "
                    f"tail_lpips={record['tail_lpips']:.5f} "
                    f"time={record['total_time_s']:.1f}s",
                    flush=True,
                )
                if hasattr(vae.model, "clear_cache"):
                    vae.model.clear_cache()
                del latent, pixels, prediction_u8, frame
                torch.cuda.empty_cache()
        del reference_latent, reference_u8

    fields = sorted({key for record in records for key in record if key != "local_errors"})
    flattened = [{key: row.get(key) for key in fields} for row in records]
    write_csv(args.output_dir / "interventions.csv", flattened, fields)
    aggregate(records, args.output_dir)
    manifest["status"] = "complete"
    base.atomic_json(args.output_dir / "manifest.json", manifest)
    print(f"[complete] results -> {args.output_dir}", flush=True)


if __name__ == "__main__":
    main()