Commit
•
82cac35
1
Parent(s):
f1bbd17
preload only default model
Browse files- handler.py +28 -28
handler.py
CHANGED
|
@@ -29,10 +29,10 @@ CONTROLNET_MAPPING = {
|
|
| 29 |
|
| 30 |
SD_ID_MAPPING = {
|
| 31 |
"default": "Lykon/dreamshaper-8",
|
| 32 |
-
"dreamshaper": "stablediffusionapi/dreamshaper-xl",
|
| 33 |
-
"juggernaut": "stablediffusionapi/juggernaut-xl-v8",
|
| 34 |
-
"realistic":"SG161222/Realistic_Vision_V1.4",
|
| 35 |
-
"rev":"s6yx/ReV_Animated"
|
| 36 |
}
|
| 37 |
|
| 38 |
class EndpointHandler():
|
|
@@ -42,30 +42,30 @@ class EndpointHandler():
|
|
| 42 |
|
| 43 |
self.stable_diffusion_id_0 = "Lykon/dreamshaper-8"
|
| 44 |
self.dreamshaper = StableDiffusionControlNetPipeline.from_pretrained(self.stable_diffusion_id_0,
|
| 45 |
-
|
| 46 |
-
|
| 47 |
-
|
| 48 |
-
|
| 49 |
-
self.stable_diffusion_id_1 = "stablediffusionapi/dreamshaper-xl"
|
| 50 |
-
self.juggernaut = StableDiffusionControlNetPipeline.from_pretrained(self.stable_diffusion_id_1,
|
| 51 |
-
|
| 52 |
-
|
| 53 |
-
|
| 54 |
-
self.stable_diffusion_id_2 = "stablediffusionapi/juggernaut-xl-v8"
|
| 55 |
-
self.realistic = StableDiffusionControlNetPipeline.from_pretrained(self.stable_diffusion_id_2,
|
| 56 |
-
|
| 57 |
-
|
| 58 |
-
|
| 59 |
-
self.stable_diffusion_id_3 = "SG161222/Realistic_Vision_V1.4"
|
| 60 |
-
self.pipe = StableDiffusionControlNetPipeline.from_pretrained(self.stable_diffusion_id_3,
|
| 61 |
-
|
| 62 |
-
|
| 63 |
-
|
| 64 |
-
self.stable_diffusion_id_4 = "s6yx/ReV_Animated"
|
| 65 |
-
self.rev = StableDiffusionControlNetPipeline.from_pretrained(self.stable_diffusion_id_4,
|
| 66 |
-
|
| 67 |
-
|
| 68 |
-
|
| 69 |
|
| 70 |
|
| 71 |
|
|
|
|
| 29 |
|
| 30 |
SD_ID_MAPPING = {
|
| 31 |
"default": "Lykon/dreamshaper-8",
|
| 32 |
+
# "dreamshaper": "stablediffusionapi/dreamshaper-xl",
|
| 33 |
+
# "juggernaut": "stablediffusionapi/juggernaut-xl-v8",
|
| 34 |
+
# "realistic":"SG161222/Realistic_Vision_V1.4",
|
| 35 |
+
# "rev":"s6yx/ReV_Animated"
|
| 36 |
}
|
| 37 |
|
| 38 |
class EndpointHandler():
|
|
|
|
| 42 |
|
| 43 |
self.stable_diffusion_id_0 = "Lykon/dreamshaper-8"
|
| 44 |
self.dreamshaper = StableDiffusionControlNetPipeline.from_pretrained(self.stable_diffusion_id_0,
|
| 45 |
+
controlnet=self.controlnet,
|
| 46 |
+
torch_dtype=dtype,
|
| 47 |
+
safety_checker=StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker", torch_dtype=dtype)).to("cuda")
|
| 48 |
+
|
| 49 |
+
# self.stable_diffusion_id_1 = "stablediffusionapi/dreamshaper-xl"
|
| 50 |
+
# self.juggernaut = StableDiffusionControlNetPipeline.from_pretrained(self.stable_diffusion_id_1,
|
| 51 |
+
# controlnet=self.controlnet,
|
| 52 |
+
# torch_dtype=dtype,
|
| 53 |
+
# safety_checker=StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker", torch_dtype=dtype)).to("cuda")
|
| 54 |
+
# self.stable_diffusion_id_2 = "stablediffusionapi/juggernaut-xl-v8"
|
| 55 |
+
# self.realistic = StableDiffusionControlNetPipeline.from_pretrained(self.stable_diffusion_id_2,
|
| 56 |
+
# controlnet=self.controlnet,
|
| 57 |
+
# torch_dtype=dtype,
|
| 58 |
+
# safety_checker=StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker", torch_dtype=dtype)).to("cuda")
|
| 59 |
+
# self.stable_diffusion_id_3 = "SG161222/Realistic_Vision_V1.4"
|
| 60 |
+
# self.pipe = StableDiffusionControlNetPipeline.from_pretrained(self.stable_diffusion_id_3,
|
| 61 |
+
# controlnet=self.controlnet,
|
| 62 |
+
# torch_dtype=dtype,
|
| 63 |
+
# safety_checker=StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker", torch_dtype=dtype)).to("cuda")
|
| 64 |
+
# self.stable_diffusion_id_4 = "s6yx/ReV_Animated"
|
| 65 |
+
# self.rev = StableDiffusionControlNetPipeline.from_pretrained(self.stable_diffusion_id_4,
|
| 66 |
+
# controlnet=self.controlnet,
|
| 67 |
+
# torch_dtype=dtype,
|
| 68 |
+
# safety_checker=StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker", torch_dtype=dtype)).to("cuda")
|
| 69 |
|
| 70 |
|
| 71 |
|