scripts/evaluate_single_block_fppf.py

#!/usr/bin/env python3
"""Evaluate trained one-block Predictors with an FPPF rollout against FFFF.

The evaluation uses the held-out offline prompt shards.  FFFF clean latents
are decoded once and cached as 8-bit RGB reference frames.  For every trained
Predictor, chunk 0 is generated with FFFF (there is no previous chunk), while
chunks 1..6 use Full-Predictor-Predictor-Full.  PSNR, Gaussian SSIM, and
AlexNet LPIPS are computed frame by frame against the matching FFFF video.
"""

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="2")
    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
import torch.nn.functional as F
from omegaconf import OmegaConf
from safetensors import safe_open
from safetensors.torch import load_file, save_file
from torchvision.io import write_video

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

from pipeline import CausalInferencePipeline
from predictor_training.offline_data import TOKENS_PER_CHUNK
from predictor_training.single_block import (
    SingleBlockPredictor,
    initialize_predictor_block,
)
from scripts.run_single_block_init_sweep import hidden_to_flow
from utils.misc import set_seed
from utils.wan_wrapper import WanDiffusionWrapper, WanVAEWrapper
from wan.modules.model import sinusoidal_embedding_1d


LATENT_CHANNELS = 16
LATENT_HEIGHT = 60
LATENT_WIDTH = 104
FRAMES_PER_CHUNK = 3
NUM_CHUNKS = 7
NUM_DENOISING_STEPS = 4
PIXEL_FRAMES_FIRST_CHUNK = 1 + 4 * (FRAMES_PER_CHUNK - 1)
DEFAULT_PROMPT_IDS = list(range(80, 100))


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(
        "--output_dir",
        type=Path,
        default=Path("outputs/single_block_fppf_eval"),
    )
    parser.add_argument(
        "--prompt_ids", type=int, nargs="*", default=DEFAULT_PROMPT_IDS
    )
    parser.add_argument(
        "--experiments",
        nargs="*",
        default=None,
        help="Experiment directory names. Omit to evaluate all summary rows.",
    )
    parser.add_argument("--max_prompts", type=int, default=None)
    parser.add_argument("--max_experiments", 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(
        "--verify_ffff",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="Re-run FFFF once and compare its latent exactly to offline data.",
    )
    parser.add_argument(
        "--skip_lpips",
        action=argparse.BooleanOptionalAction,
        default=False,
        help="Only for quick diagnostics; formal evaluation should keep LPIPS.",
    )
    parser.add_argument(
        "--rebuild_references",
        action=argparse.BooleanOptionalAction,
        default=False,
    )
    parser.add_argument(
        "--save_videos",
        action=argparse.BooleanOptionalAction,
        default=False,
        help="Save each FPPF prediction as a 16-fps H.264 MP4 for VBench.",
    )
    args = parser.parse_args()
    if args.metric_batch_size < 1:
        parser.error("--metric_batch_size must be positive")
    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]")
    return args


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


def atomic_json(path: Path, value: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(
        json.dumps(value, indent=2, ensure_ascii=False, allow_nan=True) + "\n",
        encoding="utf-8",
    )
    os.replace(temporary, path)


def atomic_safetensors(
    path: Path, tensors: dict[str, torch.Tensor], metadata: dict[str, str]
) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".tmp")
    save_file(tensors, temporary, metadata=metadata)
    os.replace(temporary, path)


def load_prompt_metadata(dataset_root: Path, prompt_id: int) -> dict[str, Any]:
    path = dataset_root / f"prompt_{prompt_id:04d}" / "metadata.json"
    return json.loads(path.read_text(encoding="utf-8"))


def load_ffff_latent(dataset_root: Path, prompt_id: int) -> torch.Tensor:
    path = dataset_root / f"prompt_{prompt_id:04d}" / "trajectory.safetensors"
    with safe_open(path, framework="pt", device="cpu") as handle:
        chunks = [
            handle.get_tensor(f"chunk_{chunk:02d}_clean_latent")
            for chunk in range(NUM_CHUNKS)
        ]
    return torch.cat(chunks, dim=1).contiguous()


def pixels_to_u8(video: torch.Tensor) -> torch.Tensor:
    """Convert [1,T,3,H,W] pixels in [-1,1] to CPU uint8 frames."""
    return (
        ((video.squeeze(0).float() + 1.0) * 127.5)
        .round_()
        .clamp_(0, 255)
        .to(device="cpu", dtype=torch.uint8)
        .contiguous()
    )


def save_mp4(frames: torch.Tensor, path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    write_video(
        str(path), frames.permute(0, 2, 3, 1), fps=16,
        video_codec="libx264", options={"crf": "18"},
    )


@torch.inference_mode()
def prepare_reference_frames(
    *,
    vae: WanVAEWrapper,
    dataset_root: Path,
    output_dir: Path,
    prompt_ids: list[int],
    device: torch.device,
    rebuild: bool,
) -> None:
    reference_dir = output_dir / "ffff_reference_frames"
    reference_dir.mkdir(parents=True, exist_ok=True)
    for offset, prompt_id in enumerate(prompt_ids, start=1):
        destination = reference_dir / f"prompt_{prompt_id:04d}.safetensors"
        if destination.exists() and not rebuild:
            print(
                f"[reference] {offset}/{len(prompt_ids)} prompt={prompt_id} cached",
                flush=True,
            )
            continue
        latent = load_ffff_latent(dataset_root, prompt_id).to(
            device=device, dtype=torch.bfloat16
        )
        started = time.perf_counter()
        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
            pixels = vae.decode_to_pixel(latent, use_cache=False)
        frames = pixels_to_u8(pixels)
        atomic_safetensors(
            destination,
            {"frames": frames},
            {
                "reference": "FFFF",
                "prompt_id": str(prompt_id),
                "range": "uint8_0_255",
                "layout": "TCHW",
            },
        )
        if hasattr(vae.model, "clear_cache"):
            vae.model.clear_cache()
        del latent, pixels, frames
        torch.cuda.empty_cache()
        print(
            f"[reference] {offset}/{len(prompt_ids)} prompt={prompt_id} "
            f"decoded={time.perf_counter() - started:.1f}s",
            flush=True,
        )


def load_reference_frames(output_dir: Path, prompt_id: int) -> torch.Tensor:
    path = output_dir / "ffff_reference_frames" / f"prompt_{prompt_id:04d}.safetensors"
    with safe_open(path, framework="pt", device="cpu") as handle:
        return handle.get_tensor("frames")


def discover_experiments(
    sweep_dir: Path,
    requested: list[str] | None,
    max_experiments: int | None,
) -> list[dict[str, Any]]:
    summary_path = sweep_dir / "summary.csv"
    with summary_path.open("r", encoding="utf-8", newline="") as handle:
        rows = list(csv.DictReader(handle))
    by_name = {row["name"]: row for row in rows}
    names = list(by_name) if requested is None else requested
    unknown = [name for name in names if name not in by_name]
    if unknown:
        raise KeyError(f"Unknown sweep experiments: {unknown}")
    if max_experiments is not None:
        names = names[:max_experiments]

    experiments = []
    for name in names:
        run_dir = sweep_dir / name
        config = json.loads((run_dir / "config.json").read_text(encoding="utf-8"))
        weights = run_dir / "predictor_final.safetensors"
        if not weights.exists():
            raise FileNotFoundError(weights)
        row = by_name[name]
        experiment = {
                "name": name,
                "initialization_method": config["initialization_method"],
                "source_layer": int(config["source_layer"]),
                "weights": weights,
                "gate_mode": config.get("gate_mode", "baseline"),
                "gate_hidden_dim": int(config.get("gate_hidden_dim", 128)),
                "gate_initial_bias": float(
                    config.get("gate_initial_bias", 4.6)
                ),
                "gate_floor": float(config.get("gate_floor", 0.0)),
                "constant_gate": float(config.get("constant_gate", 1.0)),
                "gate_override": config.get("gate_override"),
                "offline_final_val_flow_mse": float(row["final_val_flow_mse"]),
                "offline_final_val_hidden_mse": float(row["final_val_hidden_mse"]),
            }
        experiments.append(experiment)
        if experiment["gate_mode"] == "learned":
            training_metrics = json.loads(
                (run_dir / "metrics.json").read_text(encoding="utf-8")
            )
            gate_mean = float(training_metrics["evaluations"][-1]["gate_mean"])
            experiments.append(
                {
                    **experiment,
                    "name": f"{name}_constant_mean",
                    "gate_override": gate_mean,
                    "constant_gate": gate_mean,
                }
            )
    return experiments


def build_pipeline(
    config: Any,
    checkpoint_path: Path,
    vae: WanVAEWrapper,
    device: torch.device,
) -> CausalInferencePipeline:
    generator = WanDiffusionWrapper(
        **getattr(config, "model_kwargs", {}), is_causal=True
    )
    pipeline = CausalInferencePipeline(
        config,
        device=device,
        generator=generator,
        text_encoder=torch.nn.Identity(),
        vae=vae,
    )
    checkpoint = torch.load(
        checkpoint_path, map_location="cpu", weights_only=False, mmap=True
    )
    if set(checkpoint) != {"generator_ema"}:
        raise KeyError(f"Unexpected Teacher checkpoint keys: {sorted(checkpoint)}")
    pipeline.generator.load_state_dict(checkpoint["generator_ema"], strict=True)
    del checkpoint
    pipeline.to(dtype=torch.bfloat16)
    pipeline.generator.to(device=device)
    pipeline.eval()
    pipeline.generator.requires_grad_(False)
    return pipeline


def reset_kv_and_load_cross_cache(
    pipeline: CausalInferencePipeline,
    dataset_root: Path,
    prompt_id: int,
    device: torch.device,
) -> None:
    if pipeline.kv_cache1 is None:
        pipeline._initialize_kv_cache(1, torch.bfloat16, device)
        pipeline._initialize_crossattn_cache(1, torch.bfloat16, device)
    for cache in pipeline.kv_cache1:
        cache["global_end_index"].zero_()
        cache["local_end_index"].zero_()

    cross_path = (
        dataset_root / f"prompt_{prompt_id:04d}" / "cross_attention.safetensors"
    )
    with safe_open(cross_path, framework="pt", device="cpu") as handle:
        for layer, cache in enumerate(pipeline.crossattn_cache):
            cache["k"] = handle.get_tensor(f"block_{layer:02d}_k").to(
                device=device, dtype=torch.bfloat16
            )
            cache["v"] = handle.get_tensor(f"block_{layer:02d}_v").to(
                device=device, dtype=torch.bfloat16
            )
            cache["is_init"] = True


class FinalHiddenCapture:
    def __init__(self, teacher: torch.nn.Module) -> None:
        self.enabled = False
        self.value: torch.Tensor | None = None
        self.handle = teacher.head.register_forward_pre_hook(self._hook)

    def close(self) -> None:
        self.handle.remove()

    def _hook(
        self, _module: torch.nn.Module, inputs: tuple[torch.Tensor, ...]
    ) -> None:
        if self.enabled:
            if self.value is not None:
                raise RuntimeError("Teacher head was called twice in one Full step")
            self.value = inputs[0].detach()

    def start(self) -> None:
        self.value = None
        self.enabled = True

    def finish(self) -> torch.Tensor:
        self.enabled = False
        if self.value is None:
            raise RuntimeError("Teacher final hidden was not captured")
        value = self.value
        self.value = None
        return value


def load_predictor(
    teacher: torch.nn.Module,
    experiment: dict[str, Any],
    device: torch.device,
) -> SingleBlockPredictor:
    source_layer = int(experiment["source_layer"])
    block = initialize_predictor_block(
        teacher.blocks[source_layer], "teacher_full"
    )
    predictor = SingleBlockPredictor(
        block=block,
        dim=teacher.dim,
        gradient_checkpointing=False,
        input_variant=experiment.get("input_variant", "self_forcing"),
        gate_mode=experiment.get("gate_mode", "baseline"),
        gate_hidden_dim=int(experiment.get("gate_hidden_dim", 128)),
        gate_initial_bias=float(experiment.get("gate_initial_bias", 4.6)),
        gate_floor=float(experiment.get("gate_floor", 0.0)),
        constant_gate=float(experiment.get("constant_gate", 1.0)),
        atc_previous_scope=experiment.get("atc_previous_scope", "chunk"),
        atc_freq_dim=int(experiment.get("atc_freq_dim", 256)),
        atc_mlp_hidden_dim=int(experiment.get("atc_mlp_hidden_dim", 3072)),
        atc_gate_hidden_dim=int(experiment.get("atc_gate_hidden_dim", 512)),
        atc_transport_residual_scale=float(
            experiment.get("atc_transport_residual_scale", 0.1)
        ),
        atc_gate_initial_probability=float(
            experiment.get("atc_gate_initial_probability", 0.3)
        ),
        atc_collect_diagnostics=bool(
            experiment.get("atc_collect_diagnostics", False)
        ),
    )
    state = load_file(str(experiment["weights"]), device="cpu")
    predictor.load_state_dict(state, strict=True)
    if experiment.get("gate_override") is not None:
        predictor.fusion.gate_override = float(experiment["gate_override"])
    predictor.to(device=device)
    predictor.eval().requires_grad_(False)
    return predictor


@torch.inference_mode()
def predictor_step(
    *,
    predictor: SingleBlockPredictor,
    teacher: torch.nn.Module,
    noisy_input: torch.Tensor,
    timestep: torch.Tensor,
    anchor_hidden: torch.Tensor,
    previous_hidden: torch.Tensor,
    history_cache: dict[str, torch.Tensor],
    cross_cache: dict[str, torch.Tensor],
    current_start: int,
    anchor_timestep: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        current_tokens = teacher.patch_embedding(
            noisy_input.permute(0, 2, 1, 3, 4)
        ).flatten(2).transpose(1, 2)
        time_embedding = teacher.time_embedding(
            sinusoidal_embedding_1d(
                teacher.freq_dim, timestep.flatten()
            ).type_as(current_tokens)
        )
        timestep_modulation = teacher.time_projection(
            time_embedding
        ).unflatten(1, (6, teacher.dim)).unflatten(
            dim=0, sizes=timestep.shape
        )
        head_embedding = time_embedding.unflatten(
            dim=0, sizes=timestep.shape
        ).unsqueeze(2)
        condition_per_frame = time_embedding.unflatten(
            dim=0, sizes=timestep.shape
        )
        condition_tokens = (
            condition_per_frame[:, :, None, :]
            .expand(
                timestep.shape[0],
                timestep.shape[1],
                30 * 52,
                teacher.dim,
            )
            .reshape(timestep.shape[0], -1, teacher.dim)
        )
        anchor_distance = None
        if predictor.input_variant == "atc":
            if anchor_timestep is None:
                raise ValueError("ATC inference requires anchor_timestep")
            anchor_distance = (
                timestep.float() - anchor_timestep.float()
            ).abs().mean(dim=1)
        grid_sizes = torch.tensor(
            [[FRAMES_PER_CHUNK, 30, 52]], dtype=torch.long, device="cpu"
        )
        history_length = int(history_cache["local_end_index"].item())
        pred_hidden = predictor(
            current_tokens=current_tokens,
            anchor_hidden=anchor_hidden,
            previous_hidden=previous_hidden,
            timestep_modulation=timestep_modulation,
            grid_sizes=grid_sizes,
            freqs=teacher.freqs,
            history_k=history_cache["k"][:, :history_length],
            history_v=history_cache["v"][:, :history_length],
            cross_k=cross_cache["k"],
            cross_v=cross_cache["v"],
            current_start=current_start,
            condition_tokens=condition_tokens,
            anchor_distance=anchor_distance,
        )
        pred_flow = hidden_to_flow(
            pred_hidden, head_embedding, grid_sizes, teacher
        )
    return pred_hidden, pred_flow, current_tokens


@torch.inference_mode()
def generate_rollout(
    *,
    pipeline: CausalInferencePipeline,
    dataset_root: Path,
    prompt_id: int,
    generation_seed: int,
    device: torch.device,
    predictor: SingleBlockPredictor | None,
    source_layer: int | None,
    schedule: str,
) -> tuple[torch.Tensor, dict[str, float | int]]:
    if schedule not in {"FFFF", "FPPF"}:
        raise ValueError(schedule)
    if schedule == "FPPF" and (predictor is None or source_layer is None):
        raise ValueError("FPPF requires a Predictor and source layer")

    reset_kv_and_load_cross_cache(pipeline, dataset_root, prompt_id, device)
    set_seed(generation_seed)
    noise = torch.randn(
        1,
        NUM_CHUNKS * FRAMES_PER_CHUNK,
        LATENT_CHANNELS,
        LATENT_HEIGHT,
        LATENT_WIDTH,
        dtype=torch.bfloat16,
        device=device,
    )
    teacher = pipeline.generator.model
    text_dim = int(teacher.text_embedding[0].in_features)
    conditional_dict = {
        "prompt_embeds": torch.zeros(
            1, 1, text_dim, 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
    capture = FinalHiddenCapture(teacher)
    full_calls = 0
    predictor_calls = 0
    started = time.perf_counter()

    try:
        for chunk in range(NUM_CHUNKS):
            noisy_input = noise[
                :, chunk * FRAMES_PER_CHUNK : (chunk + 1) * FRAMES_PER_CHUNK
            ]
            current_hidden: list[torch.Tensor | None] = [None] * 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, FRAMES_PER_CHUNK], dtype=torch.int64, device=device
                ) * current_timestep
                use_predictor = schedule == "FPPF" and chunk > 0 and step in {1, 2}

                if use_predictor:
                    anchor_hidden = current_hidden[step - 1]
                    assert anchor_hidden is not None
                    assert previous_chunk_hidden is not None
                    previous_hidden = previous_chunk_hidden[step]
                    assert previous_hidden is not None
                    history = pipeline.kv_cache1[int(source_layer)]
                    cross = pipeline.crossattn_cache[int(source_layer)]
                    pred_hidden, flow, _ = 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 * TOKENS_PER_CHUNK,
                        anchor_timestep=(
                            torch.ones_like(timestep) * timesteps[step - 1]
                        ),
                    )
                    denoised_pred = pipeline.generator._convert_flow_pred_to_x0(
                        flow_pred=flow.flatten(0, 1),
                        xt=noisy_input.flatten(0, 1),
                        timestep=timestep.flatten(0, 1),
                    ).unflatten(0, flow.shape[:2])
                    current_hidden[step] = pred_hidden
                    predictor_calls += 1
                else:
                    capture.start()
                    _, denoised_pred = pipeline.generator(
                        noisy_image_or_video=noisy_input,
                        conditional_dict=conditional_dict,
                        timestep=timestep,
                        kv_cache=pipeline.kv_cache1,
                        crossattn_cache=pipeline.crossattn_cache,
                        current_start=chunk * TOKENS_PER_CHUNK,
                    )
                    current_hidden[step] = capture.finish()
                    full_calls += 1

                if step < NUM_DENOISING_STEPS - 1:
                    next_timestep = timesteps[step + 1]
                    denoised_flat = denoised_pred.flatten(0, 1)
                    noisy_input = pipeline.scheduler.add_noise(
                        denoised_flat,
                        torch.randn_like(denoised_flat),
                        next_timestep
                        * torch.ones(
                            [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("Denoising loop 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=conditional_dict,
                timestep=context_timestep,
                kv_cache=pipeline.kv_cache1,
                crossattn_cache=pipeline.crossattn_cache,
                current_start=chunk * 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,
    }


def gaussian_kernel(
    device: torch.device, dtype: torch.dtype, channels: int = 3
) -> torch.Tensor:
    coordinates = torch.arange(11, device=device, dtype=dtype) - 5
    kernel_1d = torch.exp(-(coordinates.square()) / (2 * 1.5**2))
    kernel_1d /= kernel_1d.sum()
    kernel_2d = torch.outer(kernel_1d, kernel_1d)
    return kernel_2d.expand(channels, 1, 11, 11).contiguous()


def ssim_per_frame(
    reference: torch.Tensor, prediction: torch.Tensor, kernel: torch.Tensor
) -> torch.Tensor:
    channels = reference.shape[1]
    mu_x = F.conv2d(reference, kernel, groups=channels)
    mu_y = F.conv2d(prediction, kernel, groups=channels)
    mu_x2 = mu_x.square()
    mu_y2 = mu_y.square()
    mu_xy = mu_x * mu_y
    sigma_x2 = F.conv2d(reference.square(), kernel, groups=channels) - mu_x2
    sigma_y2 = F.conv2d(prediction.square(), kernel, groups=channels) - mu_y2
    sigma_xy = F.conv2d(reference * prediction, kernel, groups=channels) - mu_xy
    c1 = 0.01**2
    c2 = 0.03**2
    score = ((2 * mu_xy + c1) * (2 * sigma_xy + c2)) / (
        (mu_x2 + mu_y2 + c1) * (sigma_x2 + sigma_y2 + c2)
    )
    return score.mean(dim=(1, 2, 3))


@torch.inference_mode()
def frame_metrics(
    *,
    reference_u8: torch.Tensor,
    prediction_u8: torch.Tensor,
    lpips_model: torch.nn.Module | None,
    batch_size: int,
    device: torch.device,
) -> dict[str, Any]:
    if reference_u8.shape != prediction_u8.shape:
        raise ValueError(
            f"Reference/prediction shapes differ: {reference_u8.shape}, "
            f"{prediction_u8.shape}"
        )
    kernel = gaussian_kernel(device, torch.float32)
    psnr_values: list[float] = []
    mse_values: list[float] = []
    ssim_values: list[float] = []
    lpips_values: list[float] = []
    for start in range(0, reference_u8.shape[0], batch_size):
        end = min(start + batch_size, reference_u8.shape[0])
        reference = reference_u8[start:end].to(
            device=device, dtype=torch.float32
        ) / 255.0
        prediction = prediction_u8[start:end].to(
            device=device, dtype=torch.float32
        ) / 255.0
        mse = (reference - prediction).square().mean(dim=(1, 2, 3))
        psnr = -10.0 * torch.log10(mse.clamp_min(1e-12))
        ssim = ssim_per_frame(reference, prediction, kernel)
        mse_values.extend(float(value) for value in mse.cpu())
        psnr_values.extend(float(value) for value in psnr.cpu())
        ssim_values.extend(float(value) for value in ssim.cpu())
        if lpips_model is not None:
            distance = lpips_model(
                reference.mul(2).sub(1), prediction.mul(2).sub(1)
            ).flatten()
            lpips_values.extend(float(value) for value in distance.cpu())
        del reference, prediction, mse, psnr, ssim
    global_mse = sum(mse_values) / len(mse_values)
    rollout_mse = sum(mse_values[PIXEL_FRAMES_FIRST_CHUNK:]) / len(
        mse_values[PIXEL_FRAMES_FIRST_CHUNK:]
    )
    return {
        "mse_per_frame": mse_values,
        "psnr_per_frame": psnr_values,
        "ssim_per_frame": ssim_values,
        "lpips_per_frame": lpips_values,
        "pixel_mse": global_mse,
        "psnr": -10.0 * math.log10(max(global_mse, 1e-12)),
        "psnr_frame_mean": sum(psnr_values) / len(psnr_values),
        "ssim": sum(ssim_values) / len(ssim_values),
        "lpips": (
            sum(lpips_values) / len(lpips_values)
            if lpips_values
            else None
        ),
        "rollout_start_frame": PIXEL_FRAMES_FIRST_CHUNK,
        "rollout_pixel_mse": rollout_mse,
        "rollout_psnr": -10.0 * math.log10(max(rollout_mse, 1e-12)),
        "rollout_ssim": sum(ssim_values[PIXEL_FRAMES_FIRST_CHUNK:])
        / len(ssim_values[PIXEL_FRAMES_FIRST_CHUNK:]),
        "rollout_lpips": (
            sum(lpips_values[PIXEL_FRAMES_FIRST_CHUNK:])
            / len(lpips_values[PIXEL_FRAMES_FIRST_CHUNK:])
            if lpips_values
            else None
        ),
        "num_frames": len(psnr_values),
    }


def mean_std(values: list[float]) -> tuple[float, float]:
    mean = sum(values) / len(values)
    variance = sum((value - mean) ** 2 for value in values) / len(values)
    return mean, math.sqrt(variance)


def aggregate_prompt_results(
    experiment: dict[str, Any], prompt_results: list[dict[str, Any]]
) -> dict[str, Any]:
    mse_frames = [
        value
        for result in prompt_results
        for value in result["mse_per_frame"]
    ]
    psnr_frames = [
        value
        for result in prompt_results
        for value in result["psnr_per_frame"]
    ]
    ssim_frames = [
        value
        for result in prompt_results
        for value in result["ssim_per_frame"]
    ]
    lpips_frames = [
        value
        for result in prompt_results
        for value in result["lpips_per_frame"]
    ]
    rollout_mse_frames = [
        value
        for result in prompt_results
        for value in result["mse_per_frame"][PIXEL_FRAMES_FIRST_CHUNK:]
    ]
    rollout_ssim_frames = [
        value
        for result in prompt_results
        for value in result["ssim_per_frame"][PIXEL_FRAMES_FIRST_CHUNK:]
    ]
    rollout_lpips_frames = [
        value
        for result in prompt_results
        for value in result["lpips_per_frame"][PIXEL_FRAMES_FIRST_CHUNK:]
    ]
    pixel_mse = sum(mse_frames) / len(mse_frames)
    psnr_frame_mean, psnr_std = mean_std(psnr_frames)
    ssim, ssim_std = mean_std(ssim_frames)
    if lpips_frames:
        lpips_mean, lpips_std = mean_std(lpips_frames)
    else:
        lpips_mean, lpips_std = None, None
    rollout_pixel_mse = sum(rollout_mse_frames) / len(rollout_mse_frames)
    rollout_ssim = sum(rollout_ssim_frames) / len(rollout_ssim_frames)
    rollout_lpips = (
        sum(rollout_lpips_frames) / len(rollout_lpips_frames)
        if rollout_lpips_frames
        else None
    )
    return {
        "status": "complete",
        "name": experiment["name"],
        "initialization_method": experiment["initialization_method"],
        "source_layer": experiment["source_layer"],
        "offline_final_val_flow_mse": experiment["offline_final_val_flow_mse"],
        "offline_final_val_hidden_mse": experiment[
            "offline_final_val_hidden_mse"
        ],
        "schedule": "chunk0=FFFF; chunks1-6=FPPF",
        "reference": "matching FFFF, same prompt and seed",
        "pixel_domain": "VAE-decoded RGB, rounded to uint8",
        "aggregation": (
            "PSNR from global pixel MSE; SSIM/LPIPS mean over decoded frames"
        ),
        "num_prompts": len(prompt_results),
        "num_frames": len(psnr_frames),
        "pixel_mse": pixel_mse,
        "psnr": -10.0 * math.log10(max(pixel_mse, 1e-12)),
        "psnr_frame_mean": psnr_frame_mean,
        "psnr_frame_std": psnr_std,
        "ssim": ssim,
        "ssim_frame_std": ssim_std,
        "lpips": lpips_mean,
        "lpips_frame_std": lpips_std,
        "rollout_start_frame": PIXEL_FRAMES_FIRST_CHUNK,
        "rollout_num_frames": len(rollout_mse_frames),
        "rollout_pixel_mse": rollout_pixel_mse,
        "rollout_psnr": -10.0
        * math.log10(max(rollout_pixel_mse, 1e-12)),
        "rollout_ssim": rollout_ssim,
        "rollout_lpips": rollout_lpips,
        "mean_generation_time_s": sum(
            result["generation_time_s"] for result in prompt_results
        )
        / len(prompt_results),
        "full_calls_per_prompt": prompt_results[0]["full_calls"],
        "predictor_calls_per_prompt": prompt_results[0]["predictor_calls"],
        "prompt_ids": [result["prompt_id"] for result in prompt_results],
    }


def write_summary(
    output_dir: Path, experiments: list[dict[str, Any]]
) -> None:
    rows: list[dict[str, Any]] = []
    for experiment in experiments:
        path = output_dir / experiment["name"] / "metrics.json"
        if not path.exists():
            continue
        metrics = json.loads(path.read_text(encoding="utf-8"))
        if metrics.get("status") != "complete":
            continue
        rows.append(
            {
                "name": metrics["name"],
                "initialization_method": metrics["initialization_method"],
                "source_layer": metrics["source_layer"],
                "num_prompts": metrics["num_prompts"],
                "num_frames": metrics["num_frames"],
                "psnr": metrics["psnr"],
                "ssim": metrics["ssim"],
                "lpips": metrics["lpips"],
                "rollout_psnr": metrics["rollout_psnr"],
                "rollout_ssim": metrics["rollout_ssim"],
                "rollout_lpips": metrics["rollout_lpips"],
                "offline_final_val_flow_mse": metrics[
                    "offline_final_val_flow_mse"
                ],
                "mean_generation_time_s": metrics["mean_generation_time_s"],
            }
        )
    rows.sort(key=lambda row: float(row["lpips"] or math.inf))
    fields = [
        "name",
        "initialization_method",
        "source_layer",
        "num_prompts",
        "num_frames",
        "psnr",
        "ssim",
        "lpips",
        "rollout_psnr",
        "rollout_ssim",
        "rollout_lpips",
        "offline_final_val_flow_mse",
        "mean_generation_time_s",
    ]
    destination = output_dir / "summary.csv"
    temporary = destination.with_suffix(".csv.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, destination)


def load_completed_prompt_results(
    run_dir: Path, prompt_ids: list[int]
) -> list[dict[str, Any]]:
    results = []
    for prompt_id in prompt_ids:
        path = run_dir / "per_prompt" / f"prompt_{prompt_id:04d}.json"
        if path.exists():
            results.append(json.loads(path.read_text(encoding="utf-8")))
    return results


def main() -> None:
    args = parse_args()
    args.config_path = resolve(args.config_path)
    args.checkpoint_path = resolve(args.checkpoint_path)
    args.dataset_root = resolve(args.dataset_root)
    args.sweep_dir = resolve(args.sweep_dir)
    args.output_dir = resolve(args.output_dir)
    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]
    experiments = discover_experiments(
        args.sweep_dir, args.experiments, args.max_experiments
    )
    device = torch.device("cuda")
    torch.set_grad_enabled(False)
    set_seed(args.generation_seed)

    config = OmegaConf.merge(
        OmegaConf.load(REPO_ROOT / "configs/default_config.yaml"),
        OmegaConf.load(args.config_path),
    )
    manifest = {
        "status": "running",
        "gpu": str(args.gpu),
        "config_path": str(args.config_path),
        "checkpoint_path": str(args.checkpoint_path),
        "dataset_root": str(args.dataset_root),
        "sweep_dir": str(args.sweep_dir),
        "prompt_ids": prompt_ids,
        "generation_seed_reset_per_prompt": args.generation_seed,
        "experiments": [item["name"] for item in experiments],
        "fppf_definition": "chunk0=FFFF; chunks1-6=FPPF",
        "reference": "offline FFFF clean latents from the same prompt/seed",
        "metrics": {
            "psnr": "RGB PSNR from global pixel MSE",
            "ssim": "11x11 Gaussian sigma=1.5 RGB SSIM, then frame mean",
            "lpips": "AlexNet LPIPS on RGB [-1,1], then frame mean",
            "pixel_quantization": "both inputs rounded to uint8",
            "rollout_only": (
                "also reported for decoded frames 9..80 after the FFFF-only "
                "first chunk"
            ),
        },
    }
    atomic_json(args.output_dir / "manifest.json", manifest)

    print("[setup] loading VAE and preparing FFFF reference frames", flush=True)
    vae = WanVAEWrapper().to(device=device, dtype=torch.bfloat16).eval()
    prepare_reference_frames(
        vae=vae,
        dataset_root=args.dataset_root,
        output_dir=args.output_dir,
        prompt_ids=prompt_ids,
        device=device,
        rebuild=args.rebuild_references,
    )

    print("[setup] loading frozen generator_ema", flush=True)
    pipeline = build_pipeline(config, args.checkpoint_path, vae, device)
    teacher = pipeline.generator.model
    lpips_model = None
    if not args.skip_lpips:
        print("[setup] loading AlexNet LPIPS", flush=True)
        lpips_model = lpips.LPIPS(net="alex", verbose=False).to(device).eval()
        lpips_model.requires_grad_(False)

    if args.verify_ffff:
        prompt_id = prompt_ids[0]
        print(f"[verify] reproducing FFFF prompt={prompt_id}", flush=True)
        reproduced, counts = generate_rollout(
            pipeline=pipeline,
            dataset_root=args.dataset_root,
            prompt_id=prompt_id,
            generation_seed=args.generation_seed,
            device=device,
            predictor=None,
            source_layer=None,
            schedule="FFFF",
        )
        expected = load_ffff_latent(args.dataset_root, prompt_id).to(
            device=device, dtype=torch.bfloat16
        )
        difference = reproduced.float() - expected.float()
        verification = {
            "prompt_id": prompt_id,
            "max_abs_latent_error": float(difference.abs().max()),
            "latent_mse": float(difference.square().mean()),
            **counts,
        }
        atomic_json(args.output_dir / "ffff_reproduction.json", verification)
        print(f"[verify] {verification}", flush=True)
        if verification["max_abs_latent_error"] > 1e-3:
            raise RuntimeError(
                "FFFF reproduction differs from offline reference; refusing "
                "to evaluate FPPF with unmatched randomness/caches"
            )
        del reproduced, expected, difference
        torch.cuda.empty_cache()

    for experiment_index, experiment in enumerate(experiments, start=1):
        run_dir = args.output_dir / experiment["name"]
        run_dir.mkdir(parents=True, exist_ok=True)
        metrics_path = run_dir / "metrics.json"
        if metrics_path.exists():
            existing = json.loads(metrics_path.read_text(encoding="utf-8"))
            if (
                existing.get("status") == "complete"
                and existing.get("prompt_ids") == prompt_ids
                and (args.skip_lpips or existing.get("lpips") is not None)
            ):
                print(
                    f"[run] {experiment_index}/{len(experiments)} "
                    f"skip complete {experiment['name']}",
                    flush=True,
                )
                continue

        print(
            f"[run] {experiment_index}/{len(experiments)} "
            f"{experiment['name']} source={experiment['source_layer']}",
            flush=True,
        )
        predictor = load_predictor(teacher, experiment, device)
        existing_results = {
            result["prompt_id"]: result
            for result in load_completed_prompt_results(run_dir, prompt_ids)
        }

        for prompt_index, prompt_id in enumerate(prompt_ids, start=1):
            video_path = run_dir / "videos" / f"prompt_{prompt_id:04d}.mp4"
            if prompt_id in existing_results and (
                not args.save_videos or video_path.exists()
            ):
                print(
                    f"[prompt] {experiment['name']} {prompt_index}/{len(prompt_ids)} "
                    f"id={prompt_id} cached",
                    flush=True,
                )
                continue
            started = time.perf_counter()
            latent, counts = generate_rollout(
                pipeline=pipeline,
                dataset_root=args.dataset_root,
                prompt_id=prompt_id,
                generation_seed=args.generation_seed,
                device=device,
                predictor=predictor,
                source_layer=experiment["source_layer"],
                schedule="FPPF",
            )
            with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
                pixels = vae.decode_to_pixel(latent, use_cache=False)
            prediction_u8 = pixels_to_u8(pixels)
            if args.save_videos:
                save_mp4(prediction_u8, video_path)
            reference_u8 = load_reference_frames(args.output_dir, prompt_id)
            metrics = frame_metrics(
                reference_u8=reference_u8,
                prediction_u8=prediction_u8,
                lpips_model=lpips_model,
                batch_size=args.metric_batch_size,
                device=device,
            )
            prompt_result = {
                "prompt_id": prompt_id,
                "prompt": load_prompt_metadata(args.dataset_root, prompt_id)[
                    "prompt"
                ],
                **counts,
                **metrics,
                "total_time_s": time.perf_counter() - started,
            }
            atomic_json(
                run_dir / "per_prompt" / f"prompt_{prompt_id:04d}.json",
                prompt_result,
            )
            existing_results[prompt_id] = prompt_result
            print(
                f"[prompt] {experiment['name']} {prompt_index}/{len(prompt_ids)} "
                f"id={prompt_id} psnr={metrics['psnr']:.4f} "
                f"ssim={metrics['ssim']:.6f} "
                f"lpips={metrics['lpips']} "
                f"time={prompt_result['total_time_s']:.1f}s",
                flush=True,
            )
            if hasattr(vae.model, "clear_cache"):
                vae.model.clear_cache()
            del latent, pixels, prediction_u8, reference_u8
            torch.cuda.empty_cache()

        prompt_results = [existing_results[prompt_id] for prompt_id in prompt_ids]
        aggregate = aggregate_prompt_results(experiment, prompt_results)
        atomic_json(metrics_path, aggregate)
        write_summary(args.output_dir, experiments)
        print(
            f"[result] {experiment['name']} psnr={aggregate['psnr']:.4f} "
            f"ssim={aggregate['ssim']:.6f} lpips={aggregate['lpips']}",
            flush=True,
        )
        del predictor
        torch.cuda.empty_cache()

    manifest["status"] = "complete"
    atomic_json(args.output_dir / "manifest.json", manifest)
    write_summary(args.output_dir, experiments)
    print(
        f"[complete] {len(experiments)} experiments -> "
        f"{args.output_dir / 'summary.csv'}",
        flush=True,
    )


if __name__ == "__main__":
    main()