Sync explorer_app.py and clip_utils.py from main repo

scripts/explorer_app.py

@@ -516,6 +516,43 @@ def render_zoomed_patch(img_idx, heatmap_16x16, size=THUMB, pg=None):
     return img.crop((x0, y0, x1, y1)).resize((size, size), Image.BILINEAR)
 
 
+def _load_image_from_ds(ds, img_i):
+    """Like load_image() but uses the given dataset's image_paths."""
+    path = ds['image_paths'][img_i]
+    fname = os.path.basename(path)
+    for base in [args.image_dir] + ([args.extra_image_dir] if args.extra_image_dir else []):
+        candidate = os.path.join(base, fname)
+        if os.path.exists(candidate):
+            return Image.open(candidate).convert("RGB")
+    return Image.open(path).convert("RGB")
+
+
+def _render_overlay_from_ds(ds, feat, slot, size=THUMB, alpha=None):
+    """Return (PIL overlay image, caption) for ds/feat/slot, or None on failure."""
+    if alpha is None:
+        alpha = heatmap_alpha_slider.value
+    try:
+        img_i = int(ds['top_img_idx'][feat, slot].item())
+        if img_i < 0:
+            return None
+        plain = _load_image_from_ds(ds, img_i).resize((size, size), Image.BILINEAR)
+        hm_tensor = ds.get('top_heatmaps')
+        if hm_tensor is not None:
+            pg = ds.get('heatmap_patch_grid', 16)
+            hmap = hm_tensor[feat, slot].float().numpy().reshape(pg, pg)
+            img_arr = np.array(plain).astype(np.float32) / 255.0
+            hmap_up = cv2.resize(hmap, (size, size), interpolation=cv2.INTER_CUBIC)
+            hmax = hmap_up.max()
+            hmap_norm = hmap_up / hmax if hmax > 0 else hmap_up
+            overlay = ALPHA_JET(hmap_norm)
+            ov_a = overlay[:, :, 3:4] * alpha
+            blended = np.clip((img_arr * (1 - ov_a) + overlay[:, :, :3] * ov_a) * 255, 0, 255).astype(np.uint8)
+            return Image.fromarray(blended), f"img {img_i}"
+        return plain, f"img {img_i}"
+    except Exception:
+        return None
+
+
 def pil_to_data_url(img):
     buf = io.BytesIO()
     img.save(buf, format="JPEG", quality=85)
@@ -612,6 +649,48 @@ def make_compare_aggregations_html(top_infos, mean_infos, p75_infos, feat, n_eac
     return html
 
 
+def make_cross_sae_comparison_html(ds_a, feat_a, ds_b, feat_b, n=4, size=160):
+    """
+    Two side-by-side 2×2 grids: left = SAE A / feat_a, right = SAE B / feat_b.
+    """
+    def _collect(ds, feat):
+        items = []
+        for slot in range(min(n, ds['top_img_idx'].shape[1])):
+            result = _render_overlay_from_ds(ds, feat, slot, size=size)
+            if result:
+                items.append(result)
+            if len(items) == n:
+                break
+        return items
+
+    items_a = _collect(ds_a, feat_a)
+    items_b = _collect(ds_b, feat_b)
+
+    def _grid_html(items, label, color):
+        header = (f'<div style="background:{color};color:#fff;font-size:11px;font-weight:bold;'
+                  f'text-align:center;padding:4px;border-radius:4px;margin-bottom:6px">{label}</div>')
+        grid = '<div style="display:grid;grid-template-columns:repeat(2,{s}px);gap:4px">'.format(s=size)
+        for img, cap in items:
+            url = pil_to_data_url(img)
+            grid += (f'<div style="text-align:center">'
+                     f'<img src="{url}" width="{size}" height="{size}"'
+                     f' style="border:1px solid #ccc;border-radius:3px;display:block"/>'
+                     f'<div style="font-size:9px;color:#555;margin-top:2px">{cap}</div></div>')
+        grid += '</div>'
+        return header + grid
+
+    label_a = f"{ds_a['label']} — feat {feat_a}"
+    label_b = f"{ds_b['label']} — feat {feat_b}"
+    col_a = _grid_html(items_a, label_a, "#2563a8")
+    col_b = _grid_html(items_b, label_b, "#b85c00")
+
+    return (
+        '<div style="display:flex;gap:16px;padding:8px;background:#fafafa;'
+        'border:1px solid #ddd;border-radius:6px">'
+        + col_a + col_b + '</div>'
+    )
+
+
 # ---------- UMAP data source ----------
 # live_mask / live_indices / freq / mean_act / log_freq / umap_backup are all
 # already set by _apply_dataset_globals(0) above — just build the source from them.
@@ -1795,7 +1874,40 @@ middle_panel = column(
     p75_heatmap_div, p75_zoom_div,
 )
 
-right_panel = column(summary_section, patch_section, clip_section)
+# --- Cross-SAE comparison section ---
+cmp_ds_a = Select(title="SAE A:", value="0",
+    options=[(str(i), ds['label']) for i, ds in enumerate(_all_datasets)])
+cmp_feat_a = TextInput(title="Feature (SAE A):", value="0", width=100)
+cmp_ds_b = Select(title="SAE B:", value=str(min(1, len(_all_datasets)-1)),
+    options=[(str(i), ds['label']) for i, ds in enumerate(_all_datasets)])
+cmp_feat_b = TextInput(title="Feature (SAE B):", value="0", width=100)
+cmp_btn = Button(label="Generate Comparison", button_type="primary", width=200)
+cmp_output_div = Div(text="", width=400)
+
+def _on_cmp_generate():
+    try:
+        idx_a = int(cmp_ds_a.value)
+        idx_b = int(cmp_ds_b.value)
+        fa = int(cmp_feat_a.value)
+        fb = int(cmp_feat_b.value)
+        _ensure_loaded(idx_a)
+        _ensure_loaded(idx_b)
+        ds_a = _all_datasets[idx_a]
+        ds_b = _all_datasets[idx_b]
+        cmp_output_div.text = make_cross_sae_comparison_html(ds_a, fa, ds_b, fb)
+    except Exception as e:
+        cmp_output_div.text = f'<p style="color:red">Error: {e}</p>'
+
+cmp_btn.on_click(lambda: _on_cmp_generate())
+
+cmp_section = _make_collapsible("Cross-SAE Comparison", column(
+    row(cmp_ds_a, cmp_feat_a),
+    row(cmp_ds_b, cmp_feat_b),
+    cmp_btn,
+    cmp_output_div,
+))
+
+right_panel = column(summary_section, patch_section, clip_section, cmp_section)
 
 layout = row(left_panel, middle_panel, right_panel)
 curdoc().add_root(layout)