Commit
·
5893a8e
1
Parent(s):
5b3371e
add default model
Browse files- 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 |
|