






'''simple docstring'''




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       __lowercase					:

	def __init__(      self   :						List[Any]  ,					UpperCAmelCase_   :						Optional[Any]  ,					UpperCAmelCase_   :						Optional[int]=13  ,					UpperCAmelCase_   :						Union[str, Any]=7  ,					UpperCAmelCase_   :						List[Any]=True  ,					UpperCAmelCase_   :						int=True  ,					UpperCAmelCase_   :						int=True  ,					UpperCAmelCase_   :						Union[str, Any]=True  ,					UpperCAmelCase_   :						List[Any]=99  ,					UpperCAmelCase_   :						str=64  ,					UpperCAmelCase_   :						Union[str, Any]=32  ,					UpperCAmelCase_   :						List[Any]=5  ,					UpperCAmelCase_   :						Optional[Any]=4  ,					UpperCAmelCase_   :						Optional[int]=37  ,					UpperCAmelCase_   :						Any="gelu"  ,					UpperCAmelCase_   :						Union[str, Any]=0.1  ,					UpperCAmelCase_   :						Any=0.1  ,					UpperCAmelCase_   :						Any=512  ,					UpperCAmelCase_   :						Dict=16  ,					UpperCAmelCase_   :						Any=2  ,					UpperCAmelCase_   :						Optional[Any]=0.02  ,					UpperCAmelCase_   :						List[Any]=3  ,					UpperCAmelCase_   :						Dict=4  ,					UpperCAmelCase_   :						List[Any]=None  ,					):
			UpperCamelCase__				:					List[str]					=					parent
			UpperCamelCase__				:					Optional[Any]					=					batch_size
			UpperCamelCase__				:					Tuple					=					seq_length
			UpperCamelCase__				:					Tuple					=					is_training
			UpperCamelCase__				:					Optional[Any]					=					use_input_mask
			UpperCamelCase__				:					Optional[int]					=					use_token_type_ids
			UpperCamelCase__				:					List[Any]					=					use_labels
			UpperCamelCase__				:					Any					=					vocab_size
			UpperCamelCase__				:					List[str]					=					hidden_size
			UpperCamelCase__				:					Dict					=					embedding_size
			UpperCamelCase__				:					Union[str, Any]					=					num_hidden_layers
			UpperCamelCase__				:					str					=					num_attention_heads
			UpperCamelCase__				:					str					=					intermediate_size
			UpperCamelCase__				:					List[str]					=					hidden_act
			UpperCamelCase__				:					Optional[int]					=					hidden_dropout_prob
			UpperCamelCase__				:					Dict					=					attention_probs_dropout_prob
			UpperCamelCase__				:					str					=					max_position_embeddings
			UpperCamelCase__				:					List[str]					=					type_vocab_size
			UpperCamelCase__				:					Union[str, Any]					=					type_sequence_label_size
			UpperCamelCase__				:					List[Any]					=					initializer_range
			UpperCamelCase__				:					Any					=					num_labels
			UpperCamelCase__				:					List[str]					=					num_choices
			UpperCamelCase__				:					List[str]					=					scope

	def      __UpperCamelCase							(      self   :						str):
			UpperCamelCase__				:					List[Any]					=					ids_tensor([self.batch_size, self.seq_length]  ,					self.vocab_size)

			UpperCamelCase__				:					List[str]					=					None
			if self.use_input_mask:
					UpperCamelCase__				:					Tuple					=					random_attention_mask([self.batch_size, self.seq_length])

			UpperCamelCase__				:					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[Any]					=					None
			UpperCamelCase__				:					List[str]					=					None
			if self.use_labels:
					UpperCamelCase__				:					str					=					ids_tensor([self.batch_size]  ,					self.type_sequence_label_size)
					UpperCamelCase__				:					List[str]					=					ids_tensor([self.batch_size, self.seq_length]  ,					self.num_labels)
					UpperCamelCase__				:					int					=					ids_tensor([self.batch_size]  ,					self.num_choices)

			UpperCamelCase__				:					List[str]					=					self.get_config()

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

	def      __UpperCamelCase							(      self   :						Dict):
			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=UpperCAmelCase_  ,					initializer_range=self.initializer_range  ,					)

	def      __UpperCamelCase							(      self   :						Any  ,					UpperCAmelCase_   :						Optional[int]  ,					UpperCAmelCase_   :						Any  ,					UpperCAmelCase_   :						List[Any]  ,					UpperCAmelCase_   :						Optional[int]  ,					UpperCAmelCase_   :						List[Any]  ,					UpperCAmelCase_   :						Dict  ,					UpperCAmelCase_   :						Any):
			UpperCamelCase__				:					Dict					=					MegatronBertModel(config=UpperCAmelCase_)
			model.to(UpperCAmelCase_)
			model.eval()
			UpperCamelCase__				:					Optional[int]					=					model(UpperCAmelCase_  ,					attention_mask=UpperCAmelCase_  ,					token_type_ids=UpperCAmelCase_)
			UpperCamelCase__				:					Any					=					model(UpperCAmelCase_  ,					token_type_ids=UpperCAmelCase_)
			UpperCamelCase__				:					Dict					=					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_   :						List[str]  ,					UpperCAmelCase_   :						str  ,					UpperCAmelCase_   :						str  ,					UpperCAmelCase_   :						Any  ,					UpperCAmelCase_   :						List[Any]  ,					UpperCAmelCase_   :						Any  ,					UpperCAmelCase_   :						Optional[int]):
			UpperCamelCase__				:					List[str]					=					MegatronBertForMaskedLM(config=UpperCAmelCase_)
			model.to(UpperCAmelCase_)
			model.eval()
			UpperCamelCase__				:					Optional[int]					=					model(UpperCAmelCase_  ,					attention_mask=UpperCAmelCase_  ,					token_type_ids=UpperCAmelCase_  ,					labels=UpperCAmelCase_)
			self.parent.assertEqual(result.logits.shape  ,					(self.batch_size, self.seq_length, self.vocab_size))

	def      __UpperCamelCase							(      self   :						Optional[Any]  ,					UpperCAmelCase_   :						Optional[int]  ,					UpperCAmelCase_   :						Tuple  ,					UpperCAmelCase_   :						Optional[int]  ,					UpperCAmelCase_   :						Union[str, Any]  ,					UpperCAmelCase_   :						List[str]  ,					UpperCAmelCase_   :						Tuple  ,					UpperCAmelCase_   :						List[str]):
			UpperCamelCase__				:					str					=					MegatronBertForCausalLM(config=UpperCAmelCase_)
			model.to(UpperCAmelCase_)
			model.eval()
			UpperCamelCase__				:					List[Any]					=					model(UpperCAmelCase_  ,					attention_mask=UpperCAmelCase_  ,					token_type_ids=UpperCAmelCase_  ,					labels=UpperCAmelCase_)
			self.parent.assertEqual(result.logits.shape  ,					(self.batch_size, self.seq_length, self.vocab_size))

	def      __UpperCamelCase							(      self   :						int  ,					UpperCAmelCase_   :						Tuple  ,					UpperCAmelCase_   :						int  ,					UpperCAmelCase_   :						Tuple  ,					UpperCAmelCase_   :						int  ,					UpperCAmelCase_   :						int  ,					UpperCAmelCase_   :						Tuple  ,					UpperCAmelCase_   :						Optional[Any]):
			UpperCamelCase__				:					Optional[Any]					=					MegatronBertForNextSentencePrediction(config=UpperCAmelCase_)
			model.to(UpperCAmelCase_)
			model.eval()
			UpperCamelCase__				:					int					=					model(
			    UpperCAmelCase_  ,					attention_mask=UpperCAmelCase_  ,					token_type_ids=UpperCAmelCase_  ,					labels=UpperCAmelCase_  ,					)
			self.parent.assertEqual(result.logits.shape  ,					(self.batch_size, 2))

	def      __UpperCamelCase							(      self   :						Optional[int]  ,					UpperCAmelCase_   :						List[Any]  ,					UpperCAmelCase_   :						Dict  ,					UpperCAmelCase_   :						Optional[Any]  ,					UpperCAmelCase_   :						Tuple  ,					UpperCAmelCase_   :						List[Any]  ,					UpperCAmelCase_   :						List[Any]  ,					UpperCAmelCase_   :						str):
			UpperCamelCase__				:					Dict					=					MegatronBertForPreTraining(config=UpperCAmelCase_)
			model.to(UpperCAmelCase_)
			model.eval()
			UpperCamelCase__				:					Optional[Any]					=					model(
			    UpperCAmelCase_  ,					attention_mask=UpperCAmelCase_  ,					token_type_ids=UpperCAmelCase_  ,					labels=UpperCAmelCase_  ,					next_sentence_label=UpperCAmelCase_  ,					)
			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      __UpperCamelCase							(      self   :						Optional[int]  ,					UpperCAmelCase_   :						Optional[Any]  ,					UpperCAmelCase_   :						Any  ,					UpperCAmelCase_   :						Tuple  ,					UpperCAmelCase_   :						Union[str, Any]  ,					UpperCAmelCase_   :						int  ,					UpperCAmelCase_   :						Union[str, Any]  ,					UpperCAmelCase_   :						str):
			UpperCamelCase__				:					str					=					MegatronBertForQuestionAnswering(config=UpperCAmelCase_)
			model.to(UpperCAmelCase_)
			model.eval()
			UpperCamelCase__				:					List[str]					=					model(
			    UpperCAmelCase_  ,					attention_mask=UpperCAmelCase_  ,					token_type_ids=UpperCAmelCase_  ,					start_positions=UpperCAmelCase_  ,					end_positions=UpperCAmelCase_  ,					)
			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      __UpperCamelCase							(      self   :						Any  ,					UpperCAmelCase_   :						Any  ,					UpperCAmelCase_   :						Optional[int]  ,					UpperCAmelCase_   :						str  ,					UpperCAmelCase_   :						Union[str, Any]  ,					UpperCAmelCase_   :						Union[str, Any]  ,					UpperCAmelCase_   :						Optional[int]  ,					UpperCAmelCase_   :						List[Any]):
			UpperCamelCase__				:					Optional[Any]					=					self.num_labels
			UpperCamelCase__				:					Optional[int]					=					MegatronBertForSequenceClassification(UpperCAmelCase_)
			model.to(UpperCAmelCase_)
			model.eval()
			UpperCamelCase__				:					Optional[int]					=					model(UpperCAmelCase_  ,					attention_mask=UpperCAmelCase_  ,					token_type_ids=UpperCAmelCase_  ,					labels=UpperCAmelCase_)
			self.parent.assertEqual(result.logits.shape  ,					(self.batch_size, self.num_labels))

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

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





	def      __UpperCamelCase							(      self   :						Union[str, Any]):
			UpperCamelCase__				:					Dict					=					self.prepare_config_and_inputs()
			(
			    (
			    UpperCamelCase__
			),	(
			    UpperCamelCase__
			),	(
			    UpperCamelCase__
			),	(
			    UpperCamelCase__
			),	(
			    UpperCamelCase__
			),	(
			    UpperCamelCase__
			),	(
			    UpperCamelCase__
			),	
			)				:					Dict					=					config_and_inputs
			UpperCamelCase__				:					Dict					=					{'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
			return config, inputs_dict







@require_torch
class       __lowercase					(__lowerCamelCase   ,				__lowerCamelCase   ,				unittest.TestCase    ):
	_lowerCamelCase              =  (
	    (
	        MegatronBertModel,
	        MegatronBertForMaskedLM,
	        MegatronBertForCausalLM,
	        MegatronBertForMultipleChoice,
	        MegatronBertForNextSentencePrediction,
	        MegatronBertForPreTraining,
	        MegatronBertForQuestionAnswering,
	        MegatronBertForSequenceClassification,
	        MegatronBertForTokenClassification,
	    )
	    if is_torch_available()
	    else ()
	)
	_lowerCamelCase              =  (
	    {
	        '''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 {}
	)
	_lowerCamelCase              =  True
	# test_resize_embeddings = False
	_lowerCamelCase              =  False

	def      __UpperCamelCase							(      self   :						Optional[Any]  ,					UpperCAmelCase_   :						str  ,					UpperCAmelCase_   :						str  ,					UpperCAmelCase_   :						int=False):
			UpperCamelCase__				:					int					=					super()._prepare_for_class(UpperCAmelCase_  ,					UpperCAmelCase_  ,					return_labels=UpperCAmelCase_)

			if return_labels:
					if model_class in get_values(UpperCAmelCase_):
							UpperCamelCase__				:					int					=					torch.zeros(
							    (self.model_tester.batch_size, self.model_tester.seq_length)  ,					dtype=torch.long  ,					device=UpperCAmelCase_)
							UpperCamelCase__				:					Optional[int]					=					torch.zeros(
							    self.model_tester.batch_size  ,					dtype=torch.long  ,					device=UpperCAmelCase_)
			return inputs_dict

	def      __UpperCamelCase							(      self   :						Union[str, Any]):
			UpperCamelCase__				:					List[str]					=					MegatronBertModelTester(self)
			UpperCamelCase__				:					Optional[Any]					=					ConfigTester(self  ,					config_class=UpperCAmelCase_  ,					hidden_size=37)

	def      __UpperCamelCase							(      self   :						List[Any]):
			self.config_tester.run_common_tests()

	def      __UpperCamelCase							(      self   :						str):
			UpperCamelCase__				:					int					=					self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_megatron_bert_model(*UpperCAmelCase_)

	def      __UpperCamelCase							(      self   :						Tuple):
			UpperCamelCase__				:					Optional[Any]					=					self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_megatron_bert_for_masked_lm(*UpperCAmelCase_)

	def      __UpperCamelCase							(      self   :						Dict):
			UpperCamelCase__				:					str					=					self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*UpperCAmelCase_)

	def      __UpperCamelCase							(      self   :						Optional[Any]):
			UpperCamelCase__				:					Optional[int]					=					self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*UpperCAmelCase_)

	def      __UpperCamelCase							(      self   :						Union[str, Any]):
			UpperCamelCase__				:					Optional[Any]					=					self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_megatron_bert_for_pretraining(*UpperCAmelCase_)

	def      __UpperCamelCase							(      self   :						str):
			UpperCamelCase__				:					str					=					self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_megatron_bert_for_question_answering(*UpperCAmelCase_)

	def      __UpperCamelCase							(      self   :						Tuple):
			UpperCamelCase__				:					List[str]					=					self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*UpperCAmelCase_)





	def      __UpperCamelCase							(      self   :						Optional[Any]):
			UpperCamelCase__				:					str					=					self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_megatron_bert_for_token_classification(*UpperCAmelCase_)







def       __UpperCAmelCase						(  lowerCamelCase_)							->     Optional[int]:
		return torch.tensor(
		    lowerCamelCase_     ,      dtype=torch.long     ,      device=lowerCamelCase_     ,      )


lowerCAmelCase__          =						1e-4







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

	@slow
	@unittest.skip('Model is not available.')
	def      __UpperCamelCase							(      self   :						Optional[int]):
			UpperCamelCase__				:					List[Any]					=					'nvidia/megatron-bert-uncased-345m'
			if "MYDIR" in os.environ:
					UpperCamelCase__				:					List[Any]					=					os.path.join(os.environ['MYDIR']  ,					UpperCAmelCase_)
			UpperCamelCase__				:					Optional[int]					=					MegatronBertModel.from_pretrained(UpperCAmelCase_)
			model.to(UpperCAmelCase_)
			model.half()
			UpperCamelCase__				:					Optional[int]					=					_long_tensor([[101, 7_110, 1_005, 1_056, 2_023, 11_333, 17_413, 1_029, 102]])
			with torch.no_grad():
					UpperCamelCase__				:					List[str]					=					model(UpperCAmelCase_)[0]
			UpperCamelCase__				:					Union[str, Any]					=					torch.Size((1, 9, 1_024))
			self.assertEqual(output.shape  ,					UpperCAmelCase_)

			UpperCamelCase__				:					Any					=					[-0.60_40, -0.25_17, -0.10_25, 0.34_20, -0.67_58, -0.00_17, -0.10_89, -0.19_90, 0.57_28]
			for ii in range(3):
					for jj in range(3):
							UpperCamelCase__				:					Tuple					=					output[0, ii, jj]
							UpperCamelCase__				:					Optional[Any]					=					expected[3 * ii + jj]
							UpperCamelCase__				:					Optional[Any]					=					'ii={} jj={} a={} b={}'.format(UpperCAmelCase_  ,					UpperCAmelCase_  ,					UpperCAmelCase_  ,					UpperCAmelCase_)
							self.assertTrue(math.isclose(UpperCAmelCase_  ,					UpperCAmelCase_  ,					rel_tol=UpperCAmelCase_  ,					abs_tol=UpperCAmelCase_)  ,					msg=UpperCAmelCase_)












'''simple docstring'''




import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import TransformeraDModel, VQDiffusionPipeline, VQDiffusionScheduler, VQModel
from diffusers.pipelines.vq_diffusion.pipeline_vq_diffusion import LearnedClassifierFreeSamplingEmbeddings
from diffusers.utils import load_numpy, slow, torch_device
from diffusers.utils.testing_utils import require_torch_gpu


lowerCAmelCase__          =						False







class       __lowercase					(unittest.TestCase    ):

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

	@property
	def      __UpperCamelCase							(      self   :						int):
			return 12

	@property
	def      __UpperCamelCase							(      self   :						Tuple):
			return 12

	@property
	def      __UpperCamelCase							(      self   :						Dict):
			return 32

	@property
	def      __UpperCamelCase							(      self   :						Optional[int]):
			torch.manual_seed(0)
			UpperCamelCase__				:					List[Any]					=					VQModel(
			    block_out_channels=[32, 64]  ,					in_channels=3  ,					out_channels=3  ,					down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']  ,					up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']  ,					latent_channels=3  ,					num_vq_embeddings=self.num_embed  ,					vq_embed_dim=3  ,					)
			return model

	@property
	def      __UpperCamelCase							(      self   :						Optional[Any]):
			UpperCamelCase__				:					Any					=					CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip')
			return tokenizer

	@property
	def      __UpperCamelCase							(      self   :						List[str]):
			torch.manual_seed(0)
			UpperCamelCase__				:					Optional[int]					=					CLIPTextConfig(
			    bos_token_id=0  ,					eos_token_id=2  ,					hidden_size=self.text_embedder_hidden_size  ,					intermediate_size=37  ,					layer_norm_eps=1e-05  ,					num_attention_heads=4  ,					num_hidden_layers=5  ,					pad_token_id=1  ,					vocab_size=1_000  ,					)
			return CLIPTextModel(UpperCAmelCase_)

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

			UpperCamelCase__				:					List[Any]					=					12
			UpperCamelCase__				:					Dict					=					12

			UpperCamelCase__				:					Union[str, Any]					=					{
			    'attention_bias': True,
			    'cross_attention_dim': 32,
			    'attention_head_dim': height * width,
			    'num_attention_heads': 1,
			    'num_vector_embeds': self.num_embed,
			    'num_embeds_ada_norm': self.num_embeds_ada_norm,
			    'norm_num_groups': 32,
			    'sample_size': width,
			    'activation_fn': 'geglu-approximate',
			}

			UpperCamelCase__				:					Tuple					=					TransformeraDModel(**UpperCAmelCase_)
			return model

	def      __UpperCamelCase							(      self   :						int):
			UpperCamelCase__				:					List[Any]					=					'cpu'

			UpperCamelCase__				:					List[str]					=					self.dummy_vqvae
			UpperCamelCase__				:					List[str]					=					self.dummy_text_encoder
			UpperCamelCase__				:					Optional[int]					=					self.dummy_tokenizer
			UpperCamelCase__				:					List[str]					=					self.dummy_transformer
			UpperCamelCase__				:					Dict					=					VQDiffusionScheduler(self.num_embed)
			UpperCamelCase__				:					List[Any]					=					LearnedClassifierFreeSamplingEmbeddings(learnable=UpperCAmelCase_)

			UpperCamelCase__				:					int					=					VQDiffusionPipeline(
			    vqvae=UpperCAmelCase_  ,					text_encoder=UpperCAmelCase_  ,					tokenizer=UpperCAmelCase_  ,					transformer=UpperCAmelCase_  ,					scheduler=UpperCAmelCase_  ,					learned_classifier_free_sampling_embeddings=UpperCAmelCase_  ,					)
			UpperCamelCase__				:					Optional[Any]					=					pipe.to(UpperCAmelCase_)
			pipe.set_progress_bar_config(disable=UpperCAmelCase_)

			UpperCamelCase__				:					Optional[Any]					=					'teddy bear playing in the pool'

			UpperCamelCase__				:					Dict					=					torch.Generator(device=UpperCAmelCase_).manual_seed(0)
			UpperCamelCase__				:					Any					=					pipe([prompt]  ,					generator=UpperCAmelCase_  ,					num_inference_steps=2  ,					output_type='np')
			UpperCamelCase__				:					Optional[Any]					=					output.images

			UpperCamelCase__				:					int					=					torch.Generator(device=UpperCAmelCase_).manual_seed(0)
			UpperCamelCase__				:					Any					=					pipe(
			    [prompt]  ,					generator=UpperCAmelCase_  ,					output_type='np'  ,					return_dict=UpperCAmelCase_  ,					num_inference_steps=2)[0]

			UpperCamelCase__				:					Optional[Any]					=					image[0, -3:, -3:, -1]
			UpperCamelCase__				:					Any					=					image_from_tuple[0, -3:, -3:, -1]

			assert image.shape == (1, 24, 24, 3)

			UpperCamelCase__				:					Any					=					np.array([0.65_51, 0.61_68, 0.50_08, 0.56_76, 0.56_59, 0.42_95, 0.60_73, 0.55_99, 0.49_92])

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





	def      __UpperCamelCase							(      self   :						Optional[int]):
			UpperCamelCase__				:					Optional[int]					=					'cpu'

			UpperCamelCase__				:					str					=					self.dummy_vqvae
			UpperCamelCase__				:					Any					=					self.dummy_text_encoder
			UpperCamelCase__				:					List[Any]					=					self.dummy_tokenizer
			UpperCamelCase__				:					Dict					=					self.dummy_transformer
			UpperCamelCase__				:					Optional[Any]					=					VQDiffusionScheduler(self.num_embed)
			UpperCamelCase__				:					Optional[Any]					=					LearnedClassifierFreeSamplingEmbeddings(
			    learnable=UpperCAmelCase_  ,					hidden_size=self.text_embedder_hidden_size  ,					length=tokenizer.model_max_length)

			UpperCamelCase__				:					str					=					VQDiffusionPipeline(
			    vqvae=UpperCAmelCase_  ,					text_encoder=UpperCAmelCase_  ,					tokenizer=UpperCAmelCase_  ,					transformer=UpperCAmelCase_  ,					scheduler=UpperCAmelCase_  ,					learned_classifier_free_sampling_embeddings=UpperCAmelCase_  ,					)
			UpperCamelCase__				:					str					=					pipe.to(UpperCAmelCase_)
			pipe.set_progress_bar_config(disable=UpperCAmelCase_)

			UpperCamelCase__				:					List[Any]					=					'teddy bear playing in the pool'

			UpperCamelCase__				:					Union[str, Any]					=					torch.Generator(device=UpperCAmelCase_).manual_seed(0)
			UpperCamelCase__				:					Any					=					pipe([prompt]  ,					generator=UpperCAmelCase_  ,					num_inference_steps=2  ,					output_type='np')
			UpperCamelCase__				:					int					=					output.images

			UpperCamelCase__				:					List[Any]					=					torch.Generator(device=UpperCAmelCase_).manual_seed(0)
			UpperCamelCase__				:					Optional[Any]					=					pipe(
			    [prompt]  ,					generator=UpperCAmelCase_  ,					output_type='np'  ,					return_dict=UpperCAmelCase_  ,					num_inference_steps=2)[0]

			UpperCamelCase__				:					Union[str, Any]					=					image[0, -3:, -3:, -1]
			UpperCamelCase__				:					Dict					=					image_from_tuple[0, -3:, -3:, -1]

			assert image.shape == (1, 24, 24, 3)

			UpperCamelCase__				:					str					=					np.array([0.66_93, 0.60_75, 0.49_59, 0.57_01, 0.55_83, 0.43_33, 0.61_71, 0.56_84, 0.49_88])

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







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

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





	def      __UpperCamelCase							(      self   :						List[Any]):
			UpperCamelCase__				:					Optional[Any]					=					load_numpy(
			    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
			    '/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy')

			UpperCamelCase__				:					List[Any]					=					VQDiffusionPipeline.from_pretrained('microsoft/vq-diffusion-ithq')
			UpperCamelCase__				:					Any					=					pipeline.to(UpperCAmelCase_)
			pipeline.set_progress_bar_config(disable=UpperCAmelCase_)

			# requires GPU generator for gumbel softmax
			# don't use GPU generator in tests though
			UpperCamelCase__				:					Optional[int]					=					torch.Generator(device=UpperCAmelCase_).manual_seed(0)
			UpperCamelCase__				:					int					=					pipeline(
			    'teddy bear playing in the pool'  ,					num_images_per_prompt=1  ,					generator=UpperCAmelCase_  ,					output_type='np'  ,					)

			UpperCamelCase__				:					int					=					output.images[0]

			assert image.shape == (256, 256, 3)
			assert np.abs(expected_image - image).max() < 2.0




