scripts/generate_vbench8_extended_naive_baselines.py

#!/usr/bin/env python3
"""Generate Extended-251 videos for timestep-skipping and velocity-cache baselines.

For every prompt, an original four-step FFFF rollout is generated first with
the same seed and used as the PSNR/SSIM/LPIPS reference. Timestep-skipping
policies shorten the denoising schedule itself. Velocity-cache policies retain
the original four timesteps and, at each R, reuse the velocity from the most
recent F while converting it to x0 using the current noisy latent and timestep.
The cache pattern is applied independently to every autoregressive chunk.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import sys
import tempfile
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="0")
    args, _ = parser.parse_known_args()
    os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu)
    os.environ.setdefault("MPLCONFIGDIR", tempfile.mkdtemp(prefix="self_forcing_mpl_"))
    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 scripts.naive_vbench_policies import (
    EVALUATION_STRATEGY_NAMES,
    POLICIES_BY_NAME,
    REFERENCE_POLICY,
    NaivePolicy,
)
from scripts.vbench8_protocol import SUITE_COUNTS
from utils.misc import set_seed
from utils.wan_wrapper import WanTextEncoder, WanVAEWrapper


NUM_CHUNKS = base.NUM_CHUNKS
FRAMES_PER_CHUNK = base.FRAMES_PER_CHUNK
MAPPING_DEFAULT = REPO_ROOT / "assets/vbench8_extended_subset_mapping.json"
OUTPUT_DEFAULT = REPO_ROOT / "evaluation_runs/vbench8_extended_naive_baselines"


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


def atomic_video(frames: torch.Tensor, path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_name(path.name + ".partial.mp4")
    if temporary.exists():
        temporary.unlink()
    base.save_mp4(frames, temporary)
    os.replace(temporary, path)


def read_mapping(path: Path) -> list[dict[str, Any]]:
    value = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(value, list) or len(value) != 251:
        raise ValueError(f"Expected a 251-row mapping list: {path}")
    counts = {suite: 0 for suite in SUITE_COUNTS}
    globals_seen: set[int] = set()
    suite_seen: dict[str, set[int]] = {suite: set() for suite in SUITE_COUNTS}
    for row in value:
        suite = str(row["prompt_suite"])
        global_index = int(row["global_index"])
        suite_index = int(row["suite_index"])
        if suite not in counts:
            raise ValueError(f"Unknown prompt suite: {suite}")
        if global_index in globals_seen or suite_index in suite_seen[suite]:
            raise ValueError(
                f"Duplicate mapping index: {global_index} / {suite}:{suite_index}"
            )
        if not str(row["extended_prompt"]).strip():
            raise ValueError(f"Empty prompt at global index {global_index}")
        counts[suite] += 1
        globals_seen.add(global_index)
        suite_seen[suite].add(suite_index)
    if counts != SUITE_COUNTS:
        raise ValueError(f"Unexpected prompt-suite counts: {counts}")
    return sorted(value, key=lambda row: int(row["global_index"]))


def reset_runtime_caches(pipeline: Any, 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_()
    for cache in pipeline.crossattn_cache:
        cache["is_init"] = False


def warped_timesteps(pipeline: Any, raw_timesteps: tuple[int, ...]) -> torch.Tensor:
    raw = torch.tensor(raw_timesteps, dtype=torch.long)
    if not bool(pipeline.args.warp_denoising_step):
        return raw
    scheduler_timesteps = torch.cat(
        (pipeline.scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32))
    )
    return scheduler_timesteps[1000 - raw]


def policy_timesteps(pipeline: Any, policy: NaivePolicy) -> torch.Tensor:
    if policy.family in {"full", "velocity_reuse"}:
        actual = pipeline.denoising_step_list.detach().cpu()
    else:
        actual = warped_timesteps(pipeline, policy.raw_timesteps)
    if len(actual) != policy.denoise_steps:
        raise ValueError(
            f"{policy.name}: expected {policy.denoise_steps} timesteps, got {len(actual)}"
        )
    return actual


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


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


@torch.inference_mode()
def generate_rollout(
    *,
    pipeline: Any,
    conditional_dict: dict[str, torch.Tensor],
    seed: int,
    device: torch.device,
    policy: NaivePolicy,
) -> tuple[torch.Tensor, dict[str, Any]]:
    reset_runtime_caches(pipeline, device)
    set_seed(seed)
    noise = torch.randn(
        1,
        NUM_CHUNKS * FRAMES_PER_CHUNK,
        base.LATENT_CHANNELS,
        base.LATENT_HEIGHT,
        base.LATENT_WIDTH,
        dtype=torch.bfloat16,
        device=device,
    )
    timesteps_cpu = policy_timesteps(pipeline, policy)
    actual_timestep_values = [float(value) for value in timesteps_cpu]
    timesteps = timesteps_cpu.to(device=device)
    output_chunks: list[torch.Tensor] = []
    decisions: list[dict[str, Any]] = []
    full_calls = 0
    reuse_calls = 0
    timing_events: dict[str, list[tuple[torch.cuda.Event, torch.cuda.Event]]] = {
        "full_dit": [],
        "context_dit": [],
    }
    started = time.perf_counter()

    for chunk in range(NUM_CHUNKS):
        noisy_input = noise[
            :, chunk * FRAMES_PER_CHUNK : (chunk + 1) * FRAMES_PER_CHUNK
        ]
        cached_flow: torch.Tensor | None = None
        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
            action = (
                policy.cache_pattern[step]
                if policy.family == "velocity_reuse" and policy.cache_pattern is not None
                else "F"
            )
            if action == "R":
                if cached_flow is None:
                    raise RuntimeError(
                        f"{policy.name}: reuse requested before a full step in chunk {chunk}"
                    )
                denoised_pred = pipeline.generator._convert_flow_pred_to_x0(
                    flow_pred=cached_flow.flatten(0, 1),
                    xt=noisy_input.flatten(0, 1),
                    timestep=timestep.flatten(0, 1),
                ).unflatten(0, cached_flow.shape[:2])
                reuse_calls += 1
            else:
                full_events = start_timing()
                flow, 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,
                )
                finish_timing(timing_events, "full_dit", full_events)
                cached_flow = flow.detach()
                full_calls += 1
            decisions.append(
                {
                    "chunk": chunk,
                    "step": step,
                    "action": action,
                    "timestep": actual_timestep_values[step],
                }
            )
            if step < len(timesteps) - 1:
                if denoised_pred is None:
                    raise RuntimeError("Denoising step did not produce 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(
                        [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(f"{policy.name}: chunk {chunk} produced no clean latent")
        output_chunks.append(denoised_pred)
        context_events = start_timing()
        pipeline.generator(
            noisy_image_or_video=denoised_pred,
            conditional_dict=conditional_dict,
            timestep=torch.ones_like(timestep) * pipeline.args.context_noise,
            kv_cache=pipeline.kv_cache1,
            crossattn_cache=pipeline.crossattn_cache,
            current_start=chunk * base.TOKENS_PER_CHUNK,
        )
        finish_timing(timing_events, "context_dit", context_events)

    torch.cuda.synchronize()
    elapsed = {
        category: sum(start.elapsed_time(end) for start, end in events)
        for category, events in timing_events.items()
    }
    expected_full = NUM_CHUNKS * policy.full_calls_per_chunk
    expected_reuse = NUM_CHUNKS * policy.reuse_calls_per_chunk
    if (full_calls, reuse_calls) != (expected_full, expected_reuse):
        raise RuntimeError(
            f"{policy.name}: got calls F={full_calls}, R={reuse_calls}; "
            f"expected F={expected_full}, R={expected_reuse}"
        )
    return torch.cat(output_chunks, dim=1), {
        "generation_time_s": time.perf_counter() - started,
        "full_calls": full_calls,
        "reuse_calls": reuse_calls,
        "predictor_calls": 0,
        "accepted_predictor_calls": 0,
        "rejected_predictor_calls": 0,
        "full_dit_time_ms": elapsed["full_dit"],
        "predictor_time_ms": 0.0,
        "confidence_head_time_ms": 0.0,
        "context_dit_time_ms": elapsed["context_dit"],
        "actual_dit_time_ms": elapsed["full_dit"] + elapsed["context_dit"],
        "model_path_time_ms": elapsed["full_dit"],
        "raw_timesteps": list(policy.raw_timesteps),
        "actual_timesteps": actual_timestep_values,
        "decisions": decisions,
    }


def compact_metrics(value: dict[str, Any]) -> dict[str, Any]:
    keys = (
        "pixel_mse",
        "psnr",
        "psnr_frame_mean",
        "ssim",
        "lpips",
        "rollout_start_frame",
        "rollout_pixel_mse",
        "rollout_psnr",
        "rollout_ssim",
        "rollout_lpips",
        "num_frames",
    )
    return {key: value[key] for key in keys}


def identity_metrics(num_frames: int) -> dict[str, Any]:
    return {
        "pixel_mse": 0.0,
        "psnr": 120.0,
        "psnr_frame_mean": 120.0,
        "ssim": 1.0,
        "lpips": 0.0,
        "rollout_start_frame": base.PIXEL_FRAMES_FIRST_CHUNK,
        "rollout_pixel_mse": 0.0,
        "rollout_psnr": 120.0,
        "rollout_ssim": 1.0,
        "rollout_lpips": 0.0,
        "num_frames": num_frames,
    }


def load_models(device: torch.device) -> tuple[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(REPO_ROOT / "configs/self_forcing_sid.yaml"),
    )
    print("[setup] loading frozen Self-Forcing Teacher", flush=True)
    pipeline = base.build_pipeline(
        config, REPO_ROOT / "checkpoints/self_forcing_dmd.pt", vae, device
    )
    text_encoder = WanTextEncoder().to(device=device, dtype=torch.bfloat16).eval()
    text_encoder.requires_grad_(False)
    print("[setup] loading LPIPS", flush=True)
    lpips_model = lpips.LPIPS(net="alex", verbose=False).to(device).eval()
    lpips_model.requires_grad_(False)
    return vae, pipeline, text_encoder, lpips_model


def artifact_paths(
    output_root: Path, mapping_row: dict[str, Any], strategy: str
) -> tuple[Path, Path]:
    suite = str(mapping_row["prompt_suite"])
    suite_index = int(mapping_row["suite_index"])
    global_index = int(mapping_row["global_index"])
    return (
        output_root / "generated_videos" / strategy / suite / f"{suite_index:03d}.mp4",
        output_root
        / "generation_metrics/per_prompt"
        / strategy
        / f"global_{global_index:04d}.json",
    )


def is_prompt_complete(
    output_root: Path, mapping_row: dict[str, Any], policies: tuple[NaivePolicy, ...]
) -> bool:
    for policy in (REFERENCE_POLICY, *policies):
        video, record = artifact_paths(output_root, mapping_row, policy.name)
        if not video.is_file() or not record.is_file():
            return False
    return True


def write_result(
    *,
    output_root: Path,
    mapping_row: dict[str, Any],
    policy: NaivePolicy,
    seed: int,
    diagnostic: dict[str, Any],
    pixel_metrics: dict[str, Any],
    frames: torch.Tensor,
) -> None:
    video_path, record_path = artifact_paths(output_root, mapping_row, policy.name)
    atomic_video(frames, video_path)
    diagnostic_without_decisions = {
        key: value for key, value in diagnostic.items() if key != "decisions"
    }
    record = {
        "status": "complete",
        "strategy": policy.name,
        "label": policy.label,
        "policy": policy.to_dict(),
        "global_index": int(mapping_row["global_index"]),
        "prompt_suite": str(mapping_row["prompt_suite"]),
        "suite_index": int(mapping_row["suite_index"]),
        "prompt": str(mapping_row["extended_prompt"]),
        "seed": seed,
        "generation": diagnostic_without_decisions,
        "decisions": diagnostic["decisions"],
        "pixel_metrics_vs_ffff": pixel_metrics,
        "video": str(video_path.relative_to(output_root)),
    }
    atomic_json(record_path, record)


def run_prompt(
    *,
    mapping_row: dict[str, Any],
    policies: tuple[NaivePolicy, ...],
    output_root: Path,
    seed: int,
    vae: Any,
    pipeline: Any,
    text_encoder: Any,
    lpips_model: Any,
    device: torch.device,
) -> None:
    global_index = int(mapping_row["global_index"])
    suite = str(mapping_row["prompt_suite"])
    suite_index = int(mapping_row["suite_index"])
    prompt = str(mapping_row["extended_prompt"])
    print(f"[prompt] global={global_index} suite={suite}/{suite_index}", flush=True)
    conditional = text_encoder(text_prompts=[prompt])

    reference_latent, reference_diagnostic = generate_rollout(
        pipeline=pipeline,
        conditional_dict=conditional,
        seed=seed,
        device=device,
        policy=REFERENCE_POLICY,
    )
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        reference_pixels = vae.decode_to_pixel(reference_latent, use_cache=False)
    reference_u8 = base.pixels_to_u8(reference_pixels)
    if reference_u8.ndim != 4 or reference_u8.shape[0] != 81:
        raise RuntimeError(f"Unexpected reference shape: {tuple(reference_u8.shape)}")
    write_result(
        output_root=output_root,
        mapping_row=mapping_row,
        policy=REFERENCE_POLICY,
        seed=seed,
        diagnostic=reference_diagnostic,
        pixel_metrics=identity_metrics(int(reference_u8.shape[0])),
        frames=reference_u8,
    )

    for policy in policies:
        latent, diagnostic = generate_rollout(
            pipeline=pipeline,
            conditional_dict=conditional,
            seed=seed,
            device=device,
            policy=policy,
        )
        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)
        metrics = base.frame_metrics(
            reference_u8=reference_u8,
            prediction_u8=prediction_u8,
            lpips_model=lpips_model,
            batch_size=4,
            device=device,
        )
        write_result(
            output_root=output_root,
            mapping_row=mapping_row,
            policy=policy,
            seed=seed,
            diagnostic=diagnostic,
            pixel_metrics=compact_metrics(metrics),
            frames=prediction_u8,
        )
        print(
            f"[result] global={global_index} strategy={policy.name} "
            f"F={diagnostic['full_calls']} R={diagnostic['reuse_calls']} "
            f"psnr={metrics['psnr']:.4f} ssim={metrics['ssim']:.6f} "
            f"lpips={metrics['lpips']:.6f}",
            flush=True,
        )
        if hasattr(vae.model, "clear_cache"):
            vae.model.clear_cache()
        del latent, pixels, prediction_u8, metrics
        torch.cuda.empty_cache()

    if hasattr(vae.model, "clear_cache"):
        vae.model.clear_cache()
    del conditional, reference_latent, reference_pixels, reference_u8
    torch.cuda.empty_cache()


def write_manifest(
    *,
    output_root: Path,
    gpu: str,
    shard_index: int,
    num_shards: int,
    rows: list[dict[str, Any]],
    policies: tuple[NaivePolicy, ...],
    seed: int,
    status: str,
) -> None:
    selected = "_".join(policy.name for policy in policies)
    selection_id = hashlib.sha256(selected.encode("utf-8")).hexdigest()[:8]
    manifest_name = f"naive_gpu{gpu}_shard{shard_index:02d}_{selection_id}.json"
    atomic_json(
        output_root / "generation_metrics/manifests" / manifest_name,
        {
            "status": status,
            "physical_gpu": gpu,
            "shard_index": shard_index,
            "num_shards": num_shards,
            "global_indices": [int(row["global_index"]) for row in rows],
            "num_prompts": len(rows),
            "strategies": [REFERENCE_POLICY.name, *(policy.name for policy in policies)],
            "seed": seed,
        },
    )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--gpu", default=PHYSICAL_GPU)
    parser.add_argument("--mapping", type=Path, default=MAPPING_DEFAULT)
    parser.add_argument("--output-root", type=Path, default=OUTPUT_DEFAULT)
    parser.add_argument("--shard-index", type=int, default=0)
    parser.add_argument("--num-shards", type=int, default=1)
    parser.add_argument(
        "--strategy",
        action="append",
        choices=EVALUATION_STRATEGY_NAMES,
        default=None,
        help="Baseline to generate; repeat to select several. Default: all six.",
    )
    parser.add_argument(
        "--global-index",
        dest="global_indices",
        action="append",
        type=int,
        default=None,
        help="Process only this global mapping index; may be repeated.",
    )
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--overwrite", action="store_true")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if args.num_shards < 1 or not 0 <= args.shard_index < args.num_shards:
        raise ValueError("Invalid shard index/number of shards")
    strategy_names = args.strategy or list(EVALUATION_STRATEGY_NAMES)
    if len(set(strategy_names)) != len(strategy_names):
        raise ValueError("--strategy values must be unique")
    policies = tuple(POLICIES_BY_NAME[name] for name in strategy_names)
    mapping = read_mapping(args.mapping.resolve())
    if args.global_indices is None:
        rows = [
            row
            for position, row in enumerate(mapping)
            if position % args.num_shards == args.shard_index
        ]
    else:
        requested = set(args.global_indices)
        mapping_by_global = {int(row["global_index"]): row for row in mapping}
        unknown = sorted(requested - set(mapping_by_global))
        if unknown:
            raise ValueError(f"Unknown global indices: {unknown}")
        rows = [row for row in mapping if int(row["global_index"]) in requested]

    output_root = args.output_root.resolve()
    output_root.mkdir(parents=True, exist_ok=True)
    write_manifest(
        output_root=output_root,
        gpu=str(args.gpu),
        shard_index=args.shard_index,
        num_shards=args.num_shards,
        rows=rows,
        policies=policies,
        seed=args.seed,
        status="running",
    )
    print(
        f"[setup] physical_gpu={args.gpu} shard={args.shard_index}/{args.num_shards} "
        f"prompts={len(rows)} strategies={','.join(strategy_names)}",
        flush=True,
    )
    device = torch.device("cuda")
    torch.set_grad_enabled(False)
    vae, pipeline, text_encoder, lpips_model = load_models(device)
    completed = 0
    try:
        for row in rows:
            if not args.overwrite and is_prompt_complete(output_root, row, policies):
                completed += 1
                print(
                    f"[cached] {completed}/{len(rows)} global={row['global_index']}",
                    flush=True,
                )
                continue
            run_prompt(
                mapping_row=row,
                policies=policies,
                output_root=output_root,
                seed=args.seed,
                vae=vae,
                pipeline=pipeline,
                text_encoder=text_encoder,
                lpips_model=lpips_model,
                device=device,
            )
            completed += 1
            print(f"[progress] {completed}/{len(rows)} prompts", flush=True)
    finally:
        write_manifest(
            output_root=output_root,
            gpu=str(args.gpu),
            shard_index=args.shard_index,
            num_shards=args.num_shards,
            rows=rows,
            policies=policies,
            seed=args.seed,
            status="complete" if completed == len(rows) else "failed",
        )
    print(f"[complete] gpu={args.gpu} prompts={completed}/{len(rows)}", flush=True)


if __name__ == "__main__":
    main()