"""Render portable, fail-closed representative cross-codec visual tables.""" from __future__ import annotations import argparse import hashlib import io import json import logging import os import stat import textwrap from pathlib import Path from typing import Any from PIL import Image, ImageDraw, ImageFont LOGGER = logging.getLogger(__name__) CODECS: dict[str, tuple[str, str, str, str]] = { "jpeg_q100": ("JPEG Q100", "jpeg_q100_raw", "jpeg_q100_repair_zero", "jpeg_q100_protected_high"), "jpeg_q80": ("JPEG Q80", "jpeg_q80_raw", "jpeg_q80_repair_zero", "jpeg_q80_protected_high"), "webp_q80": ("WebP Q80", "webp_q80_raw", "webp_q80_repair_zero", "webp_q80_protected_high"), "avif_q70": ("AVIF Q70", "avif_q70_raw", "avif_q70_repair_zero", "avif_q70_protected_high"), "jxl_d1": ("JPEG XL d1", "jxl_d1_raw", "jxl_d1_repair_zero", "jxl_d1_protected_high"), } COLUMNS = (("Baseline PNG", "png_baseline"), ("WebP Lossless", "webp_lossless")) BG, PANEL, TEXT, MUTED, ACCENT, FAIL = "#0d141b", "#111b24", "#edf5fb", "#9fb4c3", "#8bd5f7", "#ffb1c6" def _font(size: int, mono: bool = False) -> ImageFont.FreeTypeFont | ImageFont.ImageFont: candidates = ( ("/System/Library/Fonts/SFNSMono.ttf", "/System/Library/Fonts/SFNS.ttf") if mono else ("/System/Library/Fonts/SFNS.ttf",) ) for candidate in candidates: if Path(candidate).is_file(): return ImageFont.truetype(candidate, size) return ImageFont.load_default() def load_prompt_ids(path: str | Path) -> list[str]: """Load an ordered prompt-id manifest without guessing missing prompts.""" value = json.loads(Path(path).read_text(encoding="utf-8")) ids = value.get("prompt_ids") if isinstance(value, dict) else value if not isinstance(ids, list) or not ids or not all(isinstance(item, (str, int)) for item in ids): raise ValueError("prompt-id manifest must contain a non-empty prompt_ids list") result = [str(item) for item in ids] if len(result) != len(set(result)): raise ValueError("prompt-id manifest contains duplicates") return result def _sample_hash(path: Path) -> str: return hashlib.sha256(path.read_bytes()).hexdigest() def _read_sample_bytes_no_follow(root: Path, relative: str) -> bytes: """Open every path component without following symlinks and return immutable bytes.""" candidate = Path(relative) if candidate.is_absolute() or not candidate.parts or any(part in {"", ".", ".."} for part in candidate.parts): raise ValueError(f"path escapes run: {relative}") required_flags = ("O_DIRECTORY", "O_NOFOLLOW", "O_NONBLOCK") if any(not hasattr(os, name) for name in required_flags): raise ValueError("safe sample opening is unsupported on this platform") no_follow = os.O_NOFOLLOW directory_flags = os.O_RDONLY | os.O_DIRECTORY | no_follow directory_fd = -1 file_fd = -1 try: directory_fd = os.open(root, directory_flags) for part in candidate.parts[:-1]: next_fd = os.open(part, directory_flags, dir_fd=directory_fd) os.close(directory_fd) directory_fd = next_fd file_fd = os.open(candidate.parts[-1], os.O_RDONLY | no_follow | os.O_NONBLOCK, dir_fd=directory_fd) if not stat.S_ISREG(os.fstat(file_fd).st_mode): raise ValueError(f"successful sample is not a regular file: {relative}") with os.fdopen(file_fd, "rb") as handle: file_fd = -1 return handle.read() except OSError as exc: raise ValueError(f"successful sample cannot be opened without following symlinks: {relative}") from exc finally: if file_fd >= 0: os.close(file_fd) if directory_fd >= 0: os.close(directory_fd) def _safe_run_path(root: Path, relative: str) -> Path: candidate = Path(relative) if candidate.is_absolute(): raise ValueError(f"path escapes run: {relative}") raw_path = root / candidate if raw_path.is_symlink(): raise ValueError(f"path is a symlink: {relative}") path = raw_path.resolve() try: path.relative_to(root) except ValueError as exc: raise ValueError(f"path escapes run: {relative}") from exc return path def load_rows(run_dir: str | Path, prompt_ids: list[str]) -> dict[tuple[str, str], dict[str, Any]]: """Load a complete generation matrix and reject unverifiable success rows.""" root = Path(run_dir).resolve() path = _safe_run_path(root, "results/raw_rows.jsonl") if path.is_symlink() or not path.is_file(): raise ValueError("raw_rows.jsonl is missing or is a symlink") rows: dict[tuple[str, str], dict[str, Any]] = {} conditions = {"png_baseline", "webp_lossless"} | {item for codec in CODECS.values() for item in codec[1:]} for line_number, line in enumerate(path.read_text(encoding="utf-8").splitlines(), 1): if not line.strip(): continue row = json.loads(line) if row.get("row_type") != "generation": continue condition, prompt_id = str(row.get("condition")), str(row.get("prompt_id")) if condition not in conditions or prompt_id not in prompt_ids: continue key = (condition, prompt_id) if key in rows: raise ValueError(f"duplicate generation row: {condition}/{prompt_id}") if row.get("success"): sample = row.get("sample_path") if not isinstance(sample, str) or not sample: raise ValueError(f"successful row has no sample path: {condition}/{prompt_id}") sample_path = _safe_run_path(root, sample) if not sample_path.is_file() or sample_path.is_symlink(): raise ValueError(f"successful sample is missing: {sample}") recorded = row.get("sample_sha256", row.get("sample_hash")) if recorded is not None and recorded != _sample_hash(sample_path): raise ValueError(f"sample hash mismatch: {sample}") else: if row.get("sample_path") or row.get("sample_sha256") or row.get("sample_hash"): raise ValueError(f"failed row contains sample provenance: {condition}/{prompt_id}") rows[key] = row missing = [ (condition, prompt_id) for condition in sorted(conditions) for prompt_id in prompt_ids if (condition, prompt_id) not in rows ] if missing: raise ValueError(f"generation matrix incomplete; first missing cell: {missing[0]}") return rows def _failure_label(row: dict[str, Any]) -> str: if row.get("blocked"): return "NO SAMPLE\n\nWEIGHT GATE\nnon-finite weights" error = str(row.get("error") or row.get("first_error") or "generation failed") return "NO SAMPLE\n\nGENERATION\n" + ("non-finite output" if "non-finite" in error.lower() else "generation failed") def _draw_cell( canvas: Image.Image, draw: ImageDraw.ImageDraw, root: Path, row: dict[str, Any], x: int, y: int, size: int ) -> None: if row.get("success"): sample_bytes = _read_sample_bytes_no_follow(root, str(row["sample_path"])) recorded_bytes = row.get("sample_bytes") if recorded_bytes is not None and (not isinstance(recorded_bytes, int) or recorded_bytes != len(sample_bytes)): raise ValueError(f"sample byte count mismatch: {row['sample_path']}") recorded_hash = row.get("sample_sha256", row.get("sample_hash")) if recorded_hash is not None and recorded_hash != hashlib.sha256(sample_bytes).hexdigest(): raise ValueError(f"sample hash mismatch: {row['sample_path']}") with Image.open(io.BytesIO(sample_bytes)) as image: canvas.paste(image.convert("RGB").resize((size, size), Image.Resampling.LANCZOS), (x, y)) draw.rectangle((x, y, x + size, y + size), outline="#263746") if row.get("clip_score") is not None: draw.text((x, y + size + 6), f"CLIP {float(row['clip_score']):.3f}", font=_font(14, True), fill=ACCENT) else: draw.rectangle((x, y, x + size, y + size), fill="#351d28", outline="#e96a91", width=2) draw.multiline_text((x + 12, y + 12), _failure_label(row), font=_font(14, True), fill=FAIL, spacing=4) def render( codec_key: str, codec: tuple[str, str, str, str], root: Path, prompts: list[dict[str, Any]], rows: dict[tuple[str, str], dict[str, Any]], output: Path, ) -> Path: """Render one codec comparison table; failed cells remain explicit panels.""" columns = COLUMNS + ( (f"{codec[0]} Raw", codec[1]), ("Finite Repair Zero", codec[2]), ("Protected High Byte", codec[3]), ) size, label_width, gap, row_height = 220, 285, 30, 292 canvas = Image.new( "RGB", (24 + label_width + len(columns) * (size + gap), 160 + len(prompts) * row_height + 70), BG ) draw = ImageDraw.Draw(canvas) draw.text((24, 20), f"PixelModel-v4 · {codec[0]} robustness", font=_font(25), fill=TEXT) draw.text((24, 57), "100 prompts · 50 steps · CFG 6.0 · representative prompts", font=_font(15, True), fill=MUTED) top = 150 for index, (title, _) in enumerate(columns): x = 24 + label_width + index * (size + gap) draw.multiline_text((x, 95), title, font=_font(16), fill=ACCENT, spacing=3) for index, prompt in enumerate(prompts): y = top + index * row_height draw.text((24, y + 5), str(prompt.get("category", "")).upper(), font=_font(14), fill=ACCENT) draw.text((24, y + 30), str(prompt["prompt_id"]), font=_font(13, True), fill=TEXT) wrapped = "\n".join(textwrap.wrap(str(prompt.get("prompt", "")), width=34)) draw.multiline_text((24, y + 58), wrapped, font=_font(13), fill=TEXT, spacing=3) prompt_id = str(prompt["prompt_id"]) for column, (_, condition) in enumerate(columns): _draw_cell( canvas, draw, root, rows[(condition, prompt_id)], 24 + label_width + column * (size + gap), y, size ) destination = output / f"visual_table_{codec_key}.png" destination.parent.mkdir(parents=True, exist_ok=True) canvas.save(destination, format="PNG", optimize=True) return destination def render_contact_sheet(tables: list[Path], output: Path, *, gap: int = 48) -> Path: """Stack full-width tables vertically so labels remain readable in browser viewers.""" if not tables: raise ValueError("at least one table is required") images: list[Image.Image] = [] for table in tables: with Image.open(table) as image: images.append(image.convert("RGB")) width = max(image.width for image in images) height = sum(image.height for image in images) + gap * (len(images) - 1) canvas = Image.new("RGB", (width, height), BG) y = 0 for index, image in enumerate(images): canvas.paste(image, ((width - image.width) // 2, y)) y += image.height if index < len(images) - 1: divider_y = y + gap // 2 ImageDraw.Draw(canvas).rectangle((24, divider_y - 1, width - 24, divider_y + 1), fill="#263746") y += gap destination = output / "visual_tables_contact_sheet.png" destination.parent.mkdir(parents=True, exist_ok=True) canvas.save(destination, format="PNG", optimize=True) return destination def main() -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--run-dir", type=Path, required=True) parser.add_argument("--prompt-ids-manifest", type=Path, required=True) parser.add_argument("--output", type=Path, required=True) args = parser.parse_args() logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s") prompt_ids = load_prompt_ids(args.prompt_ids_manifest) rows = load_rows(args.run_dir, prompt_ids) prompts = [rows[("png_baseline", prompt_id)] for prompt_id in prompt_ids] output = args.output rendered = [render(key, codec, args.run_dir.resolve(), prompts, rows, output) for key, codec in CODECS.items()] render_contact_sheet(rendered, output) LOGGER.info("rendered %d tables and a vertical contact sheet to %s", len(CODECS), output) return 0 if __name__ == "__main__": raise SystemExit(main())