Update handler.py

handler.py

@@ -1,6 +1,6 @@
 from typing import Dict, List, Any
 import torch
-from transformers import AutoProcessor, AutoModelForConditionalGeneration, BitsAndBytesConfig
+from transformers import LlavaNextProcessor, LlavaNextForConditionalGeneration, BitsAndBytesConfig
 from PIL import Image
 import requests
 from io import BytesIO
@@ -16,8 +16,8 @@ class EndpointHandler():
         )
 
         # Carrega el processador i model de forma global
-        self.processor = AutoProcessor.from_pretrained("llava-hf/llava-v1.6-mistral-7b-hf")
-        self.model = AutoModelForConditionalGeneration.from_pretrained(
+        self.processor = LlavaNextProcessor.from_pretrained("llava-hf/llava-v1.6-mistral-7b-hf")
+        self.model = LlavaNextForConditionalGeneration.from_pretrained(
             "llava-hf/llava-v1.6-mistral-7b-hf",
             quantization_config=quantization_config,
             device_map="auto"
@@ -60,4 +60,3 @@ class EndpointHandler():
             score = int(match[2])
             scores[category_name] = score
         return scores
-