import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import PriorTransformer, UnCLIPPipeline, UnCLIPScheduler, UNetaDConditionModel, UNetaDModel
from diffusers.pipelines.unclip.text_proj import UnCLIPTextProjModel
from diffusers.utils import load_numpy, nightly, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()




class 						__snake_case						(			__lowerCAmelCase			,	unittest.TestCase		):





      '''simple docstring'''






      _snake_case     		=							UnCLIPPipeline
      _snake_case     		=							TEXT_TO_IMAGE_PARAMS - {
          'negative_prompt',
          'height',
          'width',
          'negative_prompt_embeds',
          'guidance_scale',
          'prompt_embeds',
          'cross_attention_kwargs',
      }
      _snake_case     		=							TEXT_TO_IMAGE_BATCH_PARAMS
      _snake_case     		=							[
          'generator',
          'return_dict',
          'prior_num_inference_steps',
          'decoder_num_inference_steps',
          'super_res_num_inference_steps',
      ]
      _snake_case     		=							False



      @property
      def 						_SCREAMING_SNAKE_CASE     (       self  :					Optional[Any])					->Dict:


          """simple docstring"""


          return 32



      @property
      def 						_SCREAMING_SNAKE_CASE     (       self  :					List[str])					->Tuple:


          """simple docstring"""


          return 32



      @property
      def 						_SCREAMING_SNAKE_CASE     (       self  :					Optional[int])					->Optional[Any]:


          """simple docstring"""


          return self.time_input_dim



      @property
      def 						_SCREAMING_SNAKE_CASE     (       self  :					str)					->Any:


          """simple docstring"""


          return self.time_input_dim * 4



      @property
      def 						_SCREAMING_SNAKE_CASE     (       self  :					List[Any])					->Any:


          """simple docstring"""


          return 100



      @property
      def 						_SCREAMING_SNAKE_CASE     (       self  :					Union[str, Any])					->Any:


          """simple docstring"""


          _lowerCamelCase  :		List[str]			    =    CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip""")
          return tokenizer



      @property
      def 						_SCREAMING_SNAKE_CASE     (       self  :					Tuple)					->Optional[Any]:


          """simple docstring"""


          torch.manual_seed(0)
          _lowerCamelCase  :		List[Any]			    =    CLIPTextConfig(
              bos_token_id=0     ,	eos_token_id=2     ,	hidden_size=self.text_embedder_hidden_size     ,	projection_dim=self.text_embedder_hidden_size     ,	intermediate_size=37     ,	layer_norm_eps=1E-0_5     ,	num_attention_heads=4     ,	num_hidden_layers=5     ,	pad_token_id=1     ,	vocab_size=1000     ,	)
          return CLIPTextModelWithProjection(_UpperCamelCase)



      @property
      def 						_SCREAMING_SNAKE_CASE     (       self  :					str)					->List[Any]:


          """simple docstring"""


          torch.manual_seed(0)

          _lowerCamelCase  :		int			    =    {
              """num_attention_heads""": 2,
              """attention_head_dim""": 12,
              """embedding_dim""": self.text_embedder_hidden_size,
              """num_layers""": 1,
          }

          _lowerCamelCase  :		Dict			    =    PriorTransformer(**_UpperCamelCase)
          return model



      @property
      def 						_SCREAMING_SNAKE_CASE     (       self  :					List[str])					->int:


          """simple docstring"""


          torch.manual_seed(0)

          _lowerCamelCase  :		Any			    =    {
              """clip_embeddings_dim""": self.text_embedder_hidden_size,
              """time_embed_dim""": self.time_embed_dim,
              """cross_attention_dim""": self.cross_attention_dim,
          }

          _lowerCamelCase  :		int			    =    UnCLIPTextProjModel(**_UpperCamelCase)
          return model



      @property
      def 						_SCREAMING_SNAKE_CASE     (       self  :					List[Any])					->Dict:


          """simple docstring"""


          torch.manual_seed(0)

          _lowerCamelCase  :		Any			    =    {
              """sample_size""": 32,
              # RGB in channels
              """in_channels""": 3,
              # Out channels is double in channels because predicts mean and variance
              """out_channels""": 6,
              """down_block_types""": ("""ResnetDownsampleBlock2D""", """SimpleCrossAttnDownBlock2D"""),
              """up_block_types""": ("""SimpleCrossAttnUpBlock2D""", """ResnetUpsampleBlock2D"""),
              """mid_block_type""": """UNetMidBlock2DSimpleCrossAttn""",
              """block_out_channels""": (self.block_out_channels_a, self.block_out_channels_a * 2),
              """layers_per_block""": 1,
              """cross_attention_dim""": self.cross_attention_dim,
              """attention_head_dim""": 4,
              """resnet_time_scale_shift""": """scale_shift""",
              """class_embed_type""": """identity""",
          }

          _lowerCamelCase  :		Any			    =    UNetaDConditionModel(**_UpperCamelCase)
          return model



      @property
      def 						_SCREAMING_SNAKE_CASE     (       self  :					Optional[Any])					->Union[str, Any]:


          """simple docstring"""


          return {
              "sample_size": 64,
              "layers_per_block": 1,
              "down_block_types": ("ResnetDownsampleBlock2D", "ResnetDownsampleBlock2D"),
              "up_block_types": ("ResnetUpsampleBlock2D", "ResnetUpsampleBlock2D"),
              "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
              "in_channels": 6,
              "out_channels": 3,
          }



      @property
      def 						_SCREAMING_SNAKE_CASE     (       self  :					str)					->Dict:


          """simple docstring"""


          torch.manual_seed(0)

          _lowerCamelCase  :		List[Any]			    =    UNetaDModel(**self.dummy_super_res_kwargs)
          return model



      @property
      def 						_SCREAMING_SNAKE_CASE     (       self  :					Any)					->Optional[int]:


          """simple docstring"""


          torch.manual_seed(1)

          _lowerCamelCase  :		str			    =    UNetaDModel(**self.dummy_super_res_kwargs)
          return model



      def 						_SCREAMING_SNAKE_CASE     (       self  :					Optional[int])					->int:


          """simple docstring"""


          _lowerCamelCase  :		Optional[Any]			    =    self.dummy_prior
          _lowerCamelCase  :		str			    =    self.dummy_decoder
          _lowerCamelCase  :		Optional[int]			    =    self.dummy_text_proj
          _lowerCamelCase  :		Optional[int]			    =    self.dummy_text_encoder
          _lowerCamelCase  :		Any			    =    self.dummy_tokenizer
          _lowerCamelCase  :		Union[str, Any]			    =    self.dummy_super_res_first
          _lowerCamelCase  :		Optional[int]			    =    self.dummy_super_res_last

          _lowerCamelCase  :		Union[str, Any]			    =    UnCLIPScheduler(
              variance_type="""fixed_small_log"""     ,	prediction_type="""sample"""     ,	num_train_timesteps=1000     ,	clip_sample_range=5.0     ,	)

          _lowerCamelCase  :		Tuple			    =    UnCLIPScheduler(
              variance_type="""learned_range"""     ,	prediction_type="""epsilon"""     ,	num_train_timesteps=1000     ,	)

          _lowerCamelCase  :		int			    =    UnCLIPScheduler(
              variance_type="""fixed_small_log"""     ,	prediction_type="""epsilon"""     ,	num_train_timesteps=1000     ,	)

          _lowerCamelCase  :		Dict			    =    {
              """prior""": prior,
              """decoder""": decoder,
              """text_proj""": text_proj,
              """text_encoder""": text_encoder,
              """tokenizer""": tokenizer,
              """super_res_first""": super_res_first,
              """super_res_last""": super_res_last,
              """prior_scheduler""": prior_scheduler,
              """decoder_scheduler""": decoder_scheduler,
              """super_res_scheduler""": super_res_scheduler,
          }

          return components



      def 						_SCREAMING_SNAKE_CASE     (       self  :					Union[str, Any]     ,	_UpperCamelCase  :					Any     ,	_UpperCamelCase  :					Optional[Any]=0)					->Dict:


          """simple docstring"""


          if str(_UpperCamelCase).startswith("""mps"""):
              _lowerCamelCase  :		Dict			    =    torch.manual_seed(_UpperCamelCase)
          else:
              _lowerCamelCase  :		List[str]			    =    torch.Generator(device=_UpperCamelCase).manual_seed(_UpperCamelCase)
          _lowerCamelCase  :		Dict			    =    {
              """prompt""": """horse""",
              """generator""": generator,
              """prior_num_inference_steps""": 2,
              """decoder_num_inference_steps""": 2,
              """super_res_num_inference_steps""": 2,
              """output_type""": """numpy""",
          }
          return inputs



      def 						_SCREAMING_SNAKE_CASE     (       self  :					int)					->Union[str, Any]:


          """simple docstring"""


          _lowerCamelCase  :		Optional[Any]			    =    """cpu"""

          _lowerCamelCase  :		List[str]			    =    self.get_dummy_components()

          _lowerCamelCase  :		Union[str, Any]			    =    self.pipeline_class(**_UpperCamelCase)
          _lowerCamelCase  :		List[Any]			    =    pipe.to(_UpperCamelCase)

          pipe.set_progress_bar_config(disable=_UpperCamelCase)

          _lowerCamelCase  :		Any			    =    pipe(**self.get_dummy_inputs(_UpperCamelCase))
          _lowerCamelCase  :		Union[str, Any]			    =    output.images

          _lowerCamelCase  :		List[Any]			    =    pipe(
              **self.get_dummy_inputs(_UpperCamelCase)     ,	return_dict=_UpperCamelCase     ,	)[0]

          _lowerCamelCase  :		Dict			    =    image[0, -3:, -3:, -1]
          _lowerCamelCase  :		List[Any]			    =    image_from_tuple[0, -3:, -3:, -1]

          assert image.shape == (1, 64, 64, 3)

          _lowerCamelCase  :		List[str]			    =    np.array(
              [
                  0.9_9_9_7,
                  0.9_9_8_8,
                  0.0_0_2_8,
                  0.9_9_9_7,
                  0.9_9_8_4,
                  0.9_9_6_5,
                  0.0_0_2_9,
                  0.9_9_8_6,
                  0.0_0_2_5,
              ])

          assert np.abs(image_slice.flatten() - expected_slice).max() < 1E-2
          assert np.abs(image_from_tuple_slice.flatten() - expected_slice).max() < 1E-2



      def 						_SCREAMING_SNAKE_CASE     (       self  :					List[str])					->Union[str, Any]:


          """simple docstring"""


          _lowerCamelCase  :		Dict			    =    torch.device("""cpu""")




          class 						__snake_case						:





                '''simple docstring'''






                _snake_case     		=							1





          _lowerCamelCase  :		Optional[int]			    =    self.get_dummy_components()

          _lowerCamelCase  :		List[str]			    =    self.pipeline_class(**_UpperCamelCase)
          _lowerCamelCase  :		Tuple			    =    pipe.to(_UpperCamelCase)

          _lowerCamelCase  :		Union[str, Any]			    =    components["""prior"""]
          _lowerCamelCase  :		int			    =    components["""decoder"""]
          _lowerCamelCase  :		Any			    =    components["""super_res_first"""]
          _lowerCamelCase  :		str			    =    components["""tokenizer"""]
          _lowerCamelCase  :		Dict			    =    components["""text_encoder"""]

          _lowerCamelCase  :		Dict			    =    torch.Generator(device=_UpperCamelCase).manual_seed(0)
          _lowerCamelCase  :		Tuple			    =    prior.dtype
          _lowerCamelCase  :		Union[str, Any]			    =    1

          _lowerCamelCase  :		Optional[Any]			    =    (batch_size, prior.config.embedding_dim)
          _lowerCamelCase  :		int			    =    pipe.prepare_latents(
              _UpperCamelCase     ,	dtype=_UpperCamelCase     ,	device=_UpperCamelCase     ,	generator=_UpperCamelCase     ,	latents=_UpperCamelCase     ,	scheduler=DummyScheduler())
          _lowerCamelCase  :		Optional[int]			    =    (batch_size, decoder.config.in_channels, decoder.config.sample_size, decoder.config.sample_size)
          _lowerCamelCase  :		Union[str, Any]			    =    pipe.prepare_latents(
              _UpperCamelCase     ,	dtype=_UpperCamelCase     ,	device=_UpperCamelCase     ,	generator=_UpperCamelCase     ,	latents=_UpperCamelCase     ,	scheduler=DummyScheduler())

          _lowerCamelCase  :		Optional[int]			    =    (
              batch_size,
              super_res_first.config.in_channels // 2,
              super_res_first.config.sample_size,
              super_res_first.config.sample_size,
          )
          _lowerCamelCase  :		Optional[Any]			    =    pipe.prepare_latents(
              _UpperCamelCase     ,	dtype=_UpperCamelCase     ,	device=_UpperCamelCase     ,	generator=_UpperCamelCase     ,	latents=_UpperCamelCase     ,	scheduler=DummyScheduler())

          pipe.set_progress_bar_config(disable=_UpperCamelCase)

          _lowerCamelCase  :		Dict			    =    """this is a prompt example"""

          _lowerCamelCase  :		Dict			    =    torch.Generator(device=_UpperCamelCase).manual_seed(0)
          _lowerCamelCase  :		Tuple			    =    pipe(
              [prompt]     ,	generator=_UpperCamelCase     ,	prior_num_inference_steps=2     ,	decoder_num_inference_steps=2     ,	super_res_num_inference_steps=2     ,	prior_latents=_UpperCamelCase     ,	decoder_latents=_UpperCamelCase     ,	super_res_latents=_UpperCamelCase     ,	output_type="""np"""     ,	)
          _lowerCamelCase  :		Tuple			    =    output.images

          _lowerCamelCase  :		Union[str, Any]			    =    tokenizer(
              _UpperCamelCase     ,	padding="""max_length"""     ,	max_length=tokenizer.model_max_length     ,	return_tensors="""pt"""     ,	)
          _lowerCamelCase  :		Dict			    =    text_encoder(text_inputs.input_ids)
          _lowerCamelCase  :		List[Any]			    =    text_inputs.attention_mask

          _lowerCamelCase  :		Optional[int]			    =    torch.Generator(device=_UpperCamelCase).manual_seed(0)
          _lowerCamelCase  :		Optional[int]			    =    pipe(
              generator=_UpperCamelCase     ,	prior_num_inference_steps=2     ,	decoder_num_inference_steps=2     ,	super_res_num_inference_steps=2     ,	prior_latents=_UpperCamelCase     ,	decoder_latents=_UpperCamelCase     ,	super_res_latents=_UpperCamelCase     ,	text_model_output=_UpperCamelCase     ,	text_attention_mask=_UpperCamelCase     ,	output_type="""np"""     ,	)[0]

          # make sure passing text embeddings manually is identical
          assert np.abs(image - image_from_text).max() < 1E-4



      @skip_mps
      def 						_SCREAMING_SNAKE_CASE     (       self  :					Any)					->Any:


          """simple docstring"""


          _lowerCamelCase  :		str			    =    torch_device == """cpu"""

          self._test_attention_slicing_forward_pass(test_max_difference=_UpperCamelCase     ,	expected_max_diff=0.0_1)



      @skip_mps
      def 						_SCREAMING_SNAKE_CASE     (       self  :					Union[str, Any])					->Tuple:


          """simple docstring"""


          _lowerCamelCase  :		Union[str, Any]			    =    torch_device == """cpu"""
          _lowerCamelCase  :		Dict			    =    True
          _lowerCamelCase  :		List[Any]			    =    [
              """prior_num_inference_steps""",
              """decoder_num_inference_steps""",
              """super_res_num_inference_steps""",
          ]

          self._test_inference_batch_single_identical(
              test_max_difference=_UpperCamelCase     ,	relax_max_difference=_UpperCamelCase     ,	additional_params_copy_to_batched_inputs=_UpperCamelCase     ,	)



      def 						_SCREAMING_SNAKE_CASE     (       self  :					Optional[int])					->Any:


          """simple docstring"""


          _lowerCamelCase  :		List[str]			    =    [
              """prior_num_inference_steps""",
              """decoder_num_inference_steps""",
              """super_res_num_inference_steps""",
          ]

          if torch_device == "mps":
              # TODO: MPS errors with larger batch sizes
              _lowerCamelCase  :		Optional[Any]			    =    [2, 3]
              self._test_inference_batch_consistent(
                  batch_sizes=_UpperCamelCase     ,	additional_params_copy_to_batched_inputs=_UpperCamelCase     ,	)
          else:
              self._test_inference_batch_consistent(
                  additional_params_copy_to_batched_inputs=_UpperCamelCase)



      @skip_mps
      def 						_SCREAMING_SNAKE_CASE     (       self  :					Optional[int])					->Union[str, Any]:


          """simple docstring"""


          return super().test_dict_tuple_outputs_equivalent()



      @skip_mps
      def 						_SCREAMING_SNAKE_CASE     (       self  :					Optional[int])					->Any:


          """simple docstring"""


          return super().test_save_load_local()




      @skip_mps
      def 						_SCREAMING_SNAKE_CASE     (       self  :					Optional[Any])					->Optional[int]:


          """simple docstring"""


          return super().test_save_load_optional_components()





@nightly
class 						__snake_case						(			unittest.TestCase		):





      '''simple docstring'''






      def 						_SCREAMING_SNAKE_CASE     (       self  :					List[str])					->List[str]:


          """simple docstring"""


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




      def 						_SCREAMING_SNAKE_CASE     (       self  :					Dict)					->str:


          """simple docstring"""


          _lowerCamelCase  :		Dict			    =    load_numpy(
              """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
              """/unclip/karlo_v1_alpha_horse_cpu.npy""")

          _lowerCamelCase  :		Dict			    =    UnCLIPPipeline.from_pretrained("""kakaobrain/karlo-v1-alpha""")
          pipeline.set_progress_bar_config(disable=_UpperCamelCase)

          _lowerCamelCase  :		int			    =    torch.manual_seed(0)
          _lowerCamelCase  :		Optional[Any]			    =    pipeline(
              """horse"""     ,	num_images_per_prompt=1     ,	generator=_UpperCamelCase     ,	output_type="""np"""     ,	)

          _lowerCamelCase  :		List[str]			    =    output.images[0]

          assert image.shape == (256, 256, 3)
          assert np.abs(expected_image - image).max() < 1E-1





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





      '''simple docstring'''






      def 						_SCREAMING_SNAKE_CASE     (       self  :					Tuple)					->List[str]:


          """simple docstring"""


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



      def 						_SCREAMING_SNAKE_CASE     (       self  :					int)					->Dict:


          """simple docstring"""


          _lowerCamelCase  :		Dict			    =    load_numpy(
              """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
              """/unclip/karlo_v1_alpha_horse_fp16.npy""")

          _lowerCamelCase  :		Any			    =    UnCLIPPipeline.from_pretrained("""kakaobrain/karlo-v1-alpha"""     ,	torch_dtype=torch.floataa)
          _lowerCamelCase  :		Union[str, Any]			    =    pipeline.to(_UpperCamelCase)
          pipeline.set_progress_bar_config(disable=_UpperCamelCase)

          _lowerCamelCase  :		Optional[int]			    =    torch.Generator(device="""cpu""").manual_seed(0)
          _lowerCamelCase  :		Optional[Any]			    =    pipeline(
              """horse"""     ,	generator=_UpperCamelCase     ,	output_type="""np"""     ,	)

          _lowerCamelCase  :		Optional[int]			    =    output.images[0]

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

          assert_mean_pixel_difference(_UpperCamelCase     ,	_UpperCamelCase)




      def 						_SCREAMING_SNAKE_CASE     (       self  :					Any)					->int:


          """simple docstring"""


          torch.cuda.empty_cache()
          torch.cuda.reset_max_memory_allocated()
          torch.cuda.reset_peak_memory_stats()

          _lowerCamelCase  :		Optional[int]			    =    UnCLIPPipeline.from_pretrained("""kakaobrain/karlo-v1-alpha"""     ,	torch_dtype=torch.floataa)
          _lowerCamelCase  :		Dict			    =    pipe.to(_UpperCamelCase)
          pipe.set_progress_bar_config(disable=_UpperCamelCase)
          pipe.enable_attention_slicing()
          pipe.enable_sequential_cpu_offload()

          _lowerCamelCase  :		List[str]			    =    pipe(
              """horse"""     ,	num_images_per_prompt=1     ,	prior_num_inference_steps=2     ,	decoder_num_inference_steps=2     ,	super_res_num_inference_steps=2     ,	output_type="""np"""     ,	)

          _lowerCamelCase  :		Optional[int]			    =    torch.cuda.max_memory_allocated()
          # make sure that less than 7 GB is allocated
          assert mem_bytes < 7 * 10**9





from __future__ import annotations







def     A__	(       __A ,			__A						):







    '''simple docstring'''
    _lowerCamelCase  :		List[str]			    =    [None] * (len(__A						) + len(__A						))

    _lowerCamelCase  :		Optional[int]			    =    0  # pointer to current index for left Half
    _lowerCamelCase  :		List[str]			    =    0  # pointer to current index for the right Half
    _lowerCamelCase  :		List[str]			    =    0  # pointer to current index for the sorted array Half

    while pointera < len(__A						) and pointera < len(__A						):
        if left_half[pointera] < right_half[pointera]:
            _lowerCamelCase  :		Optional[Any]			    =    left_half[pointera]
            pointera += 1
            index += 1
        else:
            _lowerCamelCase  :		List[Any]			    =    right_half[pointera]
            pointera += 1
            index += 1
    while pointera < len(__A						):
        _lowerCamelCase  :		int			    =    left_half[pointera]
        pointera += 1
        index += 1

    while pointera < len(__A						):
        _lowerCamelCase  :		Optional[Any]			    =    right_half[pointera]
        pointera += 1
        index += 1

    return sorted_array







def     A__	(       __A						):







    '''simple docstring'''
    if len(__A						) <= 1:
        return array
    # the actual formula to calculate the middle element = left + (right - left) // 2
    # this avoids integer overflow in case of large N
    _lowerCamelCase  :		int			    =    0 + (len(__A						) - 0) // 2

    # Split the array into halves till the array length becomes equal to One
    # merge the arrays of single length returned by mergeSort function and
    # pass them into the merge arrays function which merges the array
    _lowerCamelCase  :		List[str]			    =    array[:middle]
    _lowerCamelCase  :		Optional[int]			    =    array[middle:]

    return merge(merge_sort(__A						) ,			merge_sort(__A						)						)


if __name__ == "__main__":
    import doctest

    doctest.testmod()




