Spaces:
Runtime error
Runtime error
| # 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 unittest | |
| from collections import OrderedDict | |
| import torch | |
| from diffusers import ( | |
| AutoPipelineForImage2Image, | |
| AutoPipelineForInpainting, | |
| AutoPipelineForText2Image, | |
| ControlNetModel, | |
| ) | |
| from diffusers.pipelines.auto_pipeline import ( | |
| AUTO_IMAGE2IMAGE_PIPELINES_MAPPING, | |
| AUTO_INPAINT_PIPELINES_MAPPING, | |
| AUTO_TEXT2IMAGE_PIPELINES_MAPPING, | |
| ) | |
| from diffusers.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 | |
| 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)) | |