






"""simple docstring"""
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						(    __UpperCAmelCase   ,  unittest.TestCase       ):
						SCREAMING_SNAKE_CASE_:					Dict 					=				ShapEPipeline
						SCREAMING_SNAKE_CASE_:					List[str] 					=				['prompt']
						SCREAMING_SNAKE_CASE_:					Dict 					=				['prompt']
						SCREAMING_SNAKE_CASE_:					Dict 					=				[
						    'num_images_per_prompt',
						    'num_inference_steps',
						    'generator',
						    'latents',
						    'guidance_scale',
						    'frame_size',
						    'output_type',
						    'return_dict',
						]
						SCREAMING_SNAKE_CASE_:					Dict 					=				False







						@property
						def 			A      (    self							)		->							List[str]:
										return 3_2







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







						@property
						def 			A      (    self							)		->							Any:
										return self.time_input_dim * 4







						@property
						def 			A      (    self							)		->							int:
										return 8







						@property
						def 			A      (    self							)		->							Any:
										_SCREAMING_SNAKE_CASE   :		Union[str, Any]       		=    CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'							)
										return tokenizer







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







						@property
						def 			A      (    self							)		->							int:
										torch.manual_seed(0							)

										_SCREAMING_SNAKE_CASE   :		Dict       		=    {
										    'num_attention_heads': 2,
										    'attention_head_dim': 1_6,
										    'embedding_dim': self.time_input_dim,
										    'num_embeddings': 3_2,
										    '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,
										}

										_SCREAMING_SNAKE_CASE   :		Union[str, Any]       		=    PriorTransformer(**lowerCAmelCase_							)
										return model







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

										_SCREAMING_SNAKE_CASE   :		List[str]       		=    {
										    'param_shapes': (
										        (self.renderer_dim, 9_3),
										        (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': 1_2,
										    'background': (
										        0.1,
										        0.1,
										        0.1,
										    ),
										}
										_SCREAMING_SNAKE_CASE   :		List[Any]       		=    ShapERenderer(**lowerCAmelCase_							)
										return model







						def 			A      (    self							)		->							str:
										_SCREAMING_SNAKE_CASE   :		List[Any]       		=    self.dummy_prior
										_SCREAMING_SNAKE_CASE   :		Optional[Any]       		=    self.dummy_text_encoder
										_SCREAMING_SNAKE_CASE   :		Dict       		=    self.dummy_tokenizer
										_SCREAMING_SNAKE_CASE   :		List[Any]       		=    self.dummy_renderer

										_SCREAMING_SNAKE_CASE   :		List[str]       		=    HeunDiscreteScheduler(
										    beta_schedule='exp'   ,				num_train_timesteps=1_0_2_4   ,				prediction_type='sample'   ,				use_karras_sigmas=lowerCAmelCase_   ,				clip_sample=lowerCAmelCase_   ,				clip_sample_range=1.0   ,				)
										_SCREAMING_SNAKE_CASE   :		List[str]       		=    {
										    'prior': prior,
										    'text_encoder': text_encoder,
										    'tokenizer': tokenizer,
										    'renderer': renderer,
										    'scheduler': scheduler,
										}

										return components







						def 			A      (    self   ,				lowerCAmelCase_   ,				lowerCAmelCase_=0							)		->							List[Any]:
										if str(lowerCAmelCase_							).startswith('mps'							):
														_SCREAMING_SNAKE_CASE   :		Any       		=    torch.manual_seed(lowerCAmelCase_							)
										else:
														_SCREAMING_SNAKE_CASE   :		Union[str, Any]       		=    torch.Generator(device=lowerCAmelCase_							).manual_seed(lowerCAmelCase_							)
										_SCREAMING_SNAKE_CASE   :		str       		=    {
										    'prompt': 'horse',
										    'generator': generator,
										    'num_inference_steps': 1,
										    'frame_size': 3_2,
										    'output_type': 'np',
										}
										return inputs







						def 			A      (    self							)		->							List[str]:
										_SCREAMING_SNAKE_CASE   :		Dict       		=    'cpu'

										_SCREAMING_SNAKE_CASE   :		Tuple       		=    self.get_dummy_components()

										_SCREAMING_SNAKE_CASE   :		Tuple       		=    self.pipeline_class(**lowerCAmelCase_							)
										_SCREAMING_SNAKE_CASE   :		List[str]       		=    pipe.to(lowerCAmelCase_							)

										pipe.set_progress_bar_config(disable=lowerCAmelCase_							)

										_SCREAMING_SNAKE_CASE   :		Any       		=    pipe(**self.get_dummy_inputs(lowerCAmelCase_							)							)
										_SCREAMING_SNAKE_CASE   :		List[Any]       		=    output.images[0]
										_SCREAMING_SNAKE_CASE   :		str       		=    image[0, -3:, -3:, -1]

										assert image.shape == (2_0, 3_2, 3_2, 3)

										_SCREAMING_SNAKE_CASE   :		int       		=    np.array(
										    [
										        0.00_039_216,
										        0.00_039_216,
										        0.00_039_216,
										        0.00_039_216,
										        0.00_039_216,
										        0.00_039_216,
										        0.00_039_216,
										        0.00_039_216,
										        0.00_039_216,
										    ]							)

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







						def 			A      (    self							)		->							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 			A      (    self							)		->							Optional[int]:
										_SCREAMING_SNAKE_CASE   :		Dict       		=    torch_device == 'cpu'
										_SCREAMING_SNAKE_CASE   :		str       		=    True

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







						def 			A      (    self							)		->							Dict:
										_SCREAMING_SNAKE_CASE   :		Optional[Any]       		=    self.get_dummy_components()
										_SCREAMING_SNAKE_CASE   :		List[Any]       		=    self.pipeline_class(**lowerCAmelCase_							)
										_SCREAMING_SNAKE_CASE   :		int       		=    pipe.to(lowerCAmelCase_							)
										pipe.set_progress_bar_config(disable=lowerCAmelCase_							)

										_SCREAMING_SNAKE_CASE   :		Union[str, Any]       		=    1
										_SCREAMING_SNAKE_CASE   :		Optional[int]       		=    2

										_SCREAMING_SNAKE_CASE   :		List[str]       		=    self.get_dummy_inputs(lowerCAmelCase_							)

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

										_SCREAMING_SNAKE_CASE   :		Tuple       		=    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       ):







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







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

										_SCREAMING_SNAKE_CASE   :		int       		=    torch.Generator(device=lowerCAmelCase_							).manual_seed(0							)

										_SCREAMING_SNAKE_CASE   :		int       		=    pipe(
										    'a shark'   ,				generator=lowerCAmelCase_   ,				guidance_scale=15.0   ,				num_inference_steps=6_4   ,				frame_size=6_4   ,				output_type='np'   ,				).images[0]

										assert images.shape == (2_0, 6_4, 6_4, 3)

										assert_mean_pixel_difference(lowerCAmelCase_   ,				lowerCAmelCase_							)









"""simple docstring"""
def 				lowercase__		(					lowerCamelCase,						lowerCamelCase       ):
				return numa ^ numa < 0


if __name__ == "__main__":
			import doctest

			doctest.testmod()

