a

handler.py

@@ -1,16 +1,21 @@
 from typing import Dict, List, Any
 from transformers import pipeline
-from diffusers import AutoPipelineForText2Image
 import torch
 import base64
 from io import BytesIO
 from PIL import Image
+from diffusers import StableDiffusionXLImg2ImgPipeline
+from diffusers.utils import load_image
+
 
 
 class EndpointHandler():
-    def __init__(self, path=""):        
-        self.pipe = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16")
+    def __init__(self, path=""):                
+        self.pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16")
         self.pipe.to("cuda")
+        self.pipe.unet = torch.compile(self.pipe.unet, mode="reduce-overhead", fullgraph=True)
+        self.pipe.upcast_vae()
+
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
@@ -23,7 +28,6 @@ class EndpointHandler():
         # get inputs
         inputs = data.pop("inputs", data)
         encoded_image = data.pop("image", None)
-        encoded_mask_image = data.pop("mask_image", None)
         
         # hyperparamters
         num_inference_steps = data.pop("num_inference_steps", 25)
@@ -31,20 +35,30 @@ class EndpointHandler():
         negative_prompt = data.pop("negative_prompt", None)
         height = data.pop("height", None)
         width = data.pop("width", None)
+
+        strength = data.pop("strength", 0.7)
+        denoising_start = data.pop("denoising_start_step", 0)
+        denoising_end = data.pop("denoising_start_step", 0)
+        num_images_per_prompt = data.pop("num_images_per_prompt", 1)
+        aesthetic_score = data.pop("aesthetic_score", 0.6)
+
         
         # process image
-        if encoded_image is not None and encoded_mask_image is not None:
-            image = self.decode_base64_image(encoded_image)
-            mask_image = self.decode_base64_image(encoded_mask_image)
+        if encoded_image is not None:
+            image = self.decode_base64_image(encoded_image)            
         else:
             image = None
-            mask_image = None 
+            
         
         # run inference pipeline
         out = self.pipe(inputs, 
                         image=image,             
-                        mask_image=mask_image,             
+                        strenght=strength,
                         num_inference_steps=num_inference_steps,
+                        denoising_start_step=denoising_start,
+                        denoising_end_step=denoising_end,
+                        num_images_per_prompt=num_images_per_prompt,
+                        aesthetic_score=aesthetic_score,
                         guidance_scale=guidance_scale,
                         num_images_per_prompt=1,
                         negative_prompt=negative_prompt,
@@ -60,6 +74,4 @@ class EndpointHandler():
         base64_image = base64.b64decode(image_string)
         buffer = BytesIO(base64_image)
         image = Image.open(buffer)
-        return image        
-        
-        
+        return image