




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        UpperCamelCase__   (	__SCREAMING_SNAKE_CASE							,				unittest.TestCase   ):
					__SCREAMING_SNAKE_CASE       :  List[str]						=   ShapEImgaImgPipeline
					__SCREAMING_SNAKE_CASE       :  str						=   ['image']
					__SCREAMING_SNAKE_CASE       :  Any						=   ['image']
					__SCREAMING_SNAKE_CASE       :  Optional[Any]						=   [
					    'num_images_per_prompt',
					    'num_inference_steps',
					    'generator',
					    'latents',
					    'guidance_scale',
					    'frame_size',
					    'output_type',
					    'return_dict',
					]
					__SCREAMING_SNAKE_CASE       :  List[str]						=   False

					@property
					def   UpperCAmelCase__       (     self       :   str       ):






										'''simple docstring'''






										return 32

					@property
					def   UpperCAmelCase__       (     self       :   str       ):






										'''simple docstring'''






										return 32

					@property
					def   UpperCAmelCase__       (     self       :   List[Any]       ):






										'''simple docstring'''






										return self.time_input_dim * 4

					@property
					def   UpperCAmelCase__       (     self       :   List[str]       ):






										'''simple docstring'''






										return 8

					@property
					def   UpperCAmelCase__       (     self       :   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(_a       )
										return model

					@property
					def   UpperCAmelCase__       (     self       :   str       ):






										'''simple docstring'''






										lowercase_  							=			CLIPImageProcessor(
										    crop_size=224							,					do_center_crop=_a							,					do_normalize=_a							,					do_resize=_a							,					image_mean=[0.48_145_466, 0.4_578_275, 0.40_821_073]							,					image_std=[0.26_862_954, 0.26_130_258, 0.27_577_711]							,					resample=3							,					size=224							,					)

										return image_processor

					@property
					def   UpperCAmelCase__       (     self       :   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(**_a       )
										return model

					@property
					def   UpperCAmelCase__       (     self       :   Union[str, 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(**_a       )
										return model

					def   UpperCAmelCase__       (     self       :   List[str]       ):






										'''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=_a							,					clip_sample=_a							,					clip_sample_range=1.0							,					)
										lowercase_  							=			{
										    """prior""": prior,
										    """image_encoder""": image_encoder,
										    """image_processor""": image_processor,
										    """renderer""": renderer,
										    """scheduler""": scheduler,
										}

										return components

					def   UpperCAmelCase__       (     self       :   int							,					UpperCamelCase__       :   Dict							,					UpperCamelCase__       :   List[str]=0       ):






										'''simple docstring'''






										lowercase_  							=			floats_tensor((1, 3, 64, 64)							,					rng=random.Random(_a       )       ).to(_a       )

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

					def   UpperCAmelCase__       (     self       :   Optional[Any]       ):






										'''simple docstring'''






										lowercase_  							=			"""cpu"""

										lowercase_  							=			self.get_dummy_components()

										lowercase_  							=			self.pipeline_class(**_a       )
										lowercase_  							=			pipe.to(_a       )

										pipe.set_progress_bar_config(disable=_a       )

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

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

										lowercase_  							=			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   UpperCAmelCase__       (     self       :   Optional[int]       ):






										'''simple docstring'''






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

					def   UpperCAmelCase__       (     self       :   List[Any]       ):






										'''simple docstring'''






										lowercase_  							=			torch_device == """cpu"""
										lowercase_  							=			True
										self._test_inference_batch_single_identical(
										    batch_size=2							,					test_max_difference=_a							,					relax_max_difference=_a							,					)







					def   UpperCAmelCase__       (     self       :   Union[str, Any]       ):






										'''simple docstring'''






										lowercase_  							=			self.get_dummy_components()
										lowercase_  							=			self.pipeline_class(**_a       )
										lowercase_  							=			pipe.to(_a       )
										pipe.set_progress_bar_config(disable=_a       )

										lowercase_  							=			1
										lowercase_  							=			2

										lowercase_  							=			self.get_dummy_inputs(_a       )

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

										lowercase_  							=			pipe(**_a							,					num_images_per_prompt=_a       )[0]

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







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

					def   UpperCAmelCase__       (     self       :   Optional[Any]       ):






										'''simple docstring'''






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







					def   UpperCAmelCase__       (     self       :   List[Any]       ):






										'''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(_a       )
										pipe.set_progress_bar_config(disable=_a       )

										lowercase_  							=			torch.Generator(device=_a       ).manual_seed(0       )

										lowercase_  							=			pipe(
										    _a							,					generator=_a							,					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(_a							,					_a       )






def 			UpperCAmelCase_				(      UpperCAmelCase__  ):
					return " ".join(input_str.split()[::-1]  )


if __name__ == "__main__":
							import doctest

							doctest.testmod()


