updated.

handler.py

@@ -5,42 +5,34 @@ from PIL import Image
 from webcolors import CSS3_HEX_TO_NAMES, hex_to_rgb
 from scipy.spatial import KDTree
 import io
+import base64
+from io import BytesIO
+from PIL import Image
+import json
 
-class EndpointHandler:
+class EndpointHandler():
     
     def __init__(self, path=""):
         model_id = "timbrooks/instruct-pix2pix"
         self.model = StableDiffusionInstructPix2PixPipeline.from_pretrained(path, torch_dtype=torch.float16, safety_checker=None).to("cuda")
         self.model.scheduler = EulerAncestralDiscreteScheduler.from_config(self.model.scheduler.config)
 
-    def __call__(self, data: Dict[str, Any]) -> Image:
-        """
-        Args:
-            data (dict):
-                The payload with the input image, text prompt, color code, and optional parameters.
-        """
-            
-        input_image = process_image_base64(data['inputs']['base64_image'])
-        
-        text_prompt = data['inputs']['text_prompt']
-        color_code = data['inputs']['color_code']
-        
-        guidance_scale = data['inputs'].get('guidance_scale', 7.5)
-        image_guidance_scale = data['inputs'].get('image_guidance_scale', 1.5)
+    def __call__(self, data):
+
+        info=data['inputs']
+        image=info.pop("image",data)
+        prompt=info.pop("text",data)
+        image=base64.b64decode(image)
+        raw_images = Image.open(BytesIO(image)).convert('RGB')
         
-        result_prompt = self.build_prompt(text_prompt, color_code)
+        images = self.pipe(prompt, image=raw_images, num_inference_steps=25, image_guidance_scale=1).images
+        img=images[0]
+        img.save("./1.png")
+        with open('./1.png','rb') as img_file:
+            encoded_string = base64.b64encode(img_file.read()).decode('utf-8')
         
-        with torch.autocast("cuda"):
-            images = self.model(
-                result_prompt,
-                image=input_image,
-                num_inference_steps=50,
-                guidance_scale=guidance_scale,
-                image_guidance_scale=image_guidance_scale,
-            ).images
-            
-        return images
-
+        return {'image':encoded_string}
+    """
     def process_image_base64(self, base64_image_data):
         # Decode base64 data to bytes
         image_bytes = base64.b64decode(base64_image_data)
@@ -69,3 +61,8 @@ class EndpointHandler:
         kdt_db = KDTree(rgb_values)
         distance, index = kdt_db.query(rgb_tuple)
         return names[index]
+    """
+if __name__=="__main__":
+    my_handler = EndpointHandler(path='.')
+
+