| |
| """Generate ATC FPPF/FPPP videos against the existing Extended-251 FFFF run.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import sys |
| from pathlib import Path |
| from typing import Any |
|
|
| 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 generate_vbench8_extended_strategies as extended |
|
|
| import lpips |
| import torch |
| from omegaconf import OmegaConf |
| from safetensors import safe_open |
|
|
| from scripts import evaluate_single_block_fppf as base |
| from scripts.generate_vbench8_extended_disca import artifact_paths, read_u8_video |
| from utils.wan_wrapper import WanTextEncoder, WanVAEWrapper |
|
|
|
|
| MAPPING_DEFAULT = REPO_ROOT / "assets/vbench8_extended_subset_mapping.json" |
| REFERENCE_DEFAULT = ( |
| REPO_ROOT / "evaluation_runs/vbench8_extended_stage1_step2000_20260901" |
| ) |
| OUTPUT_DEFAULT = ( |
| REPO_ROOT / "evaluation_runs/vbench8_extended_atc_stage1_step2000_20260901" |
| ) |
| PATTERN_TO_STEPS = {"FPPF": [1, 2], "FPPP": [1, 2, 3]} |
|
|
|
|
| def read_predictor_config(weights: Path) -> dict[str, Any]: |
| with safe_open(weights, framework="pt", device="cpu") as handle: |
| metadata = handle.metadata() or {} |
| raw = metadata.get("predictor_config") |
| if raw is None: |
| raise ValueError(f"ATC checkpoint has no predictor_config metadata: {weights}") |
| config = json.loads(raw) |
| if config.get("input_variant") != "atc": |
| raise ValueError(f"Checkpoint is not ATC: {config}") |
| return config |
|
|
|
|
| def strategy_configs( |
| scope: str, patterns: list[str] | tuple[str, ...] |
| ) -> list[dict[str, Any]]: |
| return [ |
| { |
| "name": f"atc_{scope}_{pattern.lower()}", |
| "policy": "fppf", |
| "pattern": pattern, |
| "candidate_steps": PATTERN_TO_STEPS[pattern], |
| "allow_chunk0_predictor": False, |
| "beta": None, |
| "threshold": None, |
| "target_accepts": 6 * len(PATTERN_TO_STEPS[pattern]), |
| "head": None, |
| "input_variant": "atc", |
| "atc_previous_scope": scope, |
| } |
| for pattern in patterns |
| ] |
|
|
|
|
| def load_models( |
| weights: Path, scope: str, device: torch.device |
| ) -> tuple[Any, Any, Any, Any, None, None, Any]: |
| predictor_config = read_predictor_config(weights) |
| if predictor_config.get("atc_previous_scope") != scope: |
| raise ValueError( |
| f"Checkpoint scope={predictor_config.get('atc_previous_scope')} " |
| f"does not match requested scope={scope}" |
| ) |
| 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(f"[setup] loading Teacher and ATC scope={scope}", flush=True) |
| pipeline = base.build_pipeline( |
| config, REPO_ROOT / "checkpoints/self_forcing_dmd.pt", vae, device |
| ) |
| predictor = base.load_predictor( |
| pipeline.generator.model, |
| { |
| **predictor_config, |
| "weights": weights, |
| "atc_collect_diagnostics": False, |
| }, |
| 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, predictor, text_encoder, None, None, lpips_model |
|
|
|
|
| def is_complete( |
| output_root: Path, |
| row: dict[str, Any], |
| strategies: list[dict[str, Any]], |
| ) -> bool: |
| return all( |
| all(path.is_file() for path in artifact_paths(output_root, row, config["name"])) |
| for config in strategies |
| ) |
|
|
|
|
| def run_prompt( |
| *, |
| row: dict[str, Any], |
| strategies: list[dict[str, Any]], |
| reference_root: Path, |
| output_root: Path, |
| seed: int, |
| models: tuple[Any, ...], |
| device: torch.device, |
| ) -> None: |
| global_index = int(row["global_index"]) |
| suite = str(row["prompt_suite"]) |
| suite_index = int(row["suite_index"]) |
| prompt = str(row["extended_prompt"]) |
| reference_path = ( |
| reference_root / "generated_videos/ffff" / suite / f"{suite_index:03d}.mp4" |
| ) |
| print(f"[prompt] global={global_index} suite={suite}/{suite_index}", flush=True) |
| conditional = models[3](text_prompts=[prompt]) |
| for strategy in strategies: |
| name = str(strategy["name"]) |
| video_path, record_path = artifact_paths(output_root, row, name) |
| if video_path.is_file() and record_path.is_file(): |
| print(f"[cached-strategy] global={global_index} strategy={name}", flush=True) |
| continue |
| latent, diagnostic = extended.generate_rollout( |
| pipeline=models[1], |
| conditional_dict=conditional, |
| seed=seed, |
| device=device, |
| predictor=models[2], |
| head=None, |
| config=strategy, |
| ) |
| with torch.autocast(device_type="cuda", dtype=torch.bfloat16): |
| pixels = models[0].decode_to_pixel(latent, use_cache=False) |
| prediction_u8 = base.pixels_to_u8(pixels) |
| extended.atomic_video(prediction_u8, video_path) |
| reference_decoded = read_u8_video(reference_path) |
| prediction_decoded = read_u8_video(video_path) |
| metrics = base.frame_metrics( |
| reference_u8=reference_decoded, |
| prediction_u8=prediction_decoded, |
| lpips_model=models[6], |
| batch_size=4, |
| device=device, |
| ) |
| extended.atomic_json( |
| record_path, |
| { |
| "status": "complete", |
| "strategy": name, |
| "policy": "fppf", |
| "pattern": strategy["pattern"], |
| "input_variant": "atc", |
| "atc_previous_scope": strategy["atc_previous_scope"], |
| "candidate_steps": strategy["candidate_steps"], |
| "allow_chunk0_predictor": False, |
| "global_index": global_index, |
| "prompt_suite": suite, |
| "suite_index": suite_index, |
| "prompt": prompt, |
| "seed": seed, |
| "generation": { |
| key: value |
| for key, value in diagnostic.items() |
| if key != "decisions" |
| }, |
| "decisions": diagnostic["decisions"], |
| "pixel_metrics_vs_ffff": extended.compact_metrics(metrics), |
| "pixel_metric_input": ( |
| "MP4-decoded uint8 RGB, all 81 frames on both sides" |
| ), |
| "reference_video": str(reference_path), |
| "video": str(video_path.relative_to(output_root)), |
| }, |
| ) |
| print( |
| f"[result] global={global_index} strategy={name} " |
| f"accept={diagnostic['accepted_predictor_calls']} " |
| f"psnr={metrics['psnr']:.4f} ssim={metrics['ssim']:.6f} " |
| f"lpips={metrics['lpips']:.6f}", |
| flush=True, |
| ) |
| if hasattr(models[0].model, "clear_cache"): |
| models[0].model.clear_cache() |
| del latent, pixels, prediction_u8, reference_decoded, prediction_decoded |
| torch.cuda.empty_cache() |
| del conditional |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--gpu", default=extended.PHYSICAL_GPU) |
| parser.add_argument("--scope", choices=("chunk", "last_frame"), required=True) |
| parser.add_argument("--predictor-weights", type=Path, required=True) |
| parser.add_argument("--mapping", type=Path, default=MAPPING_DEFAULT) |
| parser.add_argument("--reference-root", type=Path, default=REFERENCE_DEFAULT) |
| parser.add_argument("--output-root", type=Path, default=OUTPUT_DEFAULT) |
| parser.add_argument("--pattern", action="append", choices=tuple(PATTERN_TO_STEPS)) |
| parser.add_argument("--shard-index", type=int, default=0) |
| parser.add_argument("--num-shards", type=int, default=1) |
| parser.add_argument("--global-index", action="append", type=int, default=None) |
| parser.add_argument("--seed", type=int, default=0) |
| 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") |
| weights = args.predictor_weights.resolve() |
| mapping = extended.read_mapping(args.mapping.resolve()) |
| if args.global_index is None: |
| rows = [ |
| row for index, row in enumerate(mapping) |
| if index % args.num_shards == args.shard_index |
| ] |
| else: |
| requested = set(args.global_index) |
| rows = [row for row in mapping if int(row["global_index"]) in requested] |
| if len(rows) != len(requested): |
| raise ValueError("Unknown global index") |
| patterns = args.pattern or ["FPPF", "FPPP"] |
| strategies = strategy_configs(args.scope, patterns) |
| reference_root = args.reference_root.resolve() |
| if len(list((reference_root / "generated_videos/ffff").rglob("*.mp4"))) != 251: |
| raise ValueError("Incomplete FFFF reference") |
| output_root = args.output_root.resolve() |
| output_root.mkdir(parents=True, exist_ok=True) |
| extended.write_shard_manifest( |
| output_root, |
| str(args.gpu), |
| args.shard_index, |
| args.num_shards, |
| rows, |
| strategies, |
| "running", |
| ) |
| device = torch.device("cuda") |
| torch.set_grad_enabled(False) |
| models = load_models(weights, args.scope, device) |
| completed = 0 |
| try: |
| for row in rows: |
| if is_complete(output_root, row, strategies): |
| completed += 1 |
| print(f"[cached] {completed}/{len(rows)} global={row['global_index']}", flush=True) |
| continue |
| run_prompt( |
| row=row, |
| strategies=strategies, |
| reference_root=reference_root, |
| output_root=output_root, |
| seed=args.seed, |
| models=models, |
| device=device, |
| ) |
| completed += 1 |
| print(f"[progress] {completed}/{len(rows)} prompts", flush=True) |
| finally: |
| extended.write_shard_manifest( |
| output_root, |
| str(args.gpu), |
| args.shard_index, |
| args.num_shards, |
| rows, |
| strategies, |
| "complete" if completed == len(rows) else "failed", |
| ) |
| print(f"[complete] gpu={args.gpu} scope={args.scope} prompts={completed}", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
