






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 					_snake_case						(			UpperCAmelCase_       ,    unittest.TestCase ):
	_A          =   ShapEPipeline
	_A          =   ['prompt']
	_A          =   ['prompt']
	_A          =   [
	    'num_images_per_prompt',
	    'num_inference_steps',
	    'generator',
	    'latents',
	    'guidance_scale',
	    'frame_size',
	    'output_type',
	    'return_dict',
	]
	_A          =   False






	@property
	def    lowerCAmelCase_							(		self       )						->    Dict:
					return 32






	@property
	def    lowerCAmelCase_							(		self       )						->    int:
					return 32






	@property
	def    lowerCAmelCase_							(		self       )						->    List[str]:
					return self.time_input_dim * 4






	@property
	def    lowerCAmelCase_							(		self       )						->    str:
					return 8






	@property
	def    lowerCAmelCase_							(		self       )						->    Tuple:
					snake_case__      :str         =   CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"       )
					return tokenizer






	@property
	def    lowerCAmelCase_							(		self       )						->    Optional[Any]:
					torch.manual_seed(0       )
					snake_case__      :Optional[int]         =   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(_lowercase       )






	@property
	def    lowerCAmelCase_							(		self       )						->    Any:
					torch.manual_seed(0       )

					snake_case__      :Union[str, Any]         =   {
					    '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,
					}

					snake_case__      :Union[str, Any]         =   PriorTransformer(**_lowercase       )
					return model






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

					snake_case__      :List[str]         =   {
					    '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,
					    ),
					}
					snake_case__      :Dict         =   ShapERenderer(**_lowercase       )
					return model






	def    lowerCAmelCase_							(		self       )						->    Tuple:
					snake_case__      :Optional[int]         =   self.dummy_prior
					snake_case__      :Optional[Any]         =   self.dummy_text_encoder
					snake_case__      :Optional[int]         =   self.dummy_tokenizer
					snake_case__      :List[Any]         =   self.dummy_renderer

					snake_case__      :Any         =   HeunDiscreteScheduler(
					    beta_schedule="exp"					,num_train_timesteps=1_024					,prediction_type="sample"					,use_karras_sigmas=_lowercase					,clip_sample=_lowercase					,clip_sample_range=1.0					,)
					snake_case__      :str         =   {
					    'prior': prior,
					    'text_encoder': text_encoder,
					    'tokenizer': tokenizer,
					    'renderer': renderer,
					    'scheduler': scheduler,
					}

					return components






	def    lowerCAmelCase_							(		self					,UpperCamelCase					,UpperCamelCase=0       )						->    Tuple:
					if str(_lowercase       ).startswith("mps"       ):
									snake_case__      :Union[str, Any]         =   torch.manual_seed(_lowercase       )
					else:
									snake_case__      :int         =   torch.Generator(device=_lowercase       ).manual_seed(_lowercase       )
					snake_case__      :Union[str, Any]         =   {
					    'prompt': 'horse',
					    'generator': generator,
					    'num_inference_steps': 1,
					    'frame_size': 32,
					    'output_type': 'np',
					}
					return inputs






	def    lowerCAmelCase_							(		self       )						->    Dict:
					snake_case__      :Optional[Any]         =   'cpu'

					snake_case__      :List[Any]         =   self.get_dummy_components()

					snake_case__      :Optional[Any]         =   self.pipeline_class(**_lowercase       )
					snake_case__      :List[str]         =   pipe.to(_lowercase       )

					pipe.set_progress_bar_config(disable=_lowercase       )

					snake_case__      :Dict         =   pipe(**self.get_dummy_inputs(_lowercase       )       )
					snake_case__      :Dict         =   output.images[0]
					snake_case__      :Optional[int]         =   image[0, -3:, -3:, -1]

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

					snake_case__      :int         =   np.array(
					    [
					        0.00039216,
					        0.00039216,
					        0.00039216,
					        0.00039216,
					        0.00039216,
					        0.00039216,
					        0.00039216,
					        0.00039216,
					        0.00039216,
					    ]       )

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






	def    lowerCAmelCase_							(		self       )						->    List[str]:
					# NOTE: Larger batch sizes cause this test to timeout, only test on smaller batches
					self._test_inference_batch_consistent(batch_sizes=[1, 2]       )






	def    lowerCAmelCase_							(		self       )						->    Union[str, Any]:
					snake_case__      :str         =   torch_device == 'cpu'
					snake_case__      :str         =   True

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






	def    lowerCAmelCase_							(		self       )						->    Optional[Any]:
					snake_case__      :int         =   self.get_dummy_components()
					snake_case__      :Any         =   self.pipeline_class(**_lowercase       )
					snake_case__      :Optional[int]         =   pipe.to(_lowercase       )
					pipe.set_progress_bar_config(disable=_lowercase       )

					snake_case__      :List[Any]         =   1
					snake_case__      :int         =   2

					snake_case__      :Union[str, Any]         =   self.get_dummy_inputs(_lowercase       )

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

					snake_case__      :Tuple         =   pipe(**_lowercase					,num_images_per_prompt=_lowercase       )[0]

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




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






	def    lowerCAmelCase_							(		self       )						->    Tuple:
					# clean up the VRAM after each test
					super().tearDown()
					gc.collect()
					torch.cuda.empty_cache()






	def    lowerCAmelCase_							(		self       )						->    Union[str, Any]:
					snake_case__      :str         =   load_numpy(
					    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
					    "/shap_e/test_shap_e_np_out.npy"       )
					snake_case__      :int         =   ShapEPipeline.from_pretrained("openai/shap-e"       )
					snake_case__      :str         =   pipe.to(_lowercase       )
					pipe.set_progress_bar_config(disable=_lowercase       )

					snake_case__      :Optional[Any]         =   torch.Generator(device=_lowercase       ).manual_seed(0       )

					snake_case__      :Union[str, Any]         =   pipe(
					    "a shark"					,generator=_lowercase					,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(_lowercase					,_lowercase       )












def 					lowercase_ (    __snake_case     :							str      )     -> list:





				'''simple docstring'''


				return [
				    txt[:a] + txt[a].upper() + txt[a + 1 :]
				    for a in range(len(__snake_case      )      )
				    if txt[a].isalpha()
				]


if __name__ == "__main__":
	__import__("doctest").testmod()