Update handler.py

handler.py

@@ -1,4 +1,4 @@
-from typing import  Dict, List, Any
+from typing import Dict, List, Any
 import base64
 from PIL import Image
 from io import BytesIO
@@ -6,12 +6,11 @@ from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, Autoen
 import torch
 from diffusers.utils import load_image
 
-
-
 import numpy as np
 import cv2
 import controlnet_hinter
-# ADDED AUTO PIPE, next try replacing
+
+# ADDED AUTO PIPE
 # set device
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 if device.type != 'cuda':
@@ -55,20 +54,20 @@ 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)
- 
+        self.controlnet = ControlNetModel.from_pretrained(CONTROLNET_MAPPING[self.control_type]["model_id"], torch_dtype=dtype).to(device)
+
         # 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, 
+                                                                      torch_dtype=dtype,
                                                                       safety_checker=None).to(device)
- 
+
         # Define Generator with seed
         # COMMENTED self.generator = torch.Generator(device="cpu").manual_seed(3)
 
@@ -80,6 +79,18 @@ class EndpointHandler():
         prompt = data.pop("inputs", None)
         image = data.pop("image", None)
         controlnet_type = data.pop("controlnet_type", None)
+        stablediffusion_id = data.pop("stablediffusionid", None)  # Get the stablediffusionid from the request data
+
+        if stablediffusion_id is not None and stablediffusion_id != self.stable_diffusion_id:
+            # Change the Stable Diffusion model to the new model ID
+            self.stable_diffusion_id = stablediffusion_id
+            # Reinitialize the pipeline with the new model ID
+            self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
+                self.stable_diffusion_id,
+                controlnet=self.controlnet,
+                torch_dtype=dtype,
+                safety_checker=None
+            ).to(device)
 
         # Check if neither prompt nor image is provided
         if prompt is None and image is None:
@@ -93,9 +104,8 @@ class EndpointHandler():
                                                               torch_dtype=dtype).to(device)
             self.pipe.controlnet = self.controlnet
 
-
-        # hyperparamters
-        negatice_prompt = data.pop("negative_prompt", None)
+        # hyperparameters
+        negative_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)
@@ -118,12 +128,11 @@ class EndpointHandler():
             height=height,
             width=width,
             controlnet_conditioning_scale=controlnet_conditioning_scale,
-            guess_mode=True, 
-
+            guess_mode=True,
         )
 
-        #generator=self.generator COMMENTED from self.pipe
-        # return first generate PIL image
+        # generator=self.generator COMMENTED from self.pipe
+        # return the first generated PIL image
         return out.images[0]
 
     # helper to decode input image