


'''simple docstring'''
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 					__lowerCAmelCase   (   lowercase				,			unittest.TestCase						):

						"""simple docstring"""






						_UpperCAmelCase		:       Optional[int]   						=ShapEImgaImgPipeline
						_UpperCAmelCase		:       Union[str, Any]   						=["image"]
						_UpperCAmelCase		:       Tuple   						=["image"]
						_UpperCAmelCase		:       Tuple   						=[
						    "num_images_per_prompt",
						    "num_inference_steps",
						    "generator",
						    "latents",
						    "guidance_scale",
						    "frame_size",
						    "output_type",
						    "return_dict",
						]
						_UpperCAmelCase		:       str   						=False






						@property
						def 				_UpperCAmelCase       (  self							:      Optional[int]	):
													return 32






						@property
						def 				_UpperCAmelCase       (  self							:      List[str]	):
													return 32






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






						@property
						def 				_UpperCAmelCase       (  self							:      List[str]	):
													return 8






						@property
						def 				_UpperCAmelCase       (  self							:      int	):
													torch.manual_seed(0	)
													A_			    =		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 ,  )

													A_			    =		CLIPVisionModel(lowerCAmelCase	)
													return model






						@property
						def 				_UpperCAmelCase       (  self							:      Dict	):
													A_			    =		CLIPImageProcessor(
													    crop_size=2_24 ,  do_center_crop=lowerCAmelCase ,  do_normalize=lowerCAmelCase ,  do_resize=lowerCAmelCase ,  image_mean=[0.4_8_1_4_5_4_6_6, 0.4_5_7_8_2_7_5, 0.4_0_8_2_1_0_7_3] ,  image_std=[0.2_6_8_6_2_9_5_4, 0.2_6_1_3_0_2_5_8, 0.2_7_5_7_7_7_1_1] ,  resample=3 ,  size=2_24 ,  )

													return image_processor






						@property
						def 				_UpperCAmelCase       (  self							:      Optional[Any]	):
													torch.manual_seed(0	)

													A_			    =		{
													    "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,
													}

													A_			    =		PriorTransformer(**lowerCAmelCase	)
													return model






						@property
						def 				_UpperCAmelCase       (  self							:      str	):
													torch.manual_seed(0	)

													A_			    =		{
													    "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_			    =		ShapERenderer(**lowerCAmelCase	)
													return model






						def 				_UpperCAmelCase       (  self							:      Optional[Any]	):
													A_			    =		self.dummy_prior
													A_			    =		self.dummy_image_encoder
													A_			    =		self.dummy_image_processor
													A_			    =		self.dummy_renderer

													A_			    =		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_			    =		{
													    "prior": prior,
													    "image_encoder": image_encoder,
													    "image_processor": image_processor,
													    "renderer": renderer,
													    "scheduler": scheduler,
													}

													return components






						def 				_UpperCAmelCase       (  self							:      str ,  lowerCAmelCase							:      Any ,  lowerCAmelCase							:      Dict=0	):
													A_			    =		floats_tensor((1, 3, 64, 64) ,  rng=random.Random(lowerCAmelCase	)	).to(lowerCAmelCase	)

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






						def 				_UpperCAmelCase       (  self							:      int	):
													A_			    =		"cpu"

													A_			    =		self.get_dummy_components()

													A_			    =		self.pipeline_class(**lowerCAmelCase	)
													A_			    =		pipe.to(lowerCAmelCase	)

													pipe.set_progress_bar_config(disable=lowerCAmelCase	)

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

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

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






						def 				_UpperCAmelCase       (  self							:      Tuple	):
													A_			    =		torch_device == "cpu"
													A_			    =		True
													self._test_inference_batch_single_identical(
													    batch_size=2 ,  test_max_difference=lowerCAmelCase ,  relax_max_difference=lowerCAmelCase ,  )






						def 				_UpperCAmelCase       (  self							:      Optional[Any]	):
													A_			    =		self.get_dummy_components()
													A_			    =		self.pipeline_class(**lowerCAmelCase	)
													A_			    =		pipe.to(lowerCAmelCase	)
													pipe.set_progress_bar_config(disable=lowerCAmelCase	)

													A_			    =		1
													A_			    =		2

													A_			    =		self.get_dummy_inputs(lowerCAmelCase	)

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

													A_			    =		pipe(**lowerCAmelCase ,  num_images_per_prompt=lowerCAmelCase	)[0]

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






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

						"""simple docstring"""






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






						def 				_UpperCAmelCase       (  self							:      Tuple	):
													A_			    =		load_image(
													    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" "/shap_e/corgi.png"	)
													A_			    =		load_numpy(
													    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
													    "/shap_e/test_shap_e_img2img_out.npy"	)
													A_			    =		ShapEImgaImgPipeline.from_pretrained("openai/shap-e-img2img"	)
													A_			    =		pipe.to(lowerCAmelCase	)
													pipe.set_progress_bar_config(disable=lowerCAmelCase	)

													A_			    =		torch.Generator(device=lowerCAmelCase	).manual_seed(0	)

													A_			    =		pipe(
													    lowerCAmelCase ,  generator=lowerCAmelCase ,  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(lowerCAmelCase ,  lowerCAmelCase	)









'''simple docstring'''
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


__SCREAMING_SNAKE_CASE					:				List[str]									=		logging.get_logger(__name__)






class 					__lowerCAmelCase   (   lowercase						):

						"""simple docstring"""






						_UpperCAmelCase		:       Union[str, Any]   						=["pixel_values"]






						def __init__(  self							:      Optional[int] ,  lowerCAmelCase							:      bool = True ,  lowerCAmelCase							:      Dict[str, int] = None ,  lowerCAmelCase							:      PILImageResampling = PILImageResampling.BICUBIC ,  lowerCAmelCase							:      bool = True ,  lowerCAmelCase							:      Union[int, float] = 1 / 2_55 ,  lowerCAmelCase							:      bool = True ,  lowerCAmelCase							:      Optional[Union[float, List[float]]] = None ,  lowerCAmelCase							:      Optional[Union[float, List[float]]] = None ,  lowerCAmelCase							:      bool = True ,  **lowerCAmelCase							:      List[Any] ,  ):
													super().__init__(**lowerCAmelCase	)
													A_			    =		size if size is not None else {"height": 3_84, "width": 3_84}
													A_			    =		get_size_dict(lowerCAmelCase ,  default_to_square=lowerCAmelCase	)

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






						def 				_UpperCAmelCase       (  self							:      Optional[Any] ,  lowerCAmelCase							:      np.ndarray ,  lowerCAmelCase							:      Dict[str, int] ,  lowerCAmelCase							:      PILImageResampling = PILImageResampling.BICUBIC ,  lowerCAmelCase							:      Optional[Union[str, ChannelDimension]] = None ,  **lowerCAmelCase							:      Tuple ,  ):
													A_			    =		get_size_dict(lowerCAmelCase ,  default_to_square=lowerCAmelCase	)
													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()}"	)
													A_			    =		(size["height"], size["width"])
													return resize(lowerCAmelCase ,  size=lowerCAmelCase ,  resample=lowerCAmelCase ,  data_format=lowerCAmelCase ,  **lowerCAmelCase	)






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






						def 				_UpperCAmelCase       (  self							:      Optional[int] ,  lowerCAmelCase							:      np.ndarray ,  lowerCAmelCase							:      Union[float, List[float]] ,  lowerCAmelCase							:      Union[float, List[float]] ,  lowerCAmelCase							:      Optional[Union[str, ChannelDimension]] = None ,  **lowerCAmelCase							:      Dict ,  ):
													return normalize(lowerCAmelCase ,  mean=lowerCAmelCase ,  std=lowerCAmelCase ,  data_format=lowerCAmelCase ,  **lowerCAmelCase	)






						def 				_UpperCAmelCase       (  self							:      Tuple ,  lowerCAmelCase							:      ImageInput ,  lowerCAmelCase							:      Optional[bool] = None ,  lowerCAmelCase							:      Optional[Dict[str, int]] = None ,  lowerCAmelCase							:      PILImageResampling = None ,  lowerCAmelCase							:      Optional[bool] = None ,  lowerCAmelCase							:      Optional[float] = None ,  lowerCAmelCase							:      Optional[bool] = None ,  lowerCAmelCase							:      Optional[Union[float, List[float]]] = None ,  lowerCAmelCase							:      Optional[Union[float, List[float]]] = None ,  lowerCAmelCase							:      Optional[Union[str, TensorType]] = None ,  lowerCAmelCase							:      bool = None ,  lowerCAmelCase							:      ChannelDimension = ChannelDimension.FIRST ,  **lowerCAmelCase							:      List[str] ,  ):
													A_			    =		do_resize if do_resize is not None else self.do_resize
													A_			    =		resample if resample is not None else self.resample
													A_			    =		do_rescale if do_rescale is not None else self.do_rescale
													A_			    =		rescale_factor if rescale_factor is not None else self.rescale_factor
													A_			    =		do_normalize if do_normalize is not None else self.do_normalize
													A_			    =		image_mean if image_mean is not None else self.image_mean
													A_			    =		image_std if image_std is not None else self.image_std
													A_			    =		do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb

													A_			    =		size if size is not None else self.size
													A_			    =		get_size_dict(lowerCAmelCase ,  default_to_square=lowerCAmelCase	)

													A_			    =		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_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:
																				A_			    =		[convert_to_rgb(lowerCAmelCase	) for image in images]

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

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

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

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

													A_			    =		[to_channel_dimension_format(lowerCAmelCase ,  lowerCAmelCase	) for image in images]

													A_			    =		BatchFeature(data={"pixel_values": images} ,  tensor_type=lowerCAmelCase	)

													return encoded_outputs





