scripts/build_report_pdf.py

"""Build the detailed W4A8 SVDQuant PDF report:
  report/QUANT_REPORT.md  (body, detailed)
  + inline analysis figures (outputs/quant_report_assets/*.png)
  + APPENDIX: every cell's 8-probe teacher|quant montage (all 12 cells)
-> report/QUANT_REPORT.pdf

Run: source .venv/bin/activate && PYTHONPATH=. python scripts/build_report_pdf.py
"""
import os
import markdown
from PIL import Image, ImageDraw, ImageFont
from weasyprint import HTML

AST = "outputs/quant_report_assets"
APP = f"{AST}/appendix"
os.makedirs(APP, exist_ok=True)

RANKS = [16, 32, 64]
VARIANTS = ["plain", "plain_refine", "whiten", "whiten_refine"]
VLABEL = {"plain": "plain", "plain_refine": "plain+refine",
          "whiten": "whiten", "whiten_refine": "whiten+refine"}
SMALLER = {16: "3.67×", 32: "3.59×", 64: "3.43×"}
PROBES = ["storefront text", "mountain lake", "fisherman portrait", "neon street",
          "chalkboard 'FRESH COFFEE'", "breakfast flat-lay (3 eggs/2 bacon)",
          "hand / five fingers", "dewy spiderweb macro"]
# full metrics per cell: eval-loss, vel-relerr, wrecon, orecon('' if n/a)
M = {
    (16, "plain"): (0.0620, 0.2235, 0.1269, ""), (16, "plain_refine"): (0.0655, 0.2297, 0.1251, ""),
    (16, "whiten"): (0.0656, 0.2299, 0.1290, 0.0733), (16, "whiten_refine"): (0.0556, 0.2117, 0.1273, 0.0680),
    (32, "plain"): (0.0586, 0.2174, 0.1224, ""), (32, "plain_refine"): (0.0574, 0.2151, 0.1193, ""),
    (32, "whiten"): (0.0545, 0.2095, 0.1257, 0.0719), (32, "whiten_refine"): (0.0476, 0.1959, 0.1226, 0.0646),
    (64, "plain"): (0.0487, 0.1980, 0.1163, ""), (64, "plain_refine"): (0.0446, 0.1896, 0.1110, ""),
    (64, "whiten"): (0.0588, 0.2177, 0.1209, 0.0695), (64, "whiten_refine"): (0.0451, 0.1907, 0.1155, 0.0595),
}
BEST = (64, "plain_refine")


def _font(sz):
    for p in ["/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
              "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf"]:
        if os.path.exists(p):
            try:
                return ImageFont.truetype(p, sz)
            except Exception:
                pass
    return ImageFont.load_default()


def cell_grid_montage(rank, variant, out_path, ncols=2, cell_w=560, label_h=26):
    """8 probe pairs (teacher|quant) for one cell, arranged ncols-wide, each labeled."""
    pairs = []
    for i in range(8):
        p = f"outputs/abl_c300_r{rank}_{variant}/eval/cmp_{i}.png"
        if os.path.exists(p):
            pairs.append((PROBES[i], Image.open(p)))
    if not pairs:
        return False
    # scale each pair to cell_w
    scaled = [(lab, im.resize((cell_w, int(im.height * cell_w / im.width)))) for lab, im in pairs]
    row_h = max(s.height for _, s in scaled) + label_h
    nrows = (len(scaled) + ncols - 1) // ncols
    W = ncols * cell_w + (ncols + 1) * 8
    H = nrows * row_h + (nrows + 1) * 8
    canvas = Image.new("RGB", (W, H), "white")
    d = ImageDraw.Draw(canvas)
    f = _font(15)
    for idx, (lab, s) in enumerate(scaled):
        r, c = divmod(idx, ncols)
        x = 8 + c * (cell_w + 8)
        y = 8 + r * (row_h + 8)
        d.text((x + 2, y + 4), f"{idx}. {lab}   (teacher | quant)", fill="black", font=f)
        canvas.paste(s, (x, y + label_h))
    canvas.save(out_path)
    return True


print("generating per-cell appendix montages...")
cell_imgs = {}
for r in RANKS:
    for v in VARIANTS:
        op = f"{APP}/cell_r{r}_{v}.png"
        if cell_grid_montage(r, v, op):
            cell_imgs[(r, v)] = op
# mechanism-ablation baselines (rank-0; non-grid dir tags). (dir_tag, title, eval, vel, wrecon-mean, wrecon-max)
BASELINES = [
    ("nosvd_nosmooth", "RTN W4A8 — no smooth (s=1), no SVD", 0.0573, 0.2149, 0.1112, 0.1504),
    ("nosvd", "SmoothQuant W4A8 — rank-0 (α=0.5), no SVD", 0.0729, 0.2424, 0.1356, 0.2633),
]
base_imgs = {}
for tag, *_ in BASELINES:
    op = f"{APP}/cell_r0_{tag}.png"
    if cell_grid_montage(0, tag, op):
        base_imgs[tag] = op
print(f"  wrote {len(cell_imgs)} grid + {len(base_imgs)} baseline montages")

# --------------------------------------------------------------------- assemble HTML
md_body = open("report/QUANT_REPORT.md").read()
# strip the pandoc %-title metadata lines so they don't render as a literal paragraph
md_lines = [ln for ln in md_body.splitlines() if not ln.startswith("% ")]
body_html = markdown.markdown("\n".join(md_lines), extensions=["tables", "fenced_code", "sane_lists"])

FIGS = [
    ("mechanism_ladder.png", "Figure 0. Mechanism ablation — RTN (no smooth, no SVD) → SmoothQuant rank-0 → +SVD plain r16/r32/r64 → grid best. ⚠️ SmoothQuant (α=0.5) HURTS at W4A8: the RTN floor (0.0573) beats it (0.0729) AND the smoothed r16/r32 SVD cells."),
    ("grid_4x3.png", "Figure 1. Full method×rank grid — eval velocity-loss for all 12 cells (300-img calib)."),
    ("deltas_vs_plain.png", "Figure 2. Each upgrade's % change vs plain SVD, by rank. Whitening alone is non-monotonic; refine reliable; combo strongest at low rank."),
    ("cross_track.png", "Figure 3. Cross-track context: W4A8 SVDQuant vs the shelved block-surgery frontier on the same metric."),
    ("byte_compute_budget.png", "Figure 4. Where the bytes/compute go: the 4-bit residual dominates size; the low-rank branch is ~free on compute (FLOPs ∝ r/in)."),
    ("montage_method_r64_text.png", "Figure 5. Method comparison at rank-64 on the in-image-text probe (teacher | quant), all four variants."),
    ("montage_best_allprompts.png", "Figure 6. Grid-best cell (r64 plain+refine, 0.0446) vs teacher across all 8 probes."),
]
figs_html = "<h1>Figures</h1>\n"
for fn, cap in FIGS:
    p = f"{AST}/{fn}"
    if os.path.exists(p):
        figs_html += f'<figure class="fig"><img src="{p}"/><figcaption>{cap}</figcaption></figure>\n'

# appendix: every cell, one per page, with its metrics
app_html = '<h1 class="appendix">Appendix A — Per-cell visual montages (all 12 cells × 8 probes)</h1>\n'
app_html += ("<p>Each cell built one-at-a-time on the same fixed 300-image calibration set. "
             "Left = teacher (bf16, 4-step), right = W4A8 quantized student, same prompt and seed. "
             "Probes 0–3 are the original set; 4–7 were added for richer evaluation.</p>\n")
# baselines first (mechanism ablation)
if base_imgs:
    app_html += '<h2>A.0 — Mechanism-ablation baselines (rank-0)</h2>\n'
    for tag, title, el, vr, wm, wx in BASELINES:
        if tag not in base_imgs:
            continue
        cap = (f"{title} (~3.76× smaller)<br/>eval-loss <b>{el:.4f}</b> · vel-relerr {vr:.4f} · "
               f"wrecon mean {wm:.4f} / max {wx:.4f}")
        app_html += (f'<div class="cell"><h2>{title}</h2><p class="metrics">{cap}</p>'
                     f'<img src="{base_imgs[tag]}"/></div>\n')
for r in RANKS:
    for v in VARIANTS:
        if (r, v) not in cell_imgs:
            continue
        el, vr, wr, orc = M[(r, v)]
        star = "  ★ grid best" if (r, v) == BEST else ""
        orc_s = f" · orecon {orc:.4f}" if orc != "" else ""
        cap = (f"rank {r} · {VLABEL[v]} ({SMALLER[r]} smaller){star}<br/>"
               f"eval-loss <b>{el:.4f}</b> · vel-relerr {vr:.4f} · wrecon {wr:.4f}{orc_s}")
        app_html += (f'<div class="cell"><h2>r{r} — {VLABEL[v]}</h2>'
                     f'<p class="metrics">{cap}</p>'
                     f'<img src="{cell_imgs[(r, v)]}"/></div>\n')

CSS = """
@page { size: A4; margin: 1.6cm 1.4cm; @bottom-center { content: counter(page); font-size: 9px; color:#888; } }
body { font-family: 'DejaVu Sans', sans-serif; font-size: 10.5px; line-height: 1.45; color:#1a1a1a; }
h1 { font-size: 19px; border-bottom: 2px solid #333; padding-bottom:3px; margin-top: 18px; }
h1.appendix, h1.fightitle { page-break-before: always; }
h2 { font-size: 14px; color:#222; margin-top: 14px; }
h3 { font-size: 12px; color:#333; }
table { border-collapse: collapse; width: 100%; font-size: 9px; margin: 8px 0; }
th, td { border: 1px solid #bbb; padding: 3px 5px; text-align: left; }
th { background: #f0f0f0; }
code { background:#f4f4f4; padding:0 2px; font-family:'DejaVu Sans Mono', monospace; font-size:9px; }
pre { background:#f6f6f6; padding:6px; border-radius:3px; font-size:8.5px; overflow-wrap:anywhere; white-space:pre-wrap; }
figure.fig { page-break-inside: avoid; text-align:center; margin: 10px 0; }
figure.fig img { max-width: 100%; border:1px solid #ddd; }
figcaption { font-size: 9px; color:#555; margin-top:3px; }
.cell { page-break-before: always; page-break-inside: avoid; }
.cell img { max-width: 100%; border:1px solid #ddd; }
.metrics { font-size: 10px; color:#333; margin: 2px 0 6px; }
.title { text-align:center; margin-top: 4cm; }
.title h1 { font-size: 26px; border:0; }
.title p { color:#555; font-size: 12px; }
"""

title = """
<div class="title">
  <h1>FLUX.2 [klein] 4B → W4A8 SVDQuant</h1>
  <p><b>Methodology &amp; full 4×3 method×rank ablation grid</b></p>
  <p>Our own fake-quant SVDQuant + activation-aware whitening + iterative refinement</p>
  <p>2026-06-01 · 1× A100-80GB · 300-image calibration · held-out velocity-matching metric</p>
</div>
<div style="page-break-after: always;"></div>
"""

html = (f"<html><head><meta charset='utf-8'><style>{CSS}</style></head><body>"
        f"{title}{body_html}{figs_html}{app_html}</body></html>")
HTML(string=html, base_url=".").write_pdf("report/QUANT_REPORT.pdf")
sz = os.path.getsize("report/QUANT_REPORT.pdf") / 1e6
print(f"wrote report/QUANT_REPORT.pdf ({sz:.1f} MB)")