File size: 7,667 Bytes
2552774 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | #!/usr/bin/env python3
"""Analyze MOSSStripeProbe recordings against final-image banding scores.
comfyui_venv/bin/python workflows/moss/probe_analysis.py output/moss/probe
Expects <setup>_0000N_.png images and probe_<setup>_s<seed>.npz recordings for setups sft, fix, base,
seeds 11-16 (image N = seed 10 + N). Writes probe_*.png figures and probe_summary.json next to them.
"""
import json, os, sys
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import bandviz
D = sys.argv[1]
SETUPS = [("sft", "Full SFT (bands)"), ("fix", "SFT + 21-49 LoRA @ -0.5 (Kijai fix)"), ("base", "Base")]
SEEDS = range(11, 17)
KIND = ["lens", "layer adds (lens)", "residual stream"]
SURFACE, INK, INK2, MUTED, GRID = "#fcfcfb", "#0b0b0b", "#52514e", "#898781", "#e1e0d9"
def rel(p):
"""Spectrum power relative to its median over non-DC bins, in log10."""
return np.log10(p / np.median(p[..., 1:], axis=-1, keepdims=True))
runs = []
for setup, _ in SETUPS:
for seed in SEEDS:
z = np.load(os.path.join(D, f"probe_{setup}_s{seed}.npz"), allow_pickle=True)
cond = np.array([c == "[0]" for c in z["cou"]])
spec = np.stack([np.asarray(s[0], dtype=np.float64) for s in z["spec"]])[cond] # (steps, layers, 3, bins)
img = os.path.join(D, f"{setup}_{seed - 10:05d}_.png")
runs.append({"setup": setup, "seed": seed, "spec": spec, "sigma": z["sigma"][cond], "score": float(bandviz.analyze(img)["score"])})
scores = np.array([r["score"] for r in runs])
R = np.stack([rel(r["spec"]) for r in runs]) # (runs, steps, layers, 3, bins)
nb = R.shape[-1]
W_TOK = (nb - 1) * 2
def corr(x, y):
x = x - x.mean(0); y = y - y.mean()
return (x * y.reshape(-1, *[1] * (x.ndim - 1))).sum(0) / np.sqrt((x ** 2).sum(0) * (y ** 2).sum() + 1e-12)
# 1. Which token frequency carries the banding: final-layer lens (= what the model outputs), late steps.
final = R[:, -8:, -1, 0, :].mean(1) # (runs, bins)
r_bin = corr(final, scores)
sig = int(np.argmax(r_bin[1:])) + 1
print("image scores:", {f"{r['setup']}_s{r['seed']}": round(r["score"], 1) for r in runs})
print("signature bin", sig, f"({sig} cycles across {W_TOK} tokens), r =", round(float(r_bin[sig]), 3))
fig, ax = plt.subplots(1, 2, figsize=(15, 4.6), facecolor=SURFACE)
for a in ax:
a.set_facecolor(SURFACE); a.grid(True, color=GRID); a.set_axisbelow(True)
for s in ("top", "right"): a.spines[s].set_visible(False)
ax[0].bar(range(nb), r_bin, color="#2a78d6")
ax[0].bar([sig], [r_bin[sig]], color=INK)
ax[0].set_title("How well each stripe frequency in the model's output predicts the final image's banding score\n(correlation across all 18 runs; black = strongest)", fontsize=11, loc="left", color=INK)
ax[0].set_xlabel(f"stripe frequency in the token grid (cycles across {W_TOK} tokens)", color=MUTED)
ax[0].set_ylabel("correlation r", color=MUTED)
for (setup, name), c in zip(SETUPS, ["#d03b3b", "#2a78d6", "#898781"]):
ax[1].plot(range(1, nb), final[[r["setup"] == setup for r in runs]].mean(0)[1:], color=c, linewidth=2, label=name)
ax[1].axvline(sig, color=INK, linewidth=1, linestyle=(0, (4, 3)))
ax[1].set_title("Output stripe spectrum, mean over 6 seeds (log power ÷ median)", fontsize=11, loc="left", color=INK)
ax[1].set_xlabel(f"cycles across {W_TOK} tokens", color=MUTED); ax[1].legend(frameon=False)
fig.tight_layout(); fig.savefig(os.path.join(D, "probe_1_signature.png"), dpi=110, facecolor=SURFACE); plt.close(fig)
# 2. Heatmaps: layer x step of signature power, per setup, for lens state and per-layer additions.
steps = R.shape[1]
fig, axes = plt.subplots(2, 3, figsize=(17, 10), facecolor=SURFACE, sharey=True)
for row, kind in enumerate([0, 1]):
vals = [R[[r["setup"] == s for r in runs]][..., kind, sig].mean(0) for s, _ in SETUPS] # (steps, layers)
vmin, vmax = min(v.min() for v in vals), max(v.max() for v in vals)
for col, ((s, name), v) in enumerate(zip(SETUPS, vals)):
a = axes[row, col]
im = a.imshow(v.T, aspect="auto", origin="lower", cmap="magma", vmin=vmin, vmax=vmax, extent=(0.5, steps + 0.5, -0.5, v.shape[1] - 0.5))
a.set_title(f"{name}\n{KIND[kind]}", fontsize=10.5, loc="left", color=INK)
a.set_xlabel("sampling step (noisy → clean)", color=MUTED)
if col == 0: a.set_ylabel("layer", color=MUTED)
fig.colorbar(im, ax=axes[row].tolist(), shrink=0.85, label=f"log stripe power at bin {sig} (÷ median)")
fig.suptitle(f"Where the stripe frequency lives: layer × step, mean of 6 seeds (top: what each layer's output would decode to; bottom: what each layer adds)", x=0.01, ha="left", fontsize=12.5, color=INK)
fig.savefig(os.path.join(D, "probe_2_heatmaps.png"), dpi=110, facecolor=SURFACE); plt.close(fig)
# 3. Difference and correlation maps: SFT minus base, fix minus SFT, and r(cell, image score) across all runs.
sel = lambda s: np.array([r["setup"] == s for r in runs])
add = R[..., 1, sig] # (runs, steps, layers)
maps = [("Full SFT minus Base: extra stripe energy each layer adds", add[sel("sft")].mean(0) - add[sel("base")].mean(0), "RdBu_r", True),
("Kijai fix minus Full SFT: what the fix removes", add[sel("fix")].mean(0) - add[sel("sft")].mean(0), "RdBu_r", True),
("Correlation with final banding score (all 18 runs)", corr(add, scores), "RdBu_r", True)]
fig, axes = plt.subplots(1, 3, figsize=(17, 6), facecolor=SURFACE, sharey=True)
for a, (title, v, cmap, sym) in zip(axes, maps):
lim = np.nanmax(np.abs(v))
im = a.imshow(v.T, aspect="auto", origin="lower", cmap=cmap, vmin=-lim, vmax=lim, extent=(0.5, steps + 0.5, -0.5, v.shape[1] - 0.5))
a.set_title(title, fontsize=10.5, loc="left", color=INK); a.set_xlabel("sampling step (noisy → clean)", color=MUTED)
fig.colorbar(im, ax=a, shrink=0.85)
axes[0].set_ylabel("layer", color=MUTED)
fig.suptitle(f"The dials: per-layer stripe additions at bin {sig}", x=0.01, ha="left", fontsize=12.5, color=INK)
fig.tight_layout(); fig.savefig(os.path.join(D, "probe_3_dials.png"), dpi=110, facecolor=SURFACE); plt.close(fig)
# 4. Per-layer profile: correlation with score, averaged over steps, plus per-setup addition profiles.
r_layer = corr(add.mean(1), scores) # (layers,)
fig, ax = plt.subplots(1, 2, figsize=(15, 4.6), facecolor=SURFACE)
for a in ax:
a.set_facecolor(SURFACE); a.grid(True, color=GRID); a.set_axisbelow(True)
for s in ("top", "right"): a.spines[s].set_visible(False)
ax[0].bar(range(len(r_layer)), r_layer, color="#2a78d6")
ax[0].set_title("Per layer: how well its stripe addition (all steps) predicts the final banding score", fontsize=11, loc="left", color=INK)
ax[0].set_xlabel("layer", color=MUTED); ax[0].set_ylabel("correlation r", color=MUTED)
for (s, name), c in zip(SETUPS, ["#d03b3b", "#2a78d6", "#898781"]):
ax[1].plot(add[sel(s)].mean((0, 1)), color=c, linewidth=2, label=name)
ax[1].set_title(f"Mean stripe addition per layer at bin {sig} (log ÷ median)", fontsize=11, loc="left", color=INK)
ax[1].set_xlabel("layer", color=MUTED); ax[1].legend(frameon=False)
fig.tight_layout(); fig.savefig(os.path.join(D, "probe_4_layers.png"), dpi=110, facecolor=SURFACE); plt.close(fig)
top = np.argsort(-r_layer)[:8]
json.dump({"signature_bin": sig, "r_bin": r_bin.round(3).tolist(), "r_layer": r_layer.round(3).tolist(),
"top_layers": [(int(i), round(float(r_layer[i]), 3)) for i in top],
"scores": {f"{r['setup']}_s{r['seed']}": round(r["score"], 1) for r in runs}},
open(os.path.join(D, "probe_summary.json"), "w"), indent=1)
print("top layers by r:", [(int(i), round(float(r_layer[i]), 2)) for i in top])
|