use LlavaForConditionalGeneration to load instead

handler.py

@@ -3,9 +3,17 @@ from transformers import pipeline
 from PIL import Image    
 import requests
 
+import torch
+from transformers import AutoProcessor, LlavaForConditionalGeneration
+
 class EndpointHandler():
     def __init__(self, path="."):
-        self.pipeline = pipeline("image-to-text", model=path)
+        self.model = LlavaForConditionalGeneration.from_pretrained(
+            path, 
+            torch_dtype=torch.float16, 
+            low_cpu_mem_usage=True, 
+        ).to(0)
+        self.processor = AutoProcessor.from_pretrained(path)
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
@@ -24,8 +32,10 @@ class EndpointHandler():
         image = Image.open(requests.get(url, stream=True).raw)
 
 
+        inputs = self.processor(prompt, image, return_tensors='pt').to(0, torch.float16)
+
         # run normal prediction
-        outputs = self.pipeline(image, prompt=prompt, generate_kwargs={"max_new_tokens": 200})
-        print(outputs)
+        output = self.model.generate(**inputs, max_new_tokens=200, do_sample=False)
+        print(output)
         
-        return outputs
+        return output