






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     _SCREAMING_SNAKE_CASE       (			_a		,   unittest.TestCase      ):
	snake_case__							:     Tuple     	=  ShapEImgaImgPipeline
	snake_case__							:     Optional[Any]     	=  ["""image"""]
	snake_case__							:     Union[str, Any]     	=  ["""image"""]
	snake_case__							:     Optional[Any]     	=  [
	    """num_images_per_prompt""",
	    """num_inference_steps""",
	    """generator""",
	    """latents""",
	    """guidance_scale""",
	    """frame_size""",
	    """output_type""",
	    """return_dict""",
	]
	snake_case__							:     List[str]     	=  False





	@property
	def 						_A (		self				:					Any	):
		return 32





	@property
	def 						_A (		self				:					Any	):
		return 32





	@property
	def 						_A (		self				:					Optional[Any]	):
		return self.time_input_dim * 4





	@property
	def 						_A (		self				:					Union[str, Any]	):
		return 8





	@property
	def 						_A (		self				:					int	):
		torch.manual_seed(0	)
		UpperCamelCase     :Union[str, Any]          =	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 ,  )

		UpperCamelCase     :Optional[int]          =	CLIPVisionModel(__lowerCamelCase	)
		return model





	@property
	def 						_A (		self				:					str	):
		UpperCamelCase     :Optional[int]          =	CLIPImageProcessor(
		    crop_size=224 ,  do_center_crop=__lowerCamelCase ,  do_normalize=__lowerCamelCase ,  do_resize=__lowerCamelCase ,  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 						_A (		self				:					Tuple	):
		torch.manual_seed(0	)

		UpperCamelCase     :Dict          =	{
		    """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,
		}

		UpperCamelCase     :int          =	PriorTransformer(**__lowerCamelCase	)
		return model





	@property
	def 						_A (		self				:					Optional[int]	):
		torch.manual_seed(0	)

		UpperCamelCase     :str          =	{
		    """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,
		    ),
		}
		UpperCamelCase     :List[str]          =	ShapERenderer(**__lowerCamelCase	)
		return model





	def 						_A (		self				:					str	):
		UpperCamelCase     :int          =	self.dummy_prior
		UpperCamelCase     :Any          =	self.dummy_image_encoder
		UpperCamelCase     :Dict          =	self.dummy_image_processor
		UpperCamelCase     :List[Any]          =	self.dummy_renderer

		UpperCamelCase     :int          =	HeunDiscreteScheduler(
		    beta_schedule="""exp""" ,  num_train_timesteps=1_024 ,  prediction_type="""sample""" ,  use_karras_sigmas=__lowerCamelCase ,  clip_sample=__lowerCamelCase ,  clip_sample_range=1.0 ,  )
		UpperCamelCase     :Optional[Any]          =	{
		    """prior""": prior,
		    """image_encoder""": image_encoder,
		    """image_processor""": image_processor,
		    """renderer""": renderer,
		    """scheduler""": scheduler,
		}

		return components





	def 						_A (		self				:					int ,  __lowerCamelCase				:					int ,  __lowerCamelCase				:					Any=0	):
		UpperCamelCase     :Any          =	floats_tensor((1, 3, 64, 64) ,  rng=random.Random(__lowerCamelCase	)	).to(__lowerCamelCase	)

		if str(__lowerCamelCase	).startswith("""mps"""	):
			UpperCamelCase     :List[Any]          =	torch.manual_seed(__lowerCamelCase	)
		else:
			UpperCamelCase     :Optional[int]          =	torch.Generator(device=__lowerCamelCase	).manual_seed(__lowerCamelCase	)
		UpperCamelCase     :Optional[Any]          =	{
		    """image""": input_image,
		    """generator""": generator,
		    """num_inference_steps""": 1,
		    """frame_size""": 32,
		    """output_type""": """np""",
		}
		return inputs





	def 						_A (		self				:					List[str]	):
		UpperCamelCase     :Dict          =	"""cpu"""

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

		UpperCamelCase     :Optional[int]          =	self.pipeline_class(**__lowerCamelCase	)
		UpperCamelCase     :int          =	pipe.to(__lowerCamelCase	)

		pipe.set_progress_bar_config(disable=__lowerCamelCase	)

		UpperCamelCase     :Optional[Any]          =	pipe(**self.get_dummy_inputs(__lowerCamelCase	)	)
		UpperCamelCase     :Dict          =	output.images[0]
		UpperCamelCase     :List[Any]          =	image[0, -3:, -3:, -1]

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

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





	def 						_A (		self				:					List[Any]	):
		UpperCamelCase     :str          =	torch_device == """cpu"""
		UpperCamelCase     :int          =	True
		self._test_inference_batch_single_identical(
		    batch_size=2 ,  test_max_difference=__lowerCamelCase ,  relax_max_difference=__lowerCamelCase ,  )





	def 						_A (		self				:					List[Any]	):
		UpperCamelCase     :List[Any]          =	self.get_dummy_components()
		UpperCamelCase     :Optional[int]          =	self.pipeline_class(**__lowerCamelCase	)
		UpperCamelCase     :List[Any]          =	pipe.to(__lowerCamelCase	)
		pipe.set_progress_bar_config(disable=__lowerCamelCase	)

		UpperCamelCase     :Any          =	1
		UpperCamelCase     :int          =	2

		UpperCamelCase     :Union[str, Any]          =	self.get_dummy_inputs(__lowerCamelCase	)

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

		UpperCamelCase     :Optional[int]          =	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      ):





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





	def 						_A (		self				:					Any	):
		UpperCamelCase     :Optional[Any]          =	load_image(
		    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main""" """/shap_e/corgi.png"""	)
		UpperCamelCase     :Any          =	load_numpy(
		    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
		    """/shap_e/test_shap_e_img2img_out.npy"""	)
		UpperCamelCase     :Union[str, Any]          =	ShapEImgaImgPipeline.from_pretrained("""openai/shap-e-img2img"""	)
		UpperCamelCase     :List[str]          =	pipe.to(__lowerCamelCase	)
		pipe.set_progress_bar_config(disable=__lowerCamelCase	)

		UpperCamelCase     :Optional[Any]          =	torch.Generator(device=__lowerCamelCase	).manual_seed(0	)

		UpperCamelCase     :Optional[int]          =	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, Tuple, 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_torch_available, is_torch_tensor, logging


if is_torch_available():
			import torch


__A						   =			logging.get_logger(__name__)





class     snake_case	(				__snake_case  ):
				SCREAMING_SNAKE_CASE_   :      Dict         =  ["""pixel_values"""]

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



						'''simple docstring'''


						super().__init__(**UpperCamelCase__)
						__lowerCAmelCase:       int          =			size if size is not None else {"shortest_edge": 2_5_6}
						__lowerCAmelCase:       str          =			get_size_dict(UpperCamelCase__     ,  default_to_square=UpperCamelCase__)
						__lowerCAmelCase:       Any          =			crop_size if crop_size is not None else {"height": 2_2_4, "width": 2_2_4}
						__lowerCAmelCase:       Optional[Any]          =			get_size_dict(UpperCamelCase__     ,  param_name="crop_size")
						__lowerCAmelCase:       str          =			do_resize
						__lowerCAmelCase:       Any          =			size
						__lowerCAmelCase:       Dict          =			resample
						__lowerCAmelCase:       Tuple          =			do_center_crop
						__lowerCAmelCase:       str          =			crop_size
						__lowerCAmelCase:       List[Any]          =			do_rescale
						__lowerCAmelCase:       int          =			rescale_factor
						__lowerCAmelCase:       List[Any]          =			do_normalize
						__lowerCAmelCase:       Any          =			image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
						__lowerCAmelCase:       Optional[Any]          =			image_std if image_std is not None else IMAGENET_STANDARD_STD

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



						'''simple docstring'''


						__lowerCAmelCase:       int          =			get_size_dict(UpperCamelCase__     ,  default_to_square=UpperCamelCase__)
						if "shortest_edge" not in size:
								raise ValueError(f"The `size` parameter must contain the key `shortest_edge`. Got {size.keys()}")
						__lowerCAmelCase:       Optional[Any]          =			get_resize_output_image_size(UpperCamelCase__     ,  size=size["shortest_edge"]     ,  default_to_square=UpperCamelCase__)
						return resize(UpperCamelCase__     ,  size=UpperCamelCase__     ,  resample=UpperCamelCase__     ,  data_format=UpperCamelCase__     ,  **UpperCamelCase__)

				def 							lowercase_     (	self							:				Optional[int]     ,  UpperCamelCase__							:				np.ndarray     ,  UpperCamelCase__							:				Dict[str, int]     ,  UpperCamelCase__							:				Optional[Union[str, ChannelDimension]] = None     ,  **UpperCamelCase__							:				Tuple     ,  )->   np.ndarray:



						'''simple docstring'''


						__lowerCAmelCase:       Union[str, Any]          =			get_size_dict(UpperCamelCase__)
						if "height" not in size or "width" not in size:
								raise ValueError(f"The `size` parameter must contain the keys `height` and `width`. Got {size.keys()}")
						return center_crop(UpperCamelCase__     ,  size=(size["height"], size["width"])     ,  data_format=UpperCamelCase__     ,  **UpperCamelCase__)

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



						'''simple docstring'''


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

				def 							lowercase_     (	self							:				Dict     ,  UpperCamelCase__							:				np.ndarray     ,  UpperCamelCase__							:				Union[float, List[float]]     ,  UpperCamelCase__							:				Union[float, List[float]]     ,  UpperCamelCase__							:				Optional[Union[str, ChannelDimension]] = None     ,  **UpperCamelCase__							:				Optional[Any]     ,  )->   np.ndarray:



						'''simple docstring'''


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

				def 							lowercase_     (	self							:				Union[str, Any]     ,  UpperCamelCase__							:				ImageInput     ,  UpperCamelCase__							:				Optional[bool] = None     ,  UpperCamelCase__							:				Dict[str, int] = None     ,  UpperCamelCase__							:				PILImageResampling = None     ,  UpperCamelCase__							:				bool = None     ,  UpperCamelCase__							:				Dict[str, int] = 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__							:				Union[str, ChannelDimension] = ChannelDimension.FIRST     ,  **UpperCamelCase__							:				Optional[Any]     ,  )->   Dict:



						'''simple docstring'''


						__lowerCAmelCase:       Any          =			do_resize if do_resize is not None else self.do_resize
						__lowerCAmelCase:       str          =			size if size is not None else self.size
						__lowerCAmelCase:       Tuple          =			get_size_dict(UpperCamelCase__     ,  default_to_square=UpperCamelCase__)
						__lowerCAmelCase:       List[str]          =			resample if resample is not None else self.resample
						__lowerCAmelCase:       str          =			do_center_crop if do_center_crop is not None else self.do_center_crop
						__lowerCAmelCase:       Tuple          =			crop_size if crop_size is not None else self.crop_size
						__lowerCAmelCase:       List[Any]          =			get_size_dict(UpperCamelCase__     ,  param_name="crop_size")
						__lowerCAmelCase:       List[Any]          =			do_rescale if do_rescale is not None else self.do_rescale
						__lowerCAmelCase:       Optional[int]          =			rescale_factor if rescale_factor is not None else self.rescale_factor
						__lowerCAmelCase:       Optional[int]          =			do_normalize if do_normalize is not None else self.do_normalize
						__lowerCAmelCase:       Union[str, Any]          =			image_mean if image_mean is not None else self.image_mean
						__lowerCAmelCase:       Tuple          =			image_std if image_std is not None else self.image_std

						__lowerCAmelCase:       Union[str, Any]          =			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:
								raise ValueError("Size must be specified if do_resize is True.")

						if do_center_crop and crop_size is None:
								raise ValueError("Crop size must be specified if do_center_crop 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.")

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

						if do_resize:
								__lowerCAmelCase:       Union[str, Any]          =			[self.resize(image=UpperCamelCase__     ,  size=UpperCamelCase__     ,  resample=UpperCamelCase__) for image in images]

						if do_center_crop:
								__lowerCAmelCase:       Optional[Any]          =			[self.center_crop(image=UpperCamelCase__     ,  size=UpperCamelCase__) for image in images]

						if do_rescale:
								__lowerCAmelCase:       Optional[Any]          =			[self.rescale(image=UpperCamelCase__     ,  scale=UpperCamelCase__) for image in images]

						if do_normalize:
								__lowerCAmelCase:       List[str]          =			[self.normalize(image=UpperCamelCase__     ,  mean=UpperCamelCase__     ,  std=UpperCamelCase__) for image in images]

						__lowerCAmelCase:       Optional[Any]          =			[to_channel_dimension_format(UpperCamelCase__     ,  UpperCamelCase__) for image in images]

						__lowerCAmelCase:       List[str]          =			{"pixel_values": images}
						return BatchFeature(data=UpperCamelCase__     ,  tensor_type=UpperCamelCase__)




				def 							lowercase_     (	self							:				int     ,  UpperCamelCase__							:				List[str]     ,  UpperCamelCase__							:				List[Tuple] = None)->   Dict:



						'''simple docstring'''


						__lowerCAmelCase:       Union[str, Any]          =			outputs.logits

						# Resize logits and compute semantic segmentation maps
						if target_sizes is not None:
								if len(UpperCamelCase__) != len(UpperCamelCase__):
										raise ValueError(
										    "Make sure that you pass in as many target sizes as the batch dimension of the logits")

								if is_torch_tensor(UpperCamelCase__):
										__lowerCAmelCase:       Optional[int]          =			target_sizes.numpy()

								__lowerCAmelCase:       List[Any]          =			[]

								for idx in range(len(UpperCamelCase__)):
										__lowerCAmelCase:       List[str]          =			torch.nn.functional.interpolate(
										    logits[idx].unsqueeze(dim=0)     ,  size=target_sizes[idx]     ,  mode="bilinear"     ,  align_corners=UpperCamelCase__)
										__lowerCAmelCase:       Optional[Any]          =			resized_logits[0].argmax(dim=0)
										semantic_segmentation.append(UpperCamelCase__)
						else:
								__lowerCAmelCase:       Tuple          =			logits.argmax(dim=1)
								__lowerCAmelCase:       Optional[Any]          =			[semantic_segmentation[i] for i in range(semantic_segmentation.shape[0])]

						return semantic_segmentation


