test

handler.py

@@ -61,8 +61,8 @@ CONTROLNET_MAPPING = {
 class EndpointHandler():
     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 = "depth"
+        self.controlnet = ControlNetModel.from_pretrained(CONTROLNET_MAPPING[self.control_type]["model_id"],torch_dtype=dtype).to(device)
 
         #processor = AutoProcessor.from_pretrained("CompVis/stable-diffusion-safety-checker")
 
@@ -71,10 +71,10 @@ class EndpointHandler():
         #self.stable_diffusion_id = "runwayml/stable-diffusion-v1-5"
         self.stable_diffusion_id = "Lykon/dreamshaper-8"
 
-        #self.pipe = StableDiffusionControlNetPipeline.from_pretrained(self.stable_diffusion_id,
-        #                                                              controlnet=self.controlnet, 
-        #                                                              torch_dtype=dtype,
-        #                                                              safety_checker=StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker", torch_dtype=dtype)).to("cuda")
+        self.pipe = StableDiffusionControlNetPipeline.from_pretrained(self.stable_diffusion_id,
+                                                                      #controlnet=self.controlnet, 
+                                                                      torch_dtype=dtype,
+                                                                      safety_checker=StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker", torch_dtype=dtype)).to("cuda")
         # Define Generator with seed
         self.generator = torch.Generator(device=device.type).manual_seed(3)
 
@@ -92,12 +92,12 @@ class EndpointHandler():
             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")
-        #    self.control_type = controlnet_type
-        #    self.controlnet = ControlNetModel.from_pretrained(CONTROLNET_MAPPING[self.control_type]["model_id"],
-       #                                                       torch_dtype=dtype).to(device)
-        #    self.pipe.controlnet = self.controlnet
+        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")
+            self.control_type = controlnet_type
+            self.controlnet = ControlNetModel.from_pretrained(CONTROLNET_MAPPING[self.control_type]["model_id"],
+                                                              torch_dtype=dtype).to(device)
+            self.pipe.controlnet = self.controlnet
         
         
         targets = [self.pipe.vae, self.pipe.unet]