Visual steerability dashboard (CPU PaliGemma + L10 SAE)

.gitattributes

@@ -1,35 +1,4 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+examples/apple1.jpg filter=lfs diff=lfs merge=lfs -text
+examples/apple2.jpg filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,13 +1,18 @@
 ---
 title: Visual Steerability
-emoji: 🏆
-colorFrom: green
-colorTo: gray
+emoji: 🎛️
+colorFrom: indigo
+colorTo: red
 sdk: gradio
-sdk_version: 6.19.0
-python_version: '3.13'
+sdk_version: 5.9.1
 app_file: app.py
 pinned: false
+short_description: Steer PaliGemma's vision model from the inside (SAE)
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# 🎛️ Visual Steerability
+
+Pass an image → see its **L10 SAE neuron activations** (PaliGemma's SigLIP-tower internal
+representation) → **clamp a neuron's latent** → watch the caption change.
+
+Companion demo to the explainer **[The Steerability Spectrum of Visual Attributes](https://sumityadav.com.np/study/steerability-spectrum/)**.

app.py

@@ -0,0 +1,178 @@
+"""Visual Steerability — steer PaliGemma's vision model from the inside (ZeroGPU Space).
+
+Pass an image -> see its L10 SAE neuron activations (the model's internal representation) ->
+clamp a neuron's latent -> watch PaliGemma's caption change. Companion to the explainer at
+https://sumityadav.com.np/study/steerability-spectrum/
+"""
+import os, json
+import numpy as np
+import torch, torch.nn as nn
+import plotly.graph_objects as go
+import gradio as gr
+from PIL import Image as PILImage
+from transformers import AutoProcessor, PaliGemmaForConditionalGeneration
+
+try:
+    import spaces  # ZeroGPU
+    GPU = spaces.GPU
+except Exception:                                    # local / non-ZeroGPU fallback (no-op decorator)
+    def GPU(*a, **k):
+        if a and callable(a[0]):                     # used bare: @GPU
+            return a[0]
+        def deco(f):                                 # used with args: @GPU(duration=...)
+            return f
+        return deco
+
+MODEL = "google/paligemma-3b-mix-224"
+SAE_PT = "sae_l10_coco.pt"
+DASH = "dash_data.npz"
+HF_TOKEN = os.environ.get("HF_TOKEN")
+
+
+class TopKSAE(nn.Module):
+    def __init__(self, d, m, k):
+        super().__init__(); self.k = k
+        self.b_pre = nn.Parameter(torch.zeros(d)); self.W_enc = nn.Parameter(torch.zeros(d, m))
+        self.b_enc = nn.Parameter(torch.zeros(m)); self.W_dec = nn.Parameter(torch.zeros(m, d))
+
+    def encode(self, x):
+        pre = (x - self.b_pre) @ self.W_enc + self.b_enc
+        v, i = pre.topk(self.k, dim=-1)
+        a = torch.zeros_like(pre); a.scatter_(-1, i, torch.relu(v)); return a
+
+
+# ---- load (CPU at import; moved to GPU inside @GPU fns on ZeroGPU) ----
+print("loading SAE + map …")
+ckpt = torch.load(SAE_PT, map_location="cpu", weights_only=False)
+D, M, K, L = ckpt["d"], ckpt["m"], ckpt["k"], ckpt["l"]
+SN = float(ckpt["scale_norm"])
+sae = TopKSAE(D, M, K); sae.load_state_dict(ckpt["state_dict"]); sae.eval()
+mu = torch.tensor(ckpt["mu"], dtype=torch.float32)
+_d = np.load(DASH, allow_pickle=True)
+COORDS = _d["coords"]; LAT = _d["latents"].astype(np.float32)
+
+print("loading PaliGemma … (first GPU call moves it to cuda)")
+model = PaliGemmaForConditionalGeneration.from_pretrained(MODEL, torch_dtype=torch.bfloat16, token=HF_TOKEN).eval()
+proc = AutoProcessor.from_pretrained(MODEL, token=HF_TOKEN)
+tok = proc.tokenizer
+vtower = lambda: model.model.vision_tower
+block = model.model.vision_tower.encoder.layers[L - 1]
+_state = {"on": False, "feat": 0, "strength": 0.0, "moved": False}
+
+
+def _ensure(dev):
+    if not _state["moved"] and dev == "cuda":
+        model.to("cuda"); sae.to("cuda")
+        globals()["mu"] = mu.to("cuda"); _state["moved"] = True
+
+
+def _hook(mod, inp, out):
+    if not _state["on"]:
+        return out
+    h = out[0] if isinstance(out, tuple) else out
+    hf = h.float(); hn = (hf - mu) * SN
+    b, P, d = hn.shape
+    acts = sae.encode(hn.reshape(b * P, d))
+    cur = acts[:, _state["feat"]]
+    dn = (_state["strength"] - cur).unsqueeze(-1) * sae.W_dec[_state["feat"]]
+    h2 = (hf + (dn / SN).reshape(b, P, d)).to(h.dtype)
+    return (h2,) + tuple(out[1:]) if isinstance(out, tuple) else h2
+block.register_forward_hook(_hook)
+
+
+@torch.no_grad()
+def _latent(pil, dev):
+    pv = proc.image_processor([pil], return_tensors="pt")["pixel_values"].to(dev, torch.bfloat16)
+    hs = vtower()(pixel_values=pv, output_hidden_states=True).hidden_states[L].float()
+    acts = sae.encode(((hs[0] - mu) * SN))
+    lat = acts.mean(0).detach().cpu().numpy()
+    return lat
+
+
+@torch.no_grad()
+def _caption(pil, dev, feat=0, strength=0.0, steer=False):
+    _state["on"] = steer; _state["feat"] = int(feat); _state["strength"] = float(strength)
+    inp = proc(text="describe the object in the image", images=pil, return_tensors="pt").to(dev, torch.bfloat16)
+    n = inp["input_ids"].shape[1]
+    g = model.generate(**inp, max_new_tokens=28, do_sample=False)
+    _state["on"] = False
+    return tok.decode(g[0][n:], skip_special_tokens=True).strip()
+
+
+def scatter(coord, feat):
+    feat = int(feat); c = LAT[:, feat]
+    fig = go.Figure()
+    fig.add_trace(go.Scatter(x=COORDS[:, 0], y=COORDS[:, 1], mode="markers",
+        marker=dict(size=7, color=c, colorscale="Viridis", showscale=True, colorbar=dict(title="#%d" % feat), opacity=.55),
+        hovertemplate="neuron #%d = %%{marker.color:.2f}<extra></extra>" % feat, name="500 ref"))
+    if coord is not None:
+        fig.add_trace(go.Scatter(x=[float(coord[0])], y=[float(coord[1])], mode="markers",
+            marker=dict(size=22, color="red", symbol="star", line=dict(width=1.4, color="black")),
+            name="your image", hovertemplate="your image<extra></extra>"))
+    fig.update_layout(title="L10 SAE latent map — your image (★), colored by neuron #%d" % feat,
+        xaxis_title="PC1", yaxis_title="PC2", height=420, margin=dict(l=8, r=8, t=44, b=8),
+        legend=dict(orientation="h", y=1.04, x=0))
+    return fig
+
+
+def neuron_bar(lat):
+    idx = np.argsort(-lat)[:10][::-1]
+    fig = go.Figure(go.Bar(x=lat[idx], y=["#%d" % i for i in idx], orientation="h",
+        marker_color="#4f46e5", hovertemplate="neuron #%{y} = %{x:.2f}<extra></extra>"))
+    fig.update_layout(title="this image's strongest L10 SAE neurons", height=300,
+        xaxis_title="activation", margin=dict(l=8, r=8, t=44, b=8), yaxis=dict(type="category"))
+    return fig
+
+
+@GPU(duration=70)
+def analyze(pil):
+    if pil is None:
+        return "Upload or pick an image.", None, None, gr.update(choices=[], value=None), None
+    dev = "cuda" if torch.cuda.is_available() else "cpu"; _ensure(dev)
+    lat = _latent(pil, dev)
+    base = _caption(pil, dev, steer=False)
+    top = np.argsort(-lat)[:8]
+    choices = [("#%d  (act %.1f)" % (int(t), lat[t]), int(t)) for t in top]
+    coord = (lat - _d["pca_mean"]) @ _d["pca_comp"].T
+    return base, scatter(coord, int(top[0])), neuron_bar(lat), gr.update(choices=choices, value=int(top[0])), coord.tolist()
+
+
+@GPU(duration=70)
+def steer(pil, neuron, strength):
+    if pil is None or neuron is None:
+        return "Run an image first."
+    dev = "cuda" if torch.cuda.is_available() else "cpu"; _ensure(dev)
+    if float(strength) == 0:
+        return "Move the slider to clamp this neuron."
+    return _caption(pil, dev, feat=int(neuron), strength=float(strength), steer=True)
+
+
+def recolor(coord, neuron):
+    if coord is None or neuron is None:
+        return None
+    return scatter(np.array(coord), int(neuron))
+
+
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="indigo"), title="Visual Steerability") as demo:
+    gr.Markdown("# 🎛️ Steer PaliGemma's vision model from the inside\n"
+        "Pass an image → see its **L10 SAE neuron activations** (the model's internal representation) → "
+        "**clamp a neuron** and watch the caption change. "
+        "Read the explainer: [The Steerability Spectrum](https://sumityadav.com.np/study/steerability-spectrum/).")
+    st_coord = gr.State()
+    with gr.Row():
+        with gr.Column(scale=4):
+            img = gr.Image(type="pil", label="image", height=300)
+            run = gr.Button("① Caption + read internal representation", variant="primary")
+            base = gr.Textbox(label="caption", lines=2, interactive=False)
+            neuron = gr.Radio(choices=[], label="② pick a neuron to steer", interactive=True)
+            strength = gr.Slider(-40, 40, value=0, step=1, label="③ clamp strength → re-caption on release")
+            out = gr.Textbox(label="steered caption", lines=2, interactive=False)
+        with gr.Column(scale=5):
+            plot = gr.Plot(label="SAE latent map")
+            bar = gr.Plot(label="internal representation")
+    gr.Examples(examples=[["examples/" + f] for f in sorted(os.listdir("examples"))], inputs=[img], label="examples")
+    run.click(analyze, [img], [base, plot, bar, neuron, st_coord])
+    neuron.change(recolor, [st_coord, neuron], [plot])
+    strength.release(steer, [img, neuron, strength], [out])
+
+demo.queue(max_size=12).launch()

dash_data.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f0b8c917c09dc9910c3fc78f31858787152a62149b0415e8fd19f4a1aa7de92e
+size 3147174

requirements.txt

@@ -0,0 +1,8 @@
+torch
+transformers
+accelerate
+sentencepiece
+pillow
+numpy
+plotly
+huggingface_hub

sae_l10_coco.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d2634de709e576ef1e2c75ab66d418096cad61875d91bf0a83ed488704254203
+size 84985487