yamildiego commited on
Commit
5893a8e
·
1 Parent(s): 5b3371e

add default model

Browse files
Files changed (1) hide show
  1. handler.py +5 -0
handler.py CHANGED
@@ -36,6 +36,11 @@ class EndpointHandler():
36
  controlnet=self.controlnet,
37
  torch_dtype=dtype,
38
  safety_checker=StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker", torch_dtype=dtype)).to("cuda")
 
 
 
 
 
39
  # Define Generator with seed
40
  self.generator = torch.Generator(device=device.type).manual_seed(3)
41
 
 
36
  controlnet=self.controlnet,
37
  torch_dtype=dtype,
38
  safety_checker=StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker", torch_dtype=dtype)).to("cuda")
39
+ self.stable_diffusion_id_1 = "runwayml/stable-diffusion-v1-5"
40
+ self.runwayml = StableDiffusionControlNetPipeline.from_pretrained(self.stable_diffusion_id_1,
41
+ controlnet=self.controlnet,
42
+ torch_dtype=dtype,
43
+ safety_checker=StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker", torch_dtype=dtype)).to("cuda")
44
  # Define Generator with seed
45
  self.generator = torch.Generator(device=device.type).manual_seed(3)
46