Update handler.py

handler.py

@@ -1,6 +1,9 @@
 import os
 from diffusers import AutoPipelineForText2Image
 import torch
+from PIL import Image
+from io import BytesIO
+import base64
 
 class EndpointHandler:
     def __init__(self, path: str = ""):
@@ -27,6 +30,16 @@ class EndpointHandler:
         prompt = data.get("inputs", None)
         if not prompt:
             raise ValueError("No prompt provided in the input")
+
         with torch.no_grad():
             images = self.pipeline(prompt).images
-        return {"image": images[0]}
+
+        # Get the first generated image
+        pil_image = images[0]
+
+        # Convert the image to bytes to return as the serialized format (for instance, base64)
+        buffered = BytesIO()
+        pil_image.save(buffered, format="PNG")
+        img_bytes = buffered.getvalue()
+
+        return img_bytes  # Return the image bytes directly for Hugging Face to serialize