From df2b548e893ccb8a888467c2508756680df22821 Mon Sep 17 00:00:00 2001 From: Pedro Cuenca Date: Sun, 25 Dec 2022 21:58:45 +0100 Subject: [PATCH] Make safety_checker optional in more pipelines (#1796) * Make safety_checker optional in more pipelines. * Remove inappropriate comment in inpaint pipeline. * InPaint Test: set feature_extractor to None. * Remove import * img2img test: set feature_extractor to None. * inpaint sd2 test: set feature_extractor to None. Co-authored-by: Suraj Patil --- .../alt_diffusion/pipeline_alt_diffusion_img2img.py | 2 +- .../pipelines/paint_by_example/pipeline_paint_by_example.py | 2 ++ .../pipeline_stable_diffusion_image_variation.py | 2 ++ .../stable_diffusion/pipeline_stable_diffusion_img2img.py | 2 +- .../stable_diffusion/pipeline_stable_diffusion_inpaint.py | 2 +- .../stable_diffusion/test_stable_diffusion_img2img.py | 5 ++--- .../stable_diffusion/test_stable_diffusion_inpaint.py | 5 ++--- .../stable_diffusion_2/test_stable_diffusion_inpaint.py | 5 ++--- 8 files changed, 13 insertions(+), 12 deletions(-) diff --git a/src/diffusers/pipelines/alt_diffusion/pipeline_alt_diffusion_img2img.py b/src/diffusers/pipelines/alt_diffusion/pipeline_alt_diffusion_img2img.py index 79c29a98c3..6d712cac9e 100644 --- a/src/diffusers/pipelines/alt_diffusion/pipeline_alt_diffusion_img2img.py +++ b/src/diffusers/pipelines/alt_diffusion/pipeline_alt_diffusion_img2img.py @@ -92,7 +92,7 @@ class AltDiffusionImg2ImgPipeline(DiffusionPipeline): feature_extractor ([`CLIPFeatureExtractor`]): Model that extracts features from generated images to be used as inputs for the `safety_checker`. """ - _optional_components = ["safety_checker"] + _optional_components = ["safety_checker", "feature_extractor"] def __init__( self, diff --git a/src/diffusers/pipelines/paint_by_example/pipeline_paint_by_example.py b/src/diffusers/pipelines/paint_by_example/pipeline_paint_by_example.py index 332165e1fa..f9de7d9110 100644 --- a/src/diffusers/pipelines/paint_by_example/pipeline_paint_by_example.py +++ b/src/diffusers/pipelines/paint_by_example/pipeline_paint_by_example.py @@ -161,6 +161,8 @@ class PaintByExamplePipeline(DiffusionPipeline): feature_extractor ([`CLIPFeatureExtractor`]): Model that extracts features from generated images to be used as inputs for the `safety_checker`. """ + # TODO: feature_extractor is required to encode initial images (if they are in PIL format), + # we should give a descriptive message if the pipeline doesn't have one. _optional_components = ["safety_checker"] def __init__( diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_image_variation.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_image_variation.py index 3bc0a8cfa9..975afddfe0 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_image_variation.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_image_variation.py @@ -65,6 +65,8 @@ class StableDiffusionImageVariationPipeline(DiffusionPipeline): feature_extractor ([`CLIPFeatureExtractor`]): Model that extracts features from generated images to be used as inputs for the `safety_checker`. """ + # TODO: feature_extractor is required to encode images (if they are in PIL format), + # we should give a descriptive message if the pipeline doesn't have one. _optional_components = ["safety_checker"] def __init__( diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py index 2a7def75fa..857ad6f3d5 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py @@ -90,7 +90,7 @@ class StableDiffusionImg2ImgPipeline(DiffusionPipeline): feature_extractor ([`CLIPFeatureExtractor`]): Model that extracts features from generated images to be used as inputs for the `safety_checker`. """ - _optional_components = ["safety_checker"] + _optional_components = ["safety_checker", "feature_extractor"] # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.__init__ def __init__( diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py index deec66a976..985b138dc0 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py @@ -166,7 +166,7 @@ class StableDiffusionInpaintPipeline(DiffusionPipeline): feature_extractor ([`CLIPFeatureExtractor`]): Model that extracts features from generated images to be used as inputs for the `safety_checker`. """ - _optional_components = ["safety_checker"] + _optional_components = ["safety_checker", "feature_extractor"] def __init__( self, diff --git a/tests/pipelines/stable_diffusion/test_stable_diffusion_img2img.py b/tests/pipelines/stable_diffusion/test_stable_diffusion_img2img.py index 53df7f8571..fe6dc729e6 100644 --- a/tests/pipelines/stable_diffusion/test_stable_diffusion_img2img.py +++ b/tests/pipelines/stable_diffusion/test_stable_diffusion_img2img.py @@ -31,7 +31,7 @@ from diffusers import ( ) from diffusers.utils import floats_tensor, load_image, load_numpy, nightly, slow, torch_device from diffusers.utils.testing_utils import require_torch_gpu -from transformers import CLIPImageProcessor, CLIPTextConfig, CLIPTextModel, CLIPTokenizer +from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer from ...test_pipelines_common import PipelineTesterMixin @@ -78,7 +78,6 @@ class StableDiffusionImg2ImgPipelineFastTests(PipelineTesterMixin, unittest.Test ) text_encoder = CLIPTextModel(text_encoder_config) tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") - feature_extractor = CLIPImageProcessor(crop_size=32, size=32) components = { "unet": unet, @@ -87,7 +86,7 @@ class StableDiffusionImg2ImgPipelineFastTests(PipelineTesterMixin, unittest.Test "text_encoder": text_encoder, "tokenizer": tokenizer, "safety_checker": None, - "feature_extractor": feature_extractor, + "feature_extractor": None, } return components diff --git a/tests/pipelines/stable_diffusion/test_stable_diffusion_inpaint.py b/tests/pipelines/stable_diffusion/test_stable_diffusion_inpaint.py index bd4c036934..c763eff2b3 100644 --- a/tests/pipelines/stable_diffusion/test_stable_diffusion_inpaint.py +++ b/tests/pipelines/stable_diffusion/test_stable_diffusion_inpaint.py @@ -32,7 +32,7 @@ from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_inpaint impo from diffusers.utils import floats_tensor, load_image, load_numpy, nightly, slow, torch_device from diffusers.utils.testing_utils import require_torch_gpu from PIL import Image -from transformers import CLIPImageProcessor, CLIPTextConfig, CLIPTextModel, CLIPTokenizer +from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer from ...test_pipelines_common import PipelineTesterMixin @@ -79,7 +79,6 @@ class StableDiffusionInpaintPipelineFastTests(PipelineTesterMixin, unittest.Test ) text_encoder = CLIPTextModel(text_encoder_config) tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") - feature_extractor = CLIPImageProcessor(crop_size=32, size=32) components = { "unet": unet, @@ -88,7 +87,7 @@ class StableDiffusionInpaintPipelineFastTests(PipelineTesterMixin, unittest.Test "text_encoder": text_encoder, "tokenizer": tokenizer, "safety_checker": None, - "feature_extractor": feature_extractor, + "feature_extractor": None, } return components diff --git a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_inpaint.py b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_inpaint.py index 29e14264c0..88157f22de 100644 --- a/tests/pipelines/stable_diffusion_2/test_stable_diffusion_inpaint.py +++ b/tests/pipelines/stable_diffusion_2/test_stable_diffusion_inpaint.py @@ -24,7 +24,7 @@ from diffusers import AutoencoderKL, PNDMScheduler, StableDiffusionInpaintPipeli from diffusers.utils import floats_tensor, load_image, load_numpy, torch_device from diffusers.utils.testing_utils import require_torch_gpu, slow from PIL import Image -from transformers import CLIPImageProcessor, CLIPTextConfig, CLIPTextModel, CLIPTokenizer +from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer from ...test_pipelines_common import PipelineTesterMixin @@ -78,7 +78,6 @@ class StableDiffusion2InpaintPipelineFastTests(PipelineTesterMixin, unittest.Tes ) text_encoder = CLIPTextModel(text_encoder_config) tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") - feature_extractor = CLIPImageProcessor(crop_size=32, size=32) components = { "unet": unet, @@ -87,7 +86,7 @@ class StableDiffusion2InpaintPipelineFastTests(PipelineTesterMixin, unittest.Tes "text_encoder": text_encoder, "tokenizer": tokenizer, "safety_checker": None, - "feature_extractor": feature_extractor, + "feature_extractor": None, } return components