scripts/release_tools.py

from __future__ import annotations

import copy
import math
import statistics
from pathlib import Path
from typing import Any, Iterable

from PIL import Image, ImageDraw, ImageFont


PAIRING_FIELDS = (
    "prompt_id",
    "prompt",
    "seed",
    "width",
    "height",
    "num_inference_steps",
    "guidance_scale",
)


def json_safe_loading_info(loading_info: dict[str, Any]) -> dict[str, Any]:
    """Normalize Hugging Face loading diagnostics for stable JSON output."""

    result: dict[str, Any] = {}
    for key, value in loading_info.items():
        if isinstance(value, set):
            result[key] = sorted(value)
        elif isinstance(value, tuple):
            result[key] = list(value)
        else:
            result[key] = value
    return result


def rewrite_modular_component_sources(
    model_index: dict[str, Any], source: str
) -> dict[str, Any]:
    """Point component specs in a modular model index at one release tree."""

    rewritten = copy.deepcopy(model_index)
    for value in rewritten.values():
        if (
            isinstance(value, list)
            and len(value) >= 3
            and isinstance(value[0], str)
            and value[0] in {"diffusers", "transformers"}
            and isinstance(value[2], dict)
            and "subfolder" in value[2]
        ):
            value[2]["pretrained_model_name_or_path"] = source
            value[2]["revision"] = None
    return rewritten


def validate_paired_records(
    original: list[dict[str, Any]], orbitquant: list[dict[str, Any]]
) -> None:
    if len(original) != len(orbitquant):
        raise ValueError(
            f"paired record count differs: original={len(original)}, "
            f"orbitquant={len(orbitquant)}"
        )
    for index, (original_row, quantized_row) in enumerate(zip(original, orbitquant)):
        for field in PAIRING_FIELDS:
            if original_row.get(field) != quantized_row.get(field):
                raise ValueError(
                    f"paired record {index} differs in {field}: "
                    f"{original_row.get(field)!r} != {quantized_row.get(field)!r}"
                )


def summarize_metrics(rows: list[dict[str, Any]]) -> dict[str, Any]:
    if not rows:
        raise ValueError("cannot summarize an empty metric set")
    seconds = [float(row["generation_seconds"]) for row in rows]
    hot = seconds[1:] or seconds
    result: dict[str, Any] = {
        "generated_samples": len(rows),
        "first_generation_seconds": seconds[0],
        "hot_generation_mean_seconds": statistics.fmean(hot),
        "hot_generation_median_seconds": statistics.median(hot),
        "generation_mean_seconds": statistics.fmean(seconds),
        "generation_median_seconds": statistics.median(seconds),
    }
    for field in ("gpu_peak_mb", "torch_peak_mb"):
        values = [row.get(field) for row in rows if row.get(field) is not None]
        result[field] = max(values) if values else None
    return result


def _font(label_height: int) -> ImageFont.ImageFont | ImageFont.FreeTypeFont:
    size = max(10, label_height // 3)
    for name in ("DejaVuSans-Bold.ttf", "Arial Bold.ttf", "Arial.ttf"):
        try:
            return ImageFont.truetype(name, size=size)
        except OSError:
            continue
    return ImageFont.load_default()


def _open_native_rgb(path: Path, tile_size: tuple[int, int]) -> Image.Image:
    with Image.open(path) as source:
        if source.size != tile_size:
            raise ValueError(
                f"{path} has size {source.size}; expected native tile size {tile_size}"
            )
        return source.convert("RGB")


def create_full_resolution_matrix(
    pairs: Iterable[dict[str, Any]],
    output_path: str | Path,
    *,
    tile_size: tuple[int, int] = (2048, 2048),
    prompt_pairs_per_row: int = 2,
    label_height: int = 96,
) -> dict[str, Any]:
    records = list(pairs)
    if not records:
        raise ValueError("at least one image pair is required")
    if prompt_pairs_per_row <= 0:
        raise ValueError("prompt_pairs_per_row must be positive")
    if label_height <= 0:
        raise ValueError("label_height must be positive")

    tile_width, tile_height = tile_size
    row_count = math.ceil(len(records) / prompt_pairs_per_row)
    matrix_width = prompt_pairs_per_row * 2 * tile_width
    row_height = label_height + tile_height
    matrix_height = row_count * row_height
    matrix = Image.new("RGB", (matrix_width, matrix_height), "#111111")
    draw = ImageDraw.Draw(matrix)
    font = _font(label_height)

    for index, record in enumerate(records):
        row = index // prompt_pairs_per_row
        group = index % prompt_pairs_per_row
        x = group * 2 * tile_width
        label_y = row * row_height
        image_y = label_y + label_height
        original = _open_native_rgb(Path(record["original"]), tile_size)
        quantized = _open_native_rgb(Path(record["orbitquant"]), tile_size)
        matrix.paste(original, (x, image_y))
        matrix.paste(quantized, (x + tile_width, image_y))
        seed = record.get("seed", "-")
        title = str(record.get("title", f"Prompt {index + 1}"))
        draw.text(
            (x + 12, label_y + 4),
            f"{index + 1:02d} {title} · BF16 · seed {seed}",
            fill="white",
            font=font,
        )
        draw.text(
            (x + tile_width + 12, label_y + 4),
            f"{index + 1:02d} {title} · OrbitQuant W4A4 · seed {seed}",
            fill="white",
            font=font,
        )

    destination = Path(output_path)
    destination.parent.mkdir(parents=True, exist_ok=True)
    if destination.suffix.lower() == ".webp":
        matrix.save(destination, format="WEBP", lossless=True, quality=100, method=6)
    else:
        matrix.save(destination)
    return {
        "matrix_path": str(destination),
        "matrix_size": [matrix_width, matrix_height],
        "tile_size": [tile_width, tile_height],
        "prompt_count": len(records),
        "prompt_pairs_per_row": prompt_pairs_per_row,
        "label_height": label_height,
        "resized": False,
    }