


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




 '''simple docstring'''
 __lowercase					:		Optional[int]     				=    StableDiffusionControlNetImgaImgPipeline
 __lowercase					:		Dict     				=    TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'''height''', '''width'''}
 __lowercase					:		Any     				=    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
 __lowercase					:		Dict     				=    IMAGE_TO_IMAGE_IMAGE_PARAMS.union({'''control_image'''}					)
 __lowercase					:		Union[str, Any]     				=    IMAGE_TO_IMAGE_IMAGE_PARAMS


 def 			snake_case__ (   self    ):

  """simple docstring"""







  torch.manual_seed(0    )
  __A	:		List[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    )
  __A	:		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)   ,   )
  torch.manual_seed(0    )
  __A	:		List[str]	  =      DDIMScheduler(
      beta_start=0.0_0_0_8_5   ,   beta_end=0.0_1_2   ,   beta_schedule='scaled_linear'   ,   clip_sample=__lowercase   ,   set_alpha_to_one=__lowercase   ,   )
  torch.manual_seed(0    )
  __A	:		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    )
  __A	:		Dict	  =      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=1_000   ,   )
  __A	:		Optional[int]	  =      CLIPTextModel(__lowercase    )
  __A	:		Optional[int]	  =      CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'    )

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


 def 			snake_case__ (   self   ,   __lowercase   ,   __lowercase=0    ):

  """simple docstring"""







  if str(__lowercase    ).startswith('mps'    ):
   __A	:		Optional[Any]	  =      torch.manual_seed(__lowercase    )
  else:
   __A	:		Union[str, Any]	  =      torch.Generator(device=__lowercase    ).manual_seed(__lowercase    )

  __A	:		Union[str, Any]	  =      2
  __A	:		Dict	  =      randn_tensor(
      (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)   ,   generator=__lowercase   ,   device=torch.device(__lowercase    )   ,   )
  __A	:		Any	  =      floats_tensor(control_image.shape   ,   rng=random.Random(__lowercase    )    ).to(__lowercase    )
  __A	:		Dict	  =      image.cpu().permute(0   ,   2   ,   3   ,   1    )[0]
  __A	:		Union[str, Any]	  =      Image.fromarray(np.uinta(__lowercase    )    ).convert('RGB'    ).resize((64, 64)    )
  __A	:		Union[str, 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 			snake_case__ (   self    ):

  """simple docstring"""







  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 			snake_case__ (   self    ):

  """simple docstring"""







  self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3    )







 def 			snake_case__ (   self    ):

  """simple docstring"""







  self._test_inference_batch_single_identical(expected_max_diff=2E-3    )


class     _lowerCamelCase       (     lowerCamelCase__     ,		lowerCamelCase__     ,		unittest.TestCase					):




 '''simple docstring'''
 __lowercase					:		str     				=    StableDiffusionControlNetImgaImgPipeline
 __lowercase					:		List[Any]     				=    TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'''height''', '''width'''}
 __lowercase					:		Optional[int]     				=    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
 __lowercase					:		Dict     				=    frozenset([]					)  # TO_DO: add image_params once refactored VaeImageProcessor.preprocess


 def 			snake_case__ (   self    ):

  """simple docstring"""







  torch.manual_seed(0    )
  __A	:		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    )

  def init_weights(__lowercase    ):
   if isinstance(__lowercase   ,   torch.nn.Convad    ):
    torch.nn.init.normal(m.weight    )
    m.bias.data.fill_(1.0    )

  __A	:		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)   ,   )
  controlneta.controlnet_down_blocks.apply(__lowercase    )

  torch.manual_seed(0    )
  __A	:		Union[str, 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(__lowercase    )

  torch.manual_seed(0    )
  __A	:		Optional[int]	  =      DDIMScheduler(
      beta_start=0.0_0_0_8_5   ,   beta_end=0.0_1_2   ,   beta_schedule='scaled_linear'   ,   clip_sample=__lowercase   ,   set_alpha_to_one=__lowercase   ,   )
  torch.manual_seed(0    )
  __A	:		Optional[int]	  =      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    )
  __A	:		Union[str, 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=1_000   ,   )
  __A	:		Dict	  =      CLIPTextModel(__lowercase    )
  __A	:		Optional[int]	  =      CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'    )

  __A	:		Optional[Any]	  =      MultiControlNetModel([controlneta, controlneta]    )

  __A	:		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 			snake_case__ (   self   ,   __lowercase   ,   __lowercase=0    ):

  """simple docstring"""







  if str(__lowercase    ).startswith('mps'    ):
   __A	:		List[Any]	  =      torch.manual_seed(__lowercase    )
  else:
   __A	:		Any	  =      torch.Generator(device=__lowercase    ).manual_seed(__lowercase    )

  __A	:		Dict	  =      2

  __A	:		Tuple	  =      [
      randn_tensor(
          (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)   ,   generator=__lowercase   ,   device=torch.device(__lowercase    )   ,   ),
      randn_tensor(
          (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)   ,   generator=__lowercase   ,   device=torch.device(__lowercase    )   ,   ),
  ]

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

  """simple docstring"""







  __A	:		str	  =      self.get_dummy_components()
  __A	:		List[str]	  =      self.pipeline_class(**__lowercase    )
  pipe.to(__lowercase    )

  __A	:		str	  =      1_0.0
  __A	:		List[Any]	  =      4

  __A	:		Tuple	  =      self.get_dummy_inputs(__lowercase    )
  __A	:		int	  =      steps
  __A	:		Any	  =      scale
  __A	:		List[Any]	  =      pipe(**__lowercase    )[0]

  __A	:		Optional[Any]	  =      self.get_dummy_inputs(__lowercase    )
  __A	:		Union[str, Any]	  =      steps
  __A	:		Union[str, Any]	  =      scale
  __A	:		Tuple	  =      pipe(**__lowercase   ,   control_guidance_start=0.1   ,   control_guidance_end=0.2    )[0]

  __A	:		Any	  =      self.get_dummy_inputs(__lowercase    )
  __A	:		List[Any]	  =      steps
  __A	:		Tuple	  =      scale
  __A	:		Any	  =      pipe(**__lowercase   ,   control_guidance_start=[0.1, 0.3]   ,   control_guidance_end=[0.2, 0.7]    )[0]

  __A	:		Optional[Any]	  =      self.get_dummy_inputs(__lowercase    )
  __A	:		str	  =      steps
  __A	:		Dict	  =      scale
  __A	:		str	  =      pipe(**__lowercase   ,   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 			snake_case__ (   self    ):

  """simple docstring"""







  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 			snake_case__ (   self    ):

  """simple docstring"""







  self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3    )


 def 			snake_case__ (   self    ):

  """simple docstring"""







  self._test_inference_batch_single_identical(expected_max_diff=2E-3    )







 def 			snake_case__ (   self    ):

  """simple docstring"""







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


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




 '''simple docstring'''


 def 			snake_case__ (   self    ):

  """simple docstring"""







  super().tearDown()
  gc.collect()
  torch.cuda.empty_cache()







 def 			snake_case__ (   self    ):

  """simple docstring"""







  __A	:		Tuple	  =      ControlNetModel.from_pretrained('lllyasviel/sd-controlnet-canny'    )

  __A	:		int	  =      StableDiffusionControlNetImgaImgPipeline.from_pretrained(
      'runwayml/stable-diffusion-v1-5'   ,   safety_checker=__lowercase   ,   controlnet=__lowercase    )
  pipe.enable_model_cpu_offload()
  pipe.set_progress_bar_config(disable=__lowercase    )

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

  __A	:		Any	  =      pipe(
      __lowercase   ,   __lowercase   ,   control_image=__lowercase   ,   generator=__lowercase   ,   output_type='np'   ,   num_inference_steps=50   ,   strength=0.6   ,   )

  __A	:		Dict	  =      output.images[0]

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

  __A	:		Union[str, Any]	  =      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













def   __snake_case    (       _lowerCAmelCase			:			int						)			->    bool:
   if num < 0:
      return False

   A_       :					int        =      num
   A_       :					int        =      0
   while num > 0:
      A_       :					str        =      rev_num * 10 + (num % 10)
      num //= 10

   return num_copy == rev_num


if __name__ == "__main__":
 import doctest

 doctest.testmod()
