Update handler.py

handler.py

@@ -17,6 +17,8 @@ from transformers.utils import logging
 logger = logging.get_logger(__name__)
 logging.set_verbosity_info()
 
+BASE_MODEL_ID = "mistral-community/pixtral-12b"
+
 DEFAULT_PROMPT = (
     "Here is a picture showing some food waste.\n\n"
     "Task: Provide a list of the food waste items visible in the picture.\n"
@@ -254,13 +256,13 @@ class EndpointHandler:
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         logger.info("Initializing EndpointHandler on device: %s", self.device)
         self.processor = AutoProcessor.from_pretrained(
-            path,
+            BASE_MODEL_ID,
             trust_remote_code=True,
         )
 
         dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
         self.model = LlavaForConditionalGeneration.from_pretrained(
-            path,
+            BASE_MODEL_ID,
             torch_dtype=dtype,
             low_cpu_mem_usage=True,
             device_map={"": self.device},