#!/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()