Update app.py

app.py

@@ -20,7 +20,7 @@ def get_attention_models() -> List[str]:
     attention_models = [model for model in all_models if any([model.lower().startswith(p) for p in ('vit', 'deit', 'beit', 'eva')])]
     return attention_models
 
-def load_model(model_name: str) -> Tuple[torch.nn.Module, AttentionExtractor]:
+def load_model(model_name: str) -> Tuple[torch.nn.Module, AttentionExtract]:
     """Load a model from timm and prepare it for attention extraction."""
     timm.layers.set_fused_attn(False)
     model = create_model(model_name, pretrained=True)
@@ -28,7 +28,11 @@ def load_model(model_name: str) -> Tuple[torch.nn.Module, AttentionExtractor]:
     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]:
+def process_image(
+        image: Image.Image,
+        model: torch.nn.Module,
+        extractor: AttentionExtract
+) -> Dict[str, torch.Tensor]:
     """Process the input image and get the attention maps."""
     # Get the correct transform for the model
     config = model.pretrained_cfg