






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      ):
      snake_case_	=				ShapEPipeline
      snake_case_	=				['''prompt''']
      snake_case_	=				['''prompt''']
      snake_case_	=				[
          '''num_images_per_prompt''',
          '''num_inference_steps''',
          '''generator''',
          '''latents''',
          '''guidance_scale''',
          '''frame_size''',
          '''output_type''',
          '''return_dict''',
      ]
      snake_case_	=				False


      @property
      def 			__UpperCAmelCase		(						self						)->						Tuple:
           return 3_2


      @property
      def 			__UpperCAmelCase		(						self						)->						List[Any]:
           return 3_2


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


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


      @property
      def 			__UpperCAmelCase		(						self						)->						Tuple:
           __lowerCAmelCase             =       CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"						)
           return tokenizer


      @property
      def 			__UpperCAmelCase		(						self						)->						Tuple:
           torch.manual_seed(0						)
           __lowerCAmelCase             =       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=3_7       ,    layer_norm_eps=1e-05       ,    num_attention_heads=4       ,    num_hidden_layers=5       ,    pad_token_id=1       ,    vocab_size=1_0_0_0       ,    )
           return CLIPTextModelWithProjection(_lowerCAmelCase						)


      @property
      def 			__UpperCAmelCase		(						self						)->						Tuple:
           torch.manual_seed(0						)

           __lowerCAmelCase             =       {
               '''num_attention_heads''': 2,
               '''attention_head_dim''': 1_6,
               '''embedding_dim''': self.time_input_dim,
               '''num_embeddings''': 3_2,
               '''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             =       PriorTransformer(**_lowerCAmelCase						)
           return model


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

           __lowerCAmelCase             =       {
               '''param_shapes''': (
                   (self.renderer_dim, 9_3),
                   (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''': 1_2,
               '''background''': (
                   0.1,
                   0.1,
                   0.1,
               ),
           }
           __lowerCAmelCase             =       ShapERenderer(**_lowerCAmelCase						)
           return model


      def 			__UpperCAmelCase		(						self						)->						str:
           __lowerCAmelCase             =       self.dummy_prior
           __lowerCAmelCase             =       self.dummy_text_encoder
           __lowerCAmelCase             =       self.dummy_tokenizer
           __lowerCAmelCase             =       self.dummy_renderer

           __lowerCAmelCase             =       HeunDiscreteScheduler(
               beta_schedule="exp"       ,    num_train_timesteps=1_0_2_4       ,    prediction_type="sample"       ,    use_karras_sigmas=_lowerCAmelCase       ,    clip_sample=_lowerCAmelCase       ,    clip_sample_range=1.0       ,    )
           __lowerCAmelCase             =       {
               '''prior''': prior,
               '''text_encoder''': text_encoder,
               '''tokenizer''': tokenizer,
               '''renderer''': renderer,
               '''scheduler''': scheduler,
           }

           return components


      def 			__UpperCAmelCase		(						self       ,    __UpperCamelCase       ,    __UpperCamelCase=0						)->						List[Any]:




           if str(_lowerCAmelCase						).startswith("mps"						):
                __lowerCAmelCase             =       torch.manual_seed(_lowerCAmelCase						)
           else:
                __lowerCAmelCase             =       torch.Generator(device=_lowerCAmelCase						).manual_seed(_lowerCAmelCase						)
           __lowerCAmelCase             =       {
               '''prompt''': '''horse''',
               '''generator''': generator,
               '''num_inference_steps''': 1,
               '''frame_size''': 3_2,
               '''output_type''': '''np''',
           }
           return inputs


      def 			__UpperCAmelCase		(						self						)->						Dict:
           __lowerCAmelCase             =       '''cpu'''

           __lowerCAmelCase             =       self.get_dummy_components()

           __lowerCAmelCase             =       self.pipeline_class(**_lowerCAmelCase						)
           __lowerCAmelCase             =       pipe.to(_lowerCAmelCase						)

           pipe.set_progress_bar_config(disable=_lowerCAmelCase						)

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

           assert image.shape == (2_0, 3_2, 3_2, 3)

           __lowerCAmelCase             =       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						)->						int:
           self._test_inference_batch_consistent(batch_sizes=[1, 2]						)


      def 			__UpperCAmelCase		(						self						)->						List[Any]:
           __lowerCAmelCase             =       torch_device == '''cpu'''
           __lowerCAmelCase             =       True

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


      def 			__UpperCAmelCase		(						self						)->						List[Any]:
           __lowerCAmelCase             =       self.get_dummy_components()
           __lowerCAmelCase             =       self.pipeline_class(**_lowerCAmelCase						)
           __lowerCAmelCase             =       pipe.to(_lowerCAmelCase						)
           pipe.set_progress_bar_config(disable=_lowerCAmelCase						)

           __lowerCAmelCase             =       1
           __lowerCAmelCase             =       2

           __lowerCAmelCase             =       self.get_dummy_inputs(_lowerCAmelCase						)

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

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


      def 			__UpperCAmelCase		(						self						)->						List[str]:
           super().tearDown()
           gc.collect()
           torch.cuda.empty_cache()


      def 			__UpperCAmelCase		(						self						)->						List[str]:
           __lowerCAmelCase             =       load_numpy(
               "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
               "/shap_e/test_shap_e_np_out.npy"						)
           __lowerCAmelCase             =       ShapEPipeline.from_pretrained("openai/shap-e"						)
           __lowerCAmelCase             =       pipe.to(_lowerCAmelCase						)
           pipe.set_progress_bar_config(disable=_lowerCAmelCase						)

           __lowerCAmelCase             =       torch.Generator(device=_lowerCAmelCase						).manual_seed(0						)

           __lowerCAmelCase             =       pipe(
               "a shark"       ,    generator=_lowerCAmelCase       ,    guidance_scale=1_5.0       ,    num_inference_steps=6_4       ,    frame_size=6_4       ,    output_type="np"       ,    ).images[0]

           assert images.shape == (2_0, 6_4, 6_4, 3)

           assert_mean_pixel_difference(_lowerCAmelCase       ,    _lowerCAmelCase						)













from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


if is_vision_available():
							import PIL


lowerCAmelCase							:       str					 =	logging.get_logger(__name__)





class 				a			(      __lowercase			):
				SCREAMING_SNAKE_CASE__     :   Any		      =     ['''pixel_values''']






				def __init__(  self			, _lowerCAmelCase = True			, _lowerCAmelCase = None			, _lowerCAmelCase = None			, _lowerCAmelCase = PILImageResampling.BILINEAR			, _lowerCAmelCase = True			, _lowerCAmelCase = 1 / 255			, _lowerCAmelCase = True			, _lowerCAmelCase = None			, _lowerCAmelCase = None			, **_lowerCAmelCase			, ):

										"""simple docstring"""




										super().__init__(**_lowerCAmelCase					)
										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   size if size is not None else {'''shortest_edge''': 384}
										__SCREAMING_SNAKE_CASE:				List[str]           =   get_size_dict(_lowerCAmelCase			, default_to_square=_lowerCAmelCase					)

										__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   do_resize
										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   size
										# Default value set here for backwards compatibility where the value in config is None
										__SCREAMING_SNAKE_CASE:				List[str]           =   crop_pct if crop_pct is not None else 224 / 256
										__SCREAMING_SNAKE_CASE:				List[Any]           =   resample
										__SCREAMING_SNAKE_CASE:				Any           =   do_rescale
										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   rescale_factor
										__SCREAMING_SNAKE_CASE:				Any           =   do_normalize
										__SCREAMING_SNAKE_CASE:				Optional[int]           =   image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
										__SCREAMING_SNAKE_CASE:				Optional[int]           =   image_std if image_std is not None else IMAGENET_STANDARD_STD






				def 			snake_case_    (  self			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase = PILImageResampling.BICUBIC			, _lowerCAmelCase = None			, **_lowerCAmelCase			, ):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				List[str]           =   get_size_dict(_lowerCAmelCase			, default_to_square=_lowerCAmelCase					)
										if "shortest_edge" not in size:
																raise ValueError(f"""Size dictionary must contain 'shortest_edge' key. Got {size.keys()}"""					)
										__SCREAMING_SNAKE_CASE:				str           =   size['''shortest_edge''']

										if shortest_edge < 384:
																# maintain same ratio, resizing shortest edge to shortest_edge/crop_pct
																__SCREAMING_SNAKE_CASE:				Dict           =   int(shortest_edge / crop_pct					)
																__SCREAMING_SNAKE_CASE:				Any           =   get_resize_output_image_size(_lowerCAmelCase			, size=_lowerCAmelCase			, default_to_square=_lowerCAmelCase					)
																__SCREAMING_SNAKE_CASE:				Optional[Any]           =   resize(image=_lowerCAmelCase			, size=_lowerCAmelCase			, resample=_lowerCAmelCase			, data_format=_lowerCAmelCase			, **_lowerCAmelCase					)
																# then crop to (shortest_edge, shortest_edge)
																return center_crop(image=_lowerCAmelCase			, size=(shortest_edge, shortest_edge)			, data_format=_lowerCAmelCase			, **_lowerCAmelCase					)
										else:
																# warping (no cropping) when evaluated at 384 or larger
																return resize(
																    _lowerCAmelCase			, size=(shortest_edge, shortest_edge)			, resample=_lowerCAmelCase			, data_format=_lowerCAmelCase			, **_lowerCAmelCase					)






				def 			snake_case_    (  self			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase = None			, **_lowerCAmelCase			, ):

										"""simple docstring"""




										return rescale(_lowerCAmelCase			, scale=_lowerCAmelCase			, data_format=_lowerCAmelCase			, **_lowerCAmelCase					)






				def 			snake_case_    (  self			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase = None			, **_lowerCAmelCase			, ):

										"""simple docstring"""




										return normalize(_lowerCAmelCase			, mean=_lowerCAmelCase			, std=_lowerCAmelCase			, data_format=_lowerCAmelCase			, **_lowerCAmelCase					)






				def 			snake_case_    (  self			, _lowerCAmelCase			, _lowerCAmelCase = None			, _lowerCAmelCase = None			, _lowerCAmelCase = None			, _lowerCAmelCase = None			, _lowerCAmelCase = None			, _lowerCAmelCase = None			, _lowerCAmelCase = None			, _lowerCAmelCase = None			, _lowerCAmelCase = None			, _lowerCAmelCase = None			, _lowerCAmelCase = ChannelDimension.FIRST			, **_lowerCAmelCase			, ):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				List[Any]           =   do_resize if do_resize is not None else self.do_resize
										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   crop_pct if crop_pct is not None else self.crop_pct
										__SCREAMING_SNAKE_CASE:				Tuple           =   resample if resample is not None else self.resample
										__SCREAMING_SNAKE_CASE:				Optional[int]           =   do_rescale if do_rescale is not None else self.do_rescale
										__SCREAMING_SNAKE_CASE:				Any           =   rescale_factor if rescale_factor is not None else self.rescale_factor
										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   do_normalize if do_normalize is not None else self.do_normalize
										__SCREAMING_SNAKE_CASE:				Optional[int]           =   image_mean if image_mean is not None else self.image_mean
										__SCREAMING_SNAKE_CASE:				str           =   image_std if image_std is not None else self.image_std

										__SCREAMING_SNAKE_CASE:				List[Any]           =   size if size is not None else self.size
										__SCREAMING_SNAKE_CASE:				Dict           =   get_size_dict(_lowerCAmelCase			, default_to_square=_lowerCAmelCase					)

										__SCREAMING_SNAKE_CASE:				List[str]           =   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_resize and size is None or resample is None:
																raise ValueError('''Size and resample must be specified if do_resize is True.'''					)

										if do_resize and size["shortest_edge"] < 384 and crop_pct is None:
																raise ValueError('''crop_pct must be specified if size < 384.'''					)

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

										if do_normalize and (image_mean is None or image_std is None):
																raise ValueError('''Image mean and std must be specified if do_normalize is True.'''					)

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

										if do_resize:
																__SCREAMING_SNAKE_CASE:				str           =   [self.resize(image=_lowerCAmelCase			, size=_lowerCAmelCase			, crop_pct=_lowerCAmelCase			, resample=_lowerCAmelCase					) for image in images]

										if do_rescale:
																__SCREAMING_SNAKE_CASE:				Any           =   [self.rescale(image=_lowerCAmelCase			, scale=_lowerCAmelCase					) for image in images]

										if do_normalize:
																__SCREAMING_SNAKE_CASE:				int           =   [self.normalize(image=_lowerCAmelCase			, mean=_lowerCAmelCase			, std=_lowerCAmelCase					) for image in images]

										__SCREAMING_SNAKE_CASE:				str           =   [to_channel_dimension_format(_lowerCAmelCase			, _lowerCAmelCase					) for image in images]

										__SCREAMING_SNAKE_CASE:				List[str]           =   {'''pixel_values''': images}
										return BatchFeature(data=_lowerCAmelCase			, tensor_type=_lowerCAmelCase					)



