fixed llava

app.py

@@ -208,13 +208,13 @@ 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
+        language_model_max_layer = 32
+        language_model_best_layer = 24
 
         res = (
             gr.Dropdown(choices=["Visualization only", "answer + visualization"], value="Visualization only", 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.Slider(minimum=1, maximum=32, value=24, step=1, label="visualization layers min"),
+            gr.Slider(minimum=1, maximum=32, value=24, 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")
         )

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, PaliGemmaForConditionalGeneration
+from transformers import AutoConfig, AutoModelForCausalLM, LlavaForConditionalGeneration, LlavaNextForConditionalGeneration, LlavaNextProcessor, AutoProcessor, PaliGemmaForConditionalGeneration
 from transformers import CLIPProcessor, CLIPModel
 from janus.models import MultiModalityCausalLM, VLChatProcessor
 
@@ -123,13 +123,13 @@ class LLaVA_Utils(Model_Utils):
         model_path = "llava-hf/llava-v1.6-mistral-7b-hf"
         config = AutoConfig.from_pretrained(model_path)
 
-        self.vl_gpt = LlavaForConditionalGeneration.from_pretrained(model_path,
+        self.vl_gpt = LlavaNextForConditionalGeneration.from_pretrained(model_path,
                                                     low_cpu_mem_usage=True,
                                                     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.processor = LlavaNextProcessor.from_pretrained(model_path)
         self.tokenizer = self.processor.tokenizer
         
         return self.vl_gpt, self.tokenizer