Spaces:
No application file
No application file
| # coding=utf-8 | |
| # Copyright 2023 HuggingFace Inc. | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| import gc | |
| import os | |
| import shutil | |
| import unittest | |
| from collections import OrderedDict | |
| from pathlib import Path | |
| import torch | |
| from diffusers import ( | |
| AutoPipelineForImage2Image, | |
| AutoPipelineForInpainting, | |
| AutoPipelineForText2Image, | |
| ControlNetModel, | |
| DiffusionPipeline, | |
| ) | |
| from diffusers.pipelines.auto_pipeline import ( | |
| AUTO_IMAGE2IMAGE_PIPELINES_MAPPING, | |
| AUTO_INPAINT_PIPELINES_MAPPING, | |
| AUTO_TEXT2IMAGE_PIPELINES_MAPPING, | |
| ) | |
| from diffusers.utils.testing_utils import slow | |
| PRETRAINED_MODEL_REPO_MAPPING = OrderedDict( | |
| [ | |
| ("stable-diffusion", "runwayml/stable-diffusion-v1-5"), | |
| ("if", "DeepFloyd/IF-I-XL-v1.0"), | |
| ("kandinsky", "kandinsky-community/kandinsky-2-1"), | |
| ("kandinsky22", "kandinsky-community/kandinsky-2-2-decoder"), | |
| ] | |
| ) | |
| class AutoPipelineFastTest(unittest.TestCase): | |
| def test_from_pipe_consistent(self): | |
| pipe = AutoPipelineForText2Image.from_pretrained( | |
| "hf-internal-testing/tiny-stable-diffusion-pipe", requires_safety_checker=False | |
| ) | |
| original_config = dict(pipe.config) | |
| pipe = AutoPipelineForImage2Image.from_pipe(pipe) | |
| assert dict(pipe.config) == original_config | |
| pipe = AutoPipelineForText2Image.from_pipe(pipe) | |
| assert dict(pipe.config) == original_config | |
| def test_from_pipe_override(self): | |
| pipe = AutoPipelineForText2Image.from_pretrained( | |
| "hf-internal-testing/tiny-stable-diffusion-pipe", requires_safety_checker=False | |
| ) | |
| pipe = AutoPipelineForImage2Image.from_pipe(pipe, requires_safety_checker=True) | |
| assert pipe.config.requires_safety_checker is True | |
| pipe = AutoPipelineForText2Image.from_pipe(pipe, requires_safety_checker=True) | |
| assert pipe.config.requires_safety_checker is True | |
| def test_from_pipe_consistent_sdxl(self): | |
| pipe = AutoPipelineForImage2Image.from_pretrained( | |
| "hf-internal-testing/tiny-stable-diffusion-xl-pipe", | |
| requires_aesthetics_score=True, | |
| force_zeros_for_empty_prompt=False, | |
| ) | |
| original_config = dict(pipe.config) | |
| pipe = AutoPipelineForText2Image.from_pipe(pipe) | |
| pipe = AutoPipelineForImage2Image.from_pipe(pipe) | |
| assert dict(pipe.config) == original_config | |
| def test_kwargs_local_files_only(self): | |
| repo = "hf-internal-testing/tiny-stable-diffusion-torch" | |
| tmpdirname = DiffusionPipeline.download(repo) | |
| tmpdirname = Path(tmpdirname) | |
| # edit commit_id to so that it's not the latest commit | |
| commit_id = tmpdirname.name | |
| new_commit_id = commit_id + "hug" | |
| ref_dir = tmpdirname.parent.parent / "refs/main" | |
| with open(ref_dir, "w") as f: | |
| f.write(new_commit_id) | |
| new_tmpdirname = tmpdirname.parent / new_commit_id | |
| os.rename(tmpdirname, new_tmpdirname) | |
| try: | |
| AutoPipelineForText2Image.from_pretrained(repo, local_files_only=True) | |
| except OSError: | |
| assert False, "not able to load local files" | |
| shutil.rmtree(tmpdirname.parent.parent) | |
| def test_from_pipe_controlnet_text2img(self): | |
| pipe = AutoPipelineForText2Image.from_pretrained("hf-internal-testing/tiny-stable-diffusion-pipe") | |
| controlnet = ControlNetModel.from_pretrained("hf-internal-testing/tiny-controlnet") | |
| pipe = AutoPipelineForText2Image.from_pipe(pipe, controlnet=controlnet) | |
| assert pipe.__class__.__name__ == "StableDiffusionControlNetPipeline" | |
| assert "controlnet" in pipe.components | |
| pipe = AutoPipelineForText2Image.from_pipe(pipe, controlnet=None) | |
| assert pipe.__class__.__name__ == "StableDiffusionPipeline" | |
| assert "controlnet" not in pipe.components | |
| def test_from_pipe_controlnet_img2img(self): | |
| pipe = AutoPipelineForImage2Image.from_pretrained("hf-internal-testing/tiny-stable-diffusion-pipe") | |
| controlnet = ControlNetModel.from_pretrained("hf-internal-testing/tiny-controlnet") | |
| pipe = AutoPipelineForImage2Image.from_pipe(pipe, controlnet=controlnet) | |
| assert pipe.__class__.__name__ == "StableDiffusionControlNetImg2ImgPipeline" | |
| assert "controlnet" in pipe.components | |
| pipe = AutoPipelineForImage2Image.from_pipe(pipe, controlnet=None) | |
| assert pipe.__class__.__name__ == "StableDiffusionImg2ImgPipeline" | |
| assert "controlnet" not in pipe.components | |
| def test_from_pipe_controlnet_inpaint(self): | |
| pipe = AutoPipelineForInpainting.from_pretrained("hf-internal-testing/tiny-stable-diffusion-torch") | |
| controlnet = ControlNetModel.from_pretrained("hf-internal-testing/tiny-controlnet") | |
| pipe = AutoPipelineForInpainting.from_pipe(pipe, controlnet=controlnet) | |
| assert pipe.__class__.__name__ == "StableDiffusionControlNetInpaintPipeline" | |
| assert "controlnet" in pipe.components | |
| pipe = AutoPipelineForInpainting.from_pipe(pipe, controlnet=None) | |
| assert pipe.__class__.__name__ == "StableDiffusionInpaintPipeline" | |
| assert "controlnet" not in pipe.components | |
| def test_from_pipe_controlnet_new_task(self): | |
| pipe_text2img = AutoPipelineForText2Image.from_pretrained("hf-internal-testing/tiny-stable-diffusion-torch") | |
| controlnet = ControlNetModel.from_pretrained("hf-internal-testing/tiny-controlnet") | |
| pipe_control_img2img = AutoPipelineForImage2Image.from_pipe(pipe_text2img, controlnet=controlnet) | |
| assert pipe_control_img2img.__class__.__name__ == "StableDiffusionControlNetImg2ImgPipeline" | |
| assert "controlnet" in pipe_control_img2img.components | |
| pipe_inpaint = AutoPipelineForInpainting.from_pipe(pipe_control_img2img, controlnet=None) | |
| assert pipe_inpaint.__class__.__name__ == "StableDiffusionInpaintPipeline" | |
| assert "controlnet" not in pipe_inpaint.components | |
| # testing `from_pipe` for text2img controlnet | |
| ## 1. from a different controlnet pipe, without controlnet argument | |
| pipe_control_text2img = AutoPipelineForText2Image.from_pipe(pipe_control_img2img) | |
| assert pipe_control_text2img.__class__.__name__ == "StableDiffusionControlNetPipeline" | |
| assert "controlnet" in pipe_control_text2img.components | |
| ## 2. from a different controlnet pipe, with controlnet argument | |
| pipe_control_text2img = AutoPipelineForText2Image.from_pipe(pipe_control_img2img, controlnet=controlnet) | |
| assert pipe_control_text2img.__class__.__name__ == "StableDiffusionControlNetPipeline" | |
| assert "controlnet" in pipe_control_text2img.components | |
| ## 3. from same controlnet pipeline class, with a different controlnet component | |
| pipe_control_text2img = AutoPipelineForText2Image.from_pipe(pipe_control_text2img, controlnet=controlnet) | |
| assert pipe_control_text2img.__class__.__name__ == "StableDiffusionControlNetPipeline" | |
| assert "controlnet" in pipe_control_text2img.components | |
| # testing from_pipe for inpainting | |
| ## 1. from a different controlnet pipeline class | |
| pipe_control_inpaint = AutoPipelineForInpainting.from_pipe(pipe_control_img2img) | |
| assert pipe_control_inpaint.__class__.__name__ == "StableDiffusionControlNetInpaintPipeline" | |
| assert "controlnet" in pipe_control_inpaint.components | |
| ## from a different controlnet pipe, with a different controlnet | |
| pipe_control_inpaint = AutoPipelineForInpainting.from_pipe(pipe_control_img2img, controlnet=controlnet) | |
| assert pipe_control_inpaint.__class__.__name__ == "StableDiffusionControlNetInpaintPipeline" | |
| assert "controlnet" in pipe_control_inpaint.components | |
| ## from same controlnet pipe, with a different controlnet | |
| pipe_control_inpaint = AutoPipelineForInpainting.from_pipe(pipe_control_inpaint, controlnet=controlnet) | |
| assert pipe_control_inpaint.__class__.__name__ == "StableDiffusionControlNetInpaintPipeline" | |
| assert "controlnet" in pipe_control_inpaint.components | |
| # testing from_pipe from img2img controlnet | |
| ## from a different controlnet pipe, without controlnet argument | |
| pipe_control_img2img = AutoPipelineForImage2Image.from_pipe(pipe_control_text2img) | |
| assert pipe_control_img2img.__class__.__name__ == "StableDiffusionControlNetImg2ImgPipeline" | |
| assert "controlnet" in pipe_control_img2img.components | |
| # from a different controlnet pipe, with a different controlnet component | |
| pipe_control_img2img = AutoPipelineForImage2Image.from_pipe(pipe_control_text2img, controlnet=controlnet) | |
| assert pipe_control_img2img.__class__.__name__ == "StableDiffusionControlNetImg2ImgPipeline" | |
| assert "controlnet" in pipe_control_img2img.components | |
| # from same controlnet pipeline class, with a different controlnet | |
| pipe_control_img2img = AutoPipelineForImage2Image.from_pipe(pipe_control_img2img, controlnet=controlnet) | |
| assert pipe_control_img2img.__class__.__name__ == "StableDiffusionControlNetImg2ImgPipeline" | |
| assert "controlnet" in pipe_control_img2img.components | |
| class AutoPipelineIntegrationTest(unittest.TestCase): | |
| def test_pipe_auto(self): | |
| for model_name, model_repo in PRETRAINED_MODEL_REPO_MAPPING.items(): | |
| # test txt2img | |
| pipe_txt2img = AutoPipelineForText2Image.from_pretrained( | |
| model_repo, variant="fp16", torch_dtype=torch.float16 | |
| ) | |
| self.assertIsInstance(pipe_txt2img, AUTO_TEXT2IMAGE_PIPELINES_MAPPING[model_name]) | |
| pipe_to = AutoPipelineForText2Image.from_pipe(pipe_txt2img) | |
| self.assertIsInstance(pipe_to, AUTO_TEXT2IMAGE_PIPELINES_MAPPING[model_name]) | |
| pipe_to = AutoPipelineForImage2Image.from_pipe(pipe_txt2img) | |
| self.assertIsInstance(pipe_to, AUTO_IMAGE2IMAGE_PIPELINES_MAPPING[model_name]) | |
| if "kandinsky" not in model_name: | |
| pipe_to = AutoPipelineForInpainting.from_pipe(pipe_txt2img) | |
| self.assertIsInstance(pipe_to, AUTO_INPAINT_PIPELINES_MAPPING[model_name]) | |
| del pipe_txt2img, pipe_to | |
| gc.collect() | |
| # test img2img | |
| pipe_img2img = AutoPipelineForImage2Image.from_pretrained( | |
| model_repo, variant="fp16", torch_dtype=torch.float16 | |
| ) | |
| self.assertIsInstance(pipe_img2img, AUTO_IMAGE2IMAGE_PIPELINES_MAPPING[model_name]) | |
| pipe_to = AutoPipelineForText2Image.from_pipe(pipe_img2img) | |
| self.assertIsInstance(pipe_to, AUTO_TEXT2IMAGE_PIPELINES_MAPPING[model_name]) | |
| pipe_to = AutoPipelineForImage2Image.from_pipe(pipe_img2img) | |
| self.assertIsInstance(pipe_to, AUTO_IMAGE2IMAGE_PIPELINES_MAPPING[model_name]) | |
| if "kandinsky" not in model_name: | |
| pipe_to = AutoPipelineForInpainting.from_pipe(pipe_img2img) | |
| self.assertIsInstance(pipe_to, AUTO_INPAINT_PIPELINES_MAPPING[model_name]) | |
| del pipe_img2img, pipe_to | |
| gc.collect() | |
| # test inpaint | |
| if "kandinsky" not in model_name: | |
| pipe_inpaint = AutoPipelineForInpainting.from_pretrained( | |
| model_repo, variant="fp16", torch_dtype=torch.float16 | |
| ) | |
| self.assertIsInstance(pipe_inpaint, AUTO_INPAINT_PIPELINES_MAPPING[model_name]) | |
| pipe_to = AutoPipelineForText2Image.from_pipe(pipe_inpaint) | |
| self.assertIsInstance(pipe_to, AUTO_TEXT2IMAGE_PIPELINES_MAPPING[model_name]) | |
| pipe_to = AutoPipelineForImage2Image.from_pipe(pipe_inpaint) | |
| self.assertIsInstance(pipe_to, AUTO_IMAGE2IMAGE_PIPELINES_MAPPING[model_name]) | |
| pipe_to = AutoPipelineForInpainting.from_pipe(pipe_inpaint) | |
| self.assertIsInstance(pipe_to, AUTO_INPAINT_PIPELINES_MAPPING[model_name]) | |
| del pipe_inpaint, pipe_to | |
| gc.collect() | |
| def test_from_pipe_consistent(self): | |
| for model_name, model_repo in PRETRAINED_MODEL_REPO_MAPPING.items(): | |
| if model_name in ["kandinsky", "kandinsky22"]: | |
| auto_pipes = [AutoPipelineForText2Image, AutoPipelineForImage2Image] | |
| else: | |
| auto_pipes = [AutoPipelineForText2Image, AutoPipelineForImage2Image, AutoPipelineForInpainting] | |
| # test from_pretrained | |
| for pipe_from_class in auto_pipes: | |
| pipe_from = pipe_from_class.from_pretrained(model_repo, variant="fp16", torch_dtype=torch.float16) | |
| pipe_from_config = dict(pipe_from.config) | |
| for pipe_to_class in auto_pipes: | |
| pipe_to = pipe_to_class.from_pipe(pipe_from) | |
| self.assertEqual(dict(pipe_to.config), pipe_from_config) | |
| del pipe_from, pipe_to | |
| gc.collect() | |
| def test_controlnet(self): | |
| # test from_pretrained | |
| model_repo = "runwayml/stable-diffusion-v1-5" | |
| controlnet_repo = "lllyasviel/sd-controlnet-canny" | |
| controlnet = ControlNetModel.from_pretrained(controlnet_repo, torch_dtype=torch.float16) | |
| pipe_txt2img = AutoPipelineForText2Image.from_pretrained( | |
| model_repo, controlnet=controlnet, torch_dtype=torch.float16 | |
| ) | |
| self.assertIsInstance(pipe_txt2img, AUTO_TEXT2IMAGE_PIPELINES_MAPPING["stable-diffusion-controlnet"]) | |
| pipe_img2img = AutoPipelineForImage2Image.from_pretrained( | |
| model_repo, controlnet=controlnet, torch_dtype=torch.float16 | |
| ) | |
| self.assertIsInstance(pipe_img2img, AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["stable-diffusion-controlnet"]) | |
| pipe_inpaint = AutoPipelineForInpainting.from_pretrained( | |
| model_repo, controlnet=controlnet, torch_dtype=torch.float16 | |
| ) | |
| self.assertIsInstance(pipe_inpaint, AUTO_INPAINT_PIPELINES_MAPPING["stable-diffusion-controlnet"]) | |
| # test from_pipe | |
| for pipe_from in [pipe_txt2img, pipe_img2img, pipe_inpaint]: | |
| pipe_to = AutoPipelineForText2Image.from_pipe(pipe_from) | |
| self.assertIsInstance(pipe_to, AUTO_TEXT2IMAGE_PIPELINES_MAPPING["stable-diffusion-controlnet"]) | |
| self.assertEqual(dict(pipe_to.config), dict(pipe_txt2img.config)) | |
| pipe_to = AutoPipelineForImage2Image.from_pipe(pipe_from) | |
| self.assertIsInstance(pipe_to, AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["stable-diffusion-controlnet"]) | |
| self.assertEqual(dict(pipe_to.config), dict(pipe_img2img.config)) | |
| pipe_to = AutoPipelineForInpainting.from_pipe(pipe_from) | |
| self.assertIsInstance(pipe_to, AUTO_INPAINT_PIPELINES_MAPPING["stable-diffusion-controlnet"]) | |
| self.assertEqual(dict(pipe_to.config), dict(pipe_inpaint.config)) | |