add accumulate methods: Sum or Mult

app.py

@@ -56,7 +56,7 @@ def multimodal_understanding(model_type,
                              activation_map_method, 
                              visual_method, 
                              image, question, seed, top_p, temperature, target_token_idx,
-                             visualization_layer_min, visualization_layer_max, focus, response_type, chart_type):
+                             visualization_layer_min, visualization_layer_max, focus, response_type, chart_type, accumulate_method):
     # Clear CUDA cache before generating
     gc.collect()
     if torch.cuda.is_available():
@@ -160,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_method, focus)
+                        cam_tensors, grid_size, start = gradcam.generate_cam(prepare_inputs, tokenizer, temperature, top_p, i, visual_method, focus, accumulate_method)
                         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])
@@ -168,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_method, focus)
+                cam_tensors, grid_size, start = gradcam.generate_cam(prepare_inputs, tokenizer, temperature, top_p, target_token_idx, visual_method, focus, accumulate_method)
                 if target_token_idx != -1:
                     input_text_decoded = input_ids_decoded[start + target_token_idx]
                     for i, cam_tensor in enumerate(cam_tensors):
@@ -379,7 +379,7 @@ 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")
-            # activation_function = gr.Dropdown(choices=["softmax", "sigmoid"], value="softmax", label="activation function")
+            accumulate_method = gr.Dropdown(choices=["sum", "mult"], value="sum", label="layers accumulate method")
             visual_method = gr.Dropdown(choices=["CLS", "max", "avg"], value="CLS", label="visual pooling method")
             
 
@@ -512,7 +512,7 @@ with gr.Blocks() as demo:
     understanding_button.click(
         multimodal_understanding,
         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],
+                visualization_layers_min, visualization_layers_max, focus, response_type, chart_type, accumulate_method],
         outputs=[understanding_output, activation_map_output, understanding_target_token_decoded_output]
     )
     

demo/visualization.py

@@ -196,7 +196,7 @@ class Visualization:
             cam_sum_lst.append(cam_sum)
         return cam_sum_lst, grid_size
             
-    def process_multiple_withsum(self, cams, start_idx, images_seq_mask, normalize=False):
+    def process_multiple_acc(self, cams, start_idx, images_seq_mask, normalize=False, accumulate_method="sum"):
         cam_sum_lst = []
         for i in range(start_idx, cams[0].shape[1]):
             cam_sum = None
@@ -217,7 +217,10 @@ class Visualization:
                 if cam_sum == None:
                     cam_sum = cam_reshaped
                 else:
-                    cam_sum += cam_reshaped
+                    if accumulate_method == "sum":
+                        cam_sum += cam_reshaped
+                    elif accumulate_method == "mult":
+                        cam_sum *= cam_reshaped + 1
 
             cam_sum = (cam_sum - cam_sum.min()) / (cam_sum.max() - cam_sum.min())
             cam_sum_lst.append(cam_sum)
@@ -316,7 +319,7 @@ class VisualizationJanus(Visualization):
             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_method="softmax", focus="Visual Encoder"):
+    def generate_cam(self, input_tensor, tokenizer, temperature, top_p, target_token_idx=None, visual_method="softmax", focus="Visual Encoder", accumulate_method="sum"):
         
         self.setup_grads()
 
@@ -368,7 +371,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_method="softmax", focus="Visual Encoder"):
+    def generate_cam(self, inputs, tokenizer, temperature, top_p, target_token_idx=None, visual_method="softmax", focus="Visual Encoder", accumulate_method="sum"):
         
         self.setup_grads()
         self.forward_backward(inputs)
@@ -388,7 +391,7 @@ class VisualizationLLaVA(Visualization):
         # Aggregate activations and gradients from ALL layers
         start_idx = last + 1
         cams = self.attn_guided_cam()
-        cam_sum_lst, grid_size = self.process_multiple_withsum(cams, start_idx, images_seq_mask)
+        cam_sum_lst, grid_size = self.process_multiple_acc(cams, start_idx, images_seq_mask, accumulate_method=accumulate_method)
 
         return cam_sum_lst, grid_size, start_idx
 
@@ -424,7 +427,7 @@ class VisualizationChartGemma(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_method="softmax", focus="Visual Encoder"):
+    def generate_cam(self, inputs, tokenizer, temperature, top_p, target_token_idx=None, visual_method="softmax", focus="Visual Encoder", accumulate_method="sum"):
         
         # Forward pass
         self.setup_grads()
@@ -453,7 +456,7 @@ class VisualizationChartGemma(Visualization):
         elif focus == "Language Model":
 
             cams = self.attn_guided_cam()
-            cam_sum_lst, grid_size = self.process_multiple_withsum(cams, start_idx, images_seq_mask)
+            cam_sum_lst, grid_size = self.process_multiple_acc(cams, start_idx, images_seq_mask, accumulate_method=accumulate_method)
 
             # cams shape: [layers, 1, seq_len, seq_len]