


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 convert_to_rgb, normalize, rescale, resize, to_channel_dimension_format
from ...image_utils import (
    OPENAI_CLIP_MEAN,
    OPENAI_CLIP_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	=							logging.get_logger(__name__)


class 	_a							(      UpperCamelCase__      ):
       _lowercase     :      Any  				=			['''pixel_values''']

       def __init__(							self:     Optional[int]   ,						UpperCamelCase_:     bool = True   ,						UpperCamelCase_:     Dict[str, int] = None   ,						UpperCamelCase_:     PILImageResampling = PILImageResampling.BICUBIC   ,						UpperCamelCase_:     bool = True   ,						UpperCamelCase_:     Union[int, float] = 1 / 255   ,						UpperCamelCase_:     bool = True   ,						UpperCamelCase_:     Optional[Union[float, List[float]]] = None   ,						UpperCamelCase_:     Optional[Union[float, List[float]]] = None   ,						UpperCamelCase_:     bool = True   ,						**UpperCamelCase_:     Optional[Any]   ,						)      ->       None:





         """simple docstring"""







         super().__init__(**UpperCamelCase_							)
         lowercase__						=				size if size is not None else {'''height''': 384, '''width''': 384}
         lowercase__						=				get_size_dict(UpperCamelCase_   ,						default_to_square=UpperCamelCase_							)

         lowercase__						=				do_resize
         lowercase__						=				size
         lowercase__						=				resample
         lowercase__						=				do_rescale
         lowercase__						=				rescale_factor
         lowercase__						=				do_normalize
         lowercase__						=				image_mean if image_mean is not None else OPENAI_CLIP_MEAN
         lowercase__						=				image_std if image_std is not None else OPENAI_CLIP_STD
         lowercase__						=				do_convert_rgb

       def 				lowerCamelCase_			(							self:     int   ,						UpperCamelCase_:     np.ndarray   ,						UpperCamelCase_:     Dict[str, int]   ,						UpperCamelCase_:     PILImageResampling = PILImageResampling.BICUBIC   ,						UpperCamelCase_:     Optional[Union[str, ChannelDimension]] = None   ,						**UpperCamelCase_:     Any   ,						)      ->       np.ndarray:





         """simple docstring"""







         lowercase__						=				get_size_dict(UpperCamelCase_   ,						default_to_square=UpperCamelCase_							)
         if "height" not in size or "width" not in size:
           raise ValueError(f'The `size` dictionary must contain the keys `height` and `width`. Got {size.keys()}'							)
         lowercase__						=				(size['''height'''], size['''width'''])
         return resize(UpperCamelCase_   ,						size=UpperCamelCase_   ,						resample=UpperCamelCase_   ,						data_format=UpperCamelCase_   ,						**UpperCamelCase_							)

       def 				lowerCamelCase_			(							self:     str   ,						UpperCamelCase_:     np.ndarray   ,						UpperCamelCase_:     Union[int, float]   ,						UpperCamelCase_:     Optional[Union[str, ChannelDimension]] = None   ,						**UpperCamelCase_:     Optional[int]   ,						)      ->       List[str]:





         """simple docstring"""







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

       def 				lowerCamelCase_			(							self:     List[Any]   ,						UpperCamelCase_:     np.ndarray   ,						UpperCamelCase_:     Union[float, List[float]]   ,						UpperCamelCase_:     Union[float, List[float]]   ,						UpperCamelCase_:     Optional[Union[str, ChannelDimension]] = None   ,						**UpperCamelCase_:     Optional[int]   ,						)      ->       np.ndarray:





         """simple docstring"""







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


       def 				lowerCamelCase_			(							self:     Optional[Any]   ,						UpperCamelCase_:     ImageInput   ,						UpperCamelCase_:     Optional[bool] = None   ,						UpperCamelCase_:     Optional[Dict[str, int]] = None   ,						UpperCamelCase_:     PILImageResampling = None   ,						UpperCamelCase_:     Optional[bool] = None   ,						UpperCamelCase_:     Optional[float] = None   ,						UpperCamelCase_:     Optional[bool] = None   ,						UpperCamelCase_:     Optional[Union[float, List[float]]] = None   ,						UpperCamelCase_:     Optional[Union[float, List[float]]] = None   ,						UpperCamelCase_:     Optional[Union[str, TensorType]] = None   ,						UpperCamelCase_:     bool = None   ,						UpperCamelCase_:     ChannelDimension = ChannelDimension.FIRST   ,						**UpperCamelCase_:     Union[str, Any]   ,						)      ->       PIL.Image.Image:





         """simple docstring"""







         lowercase__						=				do_resize if do_resize is not None else self.do_resize
         lowercase__						=				resample if resample is not None else self.resample
         lowercase__						=				do_rescale if do_rescale is not None else self.do_rescale
         lowercase__						=				rescale_factor if rescale_factor is not None else self.rescale_factor
         lowercase__						=				do_normalize if do_normalize is not None else self.do_normalize
         lowercase__						=				image_mean if image_mean is not None else self.image_mean
         lowercase__						=				image_std if image_std is not None else self.image_std
         lowercase__						=				do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb

         lowercase__						=				size if size is not None else self.size
         lowercase__						=				get_size_dict(UpperCamelCase_   ,						default_to_square=UpperCamelCase_							)

         lowercase__						=				make_list_of_images(UpperCamelCase_							)

         if not valid_images(UpperCamelCase_							):
           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_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.'''							)

         # PIL RGBA images are converted to RGB
         if do_convert_rgb:
           lowercase__						=				[convert_to_rgb(UpperCamelCase_							) for image in images]

         # All transformations expect numpy arrays.
         lowercase__						=				[to_numpy_array(UpperCamelCase_							) for image in images]

         if do_resize:
           lowercase__						=				[self.resize(image=UpperCamelCase_   ,						size=UpperCamelCase_   ,						resample=UpperCamelCase_							) for image in images]

         if do_rescale:
           lowercase__						=				[self.rescale(image=UpperCamelCase_   ,						scale=UpperCamelCase_							) for image in images]

         if do_normalize:
           lowercase__						=				[self.normalize(image=UpperCamelCase_   ,						mean=UpperCamelCase_   ,						std=UpperCamelCase_							) for image in images]

         lowercase__						=				[to_channel_dimension_format(UpperCamelCase_   ,						UpperCamelCase_							) for image in images]

         lowercase__						=				BatchFeature(data={'''pixel_values''': images}   ,						tensor_type=UpperCamelCase_							)

         return encoded_outputs






import gc
import random
import unittest

import numpy as np
import torch
from transformers import CLIPImageProcessor, CLIPVisionConfig, CLIPVisionModel

from diffusers import HeunDiscreteScheduler, PriorTransformer, ShapEImgaImgPipeline
from diffusers.pipelines.shap_e import ShapERenderer
from diffusers.utils import floats_tensor, load_image, 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							(      UpperCamelCase__					,     unittest.TestCase      ):
							_lowercase     :      Union[str, Any]  				=			ShapEImgaImgPipeline
							_lowercase     :      Optional[Any]  				=			['''image''']
							_lowercase     :      Optional[int]  				=			['''image''']
							_lowercase     :      Optional[int]  				=			[
							    '''num_images_per_prompt''',
							    '''num_inference_steps''',
							    '''generator''',
							    '''latents''',
							    '''guidance_scale''',
							    '''frame_size''',
							    '''output_type''',
							    '''return_dict''',
							]
							_lowercase     :      Tuple  				=			False

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





									"""simple docstring"""







									return 32

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





									"""simple docstring"""







									return 32

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





									"""simple docstring"""







									return self.time_input_dim * 4

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





									"""simple docstring"""







									return 8

							@property
							def 				lowerCamelCase_			(							self:     int							)      ->       Dict:





									"""simple docstring"""







									torch.manual_seed(0							)
									lowercase__						=				CLIPVisionConfig(
									    hidden_size=self.text_embedder_hidden_size   ,						image_size=64   ,						projection_dim=self.text_embedder_hidden_size   ,						intermediate_size=37   ,						num_attention_heads=4   ,						num_channels=3   ,						num_hidden_layers=5   ,						patch_size=1   ,						)

									lowercase__						=				CLIPVisionModel(UpperCamelCase_							)
									return model

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





									"""simple docstring"""







									lowercase__						=				CLIPImageProcessor(
									    crop_size=224   ,						do_center_crop=UpperCamelCase_   ,						do_normalize=UpperCamelCase_   ,						do_resize=UpperCamelCase_   ,						image_mean=[0.48145466, 0.4578275, 0.40821073]   ,						image_std=[0.26862954, 0.26130258, 0.27577711]   ,						resample=3   ,						size=224   ,						)

									return image_processor

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





									"""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''',
									    '''embedding_proj_norm_type''': '''layer''',
									    '''encoder_hid_proj_type''': None,
									    '''added_emb_type''': None,
									}

									lowercase__						=				PriorTransformer(**UpperCamelCase_							)
									return model

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





									"""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(**UpperCamelCase_							)
									return model

							def 				lowerCamelCase_			(							self:     str							)      ->       Any:





									"""simple docstring"""







									lowercase__						=				self.dummy_prior
									lowercase__						=				self.dummy_image_encoder
									lowercase__						=				self.dummy_image_processor
									lowercase__						=				self.dummy_renderer

									lowercase__						=				HeunDiscreteScheduler(
									    beta_schedule='''exp'''   ,						num_train_timesteps=1_024   ,						prediction_type='''sample'''   ,						use_karras_sigmas=UpperCamelCase_   ,						clip_sample=UpperCamelCase_   ,						clip_sample_range=1.0   ,						)
									lowercase__						=				{
									    '''prior''': prior,
									    '''image_encoder''': image_encoder,
									    '''image_processor''': image_processor,
									    '''renderer''': renderer,
									    '''scheduler''': scheduler,
									}

									return components

							def 				lowerCamelCase_			(							self:     Optional[Any]   ,						UpperCamelCase_:     Tuple   ,						UpperCamelCase_:     Optional[int]=0							)      ->       Tuple:





									"""simple docstring"""







									lowercase__						=				floats_tensor((1, 3, 64, 64)   ,						rng=random.Random(UpperCamelCase_							)							).to(UpperCamelCase_							)

									if str(UpperCamelCase_							).startswith('''mps'''							):
											lowercase__						=				torch.manual_seed(UpperCamelCase_							)
									else:
											lowercase__						=				torch.Generator(device=UpperCamelCase_							).manual_seed(UpperCamelCase_							)
									lowercase__						=				{
									    '''image''': input_image,
									    '''generator''': generator,
									    '''num_inference_steps''': 1,
									    '''frame_size''': 32,
									    '''output_type''': '''np''',
									}
									return inputs

							def 				lowerCamelCase_			(							self:     int							)      ->       str:





									"""simple docstring"""







									lowercase__						=				'''cpu'''

									lowercase__						=				self.get_dummy_components()

									lowercase__						=				self.pipeline_class(**UpperCamelCase_							)
									lowercase__						=				pipe.to(UpperCamelCase_							)

									pipe.set_progress_bar_config(disable=UpperCamelCase_							)

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

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

									lowercase__						=				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:     int							)      ->       int:





									"""simple docstring"""







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

							def 				lowerCamelCase_			(							self:     List[str]							)      ->       List[Any]:





									"""simple docstring"""







									lowercase__						=				torch_device == '''cpu'''
									lowercase__						=				True
									self._test_inference_batch_single_identical(
									    batch_size=2   ,						test_max_difference=UpperCamelCase_   ,						relax_max_difference=UpperCamelCase_   ,						)


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





									"""simple docstring"""







									lowercase__						=				self.get_dummy_components()
									lowercase__						=				self.pipeline_class(**UpperCamelCase_							)
									lowercase__						=				pipe.to(UpperCamelCase_							)
									pipe.set_progress_bar_config(disable=UpperCamelCase_							)

									lowercase__						=				1
									lowercase__						=				2

									lowercase__						=				self.get_dummy_inputs(UpperCamelCase_							)

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

									lowercase__						=				pipe(**UpperCamelCase_   ,						num_images_per_prompt=UpperCamelCase_							)[0]

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


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

							def 				lowerCamelCase_			(							self:     Optional[int]							)      ->       Optional[Any]:





									"""simple docstring"""







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


							def 				lowerCamelCase_			(							self:     str							)      ->       str:





									"""simple docstring"""







									lowercase__						=				load_image(
									    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main''' '''/shap_e/corgi.png'''							)
									lowercase__						=				load_numpy(
									    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
									    '''/shap_e/test_shap_e_img2img_out.npy'''							)
									lowercase__						=				ShapEImgaImgPipeline.from_pretrained('''openai/shap-e-img2img'''							)
									lowercase__						=				pipe.to(UpperCamelCase_							)
									pipe.set_progress_bar_config(disable=UpperCamelCase_							)

									lowercase__						=				torch.Generator(device=UpperCamelCase_							).manual_seed(0							)

									lowercase__						=				pipe(
									    UpperCamelCase_   ,						generator=UpperCamelCase_   ,						guidance_scale=3.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_							)


