




import gc
import unittest

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

from diffusers import HeunDiscreteScheduler, PriorTransformer, ShapEPipeline
from diffusers.pipelines.shap_e import ShapERenderer
from diffusers.utils import load_numpy, slow
from diffusers.utils.testing_utils import require_torch_gpu, torch_device

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference







class     __lowerCAmelCase (       __magic_name__			,				unittest.TestCase ):

			"""simple docstring"""







			snake_case_   		=       ShapEPipeline
			snake_case_   		=       ['''prompt''']
			snake_case_   		=       ['''prompt''']
			snake_case_   		=       [
			    '''num_images_per_prompt''',
			    '''num_inference_steps''',
			    '''generator''',
			    '''latents''',
			    '''guidance_scale''',
			    '''frame_size''',
			    '''output_type''',
			    '''return_dict''',
			]
			snake_case_   		=       False


			@property
			def       lowercase_    (      self							)    ->      Dict:



							'''simple docstring'''
							return 32


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



							'''simple docstring'''
							return 32


			@property
			def       lowercase_    (      self							)    ->      int:



							'''simple docstring'''
							return self.time_input_dim * 4


			@property
			def       lowercase_    (      self							)    ->      str:



							'''simple docstring'''
							return 8


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



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


			@property
			def       lowercase_    (      self							)    ->      Any:



							'''simple docstring'''
							torch.manual_seed(0							)
							__lowerCamelCase		       =    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=1_000			,						)
							return CLIPTextModelWithProjection(lowerCamelCase__							)


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



							'''simple docstring'''
							torch.manual_seed(0							)

							__lowerCamelCase		       =    {
							    'num_attention_heads': 2,
							    'attention_head_dim': 16,
							    'embedding_dim': self.time_input_dim,
							    'num_embeddings': 32,
							    'embedding_proj_dim': self.text_embedder_hidden_size,
							    'time_embed_dim': self.time_embed_dim,
							    'num_layers': 1,
							    'clip_embed_dim': self.time_input_dim * 2,
							    'additional_embeddings': 0,
							    'time_embed_act_fn': 'gelu',
							    'norm_in_type': 'layer',
							    'encoder_hid_proj_type': None,
							    'added_emb_type': None,
							}

							__lowerCamelCase		       =    PriorTransformer(**lowerCamelCase__							)
							return model


			@property
			def       lowercase_    (      self							)    ->      Dict:



							'''simple docstring'''
							torch.manual_seed(0							)

							__lowerCamelCase		       =    {
							    'param_shapes': (
							        (self.renderer_dim, 93),
							        (self.renderer_dim, 8),
							        (self.renderer_dim, 8),
							        (self.renderer_dim, 8),
							    ),
							    'd_latent': self.time_input_dim,
							    'd_hidden': self.renderer_dim,
							    'n_output': 12,
							    'background': (
							        0.1,
							        0.1,
							        0.1,
							    ),
							}
							__lowerCamelCase		       =    ShapERenderer(**lowerCamelCase__							)
							return model


			def       lowercase_    (      self							)    ->      Optional[Any]:



							'''simple docstring'''
							__lowerCamelCase		       =    self.dummy_prior
							__lowerCamelCase		       =    self.dummy_text_encoder
							__lowerCamelCase		       =    self.dummy_tokenizer
							__lowerCamelCase		       =    self.dummy_renderer

							__lowerCamelCase		       =    HeunDiscreteScheduler(
							    beta_schedule='exp'			,						num_train_timesteps=1_024			,						prediction_type='sample'			,						use_karras_sigmas=lowerCamelCase__			,						clip_sample=lowerCamelCase__			,						clip_sample_range=1.0			,						)
							__lowerCamelCase		       =    {
							    'prior': prior,
							    'text_encoder': text_encoder,
							    'tokenizer': tokenizer,
							    'renderer': renderer,
							    'scheduler': scheduler,
							}

							return components


			def       lowercase_    (      self			,						lowerCamelCase__			,						lowerCamelCase__=0							)    ->      Optional[Any]:



							'''simple docstring'''
							if str(lowerCamelCase__							).startswith('mps'							):
											__lowerCamelCase		       =    torch.manual_seed(lowerCamelCase__							)
							else:
											__lowerCamelCase		       =    torch.Generator(device=lowerCamelCase__							).manual_seed(lowerCamelCase__							)
							__lowerCamelCase		       =    {
							    'prompt': 'horse',
							    'generator': generator,
							    'num_inference_steps': 1,
							    'frame_size': 32,
							    'output_type': 'np',
							}
							return inputs


			def       lowercase_    (      self							)    ->      int:



							'''simple docstring'''
							__lowerCamelCase		       =    'cpu'

							__lowerCamelCase		       =    self.get_dummy_components()

							__lowerCamelCase		       =    self.pipeline_class(**lowerCamelCase__							)
							__lowerCamelCase		       =    pipe.to(lowerCamelCase__							)

							pipe.set_progress_bar_config(disable=lowerCamelCase__							)

							__lowerCamelCase		       =    pipe(**self.get_dummy_inputs(lowerCamelCase__							)							)
							__lowerCamelCase		       =    output.images[0]
							__lowerCamelCase		       =    image[0, -3:, -3:, -1]

							assert image.shape == (20, 32, 32, 3)

							__lowerCamelCase		       =    np.array(
							    [
							        0.00_03_92_16,
							        0.00_03_92_16,
							        0.00_03_92_16,
							        0.00_03_92_16,
							        0.00_03_92_16,
							        0.00_03_92_16,
							        0.00_03_92_16,
							        0.00_03_92_16,
							        0.00_03_92_16,
							    ]							)

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


			def       lowercase_    (      self							)    ->      Optional[Any]:



							'''simple docstring'''
							# NOTE: Larger batch sizes cause this test to timeout, only test on smaller batches
							self._test_inference_batch_consistent(batch_sizes=[1, 2]							)


			def       lowercase_    (      self							)    ->      Tuple:



							'''simple docstring'''
							__lowerCamelCase		       =    torch_device == 'cpu'
							__lowerCamelCase		       =    True

							self._test_inference_batch_single_identical(
							    batch_size=2			,						test_max_difference=lowerCamelCase__			,						relax_max_difference=lowerCamelCase__			,						)






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



							'''simple docstring'''
							__lowerCamelCase		       =    self.get_dummy_components()
							__lowerCamelCase		       =    self.pipeline_class(**lowerCamelCase__							)
							__lowerCamelCase		       =    pipe.to(lowerCamelCase__							)
							pipe.set_progress_bar_config(disable=lowerCamelCase__							)

							__lowerCamelCase		       =    1
							__lowerCamelCase		       =    2

							__lowerCamelCase		       =    self.get_dummy_inputs(lowerCamelCase__							)

							for key in inputs.keys():
											if key in self.batch_params:
															__lowerCamelCase		       =    batch_size * [inputs[key]]

							__lowerCamelCase		       =    pipe(**lowerCamelCase__			,						num_images_per_prompt=lowerCamelCase__							)[0]

							assert images.shape[0] == batch_size * num_images_per_prompt







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

			"""simple docstring"""







			def       lowercase_    (      self							)    ->      int:



							'''simple docstring'''
							# clean up the VRAM after each test
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()






			def       lowercase_    (      self							)    ->      List[Any]:



							'''simple docstring'''
							__lowerCamelCase		       =    load_numpy(
							    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
							    '/shap_e/test_shap_e_np_out.npy'							)
							__lowerCamelCase		       =    ShapEPipeline.from_pretrained('openai/shap-e'							)
							__lowerCamelCase		       =    pipe.to(lowerCamelCase__							)
							pipe.set_progress_bar_config(disable=lowerCamelCase__							)

							__lowerCamelCase		       =    torch.Generator(device=lowerCamelCase__							).manual_seed(0							)

							__lowerCamelCase		       =    pipe(
							    'a shark'			,						generator=lowerCamelCase__			,						guidance_scale=15.0			,						num_inference_steps=64			,						frame_size=64			,						output_type='np'			,						).images[0]

							assert images.shape == (20, 64, 64, 3)

							assert_mean_pixel_difference(lowerCamelCase__			,						lowerCamelCase__							)








import string

# frequency taken from https://en.wikipedia.org/wiki/Letter_frequency
__A        =    {
    "E": 1_2.7_0,
    "T": 9.0_6,
    "A": 8.1_7,
    "O": 7.5_1,
    "I": 6.9_7,
    "N": 6.7_5,
    "S": 6.3_3,
    "H": 6.0_9,
    "R": 5.9_9,
    "D": 4.2_5,
    "L": 4.0_3,
    "C": 2.7_8,
    "U": 2.7_6,
    "M": 2.4_1,
    "W": 2.3_6,
    "F": 2.2_3,
    "G": 2.0_2,
    "Y": 1.9_7,
    "P": 1.9_3,
    "B": 1.2_9,
    "V": 0.9_8,
    "K": 0.7_7,
    "J": 0.1_5,
    "X": 0.1_5,
    "Q": 0.1_0,
    "Z": 0.0_7,
}
__A        =    "ETAOINSHRDLCUMWFGYPBVKJXQZ"
__A        =    "ABCDEFGHIJKLMNOPQRSTUVWXYZ"







def 			lowerCamelCase_		(       UpperCamelCase__				:      str      )     ->       dict[str, int]:





				"""simple docstring"""


				__lowerCamelCase		       =    {letter: 0 for letter in string.ascii_uppercase}
				for letter in message.upper():
								if letter in LETTERS:
												letter_count[letter] += 1

				return letter_count







def 			lowerCamelCase_		(       UpperCamelCase__				:      tuple      )     ->       str:





				"""simple docstring"""


				return x[0]







def 			lowerCamelCase_		(       UpperCamelCase__				:      str      )     ->       str:





				"""simple docstring"""


				__lowerCamelCase		       =    get_letter_count(UpperCamelCase__      )
				__lowerCamelCase		       =    {
				    freq: [] for letter, freq in letter_to_freq.items()
				}
				for letter in LETTERS:
								freq_to_letter[letter_to_freq[letter]].append(UpperCamelCase__      )

				__lowerCamelCase		       =    {}

				for freq in freq_to_letter:
								freq_to_letter[freq].sort(key=ETAOIN.find ,     reverse=UpperCamelCase__      )
								__lowerCamelCase		       =    ''.join(freq_to_letter[freq]      )

				__lowerCamelCase		       =    list(freq_to_letter_str.items()      )
				freq_pairs.sort(key=UpperCamelCase__ ,     reverse=UpperCamelCase__      )

				__lowerCamelCase		       =    [freq_pair[1] for freq_pair in freq_pairs]

				return "".join(UpperCamelCase__      )







def 			lowerCamelCase_		(       UpperCamelCase__				:      str      )     ->       int:





				"""simple docstring"""


				__lowerCamelCase		       =    get_frequency_order(UpperCamelCase__      )
				__lowerCamelCase		       =    0
				for common_letter in ETAOIN[:6]:
								if common_letter in freq_order[:6]:
												match_score += 1

				for uncommon_letter in ETAOIN[-6:]:
								if uncommon_letter in freq_order[-6:]:
												match_score += 1

				return match_score


if __name__ == "__main__":
					import doctest

					doctest.testmod()


