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





 """simple docstring"""


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

 @property
 def       UpperCamelCase__(					self						):



     '''simple docstring'''




     return 32

 @property
 def       UpperCamelCase__(					self						):



     '''simple docstring'''




     return 32

 @property
 def       UpperCamelCase__(					self						):



     '''simple docstring'''




     return self.time_input_dim

 @property
 def       UpperCamelCase__(					self						):



     '''simple docstring'''




     return self.time_input_dim * 4

 @property
 def       UpperCamelCase__(					self						):



     '''simple docstring'''




     return 100

 @property
 def       UpperCamelCase__(					self						):



     '''simple docstring'''




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

 @property
 def       UpperCamelCase__(					self						):



     '''simple docstring'''




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

 @property
 def       UpperCamelCase__(					self						):



     '''simple docstring'''




     torch.manual_seed(0						)

     __A					:		Union[str, Any]						=  {
         '''num_attention_heads''': 2,
         '''attention_head_dim''': 12,
         '''embedding_dim''': self.text_embedder_hidden_size,
         '''num_layers''': 1,
     }

     __A					:		List[Any]						=  PriorTransformer(**__lowerCamelCase						)
     return model

 @property
 def       UpperCamelCase__(					self						):



     '''simple docstring'''




     torch.manual_seed(0						)

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

     __A					:		str						=  UnCLIPTextProjModel(**__lowerCamelCase						)
     return model

 @property
 def       UpperCamelCase__(					self						):



     '''simple docstring'''




     torch.manual_seed(0						)

     __A					:		List[str]						=  {
         '''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''',
     }

     __A					:		Union[str, Any]						=  UNetaDConditionModel(**__lowerCamelCase						)
     return model

 @property
 def       UpperCamelCase__(					self						):



     '''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						):



     '''simple docstring'''




     torch.manual_seed(0						)

     __A					:		Any						=  UNetaDModel(**self.dummy_super_res_kwargs						)
     return model

 @property
 def       UpperCamelCase__(					self						):



     '''simple docstring'''




     torch.manual_seed(1						)

     __A					:		Dict						=  UNetaDModel(**self.dummy_super_res_kwargs						)
     return model

 def       UpperCamelCase__(					self						):



     '''simple docstring'''




     __A					:		Any						=  self.dummy_prior
     __A					:		Tuple						=  self.dummy_decoder
     __A					:		List[str]						=  self.dummy_text_proj
     __A					:		int						=  self.dummy_text_encoder
     __A					:		Optional[Any]						=  self.dummy_tokenizer
     __A					:		List[str]						=  self.dummy_super_res_first
     __A					:		int						=  self.dummy_super_res_last

     __A					:		Optional[Any]						=  UnCLIPScheduler(
         variance_type='''fixed_small_log'''      ,		prediction_type='''sample'''      ,		num_train_timesteps=1000      ,		clip_sample_range=5.0      ,		)

     __A					:		Optional[int]						=  UnCLIPScheduler(
         variance_type='''learned_range'''      ,		prediction_type='''epsilon'''      ,		num_train_timesteps=1000      ,		)

     __A					:		Union[str, Any]						=  UnCLIPScheduler(
         variance_type='''fixed_small_log'''      ,		prediction_type='''epsilon'''      ,		num_train_timesteps=1000      ,		)

     __A					:		Optional[int]						=  {
         '''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      ,		__lowerCamelCase      ,		__lowerCamelCase=0						):



     '''simple docstring'''




     if str(__lowerCamelCase						).startswith('''mps'''						):
         __A					:		Any						=  torch.manual_seed(__lowerCamelCase						)
     else:
         __A					:		Tuple						=  torch.Generator(device=__lowerCamelCase						).manual_seed(__lowerCamelCase						)
     __A					:		Tuple						=  {
         '''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						):



     '''simple docstring'''




     __A					:		Union[str, Any]						=  '''cpu'''

     __A					:		Dict						=  self.get_dummy_components()

     __A					:		Union[str, Any]						=  self.pipeline_class(**__lowerCamelCase						)
     __A					:		int						=  pipe.to(__lowerCamelCase						)

     pipe.set_progress_bar_config(disable=__lowerCamelCase						)

     __A					:		Tuple						=  pipe(**self.get_dummy_inputs(__lowerCamelCase						)						)
     __A					:		List[Any]						=  output.images

     __A					:		Optional[int]						=  pipe(
         **self.get_dummy_inputs(__lowerCamelCase						)      ,		return_dict=__lowerCamelCase      ,		)[0]

     __A					:		Dict						=  image[0, -3:, -3:, -1]
     __A					:		List[str]						=  image_from_tuple[0, -3:, -3:, -1]

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

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



     '''simple docstring'''




     __A					:		List[str]						=  torch.device('''cpu'''						)






     class 	__snake_case      :





      """simple docstring"""


      _lowerCamelCase				 =  1






     __A					:		Union[str, Any]						=  self.get_dummy_components()

     __A					:		Union[str, Any]						=  self.pipeline_class(**__lowerCamelCase						)
     __A					:		Optional[Any]						=  pipe.to(__lowerCamelCase						)

     __A					:		Optional[Any]						=  components['''prior''']
     __A					:		Optional[Any]						=  components['''decoder''']
     __A					:		Optional[Any]						=  components['''super_res_first''']
     __A					:		Optional[int]						=  components['''tokenizer''']
     __A					:		Any						=  components['''text_encoder''']

     __A					:		Union[str, Any]						=  torch.Generator(device=__lowerCamelCase						).manual_seed(0						)
     __A					:		int						=  prior.dtype
     __A					:		Union[str, Any]						=  1

     __A					:		List[str]						=  (batch_size, prior.config.embedding_dim)
     __A					:		Union[str, Any]						=  pipe.prepare_latents(
         __lowerCamelCase      ,		dtype=__lowerCamelCase      ,		device=__lowerCamelCase      ,		generator=__lowerCamelCase      ,		latents=__lowerCamelCase      ,		scheduler=DummyScheduler()						)
     __A					:		List[Any]						=  (batch_size, decoder.config.in_channels, decoder.config.sample_size, decoder.config.sample_size)
     __A					:		Tuple						=  pipe.prepare_latents(
         __lowerCamelCase      ,		dtype=__lowerCamelCase      ,		device=__lowerCamelCase      ,		generator=__lowerCamelCase      ,		latents=__lowerCamelCase      ,		scheduler=DummyScheduler()						)

     __A					:		List[str]						=  (
         batch_size,
         super_res_first.config.in_channels // 2,
         super_res_first.config.sample_size,
         super_res_first.config.sample_size,
     )
     __A					:		Optional[int]						=  pipe.prepare_latents(
         __lowerCamelCase      ,		dtype=__lowerCamelCase      ,		device=__lowerCamelCase      ,		generator=__lowerCamelCase      ,		latents=__lowerCamelCase      ,		scheduler=DummyScheduler()						)

     pipe.set_progress_bar_config(disable=__lowerCamelCase						)

     __A					:		int						=  '''this is a prompt example'''

     __A					:		Dict						=  torch.Generator(device=__lowerCamelCase						).manual_seed(0						)
     __A					:		Optional[int]						=  pipe(
         [prompt]      ,		generator=__lowerCamelCase      ,		prior_num_inference_steps=2      ,		decoder_num_inference_steps=2      ,		super_res_num_inference_steps=2      ,		prior_latents=__lowerCamelCase      ,		decoder_latents=__lowerCamelCase      ,		super_res_latents=__lowerCamelCase      ,		output_type='''np'''      ,		)
     __A					:		List[Any]						=  output.images

     __A					:		Dict						=  tokenizer(
         __lowerCamelCase      ,		padding='''max_length'''      ,		max_length=tokenizer.model_max_length      ,		return_tensors='''pt'''      ,		)
     __A					:		Optional[int]						=  text_encoder(text_inputs.input_ids						)
     __A					:		Dict						=  text_inputs.attention_mask

     __A					:		List[Any]						=  torch.Generator(device=__lowerCamelCase						).manual_seed(0						)
     __A					:		List[str]						=  pipe(
         generator=__lowerCamelCase      ,		prior_num_inference_steps=2      ,		decoder_num_inference_steps=2      ,		super_res_num_inference_steps=2      ,		prior_latents=__lowerCamelCase      ,		decoder_latents=__lowerCamelCase      ,		super_res_latents=__lowerCamelCase      ,		text_model_output=__lowerCamelCase      ,		text_attention_mask=__lowerCamelCase      ,		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						):



     '''simple docstring'''




     __A					:		Dict						=  torch_device == '''cpu'''

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

 @skip_mps
 def       UpperCamelCase__(					self						):



     '''simple docstring'''




     __A					:		Optional[Any]						=  torch_device == '''cpu'''
     __A					:		List[str]						=  True
     __A					:		Union[str, Any]						=  [
         '''prior_num_inference_steps''',
         '''decoder_num_inference_steps''',
         '''super_res_num_inference_steps''',
     ]

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

 def       UpperCamelCase__(					self						):



     '''simple docstring'''




     __A					:		Dict						=  [
         '''prior_num_inference_steps''',
         '''decoder_num_inference_steps''',
         '''super_res_num_inference_steps''',
     ]

     if torch_device == "mps":
         # TODO: MPS errors with larger batch sizes
         __A					:		List[str]						=  [2, 3]
         self._test_inference_batch_consistent(
             batch_sizes=__lowerCamelCase      ,		additional_params_copy_to_batched_inputs=__lowerCamelCase      ,		)
     else:
         self._test_inference_batch_consistent(
             additional_params_copy_to_batched_inputs=__lowerCamelCase						)

 @skip_mps
 def       UpperCamelCase__(					self						):



     '''simple docstring'''




     return super().test_dict_tuple_outputs_equivalent()

 @skip_mps
 def       UpperCamelCase__(					self						):



     '''simple docstring'''




     return super().test_save_load_local()







 @skip_mps
 def       UpperCamelCase__(					self						):



     '''simple docstring'''




     return super().test_save_load_optional_components()






@nightly
class 	__snake_case      (	unittest.TestCase  ):





 """simple docstring"""


 def       UpperCamelCase__(					self						):



     '''simple docstring'''




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







 def       UpperCamelCase__(					self						):



     '''simple docstring'''




     __A					:		Optional[Any]						=  load_numpy(
         '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
         '''/unclip/karlo_v1_alpha_horse_cpu.npy'''						)

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

     __A					:		Optional[Any]						=  torch.manual_seed(0						)
     __A					:		Optional[Any]						=  pipeline(
         '''horse'''      ,		num_images_per_prompt=1      ,		generator=__lowerCamelCase      ,		output_type='''np'''      ,		)

     __A					:		Optional[Any]						=  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       UpperCamelCase__(					self						):



     '''simple docstring'''




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

 def       UpperCamelCase__(					self						):



     '''simple docstring'''




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

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

     __A					:		str						=  torch.Generator(device='''cpu'''						).manual_seed(0						)
     __A					:		int						=  pipeline(
         '''horse'''      ,		generator=__lowerCamelCase      ,		output_type='''np'''      ,		)

     __A					:		Any						=  output.images[0]

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

     assert_mean_pixel_difference(__lowerCamelCase      ,		__lowerCamelCase						)







 def       UpperCamelCase__(					self						):



     '''simple docstring'''




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

     __A					:		Optional[Any]						=  UnCLIPPipeline.from_pretrained('''kakaobrain/karlo-v1-alpha'''      ,		torch_dtype=torch.floataa						)
     __A					:		Optional[int]						=  pipe.to(__lowerCamelCase						)
     pipe.set_progress_bar_config(disable=__lowerCamelCase						)
     pipe.enable_attention_slicing()
     pipe.enable_sequential_cpu_offload()

     __A					:		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'''      ,		)

     __A					:		Any						=  torch.cuda.max_memory_allocated()
     # make sure that less than 7 GB is allocated
     assert mem_bytes < 7 * 10**9











"""simple docstring"""



import unittest

from transformers import load_tool

from .test_tools_common import ToolTesterMixin


a_				=    """
Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf originally as a company that developed a chatbot app targeted at teenagers.[2] After open-sourcing the model behind the chatbot, the company pivoted to focus on being a platform for machine learning.

In March 2021, Hugging Face raised $40 million in a Series B funding round.[3]

On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model.[4] In 2022, the workshop concluded with the announcement of BLOOM, a multilingual large language model with 176 billion parameters.[5]
"""






class 	__snake_case      (	unittest.TestCase							,							SCREAMING_SNAKE_CASE__  ):





 """simple docstring"""


 def       UpperCamelCase__(					self						):



     '''simple docstring'''




     __A					:		Optional[Any]						=  load_tool('''summarization'''						)
     self.tool.setup()
     __A					:		Tuple						=  load_tool('''summarization'''      ,		remote=__lowerCamelCase						)

 def       UpperCamelCase__(					self						):



     '''simple docstring'''




     __A					:		List[str]						=  self.tool(__lowerCamelCase						)
     self.assertEqual(
         __lowerCamelCase      ,		'''Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf. In March 2021, Hugging Face raised $40 million in a Series B funding round. On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model. In 2022, the workshop concluded with the announcement of BLOOM.'''      ,		)

 def       UpperCamelCase__(					self						):



     '''simple docstring'''




     __A					:		int						=  self.remote_tool(__lowerCamelCase						)
     self.assertEqual(
         __lowerCamelCase      ,		'''Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf. In March 2021, Hugging Face raised $40 million in a Series B funding round. On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model. In 2022, the workshop concluded with the announcement of BLOOM.'''      ,		)

 def       UpperCamelCase__(					self						):



     '''simple docstring'''




     __A					:		List[Any]						=  self.tool(text=__lowerCamelCase						)
     self.assertEqual(
         __lowerCamelCase      ,		'''Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf. In March 2021, Hugging Face raised $40 million in a Series B funding round. On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model. In 2022, the workshop concluded with the announcement of BLOOM.'''      ,		)







 def       UpperCamelCase__(					self						):



     '''simple docstring'''




     __A					:		Optional[int]						=  self.remote_tool(text=__lowerCamelCase						)
     self.assertEqual(
         __lowerCamelCase      ,		'''Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf. In March 2021, Hugging Face raised $40 million in a Series B funding round. On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model. In 2022, the workshop concluded with the announcement of BLOOM.'''      ,		)




