





"""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'''




		__lowerCAmelCase  =  ShapEPipeline
		__lowerCAmelCase  =  ['prompt']
		__lowerCAmelCase  =  ['prompt']
		__lowerCAmelCase  =  [
		    'num_images_per_prompt',
		    'num_inference_steps',
		    'generator',
		    'latents',
		    'guidance_scale',
		    'frame_size',
		    'output_type',
		    'return_dict',
		]
		__lowerCAmelCase  =  False







		@property
		def  _lowerCamelCase		(						self     ):
						return 32







		@property
		def  _lowerCamelCase		(						self     ):
						return 32







		@property
		def  _lowerCamelCase		(						self     ):
						return self.time_input_dim * 4







		@property
		def  _lowerCamelCase		(						self     ):
						return 8







		@property
		def  _lowerCamelCase		(						self     ):
						__a :						Any			       =  CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''     )
						return tokenizer







		@property
		def  _lowerCamelCase		(						self     ):
						torch.manual_seed(0     )
						__a :						Tuple			       =  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(UpperCamelCase_     )







		@property
		def  _lowerCamelCase		(						self     ):
						torch.manual_seed(0     )

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

						__a :						Tuple			       =  PriorTransformer(**UpperCamelCase_     )
						return model







		@property
		def  _lowerCamelCase		(						self     ):
						torch.manual_seed(0     )

						__a :						Union[str, 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,
						    ),
						}
						__a :						List[Any]			       =  ShapERenderer(**UpperCamelCase_     )
						return model







		def  _lowerCamelCase		(						self     ):
						__a :						Dict			       =  self.dummy_prior
						__a :						Optional[int]			       =  self.dummy_text_encoder
						__a :						List[Any]			       =  self.dummy_tokenizer
						__a :						List[str]			       =  self.dummy_renderer

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

						return components







		def  _lowerCamelCase		(						self	,			_UpperCAmelCase	,			_UpperCAmelCase=0     ):
						if str(UpperCamelCase_     ).startswith('''mps'''     ):
										__a :						int			       =  torch.manual_seed(UpperCamelCase_     )
						else:
										__a :						Optional[Any]			       =  torch.Generator(device=UpperCamelCase_     ).manual_seed(UpperCamelCase_     )
						__a :						List[str]			       =  {
						    '''prompt''': '''horse''',
						    '''generator''': generator,
						    '''num_inference_steps''': 1,
						    '''frame_size''': 32,
						    '''output_type''': '''np''',
						}
						return inputs







		def  _lowerCamelCase		(						self     ):
						__a :						Any			       =  '''cpu'''

						__a :						Any			       =  self.get_dummy_components()

						__a :						List[str]			       =  self.pipeline_class(**UpperCamelCase_     )
						__a :						Optional[int]			       =  pipe.to(UpperCamelCase_     )

						pipe.set_progress_bar_config(disable=UpperCamelCase_     )

						__a :						List[str]			       =  pipe(**self.get_dummy_inputs(UpperCamelCase_     )     )
						__a :						Optional[int]			       =  output.images[0]
						__a :						Tuple			       =  image[0, -3:, -3:, -1]

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

						__a :						List[str]			       =  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  _lowerCamelCase		(						self     ):
						# 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     ):
						__a :						int			       =  torch_device == '''cpu'''
						__a :						int			       =  True

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







		def  _lowerCamelCase		(						self     ):
						__a :						Tuple			       =  self.get_dummy_components()
						__a :						List[str]			       =  self.pipeline_class(**UpperCamelCase_     )
						__a :						Any			       =  pipe.to(UpperCamelCase_     )
						pipe.set_progress_bar_config(disable=UpperCamelCase_     )

						__a :						Optional[Any]			       =  1
						__a :						int			       =  2

						__a :						Optional[Any]			       =  self.get_dummy_inputs(UpperCamelCase_     )

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

						__a :						Dict			       =  pipe(**UpperCamelCase_	,			num_images_per_prompt=UpperCamelCase_     )[0]

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






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







		'''simple docstring'''




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







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

						__a :						str			       =  torch.Generator(device=UpperCamelCase_     ).manual_seed(0     )

						__a :						Optional[int]			       =  pipe(
						    '''a shark'''	,			generator=UpperCamelCase_	,			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(UpperCamelCase_	,			UpperCamelCase_     )






"""simple docstring"""






from dataclasses import dataclass, field
from typing import ClassVar, Dict

from ..features import Features, Value
from .base import TaskTemplate




@dataclass(frozen=_UpperCamelCase							)
class 				__lowercase			(			_UpperCamelCase							):







					'''simple docstring'''




					__lowerCAmelCase  =  field(default='''language-modeling'''   ,      metadata={'''include_in_asdict_even_if_is_default''': True}							)

					__lowerCAmelCase  =  Features({'''text''': Value('''string'''							)}							)
					__lowerCAmelCase  =  Features({}							)
					__lowerCAmelCase  =  "text"







					@property
					def  _lowerCamelCase		(						self     ):
									return {self.text_column: "text"}