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  (			lowercase_       ,				unittest.TestCase   ):

	'''simple docstring'''



	lowerCamelCase_       :     Dict          =			ShapEImgaImgPipeline
	lowerCamelCase_       :     str          =			["""image"""]
	lowerCamelCase_       :     List[Any]          =			["""image"""]
	lowerCamelCase_       :     List[str]          =			[
	    """num_images_per_prompt""",
	    """num_inference_steps""",
	    """generator""",
	    """latents""",
	    """guidance_scale""",
	    """frame_size""",
	    """output_type""",
	    """return_dict""",
	]
	lowerCamelCase_       :     Optional[int]          =			False

	@property
	def 	__UpperCAmelCase(  self       ):
			return 32

	@property
	def 	__UpperCAmelCase(  self       ):
			return 32

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

	@property
	def 	__UpperCAmelCase(  self       ):
			return 8

	@property
	def 	__UpperCAmelCase(  self       ):
			torch.manual_seed(0       )
			__A				:       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				,					)

			__A				:       str			    = CLIPVisionModel(__UpperCAmelCase       )
			return model

	@property
	def 	__UpperCAmelCase(  self       ):
			__A				:       List[str]			    = CLIPImageProcessor(
			    crop_size=224				,					do_center_crop=__UpperCAmelCase				,					do_normalize=__UpperCAmelCase				,					do_resize=__UpperCAmelCase				,					image_mean=[0.48_14_54_66, 0.4_57_82_75, 0.40_82_10_73]				,					image_std=[0.26_86_29_54, 0.26_13_02_58, 0.27_57_77_11]				,					resample=3				,					size=224				,					)

			return image_processor

	@property
	def 	__UpperCAmelCase(  self       ):
			torch.manual_seed(0       )

			__A				:       Any			    = {
			    "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				:       Dict			    = PriorTransformer(**__UpperCAmelCase       )
			return model

	@property
	def 	__UpperCAmelCase(  self       ):
			torch.manual_seed(0       )

			__A				:       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,
			    ),
			}
			__A				:       str			    = ShapERenderer(**__UpperCAmelCase       )
			return model

	def 	__UpperCAmelCase(  self       ):
			__A				:       int			    = self.dummy_prior
			__A				:       str			    = self.dummy_image_encoder
			__A				:       List[Any]			    = self.dummy_image_processor
			__A				:       Optional[Any]			    = self.dummy_renderer

			__A				:       List[Any]			    = HeunDiscreteScheduler(
			    beta_schedule="exp"				,					num_train_timesteps=1_024				,					prediction_type="sample"				,					use_karras_sigmas=__UpperCAmelCase				,					clip_sample=__UpperCAmelCase				,					clip_sample_range=1.0				,					)
			__A				:       str			    = {
			    "prior": prior,
			    "image_encoder": image_encoder,
			    "image_processor": image_processor,
			    "renderer": renderer,
			    "scheduler": scheduler,
			}

			return components

	def 	__UpperCAmelCase(  self				,					__UpperCAmelCase				,					__UpperCAmelCase=0       ):
			__A				:       Tuple			    = floats_tensor((1, 3, 64, 64)				,					rng=random.Random(__UpperCAmelCase       )       ).to(__UpperCAmelCase       )

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

	def 	__UpperCAmelCase(  self       ):
			__A				:       Dict			    = "cpu"

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

			__A				:       Union[str, Any]			    = self.pipeline_class(**__UpperCAmelCase       )
			__A				:       Any			    = pipe.to(__UpperCAmelCase       )

			pipe.set_progress_bar_config(disable=__UpperCAmelCase       )

			__A				:       Any			    = pipe(**self.get_dummy_inputs(__UpperCAmelCase       )       )
			__A				:       int			    = output.images[0]
			__A				:       int			    = image[0, -3:, -3:, -1]

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

			__A				:       List[Any]			    = np.array(
			    [
			        0.00_03_92_16,
			        0.00_03_92_16,
			        0.00_03_92_16,
			        0.00_03_92_16,
			        0.00_03_92_16,
			        0.00_03_92_16,
			        0.00_03_92_16,
			        0.00_03_92_16,
			        0.00_03_92_16,
			    ]       )

			assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1e-2

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

	def 	__UpperCAmelCase(  self       ):
			__A				:       Any			    = torch_device == "cpu"
			__A				:       Any			    = True
			self._test_inference_batch_single_identical(
			    batch_size=2				,					test_max_difference=__UpperCAmelCase				,					relax_max_difference=__UpperCAmelCase				,					)







	def 	__UpperCAmelCase(  self       ):
			__A				:       Union[str, Any]			    = self.get_dummy_components()
			__A				:       int			    = self.pipeline_class(**__UpperCAmelCase       )
			__A				:       List[Any]			    = pipe.to(__UpperCAmelCase       )
			pipe.set_progress_bar_config(disable=__UpperCAmelCase       )

			__A				:       Tuple			    = 1
			__A				:       str			    = 2

			__A				:       List[Any]			    = self.get_dummy_inputs(__UpperCAmelCase       )

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

			__A				:       Optional[Any]			    = 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   ):

	'''simple docstring'''



	def 	__UpperCAmelCase(  self       ):
			super().tearDown()
			gc.collect()
			torch.cuda.empty_cache()







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

			__A				:       int			    = torch.Generator(device=__UpperCAmelCase       ).manual_seed(0       )

			__A				:       List[str]			    = 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       )









'''simple docstring'''

import unittest

from transformers import EsmConfig, is_torch_available
from transformers.testing_utils import TestCasePlus, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
 import torch

 from transformers import EsmForMaskedLM, EsmForSequenceClassification, EsmForTokenClassification, EsmModel
 from transformers.models.esm.modeling_esm import (
     ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
     EsmEmbeddings,
     create_position_ids_from_input_ids,
 )
class     __lowerCAmelCase		:





      '''simple docstring'''






      def __init__(self       :       Optional[int]  ,		UpperCamelCase       :       Tuple  ,		UpperCamelCase       :       str=13  ,		UpperCamelCase       :       Optional[int]=7  ,		UpperCamelCase       :       Union[str, Any]=False  ,		UpperCamelCase       :       Dict=True  ,		UpperCamelCase       :       Optional[int]=False  ,		UpperCamelCase       :       Optional[Any]=True  ,		UpperCamelCase       :       Union[str, Any]=33  ,		UpperCamelCase       :       List[Any]=32  ,		UpperCamelCase       :       Dict=5  ,		UpperCamelCase       :       str=4  ,		UpperCamelCase       :       str=37  ,		UpperCamelCase       :       List[Any]="gelu"  ,		UpperCamelCase       :       List[str]=0.1  ,		UpperCamelCase       :       Dict=0.1  ,		UpperCamelCase       :       Dict=512  ,		UpperCamelCase       :       Any=16  ,		UpperCamelCase       :       List[Any]=2  ,		UpperCamelCase       :       List[Any]=0.02  ,		UpperCamelCase       :       Tuple=3  ,		UpperCamelCase       :       Optional[Any]=4  ,		UpperCamelCase       :       int=None  ,		):

             '''simple docstring'''







             lowercase__      =					parent
             lowercase__      =					batch_size
             lowercase__      =					seq_length
             lowercase__      =					is_training
             lowercase__      =					use_input_mask
             lowercase__      =					use_token_type_ids
             lowercase__      =					use_labels
             lowercase__      =					vocab_size
             lowercase__      =					hidden_size
             lowercase__      =					num_hidden_layers
             lowercase__      =					num_attention_heads
             lowercase__      =					intermediate_size
             lowercase__      =					hidden_act
             lowercase__      =					hidden_dropout_prob
             lowercase__      =					attention_probs_dropout_prob
             lowercase__      =					max_position_embeddings
             lowercase__      =					type_vocab_size
             lowercase__      =					type_sequence_label_size
             lowercase__      =					initializer_range
             lowercase__      =					num_labels
             lowercase__      =					num_choices
             lowercase__      =					scope



      def  UpperCamelCase__						(self       :       str					):

             '''simple docstring'''







             lowercase__      =					ids_tensor([self.batch_size, self.seq_length]  ,		self.vocab_size					)

             lowercase__      =					None
             if self.use_input_mask:
                    lowercase__      =					random_attention_mask([self.batch_size, self.seq_length]					)

             lowercase__      =					None
             lowercase__      =					None
             lowercase__      =					None
             if self.use_labels:
                    lowercase__      =					ids_tensor([self.batch_size]  ,		self.type_sequence_label_size					)
                    lowercase__      =					ids_tensor([self.batch_size, self.seq_length]  ,		self.num_labels					)
                    lowercase__      =					ids_tensor([self.batch_size]  ,		self.num_choices					)

             lowercase__      =					self.get_config()

             return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels



      def  UpperCamelCase__						(self       :       Optional[Any]					):

             '''simple docstring'''







             return EsmConfig(
                 vocab_size=self.vocab_size  ,		hidden_size=self.hidden_size  ,		pad_token_id=1  ,		num_hidden_layers=self.num_hidden_layers  ,		num_attention_heads=self.num_attention_heads  ,		intermediate_size=self.intermediate_size  ,		hidden_act=self.hidden_act  ,		hidden_dropout_prob=self.hidden_dropout_prob  ,		attention_probs_dropout_prob=self.attention_probs_dropout_prob  ,		max_position_embeddings=self.max_position_embeddings  ,		type_vocab_size=self.type_vocab_size  ,		initializer_range=self.initializer_range  ,		)



      def  UpperCamelCase__						(self       :       Optional[int]  ,		UpperCamelCase       :       Union[str, Any]  ,		UpperCamelCase       :       int  ,		UpperCamelCase       :       List[Any]  ,		UpperCamelCase       :       str  ,		UpperCamelCase       :       str  ,		UpperCamelCase       :       Dict					):

             '''simple docstring'''







             lowercase__      =					EsmModel(config=UpperCamelCase					)
             model.to(UpperCamelCase					)
             model.eval()
             lowercase__      =					model(UpperCamelCase  ,		attention_mask=UpperCamelCase					)
             lowercase__      =					model(UpperCamelCase					)
             lowercase__      =					model(UpperCamelCase					)

             self.parent.assertEqual(result.last_hidden_state.shape  ,		(self.batch_size, self.seq_length, self.hidden_size)					)
             self.parent.assertEqual(result.pooler_output.shape  ,		(self.batch_size, self.hidden_size)					)



      def  UpperCamelCase__						(self       :       str  ,		UpperCamelCase       :       Any  ,		UpperCamelCase       :       Tuple  ,		UpperCamelCase       :       Any  ,		UpperCamelCase       :       List[Any]  ,		UpperCamelCase       :       Optional[Any]  ,		UpperCamelCase       :       Optional[int]					):

             '''simple docstring'''







             lowercase__      =					EsmForMaskedLM(config=UpperCamelCase					)
             model.to(UpperCamelCase					)
             model.eval()
             lowercase__      =					model(UpperCamelCase  ,		attention_mask=UpperCamelCase  ,		labels=UpperCamelCase					)
             self.parent.assertEqual(result.logits.shape  ,		(self.batch_size, self.seq_length, self.vocab_size)					)



      def  UpperCamelCase__						(self       :       Any  ,		UpperCamelCase       :       str  ,		UpperCamelCase       :       List[Any]  ,		UpperCamelCase       :       Tuple  ,		UpperCamelCase       :       int  ,		UpperCamelCase       :       Union[str, Any]  ,		UpperCamelCase       :       List[Any]					):

             '''simple docstring'''







             lowercase__      =					self.num_labels
             lowercase__      =					EsmForTokenClassification(config=UpperCamelCase					)
             model.to(UpperCamelCase					)
             model.eval()
             lowercase__      =					model(UpperCamelCase  ,		attention_mask=UpperCamelCase  ,		labels=UpperCamelCase					)
             self.parent.assertEqual(result.logits.shape  ,		(self.batch_size, self.seq_length, self.num_labels)					)






      def  UpperCamelCase__						(self       :       Optional[Any]					):

             '''simple docstring'''







             lowercase__      =					self.prepare_config_and_inputs()
             (
                 (
                 lowercase__
             )     ,(
                 lowercase__
             )     ,(
                 lowercase__
             )     ,(
                 lowercase__
             )     ,(
                 lowercase__
             )     ,(
                 lowercase__
             )     ,
             )      =					config_and_inputs
             lowercase__      =					{'''input_ids''': input_ids, '''attention_mask''': input_mask}
             return config, inputs_dict






@require_torch
class     __lowerCAmelCase		(lowercase_    ,	lowercase_    ,	unittest.TestCase	):





      '''simple docstring'''






      lowerCAmelCase__   :	Dict       =				False

      lowerCAmelCase__   :	Optional[Any]       =				(
          (
              EsmForMaskedLM,
              EsmModel,
              EsmForSequenceClassification,
              EsmForTokenClassification,
          )
          if is_torch_available()
          else ()
      )
      lowerCAmelCase__   :	Optional[int]       =				()
      lowerCAmelCase__   :	Any       =				(
          {
              """feature-extraction""": EsmModel,
              """fill-mask""": EsmForMaskedLM,
              """text-classification""": EsmForSequenceClassification,
              """token-classification""": EsmForTokenClassification,
              """zero-shot""": EsmForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )
      lowerCAmelCase__   :	Any       =				True



      def  UpperCamelCase__						(self       :       List[Any]					):

             '''simple docstring'''







             lowercase__      =					EsmModelTester(self					)
             lowercase__      =					ConfigTester(self  ,		config_class=UpperCamelCase  ,		hidden_size=37					)



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

             '''simple docstring'''







             self.config_tester.run_common_tests()



      def  UpperCamelCase__						(self       :       Optional[Any]					):

             '''simple docstring'''







             lowercase__      =					self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_model(*UpperCamelCase					)



      def  UpperCamelCase__						(self       :       Any					):

             '''simple docstring'''







             lowercase__      =					self.model_tester.prepare_config_and_inputs()
             for type in ["absolute", "relative_key", "relative_key_query"]:
                    lowercase__      =					type
                    self.model_tester.create_and_check_model(*UpperCamelCase					)



      def  UpperCamelCase__						(self       :       Optional[int]					):

             '''simple docstring'''







             lowercase__      =					self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_masked_lm(*UpperCamelCase					)



      def  UpperCamelCase__						(self       :       List[Any]					):

             '''simple docstring'''







             lowercase__      =					self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_token_classification(*UpperCamelCase					)



      @slow
      def  UpperCamelCase__						(self       :       Union[str, Any]					):

             '''simple docstring'''







             for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                    lowercase__      =					EsmModel.from_pretrained(UpperCamelCase					)
                    self.assertIsNotNone(UpperCamelCase					)



      def  UpperCamelCase__						(self       :       Optional[int]					):

             '''simple docstring'''







             lowercase__      =					self.model_tester.prepare_config_and_inputs()[0]
             lowercase__      =					EsmEmbeddings(config=UpperCamelCase					)

             lowercase__      =					torch.as_tensor([[12, 31, 13, model.padding_idx]]					)
             lowercase__      =					torch.as_tensor(
                 [
                     [
                         0 + model.padding_idx + 1,
                         1 + model.padding_idx + 1,
                         2 + model.padding_idx + 1,
                         model.padding_idx,
                     ]
                 ]					)
             lowercase__      =					create_position_ids_from_input_ids(UpperCamelCase  ,		model.padding_idx					)
             self.assertEqual(position_ids.shape  ,		expected_positions.shape					)
             self.assertTrue(torch.all(torch.eq(UpperCamelCase  ,		UpperCamelCase					)					)					)



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

             '''simple docstring'''







             lowercase__      =					self.model_tester.prepare_config_and_inputs()[0]
             lowercase__      =					EsmEmbeddings(config=UpperCamelCase					)

             lowercase__      =					torch.empty(2  ,		4  ,		30					)
             lowercase__      =					[
                 0 + embeddings.padding_idx + 1,
                 1 + embeddings.padding_idx + 1,
                 2 + embeddings.padding_idx + 1,
                 3 + embeddings.padding_idx + 1,
             ]
             lowercase__      =					torch.as_tensor([expected_single_positions, expected_single_positions]					)
             lowercase__      =					embeddings.create_position_ids_from_inputs_embeds(UpperCamelCase					)
             self.assertEqual(position_ids.shape  ,		expected_positions.shape					)
             self.assertTrue(torch.all(torch.eq(UpperCamelCase  ,		UpperCamelCase					)					)					)



      @unittest.skip('''Esm does not support embedding resizing'''					)
      def  UpperCamelCase__						(self       :       Optional[int]					):

             '''simple docstring'''







             pass



      @unittest.skip('''Esm does not support embedding resizing'''					)
      def  UpperCamelCase__						(self       :       Optional[Any]					):

             '''simple docstring'''







             pass






      @unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''					)
      def  UpperCamelCase__						(self       :       Optional[int]					):

             '''simple docstring'''







             pass






@require_torch
class     __lowerCAmelCase		(lowercase_	):





      '''simple docstring'''






      @slow
      def  UpperCamelCase__						(self       :       Optional[Any]					):

             '''simple docstring'''







             with torch.no_grad():
                    lowercase__      =					EsmForMaskedLM.from_pretrained('''facebook/esm2_t6_8M_UR50D'''					)
                    model.eval()
                    lowercase__      =					torch.tensor([[0, 1, 2, 3, 4, 5]]					)
                    lowercase__      =					model(UpperCamelCase					)[0]

                    lowercase__      =					33

                    lowercase__      =					torch.Size((1, 6, vocab_size)					)
                    self.assertEqual(output.shape  ,		UpperCamelCase					)

                    lowercase__      =					torch.tensor(
                        [[[8.92_15, -10.58_98, -6.46_71], [-6.39_67, -13.91_14, -1.12_12], [-7.78_12, -13.95_16, -3.74_06]]]					)
                    self.assertTrue(torch.allclose(output[:, :3, :3]  ,		UpperCamelCase  ,		atol=1E-4					)					)






      @slow
      def  UpperCamelCase__						(self       :       Optional[Any]					):

             '''simple docstring'''







             with torch.no_grad():
                    lowercase__      =					EsmModel.from_pretrained('''facebook/esm2_t6_8M_UR50D'''					)
                    model.eval()

                    lowercase__      =					torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]					)
                    lowercase__      =					model(UpperCamelCase					)[0]
                    # compare the actual values for a slice.
                    lowercase__      =					torch.tensor(
                        [[[0.14_44, 0.54_13, 0.32_48], [0.30_34, 0.00_53, 0.31_08], [0.32_28, -0.24_99, 0.34_15]]]					)
                    self.assertTrue(torch.allclose(output[:, :3, :3]  ,		UpperCamelCase  ,		atol=1E-4					)					)




