Commit
·
fddc5d7
1
Parent(s):
82cac35
fix model default
Browse files- handler.py +4 -4
handler.py
CHANGED
|
@@ -27,13 +27,13 @@ CONTROLNET_MAPPING = {
|
|
| 27 |
}
|
| 28 |
|
| 29 |
|
| 30 |
-
SD_ID_MAPPING = {
|
| 31 |
-
|
| 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():
|
| 39 |
def __init__(self, path=""):
|
|
@@ -41,7 +41,7 @@ class EndpointHandler():
|
|
| 41 |
self.controlnet = ControlNetModel.from_pretrained(CONTROLNET_MAPPING[self.control_type]["model_id"],torch_dtype=dtype).to(device)
|
| 42 |
|
| 43 |
self.stable_diffusion_id_0 = "Lykon/dreamshaper-8"
|
| 44 |
-
self.
|
| 45 |
controlnet=self.controlnet,
|
| 46 |
torch_dtype=dtype,
|
| 47 |
safety_checker=StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker", torch_dtype=dtype)).to("cuda")
|
|
|
|
| 27 |
}
|
| 28 |
|
| 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():
|
| 39 |
def __init__(self, path=""):
|
|
|
|
| 41 |
self.controlnet = ControlNetModel.from_pretrained(CONTROLNET_MAPPING[self.control_type]["model_id"],torch_dtype=dtype).to(device)
|
| 42 |
|
| 43 |
self.stable_diffusion_id_0 = "Lykon/dreamshaper-8"
|
| 44 |
+
self.default = 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")
|