Update handler.py

handler.py

@@ -1,9 +1,12 @@
+%%writefile handler.py
 from typing import  Dict, List, Any
 import base64
 from PIL import Image
 from io import BytesIO
-from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
+from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, AutoencoderKL, StableDiffusionXLControlNetPipeline
 import torch
+from diffusers.utils import load_image
+
 
 
 import numpy as np
@@ -53,21 +56,22 @@ CONTROLNET_MAPPING = {
     }
 }
 
-
+ 
 class EndpointHandler():
     def __init__(self, path=""):
         # define default controlnet id and load controlnet
         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"
-        self.pipe = StableDiffusionControlNetPipeline.from_pretrained(self.stable_diffusion_id, 
-                                                                      controlnet=self.controlnet, 
-                                                                      torch_dtype=dtype,
+ 
+        # Load StableDiffusionControlNetPipeline
+        self.stable_diffusion_id = "stablediffusionapi/disney-pixar-cartoon"
+        self.pipe = StableDiffusionControlNetPipeline.from_pretrained(self.stable_diffusion_id,
+                                                                      controlnet=self.controlnet,
+                                                                      torch_dtype=dtype, 
                                                                       safety_checker=None).to(device)
+ 
         # Define Generator with seed
-        self.generator = torch.Generator(device="cpu").manual_seed(3)
+        # COMMENTED self.generator = torch.Generator(device="cpu").manual_seed(3)
 
     def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
         """
@@ -77,11 +81,11 @@ class EndpointHandler():
         prompt = data.pop("inputs", None)
         image = data.pop("image", None)
         controlnet_type = data.pop("controlnet_type", None)
-        
+
         # 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")
@@ -89,41 +93,43 @@ class EndpointHandler():
             self.controlnet = ControlNetModel.from_pretrained(CONTROLNET_MAPPING[self.control_type]["model_id"],
                                                               torch_dtype=dtype).to(device)
             self.pipe.controlnet = self.controlnet
-        
-        
+
+
         # hyperparamters
-        num_inference_steps = data.pop("num_inference_steps", 30)
-        guidance_scale = data.pop("guidance_scale", 7.5)
+        negatice_prompt = data.pop("negative_prompt", None)
+        num_inference_steps = data.pop("num_inference_steps", 150)
+        guidance_scale = data.pop("guidance_scale", 5)
         negative_prompt = data.pop("negative_prompt", None)
         height = data.pop("height", None)
         width = data.pop("width", None)
         controlnet_conditioning_scale = data.pop("controlnet_conditioning_scale", 1.0)
-        
+
         # process image
         image = self.decode_base64_image(image)
         control_image = CONTROLNET_MAPPING[self.control_type]["hinter"](image)
-        
+
         # run inference pipeline
         out = self.pipe(
-            prompt=prompt, 
+            prompt=prompt,
             negative_prompt=negative_prompt,
             image=control_image,
-            num_inference_steps=num_inference_steps, 
+            num_inference_steps=num_inference_steps,
             guidance_scale=guidance_scale,
             num_images_per_prompt=1,
             height=height,
             width=width,
             controlnet_conditioning_scale=controlnet_conditioning_scale,
-            generator=self.generator
+            guess_mode=True, 
+
         )
 
-        
+        #generator=self.generator COMMENTED from self.pipe
         # 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
+        return image