Update handler.py

handler.py

@@ -14,7 +14,7 @@ class EndpointHandler():
         ).to(0)
         processor = AutoProcessor.from_pretrained(model_id)
 
-    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+    def __call__(self, data: Dict[str, Any]):
         parameters = data.pop("inputs",data)
         inputs = data.pop("inputs", data)
         if parameters is not None:
@@ -24,22 +24,4 @@ class EndpointHandler():
             output = model.generate(**inputs, max_new_tokens=200, do_sample=False)
         return output
 
-    
-prompt = "USER: <image>\nWhat are these?\nASSISTANT:"
-image_file = "http://images.cocodataset.org/val2017/000000039769.jpg"
-
-model = LlavaForConditionalGeneration.from_pretrained(
-    model_id, 
-    torch_dtype=torch.float16, 
-    low_cpu_mem_usage=True, 
-).to(0)
-
-processor = AutoProcessor.from_pretrained(model_id)
-
-
-raw_image = Image.open(requests.get(image_file, stream=True).raw)
-inputs = processor(prompt, raw_image, return_tensors='pt').to(0, torch.float16)
-
-output = model.generate(**inputs, max_new_tokens=200, do_sample=False)
-print(processor.decode(output[0][2:], skip_special_tokens=True))
-
+