

# 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   lowercase_   (  __lowercase       ,							__lowercase       ,							__lowercase       ,							unittest.TestCase					):
      UpperCamelCase_      :		List[str]										=    StableDiffusionControlNetImgaImgPipeline
      UpperCamelCase_      :		Tuple										=    TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"}
      UpperCamelCase_      :		Optional[Any]										=    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
      UpperCamelCase_      :		Dict										=    IMAGE_TO_IMAGE_IMAGE_PARAMS.union({"control_image"}					)
      UpperCamelCase_      :		Optional[Any]										=    IMAGE_TO_IMAGE_IMAGE_PARAMS



      def 						UpperCamelCase_		(      self				:    List[Any]		)   ->					int:
        torch.manual_seed(0		)
        _snake_case						=							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		)
        _snake_case						=							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		)
        _snake_case						=							DDIMScheduler(
            beta_start=0.0_0085      , beta_end=0.012      , beta_schedule='''scaled_linear'''      , clip_sample=A__      , set_alpha_to_one=A__      , )
        torch.manual_seed(0		)
        _snake_case						=							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		)
        _snake_case						=							CLIPTextConfig(
            bos_token_id=0      , eos_token_id=2      , hidden_size=32      , intermediate_size=37      , layer_norm_eps=1e-05      , num_attention_heads=4      , num_hidden_layers=5      , pad_token_id=1      , vocab_size=1000      , )
        _snake_case						=							CLIPTextModel(A__		)
        _snake_case						=							CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''		)

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



      def 						UpperCamelCase_		(      self				:    Optional[Any]      , A__				:    List[Any]      , A__				:    Dict=0		)   ->					Optional[int]:
        if str(A__		).startswith('''mps'''		):
          _snake_case						=							torch.manual_seed(A__		)
        else:
          _snake_case						=							torch.Generator(device=A__		).manual_seed(A__		)

        _snake_case						=							2
        _snake_case						=							randn_tensor(
            (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)      , generator=A__      , device=torch.device(A__		)      , )
        _snake_case						=							floats_tensor(control_image.shape      , rng=random.Random(A__		)		).to(A__		)
        _snake_case						=							image.cpu().permute(0      , 2      , 3      , 1		)[0]
        _snake_case						=							Image.fromarray(np.uinta(A__		)		).convert('''RGB'''		).resize((64, 64)		)
        _snake_case						=							{
            '''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 						UpperCamelCase_		(      self				:    Union[str, Any]		)   ->					List[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 						UpperCamelCase_		(      self				:    Optional[int]		)   ->					List[Any]:
        self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3		)







      def 						UpperCamelCase_		(      self				:    Any		)   ->					Any:
        self._test_inference_batch_single_identical(expected_max_diff=2e-3		)



class   lowercase_   (  __lowercase       ,							__lowercase       ,							unittest.TestCase					):
      UpperCamelCase_      :		int										=    StableDiffusionControlNetImgaImgPipeline
      UpperCamelCase_      :		Union[str, Any]										=    TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"}
      UpperCamelCase_      :		Union[str, Any]										=    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
      UpperCamelCase_      :		Dict										=    frozenset([]					)  # TO_DO: add image_params once refactored VaeImageProcessor.preprocess



      def 						UpperCamelCase_		(      self				:    str		)   ->					Optional[int]:
        torch.manual_seed(0		)
        _snake_case						=							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(A__				:    Optional[int]		):
          if isinstance(A__      , torch.nn.Convad		):
            torch.nn.init.normal(m.weight		)
            m.bias.data.fill_(1.0		)

        _snake_case						=							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(A__		)

        torch.manual_seed(0		)
        _snake_case						=							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(A__		)

        torch.manual_seed(0		)
        _snake_case						=							DDIMScheduler(
            beta_start=0.0_0085      , beta_end=0.012      , beta_schedule='''scaled_linear'''      , clip_sample=A__      , set_alpha_to_one=A__      , )
        torch.manual_seed(0		)
        _snake_case						=							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		)
        _snake_case						=							CLIPTextConfig(
            bos_token_id=0      , eos_token_id=2      , hidden_size=32      , intermediate_size=37      , layer_norm_eps=1e-05      , num_attention_heads=4      , num_hidden_layers=5      , pad_token_id=1      , vocab_size=1000      , )
        _snake_case						=							CLIPTextModel(A__		)
        _snake_case						=							CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''		)

        _snake_case						=							MultiControlNetModel([controlneta, controlneta]		)

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



      def 						UpperCamelCase_		(      self				:    str      , A__				:    Union[str, Any]      , A__				:    Union[str, Any]=0		)   ->					str:
        if str(A__		).startswith('''mps'''		):
          _snake_case						=							torch.manual_seed(A__		)
        else:
          _snake_case						=							torch.Generator(device=A__		).manual_seed(A__		)

        _snake_case						=							2

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

        _snake_case						=							floats_tensor(control_image[0].shape      , rng=random.Random(A__		)		).to(A__		)
        _snake_case						=							image.cpu().permute(0      , 2      , 3      , 1		)[0]
        _snake_case						=							Image.fromarray(np.uinta(A__		)		).convert('''RGB'''		).resize((64, 64)		)
        _snake_case						=							{
            '''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 						UpperCamelCase_		(      self				:    List[str]		)   ->					List[Any]:
        _snake_case						=							self.get_dummy_components()
        _snake_case						=							self.pipeline_class(**A__		)
        pipe.to(A__		)

        _snake_case						=							10.0
        _snake_case						=							4

        _snake_case						=							self.get_dummy_inputs(A__		)
        _snake_case						=							steps
        _snake_case						=							scale
        _snake_case						=							pipe(**A__		)[0]

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

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

        _snake_case						=							self.get_dummy_inputs(A__		)
        _snake_case						=							steps
        _snake_case						=							scale
        _snake_case						=							pipe(**A__      , 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 						UpperCamelCase_		(      self				:    Optional[int]		)   ->					Optional[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 						UpperCamelCase_		(      self				:    Tuple		)   ->					Dict:
        self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3		)



      def 						UpperCamelCase_		(      self				:    int		)   ->					Dict:
        self._test_inference_batch_single_identical(expected_max_diff=2e-3		)







      def 						UpperCamelCase_		(      self				:    List[Any]		)   ->					Tuple:
        _snake_case						=							self.get_dummy_components()
        _snake_case						=							self.pipeline_class(**A__		)
        pipe.to(A__		)
        pipe.set_progress_bar_config(disable=A__		)
        with tempfile.TemporaryDirectory() as tmpdir:
          try:
            # save_pretrained is not implemented for Multi-ControlNet
            pipe.save_pretrained(A__		)
          except NotImplementedError:
            pass



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



      def 						UpperCamelCase_		(      self				:    Dict		)   ->					List[Any]:
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()







      def 						UpperCamelCase_		(      self				:    int		)   ->					Dict:
        _snake_case						=							ControlNetModel.from_pretrained('''lllyasviel/sd-controlnet-canny'''		)

        _snake_case						=							StableDiffusionControlNetImgaImgPipeline.from_pretrained(
            '''runwayml/stable-diffusion-v1-5'''      , safety_checker=A__      , controlnet=A__		)
        pipe.enable_model_cpu_offload()
        pipe.set_progress_bar_config(disable=A__		)

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

        _snake_case						=							pipe(
            A__      , A__      , control_image=A__      , generator=A__      , output_type='''np'''      , num_inference_steps=50      , strength=0.6      , )

        _snake_case						=							output.images[0]

        assert image.shape == (512, 512, 3)

        _snake_case						=							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




__A      		=     '''ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/'''



def 				snake_case_(_UpperCamelCase					)							->     bytes:



  """simple docstring"""






  if not isinstance(_UpperCamelCase			,	_UpperCamelCase					):
    _snake_case						=							F"""a bytes-like object is required, not '{data.__class__.__name__}'"""
    raise TypeError(_UpperCamelCase					)

  _snake_case						=							''''''.join(bin(_UpperCamelCase					)[2:].zfill(8					) for byte in data					)

  _snake_case						=							len(_UpperCamelCase					) % 6 != 0

  if padding_needed:
    # The padding that will be added later
    _snake_case						=							b'''=''' * ((6 - len(_UpperCamelCase					) % 6) // 2)

    # Append binary_stream with arbitrary binary digits (0's by default) to make its
    # length a multiple of 6.
    binary_stream += "0" * (6 - len(_UpperCamelCase					) % 6)
  else:
    _snake_case						=							b''''''

  # Encode every 6 binary digits to their corresponding Base64 character
  return (
      "".join(
          B64_CHARSET[int(binary_stream[index : index + 6]			,	2					)]
          for index in range(0			,	len(_UpperCamelCase					)			,	6					)					).encode()
      + padding
  )



def 				snake_case_(_UpperCamelCase					)							->     bytes:



  """simple docstring"""






  if not isinstance(_UpperCamelCase			,	_UpperCamelCase					) and not isinstance(_UpperCamelCase			,	_UpperCamelCase					):
    _snake_case						=							(
        '''argument should be a bytes-like object or ASCII string, '''
        F"""not '{encoded_data.__class__.__name__}'"""
    )
    raise TypeError(_UpperCamelCase					)

  # In case encoded_data is a bytes-like object, make sure it contains only
  # ASCII characters so we convert it to a string object
  if isinstance(_UpperCamelCase			,	_UpperCamelCase					):
    try:
      _snake_case						=							encoded_data.decode('''utf-8'''					)
    except UnicodeDecodeError:
      raise ValueError('''base64 encoded data should only contain ASCII characters'''					)

  _snake_case						=							encoded_data.count('''='''					)

  # Check if the encoded string contains non base64 characters
  if padding:
    assert all(
        char in B64_CHARSET for char in encoded_data[:-padding]					), "Invalid base64 character(s) found."
  else:
    assert all(
        char in B64_CHARSET for char in encoded_data					), "Invalid base64 character(s) found."

  # Check the padding
  assert len(_UpperCamelCase					) % 4 == 0 and padding < 3, "Incorrect padding"

  if padding:
    # Remove padding if there is one
    _snake_case						=							encoded_data[:-padding]

    _snake_case						=							''''''.join(
        bin(B64_CHARSET.index(_UpperCamelCase					)					)[2:].zfill(6					) for char in encoded_data					)[: -padding * 2]
  else:
    _snake_case						=							''''''.join(
        bin(B64_CHARSET.index(_UpperCamelCase					)					)[2:].zfill(6					) for char in encoded_data					)

  _snake_case						=							[
      int(binary_stream[index : index + 8]			,	2					)
      for index in range(0			,	len(_UpperCamelCase					)			,	8					)
  ]

  return bytes(_UpperCamelCase					)


if __name__ == "__main__":
 import doctest

 doctest.testmod()

