add visual method changing: functions and pooling aggregated

app.py

@@ -5,6 +5,7 @@ from janus.models import MultiModalityCausalLM, VLChatProcessor
 from janus.utils.io import load_pil_images
 from demo.visualization import generate_gradcam, VisualizationJanus, VisualizationClip, VisualizationChartGemma, VisualizationLLaVA
 from demo.model_utils import Clip_Utils, Janus_Utils, LLaVA_Utils, ChartGemma_Utils, add_title_to_image
+from demo.modified_attn import ModifiedLlamaAttention, ModifiedGemmaAttention
 
 import numpy as np
 import matplotlib.pyplot as plt
@@ -53,7 +54,7 @@ def clean():
 @spaces.GPU(duration=120)
 def multimodal_understanding(model_type, 
                              activation_map_method, 
-                             visual_pooling_method, 
+                             visual_method, 
                              image, question, seed, top_p, temperature, target_token_idx,
                              visualization_layer_min, visualization_layer_max, focus, response_type, chart_type):
     # Clear CUDA cache before generating
@@ -83,7 +84,7 @@ def multimodal_understanding(model_type,
             else:
                 target_layers = [all_layers[visualization_layer_min-1]]
             grad_cam = VisualizationClip(clip_utils.model, target_layers)
-            cam, outputs, grid_size = grad_cam.generate_cam(inputs, target_token_idx=0, visual_pooling_method=visual_pooling_method)
+            cam, outputs, grid_size = grad_cam.generate_cam(inputs, target_token_idx=0, visual_method=visual_method)
             cam = cam.to("cpu")
             cam = [generate_gradcam(cam, image, size=(224, 224))]
             grad_cam.remove_hooks()
@@ -144,7 +145,7 @@ def multimodal_understanding(model_type,
             cam = []
             if focus == "Visual Encoder":
                 if target_token_idx != -1:
-                    cam_tensors, grid_size, start = gradcam.generate_cam(prepare_inputs, tokenizer, temperature, top_p, target_token_idx, visual_pooling_method, focus)
+                    cam_tensors, grid_size, start = gradcam.generate_cam(prepare_inputs, tokenizer, temperature, top_p, target_token_idx, visual_method, focus)
                     cam_grid = cam_tensors.reshape(grid_size, grid_size)
                     cam_i = generate_gradcam(cam_grid, image)
                     cam_i = add_title_to_image(cam_i, input_ids_decoded[start + target_token_idx])
@@ -159,7 +160,7 @@ def multimodal_understanding(model_type,
                             gradcam = VisualizationLLaVA(vl_gpt, target_layers)
                         elif model_name.split('-')[0] == "ChartGemma":
                             gradcam = VisualizationChartGemma(vl_gpt, target_layers)
-                        cam_tensors, grid_size, start = gradcam.generate_cam(prepare_inputs, tokenizer, temperature, top_p, i, visual_pooling_method, focus)
+                        cam_tensors, grid_size, start = gradcam.generate_cam(prepare_inputs, tokenizer, temperature, top_p, i, visual_method, focus)
                         cam_grid = cam_tensors.reshape(grid_size, grid_size)
                         cam_i = generate_gradcam(cam_grid, image)
                         cam_i = add_title_to_image(cam_i, input_ids_decoded[start + i])
@@ -167,7 +168,7 @@ def multimodal_understanding(model_type,
                         gradcam.remove_hooks()
                         i += 1
             else:
-                cam_tensors, grid_size, start = gradcam.generate_cam(prepare_inputs, tokenizer, temperature, top_p, target_token_idx, visual_pooling_method, focus)
+                cam_tensors, grid_size, start = gradcam.generate_cam(prepare_inputs, tokenizer, temperature, top_p, target_token_idx, visual_method, focus)
                 if target_token_idx != -1:
                     input_text_decoded = input_ids_decoded[start + target_token_idx]
                     for i, cam_tensor in enumerate(cam_tensors):
@@ -190,13 +191,11 @@ def multimodal_understanding(model_type,
 
     # Collect Results
     RESULTS_ROOT = "./results"
-    FILES_ROOT = f"{RESULTS_ROOT}/{model_name}/{focus}/{chart_type}/layer{visualization_layer_min}-{visualization_layer_max}"
+    FILES_ROOT = f"{RESULTS_ROOT}/{model_name}/{focus}/{visual_method}/{chart_type}/layer{visualization_layer_min}-{visualization_layer_max}/{'all_tokens' if target_token_idx == -1 else f'--{input_ids_decoded[start + target_token_idx]}--'}"
     os.makedirs(FILES_ROOT, exist_ok=True)
-    if focus == "Visual Encoder":
-        cam[0].save(f"{FILES_ROOT}/{visual_pooling_method}.png")
-    else:
-        for i, cam_p in enumerate(cam):
-            cam_p.save(f"{FILES_ROOT}/{i}.png")
+    
+    for i, cam_p in enumerate(cam):
+        cam_p.save(f"{FILES_ROOT}/{i}.png")
             
     with open(f"{FILES_ROOT}/input_text_decoded.txt", "w") as f:
         f.write(input_text_decoded)
@@ -218,36 +217,58 @@ def multimodal_understanding(model_type,
 def model_slider_change(model_type):
     global model_utils, vl_gpt, tokenizer, clip_utils, model_name, language_model_max_layer, language_model_best_layer, vision_model_best_layer
     model_name = model_type
+
+
+    encoder_only_res = [
+        gr.Dropdown(choices=["Visualization only"], value="Visualization only", label="response_type"),
+        gr.Dropdown(choices=["Visual Encoder"], value="Visual Encoder", label="focus"),
+        gr.Dropdown(choices=["GradCAM"], value="GradCAM", label="activation map type"),
+        gr.Dropdown(choices=["CLS", "max", "avg"], value="CLS", label="visual pooling method")
+    ]
+
+    visual_res = [
+        gr.Dropdown(choices=["Visualization only", "answer + visualization"], value="Visualization only", label="response_type"),
+        gr.Dropdown(choices=["Visual Encoder"], value="Visual Encoder", label="focus"),
+        gr.Dropdown(choices=["GradCAM"], value="GradCAM", label="activation map type"),
+        gr.Dropdown(choices=["softmax", "sigmoid"], value="softmax", label="activation function")
+    ]
+
+    language_res = [
+        gr.Dropdown(choices=["Visualization only", "answer + visualization"], value="answer + visualization", label="response_type"),
+        gr.Dropdown(choices=["Language Model"], value="Language Model", label="focus"),
+        gr.Dropdown(choices=["GradCAM"], value="GradCAM", label="activation map type"),
+        gr.Dropdown(choices=["softmax", "sigmoid"], value="softmax", label="activation function")
+    ]
+
+
     if model_type == "Clip":
         clean()
         set_seed()
         clip_utils = Clip_Utils()
         clip_utils.init_Clip()
-        res = (
-            gr.Dropdown(choices=["Visualization only"], value="Visualization only", label="response_type"),
+        sliders = [
             gr.Slider(minimum=1, maximum=12, value=12, step=1, label="visualization layers min"),
             gr.Slider(minimum=1, maximum=12, value=12, step=1, label="visualization layers max"),
-            gr.Dropdown(choices=["Visual Encoder"], value="Visual Encoder", label="focus"),
-            gr.Dropdown(choices=["GradCAM"], value="GradCAM", label="activation map type")
-        )
-        return res
+        ]
+        return tuple(encoder_only_res + sliders)
+    
     elif model_type.split('-')[0] == "Janus":
         
         clean()
         set_seed()
         model_utils = Janus_Utils()
         vl_gpt, tokenizer = model_utils.init_Janus(model_type.split('-')[-1])
+        for layer in vl_gpt.language_model.model.layers:
+            layer.self_attn = ModifiedLlamaAttention(layer.self_attn)
+        
         language_model_max_layer = 24
         language_model_best_layer = 8
-
-        res = (
-            gr.Dropdown(choices=["Visualization only", "answer + visualization"], value="answer + visualization", label="response_type"),
+        
+        sliders = [
             gr.Slider(minimum=1, maximum=24, value=24, step=1, label="visualization layers min"),
             gr.Slider(minimum=1, maximum=24, value=24, step=1, label="visualization layers max"),
-            gr.Dropdown(choices=["Visual Encoder", "Language Model"], value="Visual Encoder", label="focus"),
-            gr.Dropdown(choices=["GradCAM"], value="GradCAM", label="activation map type")
-        )
-        return res
+        ]
+        return tuple(visual_res + sliders)
     
     elif model_type.split('-')[0] == "LLaVA":
         
@@ -259,32 +280,28 @@ def model_slider_change(model_type):
         language_model_max_layer = 32 if version == "1.5" else 28
         language_model_best_layer = 10
 
-        res = (
-            gr.Dropdown(choices=["Visualization only", "answer + visualization"], value="answer + visualization", label="response_type"),
+        sliders = [
             gr.Slider(minimum=1, maximum=language_model_max_layer, value=language_model_best_layer, step=1, label="visualization layers min"),
             gr.Slider(minimum=1, maximum=language_model_max_layer, value=language_model_best_layer, step=1, label="visualization layers max"),
-            gr.Dropdown(choices=["Language Model"], value="Language Model", label="focus"),
-            gr.Dropdown(choices=["GradCAM"], value="GradCAM", label="activation map type")
-        )
-        return res
+        ]
+        return tuple(language_res + sliders)
     
     elif model_type.split('-')[0] == "ChartGemma":
         clean()
         set_seed()
         model_utils = ChartGemma_Utils()
         vl_gpt, tokenizer = model_utils.init_ChartGemma()
+        for layer in vl_gpt.language_model.model.layers:
+            layer.self_attn = ModifiedGemmaAttention(layer.self_attn)
         language_model_max_layer = 18
         vision_model_best_layer = 19
         language_model_best_layer = 15
 
-        res = (
-            gr.Dropdown(choices=["Visualization only", "answer + visualization"], value="answer + visualization", label="response_type"),
+        sliders = [
             gr.Slider(minimum=1, maximum=language_model_best_layer, value=language_model_best_layer, step=1, label="visualization layers min"),
             gr.Slider(minimum=1, maximum=language_model_best_layer, value=language_model_best_layer, step=1, label="visualization layers max"),
-            gr.Dropdown(choices=["Visual Encoder", "Language Model"], value="Language Model", label="focus"),
-            gr.Dropdown(choices=["GradCAM"], value="GradCAM", label="activation map type")
-        )
-        return res
+        ]
+        return tuple(language_res + sliders)
 
     
 
@@ -362,7 +379,8 @@ with gr.Blocks() as demo:
             response_type = gr.Dropdown(choices=["Visualization only"], value="Visualization only", label="response_type")
             focus = gr.Dropdown(choices=["Visual Encoder"], value="Visual Encoder", label="focus")
             activation_map_method = gr.Dropdown(choices=["GradCAM"], value="GradCAM", label="visualization type")
-            visual_pooling_method = gr.Dropdown(choices=["CLS", "max", "avg"], value="CLS", label="visual pooling method")
+            # activation_function = gr.Dropdown(choices=["softmax", "sigmoid"], value="softmax", label="activation function")
+            visual_method = gr.Dropdown(choices=["CLS", "max", "avg"], value="CLS", label="visual pooling method")
             
 
             visualization_layers_min = gr.Slider(minimum=1, maximum=12, value=12, step=1, label="visualization layers min")
@@ -377,10 +395,11 @@ with gr.Blocks() as demo:
             inputs=model_selector, 
             outputs=[
                 response_type,
-                visualization_layers_min,
-                visualization_layers_max,
                 focus,
-                activation_map_method
+                activation_map_method,
+                visual_method,
+                visualization_layers_min,
+                visualization_layers_max
             ]
         )
         
@@ -492,7 +511,7 @@ with gr.Blocks() as demo:
         
     understanding_button.click(
         multimodal_understanding,
-        inputs=[model_selector, activation_map_method, visual_pooling_method, image_input, question_input, und_seed_input, top_p, temperature, target_token_idx, 
+        inputs=[model_selector, activation_map_method, visual_method, image_input, question_input, und_seed_input, top_p, temperature, target_token_idx, 
                 visualization_layers_min, visualization_layers_max, focus, response_type, chart_type],
         outputs=[understanding_output, activation_map_output, understanding_target_token_decoded_output]
     )

demo/model_utils.py

@@ -204,7 +204,7 @@ class ChartGemma_Utils(Model_Utils):
         self.vl_gpt = PaliGemmaForConditionalGeneration.from_pretrained(
             model_path, 
             torch_dtype=torch.float16,  
-            attn_implementation="sdpa", 
+            attn_implementation="eager", 
             output_attentions=True
         )
         self.vl_gpt, self.dtype, self.cuda_device = set_dtype_device(self.vl_gpt)

demo/modified_attn.py

@@ -0,0 +1,221 @@
+import torch
+import sys
+import typing
+import typing_extensions
+from torch import nn
+from typing import Callable, List, Optional, Tuple, Union
+from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+from transformers.cache_utils import Cache
+from transformers.models.llama.configuration_llama import LlamaConfig
+from transformers.models.llama.modeling_llama import LlamaAttention
+from transformers.models.gemma.modeling_gemma import GemmaAttention
+
+# from transformers.models.paligemma.modeling_paligemma import 
+
+if sys.version_info >= (3, 11):
+    Unpack = typing.Unpack
+else:
+    Unpack = typing_extensions.Unpack
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+    """Applies Rotary Position Embedding to the query and key tensors.
+
+    Args:
+        q (`torch.Tensor`): The query tensor.
+        k (`torch.Tensor`): The key tensor.
+        cos (`torch.Tensor`): The cosine part of the rotary embedding.
+        sin (`torch.Tensor`): The sine part of the rotary embedding.
+        position_ids (`torch.Tensor`, *optional*):
+            Deprecated and unused.
+        unsqueeze_dim (`int`, *optional*, defaults to 1):
+            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+    Returns:
+        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+    """
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: Optional[torch.Tensor],
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs,
+):
+    key_states = repeat_kv(key, module.num_key_value_groups)
+    value_states = repeat_kv(value, module.num_key_value_groups)
+
+    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+        attn_weights = attn_weights + causal_mask
+
+    attn_sigmoid_weights = nn.functional.sigmoid(attn_weights).to(query.dtype)
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+    attn_output = torch.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    return attn_output, attn_weights, attn_sigmoid_weights
+
+
+class ModifiedLlamaAttention(nn.Module):
+    def __init__(self, llama_attention_old: LlamaAttention):
+        super().__init__()
+        self.config = llama_attention_old.config
+        self.layer_idx = llama_attention_old.layer_idx
+        self.head_dim = llama_attention_old.head_dim
+        self.num_key_value_groups = llama_attention_old.num_key_value_groups
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = llama_attention_old.attention_dropout
+        self.is_causal = True
+
+        self.q_proj = llama_attention_old.q_proj
+        self.k_proj = llama_attention_old.k_proj
+        self.v_proj = llama_attention_old.v_proj
+        self.o_proj = llama_attention_old.o_proj
+
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: Tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor],
+        past_key_value: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        if past_key_value is not None:
+            # sin and cos are specific to RoPE models; cache_position needed for the static cache
+            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+        attention_interface: Callable = eager_attention_forward
+        
+
+        attn_output, attn_weights, attn_sigmoid_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            **kwargs,
+        )
+
+        self.attn_sigmoid_weights = attn_sigmoid_weights
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+
+
+
+
+
+
+
+
+
+class ModifiedGemmaAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, gemma_attention_old: GemmaAttention):
+        super().__init__()
+        self.config = gemma_attention_old.config
+        self.layer_idx = gemma_attention_old.layer_idx
+        self.head_dim = gemma_attention_old.head_dim
+        self.num_key_value_groups = gemma_attention_old.num_key_value_groups
+        self.scaling = gemma_attention_old.scaling
+        self.attention_dropout = gemma_attention_old.attention_dropout
+        self.is_causal = True
+
+        self.q_proj = gemma_attention_old.q_proj
+        self.k_proj = gemma_attention_old.k_proj
+        self.v_proj = gemma_attention_old.v_proj
+        self.o_proj = gemma_attention_old.o_proj
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: Tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor],
+        past_key_value: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        if past_key_value is not None:
+            # sin and cos are specific to RoPE models; cache_position needed for the static cache
+            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+        attention_interface: Callable = eager_attention_forward
+
+        attn_output, attn_weights, attn_sigmoid_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            **kwargs,
+        )
+
+        self.attn_sigmoid_weights = attn_sigmoid_weights
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+

demo/visualization.py

@@ -8,7 +8,7 @@ from PIL import Image
 from torch import nn
 import spaces
 from demo.modify_llama import *
-
+from demo.modified_attn import ModifiedLlamaAttention
 
 class Visualization:
     def __init__(self, model, register=True):
@@ -25,6 +25,7 @@ class Visualization:
             self.hooks.append(layer.register_backward_hook(self._backward_hook))
 
     def _forward_hook(self, module, input, output):
+        print("forward_hook: self_attn_input: ", input)
         self.activations.append(output)
 
     def _backward_hook(self, module, grad_in, grad_out):
@@ -41,6 +42,9 @@ class Visualization:
             layer.get_attn_map = types.MethodType(get_attn_map, layer)
 
     def _forward_activate_hooks(self, module, input, output):
+        print("forward_activate_hool: module: ", module)
+        print("forward_activate_hook: self_attn_input: ", input)
+
         attn_output, attn_weights = output  # Unpack outputs
         print("attn_output shape:", attn_output.shape)
         print("attn_weights shape:", attn_weights.shape)
@@ -231,15 +235,15 @@ class VisualizationClip(Visualization):
         super().__init__(model)
 
     @spaces.GPU(duration=120)
-    def forward_backward(self, input_tensor, visual_pooling_method, target_token_idx):
+    def forward_backward(self, input_tensor, visual_method, target_token_idx):
         output_full = self.model(**input_tensor)
 
         if target_token_idx is None:
             target_token_idx = torch.argmax(output_full.logits, dim=1).item()
 
-        if visual_pooling_method == "CLS":
+        if visual_method == "CLS":
             output = output_full.image_embeds
-        elif visual_pooling_method == "avg":
+        elif visual_method == "avg":
             output = self.model.visual_projection(output_full.vision_model_output.last_hidden_state).mean(dim=1)
         else:
             output, _ = self.model.visual_projection(output_full.vision_model_output.last_hidden_state).max(dim=1)
@@ -250,11 +254,11 @@ class VisualizationClip(Visualization):
         
         
     @spaces.GPU(duration=120)
-    def generate_cam(self, input_tensor, target_token_idx=None, visual_pooling_method="CLS"):
+    def generate_cam(self, input_tensor, target_token_idx=None, visual_method="CLS"):
         """ Generates Grad-CAM heatmap for ViT. """
         self.setup_grads()
         # Forward Backward pass
-        output_full = self.forward_backward(input_tensor, visual_pooling_method, target_token_idx)
+        output_full = self.forward_backward(input_tensor, visual_method, target_token_idx)
 
         cam_sum = self.grad_cam_vis()
         cam_sum, grid_size = self.process(cam_sum)
@@ -291,34 +295,33 @@ class VisualizationJanus(Visualization):
         self._modify_layers()
         self._register_hooks_activations()
 
-    def forward_backward(self, input_tensor, tokenizer, temperature, top_p, target_token_idx=None, visual_pooling_method="CLS", focus="Visual Encoder"):
+    def forward_backward(self, input_tensor, tokenizer, temperature, top_p, target_token_idx=None, visual_method="softmax", focus="Visual Encoder"):
         # Forward
         image_embeddings, inputs_embeddings, outputs = self.model(input_tensor, tokenizer, temperature, top_p)
         input_ids = input_tensor.input_ids
-
+        start_idx = 620
+        self.model.zero_grad()
         if focus == "Visual Encoder":
-            
-            start_idx = 620
-            self.model.zero_grad()
-
             loss = outputs.logits.max(dim=-1).values[0, start_idx + target_token_idx]
             loss.backward()
         
         elif focus == "Language Model":
-            self.model.zero_grad()
-            loss = outputs.logits.max(dim=-1).values.sum()
+            if target_token_idx == -1:
+                loss = outputs.logits.max(dim=-1).values.sum()
+            else:
+                loss = outputs.logits.max(dim=-1).values[0, start_idx + target_token_idx]
             loss.backward()
             
-            self.activations = [layer.get_attn_map() for layer in self.target_layers]
+            self.activations = self.activations = [layer.attn_sigmoid_weights for layer in self.target_layers] if visual_method == "sigmoid" else [layer.get_attn_map() for layer in self.target_layers]
             self.gradients = [layer.get_attn_gradients() for layer in self.target_layers]
     
     @spaces.GPU(duration=120)
-    def generate_cam(self, input_tensor, tokenizer, temperature, top_p, target_token_idx=None, visual_pooling_method="CLS", focus="Visual Encoder"):
+    def generate_cam(self, input_tensor, tokenizer, temperature, top_p, target_token_idx=None, visual_method="softmax", focus="Visual Encoder"):
         
         self.setup_grads()
 
         # Forward Backward pass
-        self.forward_backward(input_tensor, tokenizer, temperature, top_p, target_token_idx, visual_pooling_method, focus)
+        self.forward_backward(input_tensor, tokenizer, temperature, top_p, target_token_idx, visual_method, focus)
         
         start_idx = 620
         if focus == "Visual Encoder":
@@ -365,7 +368,7 @@ class VisualizationLLaVA(Visualization):
         self.gradients = [layer.get_attn_gradients() for layer in self.target_layers]
 
     @spaces.GPU(duration=120)
-    def generate_cam(self, inputs, tokenizer, temperature, top_p, target_token_idx=None, visual_pooling_method="CLS", focus="Visual Encoder"):
+    def generate_cam(self, inputs, tokenizer, temperature, top_p, target_token_idx=None, visual_method="softmax", focus="Visual Encoder"):
         
         self.setup_grads()
         self.forward_backward(inputs)
@@ -401,7 +404,7 @@ class VisualizationChartGemma(Visualization):
         self._modify_layers()
         self._register_hooks_activations()
     
-    def forward_backward(self, inputs, focus, start_idx, target_token_idx):
+    def forward_backward(self, inputs, focus, start_idx, target_token_idx, visual_method="softmax"):
         outputs_raw = self.model(**inputs, output_hidden_states=True)
         if focus == "Visual Encoder":
             
@@ -417,11 +420,11 @@ class VisualizationChartGemma(Visualization):
             else:
                 loss = outputs_raw.logits.max(dim=-1).values[0, start_idx + target_token_idx]
             loss.backward()
-            self.activations = [layer.get_attn_map() for layer in self.target_layers]
+            self.activations = [layer.attn_sigmoid_weights for layer in self.target_layers] if visual_method == "sigmoid" else [layer.get_attn_map() for layer in self.target_layers]
             self.gradients = [layer.get_attn_gradients() for layer in self.target_layers]
     
     @spaces.GPU(duration=120)
-    def generate_cam(self, inputs, tokenizer, temperature, top_p, target_token_idx=None, visual_pooling_method="CLS", focus="Visual Encoder"):
+    def generate_cam(self, inputs, tokenizer, temperature, top_p, target_token_idx=None, visual_method="softmax", focus="Visual Encoder"):
         
         # Forward pass
         self.setup_grads()
@@ -439,7 +442,7 @@ class VisualizationChartGemma(Visualization):
         start_idx = last + 1
 
 
-        self.forward_backward(inputs, focus, start_idx, target_token_idx)
+        self.forward_backward(inputs, focus, start_idx, target_token_idx, visual_method)
         if focus == "Visual Encoder":
             
             cam_sum = self.grad_cam_vis()