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    __lowerCAmelCase  (       lowerCAmelCase_,	unittest.TestCase						):



 '''simple docstring'''
 _lowerCamelCase						:	Union[str, Any]				    = UnCLIPPipeline
 _lowerCamelCase						:	Optional[int]				    = TEXT_TO_IMAGE_PARAMS - {
     """negative_prompt""",
     """height""",
     """width""",
     """negative_prompt_embeds""",
     """guidance_scale""",
     """prompt_embeds""",
     """cross_attention_kwargs""",
 }
 _lowerCamelCase						:	int				    = TEXT_TO_IMAGE_BATCH_PARAMS
 _lowerCamelCase						:	Optional[int]				    = [
     """generator""",
     """return_dict""",
     """prior_num_inference_steps""",
     """decoder_num_inference_steps""",
     """super_res_num_inference_steps""",
 ]
 _lowerCamelCase						:	Optional[int]				    = False

 @property
 def 					__snake_case					(			self       :		Union[str, Any]  ) ->							Union[str, Any]:
        return 3_2

 @property
 def 					__snake_case					(			self       :		List[Any]  ) ->							Optional[Any]:
        return 3_2

 @property
 def 					__snake_case					(			self       :		List[str]  ) ->							Optional[int]:
        return self.time_input_dim

 @property
 def 					__snake_case					(			self       :		Optional[int]  ) ->							Optional[Any]:
        return self.time_input_dim * 4

 @property
 def 					__snake_case					(			self       :		Optional[int]  ) ->							Any:
        return 1_0_0

 @property
 def 					__snake_case					(			self       :		Optional[Any]  ) ->							int:
        _UpperCamelCase     :       str  						=  CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'  )
        return tokenizer

 @property
 def 					__snake_case					(			self       :		Dict  ) ->							Tuple:
        torch.manual_seed(0  )
        _UpperCamelCase     :       Tuple  						=  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=3_7       ,			layer_norm_eps=1e-05       ,			num_attention_heads=4       ,			num_hidden_layers=5       ,			pad_token_id=1       ,			vocab_size=1_0_0_0       ,			)
        return CLIPTextModelWithProjection(__SCREAMING_SNAKE_CASE  )

 @property
 def 					__snake_case					(			self       :		str  ) ->							List[Any]:
        torch.manual_seed(0  )

        _UpperCamelCase     :       str  						=  {
            'num_attention_heads': 2,
            'attention_head_dim': 1_2,
            'embedding_dim': self.text_embedder_hidden_size,
            'num_layers': 1,
        }

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

 @property
 def 					__snake_case					(			self       :		str  ) ->							Union[str, Any]:
        torch.manual_seed(0  )

        _UpperCamelCase     :       int  						=  {
            'clip_embeddings_dim': self.text_embedder_hidden_size,
            'time_embed_dim': self.time_embed_dim,
            'cross_attention_dim': self.cross_attention_dim,
        }

        _UpperCamelCase     :       Dict  						=  UnCLIPTextProjModel(**__SCREAMING_SNAKE_CASE  )
        return model

 @property
 def 					__snake_case					(			self       :		Dict  ) ->							Dict:
        torch.manual_seed(0  )

        _UpperCamelCase     :       Optional[Any]  						=  {
            'sample_size': 3_2,
            # 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',
        }

        _UpperCamelCase     :       List[str]  						=  UNetaDConditionModel(**__SCREAMING_SNAKE_CASE  )
        return model

 @property
 def 					__snake_case					(			self       :		Tuple  ) ->							Tuple:
        return {
            "sample_size": 6_4,
            "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 					__snake_case					(			self       :		Optional[Any]  ) ->							List[Any]:
        torch.manual_seed(0  )

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

 @property
 def 					__snake_case					(			self       :		str  ) ->							Tuple:
        torch.manual_seed(1  )

        _UpperCamelCase     :       Tuple  						=  UNetaDModel(**self.dummy_super_res_kwargs  )
        return model

 def 					__snake_case					(			self       :		int  ) ->							Optional[Any]:
        _UpperCamelCase     :       Union[str, Any]  						=  self.dummy_prior
        _UpperCamelCase     :       str  						=  self.dummy_decoder
        _UpperCamelCase     :       Optional[int]  						=  self.dummy_text_proj
        _UpperCamelCase     :       List[Any]  						=  self.dummy_text_encoder
        _UpperCamelCase     :       str  						=  self.dummy_tokenizer
        _UpperCamelCase     :       List[str]  						=  self.dummy_super_res_first
        _UpperCamelCase     :       Dict  						=  self.dummy_super_res_last

        _UpperCamelCase     :       Optional[int]  						=  UnCLIPScheduler(
            variance_type='fixed_small_log'       ,			prediction_type='sample'       ,			num_train_timesteps=1_0_0_0       ,			clip_sample_range=5.0       ,			)

        _UpperCamelCase     :       int  						=  UnCLIPScheduler(
            variance_type='learned_range'       ,			prediction_type='epsilon'       ,			num_train_timesteps=1_0_0_0       ,			)

        _UpperCamelCase     :       Dict  						=  UnCLIPScheduler(
            variance_type='fixed_small_log'       ,			prediction_type='epsilon'       ,			num_train_timesteps=1_0_0_0       ,			)

        _UpperCamelCase     :       Optional[Any]  						=  {
            '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 					__snake_case					(			self       :		Union[str, Any]       ,			lowerCamelCase       :		Any       ,			lowerCamelCase       :		List[Any]=0  ) ->							Union[str, Any]:

        if str(__SCREAMING_SNAKE_CASE  ).startswith('mps'  ):
               _UpperCamelCase     :       Any  						=  torch.manual_seed(__SCREAMING_SNAKE_CASE  )
        else:
               _UpperCamelCase     :       Union[str, Any]  						=  torch.Generator(device=__SCREAMING_SNAKE_CASE  ).manual_seed(__SCREAMING_SNAKE_CASE  )
        _UpperCamelCase     :       int  						=  {
            '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 					__snake_case					(			self       :		Any  ) ->							int:
        _UpperCamelCase     :       Optional[int]  						=  'cpu'

        _UpperCamelCase     :       Dict  						=  self.get_dummy_components()

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

        pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE  )

        _UpperCamelCase     :       Optional[Any]  						=  pipe(**self.get_dummy_inputs(__SCREAMING_SNAKE_CASE  )  )
        _UpperCamelCase     :       List[Any]  						=  output.images

        _UpperCamelCase     :       Union[str, Any]  						=  pipe(
            **self.get_dummy_inputs(__SCREAMING_SNAKE_CASE  )       ,			return_dict=__SCREAMING_SNAKE_CASE       ,			)[0]

        _UpperCamelCase     :       Tuple  						=  image[0, -3:, -3:, -1]
        _UpperCamelCase     :       str  						=  image_from_tuple[0, -3:, -3:, -1]

        assert image.shape == (1, 6_4, 6_4, 3)

        _UpperCamelCase     :       int  						=  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 					__snake_case					(			self       :		List[str]  ) ->							Tuple:
        _UpperCamelCase     :       List[Any]  						=  torch.device('cpu'  )




        class    __lowerCAmelCase  :



         '''simple docstring'''
         _lowerCamelCase						:	Optional[Any]				    = 1

        _UpperCamelCase     :       Tuple  						=  self.get_dummy_components()

        _UpperCamelCase     :       List[Any]  						=  self.pipeline_class(**__SCREAMING_SNAKE_CASE  )
        _UpperCamelCase     :       Any  						=  pipe.to(__SCREAMING_SNAKE_CASE  )

        _UpperCamelCase     :       Tuple  						=  components['prior']
        _UpperCamelCase     :       Optional[int]  						=  components['decoder']
        _UpperCamelCase     :       Union[str, Any]  						=  components['super_res_first']
        _UpperCamelCase     :       int  						=  components['tokenizer']
        _UpperCamelCase     :       Tuple  						=  components['text_encoder']

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

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

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

        pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE  )

        _UpperCamelCase     :       List[str]  						=  'this is a prompt example'

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

        _UpperCamelCase     :       Tuple  						=  tokenizer(
            __SCREAMING_SNAKE_CASE       ,			padding='max_length'       ,			max_length=tokenizer.model_max_length       ,			return_tensors='pt'       ,			)
        _UpperCamelCase     :       List[str]  						=  text_encoder(text_inputs.input_ids  )
        _UpperCamelCase     :       Optional[int]  						=  text_inputs.attention_mask

        _UpperCamelCase     :       Tuple  						=  torch.Generator(device=__SCREAMING_SNAKE_CASE  ).manual_seed(0  )
        _UpperCamelCase     :       Any  						=  pipe(
            generator=__SCREAMING_SNAKE_CASE       ,			prior_num_inference_steps=2       ,			decoder_num_inference_steps=2       ,			super_res_num_inference_steps=2       ,			prior_latents=__SCREAMING_SNAKE_CASE       ,			decoder_latents=__SCREAMING_SNAKE_CASE       ,			super_res_latents=__SCREAMING_SNAKE_CASE       ,			text_model_output=__SCREAMING_SNAKE_CASE       ,			text_attention_mask=__SCREAMING_SNAKE_CASE       ,			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 					__snake_case					(			self       :		Tuple  ) ->							Optional[Any]:
        _UpperCamelCase     :       Union[str, Any]  						=  torch_device == 'cpu'

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

 @skip_mps
 def 					__snake_case					(			self       :		List[Any]  ) ->							Union[str, Any]:
        _UpperCamelCase     :       Union[str, Any]  						=  torch_device == 'cpu'
        _UpperCamelCase     :       Optional[int]  						=  True
        _UpperCamelCase     :       int  						=  [
            'prior_num_inference_steps',
            'decoder_num_inference_steps',
            'super_res_num_inference_steps',
        ]

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

 def 					__snake_case					(			self       :		Dict  ) ->							List[Any]:
        _UpperCamelCase     :       Optional[int]  						=  [
            'prior_num_inference_steps',
            'decoder_num_inference_steps',
            'super_res_num_inference_steps',
        ]

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

 @skip_mps
 def 					__snake_case					(			self       :		List[str]  ) ->							Tuple:
        return super().test_dict_tuple_outputs_equivalent()

 @skip_mps
 def 					__snake_case					(			self       :		Optional[Any]  ) ->							str:
        return super().test_save_load_local()

 @skip_mps
 def 					__snake_case					(			self       :		List[Any]  ) ->							int:
        return super().test_save_load_optional_components()




@nightly
class    __lowerCAmelCase  (       unittest.TestCase						):



 '''simple docstring'''

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

 def 					__snake_case					(			self       :		Tuple  ) ->							int:
        _UpperCamelCase     :       Dict  						=  load_numpy(
            'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
            '/unclip/karlo_v1_alpha_horse_cpu.npy'  )

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

        _UpperCamelCase     :       int  						=  torch.manual_seed(0  )
        _UpperCamelCase     :       str  						=  pipeline(
            'horse'       ,			num_images_per_prompt=1       ,			generator=__SCREAMING_SNAKE_CASE       ,			output_type='np'       ,			)

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

        assert image.shape == (2_5_6, 2_5_6, 3)
        assert np.abs(expected_image - image  ).max() < 1e-1




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



 '''simple docstring'''

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

 def 					__snake_case					(			self       :		List[str]  ) ->							Dict:
        _UpperCamelCase     :       List[str]  						=  load_numpy(
            'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
            '/unclip/karlo_v1_alpha_horse_fp16.npy'  )

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

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

        _UpperCamelCase     :       Tuple  						=  output.images[0]

        assert image.shape == (2_5_6, 2_5_6, 3)

        assert_mean_pixel_difference(__SCREAMING_SNAKE_CASE       ,			__SCREAMING_SNAKE_CASE  )

 def 					__snake_case					(			self       :		List[Any]  ) ->							int:
        torch.cuda.empty_cache()
        torch.cuda.reset_max_memory_allocated()
        torch.cuda.reset_peak_memory_stats()

        _UpperCamelCase     :       Tuple  						=  UnCLIPPipeline.from_pretrained('kakaobrain/karlo-v1-alpha'       ,			torch_dtype=torch.floataa  )
        _UpperCamelCase     :       Dict  						=  pipe.to(__SCREAMING_SNAKE_CASE  )
        pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE  )
        pipe.enable_attention_slicing()
        pipe.enable_sequential_cpu_offload()

        _UpperCamelCase     :       Any  						=  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'       ,			)

        _UpperCamelCase     :       str  						=  torch.cuda.max_memory_allocated()
        # make sure that less than 7 GB is allocated
        assert mem_bytes < 7 * 1_0**9




from __future__ import annotations


def SCREAMING_SNAKE_CASE				(     lowercase_:	int		,						lowercase_:	int		,						lowercase_:	int				):
       assert (
           c % greatest_common_divisor(lowercase_		,						lowercase_				) == 0
       )  # greatest_common_divisor(a,b) function implemented below
       ((_UpperCamelCase)						,      (_UpperCamelCase)						,      (_UpperCamelCase))     :       int  						=  extended_gcd(lowercase_		,						lowercase_				)  # extended_gcd(a,b) function implemented below
       _UpperCamelCase     :       List[Any]  						=  c / d
       return (r * x, r * y)


def SCREAMING_SNAKE_CASE				(     lowercase_:	int		,						lowercase_:	int		,						lowercase_:	int		,						lowercase_:	int = 2				):
       ((_UpperCamelCase)						,      (_UpperCamelCase))     :       Tuple  						=  diophantine(lowercase_		,						lowercase_		,						lowercase_				)  # Initial value
       _UpperCamelCase     :       List[Any]  						=  greatest_common_divisor(lowercase_		,						lowercase_				)
       _UpperCamelCase     :       Optional[int]  						=  a // d
       _UpperCamelCase     :       Optional[Any]  						=  b // d

       for i in range(lowercase_				):
              _UpperCamelCase     :       Union[str, Any]  						=  xa + i * q
              _UpperCamelCase     :       List[Any]  						=  ya - i * p
              print(lowercase_		,						lowercase_				)


def SCREAMING_SNAKE_CASE				(     lowercase_:	int		,						lowercase_:	int				):
       if a < b:
              _UpperCamelCase						,      _UpperCamelCase     :       Optional[Any]  						=  b, a

       while a % b != 0:
              _UpperCamelCase						,      _UpperCamelCase     :       Union[str, Any]  						=  b, a % b

       return b


def SCREAMING_SNAKE_CASE				(     lowercase_:	int		,						lowercase_:	int				):
       assert a >= 0 and b >= 0

       if b == 0:
              _UpperCamelCase						,      _UpperCamelCase						,      _UpperCamelCase     :       Optional[Any]  						=  a, 1, 0
       else:
              ((_UpperCamelCase)						,      (_UpperCamelCase)						,      (_UpperCamelCase))     :       Any  						=  extended_gcd(lowercase_		,						a % b				)
              _UpperCamelCase     :       Any  						=  q
              _UpperCamelCase     :       str  						=  p - q * (a // b)

       assert a % d == 0 and b % d == 0
       assert d == a * x + b * y

       return (d, x, y)


if __name__ == "__main__":
 from doctest import testmod

 testmod(name='''diophantine''', verbose=True)
 testmod(name='''diophantine_all_soln''', verbose=True)
 testmod(name='''extended_gcd''', verbose=True)
 testmod(name='''greatest_common_divisor''', verbose=True)



