


"""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 	UpperCamelCase    (					lowercase   ,			unittest.TestCase		):
				UpperCAmelCase       :   Any					=	ShapEPipeline
				UpperCAmelCase       :   Union[str, Any]					=	["""prompt"""]
				UpperCAmelCase       :   str					=	["""prompt"""]
				UpperCAmelCase       :   List[str]					=	[
				    """num_images_per_prompt""",
				    """num_inference_steps""",
				    """generator""",
				    """latents""",
				    """guidance_scale""",
				    """frame_size""",
				    """output_type""",
				    """return_dict""",
				]
				UpperCAmelCase       :   Dict					=	False



				@property
				def       _lowercase  (self       :   Tuple)					->			str:
					return 32



				@property
				def       _lowercase  (self       :   Tuple)					->			Optional[int]:
					return 32



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



				@property
				def       _lowercase  (self       :   Any)					->			Any:
					return 8



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



				@property
				def       _lowercase  (self       :   Union[str, Any])					->			List[Any]:
					torch.manual_seed(0)
					__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=37			,						layer_norm_eps=1E-05			,						num_attention_heads=4			,						num_hidden_layers=5			,						pad_token_id=1			,						vocab_size=10_00			,						)
					return CLIPTextModelWithProjection(_A)



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

					__snake_case    :							str   	=				{
					    '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    :							str   	=				PriorTransformer(**_A)
					return model



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

					__snake_case    :							Any   	=				{
					    '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    :							int   	=				ShapERenderer(**_A)
					return model



				def       _lowercase  (self       :   List[str])					->			Optional[Any]:
					__snake_case    :							Union[str, Any]   	=				self.dummy_prior
					__snake_case    :							Tuple   	=				self.dummy_text_encoder
					__snake_case    :							Optional[int]   	=				self.dummy_tokenizer
					__snake_case    :							Optional[int]   	=				self.dummy_renderer

					__snake_case    :							Optional[int]   	=				HeunDiscreteScheduler(
					    beta_schedule='exp'			,						num_train_timesteps=10_24			,						prediction_type='sample'			,						use_karras_sigmas=_A			,						clip_sample=_A			,						clip_sample_range=1.0			,						)
					__snake_case    :							Optional[int]   	=				{
					    'prior': prior,
					    'text_encoder': text_encoder,
					    'tokenizer': tokenizer,
					    'renderer': renderer,
					    'scheduler': scheduler,
					}

					return components



				def       _lowercase  (self       :   Union[str, Any]			,						_A       :   List[str]			,						_A       :   Dict=0)					->			Union[str, Any]:
					if str(_A).startswith('mps'):
						__snake_case    :							List[Any]   	=				torch.manual_seed(_A)
					else:
						__snake_case    :							Any   	=				torch.Generator(device=_A).manual_seed(_A)
					__snake_case    :							List[str]   	=				{
					    'prompt': 'horse',
					    'generator': generator,
					    'num_inference_steps': 1,
					    'frame_size': 32,
					    'output_type': 'np',
					}
					return inputs



				def       _lowercase  (self       :   Union[str, Any])					->			str:
					__snake_case    :							List[Any]   	=				'cpu'

					__snake_case    :							Optional[int]   	=				self.get_dummy_components()

					__snake_case    :							List[Any]   	=				self.pipeline_class(**_A)
					__snake_case    :							Tuple   	=				pipe.to(_A)

					pipe.set_progress_bar_config(disable=_A)

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

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

					__snake_case    :							Dict   	=				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       _lowercase  (self       :   Tuple)					->			int:
					# 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       :   int)					->			Tuple:
					__snake_case    :							List[str]   	=				torch_device == 'cpu'
					__snake_case    :							Optional[int]   	=				True

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






				def       _lowercase  (self       :   Optional[int])					->			str:
					__snake_case    :							Dict   	=				self.get_dummy_components()
					__snake_case    :							List[Any]   	=				self.pipeline_class(**_A)
					__snake_case    :							Optional[int]   	=				pipe.to(_A)
					pipe.set_progress_bar_config(disable=_A)

					__snake_case    :							Union[str, Any]   	=				1
					__snake_case    :							List[Any]   	=				2

					__snake_case    :							Dict   	=				self.get_dummy_inputs(_A)

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

					__snake_case    :							str   	=				pipe(**_A			,						num_images_per_prompt=_A)[0]

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







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



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






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

					__snake_case    :							Optional[int]   	=				torch.Generator(device=_A).manual_seed(0)

					__snake_case    :							List[Any]   	=				pipe(
					    'a shark'			,						generator=_A			,						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(_A			,						_A)






"""simple docstring"""







import math

import tensorflow as tf
from packaging import version
def 						__UpperCAmelCase       (					UpperCAmelCase_					:       Union[str, Any]						)							->    Any:

	'''simple docstring'''





	__snake_case    :							List[str]   	=				tf.convert_to_tensor(UpperCAmelCase_						)
	__snake_case    :							List[str]   	=				0.5 * (1.0 + tf.math.erf(x / tf.cast(tf.sqrt(2.0						)    ,							x.dtype						)						))

	return x * cdf
def 						__UpperCAmelCase       (					UpperCAmelCase_					:       Tuple						)							->    Optional[Any]:

	'''simple docstring'''





	__snake_case    :							Optional[int]   	=				tf.convert_to_tensor(UpperCAmelCase_						)
	__snake_case    :							Any   	=				tf.cast(math.pi    ,							x.dtype						)
	__snake_case    :							List[Any]   	=				tf.cast(0.044_715    ,							x.dtype						)
	__snake_case    :							Union[str, Any]   	=				0.5 * (1.0 + tf.tanh(tf.sqrt(2.0 / pi						) * (x + coeff * tf.pow(UpperCAmelCase_    ,							3						))						))

	return x * cdf
def 						__UpperCAmelCase       (					UpperCAmelCase_					:       Tuple						)							->    List[Any]:

	'''simple docstring'''





	__snake_case    :							Optional[int]   	=				tf.convert_to_tensor(UpperCAmelCase_						)

	return x * tf.tanh(tf.math.softplus(UpperCAmelCase_						)						)
def 						__UpperCAmelCase       (					UpperCAmelCase_					:       Tuple						)							->    List[str]:

	'''simple docstring'''





	__snake_case    :							Union[str, Any]   	=				tf.convert_to_tensor(UpperCAmelCase_						)
	__snake_case    :							Dict   	=				tf.cast(0.044_715    ,							x.dtype						)
	__snake_case    :							int   	=				tf.cast(0.7_978_845_608    ,							x.dtype						)

	return 0.5 * x * (1.0 + tf.tanh(x * coeffa * (1.0 + coeffa * x * x)						))
def 						__UpperCAmelCase       (					UpperCAmelCase_					:       int						)							->    Dict:

	'''simple docstring'''





	__snake_case    :							Any   	=				tf.convert_to_tensor(UpperCAmelCase_						)
	__snake_case    :							Dict   	=				tf.cast(1.702    ,							x.dtype						)
	return x * tf.math.sigmoid(coeff * x						)
def 						__UpperCAmelCase       (					UpperCAmelCase_					:       Union[str, Any]						)							->    List[Any]:

	'''simple docstring'''





	return tf.clip_by_value(_gelu(UpperCAmelCase_						)    ,							-10    ,							10						)
def 						__UpperCAmelCase       (					UpperCAmelCase_					:       Union[str, Any]    ,							UpperCAmelCase_					:       Any=-1						)							->    int:

	'''simple docstring'''





	__snake_case	,       __snake_case    :							str   	=				tf.split(UpperCAmelCase_    ,							2    ,							axis=UpperCAmelCase_						)
	return a * tf.math.sigmoid(UpperCAmelCase_						)


if version.parse(tf.version.VERSION) >= version.parse("2.4"):
	def 						__UpperCAmelCase       (					UpperCAmelCase_					:       Tuple						)							->    int:

		'''simple docstring'''





		return tf.keras.activations.gelu(UpperCAmelCase_    ,							approximate=UpperCAmelCase_						)

	_a  :  Tuple=				tf.keras.activations.gelu

	_a  :  Tuple=				approximate_gelu_wrap
else:
	_a  :  Any=				_gelu
	_a  :  str=				_gelu_new


_a  :  List[str]=				{
    "gelu": gelu,
    "gelu_10": gelu_aa,
    "gelu_fast": gelu_fast,
    "gelu_new": gelu_new,
    "glu": glu,
    "mish": mish,
    "quick_gelu": quick_gelu,
    "relu": tf.keras.activations.relu,
    "sigmoid": tf.keras.activations.sigmoid,
    "silu": tf.keras.activations.swish,
    "swish": tf.keras.activations.swish,
    "tanh": tf.keras.activations.tanh,
}

def 						__UpperCAmelCase       (					UpperCAmelCase_					:       Optional[Any]						)							->    List[Any]:

	'''simple docstring'''





	if activation_string in ACTaFN:
		return ACTaFN[activation_string]
	else:
		raise KeyError(F"function {activation_string} not found in ACT2FN mapping {list(ACTaFN.keys()						)}"						)





