adding other model

README.md

@@ -28,8 +28,4 @@ inference: true
 }
 ```
 
-SD_MODEL no tested 
-
-default => Lykon/dreamshaper-8
-
-default || dreamshaper || juggernaut || realistic-vision || rev-animated
+dreamshaper || dreamshaper_2

handler.py

@@ -26,49 +26,22 @@ CONTROLNET_MAPPING = {
     },
 }
 
-
-# SD_ID_MAPPING = {
-#     "default":  "Lykon/dreamshaper-8",
-#    "dreamshaper":  "stablediffusionapi/dreamshaper-xl",
-#    "juggernaut":  "stablediffusionapi/juggernaut-xl-v8",
-#    "realistic":"SG161222/Realistic_Vision_V1.4",
-#    "rev":"s6yx/ReV_Animated"
-# }
-
 class EndpointHandler():
     def __init__(self, path=""):
         self.control_type = "depth"
         self.controlnet = ControlNetModel.from_pretrained(CONTROLNET_MAPPING[self.control_type]["model_id"],torch_dtype=dtype).to(device)
 
         self.stable_diffusion_id_0 = "Lykon/dreamshaper-8"
-        self.default = StableDiffusionControlNetPipeline.from_pretrained(self.stable_diffusion_id_0,
+        self.dreamshaper = StableDiffusionControlNetPipeline.from_pretrained(self.stable_diffusion_id_0,
                                                                        controlnet=self.controlnet, 
                                                                        torch_dtype=dtype,
                                                                        safety_checker=StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker", torch_dtype=dtype)).to("cuda")
 
-        # self.stable_diffusion_id_1 = "stablediffusionapi/dreamshaper-xl"
-        # self.juggernaut = StableDiffusionControlNetPipeline.from_pretrained(self.stable_diffusion_id_1,
-        #                                                               controlnet=self.controlnet, 
-        #                                                               torch_dtype=dtype,
-        #                                                               safety_checker=StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker", torch_dtype=dtype)).to("cuda")
-        # self.stable_diffusion_id_2 = "stablediffusionapi/juggernaut-xl-v8"
-        # self.realistic = StableDiffusionControlNetPipeline.from_pretrained(self.stable_diffusion_id_2,
-        #                                                               controlnet=self.controlnet, 
-        #                                                               torch_dtype=dtype,
-        #                                                               safety_checker=StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker", torch_dtype=dtype)).to("cuda")
-        # self.stable_diffusion_id_3 = "SG161222/Realistic_Vision_V1.4"
-        # self.pipe = StableDiffusionControlNetPipeline.from_pretrained(self.stable_diffusion_id_3,
-        #                                                               controlnet=self.controlnet, 
-        #                                                               torch_dtype=dtype,
-        #                                                               safety_checker=StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker", torch_dtype=dtype)).to("cuda")
-        # self.stable_diffusion_id_4 = "s6yx/ReV_Animated"
-        # self.rev = StableDiffusionControlNetPipeline.from_pretrained(self.stable_diffusion_id_4,
-        #                                                               controlnet=self.controlnet, 
-        #                                                               torch_dtype=dtype,
-        #                                                               safety_checker=StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker", torch_dtype=dtype)).to("cuda")
-
-
-
+        self.stable_diffusion_id_1 = "lykon/dreamshaper-xl-v2-turbo"
+        self.dreamshaper_2 = StableDiffusionControlNetPipeline.from_pretrained(self.stable_diffusion_id_1,
+                                                                       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)
 
@@ -80,7 +53,7 @@ class EndpointHandler():
         
 
         # hyperparamters
-        sd_model = data.pop("sd_model", "default")
+        sd_model = data.pop("sd_model", "dreamshaper")
         prompt = data.pop("inputs", None)
         negative_prompt = data.pop("negative_prompt", None)
         image_depth_map = data.pop("image_depth_map", None)