Skip to content

Commit

Permalink
Mega community pipeline (#1561)
Browse files Browse the repository at this point in the history
* Mega community pipeline

* fix
  • Loading branch information
patrickvonplaten committed Dec 6, 2022
1 parent af04479 commit ae4112d
Showing 1 changed file with 3 additions and 0 deletions.
3 changes: 3 additions & 0 deletions examples/community/stable_diffusion_mega.py
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,7 @@ class StableDiffusionMegaPipeline(DiffusionPipeline):
feature_extractor ([`CLIPFeatureExtractor`]):
Model that extracts features from generated images to be used as inputs for the `safety_checker`.
"""
_optional_components = ["safety_checker", "feature_extractor"]

def __init__(
self,
Expand All @@ -60,6 +61,7 @@ def __init__(
scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler],
safety_checker: StableDiffusionSafetyChecker,
feature_extractor: CLIPFeatureExtractor,
requires_safety_checker: bool = True,
):
super().__init__()
if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1:
Expand All @@ -85,6 +87,7 @@ def __init__(
safety_checker=safety_checker,
feature_extractor=feature_extractor,
)
self.register_to_config(requires_safety_checker=requires_safety_checker)

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

0 comments on commit ae4112d

Please sign in to comment.