scripts/evaluate_layer17_dynamic_gate.py

#!/usr/bin/env python3
"""Evaluate chunk-aware dynamic gating for the frozen Layer-17 Predictor."""

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
from safetensors.torch import load_file

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

from predictor_training.confidence import PredictorConfidenceHead
from scripts import evaluate_layer17_chunk_impact as impact_eval
from scripts import evaluate_single_block_fppf as base
from scripts.run_single_block_init_sweep import hidden_to_flow
from utils.misc import set_seed
from utils.wan_wrapper import WanVAEWrapper
from wan.modules.model import sinusoidal_embedding_1d


BETAS = (0.0, 1.0, 1.5, 2.0)
TARGET_ACCEPTS = (4, 6, 8, 10)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--gpu", default=PHYSICAL_GPU)
    parser.add_argument("--mode", choices=("smoke", "validation", "test"), required=True)
    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(
        "--predictor_weights",
        type=Path,
        default=None,
        help=(
            "Optional direct Layer-17 Predictor weights. When supplied, this "
            "takes precedence over teacher_layer_17 in --sweep_dir."
        ),
    )
    parser.add_argument(
        "--confidence_weights", type=Path,
        default=Path(
            "outputs/layer17_confidence_teacher_forced_20260830/"
            "confidence_best.safetensors"
        ),
    )
    parser.add_argument(
        "--validation_predictions", type=Path,
        default=Path(
            "outputs/layer17_confidence_teacher_forced_20260830/"
            "validation_predictions.csv"
        ),
    )
    parser.add_argument(
        "--reference_root", type=Path, default=Path("outputs/single_block_fppf_eval")
    )
    parser.add_argument(
        "--output_root", type=Path,
        default=Path("outputs/layer17_dynamic_gate_20260830"),
    )
    parser.add_argument("--generation_seed", type=int, default=0)
    parser.add_argument("--metric_batch_size", type=int, default=4)
    parser.add_argument(
        "--candidate_steps", type=int, nargs="+", choices=(1, 2, 3),
        default=[1, 2],
    )
    parser.add_argument("--target_accepts", type=int, nargs="*", default=None)
    parser.add_argument("--selected_path", type=Path, default=None)
    parser.add_argument(
        "--config_names", nargs="*", default=None,
        help="Optional exact configuration names to run in validation/test mode.",
    )
    parser.add_argument(
        "--prompt_ids", type=int, nargs="*", default=None,
        help="Optional prompt shard; shard-level CSV/manifest files get a GPU suffix.",
    )
    parser.add_argument("--save_videos", action="store_true")
    parser.add_argument("--overwrite", action="store_true")
    parser.add_argument(
        "--skip_lpips", action=argparse.BooleanOptionalAction, default=False
    )
    args = parser.parse_args()
    for name in (
        "config_path", "checkpoint_path", "dataset_root", "sweep_dir",
        "predictor_weights",
        "confidence_weights", "validation_predictions", "reference_root",
        "output_root", "selected_path",
    ):
        value = getattr(args, name)
        if value is None:
            continue
        path = value.expanduser()
        setattr(args, name, path.resolve() if path.is_absolute() else (REPO_ROOT / path).resolve())
    args.candidate_steps = sorted(set(args.candidate_steps))
    max_accepts = 6 * len(args.candidate_steps)
    if args.target_accepts is None:
        args.target_accepts = (
            [4, 6, 8, 10] if len(args.candidate_steps) == 2
            else [6, 9, 12, 15]
        )
    if any(value < 1 or value >= max_accepts for value in args.target_accepts):
        parser.error(f"target accepts must be in [1, {max_accepts - 1}]")
    return args


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) + "\n", encoding="utf-8")
    os.replace(temporary, path)


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 quantile(values: list[float], fraction: float) -> float:
    ordered = sorted(values)
    position = fraction * (len(ordered) - 1)
    lower = int(math.floor(position))
    upper = int(math.ceil(position))
    if lower == upper:
        return ordered[lower]
    weight = position - lower
    return ordered[lower] * (1.0 - weight) + ordered[upper] * weight


def threshold_grid(
    predictions_path: Path,
    candidate_steps: list[int],
    targets: list[int],
) -> dict[tuple[float, int], float]:
    rows = list(csv.DictReader(predictions_path.open(encoding="utf-8")))
    expected = 10 * 6 * len(candidate_steps)
    if len(rows) != expected:
        raise ValueError(f"Expected {expected} validation predictions, got {len(rows)}")
    thresholds = {}
    for beta in BETAS:
        risks = []
        for row in rows:
            chunk = int(row["chunk"])
            alpha = (base.NUM_CHUNKS - 1 - chunk) / (base.NUM_CHUNKS - 2)
            local_error = float(row["predicted_hidden_nrmse"])
            risks.append(local_error * (1.0 + beta * alpha))
        max_accepts = 6 * len(candidate_steps)
        for target in targets:
            thresholds[(beta, target)] = quantile(risks, target / max_accepts)
    return thresholds


def dynamic_configs(
    predictions_path: Path,
    candidate_steps: list[int],
    targets: list[int],
) -> list[dict[str, Any]]:
    thresholds = threshold_grid(predictions_path, candidate_steps, targets)
    return [
        {
            "name": f"dynamic_b{str(beta).replace('.', 'p')}_k{target:02d}",
            "policy": "dynamic",
            "beta": beta,
            "target_accepts": target,
            "threshold": thresholds[(beta, target)],
        }
        for beta in BETAS
        for target in targets
    ]


def static_configs(targets: list[int]) -> list[dict[str, Any]]:
    return [
        {
            "name": f"static_late_k{target:02d}",
            "policy": "static_late",
            "beta": None,
            "target_accepts": target,
            "threshold": None,
        }
        for target in targets
    ]


def load_selected(path: Path) -> list[dict[str, Any]]:
    value = json.loads(path.read_text(encoding="utf-8"))
    return [
        {
            "name": row["config_name"],
            "policy": "dynamic",
            "beta": float(row["beta"]),
            "target_accepts": int(row["target_accepts"]),
            "threshold": float(row["threshold"]),
        }
        for row in value["selected_dynamic"]
    ]


@torch.no_grad()
def predictor_with_features(
    *,
    predictor: Any,
    teacher: Any,
    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(
            [[base.FRAMES_PER_CHUNK, 30, 52]], dtype=torch.long, device="cpu"
        )
        history_length = int(history_cache["local_end_index"].item())
        output = 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,
            return_features=True,
            condition_tokens=condition_tokens,
            anchor_distance=anchor_distance,
        )
        if not isinstance(output, tuple):
            raise RuntimeError("Predictor did not return internal features")
        pred_hidden, transformed = output
        pred_flow = hidden_to_flow(
            pred_hidden, head_embedding, grid_sizes, teacher
        )
    return pred_hidden, pred_flow, transformed


@torch.inference_mode()
def generate(
    *,
    pipeline: Any,
    dataset_root: Path,
    prompt_id: int,
    generation_seed: int,
    device: torch.device,
    predictor: Any,
    head: PredictorConfidenceHead,
    config: dict[str, Any],
    candidate_steps: list[int],
) -> tuple[torch.Tensor, dict[str, Any]]:
    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)
    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
        )
    }
    capture = base.FinalHiddenCapture(teacher)
    output_chunks: list[torch.Tensor] = []
    previous_chunk_hidden: list[torch.Tensor | None] | None = None
    decisions: list[dict[str, Any]] = []
    full_calls = 0
    predictor_calls = 0
    accepted_predictor_calls = 0
    timing_events: dict[str, list[tuple[torch.cuda.Event, torch.cuda.Event]]] = {
        "full_dit": [],
        "predictor": [],
        "confidence": [],
        "context_dit": [],
    }

    def start_timing() -> tuple[torch.cuda.Event, torch.cuda.Event]:
        start_event = torch.cuda.Event(enable_timing=True)
        end_event = torch.cuda.Event(enable_timing=True)
        start_event.record()
        return start_event, end_event

    def finish_timing(
        category: str,
        events: tuple[torch.cuda.Event, torch.cuda.Event],
    ) -> None:
        events[1].record()
        timing_events[category].append(events)

    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
                candidate = chunk > 0 and step in candidate_steps
                policy = str(config["policy"])
                run_predictor = False
                static_accept = False
                if candidate and policy == "dynamic":
                    run_predictor = True
                elif candidate and policy == "fppf":
                    run_predictor = True
                    static_accept = True
                elif candidate and policy == "static_late":
                    first_chunk = (
                        base.NUM_CHUNKS
                        - int(config["target_accepts"]) // len(candidate_steps)
                    )
                    static_accept = chunk >= first_chunk
                    run_predictor = static_accept

                accepted = False
                pred_hidden = None
                pred_x0 = None
                predicted_local_error = None
                risk = None
                alpha = None
                if run_predictor:
                    if previous_chunk_hidden is None:
                        raise RuntimeError("Previous chunk hidden is unavailable")
                    anchor_hidden = current_hidden[step - 1]
                    previous_hidden = previous_chunk_hidden[step]
                    if anchor_hidden is None or previous_hidden is None:
                        raise RuntimeError("Predictor inputs are unavailable")
                    predictor_events = start_timing()
                    pred_hidden, pred_flow, transformed = predictor_with_features(
                        predictor=predictor,
                        teacher=teacher,
                        noisy_input=noisy_input,
                        timestep=timestep,
                        anchor_hidden=anchor_hidden,
                        previous_hidden=previous_hidden,
                        history_cache=pipeline.kv_cache1[17],
                        cross_cache=pipeline.crossattn_cache[17],
                        current_start=chunk * base.TOKENS_PER_CHUNK,
                        anchor_timestep=(
                            torch.ones_like(timestep) * timesteps[step - 1]
                        ),
                    )
                    finish_timing("predictor", predictor_events)
                    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])
                    predictor_calls += 1
                    if policy == "dynamic":
                        chunk_position = torch.tensor(
                            [(chunk - 1) / 5.0], device=device
                        )
                        step_tensor = torch.tensor([step], dtype=torch.long, device=device)
                        confidence_events = start_timing()
                        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
                            predicted_log = head(
                                transformed_hidden=transformed,
                                pred_hidden=pred_hidden,
                                anchor_hidden=anchor_hidden,
                                chunk_position=chunk_position,
                                step_id=step_tensor,
                            )
                        finish_timing("confidence", confidence_events)
                        predicted_local_error = float(predicted_log.exp()[0])
                        alpha = (base.NUM_CHUNKS - 1 - chunk) / (base.NUM_CHUNKS - 2)
                        risk = predicted_local_error * (
                            1.0 + float(config["beta"]) * alpha
                        )
                        accepted = risk <= float(config["threshold"])
                    else:
                        accepted = static_accept

                if accepted:
                    assert pred_hidden is not None and pred_x0 is not None
                    current_hidden[step] = pred_hidden
                    denoised_pred = pred_x0
                    accepted_predictor_calls += 1
                else:
                    full_events = start_timing()
                    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 * base.TOKENS_PER_CHUNK,
                    )
                    current_hidden[step] = capture.finish()
                    finish_timing("full_dit", full_events)
                    full_calls += 1

                if candidate:
                    decisions.append(
                        {
                            "chunk": chunk,
                            "step": step,
                            "ran_predictor": run_predictor,
                            "accepted": accepted,
                            "predicted_local_error": predicted_local_error,
                            "chunk_alpha": alpha,
                            "impact_risk": risk,
                        }
                    )
                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
            context_events = start_timing()
            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 * base.TOKENS_PER_CHUNK,
            )
            finish_timing("context_dit", context_events)
            previous_chunk_hidden = current_hidden
    finally:
        capture.close()
    torch.cuda.synchronize()
    elapsed = {
        category: sum(start.elapsed_time(end) for start, end in events)
        for category, events in timing_events.items()
    }
    actual_dit_time_ms = (
        elapsed["full_dit"] + elapsed["predictor"] + elapsed["context_dit"]
    )
    return torch.cat(output_chunks, dim=1), {
        "generation_time_s": time.perf_counter() - started,
        "full_calls": full_calls,
        "predictor_calls": predictor_calls,
        "accepted_predictor_calls": accepted_predictor_calls,
        "rejected_predictor_calls": predictor_calls - accepted_predictor_calls,
        "full_dit_time_ms": elapsed["full_dit"],
        "predictor_time_ms": elapsed["predictor"],
        "confidence_head_time_ms": elapsed["confidence"],
        "context_dit_time_ms": elapsed["context_dit"],
        "actual_dit_time_ms": actual_dit_time_ms,
        "model_path_time_ms": actual_dit_time_ms + elapsed["confidence"],
        "decisions": decisions,
    }


def load_models(
    args: argparse.Namespace, device: torch.device
) -> tuple[Any, Any, Any, Any, Any]:
    print("[setup] loading VAE", flush=True)
    vae = WanVAEWrapper().to(device=device, dtype=torch.bfloat16).eval()
    config = OmegaConf.merge(
        OmegaConf.load(REPO_ROOT / "configs/default_config.yaml"),
        OmegaConf.load(args.config_path),
    )
    print("[setup] loading frozen generator and Predictor", flush=True)
    pipeline = base.build_pipeline(config, args.checkpoint_path, vae, device)
    if args.predictor_weights is not None:
        experiment = {
            "name": "direct_layer17_predictor",
            "source_layer": 17,
            "weights": args.predictor_weights,
            "gate_mode": "baseline",
        }
    else:
        experiment = base.discover_experiments(
            args.sweep_dir, ["teacher_layer_17"], None
        )[0]
    predictor = base.load_predictor(pipeline.generator.model, experiment, device)
    head = PredictorConfidenceHead(
        num_steps=max(args.candidate_steps)
    ).to(device=device).eval()
    head.load_state_dict(load_file(str(args.confidence_weights), device="cpu"), strict=True)
    head.requires_grad_(False)
    lpips_model = None
    if not args.skip_lpips:
        lpips_model = lpips.LPIPS(net="alex", verbose=False).to(device).eval()
        lpips_model.requires_grad_(False)
    return vae, pipeline, predictor, head, lpips_model


def smoke(args: argparse.Namespace, pipeline: Any, predictor: Any, head: Any, device: torch.device) -> None:
    max_accepts = 6 * len(args.candidate_steps)
    all_name = "fppf" if args.candidate_steps == [1, 2] else "fppp"
    configurations = [
        {"name": "ffff", "policy": "ffff", "beta": None, "threshold": None, "target_accepts": 0},
        {"name": "dynamic_all_fallback", "policy": "dynamic", "beta": 1.5, "threshold": -math.inf, "target_accepts": 0},
        {"name": all_name, "policy": "fppf", "beta": None, "threshold": None, "target_accepts": max_accepts},
        {"name": "dynamic_all_accept", "policy": "dynamic", "beta": 1.5, "threshold": math.inf, "target_accepts": max_accepts},
    ]
    latents = {}
    diagnostics = {}
    for config in configurations:
        latent, diagnostic = generate(
            pipeline=pipeline, dataset_root=args.dataset_root, prompt_id=80,
            generation_seed=args.generation_seed, device=device, predictor=predictor,
            head=head, config=config, candidate_steps=args.candidate_steps,
        )
        latents[config["name"]] = latent.cpu()
        diagnostics[config["name"]] = diagnostic
        print(
            f"[smoke] {config['name']} full={diagnostic['full_calls']} "
            f"pred={diagnostic['predictor_calls']} accept={diagnostic['accepted_predictor_calls']}",
            flush=True,
        )
    fallback_diff = float((latents["ffff"].float() - latents["dynamic_all_fallback"].float()).abs().max())
    accept_diff = float((latents[all_name].float() - latents["dynamic_all_accept"].float()).abs().max())
    result = {
        "status": "complete",
        "prompt_id": 80,
        "ffff_vs_all_fallback_max_abs": fallback_diff,
        "fppf_vs_all_accept_max_abs": accept_diff,
        "diagnostics": diagnostics,
    }
    atomic_json(args.output_root / "smoke.json", result)
    if fallback_diff != 0.0 or accept_diff != 0.0:
        raise RuntimeError(f"Smoke consistency failed: {result}")
    print("[smoke] exact consistency passed", flush=True)


def aggregate(records: list[dict[str, Any]], output_dir: Path) -> list[dict[str, Any]]:
    numeric = [
        "accepted_predictor_calls", "full_calls", "predictor_calls",
        "full_dit_time_ms", "predictor_time_ms", "confidence_head_time_ms",
        "context_dit_time_ms", "actual_dit_time_ms", "model_path_time_ms",
        "generation_time_s", "total_time_s", "latent_nrmse", "latent_tail_nrmse",
        "psnr", "ssim", "lpips", "tail_psnr", "tail_ssim", "tail_lpips",
    ]
    summary = []
    for name in sorted({str(row["config_name"]) for row in records}):
        selected = [row for row in records if row["config_name"] == name]
        first = selected[0]
        item = {
            "config_name": name,
            "policy": first["policy"],
            "beta": first["beta"],
            "target_accepts": first["target_accepts"],
            "threshold": first["threshold"],
            "num_prompts": len(selected),
        }
        for field in numeric:
            item[field] = sum(float(row[field]) for row in selected) / len(selected)
        summary.append(item)
    fields = [
        "config_name", "policy", "beta", "target_accepts", "threshold",
        "num_prompts", *numeric,
    ]
    write_csv(output_dir / "summary.csv", summary, fields)
    return summary


def select_validation(
    summary: list[dict[str, Any]], output_dir: Path, targets: list[int]
) -> None:
    selected_dynamic = []
    for target in targets:
        candidates = [
            row for row in summary
            if row["policy"] == "dynamic" and int(row["target_accepts"]) == target
        ]
        same_budget = [
            row for row in candidates
            if abs(float(row["accepted_predictor_calls"]) - target) <= 0.5 + 1e-8
        ]
        if not same_budget:
            closest = min(
                abs(float(row["accepted_predictor_calls"]) - target)
                for row in candidates
            )
            same_budget = [
                row for row in candidates
                if abs(abs(float(row["accepted_predictor_calls"]) - target) - closest)
                <= 1e-8
            ]
        same_budget.sort(
            key=lambda row: (
                float(row["tail_lpips"]),
                abs(float(row["accepted_predictor_calls"]) - target),
                float(row["beta"]),
            )
        )
        selected_dynamic.append(same_budget[0])
    atomic_json(
        output_dir / "selected.json",
        {
            "selection_rule": (
                "within target accepted calls +/-0.5, lowest validation tail LPIPS; "
                "then budget distance and lower beta"
            ),
            "selected_dynamic": selected_dynamic,
        },
    )


def formal(
    args: argparse.Namespace,
    vae: Any,
    pipeline: Any,
    predictor: Any,
    head: Any,
    lpips_model: Any,
    device: torch.device,
) -> None:
    split = args.mode
    split_prompt_ids = (
        list(range(80, 90)) if split == "validation" else list(range(90, 100))
    )
    prompt_ids = args.prompt_ids or split_prompt_ids
    invalid_prompt_ids = sorted(set(prompt_ids) - set(split_prompt_ids))
    if invalid_prompt_ids:
        raise ValueError(
            f"Prompt IDs {invalid_prompt_ids} are outside the {split} split"
        )
    if split == "validation":
        configurations = dynamic_configs(
            args.validation_predictions, args.candidate_steps, args.target_accepts
        ) + static_configs(args.target_accepts)
        max_accepts = 6 * len(args.candidate_steps)
        all_name = "fppf" if args.candidate_steps == [1, 2] else "fppp"
        configurations += [
            {"name": "ffff", "policy": "ffff", "beta": None, "threshold": None, "target_accepts": 0},
            {"name": all_name, "policy": "fppf", "beta": None, "threshold": None, "target_accepts": max_accepts},
        ]
    else:
        selected_path = args.selected_path or (
            args.output_root / "validation" / "selected.json"
        )
        configurations = load_selected(selected_path) + static_configs(args.target_accepts)
        max_accepts = 6 * len(args.candidate_steps)
        all_name = "fppf" if args.candidate_steps == [1, 2] else "fppp"
        configurations += [
            {"name": "ffff", "policy": "ffff", "beta": None, "threshold": None, "target_accepts": 0},
            {"name": all_name, "policy": "fppf", "beta": None, "threshold": None, "target_accepts": max_accepts},
        ]
    if args.config_names:
        requested = set(args.config_names)
        available = {str(config["name"]) for config in configurations}
        missing = requested - available
        if missing:
            raise ValueError(
                f"Unknown config_names {sorted(missing)}; available={sorted(available)}"
            )
        configurations = [
            config for config in configurations if config["name"] in requested
        ]
    output_dir = args.output_root / split
    output_dir.mkdir(parents=True, exist_ok=True)
    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,
        )
    total = len(configurations) * len(prompt_ids)
    records = []
    completed = 0
    print("[warmup] one unmeasured Full+Predictor+Head rollout", flush=True)
    warmup_config = {
        "name": "warmup",
        "policy": "dynamic",
        "beta": 1.0,
        "threshold": -math.inf,
        "target_accepts": 0,
    }
    warmup_latent, _ = generate(
        pipeline=pipeline,
        dataset_root=args.dataset_root,
        prompt_id=prompt_ids[0],
        generation_seed=args.generation_seed,
        device=device,
        predictor=predictor,
        head=head,
        config=warmup_config,
        candidate_steps=args.candidate_steps,
    )
    del warmup_latent
    torch.cuda.empty_cache()
    for config in configurations:
        for prompt_id in prompt_ids:
            destination = output_dir / "per_run" / config["name"] / f"prompt_{prompt_id:04d}.json"
            if destination.exists() and not args.overwrite:
                records.append(json.loads(destination.read_text(encoding="utf-8")))
                completed += 1
                print(f"[cached] {completed}/{total} {config['name']} p={prompt_id}", flush=True)
                continue
            started = time.perf_counter()
            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)
            latent, diagnostic = generate(
                pipeline=pipeline, dataset_root=args.dataset_root, prompt_id=prompt_id,
                generation_seed=args.generation_seed, device=device, predictor=predictor,
                head=head, config=config, candidate_steps=args.candidate_steps,
            )
            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)
            if args.save_videos:
                base.save_mp4(
                    prediction_u8,
                    output_dir / "videos" / config["name"]
                    / f"prompt_{prompt_id:04d}.mp4",
                )
            frame = base.frame_metrics(
                reference_u8=reference_u8, prediction_u8=prediction_u8,
                lpips_model=lpips_model, batch_size=args.metric_batch_size, device=device,
            )
            tail_start = impact_eval.chunk_frame_slice(1).start
            tail = impact_eval.summarize_frame_range(frame, slice(tail_start, None))
            record = {
                "config_name": config["name"],
                "policy": config["policy"],
                "beta": config["beta"],
                "target_accepts": config["target_accepts"],
                "threshold": config["threshold"],
                "prompt_id": prompt_id,
                "latent_nrmse": impact_eval.nrmse(latent, reference_latent),
                "latent_tail_nrmse": impact_eval.nrmse(
                    latent[:, base.FRAMES_PER_CHUNK:],
                    reference_latent[:, base.FRAMES_PER_CHUNK:],
                ),
                "psnr": frame["psnr"],
                "ssim": frame["ssim"],
                "lpips": frame["lpips"],
                "tail_psnr": tail["psnr"],
                "tail_ssim": tail["ssim"],
                "tail_lpips": tail["lpips"],
                **diagnostic,
                "total_time_s": time.perf_counter() - started,
            }
            atomic_json(destination, record)
            records.append(record)
            completed += 1
            print(
                f"[run] {completed}/{total} {config['name']} p={prompt_id} "
                f"accept={record['accepted_predictor_calls']} "
                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 reference_latent, reference_u8, latent, pixels, prediction_u8, frame
            torch.cuda.empty_cache()
    flat_fields = sorted({key for row in records for key in row if key != "decisions"})
    shard_suffix = f"_gpu{args.gpu}" if args.prompt_ids else ""
    write_csv(
        output_dir / f"runs{shard_suffix}.csv",
        [{key: row.get(key) for key in flat_fields} for row in records],
        flat_fields,
    )
    summary_output_dir = output_dir
    if shard_suffix:
        summary_output_dir = output_dir / f".summary_shard_gpu{args.gpu}"
        summary_output_dir.mkdir(parents=True, exist_ok=True)
    summary = aggregate(records, summary_output_dir)
    if shard_suffix:
        os.replace(
            summary_output_dir / "summary.csv",
            output_dir / f"summary{shard_suffix}.csv",
        )
        summary_output_dir.rmdir()
    if split == "validation" and not shard_suffix:
        select_validation(summary, output_dir, args.target_accepts)
    atomic_json(
        output_dir / f"manifest{shard_suffix}.json",
        {
            "status": "complete", "split": split, "prompt_ids": prompt_ids,
            "num_configs": len(configurations), "num_runs": len(records),
            "configs": configurations, "candidate_steps": args.candidate_steps,
            "target_accepts": args.target_accepts,
            "predictor_weights": (
                str(args.predictor_weights) if args.predictor_weights else None
            ),
        },
    )
    print(f"[complete] {split} -> {output_dir}", flush=True)


def main() -> None:
    args = parse_args()
    args.output_root.mkdir(parents=True, exist_ok=True)
    device = torch.device("cuda")
    torch.set_grad_enabled(False)
    set_seed(args.generation_seed)
    vae, pipeline, predictor, head, lpips_model = load_models(args, device)
    if args.mode == "smoke":
        smoke(args, pipeline, predictor, head, device)
    else:
        formal(args, vae, pipeline, predictor, head, lpips_model, device)


if __name__ == "__main__":
    main()