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

     """simple docstring"""

     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			):

       '''simple docstring'''

       return 32


     @property
     def 				_lowerCAmelCase	(   self			):

       '''simple docstring'''

       return 32


     @property
     def 				_lowerCAmelCase	(   self			):

       '''simple docstring'''

       return self.time_input_dim * 4


     @property
     def 				_lowerCAmelCase	(   self			):

       '''simple docstring'''

       return 8


     @property
     def 				_lowerCAmelCase	(   self			):

       '''simple docstring'''

       a_   :						List[Any]   =				CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""			)
       return tokenizer


     @property
     def 				_lowerCAmelCase	(   self			):

       '''simple docstring'''

       torch.manual_seed(0			)
       a_   :						str   =				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=10_00				,		)
       return CLIPTextModelWithProjection(lowerCAmelCase_			)


     @property
     def 				_lowerCAmelCase	(   self			):

       '''simple docstring'''

       torch.manual_seed(0			)

       a_   :						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,
       }

       a_   :						str   =				PriorTransformer(**lowerCAmelCase_			)
       return model


     @property
     def 				_lowerCAmelCase	(   self			):

       '''simple docstring'''

       torch.manual_seed(0			)

       a_   :						int   =				{
           """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,
           ),
       }
       a_   :						Any   =				ShapERenderer(**lowerCAmelCase_			)
       return model


     def 				_lowerCAmelCase	(   self			):

       '''simple docstring'''

       a_   :						int   =				self.dummy_prior
       a_   :						int   =				self.dummy_text_encoder
       a_   :						Optional[Any]   =				self.dummy_tokenizer
       a_   :						Optional[int]   =				self.dummy_renderer

       a_   :						Union[str, Any]   =				HeunDiscreteScheduler(
           beta_schedule="""exp"""				,		num_train_timesteps=10_24				,		prediction_type="""sample"""				,		use_karras_sigmas=lowerCAmelCase_				,		clip_sample=lowerCAmelCase_				,		clip_sample_range=1.0				,		)
       a_   :						int   =				{
           """prior""": prior,
           """text_encoder""": text_encoder,
           """tokenizer""": tokenizer,
           """renderer""": renderer,
           """scheduler""": scheduler,
       }

       return components


     def 				_lowerCAmelCase	(   self				,		lowerCAmelCase_				,		lowerCAmelCase_=0			):

       '''simple docstring'''

       if str(lowerCAmelCase_			).startswith("""mps"""			):
         a_   :						Any   =				torch.manual_seed(lowerCAmelCase_			)
       else:
         a_   :						Optional[int]   =				torch.Generator(device=lowerCAmelCase_			).manual_seed(lowerCAmelCase_			)
       a_   :						Dict   =				{
           """prompt""": """horse""",
           """generator""": generator,
           """num_inference_steps""": 1,
           """frame_size""": 32,
           """output_type""": """np""",
       }
       return inputs


     def 				_lowerCAmelCase	(   self			):

       '''simple docstring'''

       a_   :						List[str]   =				"""cpu"""

       a_   :						Dict   =				self.get_dummy_components()

       a_   :						Any   =				self.pipeline_class(**lowerCAmelCase_			)
       a_   :						int   =				pipe.to(lowerCAmelCase_			)

       pipe.set_progress_bar_config(disable=lowerCAmelCase_			)

       a_   :						Optional[Any]   =				pipe(**self.get_dummy_inputs(lowerCAmelCase_			)			)
       a_   :						int   =				output.images[0]
       a_   :						Optional[int]   =				image[0, -3:, -3:, -1]

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

       a_   :						Optional[Any]   =				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			):

       '''simple docstring'''

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


     def 				_lowerCAmelCase	(   self			):

       '''simple docstring'''

       a_   :						Tuple   =				torch_device == """cpu"""
       a_   :						Tuple   =				True

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






     def 				_lowerCAmelCase	(   self			):

       '''simple docstring'''

       a_   :						List[str]   =				self.get_dummy_components()
       a_   :						int   =				self.pipeline_class(**lowerCAmelCase_			)
       a_   :						int   =				pipe.to(lowerCAmelCase_			)
       pipe.set_progress_bar_config(disable=lowerCAmelCase_			)

       a_   :						Union[str, Any]   =				1
       a_   :						Tuple   =				2

       a_   :						Optional[int]   =				self.get_dummy_inputs(lowerCAmelCase_			)

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

       a_   :						Tuple   =				pipe(**lowerCAmelCase_				,		num_images_per_prompt=lowerCAmelCase_			)[0]

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







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

     """simple docstring"""

     def 				_lowerCAmelCase	(   self			):

       '''simple docstring'''

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






     def 				_lowerCAmelCase	(   self			):

       '''simple docstring'''

       a_   :						int   =				load_numpy(
           """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
           """/shap_e/test_shap_e_np_out.npy"""			)
       a_   :						List[str]   =				ShapEPipeline.from_pretrained("""openai/shap-e"""			)
       a_   :						Optional[Any]   =				pipe.to(lowerCAmelCase_			)
       pipe.set_progress_bar_config(disable=lowerCAmelCase_			)

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

       a_   :						List[str]   =				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 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       a_       (     SCREAMING_SNAKE_CASE__  ,    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 						A_(							self      )  ->					Optional[Any]:




				"""simple docstring"""



				return 32





	@property
	def 						A_(							self      )  ->					Dict:




				"""simple docstring"""



				return 32





	@property
	def 						A_(							self      )  ->					str:




				"""simple docstring"""



				return self.time_input_dim * 4





	@property
	def 						A_(							self      )  ->					str:




				"""simple docstring"""



				return 8





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




				"""simple docstring"""



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





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




				"""simple docstring"""



				torch.manual_seed(0      )
				SCREAMING_SNAKE_CASE_								=						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=1000      ,     )
				return CLIPTextModelWithProjection(SCREAMING_SNAKE_CASE      )





	@property
	def 						A_(							self      )  ->					List[Any]:




				"""simple docstring"""



				torch.manual_seed(0      )

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

				SCREAMING_SNAKE_CASE_								=						PriorTransformer(**SCREAMING_SNAKE_CASE      )
				return model





	@property
	def 						A_(							self      )  ->					Union[str, Any]:




				"""simple docstring"""



				torch.manual_seed(0      )

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





	def 						A_(							self      )  ->					int:




				"""simple docstring"""



				SCREAMING_SNAKE_CASE_								=						self.dummy_prior
				SCREAMING_SNAKE_CASE_								=						self.dummy_text_encoder
				SCREAMING_SNAKE_CASE_								=						self.dummy_tokenizer
				SCREAMING_SNAKE_CASE_								=						self.dummy_renderer

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

				return components





	def 						A_(							self      ,     SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE=0      )  ->					str:




				"""simple docstring"""



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





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




				"""simple docstring"""



				SCREAMING_SNAKE_CASE_								=						'cpu'

				SCREAMING_SNAKE_CASE_								=						self.get_dummy_components()

				SCREAMING_SNAKE_CASE_								=						self.pipeline_class(**SCREAMING_SNAKE_CASE      )
				SCREAMING_SNAKE_CASE_								=						pipe.to(SCREAMING_SNAKE_CASE      )

				pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE      )

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

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

				SCREAMING_SNAKE_CASE_								=						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 						A_(							self      )  ->					List[str]:




				"""simple docstring"""



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





	def 						A_(							self      )  ->					str:




				"""simple docstring"""



				SCREAMING_SNAKE_CASE_								=						torch_device == 'cpu'
				SCREAMING_SNAKE_CASE_								=						True

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





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




				"""simple docstring"""



				SCREAMING_SNAKE_CASE_								=						self.get_dummy_components()
				SCREAMING_SNAKE_CASE_								=						self.pipeline_class(**SCREAMING_SNAKE_CASE      )
				SCREAMING_SNAKE_CASE_								=						pipe.to(SCREAMING_SNAKE_CASE      )
				pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE      )

				SCREAMING_SNAKE_CASE_								=						1
				SCREAMING_SNAKE_CASE_								=						2

				SCREAMING_SNAKE_CASE_								=						self.get_dummy_inputs(SCREAMING_SNAKE_CASE      )

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

				SCREAMING_SNAKE_CASE_								=						pipe(**SCREAMING_SNAKE_CASE      ,     num_images_per_prompt=SCREAMING_SNAKE_CASE      )[0]

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




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





	def 						A_(							self      )  ->					List[str]:




				"""simple docstring"""



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





	def 						A_(							self      )  ->					List[str]:




				"""simple docstring"""



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

				SCREAMING_SNAKE_CASE_								=						torch.Generator(device=SCREAMING_SNAKE_CASE      ).manual_seed(0      )

				SCREAMING_SNAKE_CASE_								=						pipe(
				    'a shark'      ,     generator=SCREAMING_SNAKE_CASE      ,     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(SCREAMING_SNAKE_CASE      ,     SCREAMING_SNAKE_CASE      )
