





'''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 							__lowercase       (      lowerCAmelCase__ ,      unittest.TestCase						):







	'''simple docstring'''

	a    :				List[str]											=       ShapEPipeline
	a    :				str											=       ["prompt"]
	a    :				Tuple											=       ["prompt"]
	a    :				Optional[int]											=       [
	    "num_images_per_prompt",
	    "num_inference_steps",
	    "generator",
	    "latents",
	    "guidance_scale",
	    "frame_size",
	    "output_type",
	    "return_dict",
	]
	a    :				int											=       False


	@property
	def 						_UpperCAmelCase     (self )							->							Dict:







				'''simple docstring'''





				return 32


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







				'''simple docstring'''





				return 32


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







				'''simple docstring'''





				return self.time_input_dim * 4


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







				'''simple docstring'''





				return 8


	@property
	def 						_UpperCAmelCase     (self )							->							Dict:







				'''simple docstring'''





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


	@property
	def 						_UpperCAmelCase     (self )							->							Optional[int]:







				'''simple docstring'''





				torch.manual_seed(0 )
				__lowercase										=							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 )							->							Optional[int]:







				'''simple docstring'''





				torch.manual_seed(0 )

				__lowercase										=							{
				    '''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,
				}

				__lowercase										=							PriorTransformer(**_lowerCamelCase )
				return model


	@property
	def 						_UpperCAmelCase     (self )							->							Optional[int]:







				'''simple docstring'''





				torch.manual_seed(0 )

				__lowercase										=							{
				    '''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,
				    ),
				}
				__lowercase										=							ShapERenderer(**_lowerCamelCase )
				return model


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







				'''simple docstring'''





				__lowercase										=							self.dummy_prior
				__lowercase										=							self.dummy_text_encoder
				__lowercase										=							self.dummy_tokenizer
				__lowercase										=							self.dummy_renderer

				__lowercase										=							HeunDiscreteScheduler(
				    beta_schedule='''exp'''     ,num_train_timesteps=1024     ,prediction_type='''sample'''     ,use_karras_sigmas=_lowerCamelCase     ,clip_sample=_lowerCamelCase     ,clip_sample_range=1.0     ,)
				__lowercase										=							{
				    '''prior''': prior,
				    '''text_encoder''': text_encoder,
				    '''tokenizer''': tokenizer,
				    '''renderer''': renderer,
				    '''scheduler''': scheduler,
				}

				return components


	def 						_UpperCAmelCase     (self     ,_lowerCamelCase     ,_lowerCamelCase=0 )							->							Union[str, Any]:







				'''simple docstring'''





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


	def 						_UpperCAmelCase     (self )							->							Dict:







				'''simple docstring'''





				__lowercase										=							'''cpu'''

				__lowercase										=							self.get_dummy_components()

				__lowercase										=							self.pipeline_class(**_lowerCamelCase )
				__lowercase										=							pipe.to(_lowerCamelCase )

				pipe.set_progress_bar_config(disable=_lowerCamelCase )

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

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

				__lowercase										=							np.array(
				    [
				        0.0_0_0_3_9_2_1_6,
				        0.0_0_0_3_9_2_1_6,
				        0.0_0_0_3_9_2_1_6,
				        0.0_0_0_3_9_2_1_6,
				        0.0_0_0_3_9_2_1_6,
				        0.0_0_0_3_9_2_1_6,
				        0.0_0_0_3_9_2_1_6,
				        0.0_0_0_3_9_2_1_6,
				        0.0_0_0_3_9_2_1_6,
				    ] )

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


	def 						_UpperCAmelCase     (self )							->							str:







				'''simple docstring'''





				self._test_inference_batch_consistent(batch_sizes=[1, 2] )


	def 						_UpperCAmelCase     (self )							->							Optional[int]:







				'''simple docstring'''





				__lowercase										=							torch_device == '''cpu'''
				__lowercase										=							True

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


	def 						_UpperCAmelCase     (self )							->							int:







				'''simple docstring'''





				__lowercase										=							self.get_dummy_components()
				__lowercase										=							self.pipeline_class(**_lowerCamelCase )
				__lowercase										=							pipe.to(_lowerCamelCase )
				pipe.set_progress_bar_config(disable=_lowerCamelCase )

				__lowercase										=							1
				__lowercase										=							2

				__lowercase										=							self.get_dummy_inputs(_lowerCamelCase )

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

				__lowercase										=							pipe(**_lowerCamelCase     ,num_images_per_prompt=_lowerCamelCase )[0]

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






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







	'''simple docstring'''

	def 						_UpperCAmelCase     (self )							->							Any:







				'''simple docstring'''





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


	def 						_UpperCAmelCase     (self )							->							str:







				'''simple docstring'''





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

				__lowercase										=							torch.Generator(device=_lowerCamelCase ).manual_seed(0 )

				__lowercase										=							pipe(
				    '''a shark'''     ,generator=_lowerCamelCase     ,guidance_scale=1_5.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 )












'''simple docstring'''



from typing import TYPE_CHECKING

from ....utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available


_SCREAMING_SNAKE_CASE					=    {'''configuration_van''': ['''VAN_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''VanConfig''']}


try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	_SCREAMING_SNAKE_CASE					=    [
	    '''VAN_PRETRAINED_MODEL_ARCHIVE_LIST''',
	    '''VanForImageClassification''',
	    '''VanModel''',
	    '''VanPreTrainedModel''',
	]

if TYPE_CHECKING:
	from .configuration_van import VAN_PRETRAINED_CONFIG_ARCHIVE_MAP, VanConfig

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_van import (
		    VAN_PRETRAINED_MODEL_ARCHIVE_LIST,
		    VanForImageClassification,
		    VanModel,
		    VanPreTrainedModel,
		)

else:
	import sys

	_SCREAMING_SNAKE_CASE					=    _LazyModule(__name__, globals()['''__file__'''], _import_structure)





