Update handler.py

handler.py

@@ -1,9 +1,11 @@
-from typing import Dict, List, Any
+from typing import  Dict, List, Any
 import base64
 from PIL import Image
 from io import BytesIO
 from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
 import torch
+
+
 import numpy as np
 import cv2
 import controlnet_hinter
@@ -11,8 +13,7 @@ import controlnet_hinter
 # set device
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 if device.type != 'cuda':
-    raise ValueError("Need to run on GPU")
-
+    raise ValueError("need to run on GPU")
 # set mixed precision dtype
 dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] == 8 else torch.float16
 
@@ -52,14 +53,12 @@ CONTROLNET_MAPPING = {
     }
 }
 
+
 class EndpointHandler():
-    """
-    A class to handle endpoint logic.
-    """
     def __init__(self, path=""):
         # define default controlnet id and load controlnet
-        self.control_type = "depth"
-        self.controlnet = ControlNetModel.from_pretrained(controlnet_mapping[self.control_type]["model_id"], torch_dtype=dtype).to(device)
+        self.control_type = "normal"
+        self.controlnet = ControlNetModel.from_pretrained(CONTROLNET_MAPPING[self.control_type]["model_id"],torch_dtype=dtype).to(device)
         
         # Load StableDiffusionControlNetPipeline 
         self.stable_diffusion_id = "runwayml/stable-diffusion-v1-5"
@@ -70,30 +69,29 @@ class EndpointHandler():
         # Define Generator with seed
         self.generator = torch.Generator(device="cpu").manual_seed(3)
 
-    def __call__(self, data: Any) -> None:
+    def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
         """
-        Process input data and perform inference.
-        
-        :param data: A dictionary containing `inputs` and optional `image_path` field.
-        :return: None
+        :param data: A dictionary contains `inputs` and optional `image` field.
+        :return: A dictionary with `image` field contains image in base64.
         """
         prompt = data.pop("inputs", None)
-        image_path = data.pop("image_path", None)
+        image = data.pop("image", None)
         controlnet_type = data.pop("controlnet_type", None)
-
-        # Check if neither prompt nor image path is provided
-        if prompt is None and image_path is None:
-            raise ValueError("Please provide a prompt and either an image path or a base64-encoded image.")
-
+        
+        # Check if neither prompt nor image is provided
+        if prompt is None and image is None:
+            return {"error": "Please provide a prompt and base64 encoded image."}
+        
         # Check if a new controlnet is provided
         if controlnet_type is not None and controlnet_type != self.control_type:
-            print(f"Changing controlnet from {self.control_type} to {controlnet_type} using {controlnet_mapping[controlnet_type]['model_id']} model")
+            print(f"changing controlnet from {self.control_type} to {controlnet_type} using {CONTROLNET_MAPPING[controlnet_type]['model_id']} model")
             self.control_type = controlnet_type
-            self.controlnet = ControlNetModel.from_pretrained(controlnet_mapping[self.control_type]["model_id"],
+            self.controlnet = ControlNetModel.from_pretrained(CONTROLNET_MAPPING[self.control_type]["model_id"],
                                                               torch_dtype=dtype).to(device)
             self.pipe.controlnet = self.controlnet
         
-        # hyperparameters
+        
+        # hyperparamters
         num_inference_steps = data.pop("num_inference_steps", 30)
         guidance_scale = data.pop("guidance_scale", 7.5)
         negative_prompt = data.pop("negative_prompt", None)
@@ -102,14 +100,8 @@ class EndpointHandler():
         controlnet_conditioning_scale = data.pop("controlnet_conditioning_scale", 1.0)
         
         # process image
-        if image_path is not None:
-            # Load the image from the specified path
-            image = Image.open(image_path)
-        else:
-            # Decode base64-encoded image
-            image = self.decode_base64_image(data.pop("image", ""))
-        
-        control_image = controlnet_mapping[self.control_type]["hinter"](image)
+        image = self.decode_base64_image(image)
+        control_image = CONTROLNET_MAPPING[self.control_type]["hinter"](image)
         
         # run inference pipeline
         out = self.pipe(
@@ -125,28 +117,13 @@ class EndpointHandler():
             generator=self.generator
         )
 
-        # save the generated image as a JPEG file
-        output_image = out.images[0]
-        output_image.save("output.jpg", format="JPEG")
-
+        
+        # return first generate PIL image
+        return out.images[0]
+    
+    # helper to decode input image
     def decode_base64_image(self, image_string):
         base64_image = base64.b64decode(image_string)
         buffer = BytesIO(base64_image)
         image = Image.open(buffer)
-        return image
-
-# Example usage
-payload = {
-    "inputs": "Your prompt here",
-    "image_path": "path/to/your/image.jpg",
-    "controlnet_type": "depth",
-    "num_inference_steps": 30,
-    "guidance_scale": 7.5,
-    "negative_prompt": None,
-    "height": None,
-    "width": None,
-    "controlnet_conditioning_scale": 1.0,
-}
-
-handler = EndpointHandler()
-handler(payload)
+        return image