

"""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    snake_case     (    lowerCAmelCase_,							unittest.TestCase  ):
				SCREAMING_SNAKE_CASE_       :				int       				=  ShapEImgaImgPipeline
				SCREAMING_SNAKE_CASE_       :				Union[str, Any]       				=  ["""image"""]
				SCREAMING_SNAKE_CASE_       :				List[Any]       				=  ["""image"""]
				SCREAMING_SNAKE_CASE_       :				Dict       				=  [
				    """num_images_per_prompt""",
				    """num_inference_steps""",
				    """generator""",
				    """latents""",
				    """guidance_scale""",
				    """frame_size""",
				    """output_type""",
				    """return_dict""",
				]
				SCREAMING_SNAKE_CASE_       :				Optional[int]       				=  False

				@property
				def 							lowercase_					(     self							:    Tuple)->       Optional[Any]:



						'''simple docstring'''


						return 3_2

				@property
				def 							lowercase_					(     self							:    Optional[Any])->       Any:



						'''simple docstring'''


						return 3_2

				@property
				def 							lowercase_					(     self							:    List[Any])->       Optional[Any]:



						'''simple docstring'''


						return self.time_input_dim * 4

				@property
				def 							lowercase_					(     self							:    Optional[Any])->       Optional[Any]:



						'''simple docstring'''


						return 8

				@property
				def 							lowercase_					(     self							:    Any)->       Optional[Any]:



						'''simple docstring'''


						torch.manual_seed(0)
						__lowerCAmelCase:	int   				=      CLIPVisionConfig(
						    hidden_size=self.text_embedder_hidden_size					,  image_size=6_4					,  projection_dim=self.text_embedder_hidden_size					,  intermediate_size=3_7					,  num_attention_heads=4					,  num_channels=3					,  num_hidden_layers=5					,  patch_size=1					,  )

						__lowerCAmelCase:	Optional[Any]   				=      CLIPVisionModel(__SCREAMING_SNAKE_CASE)
						return model

				@property
				def 							lowercase_					(     self							:    str)->       Optional[Any]:



						'''simple docstring'''


						__lowerCAmelCase:	Optional[int]   				=      CLIPImageProcessor(
						    crop_size=2_2_4					,  do_center_crop=__SCREAMING_SNAKE_CASE					,  do_normalize=__SCREAMING_SNAKE_CASE					,  do_resize=__SCREAMING_SNAKE_CASE					,  image_mean=[0.48145466, 0.4578275, 0.40821073]					,  image_std=[0.26862954, 0.26130258, 0.27577711]					,  resample=3					,  size=2_2_4					,  )

						return image_processor

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



						'''simple docstring'''


						torch.manual_seed(0)

						__lowerCAmelCase:	Tuple   				=      {
						    "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",
						    "embedding_proj_norm_type": "layer",
						    "encoder_hid_proj_type": None,
						    "added_emb_type": None,
						}

						__lowerCAmelCase:	Optional[int]   				=      PriorTransformer(**__SCREAMING_SNAKE_CASE)
						return model

				@property
				def 							lowercase_					(     self							:    Tuple)->       Optional[int]:



						'''simple docstring'''


						torch.manual_seed(0)

						__lowerCAmelCase:	str   				=      {
						    "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:	Dict   				=      ShapERenderer(**__SCREAMING_SNAKE_CASE)
						return model

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



						'''simple docstring'''


						__lowerCAmelCase:	List[str]   				=      self.dummy_prior
						__lowerCAmelCase:	Dict   				=      self.dummy_image_encoder
						__lowerCAmelCase:	Tuple   				=      self.dummy_image_processor
						__lowerCAmelCase:	int   				=      self.dummy_renderer

						__lowerCAmelCase:	List[str]   				=      HeunDiscreteScheduler(
						    beta_schedule="exp"					,  num_train_timesteps=1_0_2_4					,  prediction_type="sample"					,  use_karras_sigmas=__SCREAMING_SNAKE_CASE					,  clip_sample=__SCREAMING_SNAKE_CASE					,  clip_sample_range=1.0					,  )
						__lowerCAmelCase:	Dict   				=      {
						    "prior": prior,
						    "image_encoder": image_encoder,
						    "image_processor": image_processor,
						    "renderer": renderer,
						    "scheduler": scheduler,
						}

						return components

				def 							lowercase_					(     self							:    Union[str, Any]					,  UpperCamelCase__							:    List[str]					,  UpperCamelCase__							:    int=0)->       Union[str, Any]:



						'''simple docstring'''


						__lowerCAmelCase:	str   				=      floats_tensor((1, 3, 6_4, 6_4)					,  rng=random.Random(__SCREAMING_SNAKE_CASE)).to(__SCREAMING_SNAKE_CASE)

						if str(__SCREAMING_SNAKE_CASE).startswith("mps"):
								__lowerCAmelCase:	Tuple   				=      torch.manual_seed(__SCREAMING_SNAKE_CASE)
						else:
								__lowerCAmelCase:	Tuple   				=      torch.Generator(device=__SCREAMING_SNAKE_CASE).manual_seed(__SCREAMING_SNAKE_CASE)
						__lowerCAmelCase:	List[str]   				=      {
						    "image": input_image,
						    "generator": generator,
						    "num_inference_steps": 1,
						    "frame_size": 3_2,
						    "output_type": "np",
						}
						return inputs

				def 							lowercase_					(     self							:    List[Any])->       Optional[Any]:



						'''simple docstring'''


						__lowerCAmelCase:	Dict   				=      "cpu"

						__lowerCAmelCase:	Optional[int]   				=      self.get_dummy_components()

						__lowerCAmelCase:	Any   				=      self.pipeline_class(**__SCREAMING_SNAKE_CASE)
						__lowerCAmelCase:	Dict   				=      pipe.to(__SCREAMING_SNAKE_CASE)

						pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE)

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

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

						__lowerCAmelCase:	List[str]   				=      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 							lowercase_					(     self							:    List[str])->       Optional[Any]:



						'''simple docstring'''


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

				def 							lowercase_					(     self							:    Tuple)->       Any:



						'''simple docstring'''


						__lowerCAmelCase:	List[Any]   				=      torch_device == "cpu"
						__lowerCAmelCase:	List[str]   				=      True
						self._test_inference_batch_single_identical(
						    batch_size=2					,  test_max_difference=__SCREAMING_SNAKE_CASE					,  relax_max_difference=__SCREAMING_SNAKE_CASE					,  )




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



						'''simple docstring'''


						__lowerCAmelCase:	List[Any]   				=      self.get_dummy_components()
						__lowerCAmelCase:	Tuple   				=      self.pipeline_class(**__SCREAMING_SNAKE_CASE)
						__lowerCAmelCase:	int   				=      pipe.to(__SCREAMING_SNAKE_CASE)
						pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE)

						__lowerCAmelCase:	List[str]   				=      1
						__lowerCAmelCase:	Dict   				=      2

						__lowerCAmelCase:	Union[str, Any]   				=      self.get_dummy_inputs(__SCREAMING_SNAKE_CASE)

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

						__lowerCAmelCase:	Dict   				=      pipe(**__SCREAMING_SNAKE_CASE					,  num_images_per_prompt=__SCREAMING_SNAKE_CASE)[0]

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





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

				def 							lowercase_					(     self							:    Union[str, Any])->       List[str]:



						'''simple docstring'''


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




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



						'''simple docstring'''


						__lowerCAmelCase:	List[Any]   				=      load_image(
						    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" "/shap_e/corgi.png")
						__lowerCAmelCase:	Any   				=      load_numpy(
						    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
						    "/shap_e/test_shap_e_img2img_out.npy")
						__lowerCAmelCase:	Optional[Any]   				=      ShapEImgaImgPipeline.from_pretrained("openai/shap-e-img2img")
						__lowerCAmelCase:	Tuple   				=      pipe.to(__SCREAMING_SNAKE_CASE)
						pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE)

						__lowerCAmelCase:	List[Any]   				=      torch.Generator(device=__SCREAMING_SNAKE_CASE).manual_seed(0)

						__lowerCAmelCase:	int   				=      pipe(
						    __SCREAMING_SNAKE_CASE					,  generator=__SCREAMING_SNAKE_CASE					,  guidance_scale=3.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(__SCREAMING_SNAKE_CASE					,  __SCREAMING_SNAKE_CASE)





"""simple docstring"""


from math import ceil






def    a__   (   __SCREAMING_SNAKE_CASE ,			__SCREAMING_SNAKE_CASE			)     ->     str:
		__lowerCAmelCase:	Tuple   				=      list(range(0 ,			__SCREAMING_SNAKE_CASE			)			)

		__lowerCAmelCase:	Optional[Any]   				=      [item for sublist in list(device_map.values()			) for item in sublist]

		# Duplicate check
		__lowerCAmelCase:	List[Any]   				=      []
		for i in device_map_blocks:
				if device_map_blocks.count(__SCREAMING_SNAKE_CASE			) > 1 and i not in duplicate_blocks:
						duplicate_blocks.append(__SCREAMING_SNAKE_CASE			)
    # Missing blocks
		__lowerCAmelCase:	Optional[Any]   				=      [i for i in blocks if i not in device_map_blocks]
		__lowerCAmelCase:	List[Any]   				=      [i for i in device_map_blocks if i not in blocks]

		if len(__SCREAMING_SNAKE_CASE			) != 0:
				raise ValueError(
				    "Duplicate attention blocks specified in device_map. Attention blocks must be specified to one device."
				    " These attention blocks were specified more than once: " + str(__SCREAMING_SNAKE_CASE			)			)
		if len(__SCREAMING_SNAKE_CASE			) != 0:
				raise ValueError(
				    "There are attention blocks for this model that are not specified in the device_map. Add these attention "
				    "blocks to a device on the device_map: " + str(__SCREAMING_SNAKE_CASE			)			)
		if len(__SCREAMING_SNAKE_CASE			) != 0:
				raise ValueError(
				    "The device_map contains more attention blocks than this model has. Remove these from the device_map:"
				    + str(__SCREAMING_SNAKE_CASE			)			)






def    a__   (   __SCREAMING_SNAKE_CASE ,			__SCREAMING_SNAKE_CASE			)     ->     str:
		__lowerCAmelCase:	List[Any]   				=      list(range(__SCREAMING_SNAKE_CASE			)			)
		__lowerCAmelCase:	Dict   				=      int(ceil(n_layers / len(__SCREAMING_SNAKE_CASE			)			)			)
		__lowerCAmelCase:	Union[str, Any]   				=      [layers[i : i + n_blocks] for i in range(0 ,			__SCREAMING_SNAKE_CASE ,			__SCREAMING_SNAKE_CASE			)]

		return dict(zip(__SCREAMING_SNAKE_CASE ,			__SCREAMING_SNAKE_CASE			)			)


