





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






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


							'''simple docstring'''



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






			@property
			def       lowerCamelCase  (self   )					->       int:


							'''simple docstring'''



							return 12






			@property
			def       lowerCamelCase  (self   )					->       Union[str, Any]:


							'''simple docstring'''



							return 12






			@property
			def       lowerCamelCase  (self   )					->       int:


							'''simple docstring'''



							return 32






			@property
			def       lowerCamelCase  (self   )					->       Dict:


							'''simple docstring'''



							torch.manual_seed(0   )
							snake_case_						:   int           =  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       lowerCamelCase  (self   )					->       int:


							'''simple docstring'''



							snake_case_						:   Optional[int]           =  CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''   )
							return tokenizer






			@property
			def       lowerCamelCase  (self   )					->       List[str]:


							'''simple docstring'''



							torch.manual_seed(0   )
							snake_case_						:   Optional[Any]           =  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=1000 ,							)
							return CLIPTextModel(snake_case_   )






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


							'''simple docstring'''



							torch.manual_seed(0   )

							snake_case_						:   Union[str, Any]           =  12
							snake_case_						:   Dict           =  12

							snake_case_						:   Optional[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""",
							}

							snake_case_						:   Union[str, Any]           =  TransformeraDModel(**snake_case_   )
							return model






			def       lowerCamelCase  (self   )					->       Optional[Any]:


							'''simple docstring'''



							snake_case_						:   List[Any]           =  """cpu"""

							snake_case_						:   Union[str, Any]           =  self.dummy_vqvae
							snake_case_						:   Dict           =  self.dummy_text_encoder
							snake_case_						:   Union[str, Any]           =  self.dummy_tokenizer
							snake_case_						:   Any           =  self.dummy_transformer
							snake_case_						:   Optional[Any]           =  VQDiffusionScheduler(self.num_embed   )
							snake_case_						:   Optional[int]           =  LearnedClassifierFreeSamplingEmbeddings(learnable=snake_case_   )

							snake_case_						:   int           =  VQDiffusionPipeline(
							    vqvae=snake_case_ ,							text_encoder=snake_case_ ,							tokenizer=snake_case_ ,							transformer=snake_case_ ,							scheduler=snake_case_ ,							learned_classifier_free_sampling_embeddings=snake_case_ ,							)
							snake_case_						:   Dict           =  pipe.to(snake_case_   )
							pipe.set_progress_bar_config(disable=snake_case_   )

							snake_case_						:   Tuple           =  """teddy bear playing in the pool"""

							snake_case_						:   Tuple           =  torch.Generator(device=snake_case_   ).manual_seed(0   )
							snake_case_						:   List[str]           =  pipe([prompt] ,							generator=snake_case_ ,							num_inference_steps=2 ,							output_type='''np'''   )
							snake_case_						:   Optional[int]           =  output.images

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

							snake_case_						:   List[Any]           =  image[0, -3:, -3:, -1]
							snake_case_						:   Dict           =  image_from_tuple[0, -3:, -3:, -1]

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

							snake_case_						:   Tuple           =  np.array([0.6_551, 0.6_168, 0.5_008, 0.5_676, 0.5_659, 0.4_295, 0.6_073, 0.5_599, 0.4_992]   )

							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       lowerCamelCase  (self   )					->       int:


							'''simple docstring'''



							snake_case_						:   Tuple           =  """cpu"""

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

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

							snake_case_						:   Dict           =  """teddy bear playing in the pool"""

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

							snake_case_						:   Dict           =  torch.Generator(device=snake_case_   ).manual_seed(0   )
							snake_case_						:   List[str]           =  pipe(
							    [prompt] ,							generator=snake_case_ ,							output_type='''np''' ,							return_dict=snake_case_ ,							num_inference_steps=2   )[0]

							snake_case_						:   List[Any]           =  image[0, -3:, -3:, -1]
							snake_case_						:   List[Any]           =  image_from_tuple[0, -3:, -3:, -1]

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

							snake_case_						:   Optional[int]           =  np.array([0.6_693, 0.6_075, 0.4_959, 0.5_701, 0.5_583, 0.4_333, 0.6_171, 0.5_684, 0.4_988]   )

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






			def       lowerCamelCase  (self   )					->       Any:


							'''simple docstring'''



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






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


							'''simple docstring'''



							snake_case_						:   Dict           =  load_numpy(
							    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
							    '''/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy'''   )

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

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

							snake_case_						:   Optional[Any]           =  output.images[0]

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







import unittest

import numpy as np

from transformers import AlbertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
	import jax.numpy as jnp

	from transformers.models.albert.modeling_flax_albert import (
	    FlaxAlbertForMaskedLM,
	    FlaxAlbertForMultipleChoice,
	    FlaxAlbertForPreTraining,
	    FlaxAlbertForQuestionAnswering,
	    FlaxAlbertForSequenceClassification,
	    FlaxAlbertForTokenClassification,
	    FlaxAlbertModel,
	)



class  __lowerCAmelCase	(      unittest.TestCase					):






			def __init__(self ,							__magic_name__ ,							__magic_name__=13 ,							__magic_name__=7 ,							__magic_name__=True ,							__magic_name__=True ,							__magic_name__=True ,							__magic_name__=True ,							__magic_name__=99 ,							__magic_name__=32 ,							__magic_name__=5 ,							__magic_name__=4 ,							__magic_name__=37 ,							__magic_name__="gelu" ,							__magic_name__=0.1 ,							__magic_name__=0.1 ,							__magic_name__=512 ,							__magic_name__=16 ,							__magic_name__=2 ,							__magic_name__=0.02 ,							__magic_name__=4 ,							)					->       Tuple:


							'''simple docstring'''



							snake_case_						:   Union[str, Any]           =  parent
							snake_case_						:   Optional[Any]           =  batch_size
							snake_case_						:   List[Any]           =  seq_length
							snake_case_						:   Tuple           =  is_training
							snake_case_						:   List[str]           =  use_attention_mask
							snake_case_						:   Any           =  use_token_type_ids
							snake_case_						:   Dict           =  use_labels
							snake_case_						:   Optional[Any]           =  vocab_size
							snake_case_						:   Dict           =  hidden_size
							snake_case_						:   List[Any]           =  num_hidden_layers
							snake_case_						:   Union[str, Any]           =  num_attention_heads
							snake_case_						:   Any           =  intermediate_size
							snake_case_						:   Optional[int]           =  hidden_act
							snake_case_						:   Optional[int]           =  hidden_dropout_prob
							snake_case_						:   Optional[Any]           =  attention_probs_dropout_prob
							snake_case_						:   Optional[int]           =  max_position_embeddings
							snake_case_						:   Optional[int]           =  type_vocab_size
							snake_case_						:   List[Any]           =  type_sequence_label_size
							snake_case_						:   Dict           =  initializer_range
							snake_case_						:   Dict           =  num_choices






			def       lowerCamelCase  (self   )					->       Tuple:


							'''simple docstring'''



							snake_case_						:   int           =  ids_tensor([self.batch_size, self.seq_length] ,							self.vocab_size   )

							snake_case_						:   Any           =  None
							if self.use_attention_mask:
											snake_case_						:   List[Any]           =  random_attention_mask([self.batch_size, self.seq_length]   )

							snake_case_						:   List[Any]           =  None
							if self.use_token_type_ids:
											snake_case_						:   Dict           =  ids_tensor([self.batch_size, self.seq_length] ,							self.type_vocab_size   )

							snake_case_						:   List[Any]           =  AlbertConfig(
							    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 ,							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=__magic_name__ ,							initializer_range=self.initializer_range ,							)

							return config, input_ids, token_type_ids, attention_mask






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


							'''simple docstring'''



							snake_case_						:   Tuple           =  self.prepare_config_and_inputs()
							snake_case_      ,							snake_case_      ,							snake_case_      ,							snake_case_						:   Optional[int]           =  config_and_inputs
							snake_case_						:   int           =  {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': attention_mask}
							return config, inputs_dict




@require_flax
class  __lowerCAmelCase	(      _a,					unittest.TestCase					):
			lowerCamelCase_					:					Optional[int]	      =				(
			    (
			        FlaxAlbertModel,
			        FlaxAlbertForPreTraining,
			        FlaxAlbertForMaskedLM,
			        FlaxAlbertForMultipleChoice,
			        FlaxAlbertForQuestionAnswering,
			        FlaxAlbertForSequenceClassification,
			        FlaxAlbertForTokenClassification,
			        FlaxAlbertForQuestionAnswering,
			    )
			    if is_flax_available()
			    else ()
			)






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


							'''simple docstring'''



							snake_case_						:   Optional[Any]           =  FlaxAlbertModelTester(self   )






			@slow
			def       lowerCamelCase  (self   )					->       Tuple:


							'''simple docstring'''



							for model_class_name in self.all_model_classes:
											snake_case_						:   Dict           =  model_class_name.from_pretrained('''albert-base-v2'''   )
											snake_case_						:   Union[str, Any]           =  model(np.ones((1, 1)   )   )
											self.assertIsNotNone(__magic_name__   )




@require_flax
class  __lowerCAmelCase	(      unittest.TestCase					):






			@slow
			def       lowerCamelCase  (self   )					->       Optional[Any]:


							'''simple docstring'''



							snake_case_						:   Optional[Any]           =  FlaxAlbertModel.from_pretrained('''albert-base-v2'''   )
							snake_case_						:   Optional[int]           =  np.array([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]   )
							snake_case_						:   Dict           =  np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]   )
							snake_case_						:   Union[str, Any]           =  model(__magic_name__ ,							attention_mask=__magic_name__   )[0]
							snake_case_						:   Tuple           =  (1, 11, 768)
							self.assertEqual(output.shape ,							__magic_name__   )
							snake_case_						:   str           =  np.array(
							    [[[-0.6_513, 1.5_035, -0.2_766], [-0.6_515, 1.5_046, -0.2_780], [-0.6_512, 1.5_049, -0.2_784]]]   )

							self.assertTrue(jnp.allclose(output[:, 1:4, 1:4] ,							__magic_name__ ,							atol=1e-4   )   )
