File tree Expand file tree Collapse file tree 1 file changed +3
-0
lines changed Expand file tree Collapse file tree 1 file changed +3
-0
lines changed Original file line number Diff line number Diff 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 ]:
You can’t perform that action at this time.
0 commit comments