File size: 4,498 Bytes
ae8ade0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | #!/usr/bin/env python3
"""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()
|