Skip to content

Commit ae4112d

Browse files
Mega community pipeline (huggingface#1561)
* Mega community pipeline * fix
1 parent af04479 commit ae4112d

File tree

1 file changed

+3
-0
lines changed

1 file changed

+3
-0
lines changed

examples/community/stable_diffusion_mega.py

+3
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ class StableDiffusionMegaPipeline(DiffusionPipeline):
5050
feature_extractor ([`CLIPFeatureExtractor`]):
5151
Model that extracts features from generated images to be used as inputs for the `safety_checker`.
5252
"""
53+
_optional_components = ["safety_checker", "feature_extractor"]
5354

5455
def __init__(
5556
self,
@@ -60,6 +61,7 @@ def __init__(
6061
scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler],
6162
safety_checker: StableDiffusionSafetyChecker,
6263
feature_extractor: CLIPFeatureExtractor,
64+
requires_safety_checker: bool = True,
6365
):
6466
super().__init__()
6567
if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1:
@@ -85,6 +87,7 @@ def __init__(
8587
safety_checker=safety_checker,
8688
feature_extractor=feature_extractor,
8789
)
90+
self.register_to_config(requires_safety_checker=requires_safety_checker)
8891

8992
@property
9093
def components(self) -> Dict[str, Any]:

0 commit comments

Comments
 (0)