



"""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        __SCREAMING_SNAKE_CASE	(      lowerCAmelCase_		,   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 		snake_case       (       self							:  str     )->					List[str]:
								return 32



					@property
					def 		snake_case       (       self							:  Any     )->					str:
								return 32



					@property
					def 		snake_case       (       self							:  Dict     )->					Dict:
								return self.time_input_dim * 4



					@property
					def 		snake_case       (       self							:  Union[str, Any]     )->					Union[str, Any]:
								return 8



					@property
					def 		snake_case       (       self							:  Optional[int]     )->					List[str]:
								lowerCamelCase__		:		Optional[int]			=CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''     )
								return tokenizer



					@property
					def 		snake_case       (       self							:  str     )->					str:
								torch.manual_seed(0     )
								lowerCamelCase__		:		Dict			=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(lowerCamelCase     )



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

								lowerCamelCase__		:		List[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,
								}

								lowerCamelCase__		:		Optional[Any]			=PriorTransformer(**lowerCamelCase     )
								return model



					@property
					def 		snake_case       (       self							:  str     )->					List[str]:
								torch.manual_seed(0     )

								lowerCamelCase__		:		Optional[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,
								    ),
								}
								lowerCamelCase__		:		Dict			=ShapERenderer(**lowerCamelCase     )
								return model



					def 		snake_case       (       self							:  Optional[int]     )->					Tuple:
								lowerCamelCase__		:		Optional[Any]			=self.dummy_prior
								lowerCamelCase__		:		Optional[int]			=self.dummy_text_encoder
								lowerCamelCase__		:		str			=self.dummy_tokenizer
								lowerCamelCase__		:		int			=self.dummy_renderer

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

								return components



					def 		snake_case       (       self							:  Optional[Any],  lowerCamelCase							:  Optional[Any],  lowerCamelCase							:  Any=0     )->					str:
								if str(lowerCamelCase     ).startswith('''mps'''     ):
											lowerCamelCase__		:		str			=torch.manual_seed(lowerCamelCase     )
								else:
											lowerCamelCase__		:		Any			=torch.Generator(device=lowerCamelCase     ).manual_seed(lowerCamelCase     )
								lowerCamelCase__		:		Optional[Any]			={
								    '''prompt''': '''horse''',
								    '''generator''': generator,
								    '''num_inference_steps''': 1,
								    '''frame_size''': 32,
								    '''output_type''': '''np''',
								}
								return inputs



					def 		snake_case       (       self							:  List[Any]     )->					List[str]:
								lowerCamelCase__		:		Tuple			='''cpu'''

								lowerCamelCase__		:		List[Any]			=self.get_dummy_components()

								lowerCamelCase__		:		int			=self.pipeline_class(**lowerCamelCase     )
								lowerCamelCase__		:		Dict			=pipe.to(lowerCamelCase     )

								pipe.set_progress_bar_config(disable=lowerCamelCase     )

								lowerCamelCase__		:		int			=pipe(**self.get_dummy_inputs(lowerCamelCase     )     )
								lowerCamelCase__		:		int			=output.images[0]
								lowerCamelCase__		:		int			=image[0, -3:, -3:, -1]

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

								lowerCamelCase__		:		Optional[Any]			=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 		snake_case       (       self							:  Optional[Any]     )->					Dict:
								# NOTE: Larger batch sizes cause this test to timeout, only test on smaller batches
								self._test_inference_batch_consistent(batch_sizes=[1, 2]     )



					def 		snake_case       (       self							:  int     )->					str:
								lowerCamelCase__		:		Optional[int]			=torch_device == '''cpu'''
								lowerCamelCase__		:		Tuple			=True

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




					def 		snake_case       (       self							:  int     )->					Dict:
								lowerCamelCase__		:		str			=self.get_dummy_components()
								lowerCamelCase__		:		Any			=self.pipeline_class(**lowerCamelCase     )
								lowerCamelCase__		:		List[str]			=pipe.to(lowerCamelCase     )
								pipe.set_progress_bar_config(disable=lowerCamelCase     )

								lowerCamelCase__		:		List[Any]			=1
								lowerCamelCase__		:		str			=2

								lowerCamelCase__		:		List[str]			=self.get_dummy_inputs(lowerCamelCase     )

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

								lowerCamelCase__		:		Tuple			=pipe(**lowerCamelCase,  num_images_per_prompt=lowerCamelCase     )[0]

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





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


					'''simple docstring'''





					def 		snake_case       (       self							:  List[str]     )->					Optional[Any]:
								# clean up the VRAM after each test
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()




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

								lowerCamelCase__		:		Union[str, Any]			=torch.Generator(device=lowerCamelCase     ).manual_seed(0     )

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









"""simple docstring"""






from typing import Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature
from ...image_transforms import get_image_size, pad, rescale, to_channel_dimension_format
from ...image_utils import ChannelDimension, ImageInput, make_list_of_images, to_numpy_array, valid_images
from ...utils import TensorType, logging


_lowercase : Tuple      						=						logging.get_logger(__name__)

class        __SCREAMING_SNAKE_CASE	(      lowerCAmelCase_   ):


					'''simple docstring'''





					_a   			=							['pixel_values']



					def __init__(       self							:  Optional[Any],  lowerCamelCase							:  bool = True,  lowerCamelCase							:  Union[int, float] = 1 / 255,  lowerCamelCase							:  bool = True,  lowerCamelCase							:  int = 8,  **lowerCamelCase							:  Tuple,  )->					None:
								super().__init__(**lowerCamelCase     )

								lowerCamelCase__		:		int			=do_rescale
								lowerCamelCase__		:		Dict			=rescale_factor
								lowerCamelCase__		:		Union[str, Any]			=do_pad
								lowerCamelCase__		:		Union[str, Any]			=pad_size



					def 		snake_case       (       self							:  int,  lowerCamelCase							:  np.ndarray,  lowerCamelCase							:  float,  lowerCamelCase							:  Optional[Union[str, ChannelDimension]] = None,  **lowerCamelCase							:  int     )->					np.ndarray:
								return rescale(lowerCamelCase,  scale=lowerCamelCase,  data_format=lowerCamelCase,  **lowerCamelCase     )



					def 		snake_case       (       self							:  Optional[Any],  lowerCamelCase							:  np.ndarray,  lowerCamelCase							:  int,  lowerCamelCase							:  Optional[Union[str, ChannelDimension]] = None     )->					List[Any]:
								lowerCamelCase__	,   lowerCamelCase__		:		Optional[int]			=get_image_size(lowerCamelCase     )
								lowerCamelCase__		:		List[str]			=(old_height // size + 1) * size - old_height
								lowerCamelCase__		:		List[str]			=(old_width // size + 1) * size - old_width

								return pad(lowerCamelCase,  ((0, pad_height), (0, pad_width)),  mode='''symmetric''',  data_format=lowerCamelCase     )




					def 		snake_case       (       self							:  List[Any],  lowerCamelCase							:  ImageInput,  lowerCamelCase							:  Optional[bool] = None,  lowerCamelCase							:  Optional[float] = None,  lowerCamelCase							:  Optional[bool] = None,  lowerCamelCase							:  Optional[int] = None,  lowerCamelCase							:  Optional[Union[str, TensorType]] = None,  lowerCamelCase							:  Union[str, ChannelDimension] = ChannelDimension.FIRST,  **lowerCamelCase							:  Union[str, Any],  )->					Dict:
								lowerCamelCase__		:		List[str]			=do_rescale if do_rescale is not None else self.do_rescale
								lowerCamelCase__		:		Tuple			=rescale_factor if rescale_factor is not None else self.rescale_factor
								lowerCamelCase__		:		str			=do_pad if do_pad is not None else self.do_pad
								lowerCamelCase__		:		int			=pad_size if pad_size is not None else self.pad_size

								lowerCamelCase__		:		Optional[int]			=make_list_of_images(lowerCamelCase     )

								if not valid_images(lowerCamelCase     ):
											raise ValueError(
											    '''Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, '''
											    '''torch.Tensor, tf.Tensor or jax.ndarray.'''     )

								if do_rescale and rescale_factor is None:
											raise ValueError('''Rescale factor must be specified if do_rescale is True.'''     )

								# All transformations expect numpy arrays.
								lowerCamelCase__		:		Tuple			=[to_numpy_array(lowerCamelCase     ) for image in images]

								if do_rescale:
											lowerCamelCase__		:		Tuple			=[self.rescale(image=lowerCamelCase,  scale=lowerCamelCase     ) for image in images]

								if do_pad:
											lowerCamelCase__		:		Tuple			=[self.pad(lowerCamelCase,  size=lowerCamelCase     ) for image in images]

								lowerCamelCase__		:		int			=[to_channel_dimension_format(lowerCamelCase,  lowerCamelCase     ) for image in images]

								lowerCamelCase__		:		Dict			={'''pixel_values''': images}
								return BatchFeature(data=lowerCamelCase,  tensor_type=lowerCamelCase     )




