feat: adds different models capability

.github/copilot-instructions.md

@@ -0,0 +1,164 @@
+# ViTViz - AI Coding Agent Instructions
+
+## Project Overview
+
+**ViTViz** is a Gradio-based web app for visualizing Vision Transformer (ViT) attention mechanisms and adversarial attacks on image classification. The app supports:
+- Custom ViT model upload (.pth files) or Hugging Face Hub models
+- Multiple adversarial attack methods (FGSM, PGD, MIM, TGR, SAGA)
+- Attention visualization via Attention Rollout and per-layer/per-head views
+- Interactive iteration-by-iteration comparison of adversarial examples
+
+## Architecture
+
+### Core Components
+- **[app.py](app.py)**: Main Gradio interface with three tabs: Basic Classification, Attention Visualization, and Adversarial Attack Analysis
+- **[utils/model_loader.py](utils/model_loader.py)**: Handles model loading from local .pth files, Hugging Face Hub, or special `hf://` URIs. Includes:
+  - `ViTConfig` dataclass for dynamic architecture configuration
+  - Automatic architecture inference from state_dict or loaded model
+  - Hugging Face → timm state_dict conversion
+- **[utils/attacks.py](utils/attacks.py)**: Custom adversarial attack implementations that capture attention maps during attack iterations
+- **[utils/visualization.py](utils/visualization.py)**: Attention extraction via forward hooks, attention rollout computation, and overlay creation with dynamic grid size inference
+- **[utils/inference.py](utils/inference.py)**: Top-k prediction logic
+- **[utils/preprocessing.py](utils/preprocessing.py)**: ImageNet-standard transforms with dynamic `img_size` support
+
+### Key Design Patterns
+
+#### Dynamic Architecture Support (ViTConfig)
+The codebase now supports multiple ViT architectures via automatic inference:
+
+```python
+from utils.model_loader import ViTConfig, infer_config_from_model, infer_config_from_state_dict
+
+# ViTConfig contains all architecture parameters
+config = ViTConfig(
+    embed_dim=768,    # 384=small, 768=base, 1024=large
+    num_heads=12,     # 6=small, 12=base, 16=large
+    num_layers=12,    # varies by model
+    patch_size=16,    # 16 or 32
+    img_size=224,     # 224, 384, etc.
+    num_classes=1000
+)
+
+# Properties computed automatically
+config.grid_size      # img_size // patch_size (e.g., 14 for 224/16)
+config.num_patches    # grid_size ** 2
+config.timm_model_name  # e.g., "vit_base_patch16_224"
+```
+
+Supported architectures (auto-detected):
+- `vit_tiny_patch16_224` (embed_dim=192, heads=3)
+- `vit_small_patch16_224` (embed_dim=384, heads=6)
+- `vit_base_patch16_224` (embed_dim=768, heads=12)
+- `vit_large_patch16_224` (embed_dim=1024, heads=16)
+- `vit_base_patch32_224` (embed_dim=768, patch_size=32, grid=7)
+
+#### Model Loading Strategy
+The codebase supports multiple model sources:
+1. **Local .pth files**: Can contain full model, `state_dict`, `model_state_dict`, or checkpoint dicts with `class_names`
+2. **Hugging Face Hub**: Use `hf-model://username/repo-name` format; automatically converts HF ViT to timm-compatible format
+3. **Special `hf://` URIs**: For CNN backbones in SAGA attacks (e.g., `hf://lucasddmc/resnet101-stanford40-actions/resnet.pth`)
+
+The main loader returns 4 values:
+```python
+model, class_names, label_source, vit_config = load_model_and_labels(model_path, None, device=DEVICE)
+# vit_config.img_size, vit_config.grid_size, etc. are now available
+```
+
+#### Attention Capture with Forward Hooks
+All attention extraction uses PyTorch forward hooks on `model.blocks[i].attn` modules. The hook calculates Q, K, V manually and captures softmax attention weights before removal. See [visualization.py](utils/visualization.py#L12-L62).
+
+**Critical**: Attention tensors are immediately moved to CPU to avoid GPU memory accumulation during iterative attacks.
+
+**Dynamic grid size**: The `_infer_grid_size_from_attentions()` function automatically detects grid size from attention tensor shapes, eliminating hardcoded 14×14 assumptions.
+
+#### Adversarial Attack Iteration Tracking
+Custom attack classes (e.g., `PGDIterations`, `SAGA`) extend torchattacks and store:
+- `attentions_per_iter`: List of attention maps per iteration (each iteration = list of layer tensors)
+- Intermediate adversarial images via `tensor_to_pil()` with ImageNet denormalization
+
+See [attacks.py](utils/attacks.py#L38-L107) for the denormalization pattern used consistently across attacks.
+
+#### Gradio State Management
+The attack tab uses multiple Gradio `State` components to cache expensive computations:
+- `cached_attentions_state`: Raw attention maps from attack iterations
+- `per_iter_rollout_masks_state`: Pre-computed rollout masks for all iterations
+- `per_iter_layer_head_masks_state`: Pre-computed masks for all layers/heads (nested structure: `[iter][layer][head]`)
+
+This avoids re-running attacks when users adjust visualization parameters (discard ratio, head fusion, alpha overlay).
+
+## Development Workflows
+
+### Running the App
+```bash
+# Activate virtual environment (if exists)
+source venv/bin/activate
+
+# Install dependencies
+pip install -r requirements.txt
+
+# Run app locally (default port 7860)
+python app.py
+
+# Run on specific port
+PORT=8080 python app.py
+```
+
+### Model Conversion
+Use [convert_model_with_classes.py](convert_model_with_classes.py) to embed class names into checkpoint files:
+```bash
+python convert_model_with_classes.py
+```
+This extracts class names from Stanford40 dataset structure (`action_name_###.jpg`) and adds them to the checkpoint as `class_names` dict.
+
+### Adding New Attack Methods
+1. Subclass from `torchattacks.Attack` in [attacks.py](utils/attacks.py)
+2. Store `self.attentions_per_iter` as a list during `forward()` calls
+3. Call `capture_outputs_and_attentions(model, x_adv)` at each iteration to extract attention
+4. Return `(final_adv_tensor, iteration_images)` where `iteration_images` includes original + all intermediate steps
+5. Add attack to dropdown in [app.py](app.py#L270-L280) and handle parameters in `run_attack()` function
+
+Example: [SAGA attack](utils/attacks.py#L774-L920) implements ViT+CNN gradient blending.
+
+## Critical Conventions
+
+### Device Management
+All code uses `DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")` pattern. Models and tensors are explicitly moved to device at load/inference time.
+
+### Epsilon in [0,1] Space
+**Important**: Adversarial perturbation epsilon is defined in the **denormalized [0,1]** image space, not the normalized space. The L∞ distance metric in attack results also uses denormalized space for user interpretability. See [app.py](app.py#L271-L280).
+
+### Class Name Handling
+Supports three sources (priority order):
+1. External labels file (currently disabled in code)
+2. Embedded in checkpoint as `class_names` dict
+3. Hugging Face `id2label` from model config
+
+Returns `None` if unavailable; UI falls back to showing class indices.
+
+### Gradio File Input Extraction
+Use `_to_path()` helper ([app.py](app.py#L45-L57)) to handle different Gradio file input formats (string, dict with 'name', object with .name attribute).
+
+### Custom CSS and Icons
+The app injects Bootstrap Icons via CDN and custom CSS for panels/tables. Icon constants (e.g., `ICON_SUCCESS`, `ICON_FAIL`) are defined at the top of [app.py](app.py#L26-L30).
+
+## External Dependencies
+
+- **timm**: ViT model architecture (`vit_base_patch16_224` is the default)
+- **torchattacks**: Base classes for adversarial attacks
+- **transformers**: Optional, for loading HF Hub models
+- **gradio**: Version 5.49.1 (specified in requirements)
+
+## Testing Strategy
+
+Currently no automated tests. Manual testing workflow:
+1. Upload model → check classification works
+2. Run attention visualization → verify heatmaps align with predicted class
+3. Run attack → verify iteration slider shows progression
+4. Toggle layer/head sliders → verify attention updates without re-running attack
+
+## Known Limitations
+
+- Supports timm ViT architectures (tiny, small, base, large) with patch sizes 16 and 32
+- No support for non-standard ViT variants (DeiT distillation token, Swin hierarchical, BEiT) without additional conversion
+- Custom CSS may break with Gradio version updates
+- No batch processing support (processes one image at a time)

app.py

@@ -6,7 +6,7 @@ from PIL import Image
 from typing import Optional, List, Tuple
 from pathlib import Path
 
-from utils.model_loader import load_model_and_labels
+from utils.model_loader import load_model_and_labels, ViTConfig
 from utils.preprocessing import get_default_transform, preprocess_image
 from utils.inference import predict_topk
 from utils.attacks import PGDIterations, FGSM, SAGA, MIFGSM, TGR
@@ -32,7 +32,6 @@ ICON_CHART = '<i class="bi bi-bar-chart-line-fill vitviz-bi" aria-hidden="true">
 ICON_RULER = '<i class="bi bi-speedometer2 vitviz-bi" aria-hidden="true"></i>'
 ICON_GEAR = '<i class="bi bi-gear-fill vitviz-bi" aria-hidden="true"></i>'
 
-transform = get_default_transform()
 # Backbone CNN opcional usado no modo "SAGA (with CNN gradient)".
 # Pode ser um caminho local (ex.: "models/resnet.pth") ou um checkpoint no Hugging Face Hub.
 RESNET_BACKBONE_SPEC = "hf://lucasddmc/resnet101-stanford40-actions/resnet.pth"
@@ -89,14 +88,15 @@ def classify_image(model_file, use_hf_vit: bool, image):
         model_path = HF_VIT_MODEL_SPEC if use_hf_vit else _to_path(model_file)
 
         # Carregar modelo e labels
-        model, class_names, label_source = load_model_and_labels(model_path, None, device=DEVICE)
+        model, class_names, label_source, vit_config = load_model_and_labels(model_path, None, device=DEVICE)
         # _print_model_heads(model)
 
-        # Processar imagem
+        # Processar imagem com transform dinâmico baseado no modelo
         if not (isinstance(image, str) or isinstance(image, Image.Image)):
             return "Please upload a valid image"
 
-        img_tensor = preprocess_image(image, transform=transform).to(DEVICE)
+        dynamic_transform = get_default_transform(img_size=vit_config.img_size)
+        img_tensor = preprocess_image(image, transform=dynamic_transform).to(DEVICE)
 
         # Inferência
         top_prob, top_idx, num_classes, probabilities = predict_topk(model, img_tensor, top_k=5)
@@ -159,11 +159,12 @@ def visualize_attention(
         
         # Carregar modelo e labels
         model_path = HF_VIT_MODEL_SPEC if use_hf_vit else _to_path(model_file)
-        model, class_names, label_source = load_model_and_labels(model_path, None, device=DEVICE)
+        model, class_names, label_source, vit_config = load_model_and_labels(model_path, None, device=DEVICE)
         # _print_model_heads(model)
         
-        # Processar imagem
-        img_tensor = preprocess_image(image, transform=transform).to(DEVICE)
+        # Processar imagem com transform dinâmico baseado no modelo
+        dynamic_transform = get_default_transform(img_size=vit_config.img_size)
+        img_tensor = preprocess_image(image, transform=dynamic_transform).to(DEVICE)
         
         # Predição
         top_prob, top_idx, num_classes, _ = predict_topk(model, img_tensor, top_k=1, device=DEVICE)
@@ -249,11 +250,12 @@ def run_attack(
         
         # Carregar modelo e labels
         model_path = HF_VIT_MODEL_SPEC if use_hf_vit else _to_path(model_file)
-        model, class_names, label_source = load_model_and_labels(model_path, None, device=DEVICE)
+        model, class_names, label_source, vit_config = load_model_and_labels(model_path, None, device=DEVICE)
         # _print_model_heads(model)
         
-        # Processar imagem
-        img_tensor = preprocess_image(image, transform=transform).to(DEVICE)
+        # Processar imagem com transform dinâmico baseado no modelo
+        dynamic_transform = get_default_transform(img_size=vit_config.img_size)
+        img_tensor = preprocess_image(image, transform=dynamic_transform).to(DEVICE)
         
         # Predição original (top-5 para comparação)
         top_prob_orig, top_idx_orig, num_classes, _ = predict_topk(model, img_tensor, top_k=5, device=DEVICE)

utils/model_loader.py

@@ -1,6 +1,7 @@
 import pickle
 import torch
 import timm
+from dataclasses import dataclass
 from typing import Optional, Tuple, Dict, Any
 
 try:
@@ -11,6 +12,128 @@ except Exception:  # pragma: no cover
 DEVICE_DEFAULT = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 
+@dataclass
+class ViTConfig:
+    """Configuração de arquitetura ViT extraída dinamicamente do modelo."""
+    embed_dim: int = 768
+    num_heads: int = 12
+    num_layers: int = 12
+    patch_size: int = 16
+    img_size: int = 224
+    num_classes: int = 1000
+    
+    @property
+    def grid_size(self) -> int:
+        """Tamanho do grid de patches (ex: 224/16 = 14)."""
+        return self.img_size // self.patch_size
+    
+    @property
+    def num_patches(self) -> int:
+        """Número total de patches (ex: 14*14 = 196)."""
+        return self.grid_size ** 2
+    
+    @property
+    def timm_model_name(self) -> str:
+        """Retorna o nome do modelo timm correspondente à configuração."""
+        # Mapeamento baseado em embed_dim e num_heads
+        size_map = {
+            (192, 3): 'tiny',
+            (384, 6): 'small', 
+            (768, 12): 'base',
+            (1024, 16): 'large',
+            (1280, 16): 'huge',
+        }
+        size = size_map.get((self.embed_dim, self.num_heads), 'base')
+        return f"vit_{size}_patch{self.patch_size}_{self.img_size}"
+
+
+def infer_config_from_model(model: torch.nn.Module) -> ViTConfig:
+    """Infere configuração ViT a partir de um modelo timm carregado."""
+    config = ViTConfig()
+    
+    # Extrair img_size e patch_size do patch_embed
+    if hasattr(model, 'patch_embed'):
+        pe = model.patch_embed
+        if hasattr(pe, 'img_size'):
+            img_size = pe.img_size
+            config.img_size = img_size[0] if isinstance(img_size, (tuple, list)) else img_size
+        if hasattr(pe, 'patch_size'):
+            patch_size = pe.patch_size
+            config.patch_size = patch_size[0] if isinstance(patch_size, (tuple, list)) else patch_size
+    
+    # Extrair num_layers, embed_dim, num_heads dos blocks
+    if hasattr(model, 'blocks') and len(model.blocks) > 0:
+        config.num_layers = len(model.blocks)
+        block = model.blocks[0]
+        if hasattr(block, 'attn'):
+            attn = block.attn
+            if hasattr(attn, 'num_heads'):
+                config.num_heads = attn.num_heads
+            if hasattr(attn, 'qkv') and hasattr(attn.qkv, 'in_features'):
+                config.embed_dim = attn.qkv.in_features
+    
+    # Extrair num_classes do head
+    if hasattr(model, 'head') and hasattr(model.head, 'out_features'):
+        config.num_classes = model.head.out_features
+    elif hasattr(model, 'head') and hasattr(model.head, 'weight'):
+        config.num_classes = model.head.weight.shape[0]
+    
+    return config
+
+
+def infer_config_from_state_dict(state_dict: Dict[str, torch.Tensor]) -> ViTConfig:
+    """Infere configuração ViT a partir de um state_dict."""
+    config = ViTConfig()
+    
+    # Inferir num_layers contando blocks
+    layer_indices = set()
+    for key in state_dict.keys():
+        if key.startswith('blocks.') and '.attn.' in key:
+            # blocks.0.attn.qkv.weight -> extrair 0
+            idx = int(key.split('.')[1])
+            layer_indices.add(idx)
+    if layer_indices:
+        config.num_layers = max(layer_indices) + 1
+    
+    # Inferir embed_dim e num_heads do primeiro bloco
+    qkv_key = 'blocks.0.attn.qkv.weight'
+    if qkv_key in state_dict:
+        qkv_weight = state_dict[qkv_key]
+        # qkv.weight shape: [3*embed_dim, embed_dim]
+        config.embed_dim = qkv_weight.shape[1]
+    
+    # Inferir num_heads do proj bias ou de forma heurística
+    proj_key = 'blocks.0.attn.proj.weight'
+    if proj_key in state_dict:
+        # proj.weight shape: [embed_dim, embed_dim]
+        embed_dim = state_dict[proj_key].shape[0]
+        # Heurística: head_dim típico é 64
+        config.num_heads = embed_dim // 64
+    
+    # Inferir num_classes do head
+    head_key = 'head.weight'
+    if head_key in state_dict:
+        config.num_classes = state_dict[head_key].shape[0]
+    
+    # Inferir patch_size e img_size do patch_embed
+    patch_proj_key = 'patch_embed.proj.weight'
+    if patch_proj_key in state_dict:
+        # shape: [embed_dim, 3, patch_size, patch_size]
+        patch_weight = state_dict[patch_proj_key]
+        config.patch_size = patch_weight.shape[2]
+    
+    # Inferir img_size do pos_embed
+    pos_embed_key = 'pos_embed'
+    if pos_embed_key in state_dict:
+        # shape: [1, num_patches+1, embed_dim]
+        num_tokens = state_dict[pos_embed_key].shape[1]
+        num_patches = num_tokens - 1  # -1 para CLS token
+        grid_size = int(num_patches ** 0.5)
+        config.img_size = grid_size * config.patch_size
+    
+    return config
+
+
 def _hf_id2label_to_class_names(id2label: Any) -> Optional[Dict[int, str]]:
     if not isinstance(id2label, dict):
         return None
@@ -77,8 +200,12 @@ def _convert_hf_vit_to_timm_state_dict(hf_sd: Dict[str, torch.Tensor], num_layer
     return out
 
 
-def load_vit_from_huggingface(model_id: str, device: Optional[torch.device] = None) -> Tuple[torch.nn.Module, Optional[Dict[int, str]]]:
-    """Carrega ViT do Hugging Face Hub e retorna um modelo timm equivalente."""
+def load_vit_from_huggingface(model_id: str, device: Optional[torch.device] = None) -> Tuple[torch.nn.Module, Optional[Dict[int, str]], ViTConfig]:
+    """Carrega ViT do Hugging Face Hub e retorna um modelo timm equivalente.
+    
+    Returns:
+        (model, class_names, config)
+    """
     if AutoModelForImageClassification is None:
         raise RuntimeError("transformers não está instalado; instale 'transformers' para carregar do Hugging Face.")
 
@@ -88,14 +215,40 @@ def load_vit_from_huggingface(model_id: str, device: Optional[torch.device] = No
     cfg = getattr(hf_model, "config", None)
     num_labels = int(getattr(cfg, "num_labels", 1000)) if cfg is not None else 1000
     num_layers = int(getattr(cfg, "num_hidden_layers", 12)) if cfg is not None else 12
+    hidden_size = int(getattr(cfg, "hidden_size", 768)) if cfg is not None else 768
+    num_heads = int(getattr(cfg, "num_attention_heads", 12)) if cfg is not None else 12
+    patch_size = int(getattr(cfg, "patch_size", 16)) if cfg is not None else 16
+    img_size = int(getattr(cfg, "image_size", 224)) if cfg is not None else 224
     class_names = _hf_id2label_to_class_names(getattr(cfg, "id2label", None)) if cfg is not None else None
 
-    timm_model = timm.create_model("vit_base_patch16_224", pretrained=False, num_classes=num_labels)
+    # Criar config dinâmico
+    vit_config = ViTConfig(
+        embed_dim=hidden_size,
+        num_heads=num_heads,
+        num_layers=num_layers,
+        patch_size=patch_size,
+        img_size=img_size,
+        num_classes=num_labels
+    )
+    
+    # Tentar encontrar o modelo timm correspondente
+    timm_name = vit_config.timm_model_name
+    try:
+        timm_model = timm.create_model(timm_name, pretrained=False, num_classes=num_labels)
+    except Exception:
+        # Fallback para vit_base_patch16_224 se o modelo não existir
+        print(f"[ViTViz] Modelo timm '{timm_name}' não encontrado, usando vit_base_patch16_224")
+        timm_model = timm.create_model("vit_base_patch16_224", pretrained=False, num_classes=num_labels)
+    
     timm_sd = _convert_hf_vit_to_timm_state_dict(hf_model.state_dict(), num_layers=num_layers)
     timm_model.load_state_dict(timm_sd, strict=False)
     timm_model = timm_model.to(device)
     timm_model.eval()
-    return timm_model, class_names
+    
+    # Atualizar config com valores reais do modelo carregado
+    vit_config = infer_config_from_model(timm_model)
+    
+    return timm_model, class_names, vit_config
 
 
 class CustomUnpickler(pickle.Unpickler):
@@ -193,59 +346,94 @@ def load_class_names_from_file(labels_file: Optional[str]) -> Optional[Dict[int,
         return None
 
 
-def build_model_from_checkpoint(checkpoint: Any, device: Optional[torch.device] = None) -> torch.nn.Module:
-    """Constroi um modelo a partir de um checkpoint que pode ser um dict, state_dict ou o próprio modelo."""
+def build_model_from_checkpoint(checkpoint: Any, device: Optional[torch.device] = None) -> Tuple[torch.nn.Module, ViTConfig]:
+    """Constroi um modelo a partir de um checkpoint que pode ser um dict, state_dict ou o próprio modelo.
+    
+    Returns:
+        (model, config) - modelo carregado e configuração inferida
+    """
     device = device or DEVICE_DEFAULT
+    config: Optional[ViTConfig] = None
+    
     if isinstance(checkpoint, dict):
         if 'model' in checkpoint:
             model = checkpoint['model']
+            config = infer_config_from_model(model)
         elif 'state_dict' in checkpoint:
             state_dict = checkpoint['state_dict']
+            config = infer_config_from_state_dict(state_dict)
             num_classes = infer_num_classes(state_dict)
-            # TODO: fazer tratamento dinâmico para timm, pytorch, etc.
-            model = timm.create_model('vit_base_patch16_224', pretrained=False, num_classes=num_classes)
+            config.num_classes = num_classes
+            # Usar arquitetura inferida
+            timm_name = config.timm_model_name
+            try:
+                model = timm.create_model(timm_name, pretrained=False, num_classes=num_classes)
+            except Exception:
+                print(f"[ViTViz] Modelo timm '{timm_name}' não encontrado, usando vit_base_patch16_224")
+                model = timm.create_model('vit_base_patch16_224', pretrained=False, num_classes=num_classes)
             model.load_state_dict(state_dict)
         elif 'model_state_dict' in checkpoint:
             # Novo formato com class_names embutidas
             state_dict = checkpoint['model_state_dict']
+            config = infer_config_from_state_dict(state_dict)
             num_classes = infer_num_classes(state_dict)
-            # TODO: fazer tratamento dinâmico para timm, pytorch, etc.
-            model = timm.create_model('vit_base_patch16_224', pretrained=False, num_classes=num_classes)
+            config.num_classes = num_classes
+            # Usar arquitetura inferida
+            timm_name = config.timm_model_name
+            try:
+                model = timm.create_model(timm_name, pretrained=False, num_classes=num_classes)
+            except Exception:
+                print(f"[ViTViz] Modelo timm '{timm_name}' não encontrado, usando vit_base_patch16_224")
+                model = timm.create_model('vit_base_patch16_224', pretrained=False, num_classes=num_classes)
             model.load_state_dict(state_dict)
         else:
             # assume dict é um state_dict
+            config = infer_config_from_state_dict(checkpoint)
             num_classes = infer_num_classes(checkpoint)
-            # TODO: fazer tratamento dinâmico para timm, pytorch, etc.
-            model = timm.create_model('vit_base_patch16_224', pretrained=False, num_classes=num_classes)
+            config.num_classes = num_classes
+            # Usar arquitetura inferida
+            timm_name = config.timm_model_name
+            try:
+                model = timm.create_model(timm_name, pretrained=False, num_classes=num_classes)
+            except Exception:
+                print(f"[ViTViz] Modelo timm '{timm_name}' não encontrado, usando vit_base_patch16_224")
+                model = timm.create_model('vit_base_patch16_224', pretrained=False, num_classes=num_classes)
             model.load_state_dict(checkpoint)
     else:
         # modelo completo salvo via torch.save(model, ...)
         model = checkpoint
+        config = infer_config_from_model(model)
 
     model = model.to(device)
     model.eval()
-    return model
+    
+    # Garantir que config está preenchido
+    if config is None:
+        config = infer_config_from_model(model)
+    
+    return model, config
 
 
 def load_model_and_labels(
     model_path: str,
     labels_file: Optional[str] = None,
     device: Optional[torch.device] = None,
-) -> Tuple[torch.nn.Module, Optional[Dict[int, str]], Optional[str]]:
+) -> Tuple[torch.nn.Module, Optional[Dict[int, str]], Optional[str], ViTConfig]:
     """
     ** Função Principal **
     Carrega modelo e, se disponível, nomes de classes.
 
-    Retorna: (model, class_names, origem_labels) onde origem_labels ∈ {"file", "checkpoint", None}
+    Retorna: (model, class_names, origem_labels, config) onde origem_labels ∈ {"file", "checkpoint", "hf", None}
         None se não houver nomes de classes disponíveis.
+        config contém a configuração da arquitetura ViT (embed_dim, num_heads, grid_size, etc.)
     """
     device = device or DEVICE_DEFAULT
 
     # Carregar diretamente do Hugging Face Hub (Transformers -> timm)
     if isinstance(model_path, str) and model_path.startswith("hf-model://"):
         model_id = model_path[len("hf-model://"):].strip("/")
-        model, class_names = load_vit_from_huggingface(model_id, device=device)
-        return model, class_names, 'hf'
+        model, class_names, config = load_vit_from_huggingface(model_id, device=device)
+        return model, class_names, 'hf', config
 
     checkpoint = load_checkpoint(model_path, device=device)
     class_names_ckpt = extract_class_names(checkpoint)
@@ -260,5 +448,5 @@ def load_model_and_labels(
     class_names = class_names_ckpt
     source = 'checkpoint' if class_names_ckpt else None
 
-    model = build_model_from_checkpoint(checkpoint, device=device)
-    return model, class_names, source
+    model, config = build_model_from_checkpoint(checkpoint, device=device)
+    return model, class_names, source, config

utils/preprocessing.py

@@ -4,13 +4,20 @@ from torchvision import transforms
 import torch
 
 
-# TODO: implementar adapters para diferentes modelos com outros tipos de classes
-
-def get_default_transform() -> transforms.Compose:
-    """Transform padrão (Resize+CenterCrop+Normalize) compatível com modelos ImageNet."""
+def get_default_transform(img_size: int = 224) -> transforms.Compose:
+    """Transform padrão (Resize+CenterCrop+Normalize) compatível com modelos ImageNet.
+    
+    Args:
+        img_size: Tamanho da imagem de entrada do modelo (default: 224)
+    
+    Returns:
+        Compose de transforms para preprocessamento
+    """
+    # Resize proporcional: 256 para 224, escala para outros tamanhos
+    resize_size = int(img_size * 256 / 224)
     return transforms.Compose([
-        transforms.Resize(256),
-        transforms.CenterCrop(224),
+        transforms.Resize(resize_size),
+        transforms.CenterCrop(img_size),
         transforms.ToTensor(),
         transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
     ])

utils/visualization.py

@@ -77,6 +77,24 @@ def extract_attention_maps(model, image: torch.Tensor) -> list:
     return attentions
 
 
+def _infer_grid_size_from_attentions(attentions_per_iter: list) -> int:
+    """Infere o tamanho do grid a partir dos tensores de atenção."""
+    if not attentions_per_iter:
+        return 14
+    for iter_attns in attentions_per_iter:
+        if not iter_attns:
+            continue
+        for layer_tensor in iter_attns:
+            if isinstance(layer_tensor, torch.Tensor) and layer_tensor.ndim == 4:
+                # shape: [B, H, T, T] onde T = num_patches + 1 (CLS)
+                num_tokens = layer_tensor.shape[-1]
+                num_patches = num_tokens - 1
+                side = int(num_patches ** 0.5)
+                if side * side == num_patches:
+                    return side
+    return 14  # fallback
+
+
 def extract_layer_head_masks(
     attentions_per_iter: list,
     layer_idx: int,
@@ -98,10 +116,14 @@ def extract_layer_head_masks(
     masks = []
     if attentions_per_iter is None or len(attentions_per_iter) == 0:
         return masks
+    
+    # Inferir grid_size dinamicamente
+    default_grid = _infer_grid_size_from_attentions(attentions_per_iter)
     eps = 1e-8
+    
     for iter_attns in attentions_per_iter:
         if not iter_attns or layer_idx < 0 or layer_idx >= len(iter_attns):
-            masks.append(np.zeros((14, 14), dtype=np.float32))
+            masks.append(np.zeros((default_grid, default_grid), dtype=np.float32))
             continue
         layer_tensor = iter_attns[layer_idx]
         if isinstance(layer_tensor, torch.Tensor):
@@ -109,7 +131,7 @@ def extract_layer_head_masks(
         else:
             att = torch.as_tensor(layer_tensor)
         if att.ndim != 4 or att.size(0) < 1 or head_idx < 0 or head_idx >= att.size(1):
-            masks.append(np.zeros((14, 14), dtype=np.float32))
+            masks.append(np.zeros((default_grid, default_grid), dtype=np.float32))
             continue
         att_head = att[0, head_idx]  # [T,T]
         vec = att_head[0] if cls_only else att_head.mean(dim=0)
@@ -117,7 +139,7 @@ def extract_layer_head_masks(
         tokens = vec_patches.numel()
         side = int(tokens ** 0.5)
         if side * side != tokens:
-            masks.append(np.zeros((14, 14), dtype=np.float32))
+            masks.append(np.zeros((default_grid, default_grid), dtype=np.float32))
             continue
         mask = vec_patches.reshape(side, side)
         mask = mask / (mask.max() + eps)
@@ -164,6 +186,18 @@ def compute_layer_head_masks_from_cached_attns(iter_attns: List[torch.Tensor], c
     """
     per_layer_head_masks: List[List[np.ndarray]] = []
     eps = 1e-8
+    
+    # Inferir grid_size do primeiro tensor válido
+    default_grid = 14
+    for layer_tensor in iter_attns:
+        if isinstance(layer_tensor, torch.Tensor) and layer_tensor.ndim == 4:
+            num_tokens = layer_tensor.shape[-1]
+            num_patches = num_tokens - 1
+            side = int(num_patches ** 0.5)
+            if side * side == num_patches:
+                default_grid = side
+                break
+    
     for li, layer_tensor in enumerate(iter_attns):
         if isinstance(layer_tensor, torch.Tensor):
             att = layer_tensor.detach().cpu()
@@ -183,7 +217,7 @@ def compute_layer_head_masks_from_cached_attns(iter_attns: List[torch.Tensor], c
             side = int(tokens ** 0.5)
             if side * side != tokens:
                 # print(f"[ViTViz][compute_layer_head_masks] Layer {li} head {h}: tokens {tokens} not square -> side={side}")
-                heads_masks.append(np.zeros((14, 14), dtype=np.float32))
+                heads_masks.append(np.zeros((default_grid, default_grid), dtype=np.float32))
                 continue
             mask = vec_patches.reshape(side, side)
             mmax = float(mask.max())
@@ -460,10 +494,13 @@ def extract_last_layer_head_masks(
     if attentions_per_iter is None or len(attentions_per_iter) == 0:
         return masks
 
+    # Inferir grid_size dinamicamente
+    default_grid = _infer_grid_size_from_attentions(attentions_per_iter)
     eps = 1e-8
+    
     for iter_attns in attentions_per_iter:
         if not iter_attns:
-            masks.append(np.zeros((14, 14), dtype=np.float32))
+            masks.append(np.zeros((default_grid, default_grid), dtype=np.float32))
             print("Atenções vazias para esta iteração.")
             continue
         # Última camada
@@ -475,7 +512,7 @@ def extract_last_layer_head_masks(
 
         # Esperado: [B, H, T, T] com B=1
         if att.ndim != 4 or att.size(0) < 1 or head_idx < 0 or head_idx >= att.size(1):
-            masks.append(np.zeros((14, 14), dtype=np.float32))
+            masks.append(np.zeros((default_grid, default_grid), dtype=np.float32))
             print("Atenção inválida na última camada.")
             continue
 
@@ -495,7 +532,7 @@ def extract_last_layer_head_masks(
         side = int(tokens ** 0.5)
         if side * side != tokens:
             # fallback: normalizar e retornar zeros coerentes
-            masks.append(np.zeros((14, 14), dtype=np.float32))
+            masks.append(np.zeros((default_grid, default_grid), dtype=np.float32))
             print("Número de patches não forma uma grade quadrada.")
             continue