scripts/evaluate_long_video_fppf.py

#!/usr/bin/env python3
"""Evaluate 2x/4x long FFFF and one-block Layer-17 FPPF rollouts."""

from __future__ import annotations

import argparse
import json
import os
import sys
import time
from pathlib import Path


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)
    return str(args.gpu)


PHYSICAL_GPU = _preparse_gpu()

import lpips
import torch
from omegaconf import OmegaConf
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 predictor_training.offline_data import TOKENS_PER_CHUNK
from scripts.evaluate_single_block_fppf import (
    DEFAULT_PROMPT_IDS,
    FRAMES_PER_CHUNK,
    LATENT_CHANNELS,
    LATENT_HEIGHT,
    LATENT_WIDTH,
    NUM_DENOISING_STEPS,
    FinalHiddenCapture,
    atomic_json,
    build_pipeline,
    discover_experiments,
    frame_metrics,
    load_predictor,
    load_prompt_metadata,
    pixels_to_u8,
    predictor_step,
    reset_kv_and_load_cross_cache,
)
from utils.misc import set_seed
from utils.wan_wrapper import WanVAEWrapper


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--gpu", default=PHYSICAL_GPU)
    parser.add_argument("--prompt_ids", type=int, nargs="+", required=True)
    parser.add_argument("--latent_lengths", type=int, nargs="+", default=[42, 84])
    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/long_video_2x4x_eval")
    )
    parser.add_argument("--metric_batch_size", type=int, default=4)
    parser.add_argument("--generation_seed", type=int, default=0)
    args = parser.parse_args()
    if any(length <= 0 or length % FRAMES_PER_CHUNK for length in args.latent_lengths):
        parser.error("Latent lengths must be positive multiples of 3")
    if any(prompt not in DEFAULT_PROMPT_IDS for prompt in args.prompt_ids):
        parser.error("This evaluation is restricted to validation prompt IDs 80..99")
    return args


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


@torch.inference_mode()
def generate_rollout(
    *, pipeline, dataset_root: Path, prompt_id: int, latent_length: int,
    generation_seed: int, device: torch.device, predictor, 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 the Predictor")
    num_chunks = latent_length // FRAMES_PER_CHUNK
    reset_kv_and_load_cross_cache(pipeline, dataset_root, prompt_id, device)
    set_seed(generation_seed)
    noise = torch.randn(
        1, latent_length, 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 = []
    previous_chunk_hidden = None
    capture = FinalHiddenCapture(teacher)
    full_calls = predictor_calls = 0
    started = time.perf_counter()
    try:
        for chunk in range(num_chunks):
            noisy_input = noise[:, chunk * 3:(chunk + 1) * 3]
            current_hidden = [None] * NUM_DENOISING_STEPS
            denoised_pred = timestep = 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:
                    pred_hidden, flow, _ = predictor_step(
                        predictor=predictor,
                        teacher=teacher,
                        noisy_input=noisy_input,
                        timestep=timestep,
                        anchor_hidden=current_hidden[step - 1],
                        previous_hidden=previous_chunk_hidden[step],
                        history_cache=pipeline.kv_cache1[source_layer],
                        cross_cache=pipeline.crossattn_cache[source_layer],
                        current_start=chunk * TOKENS_PER_CHUNK,
                    )
                    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:
                    flat = denoised_pred.flatten(0, 1)
                    noisy_input = pipeline.scheduler.add_noise(
                        flat, torch.randn_like(flat),
                        timesteps[step + 1] * torch.ones(
                            [FRAMES_PER_CHUNK], dtype=torch.long, device=device
                        ),
                    ).unflatten(0, denoised_pred.shape[:2])
            output_chunks.append(denoised_pred)
            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 * 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,
        "num_chunks": num_chunks,
    }


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"},
    )


def main() -> None:
    args = parse_args()
    for field in ("config_path", "checkpoint_path", "dataset_root", "sweep_dir", "output_dir"):
        setattr(args, field, resolve(getattr(args, field)))
    args.output_dir.mkdir(parents=True, exist_ok=True)
    atomic_json(args.output_dir / "manifest.json", {
        "status": "running", "physical_gpu": str(args.gpu),
        "prompt_ids": args.prompt_ids, "latent_lengths": args.latent_lengths,
        "methods": ["FFFF", "FPPF_teacher_layer_17"],
        "generation_seed": args.generation_seed,
    })
    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),
    )
    # The released checkpoint uses full attention over its 21-latent training
    # horizon.  Long inference keeps exactly that horizon as a rolling window;
    # within the first 21 latents this is numerically the same attention span.
    config.model_kwargs.local_attn_size = 21
    vae = WanVAEWrapper().to(device=device, dtype=torch.bfloat16).eval()
    pipeline = build_pipeline(config, args.checkpoint_path, vae, device)
    experiment = discover_experiments(
        args.sweep_dir, ["teacher_layer_17"], None
    )[0]
    predictor = load_predictor(pipeline.generator.model, experiment, device)
    lpips_model = lpips.LPIPS(net="alex", verbose=False).to(device).eval()
    lpips_model.requires_grad_(False)

    for prompt_id in args.prompt_ids:
        prompt = load_prompt_metadata(args.dataset_root, prompt_id)["prompt"]
        for latent_length in args.latent_lengths:
            run_dir = args.output_dir / f"latent_{latent_length}" / f"prompt_{prompt_id:04d}"
            result_path = run_dir / "metrics.json"
            if result_path.exists():
                existing = json.loads(result_path.read_text(encoding="utf-8"))
                if existing.get("status") == "complete":
                    print(f"[skip] latent={latent_length} prompt={prompt_id}", flush=True)
                    continue
            print(f"[run] latent={latent_length} prompt={prompt_id} FFFF", flush=True)
            reference_latent, ffff_counts = generate_rollout(
                pipeline=pipeline, dataset_root=args.dataset_root,
                prompt_id=prompt_id, latent_length=latent_length,
                generation_seed=args.generation_seed, device=device,
                predictor=None, source_layer=None, schedule="FFFF",
            )
            with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
                reference_pixels = vae.decode_to_pixel(reference_latent, use_cache=False)
            reference_u8 = pixels_to_u8(reference_pixels)
            save_mp4(reference_u8, run_dir / "ffff.mp4")
            del reference_latent, reference_pixels
            if hasattr(vae.model, "clear_cache"):
                vae.model.clear_cache()
            torch.cuda.empty_cache()

            print(f"[run] latent={latent_length} prompt={prompt_id} FPPF", flush=True)
            prediction_latent, fppf_counts = generate_rollout(
                pipeline=pipeline, dataset_root=args.dataset_root,
                prompt_id=prompt_id, latent_length=latent_length,
                generation_seed=args.generation_seed, device=device,
                predictor=predictor, source_layer=17, schedule="FPPF",
            )
            with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
                prediction_pixels = vae.decode_to_pixel(prediction_latent, use_cache=False)
            prediction_u8 = pixels_to_u8(prediction_pixels)
            save_mp4(prediction_u8, run_dir / "fppf_layer17.mp4")
            metrics = frame_metrics(
                reference_u8=reference_u8, prediction_u8=prediction_u8,
                lpips_model=lpips_model, batch_size=args.metric_batch_size,
                device=device,
            )
            atomic_json(result_path, {
                "status": "complete", "prompt_id": prompt_id, "prompt": prompt,
                "latent_length": latent_length, "decoded_frames": metrics["num_frames"],
                "reference": "FFFF same prompt/seed/noise",
                "predictor": "single_block_teacher_layer_17",
                "ffff": ffff_counts, "fppf": fppf_counts, **metrics,
            })
            print(
                f"[result] latent={latent_length} prompt={prompt_id} "
                f"psnr={metrics['psnr']:.4f} ssim={metrics['ssim']:.6f} "
                f"lpips={metrics['lpips']:.6f}", flush=True,
            )
            del prediction_latent, prediction_pixels, reference_u8, prediction_u8
            if hasattr(vae.model, "clear_cache"):
                vae.model.clear_cache()
            torch.cuda.empty_cache()

    manifest = json.loads((args.output_dir / "manifest.json").read_text(encoding="utf-8"))
    manifest["status"] = "complete"
    atomic_json(args.output_dir / "manifest.json", manifest)


if __name__ == "__main__":
    main()