#!/usr/bin/env python3 """Compute paired RGB metrics for publication candidates against Original BF16.""" from __future__ import annotations import argparse import hashlib import json import math import os from datetime import datetime, timezone from pathlib import Path from typing import Any import numpy as np from PIL import Image REFERENCE_ID = "original_bf16" EXPECTED_GRID_COUNT = 20 CHUNK_BYTES = 4 * 1024 * 1024 class MetricError(RuntimeError): pass def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--grid-manifest", type=Path, required=True) parser.add_argument("--output", type=Path, required=True) return parser.parse_args() def sha256_file(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as handle: for chunk in iter(lambda: handle.read(CHUNK_BYTES), b""): digest.update(chunk) return digest.hexdigest() def psnr(mse: float) -> float | str: if mse == 0.0: return "infinity" return 10.0 * math.log10((255.0**2) / mse) def load_rgb(record: dict[str, Any]) -> np.ndarray: path = Path(str(record["path"])).resolve(strict=True) if path.stat().st_size != int(record["bytes"]): raise MetricError(f"byte drift: {path}") if sha256_file(path) != str(record["sha256"]): raise MetricError(f"hash drift: {path}") with Image.open(path) as image: if image.size != (1024, 1024): raise MetricError(f"dimension drift: {path}: {image.size}") return np.asarray(image.convert("RGB"), dtype=np.float64) def atomic_json(path: Path, value: dict[str, Any]) -> None: path.parent.mkdir(parents=True, exist_ok=True) temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp") temporary.write_text( json.dumps(value, indent=2, sort_keys=True) + "\n", encoding="utf-8" ) os.replace(temporary, path) def main() -> int: args = parse_args() manifest_path = args.grid_manifest.expanduser().resolve(strict=True) manifest = json.loads(manifest_path.read_text(encoding="utf-8")) grids = manifest.get("grids") if not isinstance(grids, list) or len(grids) != EXPECTED_GRID_COUNT: raise MetricError(f"expected {EXPECTED_GRID_COUNT} verified grids") if manifest.get("status") != "verified" or manifest.get("profile") != "publication": raise MetricError("grid manifest is not a verified publication profile") expected_models: list[tuple[str, str]] | None = None accumulators: dict[str, dict[str, Any]] = {} prompt_ids: set[str] = set() for grid in grids: prompt_id = str(grid["prompt_id"]) if prompt_id in prompt_ids: raise MetricError(f"duplicate prompt: {prompt_id}") prompt_ids.add(prompt_id) columns = grid.get("columns") if not isinstance(columns, list): raise MetricError(f"missing columns for {prompt_id}") by_id = {str(column["model_id"]): column for column in columns} if len(by_id) != len(columns) or REFERENCE_ID not in by_id: raise MetricError(f"invalid model inventory for {prompt_id}") models = [(str(column["model_id"]), str(column["label"])) for column in columns] if expected_models is None: expected_models = models for model_id, label in models: if model_id != REFERENCE_ID: accumulators[model_id] = { "label": label, "sum_abs_error": 0.0, "sum_squared_error": 0.0, "value_count": 0, "pairs": [], } elif models != expected_models: raise MetricError(f"column order drift for {prompt_id}") reference = load_rgb(by_id[REFERENCE_ID]) for model_id, _label in models: if model_id == REFERENCE_ID: continue candidate = load_rgb(by_id[model_id]) difference = candidate - reference absolute_error = float(np.abs(difference).sum(dtype=np.float64)) squared_error = float(np.square(difference).sum(dtype=np.float64)) value_count = int(difference.size) pair_mse = squared_error / value_count accumulator = accumulators[model_id] accumulator["sum_abs_error"] += absolute_error accumulator["sum_squared_error"] += squared_error accumulator["value_count"] += value_count accumulator["pairs"].append( { "prompt_id": prompt_id, "seed": int(grid["seed"]), "mae_rgb_8bit": absolute_error / value_count, "mse_rgb_8bit": pair_mse, "psnr_db": psnr(pair_mse), } ) if expected_models is None: raise MetricError("no models found") results: list[dict[str, Any]] = [] for model_id, label in expected_models: if model_id == REFERENCE_ID: continue accumulator = accumulators[model_id] value_count = int(accumulator["value_count"]) global_mse = float(accumulator["sum_squared_error"]) / value_count global_mae = float(accumulator["sum_abs_error"]) / value_count pair_scores = [float(pair["psnr_db"]) for pair in accumulator["pairs"]] results.append( { "model_id": model_id, "label": label, "reference_model_id": REFERENCE_ID, "pair_count": len(pair_scores), "global_mae_rgb_8bit": global_mae, "global_mse_rgb_8bit": global_mse, "global_psnr_db": psnr(global_mse), "mean_pair_psnr_db": float(np.mean(pair_scores)), "median_pair_psnr_db": float(np.median(pair_scores)), "min_pair_psnr_db": min(pair_scores), "max_pair_psnr_db": max(pair_scores), "pairs": accumulator["pairs"], } ) output = { "schema_version": 1, "status": "verified_complete", "created_at": datetime.now(timezone.utc).isoformat(timespec="seconds"), "source_grid_manifest_sha256": sha256_file(manifest_path), "reference": {"model_id": REFERENCE_ID, "label": "Original BF16"}, "contract": { "prompt_count": EXPECTED_GRID_COUNT, "width": 1024, "height": 1024, "channels": "RGB", "pixel_range": "8-bit [0,255]", "same_prompt_seed_sampler_and_settings": True, }, "results": results, "limitations": [ "PSNR measures paired pixel similarity to the BF16 reference, not aesthetics or prompt adherence.", "LPIPS and human-preference metrics were not computed.", ], } atomic_json(args.output.expanduser().resolve(), output) print(json.dumps({"status": output["status"], "results": results}, indent=2)) return 0 if __name__ == "__main__": raise SystemExit(main())