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