scripts/37_assets.py

"""Build writeup assets from the eval run: a merged results table (markdown) + a comparison montage.

Reads outputs/eval/metrics.json + outputs/eval/timing_*.json. Montage shows a few prompts across
model dirs (teacher | plain | SVDQuant | ...).
Usage: python3 scripts/37_assets.py [idx1,idx2,...]   (montage prompt indices; default a category spread)
"""
import sys, json, os, glob
from PIL import Image, ImageDraw

metrics = json.load(open('outputs/eval/metrics.json'))
timing = {}
for f in glob.glob('outputs/eval/timing_*.json'):
    t = json.load(open(f)); timing[t.get('tag', os.path.basename(f))] = t

# ---- merged markdown table ----
lines = ["# Eval results — SVDQuant vs plain NVFP4 (klein-4B)\n",
         "| model | PickScore↑ | CLIP↑ | LPIPS↓(vs teacher) | PSNR↑ | FID↓(real) | FID↓(teacher) | s/img | step ms | VRAM GB |",
         "|---|---|---|---|---|---|---|---|---|---|"]
for name, v in metrics.items():
    tk = next((timing[k] for k in timing if k in name or name in k), {})
    spi = tk.get('s_per_img_median') or tk.get('s_per_img_median_batched') or '–'
    lines.append(f"| {name} | {v.get('PickScore')} | {v.get('CLIP')} | {v.get('LPIPS')} | {v.get('PSNR')} | "
                 f"{v.get('FID_vs_real')} | {v.get('FID_vs_teacher')} | {spi} | {tk.get('step_ms_median','–')} | {tk.get('peak_vram_gb','–')} |")
md = "\n".join(lines) + "\n"
open('outputs/eval/results_table.md', 'w').write(md)
print(md)

# ---- comparison montage ----
prompts = {d['idx']: d for d in json.load(open('outputs/eval/prompts.json'))}
if len(sys.argv) > 1:
    idxs = [int(x) for x in sys.argv[1].split(',')]
else:
    # one per category from the first 64 (pilot range), preferring text/people/logo
    seen, idxs = set(), []
    for i in range(64):
        c = prompts.get(i, {}).get('category')
        if c and c not in seen:
            seen.add(c); idxs.append(i)
        if len(idxs) >= 6:
            break
# model columns that exist (label, dir)
cand = [('teacher', 'outputs/eval/imgs/teacher'),
        ('plain NVFP4 (r0)', 'outputs/eval/imgs/pilot_fq0'),
        ('SVDQuant r64', 'outputs/eval/imgs/pilot_fq64'),
        ('SVDQuant r128 (fake-q)', 'outputs/eval/imgs/pilot_fq128'),
        ('SVDQuant r128 (real kernel)', 'outputs/eval/imgs/pilot_ours128_real'),
        ('BFL plain NVFP4', 'outputs/eval/imgs/bfl')]
cols = [(lbl, d) for lbl, d in cand if os.path.isdir(d) and glob.glob(os.path.join(d, '*.png'))]
S, pad, top = 320, 6, 22
W = len(cols) * (S + pad) + pad
H = len(idxs) * (S + pad) + top + pad
canvas = Image.new('RGB', (W, top + H), 'white')
dr = ImageDraw.Draw(canvas)
for ci, (lbl, _) in enumerate(cols):
    dr.text((pad + ci * (S + pad) + 4, 4), lbl, fill='black')
for ri, idx in enumerate(idxs):
    for ci, (lbl, d) in enumerate(cols):
        p = os.path.join(d, f"{idx:05d}.png")
        if not os.path.exists(p):
            continue
        im = Image.open(p).convert('RGB').resize((S, S))
        x = pad + ci * (S + pad); y = top + ri * (S + pad)
        canvas.paste(im, (x, y))
        dr.text((x + 3, y + 3), prompts[idx]['category'], fill='yellow')
canvas.save('outputs/eval/montage_compare.png')
print(f"saved montage -> outputs/eval/montage_compare.png ({len(idxs)} prompts x {len(cols)} models)")