Update handler.py

handler.py

@@ -11,7 +11,7 @@ class EndpointHandler():
         model_id = path
         self.model = LlavaForConditionalGeneration.from_pretrained(
             model_id, 
-            torch_dtype=torch.float32   , 
+            torch_dtype=torch.float32, 
             low_cpu_mem_usage=True,
             load_in_4bit=True
         )
@@ -23,9 +23,19 @@ class EndpointHandler():
 
         outputs = []
         for link in parameters:
-            raw_image = Image.open(requests.get(link, stream=True).raw)
-            inputs = self.processor(prompt, raw_image, return_tensors='pt').to(0, torch.float16)
-            output = self.model.generate(**inputs, max_new_tokens=200, do_sample=False)
-            readable = (self.processor.decode(output[0][2:], skip_special_tokens=True))
-            outputs.append(readable)
+            try:
+                # Fetch image from URL
+                response = requests.get(link, stream=True)
+                response.raise_for_status()  # Raise an exception for 4xx or 5xx status codes
+                raw_image = Image.open(response.raw)
+
+                # Process image and generate output
+                inputs = self.processor(prompt, raw_image, return_tensors='pt').to(0, torch.float32)
+                output = self.model.generate(**inputs, max_new_tokens=200, do_sample=False)
+                readable = self.processor.decode(output[0][2:], skip_special_tokens=True)
+                outputs.append(readable)
+            except Exception as e:
+                # Handle any exceptions and log the error
+                outputs.append(f"Error processing image from {link}: {str(e)}")
+        
         return outputs