Update app.py

app.py

@@ -1,37 +1,17 @@
+from typing import List, Tuple, Dict
+from collections import OrderedDict
+
 import gradio as gr
 import torch
-import timm
 import torch.nn.functional as F
-from timm.models import create_model
+import timm
 from timm.data import create_transform
+from timm.models import create_model
+from timm.utils import AttentionExtract
 from PIL import Image
 import numpy as np
 import matplotlib.pyplot as plt
-from typing import List, Tuple, Dict
-from collections import OrderedDict
 
-class AttentionExtractor:
-    def __init__(self, model: torch.nn.Module):
-        self.model = model
-        self.attention_maps = OrderedDict()
-        self._register_hooks()
-
-    def _register_hooks(self):
-        def hook_fn(module, input, output):
-            if isinstance(output, tuple):
-                self.attention_maps[module.full_name] = output[1]  # attention_probs
-            else:
-                self.attention_maps[module.full_name] = output
-
-        for name, module in self.model.named_modules():
-            # FIXME need to make more generic outside of vit
-            if name.lower().endswith('.attn_drop'):
-                module.full_name = name
-                print('hooking', name)
-                module.register_forward_hook(hook_fn)
-
-    def get_attention_maps(self) -> OrderedDict:
-        return self.attention_maps
 
 def get_attention_models() -> List[str]:
     """Get a list of timm models that have attention blocks."""
@@ -45,7 +25,7 @@ def load_model(model_name: str) -> Tuple[torch.nn.Module, AttentionExtractor]:
     timm.layers.set_fused_attn(False)
     model = create_model(model_name, pretrained=True)
     model.eval()
-    extractor = AttentionExtractor(model)
+    extractor = AttentionExtract(model, method='fx')  # can use 'hooks', can also allow specifying matching names for attention nodes or modules...
     return model, extractor
 
 def process_image(image: Image.Image, model: torch.nn.Module, extractor: AttentionExtractor) -> Dict[str, torch.Tensor]:
@@ -61,16 +41,11 @@ def process_image(image: Image.Image, model: torch.nn.Module, extractor: Attenti
         is_training=False
     )
     
-    
     # Preprocess the image
     tensor = transform(image).unsqueeze(0)
-    
-    # Forward pass
-    with torch.no_grad():
-        _ = model(tensor)
-    
+       
     # Extract attention maps
-    attention_maps = extractor.get_attention_maps()
+    attention_maps = extractor(tensor)
     
     return attention_maps