Update model.py
Browse files
model.py
CHANGED
|
@@ -12,7 +12,7 @@ from diffusers import (
|
|
| 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,7 +58,8 @@ class Model:
|
|
| 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=
|
|
|
|
| 62 |
)
|
| 63 |
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
|
| 64 |
if self.device.type == "cuda":
|
|
|
|
| 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=None, controlnet=controlnet, torch_dtype=torch.float16
|
| 62 |
+
# StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker")
|
| 63 |
)
|
| 64 |
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
|
| 65 |
if self.device.type == "cuda":
|