| |
| """Aggregate MovieBench pixel metrics and prepare/read VBench-5 inputs.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import csv |
| import json |
| import math |
| import os |
| from pathlib import Path |
|
|
|
|
| DIMENSIONS = [ |
| "subject_consistency", "background_consistency", "motion_smoothness", |
| "aesthetic_quality", "imaging_quality", |
| ] |
| METHODS = ("ffff", "fppf_step2000") |
|
|
|
|
| def atomic_json(path: Path, value) -> None: |
| path.parent.mkdir(parents=True, exist_ok=True) |
| tmp = path.with_suffix(path.suffix + ".tmp") |
| tmp.write_text(json.dumps(value, indent=2, ensure_ascii=False) + "\n", encoding="utf-8") |
| os.replace(tmp, path) |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--output_dir", type=Path, required=True) |
| args = parser.parse_args() |
| rows = [] |
| for shard in range(4): |
| directory = args.output_dir / f"shard_gpu{shard}" / "per_prompt" |
| rows.extend(json.loads(path.read_text()) for path in sorted(directory.glob("*.json"))) |
| rows.sort(key=lambda x: x["prompt_id"]) |
| ids = [row["prompt_id"] for row in rows] |
| if ids != list(range(100)): |
| raise RuntimeError(f"Expected MovieBench IDs 0..99, got {ids}") |
|
|
| mse = [x for row in rows for x in row["mse_per_frame"]] |
| ssim = [x for row in rows for x in row["ssim_per_frame"]] |
| lpips = [x for row in rows for x in row["lpips_per_frame"]] |
| rollout_mse = [x for row in rows for x in row["mse_per_frame"][9:]] |
| rollout_ssim = [x for row in rows for x in row["ssim_per_frame"][9:]] |
| rollout_lpips = [x for row in rows for x in row["lpips_per_frame"][9:]] |
| pixel_metrics = { |
| "status": "complete", "checkpoint_step": 2000, "source_layer": 17, |
| "num_prompts": 100, "num_frames": len(mse), "prompt_ids": ids, |
| "generation_prompts": "MovieGenVideoBench_extended.txt first 100", |
| "evaluation_prompts": "MovieGenVideoBench.txt first 100", |
| "schedule": "chunk0=FFFF; chunks1-6=FPPF", |
| "psnr": -10 * math.log10(sum(mse) / len(mse)), |
| "ssim": sum(ssim) / len(ssim), "lpips": sum(lpips) / len(lpips), |
| "rollout_psnr": -10 * math.log10(sum(rollout_mse) / len(rollout_mse)), |
| "rollout_ssim": sum(rollout_ssim) / len(rollout_ssim), |
| "rollout_lpips": sum(rollout_lpips) / len(rollout_lpips), |
| "mean_ffff_generation_time_s": sum(x["ffff"]["generation_time_s"] for x in rows) / 100, |
| "mean_fppf_generation_time_s": sum(x["fppf"]["generation_time_s"] for x in rows) / 100, |
| } |
| atomic_json(args.output_dir / "pixel_metrics.json", pixel_metrics) |
|
|
| for method in METHODS: |
| video_dir = args.output_dir / "vbench_inputs" / method |
| video_dir.mkdir(parents=True, exist_ok=True) |
| info = [] |
| for row in rows: |
| prompt_id = row["prompt_id"] |
| matches = list(args.output_dir.glob(f"shard_gpu*/videos/{method}/{prompt_id:05d}.mp4")) |
| if len(matches) != 1: |
| raise RuntimeError(f"Expected one {method} video for {prompt_id}: {matches}") |
| target = video_dir / f"{prompt_id:05d}.mp4" |
| if target.exists() or target.is_symlink(): |
| target.unlink() |
| target.symlink_to(matches[0].resolve()) |
| info.append({ |
| "video_list": [target.name], "prompt_en": row["original_prompt"], |
| "dimension": DIMENSIONS, |
| }) |
| atomic_json(video_dir / "full_info.json", info) |
|
|
| vbench_rows = [] |
| for method in METHODS: |
| result_dir = args.output_dir / "vbench_results" / method |
| matches = list(result_dir.glob("*_eval_results.json")) |
| if not matches: |
| continue |
| if len(matches) != 1: |
| raise RuntimeError(f"Expected one VBench result for {method}: {matches}") |
| result = json.loads(matches[0].read_text()) |
| scores = {dimension: float(result[dimension][0]) for dimension in DIMENSIONS} |
| vbench_rows.append({"method": method, **scores, "vbench5_mean": sum(scores.values()) / 5}) |
| if vbench_rows: |
| atomic_json(args.output_dir / "vbench_summary.json", vbench_rows) |
| with (args.output_dir / "vbench_summary.csv").open("w", newline="", encoding="utf-8") as f: |
| writer = csv.DictWriter(f, fieldnames=list(vbench_rows[0])) |
| writer.writeheader(); writer.writerows(vbench_rows) |
| print(json.dumps({"pixel_metrics": pixel_metrics, "vbench": vbench_rows}, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
