


'''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   __snake_case			(  _SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,unittest.TestCase):







    """simple docstring"""



    lowercase  					=   StableDiffusionControlNetImgaImgPipeline
    lowercase  					=   TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'height', 'width'}
    lowercase  					=   TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
    lowercase  					=   IMAGE_TO_IMAGE_IMAGE_PARAMS.union({'control_image'})
    lowercase  					=   IMAGE_TO_IMAGE_IMAGE_PARAMS

    def __lowercase  (  self       :    str						)  ->    Tuple:
      torch.manual_seed(0						)
      lowerCAmelCase_     : Optional[Any]     =  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_     : str     =  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_     : Optional[int]     =  DDIMScheduler(
          beta_start=0.00_085       ,		beta_end=0.012       ,		beta_schedule="""scaled_linear"""       ,		clip_sample=lowerCamelCase       ,		set_alpha_to_one=lowerCamelCase       ,		)
      torch.manual_seed(0						)
      lowerCAmelCase_     : List[Any]     =  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_     : Optional[int]     =  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=10_00       ,		)
      lowerCAmelCase_     : Tuple     =  CLIPTextModel(lowerCamelCase						)
      lowerCAmelCase_     : Any     =  CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""						)

      lowerCAmelCase_     : Union[str, Any]     =  {
          """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       :    List[Any]       ,		lowerCamelCase       :    Any       ,		lowerCamelCase       :    Dict=0						)  ->    Dict:
      if str(lowerCamelCase						).startswith("""mps"""						):
        lowerCAmelCase_     : int     =  torch.manual_seed(lowerCamelCase						)
      else:
        lowerCAmelCase_     : Optional[int]     =  torch.Generator(device=lowerCamelCase						).manual_seed(lowerCamelCase						)

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





    def __lowercase  (  self       :    str						)  ->    Optional[Any]:
      self._test_inference_batch_single_identical(expected_max_diff=2E-3						)







class   __snake_case			(  _SCREAMING_SNAKE_CASE      ,_SCREAMING_SNAKE_CASE      ,unittest.TestCase):







    """simple docstring"""



    lowercase  					=   StableDiffusionControlNetImgaImgPipeline
    lowercase  					=   TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'height', 'width'}
    lowercase  					=   TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
    lowercase  					=   frozenset([])  # TO_DO: add image_params once refactored VaeImageProcessor.preprocess

    def __lowercase  (  self       :    List[Any]						)  ->    Dict:
      torch.manual_seed(0						)
      lowerCAmelCase_     : Dict     =  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(lowerCamelCase       :    Optional[int]						):
        if isinstance(lowerCamelCase       ,		torch.nn.Convad						):
          torch.nn.init.normal(m.weight						)
          m.bias.data.fill_(1.0						)

      lowerCAmelCase_     : List[Any]     =  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(lowerCamelCase						)

      torch.manual_seed(0						)
      lowerCAmelCase_     : Tuple     =  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(lowerCamelCase						)

      torch.manual_seed(0						)
      lowerCAmelCase_     : Optional[Any]     =  DDIMScheduler(
          beta_start=0.00_085       ,		beta_end=0.012       ,		beta_schedule="""scaled_linear"""       ,		clip_sample=lowerCamelCase       ,		set_alpha_to_one=lowerCamelCase       ,		)
      torch.manual_seed(0						)
      lowerCAmelCase_     : str     =  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_     : List[Any]     =  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=10_00       ,		)
      lowerCAmelCase_     : Optional[Any]     =  CLIPTextModel(lowerCamelCase						)
      lowerCAmelCase_     : Union[str, Any]     =  CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""						)

      lowerCAmelCase_     : Union[str, Any]     =  MultiControlNetModel([controlneta, controlneta]						)

      lowerCAmelCase_     : Dict     =  {
          """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       :    Optional[int]       ,		lowerCamelCase       :    Dict       ,		lowerCamelCase       :    Optional[int]=0						)  ->    Optional[Any]:
      if str(lowerCamelCase						).startswith("""mps"""						):
        lowerCAmelCase_     : Tuple     =  torch.manual_seed(lowerCamelCase						)
      else:
        lowerCAmelCase_     : Tuple     =  torch.Generator(device=lowerCamelCase						).manual_seed(lowerCamelCase						)

      lowerCAmelCase_     : str     =  2

      lowerCAmelCase_     : Any     =  [
          randn_tensor(
              (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)       ,		generator=lowerCamelCase       ,		device=torch.device(lowerCamelCase						)       ,		),
          randn_tensor(
              (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)       ,		generator=lowerCamelCase       ,		device=torch.device(lowerCamelCase						)       ,		),
      ]

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

      lowerCAmelCase_     : Any     =  10.0
      lowerCAmelCase_     : int     =  4

      lowerCAmelCase_     : List[str]     =  self.get_dummy_inputs(lowerCamelCase						)
      lowerCAmelCase_     : Optional[Any]     =  steps
      lowerCAmelCase_     : str     =  scale
      lowerCAmelCase_     : Optional[int]     =  pipe(**lowerCamelCase						)[0]

      lowerCAmelCase_     : List[str]     =  self.get_dummy_inputs(lowerCamelCase						)
      lowerCAmelCase_     : Optional[Any]     =  steps
      lowerCAmelCase_     : List[str]     =  scale
      lowerCAmelCase_     : Dict     =  pipe(**lowerCamelCase       ,		control_guidance_start=0.1       ,		control_guidance_end=0.2						)[0]

      lowerCAmelCase_     : int     =  self.get_dummy_inputs(lowerCamelCase						)
      lowerCAmelCase_     : Union[str, Any]     =  steps
      lowerCAmelCase_     : Tuple     =  scale
      lowerCAmelCase_     : Optional[int]     =  pipe(**lowerCamelCase       ,		control_guidance_start=[0.1, 0.3]       ,		control_guidance_end=[0.2, 0.7]						)[0]

      lowerCAmelCase_     : Tuple     =  self.get_dummy_inputs(lowerCamelCase						)
      lowerCAmelCase_     : int     =  steps
      lowerCAmelCase_     : int     =  scale
      lowerCAmelCase_     : Any     =  pipe(**lowerCamelCase       ,		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       :    Tuple						)  ->    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						)  ->    List[Any]:
      self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3						)

    def __lowercase  (  self       :    int						)  ->    Tuple:
      self._test_inference_batch_single_identical(expected_max_diff=2E-3						)





    def __lowercase  (  self       :    Dict						)  ->    Tuple:
      lowerCAmelCase_     : Optional[int]     =  self.get_dummy_components()
      lowerCAmelCase_     : Any     =  self.pipeline_class(**lowerCamelCase						)
      pipe.to(lowerCamelCase						)
      pipe.set_progress_bar_config(disable=lowerCamelCase						)
      with tempfile.TemporaryDirectory() as tmpdir:
        try:
          # save_pretrained is not implemented for Multi-ControlNet
          pipe.save_pretrained(lowerCamelCase						)
        except NotImplementedError:
          pass







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







    """simple docstring"""



    def __lowercase  (  self       :    Any						)  ->    int:
      super().tearDown()
      gc.collect()
      torch.cuda.empty_cache()





    def __lowercase  (  self       :    Optional[int]						)  ->    Any:
      lowerCAmelCase_     : List[str]     =  ControlNetModel.from_pretrained("""lllyasviel/sd-controlnet-canny"""						)

      lowerCAmelCase_     : str     =  StableDiffusionControlNetImgaImgPipeline.from_pretrained(
          """runwayml/stable-diffusion-v1-5"""       ,		safety_checker=lowerCamelCase       ,		controlnet=lowerCamelCase						)
      pipe.enable_model_cpu_offload()
      pipe.set_progress_bar_config(disable=lowerCamelCase						)

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

      lowerCAmelCase_     : Dict     =  pipe(
          lowerCamelCase       ,		lowerCamelCase       ,		control_image=lowerCamelCase       ,		generator=lowerCamelCase       ,		output_type="""np"""       ,		num_inference_steps=50       ,		strength=0.6       ,		)

      lowerCAmelCase_     : Union[str, Any]     =  output.images[0]

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

      lowerCAmelCase_     : int     =  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'''






def 							UpperCamelCase_     (       A__   :  bytes   ):
  '''simple docstring'''





  return "".join([hex(A__   )[2:].zfill(2   ).upper() for byte in list(A__   )]   )







def 							UpperCamelCase_     (       A__   :  str   ):
  '''simple docstring'''





  if (len(A__   ) % 2) != 0:
    raise ValueError(
        """Base16 encoded data is invalid:
Data does not have an even number of hex digits."""   )
  # Check the character set - the standard base16 alphabet
  # is uppercase according to RFC3548 section 6
  if not set(A__   ) <= set("""0123456789ABCDEF"""   ):
    raise ValueError(
        """Base16 encoded data is invalid:
Data is not uppercase hex or it contains invalid characters."""   )
  # For every two hexadecimal digits (= a byte), turn it into an integer.
  # Then, string the result together into bytes, and return it.
  return bytes(int(data[i] + data[i + 1]				,       16   ) for i in range(0				,       len(A__   )				,       2   )   )


if __name__ == "__main__":
  import doctest

  doctest.testmod()



