Quick test.

README.md

@@ -12,3 +12,13 @@ short_description: Visualizes cosine-similarity of CLS-token with image patches
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+
+# Dynamic ViT Visualizer
+Upload a Hugging Face ViT model and an image.  
+Visualize the cosine similarity between the CLS token and image patches.
+
+**How to use:**
+1. Enter the model repo ID. 
+2. Upload an image.
+3. Get visualization of cosine-similarity between CLS-token and tokens/patches of the uploaded image. 

app.py

@@ -1,4 +1,5 @@
-# app.py
+import matplotlib.pyplot as plt
+import numpy as np
 import gradio as gr
 from transformers import AutoModel, AutoImageProcessor
 from PIL import Image
@@ -11,6 +12,26 @@ state = {
     "repo_id": None,
 }
 
+
+def plot_similarity_heatmap(sim_array: np.ndarray):
+    """
+    sim_array: 2D numpy array of shape (h, w)
+    Returns a PIL image that can be displayed in Gradio
+    """
+    fig, ax = plt.subplots(figsize=(5, 5))
+    cax = ax.imshow(sim_array, cmap='viridis')    
+    ax.set_xticks([])
+    ax.set_yticks([])
+    
+    fig.colorbar(cax)
+    fig.canvas.draw()
+    img = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
+    img = img.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+    
+    plt.close(fig)  
+    return img
+
+
 def load_model(repo_id: str, revision: str = None):
     """
     Load a Hugging Face model and processor from a repo ID.
@@ -50,6 +71,12 @@ with gr.Blocks() as demo:
     image_input = gr.Image(type="pil", label="Upload Image")
     image_output = gr.Image(label="Displayed Image")
 
+    # cos-sim visualization: 
+    # sim_array is your (h, w) numpy array
+    sim_array = np.random.normal((128, 128))
+    heatmap_img = plot_similarity_heatmap(sim_array)
+    gr.Image(value=heatmap_img, label="Cosine Similarity Heatmap")
+
     # Button clicks / image upload handlers
     load_btn.click(fn=load_model, inputs=[repo_input, revision_input], outputs=load_status)
     image_input.change(fn=display_image, inputs=image_input, outputs=image_output)