renew app, formatted visualization

app-old.py

@@ -0,0 +1,501 @@
+import gradio as gr
+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, 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
+import gc
+import os
+import spaces
+from PIL import Image
+
+def set_seed(model_seed = 42):
+    torch.manual_seed(model_seed)
+    np.random.seed(model_seed)
+    torch.cuda.manual_seed(model_seed) if torch.cuda.is_available() else None
+
+set_seed()
+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
+vision_model_best_layer = 24
+
+def clean():
+    global model_utils, vl_gpt, tokenizer, clip_utils
+    # Move models to CPU first (prevents CUDA references)
+    if 'vl_gpt' in globals() and vl_gpt is not None:
+        vl_gpt.to("cpu")
+    if 'clip_utils' in globals() and clip_utils is not None:
+        del clip_utils
+
+    # Delete all references
+    del model_utils, vl_gpt, tokenizer
+    model_utils, vl_gpt, tokenizer, clip_utils = None, None, None, None
+    gc.collect()
+
+    # Empty CUDA cache
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.ipc_collect()  # Frees inter-process CUDA memory
+    
+    # Empty MacOS Metal backend (if using Apple Silicon)
+    if torch.backends.mps.is_available():
+        torch.mps.empty_cache()
+
+# Multimodal Understanding function
+@spaces.GPU(duration=120)
+def multimodal_understanding(model_type, 
+                             activation_map_method, 
+                             visual_pooling_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
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.ipc_collect()
+
+    # set seed
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+    torch.cuda.manual_seed(seed) if torch.cuda.is_available() else None
+
+    input_text_decoded = ""
+    answer = ""
+    if model_name == "Clip":
+        
+        inputs = clip_utils.prepare_inputs([question], image)
+
+
+        if activation_map_method == "GradCAM":
+            # Generate Grad-CAM
+            all_layers = [layer.layer_norm1 for layer in clip_utils.model.vision_model.encoder.layers]
+
+            if visualization_layer_min != visualization_layer_max:
+                target_layers = all_layers[visualization_layer_min-1 : visualization_layer_max-1]
+            else:
+                target_layers = [all_layers[visualization_layer_min-1]]
+            grad_cam = AttentionGuidedCAMClip(clip_utils.model, target_layers)
+            cam, outputs, grid_size = grad_cam.generate_cam(inputs, class_idx=0, visual_pooling_method=visual_pooling_method)
+            cam = cam.to("cpu")
+            cam = [generate_gradcam(cam, image, size=(224, 224))]
+            grad_cam.remove_hooks()
+            target_token_decoded = ""
+            
+    
+
+    else:
+        
+        for param in vl_gpt.parameters():
+            param.requires_grad = True
+
+
+        prepare_inputs = model_utils.prepare_inputs(question, image)
+
+        if response_type == "answer + visualization":
+            if model_name.split('-')[0] == "Janus":
+                inputs_embeds = model_utils.generate_inputs_embeddings(prepare_inputs)
+                outputs = model_utils.generate_outputs(inputs_embeds, prepare_inputs, temperature, top_p)
+            else:
+                outputs = model_utils.generate_outputs(prepare_inputs, temperature, top_p)
+
+            sequences = outputs.sequences.cpu().tolist()
+            answer = tokenizer.decode(sequences[0], skip_special_tokens=True)
+            attention_raw = outputs.attentions
+            print("answer generated")
+
+        input_ids = prepare_inputs.input_ids[0].cpu().tolist()
+        input_ids_decoded = [tokenizer.decode([input_ids[i]]) for i in range(len(input_ids))]
+
+        if activation_map_method == "GradCAM":
+            # target_layers = vl_gpt.vision_model.vision_tower.blocks
+            if focus == "Visual Encoder":
+                if model_name.split('-')[0] == "Janus":
+                    all_layers = [block.norm1 for block in vl_gpt.vision_model.vision_tower.blocks]
+                else:
+                    all_layers = [block.layer_norm1 for block in vl_gpt.vision_tower.vision_model.encoder.layers]
+            else:
+                all_layers = [layer.self_attn for layer in vl_gpt.language_model.model.layers]
+            
+            print("layer values:", visualization_layer_min, visualization_layer_max)
+            if visualization_layer_min != visualization_layer_max:
+                print("multi layers")
+                target_layers = all_layers[visualization_layer_min-1 : visualization_layer_max]
+            else:
+                print("single layer")
+                target_layers = [all_layers[visualization_layer_min-1]]
+            
+
+            if model_name.split('-')[0] == "Janus":
+                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)
+
+            start = 0
+            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_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])
+                    cam = [cam_i]
+                else:
+                    i = 0
+                    cam = []
+                    while start + i < len(input_ids_decoded):
+                        if model_name.split('-')[0] == "Janus":
+                            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, start = gradcam.generate_cam(prepare_inputs, tokenizer, temperature, top_p, i, visual_pooling_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])
+                        cam.append(cam_i)
+                        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)
+                if target_token_idx != -1:
+                    input_text_decoded = input_ids_decoded[start + target_token_idx]
+                    for i, cam_tensor in enumerate(cam_tensors):
+                        if i == target_token_idx:
+                            cam_grid = cam_tensor.reshape(grid_size, grid_size)
+                            cam_i = generate_gradcam(cam_grid, image)
+                            cam = [add_title_to_image(cam_i, input_text_decoded)]
+                            break
+                else:
+                    cam = []
+                    for i, cam_tensor in enumerate(cam_tensors):
+                        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])
+
+                        cam.append(cam_i)
+                     
+            gradcam.remove_hooks()
+                
+
+    # Collect Results
+    RESULTS_ROOT = "./results"
+    FILES_ROOT = f"{RESULTS_ROOT}/{model_name}/{focus}/{chart_type}/layer{visualization_layer_min}-{visualization_layer_max}"
+    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")
+            
+    with open(f"{FILES_ROOT}/input_text_decoded.txt", "w") as f:
+        f.write(input_text_decoded)
+        f.close()
+
+    with open(f"{FILES_ROOT}/answer.txt", "w") as f:
+        f.write(answer)
+        f.close()
+            
+
+
+    return answer, cam, input_text_decoded
+
+
+
+
+# Gradio interface
+
+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
+    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"),
+            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
+    elif model_type.split('-')[0] == "Janus":
+        
+        clean()
+        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="answer + visualization", label="response_type"),
+            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
+    
+    elif model_type.split('-')[0] == "LLaVA":
+        
+        clean()
+        set_seed()
+        model_utils = LLaVA_Utils()
+        version = model_type.split('-')[1]
+        vl_gpt, tokenizer = model_utils.init_LLaVA(version=version)
+        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"),
+            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
+    
+    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
+        vision_model_best_layer = 19
+        language_model_best_layer = 15
+
+        res = (
+            gr.Dropdown(choices=["Visualization only", "answer + visualization"], value="answer + visualization", label="response_type"),
+            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
+
+    
+
+
+def focus_change(focus):
+    global model_name, language_model_max_layer
+    if model_name == "Clip":
+        res = (
+                gr.Dropdown(choices=["GradCAM"], value="GradCAM", label="activation map type"),
+                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")
+            )
+        return res
+
+    if focus == "Language Model":
+        if response_type.value == "answer + visualization":
+            res = (
+                gr.Dropdown(choices=["GradCAM"], value="GradCAM", label="activation map type"),
+                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="activation map type"),
+                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:
+        if model_name.split('-')[0] == "ChartGemma":
+            res = (
+                gr.Dropdown(choices=["GradCAM"], value="GradCAM", label="activation map type"),
+                gr.Slider(minimum=1, maximum=26, value=vision_model_best_layer, step=1, label="visualization layers min"), 
+                gr.Slider(minimum=1, maximum=26, value=vision_model_best_layer, step=1, label="visualization layers max")
+            )
+            return res
+        
+        else:
+            res = (
+                gr.Dropdown(choices=["GradCAM"], value="GradCAM", label="activation map type"),
+                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")
+            )
+            return res
+
+
+
+
+
+with gr.Blocks() as demo:
+    gr.Markdown(value="# Multimodal Understanding")
+    
+    with gr.Row():
+        image_input = gr.Image(height=500, label="Image")
+        activation_map_output = gr.Gallery(label="Visualization", height=500, columns=1, preview=True)
+
+    with gr.Row():
+        chart_type = gr.Textbox(label="Chart Type")
+        understanding_output = gr.Textbox(label="Answer")
+
+    with gr.Row():
+
+        with gr.Column():
+            model_selector = gr.Dropdown(choices=["Clip", "ChartGemma-3B", "Janus-Pro-1B", "Janus-Pro-7B", "LLaVA-1.5-7B"], value="Clip", label="model")
+            question_input = gr.Textbox(label="Input Prompt")
+            und_seed_input = gr.Number(label="Seed", precision=0, value=42)
+            top_p = gr.Slider(minimum=0, maximum=1, value=0.95, step=0.05, label="top_p")
+            temperature = gr.Slider(minimum=0, maximum=1, value=0.1, step=0.05, label="temperature")
+            target_token_idx = gr.Number(label="target_token_idx (-1 means all)", precision=0, value=-1)
+
+        
+        with gr.Column():
+            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")
+            
+
+            visualization_layers_min = gr.Slider(minimum=1, maximum=12, value=12, step=1, label="visualization layers min")
+            visualization_layers_max = gr.Slider(minimum=1, maximum=12, value=12, step=1, label="visualization layers max")
+
+        
+        
+
+
+        model_selector.change(
+            fn=model_slider_change, 
+            inputs=model_selector, 
+            outputs=[
+                response_type,
+                visualization_layers_min,
+                visualization_layers_max,
+                focus,
+                activation_map_method
+            ]
+        )
+        
+        focus.change(
+            fn = focus_change,
+            inputs = focus,
+            outputs=[
+                activation_map_method,
+                visualization_layers_min,
+                visualization_layers_max,
+            ]
+        )
+
+        # response_type.change(
+        #     fn = response_type_change,
+        #     inputs = response_type,
+        #     outputs = [activation_map_method]
+        # )
+
+        
+
+    understanding_button = gr.Button("Submit")
+    
+    understanding_target_token_decoded_output = gr.Textbox(label="Target Token Decoded")
+
+
+    examples_inpainting = gr.Examples(
+        label="Multimodal Understanding examples",
+        examples=[
+
+            [
+                "LineChart",
+                "What was the price of a barrel of oil in February 2020?", 
+                "images/LineChart.png" 
+            ],
+
+            [
+                "BarChart",
+                "What is the average internet speed in Japan?",
+                "images/BarChart.png"
+            ],
+
+            [
+                "StackedBar",
+                "What is the cost of peanuts in Seoul?", 
+                "images/StackedBar.png"
+            ],
+
+            [
+                "100%StackedBar",
+                "Which country has the lowest proportion of Gold medals?", 
+                "images/Stacked100.png"
+            ],
+            
+            [
+                "PieChart",
+                "What is the approximate global smartphone market share of Samsung?",
+                "images/PieChart.png"
+            ],
+            
+            [
+                "Histogram",
+                "What distance have customers traveled in the taxi the most?", 
+                "images/Histogram.png"
+            ],
+
+            [
+                "Scatterplot",
+                "True/False: There is a negative linear relationship between the height and the weight of the 85 males.", 
+                "images/Scatterplot.png"
+            ],
+
+            [
+                "AreaChart",
+                "What was the average price of pount of coffee beans in October 2019?",
+                "images/AreaChart.png"
+            ],
+
+            [
+                "StackedArea",
+                "What was the ratio of girls named 'Isla' to girls named 'Amelia' in 2012 in the UK?", 
+                "images/StackedArea.png"
+            ],
+
+            [
+                "BubbleChart",
+                "Which city's metro system has the largest number of stations?", 
+                "images/BubbleChart.png"
+            ],
+
+            [ 
+                "Choropleth",
+                "True/False: In 2020, the unemployment rate for Washington (WA) was higher than that of Wisconsin (WI).", 
+                "images/Choropleth_New.png"
+            ],
+
+            [
+                "TreeMap",
+                "True/False: eBay is nested in the Software category.", 
+                "images/TreeMap.png"
+            ]
+            
+        ],
+        inputs=[chart_type, question_input, image_input],
+    )
+    
+
+
+        
+    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, 
+                visualization_layers_min, visualization_layers_max, focus, response_type, chart_type],
+        outputs=[understanding_output, activation_map_output, understanding_target_token_decoded_output]
+    )
+    
+demo.launch(share=True)
+# demo.queue(concurrency_count=1, max_size=10).launch(server_name="0.0.0.0", server_port=37906, root_path="/path")

app.py

@@ -3,7 +3,7 @@ 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, AttentionGuidedCAMChartGemma, AttentionGuidedCAMLLaVA
+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
 
 import numpy as np
@@ -82,8 +82,8 @@ def multimodal_understanding(model_type,
                 target_layers = all_layers[visualization_layer_min-1 : visualization_layer_max-1]
             else:
                 target_layers = [all_layers[visualization_layer_min-1]]
-            grad_cam = AttentionGuidedCAMClip(clip_utils.model, target_layers)
-            cam, outputs, grid_size = grad_cam.generate_cam(inputs, class_idx=0, visual_pooling_method=visual_pooling_method)
+            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 = cam.to("cpu")
             cam = [generate_gradcam(cam, image, size=(224, 224))]
             grad_cam.remove_hooks()
@@ -134,11 +134,11 @@ def multimodal_understanding(model_type,
             
 
             if model_name.split('-')[0] == "Janus":
-                gradcam = AttentionGuidedCAMJanus(vl_gpt, target_layers)
+                gradcam = VisualizationJanus(vl_gpt, target_layers)
             elif model_name.split('-')[0] == "LLaVA":
-                gradcam = AttentionGuidedCAMLLaVA(vl_gpt, target_layers)
+                gradcam = VisualizationLLaVA(vl_gpt, target_layers)
             elif model_name.split('-')[0] == "ChartGemma":
-                gradcam = AttentionGuidedCAMChartGemma(vl_gpt, target_layers)
+                gradcam = VisualizationChartGemma(vl_gpt, target_layers)
 
             start = 0
             cam = []
@@ -154,11 +154,11 @@ def multimodal_understanding(model_type,
                     cam = []
                     while start + i < len(input_ids_decoded):
                         if model_name.split('-')[0] == "Janus":
-                            gradcam = AttentionGuidedCAMJanus(vl_gpt, target_layers)
+                            gradcam = VisualizationJanus(vl_gpt, target_layers)
                         elif model_name.split('-')[0] == "LLaVA":
-                            gradcam = AttentionGuidedCAMLLaVA(vl_gpt, target_layers)
+                            gradcam = VisualizationLLaVA(vl_gpt, target_layers)
                         elif model_name.split('-')[0] == "ChartGemma":
-                            gradcam = AttentionGuidedCAMChartGemma(vl_gpt, target_layers)
+                            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_grid = cam_tensors.reshape(grid_size, grid_size)
                         cam_i = generate_gradcam(cam_grid, image)

demo/visualization.py

@@ -0,0 +1,524 @@
+import cv2
+import numpy as np
+import types
+import torch
+import torch.nn.functional as F
+import matplotlib.pyplot as plt
+from PIL import Image
+from torch import nn
+import spaces
+from demo.modify_llama import *
+
+
+class Visualization:
+    def __init__(self, model, register=True):
+        self.model = model
+        self.gradients = []
+        self.activations = []
+        self.hooks = []
+        if register:
+            self._register_hooks()
+
+    def _register_hooks(self):
+        for layer in self.target_layers:
+            self.hooks.append(layer.register_forward_hook(self._forward_hook))
+            self.hooks.append(layer.register_backward_hook(self._backward_hook))
+
+    def _forward_hook(self, module, input, output):
+        self.activations.append(output)
+
+    def _backward_hook(self, module, grad_in, grad_out):
+        self.gradients.append(grad_out[0])
+
+    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))
+
+    
+    def remove_hooks(self):
+        for hook in self.hooks:
+            hook.remove()  
+
+    def setup_grads(self):
+        torch.autograd.set_detect_anomaly(True)
+        for param in self.model.parameters():
+            param.requires_grad = False
+        
+        for layer in self.target_layers:
+            for param in layer.parameters():
+                param.requires_grad = True
+
+    def forward_backward(self):
+        raise NotImplementedError
+    
+    def grad_cam_vis(self):
+        self.model.zero_grad()
+        cam_sum = None
+        for act, grad in zip(self.activations, self.gradients):
+
+            act = F.relu(act[0])
+            
+            grad_weights = grad.mean(dim=-1, keepdim=True)
+            
+            print("act shape", act.shape)
+            print("grad_weights shape", grad_weights.shape)
+            
+            # cam = (act * grad_weights).sum(dim=-1)
+            cam, _ = (act * grad_weights).max(dim=-1)
+
+            print("cam_shape: ", cam.shape)
+
+            # Sum across all layers
+            if cam_sum is None:
+                cam_sum = cam
+            else:
+                cam_sum += cam  
+
+        cam_sum = F.relu(cam_sum)
+        return cam_sum
+    
+
+
+    def grad_cam_llm(self, mean_inside=False):
+
+        cam_sum = None
+        for act, grad in zip(self.activations, self.gradients):
+            
+            if mean_inside:
+                act = act.mean(dim=1)
+                grad = F.relu(grad.mean(dim=1))
+                cam = act * grad
+            else:
+                cam = act * grad
+                cam = act * grad.sum(dim=1)
+            
+            print(cam.shape)
+
+            # Sum across all layers
+            if cam_sum is None:
+                cam_sum = cam
+            else:
+                cam_sum += cam  
+
+        cam_sum = F.relu(cam_sum)
+        return cam_sum
+    
+    def attention_map(self):
+        raise NotImplementedError
+    
+    def attn_guided_cam(self):
+        
+        cams = []
+        for act, grad in zip(self.activations, self.gradients):
+            print("act shape", act.shape)
+            print("grad shape", grad.shape)
+            
+            grad = F.relu(grad)
+
+            # cam = grad
+            cam = act * grad # shape: [1, heads, seq_len, seq_len]
+            cam = cam.sum(dim=1) # shape: [1, seq_len, seq_len]
+            cam = cam.to(torch.float32).detach().cpu()
+            cams.append(cam)
+        return cams
+    
+    
+    def process(self, cam_sum, thresholding=True, remove_cls=True, normalize=True):
+
+        cam_sum = cam_sum.to(torch.float32)
+
+        # thresholding
+        if thresholding:
+            percentile = torch.quantile(cam_sum, 0.2)  # Adjust threshold dynamically
+            cam_sum[cam_sum < percentile] = 0
+
+        # Remove CLS
+        if remove_cls:
+            cam_sum = cam_sum[0, 1:]
+
+        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}")
+        cam_sum = cam_sum.view(grid_size, grid_size).detach()
+
+        # Normalize
+        if normalize:
+            cam_sum = (cam_sum - cam_sum.min()) / (cam_sum.max() - cam_sum.min())
+
+        return cam_sum, grid_size
+    
+    def process_multiple(self, cam_sum, start_idx, images_seq_mask, thresholding=True, normalize=True):
+        cam_sum = cam_sum.to(torch.float32)
+        # thresholding
+        if thresholding:
+            percentile = torch.quantile(cam_sum, 0.2)  # Adjust threshold dynamically
+            cam_sum[cam_sum < percentile] = 0
+
+
+        # cam_sum shape: [1, seq_len, seq_len]
+        cam_sum_lst = []
+        cam_sum_raw = cam_sum
+        start = start_idx
+        for i in range(start, cam_sum_raw.shape[1]):
+            cam_sum = cam_sum_raw[:, i, :] # shape: [1: seq_len]
+            cam_sum = cam_sum[images_seq_mask].unsqueeze(0) # shape: [1, img_seq_len]
+            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}")
+
+            cam_sum = cam_sum.view(grid_size, grid_size)
+            if normalize:
+                cam_sum = (cam_sum - cam_sum.min()) / (cam_sum.max() - cam_sum.min())
+            cam_sum = cam_sum.detach().to("cpu")
+            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):
+        cam_sum_lst = []
+        for i in range(start_idx, cams[0].shape[1]):
+            cam_sum = None
+            for layer, cam_l in enumerate(cams):
+                cam_l_i = cam_l[0, i, :] # shape: [1: seq_len]
+
+                cam_l_i = cam_l_i[images_seq_mask].unsqueeze(0) # shape: [1, img_seq_len]
+
+                num_patches = cam_l_i.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_reshaped = cam_l_i.view(grid_size, grid_size)
+
+                if normalize:
+                    cam_reshaped = (cam_reshaped - cam_reshaped.min()) / (cam_reshaped.max() - cam_reshaped.min())
+                if cam_sum == None:
+                    cam_sum = cam_reshaped
+                else:
+                    cam_sum += cam_reshaped
+
+            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_cam(self, input_tensor, target_token_idx=None):
+        raise NotImplementedError
+
+
+
+
+class VisualizationClip(Visualization):
+    def __init__(self, model, target_layers):
+        self.target_layers = target_layers
+        super().__init__(model)
+
+    @spaces.GPU(duration=120)
+    def forward_backward(self, input_tensor, visual_pooling_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":
+            output = output_full.image_embeds
+        elif visual_pooling_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)
+
+
+        output.backward(output_full.text_embeds[target_token_idx:target_token_idx+1], retain_graph=True)
+        return output_full
+        
+        
+    @spaces.GPU(duration=120)
+    def generate_cam(self, input_tensor, target_token_idx=None, visual_pooling_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)
+
+        cam_sum = self.grad_cam_vis()
+        cam_sum, grid_size = self.process(cam_sum)
+        
+        return cam_sum, output_full, grid_size
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+class VisualizationJanus(Visualization):
+    def __init__(self, model, target_layers):
+        self.target_layers = target_layers
+        super().__init__(model)
+        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"):
+        # Forward
+        image_embeddings, inputs_embeddings, outputs = self.model(input_tensor, tokenizer, temperature, top_p)
+        input_ids = input_tensor.input_ids
+
+        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()
+            loss.backward()
+            
+            self.activations = [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"):
+        
+        self.setup_grads()
+
+        # Forward Backward pass
+        self.forward_backward(input_tensor, tokenizer, temperature, top_p, target_token_idx, visual_pooling_method, focus)
+        
+        start_idx = 620
+        if focus == "Visual Encoder":
+
+            cam_sum = self.grad_cam_vis()
+            cam_sum, grid_size = self.process(cam_sum)
+            return cam_sum, grid_size, start_idx
+
+        elif focus == "Language Model":
+            
+            cam_sum = self.grad_cam_llm(mean_inside=True)
+            
+            images_seq_mask = input_tensor.images_seq_mask
+            
+            cam_sum_lst, grid_size = self.process_multiple(cam_sum, start_idx, images_seq_mask)
+
+            return cam_sum_lst, grid_size, start_idx
+
+
+
+
+
+
+
+
+
+class VisualizationLLaVA(Visualization):
+    def __init__(self, model, target_layers):
+        self.target_layers = target_layers
+        super().__init__(model, register=False)
+        self._modify_layers()
+        self._register_hooks_activations()
+
+    def forward_backward(self, inputs):
+        # Forward pass
+        outputs_raw = self.model(**inputs)
+
+        self.model.zero_grad()
+        print("outputs_raw", outputs_raw)
+
+        loss = outputs_raw.logits.max(dim=-1).values.sum()
+        loss.backward()
+        self.activations = [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"):
+        
+        self.setup_grads()
+        self.forward_backward(inputs)
+
+        # get image masks
+        images_seq_mask = []
+        last = 0
+        for i in range(inputs["input_ids"].shape[1]):
+            decoded_token = tokenizer.decode(inputs["input_ids"][0][i].item())
+            if (decoded_token == "<image>"):
+                images_seq_mask.append(True)
+                last = i
+            else:
+                images_seq_mask.append(False)
+
+
+        # 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)
+
+        return cam_sum_lst, grid_size, start_idx
+
+
+
+
+
+
+class VisualizationChartGemma(Visualization):
+    def __init__(self, model, target_layers):
+        self.target_layers = target_layers
+        super().__init__(model, register=True)
+        self._modify_layers()
+        self._register_hooks_activations()
+    
+    def forward_backward(self, inputs, focus, start_idx, target_token_idx):
+        outputs_raw = self.model(**inputs, output_hidden_states=True)
+        if focus == "Visual Encoder":
+            
+            self.model.zero_grad()
+
+            loss = outputs_raw.logits.max(dim=-1).values[0, start_idx + target_token_idx]
+            loss.backward()
+        
+        elif focus == "Language Model":
+            self.model.zero_grad()
+            if target_token_idx == -1:
+                loss = outputs_raw.logits.max(dim=-1).values.sum()
+            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.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"):
+        
+        # Forward pass
+        self.setup_grads()
+        
+        # get image masks
+        images_seq_mask = []
+        last = 0
+        for i in range(inputs["input_ids"].shape[1]):
+            decoded_token = tokenizer.decode(inputs["input_ids"][0][i].item())
+            if (decoded_token == "<image>"):
+                images_seq_mask.append(True)
+                last = i
+            else:
+                images_seq_mask.append(False)
+        start_idx = last + 1
+
+
+        self.forward_backward(inputs, focus, start_idx, target_token_idx)
+        if focus == "Visual Encoder":
+            
+            cam_sum = self.grad_cam_vis()
+            cam_sum, grid_size = self.process(cam_sum, remove_cls=False)
+
+            return cam_sum, grid_size, start_idx
+
+        elif focus == "Language Model":
+
+            cams = self.attn_guided_cam()
+            cam_sum_lst, grid_size = self.process_multiple_withsum(cams, start_idx, images_seq_mask)
+
+            # cams shape: [layers, 1, seq_len, seq_len]
+            
+
+
+        return cam_sum_lst, grid_size, start_idx
+
+
+
+
+
+
+
+
+
+
+def generate_gradcam(
+    cam, 
+    image,
+    size = (384, 384),
+    alpha=0.5, 
+    colormap=cv2.COLORMAP_JET, 
+    aggregation='mean',
+    normalize=False
+):
+    """
+    Generates a Grad-CAM heatmap overlay on top of the input image.
+
+    Parameters:
+      attributions (torch.Tensor): A tensor of shape (C, H, W) representing the
+        intermediate activations or gradients at the target layer.
+      image (PIL.Image): The original image.
+      alpha (float): The blending factor for the heatmap overlay (default 0.5).
+      colormap (int): OpenCV colormap to apply (default cv2.COLORMAP_JET).
+      aggregation (str): How to aggregate across channels; either 'mean' or 'sum'.
+
+    Returns:
+      PIL.Image: The image overlaid with the Grad-CAM heatmap.
+    """
+    # print("Generating Grad-CAM with shape:", cam.shape)
+
+    if normalize:
+        cam_min, cam_max = cam.min(), cam.max()
+        cam = cam - cam_min
+        cam = cam / (cam_max - cam_min)
+    # Convert tensor to numpy array
+    cam = torch.nn.functional.interpolate(cam.unsqueeze(0).unsqueeze(0), size=size, mode='bilinear').squeeze()
+    cam_np = cam.squeeze().detach().cpu().numpy()
+
+    # Apply Gaussian blur for smoother heatmaps
+    cam_np = cv2.GaussianBlur(cam_np, (5,5), sigmaX=0.8)
+
+    # Resize the cam to match the image size
+    width, height = size
+    cam_resized = cv2.resize(cam_np, (width, height))
+
+    # Convert the normalized map to a heatmap (0-255 uint8)
+    heatmap = np.uint8(255 * cam_resized)
+    heatmap = cv2.applyColorMap(heatmap, colormap)
+    # OpenCV produces heatmaps in BGR, so convert to RGB for consistency
+    heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)
+
+    # Convert original image to a numpy array
+    image_np = np.array(image)
+    image_np = cv2.resize(image_np, (width, height))
+
+    # Blend the heatmap with the original image
+    overlay = cv2.addWeighted(image_np, 1 - alpha, heatmap, alpha, 0)
+
+    return Image.fromarray(overlay)
+