Update app.py

app.py

@@ -5,7 +5,7 @@
 #  - Patch grid (16x16)
 #  - Patch attention (per layer / per head / query token)
 #  - Attention rollout (layer aggregated)
-#  - PCA of patch embeddings across selected layers
+#  - PCA of patch embeddings across layers
 #  - Top-5 predictions & simple/technical explanations
 # ==========================================================
 
@@ -16,11 +16,15 @@ from typing import Any, Dict, List, Optional, Tuple
 import gradio as gr
 import numpy as np
 import torch
-from PIL import Image, ImageDraw, ImageFont
-from transformers import AutoImageProcessor, ViTModel, ViTForImageClassification
+from PIL import Image, ImageDraw
+from transformers import (
+    AutoImageProcessor,
+    ViTModel,
+    ViTForImageClassification,
+    AutoConfig,
+)
 from sklearn.decomposition import PCA
 import plotly.express as px
-import plotly.graph_objects as go
 
 warnings.filterwarnings("ignore")
 
@@ -33,24 +37,40 @@ VIT_CLF = None         # ViTForImageClassification (classification head)
 PROCESSOR = None
 
 
-# ------------------ model loader with SDPA fix ------------------
+# ------------------ model loader with SDPA -> eager fix ------------------
 def load_models():
+    """
+    Load processor + ViT base + classification head.
+    Important: create config first, set attn_implementation='eager'
+    before enabling output_attentions/output_hidden_states, then load models.
+    """
     global VIT_BASE, VIT_CLF, PROCESSOR
     if VIT_BASE is not None and VIT_CLF is not None and PROCESSOR is not None:
         return VIT_BASE, VIT_CLF, PROCESSOR
 
     PROCESSOR = AutoImageProcessor.from_pretrained(MODEL_NAME)
 
-    # base ViT (encoder) - we need hidden_states & attentions
-    base = ViTModel.from_pretrained(MODEL_NAME, output_hidden_states=True)
-    # fix attn backend so we can access attentions
-    base.config.attn_implementation = "eager"
-    base.config.output_attentions = True
-    base.config.output_hidden_states = True
+    # load config, modify before creating model
+    cfg = AutoConfig = None
+    try:
+        cfg = AutoConfig.from_pretrained(MODEL_NAME)
+    except Exception:
+        # fallback: load a default config and set minimal fields
+        from transformers import ViTConfig
+        cfg = ViTConfig.from_pretrained(MODEL_NAME)
+
+    # FORCE eager attention backend so we can extract attentions
+    # (must set attn_implementation before enabling output_attentions)
+    cfg.attn_implementation = "eager"
+    cfg.output_attentions = True
+    cfg.output_hidden_states = True
+
+    # now load the base encoder with the modified config
+    base = ViTModel.from_pretrained(MODEL_NAME, config=cfg)
     base.to(DEVICE)
     base.eval()
 
-    # classifier head for top-k labels
+    # load classifier separately (we can use default config for classifier)
     clf = ViTForImageClassification.from_pretrained(MODEL_NAME)
     clf.to(DEVICE)
     clf.eval()
@@ -74,8 +94,7 @@ def make_patch_grid_image(pil: Image.Image, patch_size: int = 16, target_size: i
 
 def make_attention_overlay(base_img: Image.Image, heat_grid: np.ndarray, cmap_alpha: float = 0.45) -> Image.Image:
     """
-    heat_grid: (G, G) values in any scale (we will normalize)
-    overlay on base_img (resized to 224x224)
+    heat_grid: (G, G) values (any scale) -> normalized then overlaid on base_img (resized to 224x224)
     """
     img = base_img.convert("RGB").resize((224, 224))
     g = np.array(heat_grid, dtype=np.float32)
@@ -87,11 +106,11 @@ def make_attention_overlay(base_img: Image.Image, heat_grid: np.ndarray, cmap_al
     else:
         g = np.zeros_like(g, dtype=np.float32)
 
-    # upsample
+    # upsample to image resolution
     heat_img = Image.fromarray((g * 255).astype("uint8"), mode="L").resize((224, 224), Image.BILINEAR)
     heat = np.array(heat_img).astype(np.float32) / 255.0
 
-    # simple colormap blue->red
+    # simple blue->red colormap
     r = heat
     gch = np.zeros_like(heat)
     b = 1.0 - heat
@@ -111,7 +130,6 @@ def compute_attention_rollout(all_attentions: List[torch.Tensor]) -> np.ndarray:
       R = prod_l (A_l_hat) where A_l_hat = A_l + I; rows normalized
     Returns rollout matrix (seq, seq)
     """
-    # convert to np arrays averaged over heads
     avg_mats = []
     for a in all_attentions:
         # a: (batch=1, heads, seq, seq)
@@ -119,14 +137,16 @@ def compute_attention_rollout(all_attentions: List[torch.Tensor]) -> np.ndarray:
         avg_mats.append(mat)
 
     seq = avg_mats[0].shape[0]
-    # add identity & normalize rows
     aug = []
     for A in avg_mats:
         A_hat = A + np.eye(seq)
-        A_hat = A_hat / A_hat.sum(axis=-1, keepdims=True)
+        # normalize rows (sum over last dim)
+        row_sums = A_hat.sum(axis=-1, keepdims=True)
+        # avoid division by zero
+        row_sums[row_sums == 0] = 1.0
+        A_hat = A_hat / row_sums
         aug.append(A_hat)
 
-    # multiply (matrix product) in order
     R = aug[0]
     for A in aug[1:]:
         R = A @ R
@@ -146,13 +166,11 @@ def layers_pca_plot(hidden_states: List[torch.Tensor], layers: List[int]) -> Any
     for li in layers:
         hs = hidden_states[li][0].detach().cpu().numpy()  # (seq, hidden)
         patches = hs[1:, :]  # remove CLS -> (N_patches, hidden)
-        # PCA to 2D
         pca = PCA(n_components=2)
         pts = pca.fit_transform(patches)
         pts_all.append(pts)
         layer_labels.append(np.array([li] * pts.shape[0]))
 
-    # combine
     coords = np.vstack(pts_all)
     labels = np.concatenate(layer_labels)
     df = {"x": coords[:, 0], "y": coords[:, 1], "layer": labels.astype(str)}
@@ -165,9 +183,7 @@ def layers_pca_plot(hidden_states: List[torch.Tensor], layers: List[int]) -> Any
 # ------------------ core analyzer ------------------
 def analyze_vit_full(img: Optional[Image.Image], simple: bool):
     if img is None:
-        return (
-            None, None, None, None, None, "", {}, {}
-        )
+        return (None, None, None, None, None, "", {})
 
     base, clf, processor = load_models()
 
@@ -180,43 +196,46 @@ def analyze_vit_full(img: Optional[Image.Image], simple: bool):
         outputs = base(**inputs)
 
     # outputs.attentions: list L tensors (batch=1, heads, seq, seq)
-    attentions = outputs.attentions  # list length L
-    hidden_states = outputs.hidden_states  # list length L+1 (including embeddings) usually
+    attentions = outputs.attentions
+    hidden_states = outputs.hidden_states
 
     L = len(attentions)
     seq_len = attentions[0].shape[-1]
     n_patches = seq_len - 1
     grid_size = int(math.sqrt(n_patches))
     if grid_size * grid_size != n_patches:
-        # fallback: compute closest integer grid
         grid_size = int(round(math.sqrt(n_patches)))
 
-    # default selections
-    default_layer = L - 1
-    default_head = 0
-    # default query token = 0 (CLS)
-    default_query = 0
-
     # Build patch grid image
     patch_grid = make_patch_grid_image(img.copy(), patch_size=16, target_size=224)
 
-    # Build per-layer per-head CLS->patch default overlay
-    # pick last layer, head 0, CLS query
-    att_np = attentions[default_layer][0].cpu().numpy()  # (heads, seq, seq)
-    cls_to_patches = att_np[default_head, 0, 1:]  # (n_patches,)
+    # default overlay: last layer, head 0, CLS query
+    last_layer = L - 1
+    head0 = 0
+    # attentions[last_layer]: shape (batch=1, heads, seq, seq)
+    att_np = attentions[last_layer][0].cpu().numpy()  # (heads, seq, seq)
+    cls_to_patches = att_np[head0, 0, 1:]  # (n_patches,)
+    if cls_to_patches.shape[0] != grid_size * grid_size:
+        tmp = np.zeros(grid_size * grid_size, dtype=np.float32)
+        nmin = min(cls_to_patches.shape[0], tmp.shape[0])
+        tmp[:nmin] = cls_to_patches[:nmin]
+        cls_to_patches = tmp
     cls_grid = cls_to_patches.reshape(grid_size, grid_size)
     attn_overlay = make_attention_overlay(img, cls_grid)
 
-    # Compute rollout
+    # Compute rollout overlay (CLS)
     rollout_mat = compute_attention_rollout(attentions)  # (seq, seq)
     rollout_cls = rollout_mat[0, 1:]
+    if rollout_cls.shape[0] != grid_size * grid_size:
+        tmp = np.zeros(grid_size * grid_size, dtype=np.float32)
+        nmin = min(rollout_cls.shape[0], tmp.shape[0])
+        tmp[:nmin] = rollout_cls[:nmin]
+        rollout_cls = tmp
     rollout_grid = rollout_cls.reshape(grid_size, grid_size)
-    rollout_overlay = make_attention_overlay(img, rollout_grid, cmap_alpha=0.5)
+    rollout_overlay = make_attention_overlay(img, rollout_grid, cmap_alpha=0.55)
 
-    # PCA multi-layer: pick a few representative layers (start, quarter, half, three-quarters, last)
-    layers_to_show = sorted(
-        list({0, max(0, L // 4), max(0, L // 2), max(0, 3 * L // 4), L - 1})
-    )
+    # PCA multi-layer: choose representative layers
+    layers_to_show = sorted(list({0, max(0, L // 4), max(0, L // 2), max(0, 3 * L // 4), L - 1}))
     pca_fig = layers_pca_plot(hidden_states, layers_to_show)
 
     # Classification top-5
@@ -248,34 +267,20 @@ def analyze_vit_full(img: Optional[Image.Image], simple: bool):
 - Attention rollout uses Abnar & Zuidema's method to accumulate attention paths across layers.
 """
 
-    # return many things + state necessary for interactive updates (layer/head/query)
     state = {
-        "attentions": [a.cpu() for a in attentions],  # store on CPU to allow slider updates
+        "attentions": [a.cpu() for a in attentions],  # move to CPU for interactive updates
         "hidden_states": [h.cpu() for h in hidden_states],
         "grid_size": grid_size,
         "num_layers": L,
         "num_heads": attentions[0].shape[1],
-        "base_image": img,  # original high-res image (we'll resize to 224 when overlaying)
+        "base_image": img,
     }
 
-    return (
-        patch_grid,
-        attn_overlay,
-        rollout_overlay,
-        pca_fig,
-        preds_text,
-        explain_md,
-        state,
-    )
+    return patch_grid, attn_overlay, rollout_overlay, pca_fig, preds_text, explain_md, state
 
 
 # ------------------ update functions for sliders / choices ------------------
 def update_layer_head_query(state: Dict[str, Any], layer_idx: int, head_idx: int, query_token: int, mode: str):
-    """
-    mode:
-      - "patch_attention": attention of query_token -> patches at (layer, head)
-      - "rollout": ignored (we will return rollout overlay)
-    """
     if not state:
         return None
 
@@ -288,25 +293,17 @@ def update_layer_head_query(state: Dict[str, Any], layer_idx: int, head_idx: int
     h = max(0, min(int(head_idx), H - 1))
     q = max(0, min(int(query_token), grid * grid))  # q in 0..n_patches (0==CLS)
 
-    # load attention for layer l: it's a CPU tensor (heads, seq, seq) already stored as state
-    att_tensor = state["attentions"][l]  # shape (heads, seq, seq) because we saved a[0] earlier
-    # ensure shape (heads, seq, seq)
-    if att_tensor.ndim == 4:  # sometimes shape might be (1, heads, seq, seq)
+    att_tensor = state["attentions"][l]  # shape (heads, seq, seq) or (1,heads,seq,seq)
+    if att_tensor.ndim == 4:
         att_tensor = att_tensor[0]
     att_np = att_tensor.numpy()  # (heads, seq, seq)
 
-    # query q -> keys: if q == 0 it's CLS; keys positions 1..seq-1 are patches
     seq = att_np.shape[-1]
-    n_patches = seq - 1
-    # column indices for keys: 1..seq-1 map to patches 0..n_patches-1
     if q >= seq:
         q = 0
 
-    # get attention vector for head h: att[h, q, 1:]
     vec = att_np[h, q, 1:]
-    # if vec shorter/longer than grid^2, adjust
     if vec.shape[0] != grid * grid:
-        # pad or trim
         tmp = np.zeros(grid * grid, dtype=np.float32)
         nmin = min(vec.shape[0], tmp.shape[0])
         tmp[:nmin] = vec[:nmin]
@@ -321,8 +318,6 @@ def get_rollout_overlay(state: Dict[str, Any]):
     if not state:
         return None
     attentions = state["attentions"]
-    # attentions list of tensors (heads, seq, seq)
-    # convert to list of (1, heads, seq, seq) for compute_attention_rollout
     mats = [a.unsqueeze(0) if a.ndim == 3 else a for a in attentions]
     R = compute_attention_rollout(mats)  # (seq, seq)
     grid = state["grid_size"]
@@ -339,9 +334,7 @@ def get_rollout_overlay(state: Dict[str, Any]):
 def update_pca_layers(state: Dict[str, Any], selected_layers: List[int]):
     if not state:
         return None
-    # hidden_states stored as list of CPU tensors (batch, seq, hidden)
     hs = state["hidden_states"]
-    # ensure layers within range
     layers = [max(0, min(int(l), len(hs) - 1)) for l in selected_layers]
     fig = layers_pca_plot(hs, layers)
     return fig
@@ -365,8 +358,7 @@ with gr.Blocks(title="ViT Full Interpretability (A+B+C)") as demo:
 
             gr.Markdown("**Attention Rollout & PCA**")
             rollout_btn = gr.Button("Refresh Rollout Overlay")
-            # PCA layers selection: simple multi-select text entry allowed (comma separated)
-            pca_layers_txt = gr.Textbox(label="PCA layers (comma separated indices, e.g. 0,3,6,11)", value="0,3,6,11,11")
+            pca_layers_txt = gr.Textbox(label="PCA layers (comma separated indices, e.g. 0,3,6,11)", value="0,3,6,11")
 
         with gr.Column(scale=1):
             gr.Markdown("### Outputs")