yamildiego commited on
Commit
fddc5d7
1 Parent(s): 82cac35

fix model default

Browse files
Files changed (1) hide show
  1. handler.py +4 -4
handler.py CHANGED
@@ -27,13 +27,13 @@ CONTROLNET_MAPPING = {
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,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.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")
 
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")