Update model.py
Browse files
model.py
CHANGED
@@ -10,10 +10,9 @@ from diffusers import (
|
|
10 |
ControlNetModel,
|
11 |
DiffusionPipeline,
|
12 |
StableDiffusionControlNetPipeline,
|
13 |
-
UniPCMultistepScheduler
|
14 |
-
StableDiffusionSafetyChecker
|
15 |
)
|
16 |
-
|
17 |
from cv_utils import resize_image
|
18 |
from preprocessor import Preprocessor
|
19 |
from settings import MAX_IMAGE_RESOLUTION, MAX_NUM_IMAGES
|
@@ -59,9 +58,8 @@ class Model:
|
|
59 |
model_id = CONTROLNET_MODEL_IDS[task_name]
|
60 |
controlnet = ControlNetModel.from_pretrained(model_id, torch_dtype=torch.float16)
|
61 |
pipe = StableDiffusionControlNetPipeline.from_pretrained(
|
62 |
-
base_model_id, safety_checker=
|
63 |
)
|
64 |
-
pipe.safety_checker = StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker")
|
65 |
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
|
66 |
if self.device.type == "cuda":
|
67 |
pipe.enable_xformers_memory_efficient_attention()
|
|
|
10 |
ControlNetModel,
|
11 |
DiffusionPipeline,
|
12 |
StableDiffusionControlNetPipeline,
|
13 |
+
UniPCMultistepScheduler
|
|
|
14 |
)
|
15 |
+
from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
|
16 |
from cv_utils import resize_image
|
17 |
from preprocessor import Preprocessor
|
18 |
from settings import MAX_IMAGE_RESOLUTION, MAX_NUM_IMAGES
|
|
|
58 |
model_id = CONTROLNET_MODEL_IDS[task_name]
|
59 |
controlnet = ControlNetModel.from_pretrained(model_id, torch_dtype=torch.float16)
|
60 |
pipe = StableDiffusionControlNetPipeline.from_pretrained(
|
61 |
+
base_model_id, safety_checker=StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker"), controlnet=controlnet, torch_dtype=torch.float16
|
62 |
)
|
|
|
63 |
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
|
64 |
if self.device.type == "cuda":
|
65 |
pipe.enable_xformers_memory_efficient_attention()
|