"""simple docstring"""

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



						"""simple docstring"""




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





						@property
						def   UpperCAmelCase_ (				self							:    Optional[int]				)				->				str:
													'''simple docstring'''



													return 32





						@property
						def   UpperCAmelCase_ (				self							:    Optional[Any]				)				->				List[str]:
													'''simple docstring'''



													return 32





						@property
						def   UpperCAmelCase_ (				self							:    str				)				->				List[Any]:
													'''simple docstring'''



													return self.time_input_dim





						@property
						def   UpperCAmelCase_ (				self							:    int				)				->				List[str]:
													'''simple docstring'''



													return self.time_input_dim * 4





						@property
						def   UpperCAmelCase_ (				self							:    Tuple				)				->				Dict:
													'''simple docstring'''



													return 100





						@property
						def   UpperCAmelCase_ (				self							:    int				)				->				Tuple:
													'''simple docstring'''



													__UpperCamelCase							=CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''				)
													return tokenizer





						@property
						def   UpperCAmelCase_ (				self							:    Tuple				)				->				Union[str, Any]:
													'''simple docstring'''



													torch.manual_seed(0				)
													__UpperCamelCase							=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-05     , num_attention_heads=4     , num_hidden_layers=5     , pad_token_id=1     , vocab_size=1000     , )
													return CLIPTextModelWithProjection(__A				)





						@property
						def   UpperCAmelCase_ (				self							:    List[Any]				)				->				int:
													'''simple docstring'''



													torch.manual_seed(0				)

													__UpperCamelCase							={
													    '''num_attention_heads''': 2,
													    '''attention_head_dim''': 12,
													    '''embedding_dim''': self.text_embedder_hidden_size,
													    '''num_layers''': 1,
													}

													__UpperCamelCase							=PriorTransformer(**__A				)
													return model





						@property
						def   UpperCAmelCase_ (				self							:    Dict				)				->				Optional[Any]:
													'''simple docstring'''



													torch.manual_seed(0				)

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

													__UpperCamelCase							=UnCLIPTextProjModel(**__A				)
													return model





						@property
						def   UpperCAmelCase_ (				self							:    Tuple				)				->				List[Any]:
													'''simple docstring'''



													torch.manual_seed(0				)

													__UpperCamelCase							={
													    '''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''',
													}

													__UpperCamelCase							=UNetaDConditionModel(**__A				)
													return model





						@property
						def   UpperCAmelCase_ (				self							:    Union[str, Any]				)				->				List[str]:
													'''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   UpperCAmelCase_ (				self							:    int				)				->				Optional[Any]:
													'''simple docstring'''



													torch.manual_seed(0				)

													__UpperCamelCase							=UNetaDModel(**self.dummy_super_res_kwargs				)
													return model





						@property
						def   UpperCAmelCase_ (				self							:    int				)				->				Optional[Any]:
													'''simple docstring'''



													torch.manual_seed(1				)

													__UpperCamelCase							=UNetaDModel(**self.dummy_super_res_kwargs				)
													return model





						def   UpperCAmelCase_ (				self							:    Optional[int]				)				->				Tuple:
													'''simple docstring'''



													__UpperCamelCase							=self.dummy_prior
													__UpperCamelCase							=self.dummy_decoder
													__UpperCamelCase							=self.dummy_text_proj
													__UpperCamelCase							=self.dummy_text_encoder
													__UpperCamelCase							=self.dummy_tokenizer
													__UpperCamelCase							=self.dummy_super_res_first
													__UpperCamelCase							=self.dummy_super_res_last

													__UpperCamelCase							=UnCLIPScheduler(
													    variance_type='''fixed_small_log'''     , prediction_type='''sample'''     , num_train_timesteps=1000     , clip_sample_range=5.0     , )

													__UpperCamelCase							=UnCLIPScheduler(
													    variance_type='''learned_range'''     , prediction_type='''epsilon'''     , num_train_timesteps=1000     , )

													__UpperCamelCase							=UnCLIPScheduler(
													    variance_type='''fixed_small_log'''     , prediction_type='''epsilon'''     , num_train_timesteps=1000     , )

													__UpperCamelCase							={
													    '''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   UpperCAmelCase_ (				self							:    List[Any]     , UpperCamelCase__							:    Dict     , UpperCamelCase__							:    int=0				)				->				int:
													'''simple docstring'''



													if str(__A				).startswith('''mps'''				):
																				__UpperCamelCase							=torch.manual_seed(__A				)
													else:
																				__UpperCamelCase							=torch.Generator(device=__A				).manual_seed(__A				)
													__UpperCamelCase							={
													    '''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   UpperCAmelCase_ (				self							:    Optional[int]				)				->				Optional[Any]:
													'''simple docstring'''



													__UpperCamelCase							='''cpu'''

													__UpperCamelCase							=self.get_dummy_components()

													__UpperCamelCase							=self.pipeline_class(**__A				)
													__UpperCamelCase							=pipe.to(__A				)

													pipe.set_progress_bar_config(disable=__A				)

													__UpperCamelCase							=pipe(**self.get_dummy_inputs(__A				)				)
													__UpperCamelCase							=output.images

													__UpperCamelCase							=pipe(
													    **self.get_dummy_inputs(__A				)     , return_dict=__A     , )[0]

													__UpperCamelCase							=image[0, -3:, -3:, -1]
													__UpperCamelCase							=image_from_tuple[0, -3:, -3:, -1]

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

													__UpperCamelCase							=np.array(
													    [
													        0.99_97,
													        0.99_88,
													        0.00_28,
													        0.99_97,
													        0.99_84,
													        0.99_65,
													        0.00_29,
													        0.99_86,
													        0.00_25,
													    ]				)

													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   UpperCAmelCase_ (				self							:    List[Any]				)				->				str:
													'''simple docstring'''



													__UpperCamelCase							=torch.device('''cpu'''				)






													class 		_lowercase :



																			"""simple docstring"""




																			lowercase__       =  1






													__UpperCamelCase							=self.get_dummy_components()

													__UpperCamelCase							=self.pipeline_class(**__A				)
													__UpperCamelCase							=pipe.to(__A				)

													__UpperCamelCase							=components['''prior''']
													__UpperCamelCase							=components['''decoder''']
													__UpperCamelCase							=components['''super_res_first''']
													__UpperCamelCase							=components['''tokenizer''']
													__UpperCamelCase							=components['''text_encoder''']

													__UpperCamelCase							=torch.Generator(device=__A				).manual_seed(0				)
													__UpperCamelCase							=prior.dtype
													__UpperCamelCase							=1

													__UpperCamelCase							=(batch_size, prior.config.embedding_dim)
													__UpperCamelCase							=pipe.prepare_latents(
													    __A     , dtype=__A     , device=__A     , generator=__A     , latents=__A     , scheduler=DummyScheduler()				)
													__UpperCamelCase							=(batch_size, decoder.config.in_channels, decoder.config.sample_size, decoder.config.sample_size)
													__UpperCamelCase							=pipe.prepare_latents(
													    __A     , dtype=__A     , device=__A     , generator=__A     , latents=__A     , scheduler=DummyScheduler()				)

													__UpperCamelCase							=(
													    batch_size,
													    super_res_first.config.in_channels // 2,
													    super_res_first.config.sample_size,
													    super_res_first.config.sample_size,
													)
													__UpperCamelCase							=pipe.prepare_latents(
													    __A     , dtype=__A     , device=__A     , generator=__A     , latents=__A     , scheduler=DummyScheduler()				)

													pipe.set_progress_bar_config(disable=__A				)

													__UpperCamelCase							='''this is a prompt example'''

													__UpperCamelCase							=torch.Generator(device=__A				).manual_seed(0				)
													__UpperCamelCase							=pipe(
													    [prompt]     , generator=__A     , prior_num_inference_steps=2     , decoder_num_inference_steps=2     , super_res_num_inference_steps=2     , prior_latents=__A     , decoder_latents=__A     , super_res_latents=__A     , output_type='''np'''     , )
													__UpperCamelCase							=output.images

													__UpperCamelCase							=tokenizer(
													    __A     , padding='''max_length'''     , max_length=tokenizer.model_max_length     , return_tensors='''pt'''     , )
													__UpperCamelCase							=text_encoder(text_inputs.input_ids				)
													__UpperCamelCase							=text_inputs.attention_mask

													__UpperCamelCase							=torch.Generator(device=__A				).manual_seed(0				)
													__UpperCamelCase							=pipe(
													    generator=__A     , prior_num_inference_steps=2     , decoder_num_inference_steps=2     , super_res_num_inference_steps=2     , prior_latents=__A     , decoder_latents=__A     , super_res_latents=__A     , text_model_output=__A     , text_attention_mask=__A     , 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   UpperCAmelCase_ (				self							:    Tuple				)				->				List[Any]:
													'''simple docstring'''



													__UpperCamelCase							=torch_device == '''cpu'''

													self._test_attention_slicing_forward_pass(test_max_difference=__A     , expected_max_diff=0.01				)





						@skip_mps
						def   UpperCAmelCase_ (				self							:    List[Any]				)				->				int:
													'''simple docstring'''



													__UpperCamelCase							=torch_device == '''cpu'''
													__UpperCamelCase							=True
													__UpperCamelCase							=[
													    '''prior_num_inference_steps''',
													    '''decoder_num_inference_steps''',
													    '''super_res_num_inference_steps''',
													]

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





						def   UpperCAmelCase_ (				self							:    Any				)				->				List[str]:
													'''simple docstring'''



													__UpperCamelCase							=[
													    '''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							=[2, 3]
																				self._test_inference_batch_consistent(
																				    batch_sizes=__A     , additional_params_copy_to_batched_inputs=__A     , )
													else:
																				self._test_inference_batch_consistent(
																				    additional_params_copy_to_batched_inputs=__A				)





						@skip_mps
						def   UpperCAmelCase_ (				self							:    List[str]				)				->				List[Any]:
													'''simple docstring'''



													return super().test_dict_tuple_outputs_equivalent()





						@skip_mps
						def   UpperCAmelCase_ (				self							:    Dict				)				->				Tuple:
													'''simple docstring'''



													return super().test_save_load_local()





						@skip_mps
						def   UpperCAmelCase_ (				self							:    List[Any]				)				->				Tuple:
													'''simple docstring'''



													return super().test_save_load_optional_components()






@nightly
class 		_lowercase (						unittest.TestCase      ):



						"""simple docstring"""




						def   UpperCAmelCase_ (				self							:    Optional[int]				)				->				Dict:
													'''simple docstring'''



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





						def   UpperCAmelCase_ (				self							:    int				)				->				Optional[Any]:
													'''simple docstring'''



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

													__UpperCamelCase							=UnCLIPPipeline.from_pretrained('''kakaobrain/karlo-v1-alpha'''				)
													pipeline.set_progress_bar_config(disable=__A				)

													__UpperCamelCase							=torch.manual_seed(0				)
													__UpperCamelCase							=pipeline(
													    '''horse'''     , num_images_per_prompt=1     , generator=__A     , output_type='''np'''     , )

													__UpperCamelCase							=output.images[0]

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






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



						"""simple docstring"""




						def   UpperCAmelCase_ (				self							:    Union[str, Any]				)				->				Optional[Any]:
													'''simple docstring'''



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





						def   UpperCAmelCase_ (				self							:    Union[str, Any]				)				->				Union[str, Any]:
													'''simple docstring'''



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

													__UpperCamelCase							=UnCLIPPipeline.from_pretrained('''kakaobrain/karlo-v1-alpha'''     , torch_dtype=torch.floataa				)
													__UpperCamelCase							=pipeline.to(__A				)
													pipeline.set_progress_bar_config(disable=__A				)

													__UpperCamelCase							=torch.Generator(device='''cpu'''				).manual_seed(0				)
													__UpperCamelCase							=pipeline(
													    '''horse'''     , generator=__A     , output_type='''np'''     , )

													__UpperCamelCase							=output.images[0]

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

													assert_mean_pixel_difference(__A     , __A				)





						def   UpperCAmelCase_ (				self							:    List[Any]				)				->				Dict:
													'''simple docstring'''



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

													__UpperCamelCase							=UnCLIPPipeline.from_pretrained('''kakaobrain/karlo-v1-alpha'''     , torch_dtype=torch.floataa				)
													__UpperCamelCase							=pipe.to(__A				)
													pipe.set_progress_bar_config(disable=__A				)
													pipe.enable_attention_slicing()
													pipe.enable_sequential_cpu_offload()

													__UpperCamelCase							=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							=torch.cuda.max_memory_allocated()
													# make sure that less than 7 GB is allocated
													assert mem_bytes < 7 * 10**9







"""simple docstring"""







import torch
import torch.nn as nn
from transformers.modeling_utils import ModuleUtilsMixin
from transformers.models.ta.modeling_ta import (
    TaBlock,
    TaConfig,
    TaLayerNorm,
)

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin






class  __A							(			_SCREAMING_SNAKE_CASE,       _SCREAMING_SNAKE_CASE,       _SCREAMING_SNAKE_CASE     ):






 """simple docstring"""




 @register_to_config
 def __init__(						self					,       __A					,       __A					,       __A					,       __A					,       __A					,       __A					,       __A					,       __A					,       __A					,       __A = False					,       )   ->	Optional[Any]:
     super().__init__()

     a											=nn.Linear(__A					,       __A					,       bias=__A		)

     a											=nn.Embedding(__A					,       __A		)
     a											=False

     a											=nn.Dropout(p=__A		)

     a											=TaConfig(
         d_model=__A					,       num_heads=__A					,       d_kv=__A					,       d_ff=__A					,       feed_forward_proj=__A					,       dropout_rate=__A					,       is_decoder=__A					,       is_encoder_decoder=__A					,       )
     a											=nn.ModuleList()
     for lyr_num in range(__A		):
         a											=TaBlock(__A		)
         self.encoders.append(__A		)

     a											=TaLayerNorm(__A		)
     a											=nn.Dropout(p=__A		)





 def  SCREAMING_SNAKE_CASE				(						self					,       __A					,       __A		)   ->	str:
     a											=self.input_proj(__A		)

     # terminal relative positional encodings
     a											=encoder_inputs.shape[1]
     a											=torch.arange(__A					,       device=encoder_inputs.device		)

     a											=encoder_inputs_mask.sum(-1		)
     a											=torch.roll(input_positions.unsqueeze(0		)					,       tuple(seq_lens.tolist()		)					,       dims=0		)
     x += self.position_encoding(__A		)

     a											=self.dropout_pre(__A		)

     # inverted the attention mask
     a											=encoder_inputs.size()
     a											=self.get_extended_attention_mask(__A					,       __A		)

     for lyr in self.encoders:
         a											=lyr(__A					,       __A		)[0]
     a											=self.layer_norm(__A		)

     return self.dropout_post(__A		), encoder_inputs_mask