| |
| """Merge per-prompt Causal FPPF pixel metrics produced by GPU shards.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import csv |
| import json |
| import math |
| from pathlib import Path |
|
|
|
|
| ROLLOUT_START_FRAME = 9 |
|
|
|
|
| def mean_std(values: list[float]) -> tuple[float, float]: |
| mean = sum(values) / len(values) |
| variance = sum((value - mean) ** 2 for value in values) / len(values) |
| return mean, math.sqrt(variance) |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--output_dir", type=Path, required=True) |
| parser.add_argument("--experiment", required=True) |
| parser.add_argument("--num_shards", type=int, default=4) |
| args = parser.parse_args() |
|
|
| prompt_results = [] |
| for shard in range(args.num_shards): |
| directory = ( |
| args.output_dir / f"shard_gpu{shard}" / args.experiment / "per_prompt" |
| ) |
| for path in sorted(directory.glob("prompt_*.json")): |
| prompt_results.append(json.loads(path.read_text(encoding="utf-8"))) |
| prompt_results.sort(key=lambda item: int(item["prompt_id"])) |
| ids = [int(item["prompt_id"]) for item in prompt_results] |
| if ids != list(range(80, 100)): |
| raise RuntimeError(f"Expected prompt IDs 80..99, got {ids}") |
|
|
| mse = [x for item in prompt_results for x in item["mse_per_frame"]] |
| psnr = [x for item in prompt_results for x in item["psnr_per_frame"]] |
| ssim = [x for item in prompt_results for x in item["ssim_per_frame"]] |
| lpips = [x for item in prompt_results for x in item["lpips_per_frame"]] |
| rollout_mse = [ |
| x for item in prompt_results for x in item["mse_per_frame"][ROLLOUT_START_FRAME:] |
| ] |
| rollout_ssim = [ |
| x for item in prompt_results for x in item["ssim_per_frame"][ROLLOUT_START_FRAME:] |
| ] |
| rollout_lpips = [ |
| x for item in prompt_results for x in item["lpips_per_frame"][ROLLOUT_START_FRAME:] |
| ] |
| pixel_mse = sum(mse) / len(mse) |
| psnr_mean, psnr_std = mean_std(psnr) |
| ssim_mean, ssim_std = mean_std(ssim) |
| lpips_mean, lpips_std = mean_std(lpips) |
| rollout_pixel_mse = sum(rollout_mse) / len(rollout_mse) |
| metrics = { |
| "status": "complete", |
| "name": args.experiment, |
| "source_layer": 17, |
| "checkpoint_step": 2000, |
| "schedule": "chunk0=FFFF; chunks1-6=FPPF", |
| "reference": "matching Causal-Forcing FFFF, same prompt and seed", |
| "pixel_domain": "VAE-decoded RGB, rounded to uint8", |
| "aggregation": "PSNR from global pixel MSE; SSIM/LPIPS mean over decoded frames", |
| "num_prompts": len(prompt_results), |
| "num_frames": len(psnr), |
| "prompt_ids": ids, |
| "pixel_mse": pixel_mse, |
| "psnr": -10.0 * math.log10(max(pixel_mse, 1e-12)), |
| "psnr_frame_mean": psnr_mean, |
| "psnr_frame_std": psnr_std, |
| "ssim": ssim_mean, |
| "ssim_frame_std": ssim_std, |
| "lpips": lpips_mean, |
| "lpips_frame_std": lpips_std, |
| "rollout_start_frame": ROLLOUT_START_FRAME, |
| "rollout_num_frames": len(rollout_mse), |
| "rollout_pixel_mse": rollout_pixel_mse, |
| "rollout_psnr": -10.0 * math.log10(max(rollout_pixel_mse, 1e-12)), |
| "rollout_ssim": sum(rollout_ssim) / len(rollout_ssim), |
| "rollout_lpips": sum(rollout_lpips) / len(rollout_lpips), |
| "mean_generation_time_s": sum(x["generation_time_s"] for x in prompt_results) |
| / len(prompt_results), |
| "full_calls_per_prompt": prompt_results[0]["full_calls"], |
| "predictor_calls_per_prompt": prompt_results[0]["predictor_calls"], |
| } |
| destination = args.output_dir / args.experiment |
| destination.mkdir(parents=True, exist_ok=True) |
| (destination / "metrics.json").write_text( |
| json.dumps(metrics, indent=2, ensure_ascii=False) + "\n", encoding="utf-8" |
| ) |
| fields = [ |
| "name", "source_layer", "checkpoint_step", "num_prompts", "num_frames", |
| "psnr", "ssim", "lpips", "rollout_psnr", "rollout_ssim", "rollout_lpips", |
| "mean_generation_time_s", |
| ] |
| with (args.output_dir / "summary.csv").open("w", encoding="utf-8", newline="") as f: |
| writer = csv.DictWriter(f, fieldnames=fields) |
| writer.writeheader() |
| writer.writerow({key: metrics[key] for key in fields}) |
| print(json.dumps(metrics, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
