





'''simple docstring'''






# This model implementation is heavily inspired by https://github.com/haofanwang/ControlNet-for-Diffusers/

import gc
import random
import tempfile
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    ControlNetModel,
    DDIMScheduler,
    StableDiffusionControlNetImgaImgPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_controlnet import MultiControlNetModel
from diffusers.utils import floats_tensor, load_image, load_numpy, randn_tensor, slow, torch_device
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
)


enable_full_determinism()




class  UpperCAmelCase     (			snake_case_							,		snake_case_							,		snake_case_							,		unittest.TestCase   ):
      _lowercase:			Optional[int]						     =   StableDiffusionControlNetImgaImgPipeline
      _lowercase:			Tuple						     =   TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'''height''', '''width'''}
      _lowercase:			str						     =   TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
      _lowercase:			Tuple						     =   IMAGE_TO_IMAGE_IMAGE_PARAMS.union({'''control_image'''}   )
      _lowercase:			Union[str, Any]						     =   IMAGE_TO_IMAGE_IMAGE_PARAMS




      def 				lowercase__    (			self			: List[str]  )				->      List[str]:
            torch.manual_seed(0  )
            _lowerCAmelCase							     =							UNetaDConditionModel(
                block_out_channels=(32, 64)			,			layers_per_block=2			,			sample_size=32			,			in_channels=4			,			out_channels=4			,			down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")			,			up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")			,			cross_attention_dim=32			,			)
            torch.manual_seed(0  )
            _lowerCAmelCase							     =							ControlNetModel(
                block_out_channels=(32, 64)			,			layers_per_block=2			,			in_channels=4			,			down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")			,			cross_attention_dim=32			,			conditioning_embedding_out_channels=(16, 32)			,			)
            torch.manual_seed(0  )
            _lowerCAmelCase							     =							DDIMScheduler(
                beta_start=0.0_00_85			,			beta_end=0.0_12			,			beta_schedule="""scaled_linear"""			,			clip_sample=__snake_case			,			set_alpha_to_one=__snake_case			,			)
            torch.manual_seed(0  )
            _lowerCAmelCase							     =							AutoencoderKL(
                block_out_channels=[32, 64]			,			in_channels=3			,			out_channels=3			,			down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]			,			up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]			,			latent_channels=4			,			)
            torch.manual_seed(0  )
            _lowerCAmelCase							     =							CLIPTextConfig(
                bos_token_id=0			,			eos_token_id=2			,			hidden_size=32			,			intermediate_size=37			,			layer_norm_eps=1E-0_5			,			num_attention_heads=4			,			num_hidden_layers=5			,			pad_token_id=1			,			vocab_size=10_00			,			)
            _lowerCAmelCase							     =							CLIPTextModel(__snake_case  )
            _lowerCAmelCase							     =							CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""  )

            _lowerCAmelCase							     =							{
                """unet""": unet,
                """controlnet""": controlnet,
                """scheduler""": scheduler,
                """vae""": vae,
                """text_encoder""": text_encoder,
                """tokenizer""": tokenizer,
                """safety_checker""": None,
                """feature_extractor""": None,
            }
            return components




      def 				lowercase__    (			self			: Any			,			__snake_case			: str			,			__snake_case			: Any=0  )				->      str:
            if str(__snake_case  ).startswith("""mps"""  ):
                  _lowerCAmelCase							     =							torch.manual_seed(__snake_case  )
            else:
                  _lowerCAmelCase							     =							torch.Generator(device=__snake_case  ).manual_seed(__snake_case  )

            _lowerCAmelCase							     =							2
            _lowerCAmelCase							     =							randn_tensor(
                (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)			,			generator=__snake_case			,			device=torch.device(__snake_case  )			,			)
            _lowerCAmelCase							     =							floats_tensor(control_image.shape			,			rng=random.Random(__snake_case  )  ).to(__snake_case  )
            _lowerCAmelCase							     =							image.cpu().permute(0			,			2			,			3			,			1  )[0]
            _lowerCAmelCase							     =							Image.fromarray(np.uinta(__snake_case  )  ).convert("""RGB"""  ).resize((64, 64)  )
            _lowerCAmelCase							     =							{
                """prompt""": """A painting of a squirrel eating a burger""",
                """generator""": generator,
                """num_inference_steps""": 2,
                """guidance_scale""": 6.0,
                """output_type""": """numpy""",
                """image""": image,
                """control_image""": control_image,
            }

            return inputs




      def 				lowercase__    (			self			: Optional[int]  )				->      List[Any]:
            return self._test_attention_slicing_forward_pass(expected_max_diff=2E-3  )




      @unittest.skipIf(
          torch_device != """cuda""" or not is_xformers_available()			,			reason="""XFormers attention is only available with CUDA and `xformers` installed"""			,			)
      def 				lowercase__    (			self			: Tuple  )				->      Optional[int]:
            self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3  )




      def 				lowercase__    (			self			: Tuple  )				->      Optional[int]:
            self._test_inference_batch_single_identical(expected_max_diff=2E-3  )




class  UpperCAmelCase     (			snake_case_							,		snake_case_							,		unittest.TestCase   ):
      _lowercase:			Any						     =   StableDiffusionControlNetImgaImgPipeline
      _lowercase:			Dict						     =   TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'''height''', '''width'''}
      _lowercase:			List[Any]						     =   TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
      _lowercase:			Any						     =   frozenset([]   )  # TO_DO: add image_params once refactored VaeImageProcessor.preprocess




      def 				lowercase__    (			self			: Optional[Any]  )				->      Union[str, Any]:
            torch.manual_seed(0  )
            _lowerCAmelCase							     =							UNetaDConditionModel(
                block_out_channels=(32, 64)			,			layers_per_block=2			,			sample_size=32			,			in_channels=4			,			out_channels=4			,			down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")			,			up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")			,			cross_attention_dim=32			,			)
            torch.manual_seed(0  )

            def init_weights(__snake_case			: Optional[Any]  ):
                  if isinstance(__snake_case			,			torch.nn.Convad  ):
                        torch.nn.init.normal(m.weight  )
                        m.bias.data.fill_(1.0  )

            _lowerCAmelCase							     =							ControlNetModel(
                block_out_channels=(32, 64)			,			layers_per_block=2			,			in_channels=4			,			down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")			,			cross_attention_dim=32			,			conditioning_embedding_out_channels=(16, 32)			,			)
            controlneta.controlnet_down_blocks.apply(__snake_case  )

            torch.manual_seed(0  )
            _lowerCAmelCase							     =							ControlNetModel(
                block_out_channels=(32, 64)			,			layers_per_block=2			,			in_channels=4			,			down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")			,			cross_attention_dim=32			,			conditioning_embedding_out_channels=(16, 32)			,			)
            controlneta.controlnet_down_blocks.apply(__snake_case  )

            torch.manual_seed(0  )
            _lowerCAmelCase							     =							DDIMScheduler(
                beta_start=0.0_00_85			,			beta_end=0.0_12			,			beta_schedule="""scaled_linear"""			,			clip_sample=__snake_case			,			set_alpha_to_one=__snake_case			,			)
            torch.manual_seed(0  )
            _lowerCAmelCase							     =							AutoencoderKL(
                block_out_channels=[32, 64]			,			in_channels=3			,			out_channels=3			,			down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]			,			up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]			,			latent_channels=4			,			)
            torch.manual_seed(0  )
            _lowerCAmelCase							     =							CLIPTextConfig(
                bos_token_id=0			,			eos_token_id=2			,			hidden_size=32			,			intermediate_size=37			,			layer_norm_eps=1E-0_5			,			num_attention_heads=4			,			num_hidden_layers=5			,			pad_token_id=1			,			vocab_size=10_00			,			)
            _lowerCAmelCase							     =							CLIPTextModel(__snake_case  )
            _lowerCAmelCase							     =							CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""  )

            _lowerCAmelCase							     =							MultiControlNetModel([controlneta, controlneta]  )

            _lowerCAmelCase							     =							{
                """unet""": unet,
                """controlnet""": controlnet,
                """scheduler""": scheduler,
                """vae""": vae,
                """text_encoder""": text_encoder,
                """tokenizer""": tokenizer,
                """safety_checker""": None,
                """feature_extractor""": None,
            }
            return components




      def 				lowercase__    (			self			: Tuple			,			__snake_case			: int			,			__snake_case			: List[str]=0  )				->      Union[str, Any]:
            if str(__snake_case  ).startswith("""mps"""  ):
                  _lowerCAmelCase							     =							torch.manual_seed(__snake_case  )
            else:
                  _lowerCAmelCase							     =							torch.Generator(device=__snake_case  ).manual_seed(__snake_case  )

            _lowerCAmelCase							     =							2

            _lowerCAmelCase							     =							[
                randn_tensor(
                    (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)			,			generator=__snake_case			,			device=torch.device(__snake_case  )			,			),
                randn_tensor(
                    (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)			,			generator=__snake_case			,			device=torch.device(__snake_case  )			,			),
            ]

            _lowerCAmelCase							     =							floats_tensor(control_image[0].shape			,			rng=random.Random(__snake_case  )  ).to(__snake_case  )
            _lowerCAmelCase							     =							image.cpu().permute(0			,			2			,			3			,			1  )[0]
            _lowerCAmelCase							     =							Image.fromarray(np.uinta(__snake_case  )  ).convert("""RGB"""  ).resize((64, 64)  )
            _lowerCAmelCase							     =							{
                """prompt""": """A painting of a squirrel eating a burger""",
                """generator""": generator,
                """num_inference_steps""": 2,
                """guidance_scale""": 6.0,
                """output_type""": """numpy""",
                """image""": image,
                """control_image""": control_image,
            }

            return inputs




      def 				lowercase__    (			self			: List[str]  )				->      Dict:
            _lowerCAmelCase							     =							self.get_dummy_components()
            _lowerCAmelCase							     =							self.pipeline_class(**__snake_case  )
            pipe.to(__snake_case  )

            _lowerCAmelCase							     =							10.0
            _lowerCAmelCase							     =							4

            _lowerCAmelCase							     =							self.get_dummy_inputs(__snake_case  )
            _lowerCAmelCase							     =							steps
            _lowerCAmelCase							     =							scale
            _lowerCAmelCase							     =							pipe(**__snake_case  )[0]

            _lowerCAmelCase							     =							self.get_dummy_inputs(__snake_case  )
            _lowerCAmelCase							     =							steps
            _lowerCAmelCase							     =							scale
            _lowerCAmelCase							     =							pipe(**__snake_case			,			control_guidance_start=0.1			,			control_guidance_end=0.2  )[0]

            _lowerCAmelCase							     =							self.get_dummy_inputs(__snake_case  )
            _lowerCAmelCase							     =							steps
            _lowerCAmelCase							     =							scale
            _lowerCAmelCase							     =							pipe(**__snake_case			,			control_guidance_start=[0.1, 0.3]			,			control_guidance_end=[0.2, 0.7]  )[0]

            _lowerCAmelCase							     =							self.get_dummy_inputs(__snake_case  )
            _lowerCAmelCase							     =							steps
            _lowerCAmelCase							     =							scale
            _lowerCAmelCase							     =							pipe(**__snake_case			,			control_guidance_start=0.4			,			control_guidance_end=[0.5, 0.8]  )[0]

            # make sure that all outputs are different
            assert np.sum(np.abs(output_a - output_a  )  ) > 1E-3
            assert np.sum(np.abs(output_a - output_a  )  ) > 1E-3
            assert np.sum(np.abs(output_a - output_a  )  ) > 1E-3




      def 				lowercase__    (			self			: int  )				->      str:
            return self._test_attention_slicing_forward_pass(expected_max_diff=2E-3  )




      @unittest.skipIf(
          torch_device != """cuda""" or not is_xformers_available()			,			reason="""XFormers attention is only available with CUDA and `xformers` installed"""			,			)
      def 				lowercase__    (			self			: Optional[Any]  )				->      Dict:
            self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3  )




      def 				lowercase__    (			self			: int  )				->      str:
            self._test_inference_batch_single_identical(expected_max_diff=2E-3  )




      def 				lowercase__    (			self			: Union[str, Any]  )				->      Optional[Any]:
            _lowerCAmelCase							     =							self.get_dummy_components()
            _lowerCAmelCase							     =							self.pipeline_class(**__snake_case  )
            pipe.to(__snake_case  )
            pipe.set_progress_bar_config(disable=__snake_case  )
            with tempfile.TemporaryDirectory() as tmpdir:
                  try:
                        # save_pretrained is not implemented for Multi-ControlNet
                        pipe.save_pretrained(__snake_case  )
                  except NotImplementedError:
                        pass




@slow
@require_torch_gpu
class  UpperCAmelCase     (			unittest.TestCase   ):




      def 				lowercase__    (			self			: Union[str, Any]  )				->      int:
            super().tearDown()
            gc.collect()
            torch.cuda.empty_cache()




      def 				lowercase__    (			self			: List[str]  )				->      Any:
            _lowerCAmelCase							     =							ControlNetModel.from_pretrained("""lllyasviel/sd-controlnet-canny"""  )

            _lowerCAmelCase							     =							StableDiffusionControlNetImgaImgPipeline.from_pretrained(
                """runwayml/stable-diffusion-v1-5"""			,			safety_checker=__snake_case			,			controlnet=__snake_case  )
            pipe.enable_model_cpu_offload()
            pipe.set_progress_bar_config(disable=__snake_case  )

            _lowerCAmelCase							     =							torch.Generator(device="""cpu"""  ).manual_seed(0  )
            _lowerCAmelCase							     =							"""evil space-punk bird"""
            _lowerCAmelCase							     =							load_image(
                """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/bird_canny.png"""  ).resize((5_12, 5_12)  )
            _lowerCAmelCase							     =							load_image(
                """https://huggingface.co/lllyasviel/sd-controlnet-canny/resolve/main/images/bird.png"""  ).resize((5_12, 5_12)  )

            _lowerCAmelCase							     =							pipe(
                __snake_case			,			__snake_case			,			control_image=__snake_case			,			generator=__snake_case			,			output_type="""np"""			,			num_inference_steps=50			,			strength=0.6			,			)

            _lowerCAmelCase							     =							output.images[0]

            assert image.shape == (5_12, 5_12, 3)

            _lowerCAmelCase							     =							load_numpy(
                """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/img2img.npy"""  )

            assert np.abs(expected_image - image  ).max() < 9E-2







"""simple docstring"""







import sys

from .dependency_versions_table import deps
from .utils.versions import require_version, require_version_core


# define which module versions we always want to check at run time
# (usually the ones defined in `install_requires` in setup.py)
#
# order specific notes:
# - tqdm must be checked before tokenizers

__SCREAMING_SNAKE_CASE    	="python tqdm regex requests packaging filelock numpy tokenizers".split()
if sys.version_info < (3, 7):
       pkgs_to_check_at_runtime.append("dataclasses")
if sys.version_info < (3, 8):
       pkgs_to_check_at_runtime.append("importlib_metadata")

for pkg in pkgs_to_check_at_runtime:
       if pkg in deps:
              if pkg == "tokenizers":
                     # must be loaded here, or else tqdm check may fail
                     from .utils import is_tokenizers_available

                     if not is_tokenizers_available():
                            continue  # not required, check version only if installed

              require_version_core(deps[pkg])
       else:
              raise ValueError(F"can't find {pkg} in {deps.keys()}, check dependency_versions_table.py")



def 			lowercase__(       __SCREAMING_SNAKE_CASE							:     int     ,			__SCREAMING_SNAKE_CASE							:     Union[str, Any]=None      ):
     require_version(deps[pkg]     ,			__SCREAMING_SNAKE_CASE      )

