

import math
import os
import unittest

from transformers import MegatronBertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_sentencepiece, require_tokenizers, 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 (
         MODEL_FOR_PRETRAINING_MAPPING,
         MegatronBertForCausalLM,
         MegatronBertForMaskedLM,
         MegatronBertForMultipleChoice,
         MegatronBertForNextSentencePrediction,
         MegatronBertForPreTraining,
         MegatronBertForQuestionAnswering,
         MegatronBertForSequenceClassification,
         MegatronBertForTokenClassification,
         MegatronBertModel,
     )





class  __UpperCAmelCase					:



   def __init__(							self							:					int,		__A							:					List[Any],		__A							:					int=1_3,		__A							:					List[Any]=7,		__A							:					str=True,		__A							:					List[Any]=True,		__A							:					int=True,		__A							:					str=True,		__A							:					List[Any]=9_9,		__A							:					Optional[int]=6_4,		__A							:					Optional[Any]=3_2,		__A							:					Any=5,		__A							:					List[Any]=4,		__A							:					str=3_7,		__A							:					int="gelu",		__A							:					Dict=0.1,		__A							:					Optional[Any]=0.1,		__A							:					Any=5_1_2,		__A							:					List[str]=1_6,		__A							:					Union[str, Any]=2,		__A							:					Dict=0.0_2,		__A							:					int=3,		__A							:					Any=4,		__A							:					List[str]=None,		):
     UpperCAmelCase	:							List[str]						=			parent
     UpperCAmelCase	:							List[Any]						=			batch_size
     UpperCAmelCase	:							Optional[Any]						=			seq_length
     UpperCAmelCase	:							Tuple						=			is_training
     UpperCAmelCase	:							Union[str, Any]						=			use_input_mask
     UpperCAmelCase	:							Dict						=			use_token_type_ids
     UpperCAmelCase	:							Dict						=			use_labels
     UpperCAmelCase	:							Optional[Any]						=			vocab_size
     UpperCAmelCase	:							Any						=			hidden_size
     UpperCAmelCase	:							List[str]						=			embedding_size
     UpperCAmelCase	:							str						=			num_hidden_layers
     UpperCAmelCase	:							Union[str, Any]						=			num_attention_heads
     UpperCAmelCase	:							str						=			intermediate_size
     UpperCAmelCase	:							Union[str, Any]						=			hidden_act
     UpperCAmelCase	:							List[str]						=			hidden_dropout_prob
     UpperCAmelCase	:							Dict						=			attention_probs_dropout_prob
     UpperCAmelCase	:							str						=			max_position_embeddings
     UpperCAmelCase	:							Any						=			type_vocab_size
     UpperCAmelCase	:							Union[str, Any]						=			type_sequence_label_size
     UpperCAmelCase	:							Dict						=			initializer_range
     UpperCAmelCase	:							Dict						=			num_labels
     UpperCAmelCase	:							List[str]						=			num_choices
     UpperCAmelCase	:							str						=			scope



   def 	__magic_name__   (							self							:					str						):
     UpperCAmelCase	:							Any						=			ids_tensor([self.batch_size, self.seq_length],		self.vocab_size						)

     UpperCAmelCase	:							Any						=			None
     if self.use_input_mask:
       UpperCAmelCase	:							Tuple						=			random_attention_mask([self.batch_size, self.seq_length]						)

     UpperCAmelCase	:							Optional[Any]						=			None
     if self.use_token_type_ids:
       UpperCAmelCase	:							int						=			ids_tensor([self.batch_size, self.seq_length],		self.type_vocab_size						)

     UpperCAmelCase	:							Union[str, Any]						=			None
     UpperCAmelCase	:							List[str]						=			None
     UpperCAmelCase	:							Tuple						=			None
     if self.use_labels:
       UpperCAmelCase	:							Union[str, Any]						=			ids_tensor([self.batch_size],		self.type_sequence_label_size						)
       UpperCAmelCase	:							Optional[Any]						=			ids_tensor([self.batch_size, self.seq_length],		self.num_labels						)
       UpperCAmelCase	:							List[str]						=			ids_tensor([self.batch_size],		self.num_choices						)

     UpperCAmelCase	:							Optional[int]						=			self.get_config()

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



   def 	__magic_name__   (							self							:					List[str]						):
     return MegatronBertConfig(
         vocab_size=self.vocab_size,		hidden_size=self.hidden_size,		num_hidden_layers=self.num_hidden_layers,		num_attention_heads=self.num_attention_heads,		intermediate_size=self.intermediate_size,		embedding_size=self.embedding_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,		is_decoder=__A,		initializer_range=self.initializer_range,		)



   def 	__magic_name__   (							self							:					Optional[Any],		__A							:					Optional[Any],		__A							:					Union[str, Any],		__A							:					int,		__A							:					str,		__A							:					List[str],		__A							:					List[Any],		__A							:					Union[str, Any]						):
     UpperCAmelCase	:							Optional[int]						=			MegatronBertModel(config=__A						)
     model.to(__A						)
     model.eval()
     UpperCAmelCase	:							str						=			model(__A,		attention_mask=__A,		token_type_ids=__A						)
     UpperCAmelCase	:							Any						=			model(__A,		token_type_ids=__A						)
     UpperCAmelCase	:							Optional[int]						=			model(__A						)

     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 	__magic_name__   (							self							:					List[Any],		__A							:					List[Any],		__A							:					List[str],		__A							:					Dict,		__A							:					Tuple,		__A							:					Optional[Any],		__A							:					int,		__A							:					Any						):
     UpperCAmelCase	:							Tuple						=			MegatronBertForMaskedLM(config=__A						)
     model.to(__A						)
     model.eval()
     UpperCAmelCase	:							Optional[int]						=			model(__A,		attention_mask=__A,		token_type_ids=__A,		labels=__A						)
     self.parent.assertEqual(result.logits.shape,		(self.batch_size, self.seq_length, self.vocab_size)						)



   def 	__magic_name__   (							self							:					Dict,		__A							:					int,		__A							:					Tuple,		__A							:					List[Any],		__A							:					Optional[Any],		__A							:					Dict,		__A							:					Optional[Any],		__A							:					Tuple						):
     UpperCAmelCase	:							Tuple						=			MegatronBertForCausalLM(config=__A						)
     model.to(__A						)
     model.eval()
     UpperCAmelCase	:							Dict						=			model(__A,		attention_mask=__A,		token_type_ids=__A,		labels=__A						)
     self.parent.assertEqual(result.logits.shape,		(self.batch_size, self.seq_length, self.vocab_size)						)



   def 	__magic_name__   (							self							:					Union[str, Any],		__A							:					Dict,		__A							:					str,		__A							:					List[Any],		__A							:					str,		__A							:					Dict,		__A							:					str,		__A							:					Optional[int]						):
     UpperCAmelCase	:							Tuple						=			MegatronBertForNextSentencePrediction(config=__A						)
     model.to(__A						)
     model.eval()
     UpperCAmelCase	:							Optional[int]						=			model(
         __A,		attention_mask=__A,		token_type_ids=__A,		labels=__A,		)
     self.parent.assertEqual(result.logits.shape,		(self.batch_size, 2)						)



   def 	__magic_name__   (							self							:					Dict,		__A							:					str,		__A							:					Any,		__A							:					Optional[int],		__A							:					str,		__A							:					List[str],		__A							:					str,		__A							:					str						):
     UpperCAmelCase	:							Any						=			MegatronBertForPreTraining(config=__A						)
     model.to(__A						)
     model.eval()
     UpperCAmelCase	:							List[Any]						=			model(
         __A,		attention_mask=__A,		token_type_ids=__A,		labels=__A,		next_sentence_label=__A,		)
     self.parent.assertEqual(result.prediction_logits.shape,		(self.batch_size, self.seq_length, self.vocab_size)						)
     self.parent.assertEqual(result.seq_relationship_logits.shape,		(self.batch_size, 2)						)



   def 	__magic_name__   (							self							:					List[str],		__A							:					Optional[int],		__A							:					Any,		__A							:					Union[str, Any],		__A							:					List[Any],		__A							:					Dict,		__A							:					str,		__A							:					str						):
     UpperCAmelCase	:							List[Any]						=			MegatronBertForQuestionAnswering(config=__A						)
     model.to(__A						)
     model.eval()
     UpperCAmelCase	:							Dict						=			model(
         __A,		attention_mask=__A,		token_type_ids=__A,		start_positions=__A,		end_positions=__A,		)
     self.parent.assertEqual(result.start_logits.shape,		(self.batch_size, self.seq_length)						)
     self.parent.assertEqual(result.end_logits.shape,		(self.batch_size, self.seq_length)						)



   def 	__magic_name__   (							self							:					Optional[int],		__A							:					str,		__A							:					int,		__A							:					Dict,		__A							:					Optional[Any],		__A							:					Dict,		__A							:					Union[str, Any],		__A							:					Optional[int]						):
     UpperCAmelCase	:							Tuple						=			self.num_labels
     UpperCAmelCase	:							str						=			MegatronBertForSequenceClassification(__A						)
     model.to(__A						)
     model.eval()
     UpperCAmelCase	:							Optional[int]						=			model(__A,		attention_mask=__A,		token_type_ids=__A,		labels=__A						)
     self.parent.assertEqual(result.logits.shape,		(self.batch_size, self.num_labels)						)



   def 	__magic_name__   (							self							:					Tuple,		__A							:					Tuple,		__A							:					List[Any],		__A							:					int,		__A							:					Tuple,		__A							:					List[str],		__A							:					Union[str, Any],		__A							:					Optional[Any]						):
     UpperCAmelCase	:							Optional[int]						=			self.num_labels
     UpperCAmelCase	:							Any						=			MegatronBertForTokenClassification(config=__A						)
     model.to(__A						)
     model.eval()
     UpperCAmelCase	:							List[str]						=			model(__A,		attention_mask=__A,		token_type_ids=__A,		labels=__A						)
     self.parent.assertEqual(result.logits.shape,		(self.batch_size, self.seq_length, self.num_labels)						)



   def 	__magic_name__   (							self							:					Dict,		__A							:					Union[str, Any],		__A							:					Optional[Any],		__A							:					int,		__A							:					int,		__A							:					Any,		__A							:					int,		__A							:					Union[str, Any]						):
     UpperCAmelCase	:							str						=			self.num_choices
     UpperCAmelCase	:							List[Any]						=			MegatronBertForMultipleChoice(config=__A						)
     model.to(__A						)
     model.eval()
     UpperCAmelCase	:							Optional[Any]						=			input_ids.unsqueeze(1						).expand(-1,		self.num_choices,		-1						).contiguous()
     UpperCAmelCase	:							Optional[Any]						=			token_type_ids.unsqueeze(1						).expand(-1,		self.num_choices,		-1						).contiguous()
     UpperCAmelCase	:							List[str]						=			input_mask.unsqueeze(1						).expand(-1,		self.num_choices,		-1						).contiguous()
     UpperCAmelCase	:							List[Any]						=			model(
         __A,		attention_mask=__A,		token_type_ids=__A,		labels=__A,		)
     self.parent.assertEqual(result.logits.shape,		(self.batch_size, self.num_choices)						)





   def 	__magic_name__   (							self							:					Dict						):
     UpperCAmelCase	:							Dict						=			self.prepare_config_and_inputs()
     (
         (
         UpperCAmelCase
     )			,						(
         UpperCAmelCase
     )			,						(
         UpperCAmelCase
     )			,						(
         UpperCAmelCase
     )			,						(
         UpperCAmelCase
     )			,						(
         UpperCAmelCase
     )			,						(
         UpperCAmelCase
     )			,						
     )	:							Optional[Any]						=			config_and_inputs
     UpperCAmelCase	:							Any						=			{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
     return config, inputs_dict






@require_torch
class  __UpperCAmelCase					(	lowerCamelCase__      ,  lowerCamelCase__      ,  unittest.TestCase						):
   UpperCamelCase       					=					(
       (
           MegatronBertModel,
           MegatronBertForMaskedLM,
           MegatronBertForCausalLM,
           MegatronBertForMultipleChoice,
           MegatronBertForNextSentencePrediction,
           MegatronBertForPreTraining,
           MegatronBertForQuestionAnswering,
           MegatronBertForSequenceClassification,
           MegatronBertForTokenClassification,
       )
       if is_torch_available()
       else ()
   )
   UpperCamelCase       					=					(
       {
           """feature-extraction""": MegatronBertModel,
           """fill-mask""": MegatronBertForMaskedLM,
           """question-answering""": MegatronBertForQuestionAnswering,
           """text-classification""": MegatronBertForSequenceClassification,
           """text-generation""": MegatronBertForCausalLM,
           """token-classification""": MegatronBertForTokenClassification,
           """zero-shot""": MegatronBertForSequenceClassification,
       }
       if is_torch_available()
       else {}
   )
   UpperCamelCase       					=					True
   # test_resize_embeddings = False
   UpperCamelCase       					=					False



   def 	__magic_name__   (							self							:					Optional[int],		__A							:					Optional[int],		__A							:					Any,		__A							:					Union[str, Any]=False						):
     UpperCAmelCase	:							Optional[Any]						=			super()._prepare_for_class(__A,		__A,		return_labels=__A						)

     if return_labels:
       if model_class in get_values(__A						):
         UpperCAmelCase	:							Union[str, Any]						=			torch.zeros(
             (self.model_tester.batch_size, self.model_tester.seq_length),		dtype=torch.long,		device=__A						)
         UpperCAmelCase	:							Dict						=			torch.zeros(
             self.model_tester.batch_size,		dtype=torch.long,		device=__A						)
     return inputs_dict



   def 	__magic_name__   (							self							:					Optional[Any]						):
     UpperCAmelCase	:							int						=			MegatronBertModelTester(self						)
     UpperCAmelCase	:							List[Any]						=			ConfigTester(self,		config_class=__A,		hidden_size=3_7						)



   def 	__magic_name__   (							self							:					int						):
     self.config_tester.run_common_tests()



   def 	__magic_name__   (							self							:					Optional[int]						):
     UpperCAmelCase	:							Tuple						=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_model(*__A						)



   def 	__magic_name__   (							self							:					Any						):
     UpperCAmelCase	:							str						=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_masked_lm(*__A						)



   def 	__magic_name__   (							self							:					Any						):
     UpperCAmelCase	:							Any						=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*__A						)



   def 	__magic_name__   (							self							:					List[str]						):
     UpperCAmelCase	:							str						=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*__A						)



   def 	__magic_name__   (							self							:					Union[str, Any]						):
     UpperCAmelCase	:							Optional[int]						=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_pretraining(*__A						)



   def 	__magic_name__   (							self							:					int						):
     UpperCAmelCase	:							Any						=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_question_answering(*__A						)



   def 	__magic_name__   (							self							:					Optional[int]						):
     UpperCAmelCase	:							Union[str, Any]						=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*__A						)





   def 	__magic_name__   (							self							:					Optional[Any]						):
     UpperCAmelCase	:							Union[str, Any]						=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_token_classification(*__A						)






def       a__					(		UpperCAmelCase							:      Tuple      )  ->						Tuple:
  return torch.tensor(
      UpperCAmelCase    ,			dtype=torch.long    ,			device=UpperCAmelCase    ,			)


_lowerCamelCase    :       Any        =							1E-4






@require_torch
@require_sentencepiece
@require_tokenizers
class  __UpperCAmelCase					(	unittest.TestCase						):



   @slow
   @unittest.skip('''Model is not available.'''						)
   def 	__magic_name__   (							self							:					List[Any]						):
     UpperCAmelCase	:							Union[str, Any]						=			'''nvidia/megatron-bert-uncased-345m'''
     if "MYDIR" in os.environ:
       UpperCAmelCase	:							int						=			os.path.join(os.environ['''MYDIR'''],		__A						)
     UpperCAmelCase	:							Any						=			MegatronBertModel.from_pretrained(__A						)
     model.to(__A						)
     model.half()
     UpperCAmelCase	:							str						=			_long_tensor([[1_0_1, 7_1_1_0, 1_0_0_5, 1_0_5_6, 2_0_2_3, 1_1_3_3_3, 1_7_4_1_3, 1_0_2_9, 1_0_2]]						)
     with torch.no_grad():
       UpperCAmelCase	:							Optional[Any]						=			model(__A						)[0]
     UpperCAmelCase	:							str						=			torch.Size((1, 9, 1_0_2_4)						)
     self.assertEqual(output.shape,		__A						)

     UpperCAmelCase	:							Dict						=			[-0.6_0_4_0, -0.2_5_1_7, -0.1_0_2_5, 0.3_4_2_0, -0.6_7_5_8, -0.0_0_1_7, -0.1_0_8_9, -0.1_9_9_0, 0.5_7_2_8]
     for ii in range(3						):
       for jj in range(3						):
         UpperCAmelCase	:							List[Any]						=			output[0, ii, jj]
         UpperCAmelCase	:							List[str]						=			expected[3 * ii + jj]
         UpperCAmelCase	:							str						=			'''ii={} jj={} a={} b={}'''.format(__A,		__A,		__A,		__A						)
         self.assertTrue(math.isclose(__A,		__A,		rel_tol=__A,		abs_tol=__A						),		msg=__A						)








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					(	lowerCamelCase__      ,  unittest.TestCase						):
   UpperCamelCase       					=					ShapEImgaImgPipeline
   UpperCamelCase       					=					["""image"""]
   UpperCamelCase       					=					["""image"""]
   UpperCamelCase       					=					[
       """num_images_per_prompt""",
       """num_inference_steps""",
       """generator""",
       """latents""",
       """guidance_scale""",
       """frame_size""",
       """output_type""",
       """return_dict""",
   ]
   UpperCamelCase       					=					False



   @property
   def 	__magic_name__   (							self							:					Optional[Any]						):
     return 3_2



   @property
   def 	__magic_name__   (							self							:					Optional[int]						):
     return 3_2



   @property
   def 	__magic_name__   (							self							:					Union[str, Any]						):
     return self.time_input_dim * 4



   @property
   def 	__magic_name__   (							self							:					Any						):
     return 8



   @property
   def 	__magic_name__   (							self							:					List[str]						):
     torch.manual_seed(0						)
     UpperCAmelCase	:							str						=			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,		)

     UpperCAmelCase	:							Any						=			CLIPVisionModel(__A						)
     return model



   @property
   def 	__magic_name__   (							self							:					int						):
     UpperCAmelCase	:							Optional[int]						=			CLIPImageProcessor(
         crop_size=2_2_4,		do_center_crop=__A,		do_normalize=__A,		do_resize=__A,		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_2_4,		)

     return image_processor



   @property
   def 	__magic_name__   (							self							:					Dict						):
     torch.manual_seed(0						)

     UpperCAmelCase	:							Any						=			{
         '''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,
     }

     UpperCAmelCase	:							List[Any]						=			PriorTransformer(**__A						)
     return model



   @property
   def 	__magic_name__   (							self							:					Union[str, Any]						):
     torch.manual_seed(0						)

     UpperCAmelCase	:							List[Any]						=			{
         '''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,
         ),
     }
     UpperCAmelCase	:							List[str]						=			ShapERenderer(**__A						)
     return model



   def 	__magic_name__   (							self							:					List[Any]						):
     UpperCAmelCase	:							str						=			self.dummy_prior
     UpperCAmelCase	:							List[str]						=			self.dummy_image_encoder
     UpperCAmelCase	:							List[Any]						=			self.dummy_image_processor
     UpperCAmelCase	:							Dict						=			self.dummy_renderer

     UpperCAmelCase	:							int						=			HeunDiscreteScheduler(
         beta_schedule='''exp''',		num_train_timesteps=1_0_2_4,		prediction_type='''sample''',		use_karras_sigmas=__A,		clip_sample=__A,		clip_sample_range=1.0,		)
     UpperCAmelCase	:							List[str]						=			{
         '''prior''': prior,
         '''image_encoder''': image_encoder,
         '''image_processor''': image_processor,
         '''renderer''': renderer,
         '''scheduler''': scheduler,
     }

     return components



   def 	__magic_name__   (							self							:					List[Any],		__A							:					Dict,		__A							:					List[Any]=0						):
     UpperCAmelCase	:							int						=			floats_tensor((1, 3, 6_4, 6_4),		rng=random.Random(__A						)						).to(__A						)

     if str(__A						).startswith('''mps'''						):
       UpperCAmelCase	:							List[str]						=			torch.manual_seed(__A						)
     else:
       UpperCAmelCase	:							Optional[Any]						=			torch.Generator(device=__A						).manual_seed(__A						)
     UpperCAmelCase	:							Dict						=			{
         '''image''': input_image,
         '''generator''': generator,
         '''num_inference_steps''': 1,
         '''frame_size''': 3_2,
         '''output_type''': '''np''',
     }
     return inputs



   def 	__magic_name__   (							self							:					str						):
     UpperCAmelCase	:							Dict						=			'''cpu'''

     UpperCAmelCase	:							Any						=			self.get_dummy_components()

     UpperCAmelCase	:							Tuple						=			self.pipeline_class(**__A						)
     UpperCAmelCase	:							Dict						=			pipe.to(__A						)

     pipe.set_progress_bar_config(disable=__A						)

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

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

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



   def 	__magic_name__   (							self							:					str						):
     UpperCAmelCase	:							Tuple						=			torch_device == '''cpu'''
     UpperCAmelCase	:							List[Any]						=			True
     self._test_inference_batch_single_identical(
         batch_size=2,		test_max_difference=__A,		relax_max_difference=__A,		)





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

     UpperCAmelCase	:							Any						=			1
     UpperCAmelCase	:							Any						=			2

     UpperCAmelCase	:							Optional[int]						=			self.get_dummy_inputs(__A						)

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

     UpperCAmelCase	:							int						=			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 	__magic_name__   (							self							:					List[Any]						):
     # clean up the VRAM after each test
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()





   def 	__magic_name__   (							self							:					Tuple						):
     UpperCAmelCase	:							List[str]						=			load_image(
         '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main''' '''/shap_e/corgi.png'''						)
     UpperCAmelCase	:							Dict						=			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	:							int						=			pipe.to(__A						)
     pipe.set_progress_bar_config(disable=__A						)

     UpperCAmelCase	:							Dict						=			torch.Generator(device=__A						).manual_seed(0						)

     UpperCAmelCase	:							List[str]						=			pipe(
         __A,		generator=__A,		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(__A,		__A						)





