better viz.

app.py

@@ -7,6 +7,7 @@ from timm import create_model
 from timm.models.layers import PatchEmbed
 from torchvision.models.feature_extraction import create_feature_extractor
 from torchvision.transforms import functional as F
+import glob
 
 cait_model = create_model("cait_xxs24_224.fb_dist_in1k", pretrained=True).eval()
 transform = timm.data.create_transform(
@@ -16,7 +17,7 @@ transform = timm.data.create_transform(
 patch_size = 16
 
 
-def create_attn_extractor(model, block_id=0):
+def create_attn_extractor(block_id=0):
     """Creates a model that produces the softmax attention scores.
     References:
         https://github.com/huggingface/pytorch-image-models/discussions/926
@@ -73,6 +74,13 @@ def generate_plot(processed_map):
     fig.tight_layout()
     return fig
 
+def serialize_images(processed_map):
+    """Serializes attention maps."""
+    for i in range(processed_map.shape[0]):
+        plt.imshow(processed_map[i].numpy())
+        plt.tile(f"Attention head: {i}")
+        plt.imsave(fname="attention_map_{i}.png")
+
 
 def generate_class_attn_map(image, block_id=0):
     """Collates the above utilities together for generating
@@ -85,7 +93,10 @@ def generate_class_attn_map(image, block_id=0):
 
     block_key = f"blocks_token_only.{block_id}.attn.softmax"
     processed_cls_attn_map = get_cls_attention_map(image_tensor, out, block_key)
-    return generate_plot(processed_cls_attn_map)
+    
+    serialize_images(processed_cls_attn_map)
+    all_attn_img_paths = sorted(glob.glob("attention_map_*.png"))
+    return all_attn_img_paths
 
 
 title = "Class Attention Maps"
@@ -97,7 +108,7 @@ iface = gr.Interface(
         gr.inputs.Image(type="pil", label="Input Image"),
         gr.Slider(0, 1, value=0, step=1, label="Block ID", info="Transformer Block ID"),
     ],
-    outputs=[gr.Plot(type="auto").style()],
+    outputs=gr.Gallery().style(grid=[2], height="auto"),
     title=title,
     article=article,
     allow_flagging="never",