


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





							"""simple docstring"""







							def __init__(      self			,						lowerCAmelCase__			,						lowerCAmelCase__=1_3			,						lowerCAmelCase__=7			,						lowerCAmelCase__=True			,						lowerCAmelCase__=True			,						lowerCAmelCase__=True			,						lowerCAmelCase__=True			,						lowerCAmelCase__=9_9			,						lowerCAmelCase__=3_2			,						lowerCAmelCase__=5			,						lowerCAmelCase__=4			,						lowerCAmelCase__=3_7			,						lowerCAmelCase__="gelu"			,						lowerCAmelCase__=0.1			,						lowerCAmelCase__=0.1			,						lowerCAmelCase__=5_1_2			,						lowerCAmelCase__=1_6			,						lowerCAmelCase__=2			,						lowerCAmelCase__=0.02			,						lowerCAmelCase__=4			,						):





										"""simple docstring"""




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






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     str					=					ids_tensor([self.batch_size, self.seq_length]			,						self.vocab_size       )

										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					None
										if self.use_attention_mask:
													SCREAMING_SNAKE_CASE_		:     Dict					=					random_attention_mask([self.batch_size, self.seq_length]       )

										SCREAMING_SNAKE_CASE_		:     Optional[int]					=					None
										if self.use_token_type_ids:
													SCREAMING_SNAKE_CASE_		:     List[str]					=					ids_tensor([self.batch_size, self.seq_length]			,						self.type_vocab_size       )

										SCREAMING_SNAKE_CASE_		:     str					=					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=lowerCAmelCase__			,						initializer_range=self.initializer_range			,						)

										return config, input_ids, token_type_ids, attention_mask






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Any					=					self.prepare_config_and_inputs()
										SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		,      SCREAMING_SNAKE_CASE_		:     Any					=					config_and_inputs
										SCREAMING_SNAKE_CASE_		:     str					=					{'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': attention_mask}
										return config, inputs_dict







@require_flax
class 	__lowercase    (__SCREAMING_SNAKE_CASE	,       unittest.TestCase  ):





							"""simple docstring"""







							_UpperCAmelCase      =						(
							    (
							        FlaxAlbertModel,
							        FlaxAlbertForPreTraining,
							        FlaxAlbertForMaskedLM,
							        FlaxAlbertForMultipleChoice,
							        FlaxAlbertForQuestionAnswering,
							        FlaxAlbertForSequenceClassification,
							        FlaxAlbertForTokenClassification,
							        FlaxAlbertForQuestionAnswering,
							    )
							    if is_flax_available()
							    else ()
							)






							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     List[str]					=					FlaxAlbertModelTester(self       )






							@slow
							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




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







@require_flax
class 	__lowercase    (unittest.TestCase  ):





							"""simple docstring"""







							@slow
							def 					UpperCamelCase__     (      self       ):





										"""simple docstring"""




										SCREAMING_SNAKE_CASE_		:     Tuple					=					FlaxAlbertModel.from_pretrained('albert-base-v2'       )
										SCREAMING_SNAKE_CASE_		:     List[Any]					=					np.array([[0, 3_4_5, 2_3_2, 3_2_8, 7_4_0, 1_4_0, 1_6_9_5, 6_9, 6_0_7_8, 1_5_8_8, 2]]       )
										SCREAMING_SNAKE_CASE_		:     Any					=					np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]       )
										SCREAMING_SNAKE_CASE_		:     Any					=					model(lowerCAmelCase__			,						attention_mask=lowerCAmelCase__       )[0]
										SCREAMING_SNAKE_CASE_		:     str					=					(1, 1_1, 7_6_8)
										self.assertEqual(output.shape			,						lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Tuple					=					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]			,						lowerCAmelCase__			,						atol=1E-4       )       )












"""simple docstring"""



import gc
import unittest

import numpy as np
import torch

from diffusers import (
    AudioDiffusionPipeline,
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    DiffusionPipeline,
    Mel,
    UNetaDConditionModel,
    UNetaDModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()






class 						__snake_case   (						unittest.TestCase	):





 """simple docstring"""


 def    UpperCamelCase__(       self    ):



     '''simple docstring'''




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

 @property
 def    UpperCamelCase__(       self    ):



     '''simple docstring'''




     torch.manual_seed(0    )
     __A       :		List[Any]							=	UNetaDModel(
         sample_size=(32, 64)   ,		in_channels=1   ,		out_channels=1   ,		layers_per_block=2   ,		block_out_channels=(128, 128)   ,		down_block_types=('''AttnDownBlock2D''', '''DownBlock2D''')   ,		up_block_types=('''UpBlock2D''', '''AttnUpBlock2D''')   ,		)
     return model

 @property
 def    UpperCamelCase__(       self    ):



     '''simple docstring'''




     torch.manual_seed(0    )
     __A       :		str							=	UNetaDConditionModel(
         sample_size=(64, 32)   ,		in_channels=1   ,		out_channels=1   ,		layers_per_block=2   ,		block_out_channels=(128, 128)   ,		down_block_types=('''CrossAttnDownBlock2D''', '''DownBlock2D''')   ,		up_block_types=('''UpBlock2D''', '''CrossAttnUpBlock2D''')   ,		cross_attention_dim=10   ,		)
     return model

 @property
 def    UpperCamelCase__(       self    ):



     '''simple docstring'''




     torch.manual_seed(0    )
     __A       :		int							=	AutoencoderKL(
         sample_size=(128, 64)   ,		in_channels=1   ,		out_channels=1   ,		latent_channels=1   ,		layers_per_block=2   ,		block_out_channels=(128, 128)   ,		down_block_types=('''DownEncoderBlock2D''', '''DownEncoderBlock2D''')   ,		up_block_types=('''UpDecoderBlock2D''', '''UpDecoderBlock2D''')   ,		)
     __A       :		int							=	UNetaDModel(
         sample_size=(64, 32)   ,		in_channels=1   ,		out_channels=1   ,		layers_per_block=2   ,		block_out_channels=(128, 128)   ,		down_block_types=('''AttnDownBlock2D''', '''DownBlock2D''')   ,		up_block_types=('''UpBlock2D''', '''AttnUpBlock2D''')   ,		)
     return vqvae, unet







 @slow
 def    UpperCamelCase__(       self    ):



     '''simple docstring'''




     __A       :		str							=	'''cpu'''  # ensure determinism for the device-dependent torch.Generator
     __A       :		Union[str, Any]							=	Mel(
         x_res=self.dummy_unet.config.sample_size[1]   ,		y_res=self.dummy_unet.config.sample_size[0]   ,		)

     __A       :		List[Any]							=	DDPMScheduler()
     __A       :		Optional[Any]							=	AudioDiffusionPipeline(vqvae=__lowerCamelCase   ,		unet=self.dummy_unet   ,		mel=__lowerCamelCase   ,		scheduler=__lowerCamelCase    )
     __A       :		Tuple							=	pipe.to(__lowerCamelCase    )
     pipe.set_progress_bar_config(disable=__lowerCamelCase    )

     __A       :		Tuple							=	torch.Generator(device=__lowerCamelCase    ).manual_seed(42    )
     __A       :		List[str]							=	pipe(generator=__lowerCamelCase   ,		steps=4    )
     __A       :		Union[str, Any]							=	output.audios[0]
     __A       :		str							=	output.images[0]

     __A       :		List[Any]							=	torch.Generator(device=__lowerCamelCase    ).manual_seed(42    )
     __A       :		Any							=	pipe(generator=__lowerCamelCase   ,		steps=4   ,		return_dict=__lowerCamelCase    )
     __A       :		int							=	output[0][0]

     assert audio.shape == (1, (self.dummy_unet.config.sample_size[1] - 1) * mel.hop_length)
     assert (
         image.height == self.dummy_unet.config.sample_size[0]
         and image.width == self.dummy_unet.config.sample_size[1]
     )
     __A       :		Tuple							=	np.frombuffer(image.tobytes()   ,		dtype='''uint8'''    )[:10]
     __A       :		str							=	np.frombuffer(image_from_tuple.tobytes()   ,		dtype='''uint8'''    )[:10]
     __A       :		Any							=	np.array([69, 255, 255, 255, 0, 0, 77, 181, 12, 127]    )

     assert np.abs(image_slice.flatten() - expected_slice    ).max() == 0
     assert np.abs(image_from_tuple_slice.flatten() - expected_slice    ).max() == 0

     __A       :		int							=	Mel(
         x_res=self.dummy_vqvae_and_unet[0].config.sample_size[1]   ,		y_res=self.dummy_vqvae_and_unet[0].config.sample_size[0]   ,		)

     __A       :		Any							=	DDIMScheduler()
     __A       :		Optional[Any]							=	self.dummy_vqvae_and_unet
     __A       :		int							=	AudioDiffusionPipeline(
         vqvae=self.dummy_vqvae_and_unet[0]   ,		unet=dummy_vqvae_and_unet[1]   ,		mel=__lowerCamelCase   ,		scheduler=__lowerCamelCase    )
     __A       :		Any							=	pipe.to(__lowerCamelCase    )
     pipe.set_progress_bar_config(disable=__lowerCamelCase    )

     np.random.seed(0    )
     __A       :		str							=	np.random.uniform(-1   ,		1   ,		((dummy_vqvae_and_unet[0].config.sample_size[1] - 1) * mel.hop_length,)    )
     __A       :		Union[str, Any]							=	torch.Generator(device=__lowerCamelCase    ).manual_seed(42    )
     __A       :		int							=	pipe(raw_audio=__lowerCamelCase   ,		generator=__lowerCamelCase   ,		start_step=5   ,		steps=10    )
     __A       :		Any							=	output.images[0]

     assert (
         image.height == self.dummy_vqvae_and_unet[0].config.sample_size[0]
         and image.width == self.dummy_vqvae_and_unet[0].config.sample_size[1]
     )
     __A       :		str							=	np.frombuffer(image.tobytes()   ,		dtype='''uint8'''    )[:10]
     __A       :		Optional[Any]							=	np.array([120, 117, 110, 109, 138, 167, 138, 148, 132, 121]    )

     assert np.abs(image_slice.flatten() - expected_slice    ).max() == 0

     __A       :		Union[str, Any]							=	self.dummy_unet_condition
     __A       :		List[str]							=	AudioDiffusionPipeline(
         vqvae=self.dummy_vqvae_and_unet[0]   ,		unet=__lowerCamelCase   ,		mel=__lowerCamelCase   ,		scheduler=__lowerCamelCase    )
     __A       :		Dict							=	pipe.to(__lowerCamelCase    )
     pipe.set_progress_bar_config(disable=__lowerCamelCase    )

     np.random.seed(0    )
     __A       :		int							=	torch.rand((1, 1, 10)    )
     __A       :		Optional[Any]							=	pipe(generator=__lowerCamelCase   ,		encoding=__lowerCamelCase    )
     __A       :		str							=	output.images[0]
     __A       :		Optional[Any]							=	np.frombuffer(image.tobytes()   ,		dtype='''uint8'''    )[:10]
     __A       :		Tuple							=	np.array([107, 103, 120, 127, 142, 122, 113, 122, 97, 111]    )

     assert np.abs(image_slice.flatten() - expected_slice    ).max() == 0






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





 """simple docstring"""


 def    UpperCamelCase__(       self    ):



     '''simple docstring'''




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







 def    UpperCamelCase__(       self    ):



     '''simple docstring'''




     __A       :		int							=	torch_device

     __A       :		List[Any]							=	DiffusionPipeline.from_pretrained('''teticio/audio-diffusion-ddim-256'''    )
     __A       :		List[Any]							=	pipe.to(__lowerCamelCase    )
     pipe.set_progress_bar_config(disable=__lowerCamelCase    )

     __A       :		Dict							=	torch.Generator(device=__lowerCamelCase    ).manual_seed(42    )
     __A       :		str							=	pipe(generator=__lowerCamelCase    )
     __A       :		Dict							=	output.audios[0]
     __A       :		List[Any]							=	output.images[0]

     assert audio.shape == (1, (pipe.unet.config.sample_size[1] - 1) * pipe.mel.hop_length)
     assert image.height == pipe.unet.config.sample_size[0] and image.width == pipe.unet.config.sample_size[1]
     __A       :		Optional[int]							=	np.frombuffer(image.tobytes()   ,		dtype='''uint8'''    )[:10]
     __A       :		Optional[Any]							=	np.array([151, 167, 154, 144, 122, 134, 121, 105, 70, 26]    )

     assert np.abs(image_slice.flatten() - expected_slice    ).max() == 0




