Skip to content

GitLab

Unverified Commit ae4112d2 authored by Patrick von Platen's avatar Patrick von Platen Committed by GitHub
Browse files

Mega community pipeline (#1561) 

* Mega community pipeline

* fix
parent af04479e
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 0 deletions
examples/community/stable_diffusion_mega.py
View file @ ae4112d2
...@@ -50,6 +50,7 @@ class StableDiffusionMegaPipeline(DiffusionPipeline): ...@@ -50,6 +50,7 @@ class StableDiffusionMegaPipeline(DiffusionPipeline):
feature_extractor ([`CLIPFeatureExtractor`]): feature_extractor ([`CLIPFeatureExtractor`]):
Model that extracts features from generated images to be used as inputs for the `safety_checker`. Model that extracts features from generated images to be used as inputs for the `safety_checker`.
""" """
_optional_components = ["safety_checker", "feature_extractor"]
def __init__( def __init__(
self, self,
...@@ -60,6 +61,7 @@ class StableDiffusionMegaPipeline(DiffusionPipeline): ...@@ -60,6 +61,7 @@ class StableDiffusionMegaPipeline(DiffusionPipeline):
scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler], scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler],
safety_checker: StableDiffusionSafetyChecker, safety_checker: StableDiffusionSafetyChecker,
feature_extractor: CLIPFeatureExtractor, feature_extractor: CLIPFeatureExtractor,
requires_safety_checker: bool = True,
): ):
super().__init__() super().__init__()
if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1: if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1:
...@@ -85,6 +87,7 @@ class StableDiffusionMegaPipeline(DiffusionPipeline): ...@@ -85,6 +87,7 @@ class StableDiffusionMegaPipeline(DiffusionPipeline):
safety_checker=safety_checker, safety_checker=safety_checker,
feature_extractor=feature_extractor, feature_extractor=feature_extractor,
) )
self.register_to_config(requires_safety_checker=requires_safety_checker)
@property @property
def components(self) -> Dict[str, Any]: def components(self) -> Dict[str, Any]:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment