Spaces:
Paused
Paused
| #!/usr/bin/env python3 | |
| """Make labelled review sheets from actual benchmark PNGs; never synthesizes outputs.""" | |
| import argparse | |
| import json | |
| from pathlib import Path | |
| from PIL import Image, ImageDraw, ImageFont | |
| def main(): | |
| parser = argparse.ArgumentParser(description=__doc__) | |
| parser.add_argument("results", type=Path) | |
| parser.add_argument("--out", type=Path) | |
| args = parser.parse_args() | |
| rows = [json.loads(line) for line in args.results.read_text().splitlines()] | |
| rows = [r for r in rows if r["status"] == "ok"] | |
| target = args.out or args.results.parent / "review" | |
| target.mkdir(parents=True, exist_ok=True) | |
| font = ImageFont.load_default(size=17) | |
| for prompt in sorted({r["trial"]["prompt_index"] for r in rows}): | |
| subset = [r for r in rows if r["trial"]["prompt_index"] == prompt] | |
| seeds = sorted({r["trial"]["seed"] for r in subset}) | |
| sheet = Image.new("RGB", (1200, len(seeds) * 440), "#f4f3ee") | |
| draw = ImageDraw.Draw(sheet) | |
| for r in subset: | |
| p = r["recipe"]["parameters"] | |
| col = [20, 30, 40].index(p["steps"]) | |
| y = seeds.index(p["seed"]) * 440 | |
| with Image.open(args.results.parent / r["image_file"]) as original: | |
| thumb = original.convert("RGBA") | |
| thumb.thumbnail((400, 400), Image.Resampling.LANCZOS) | |
| background = Image.new("RGBA", (400, 400), "#eeeeee") | |
| grid = ImageDraw.Draw(background) | |
| for gy in range(0, 400, 16): | |
| for gx in range(0, 400, 16): | |
| if (gx // 16 + gy // 16) % 2: | |
| grid.rectangle((gx, gy, gx + 15, gy + 15), fill="#d0d0d0") | |
| background.alpha_composite(thumb, ((400-thumb.width)//2, (400-thumb.height)//2)) | |
| sheet.paste(background.convert("RGB"), (col*400, y)) | |
| draw.text((col*400+8, y+405), | |
| f"{p['steps']} steps / seed {p['seed']} / {r['recipe']['metrics']['inference_seconds']:.2f}s", | |
| font=font, fill="#222222") | |
| destination = target / f"prompt-{prompt}.jpg" | |
| sheet.save(destination, quality=95) | |
| print(destination) | |
| if __name__ == "__main__": | |
| main() | |