File size: 4,349 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 104 105 106 107 108 109 | #!/usr/bin/env python3
"""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()
|