Upload handler.py

handler.py

@@ -0,0 +1,44 @@
+from typing import Dict, List, Any
+from diffusers import StableDiffusionUpscalePipeline
+import torch
+from PIL import Image
+import io
+
+class EndpointHandler:
+    def __init__(self, path: str):
+        # Load the Stable Diffusion x4 upscaler model
+        self.pipeline = StableDiffusionUpscalePipeline.from_pretrained(
+            path,
+            torch_dtype=torch.float16
+        )
+        self.pipeline.to("cuda")
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """
+        data args:
+            inputs: str - The text prompt for the upscaling.
+            image: bytes - The low-resolution image as byte data.
+        
+        Return:
+            A list of dictionaries with the upscaled image.
+        """
+        # Extract inputs and image from the payload
+        prompt = data.get("inputs", "")
+        image_bytes = data.get("image", None)
+
+        if image_bytes is None:
+            return [{"error": "No image provided"}]
+
+        # Convert the byte data to an image
+        low_res_img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+
+        # Perform upscaling
+        upscaled_image = self.pipeline(prompt=prompt, image=low_res_img).images[0]
+
+        # Save the upscaled image to a byte stream
+        byte_io = io.BytesIO()
+        upscaled_image.save(byte_io, format="PNG")
+        byte_io.seek(0)
+
+        # Return the upscaled image as byte data
+        return [{"upscaled_image": byte_io.getvalue()}]