Update handler.py

handler.py

@@ -4,6 +4,8 @@ from PIL import Image
 import requests
 import torch
 import gc
+import base64
+import io
 
 class EndpointHandler:
     def __init__(self, path=""):
@@ -22,27 +24,30 @@ class EndpointHandler:
         )
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
-        # Clear CUDA cache
         torch.cuda.empty_cache()
         gc.collect()
 
-        # Extract inputs from the request data
         inputs = data.get("inputs", {})
         image_url = inputs.get("image_url")
+        image_data = inputs.get("image")
         text_prompt = inputs.get("text_prompt", "Describe this image.")
 
-        if not image_url:
-            return [{"error": "No image_url provided in inputs"}]
+        if image_url:
+            try:
+                image = Image.open(requests.get(image_url, stream=True).raw)
+            except Exception as e:
+                return [{"error": f"Failed to load image from URL: {str(e)}"}]
+        elif image_data:
+            try:
+                image = Image.open(io.BytesIO(base64.b64decode(image_data)))
+            except Exception as e:
+                return [{"error": f"Failed to decode image data: {str(e)}"}]
+        else:
+            return [{"error": "No image_url or image data provided in inputs"}]
 
-        # Download and process the image
-        try:
-            image = Image.open(requests.get(image_url, stream=True).raw)
-            if image.mode != "RGB":
-                image = image.convert("RGB")
-        except Exception as e:
-            return [{"error": f"Failed to load image: {str(e)}"}]
+        if image.mode != "RGB":
+            image = image.convert("RGB")
 
-        # Process the image and text
         try:
             with torch.cuda.amp.autocast(enabled=True, dtype=torch.bfloat16):
                 inputs = self.processor.process(
@@ -50,21 +55,17 @@ class EndpointHandler:
                     text=text_prompt
                 )
                 
-                # Move inputs to the correct device and make a batch of size 1
                 inputs = {k: v.to(self.model.device).unsqueeze(0) for k, v in inputs.items()}
 
-                # Generate output
                 output = self.model.generate_from_batch(
                     inputs,
                     GenerationConfig(max_new_tokens=200, stop_strings="<|endoftext|>"),
                     tokenizer=self.processor.tokenizer
                 )
 
-            # Decode the generated tokens
             generated_tokens = output[0, inputs['input_ids'].size(1):]
             generated_text = self.processor.tokenizer.decode(generated_tokens, skip_special_tokens=True)
 
-            # Clear CUDA cache again
             torch.cuda.empty_cache()
             gc.collect()