add ChartGemma

app.py

@@ -3,8 +3,8 @@ import torch
 from transformers import AutoConfig, AutoModelForCausalLM
 from janus.models import MultiModalityCausalLM, VLChatProcessor
 from janus.utils.io import load_pil_images
-from demo.cam import generate_gradcam, AttentionGuidedCAMJanus, AttentionGuidedCAMClip, AttentionGuidedCAMLLaVA
-from demo.model_utils import Clip_Utils, Janus_Utils, LLaVA_Utils, add_title_to_image
+from demo.cam import generate_gradcam, AttentionGuidedCAMJanus, AttentionGuidedCAMClip, AttentionGuidedCAMChartGemma, AttentionGuidedCAMLLaVA
+from demo.model_utils import Clip_Utils, Janus_Utils, LLaVA_Utils, ChartGemma_Utils, add_title_to_image
 
 import numpy as np
 import matplotlib.pyplot as plt
@@ -22,7 +22,8 @@ clip_utils = Clip_Utils()
 clip_utils.init_Clip()
 model_utils, vl_gpt, tokenizer = None, None, None
 model_name = "Clip"
-
+language_model_max_layer = 24
+language_model_best_layer = 8
 
 def clean():
     global model_utils, vl_gpt, tokenizer, clip_utils
@@ -109,7 +110,12 @@ def multimodal_understanding(model_type,
 
         input_ids = prepare_inputs.input_ids[0].cpu().tolist()
         input_ids_decoded = [tokenizer.decode([input_ids[i]]) for i in range(len(input_ids))]
-        start=620 if model_name.split('-')[0] == "Janus" else 512
+        if model_name.split('-')[0] == "Janus":
+            start = 620 
+        elif model_name.split('-')[0] == "ChartGemma":
+            start = 1024
+        else: 
+            start = 512
 
         if saliency_map_method == "GradCAM":
             # target_layers = vl_gpt.vision_model.vision_tower.blocks
@@ -127,8 +133,13 @@ def multimodal_understanding(model_type,
                 gradcam = AttentionGuidedCAMJanus(vl_gpt, target_layers)
             elif model_name.split('-')[0] == "LLaVA":
                 gradcam = AttentionGuidedCAMLLaVA(vl_gpt, target_layers)
+            elif model_name.split('-')[0] == "ChartGemma":
+                gradcam = AttentionGuidedCAMChartGemma(vl_gpt, target_layers)
+
             cam_tensors, grid_size = gradcam.generate_cam(prepare_inputs, tokenizer, temperature, top_p, target_token_idx, visual_pooling_method, focus)
             gradcam.remove_hooks()
+
+            
             if focus == "Visual Encoder":
                 cam_grid = cam_tensors.reshape(grid_size, grid_size)
                 cam = [generate_gradcam(cam_grid, image)]
@@ -144,7 +155,7 @@ def multimodal_understanding(model_type,
                 else:
                     cam = []
                     for i, cam_tensor in enumerate(cam_tensors):
-                        cam_grid = cam_tensor.reshape(24, 24)
+                        cam_grid = cam_tensor.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])
 
@@ -158,7 +169,7 @@ def multimodal_understanding(model_type,
 # Gradio interface
 
 def model_slider_change(model_type):
-    global model_utils, vl_gpt, tokenizer, clip_utils, model_name
+    global model_utils, vl_gpt, tokenizer, clip_utils, model_name, language_model_max_layer, language_model_best_layer
     model_name = model_type
     if model_type == "Clip":
         clean()
@@ -179,6 +190,8 @@ def model_slider_change(model_type):
         set_seed()
         model_utils = Janus_Utils()
         vl_gpt, tokenizer = model_utils.init_Janus(model_type.split('-')[-1])
+        language_model_max_layer = 24
+        language_model_best_layer = 8
 
         res = (
             gr.Dropdown(choices=["Visualization only", "answer + visualization"], value="Visualization only", label="response_type"),
@@ -195,6 +208,8 @@ def model_slider_change(model_type):
         set_seed()
         model_utils = LLaVA_Utils()
         vl_gpt, tokenizer = model_utils.init_LLaVA()
+        language_model_max_layer = 24
+        language_model_best_layer = 8
 
         res = (
             gr.Dropdown(choices=["Visualization only", "answer + visualization"], value="Visualization only", label="response_type"),
@@ -204,9 +219,29 @@ def model_slider_change(model_type):
             gr.Dropdown(choices=["GradCAM"], value="GradCAM", label="saliency map type")
         )
         return res
+    
+    elif model_type.split('-')[0] == "ChartGemma":
+        clean()
+        set_seed()
+        model_utils = ChartGemma_Utils()
+        vl_gpt, tokenizer = model_utils.init_ChartGemma()
+        language_model_max_layer = 18
+        language_model_best_layer = 12
+
+        res = (
+            gr.Dropdown(choices=["Visualization only", "answer + visualization"], value="Visualization only", label="response_type"),
+            gr.Slider(minimum=1, maximum=18, value=12, step=1, label="visualization layers min"),
+            gr.Slider(minimum=1, maximum=18, value=12, step=1, label="visualization layers max"),
+            gr.Dropdown(choices=["Language Model"], value="Language Model", label="focus"),
+            gr.Dropdown(choices=["GradCAM"], value="GradCAM", label="saliency map type")
+        )
+        return res
+
+    
+
 
 def focus_change(focus):
-    global model_name
+    global model_name, language_model_max_layer
     if model_name == "Clip":
         res = (
                 gr.Dropdown(choices=["GradCAM"], value="GradCAM", label="saliency map type"),
@@ -219,15 +254,15 @@ def focus_change(focus):
         if response_type.value == "answer + visualization":
             res = (
                 gr.Dropdown(choices=["GradCAM"], value="GradCAM", label="saliency map type"),
-                gr.Slider(minimum=1, maximum=24, value=8, step=1, label="visualization layers min"), 
-                gr.Slider(minimum=1, maximum=24, value=8, step=1, label="visualization layers max")
+                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")
             )
             return res
         else:
             res = (
                 gr.Dropdown(choices=["GradCAM"], value="GradCAM", label="saliency map type"),
-                gr.Slider(minimum=1, maximum=24, value=8, step=1, label="visualization layers min"), 
-                gr.Slider(minimum=1, maximum=24, value=8, step=1, label="visualization layers max")
+                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")
             )
             return res
 
@@ -251,7 +286,7 @@ with gr.Blocks() as demo:
             saliency_map_output = gr.Gallery(label="Saliency Map", height=300, columns=1)
 
         with gr.Column():
-            model_selector = gr.Dropdown(choices=["Clip", "Janus-1B", "Janus-7B", "LLaVA-1.5-7B"], value="Clip", label="model")
+            model_selector = gr.Dropdown(choices=["Clip", "ChartGemma-2B", "Janus-1B", "Janus-7B", "LLaVA-1.5-7B"], value="Clip", label="model")
             response_type = gr.Dropdown(choices=["Visualization only"], value="Visualization only", label="response_type")
             focus = gr.Dropdown(choices=["Visual Encoder"], value="Visual Encoder", label="focus")
             saliency_map_method = gr.Dropdown(choices=["GradCAM"], value="GradCAM", label="saliency map type")

demo/cam.py

@@ -229,8 +229,8 @@ class AttentionGuidedCAMJanus(AttentionGuidedCAM):
 
 
         elif focus == "Language Model":
-            loss = self.target_layers[-1].attention_map.sum()
             self.model.zero_grad()
+            loss = outputs.logits.max(dim=-1).values.sum()
             loss.backward()
             
             self.activations = [layer.get_attn_map() for layer in self.target_layers]
@@ -429,6 +429,193 @@ class AttentionGuidedCAMLLaVA(AttentionGuidedCAM):
 
 
 
+
+
+
+
+
+
+
+
+
+
+
+
+class AttentionGuidedCAMChartGemma(AttentionGuidedCAM):
+    def __init__(self, model, target_layers):
+        self.target_layers = target_layers
+        super().__init__(model)
+        self._modify_layers()
+        self._register_hooks_activations()
+    
+    def _modify_layers(self):
+        for layer in self.target_layers:
+            setattr(layer, "attn_gradients", None)
+            setattr(layer, "attention_map", None)
+
+            layer.save_attn_gradients = types.MethodType(save_attn_gradients, layer)
+            layer.get_attn_gradients = types.MethodType(get_attn_gradients, layer)
+            layer.save_attn_map = types.MethodType(save_attn_map, layer)
+            layer.get_attn_map = types.MethodType(get_attn_map, layer)
+
+    def _forward_activate_hooks(self, module, input, output):
+        attn_output, attn_weights = output  # Unpack outputs
+        print("attn_output shape:", attn_output.shape)
+        print("attn_weights shape:", attn_weights.shape)
+        module.save_attn_map(attn_weights)
+        attn_weights.register_hook(module.save_attn_gradients)
+
+    def _register_hooks_activations(self):
+        for layer in self.target_layers:
+            if hasattr(layer, "q_proj"): # is an attention layer
+                self.hooks.append(layer.register_forward_hook(self._forward_activate_hooks))
+    
+    @spaces.GPU(duration=120)
+    def generate_cam(self, inputs, tokenizer, temperature, top_p, class_idx=None, visual_pooling_method="CLS", focus="Visual Encoder"):
+        """ Generates Grad-CAM heatmap for ViT. """
+        
+        # Forward pass
+        outputs_raw = self.model(**inputs)
+
+        image_embeddings = outputs_raw.image_hidden_states
+        inputs_embeddings = self.model.get_input_embeddings()(inputs['input_ids'])
+
+        # Pooling
+        image_embeddings_pooled = image_embeddings.mean(dim=1)
+
+        inputs_embeddings_pooled = inputs_embeddings.mean(dim=1) # end of image: 618
+        # inputs_embeddings_pooled = inputs_embeddings[
+        #     torch.arange(inputs_embeddings.shape[0], device=inputs_embeddings.device),
+        #     input_ids.to(dtype=torch.int, device=inputs_embeddings.device).argmax(dim=-1),
+        # ]
+
+
+        # Backpropagate to get gradients
+        # image_embeddings_pooled.backward(inputs_embeddings_pooled, retain_graph=True)
+        # similarity = F.cosine_similarity(image_embeddings_mean, inputs_embeddings_mean, dim=-1)
+        # similarity.backward()
+        self.model.zero_grad()
+        print(outputs_raw)
+        # loss = self.target_layers[-1].attention_map.sum()
+        loss = outputs_raw.logits.max(dim=-1).values.sum()
+        loss.backward()
+
+        # get image masks
+        image_mask = []
+        last = 0
+        for i in range(inputs["input_ids"].shape[1]):
+            decoded_token = tokenizer.decode(inputs["input_ids"][0][i].item())
+            print(decoded_token)
+            if (decoded_token == "<image>"):
+                image_mask.append(True)
+                last = i
+            else:
+                image_mask.append(False)
+
+
+        # Aggregate activations and gradients from ALL layers
+        self.activations = [layer.get_attn_map() for layer in self.target_layers]
+        self.gradients = [layer.get_attn_gradients() for layer in self.target_layers]
+        cam_sum = None
+        # Ver 1
+        # for act, grad in zip(self.activations, self.gradients):
+        #     # act = torch.sigmoid(act)
+        #     print("act:", act)
+        #     print(len(act))
+        #     print("act_shape:", act.shape)
+        #     # print("act1_shape:", act[1].shape)
+            
+        #     act = F.relu(act.mean(dim=1))
+
+
+        #     # Compute mean of gradients
+        #     print("grad:", grad)
+        #     print(len(grad))
+        #     print("grad_shape:", grad.shape)
+        #     grad_weights = grad.mean(dim=1)
+
+        #     print("act shape", act.shape)
+        #     print("grad_weights shape", grad_weights.shape)
+
+        #     cam = act * grad_weights
+        #     # cam = act
+        #     print(cam.shape)
+
+        #     # Sum across all layers
+        #     if cam_sum is None:
+        #         cam_sum = cam
+        #     else:
+        #         cam_sum += cam  
+
+        # Ver 2
+        for act, grad in zip(self.activations, self.gradients):
+
+            print("act shape", act.shape)
+            print("grad shape", grad.shape)
+
+            act = F.relu(act)
+            grad = F.relu(grad)
+
+
+            cam = act * grad # shape: [1, heads, seq_len, seq_len]
+            cam = cam.sum(dim=1) # shape: [1, seq_len, seq_len]
+
+            # Sum across all layers
+            if cam_sum is None:
+                cam_sum = cam
+            else:
+                cam_sum += cam 
+
+        cam_sum = F.relu(cam_sum)
+        cam_sum = cam_sum.to(torch.float32)
+
+        # thresholding
+        # percentile = torch.quantile(cam_sum, 0.4)  # Adjust threshold dynamically
+        # cam_sum[cam_sum < percentile] = 0
+
+        # Reshape
+        # if visual_pooling_method == "CLS":
+        # cam_sum = cam_sum[0, 1:]
+
+        # cam_sum shape: [1, seq_len, seq_len]
+        cam_sum_lst = []
+        cam_sum_raw = cam_sum
+        start_idx = 1024
+        for i in range(start_idx, cam_sum_raw.shape[1]):
+            cam_sum = cam_sum_raw[0, i, :] # shape: [1: seq_len]
+            # cam_sum_min = cam_sum.min()
+            # cam_sum_max = cam_sum.max()
+            # cam_sum = (cam_sum - cam_sum_min) / (cam_sum_max - cam_sum_min)
+            cam_sum = cam_sum[image_mask].unsqueeze(0) # shape: [1, 1024]
+            print("cam_sum shape: ", cam_sum.shape)
+            num_patches = cam_sum.shape[-1]  # Last dimension of CAM output
+            grid_size = int(num_patches ** 0.5)
+            print(f"Detected grid size: {grid_size}x{grid_size}")
+
+            # Fix the reshaping step dynamically
+            
+            cam_sum = cam_sum.view(grid_size, grid_size)
+            cam_sum = (cam_sum - cam_sum.min()) / (cam_sum.max() - cam_sum.min())
+            cam_sum_lst.append(cam_sum)
+
+
+        return cam_sum_lst, grid_size
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
 def generate_gradcam(
     cam, 
     image,

demo/model_utils.py

@@ -2,7 +2,7 @@ import torch
 import numpy as np
 import spaces
 from PIL import Image, ImageDraw, ImageFont
-from transformers import AutoConfig, AutoModelForCausalLM, LlavaForConditionalGeneration, AutoProcessor
+from transformers import AutoConfig, AutoModelForCausalLM, LlavaForConditionalGeneration, AutoProcessor, PaliGemmaForConditionalGeneration
 from transformers import CLIPProcessor, CLIPModel
 from janus.models import MultiModalityCausalLM, VLChatProcessor
 
@@ -170,6 +170,61 @@ class LLaVA_Utils(Model_Utils):
         )
 
         return outputs
+    
+
+
+
+
+class ChartGemma_Utils(Model_Utils):
+    def __init__(self):
+        super().__init__()
+
+    def init_ChartGemma(self):
+
+        model_path = "ahmed-masry/chartgemma"
+        
+
+        self.vl_gpt = PaliGemmaForConditionalGeneration.from_pretrained(
+            model_path, 
+            torch_dtype=torch.float16,  
+            attn_implementation="eager", 
+            output_attentions=True
+        )
+        self.vl_gpt, self.dtype, self.cuda_device = set_dtype_device(self.vl_gpt)
+        self.processor = AutoProcessor.from_pretrained(model_path)
+        self.tokenizer = self.processor.tokenizer
+        
+        return self.vl_gpt, self.tokenizer
+    
+    @spaces.GPU(duration=120)
+    def prepare_inputs(self, question, image):
+
+        pil_image = Image.fromarray(image)
+        prepare_inputs = self.processor(
+            images=pil_image, text=[question], return_tensors="pt"
+        ).to(self.cuda_device, dtype=self.dtype)
+
+        return prepare_inputs
+    
+    @spaces.GPU(duration=120)
+    def generate_inputs_embeddings(self, prepare_inputs):
+        return self.vl_gpt.prepare_inputs_embeds(**prepare_inputs)
+    
+    @spaces.GPU(duration=120)
+    def generate_outputs(self, prepare_inputs, temperature, top_p):
+        
+        outputs = self.vl_gpt.generate(
+            **prepare_inputs,
+            max_new_tokens=512,
+            do_sample=False if temperature == 0 else True,
+            use_cache=True,
+            return_dict_in_generate=True,
+            output_attentions=True
+        )
+
+        return outputs
+
+
 
 
 def add_title_to_image(image, title, font_size=20):