flux2distill/eval_utils.py

"""Eval helpers: parse prompts, load the student into a pipeline, build comparison grids."""
from __future__ import annotations

import json
import re

import torch
from PIL import Image, ImageDraw


def parse_prompts(path):
    """Return list of (idx, category, prompt). Strips leading [category] tags and # comments."""
    items = []
    with open(path) as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            m = re.match(r"^\[([a-zA-Z]+)\]\s*(.+)$", line)
            cat, prompt = (m.group(1), m.group(2)) if m else ("misc", line)
            items.append((len(items), cat, prompt))
    return items


def load_student(pipe, selection_path, state_path, device="cuda"):
    """Mutate pipe.transformer into the saved student (attach surrogates + load weights)."""
    from .surgery import attach_surrogates
    with open(selection_path) as f:
        sel = json.load(f)
    attach_surrogates(pipe.transformer, sel["surrogate_idx"], kind=sel.get("kind", "lowrank"),
                      rank=sel.get("rank", 512), act=sel.get("act", "gelu"),
                      heads=sel.get("heads", 4), head_dim=sel.get("head_dim", 128),
                      conv_kernel=sel.get("conv_kernel", 5), ffn_hidden=sel.get("ffn_hidden", 1024),
                      ffn_idx=sel.get("ffn_idx", None),
                      device=device, dtype=next(pipe.transformer.parameters()).dtype)
    state = torch.load(state_path, map_location=device)
    missing, unexpected = pipe.transformer.load_state_dict(state, strict=False)
    assert not unexpected, f"unexpected keys: {unexpected[:5]}"
    return pipe, sel


def label(img, text, h=22):
    """Add a caption strip below an image."""
    w = img.width
    out = Image.new("RGB", (w, img.height + h), "white")
    out.paste(img, (0, 0))
    d = ImageDraw.Draw(out)
    d.text((4, img.height + 4), text[:60], fill="black")
    return out


def side_by_side(left, right, left_label, right_label, header):
    l = label(left, left_label)
    r = label(right, right_label)
    w, hh = l.width + r.width, l.height
    canvas = Image.new("RGB", (w, hh + 20), "white")
    d = ImageDraw.Draw(canvas)
    d.text((4, 4), header[:90], fill="black")
    canvas.paste(l, (0, 20))
    canvas.paste(r, (l.width, 20))
    return canvas