






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



	'''simple docstring'''




	def 				_UpperCamelCase     (			self   )			->       List[str]:
								# clean up the VRAM after each test
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()




	@property
	def 				_UpperCamelCase     (			self   )			->       Dict:
								torch.manual_seed(0   )
								snake_case_           =				UNetaDModel(
								    sample_size=(32, 64)			,     in_channels=1			,     out_channels=1			,     layers_per_block=2			,     block_out_channels=(1_28, 1_28)			,     down_block_types=('AttnDownBlock2D', 'DownBlock2D')			,     up_block_types=('UpBlock2D', 'AttnUpBlock2D')			,     )
								return model




	@property
	def 				_UpperCamelCase     (			self   )			->       List[str]:
								torch.manual_seed(0   )
								snake_case_           =				UNetaDConditionModel(
								    sample_size=(64, 32)			,     in_channels=1			,     out_channels=1			,     layers_per_block=2			,     block_out_channels=(1_28, 1_28)			,     down_block_types=('CrossAttnDownBlock2D', 'DownBlock2D')			,     up_block_types=('UpBlock2D', 'CrossAttnUpBlock2D')			,     cross_attention_dim=10			,     )
								return model




	@property
	def 				_UpperCamelCase     (			self   )			->       Dict:
								torch.manual_seed(0   )
								snake_case_           =				AutoencoderKL(
								    sample_size=(1_28, 64)			,     in_channels=1			,     out_channels=1			,     latent_channels=1			,     layers_per_block=2			,     block_out_channels=(1_28, 1_28)			,     down_block_types=('DownEncoderBlock2D', 'DownEncoderBlock2D')			,     up_block_types=('UpDecoderBlock2D', 'UpDecoderBlock2D')			,     )
								snake_case_           =				UNetaDModel(
								    sample_size=(64, 32)			,     in_channels=1			,     out_channels=1			,     layers_per_block=2			,     block_out_channels=(1_28, 1_28)			,     down_block_types=('AttnDownBlock2D', 'DownBlock2D')			,     up_block_types=('UpBlock2D', 'AttnUpBlock2D')			,     )
								return vqvae, unet






	@slow
	def 				_UpperCamelCase     (			self   )			->       Dict:
								snake_case_           =				'cpu'  # ensure determinism for the device-dependent torch.Generator
								snake_case_           =				Mel(
								    x_res=self.dummy_unet.config.sample_size[1]			,     y_res=self.dummy_unet.config.sample_size[0]			,     )

								snake_case_           =				DDPMScheduler()
								snake_case_           =				AudioDiffusionPipeline(vqvae=a			,     unet=self.dummy_unet			,     mel=a			,     scheduler=a   )
								snake_case_           =				pipe.to(a   )
								pipe.set_progress_bar_config(disable=a   )

								snake_case_           =				torch.Generator(device=a   ).manual_seed(42   )
								snake_case_           =				pipe(generator=a			,     steps=4   )
								snake_case_           =				output.audios[0]
								snake_case_           =				output.images[0]

								snake_case_           =				torch.Generator(device=a   ).manual_seed(42   )
								snake_case_           =				pipe(generator=a			,     steps=4			,     return_dict=a   )
								snake_case_           =				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]
								)
								snake_case_           =				np.frombuffer(image.tobytes()			,     dtype='uint8'   )[:10]
								snake_case_           =				np.frombuffer(image_from_tuple.tobytes()			,     dtype='uint8'   )[:10]
								snake_case_           =				np.array([69, 2_55, 2_55, 2_55, 0, 0, 77, 1_81, 12, 1_27]   )

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

								snake_case_           =				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]			,     )

								snake_case_           =				DDIMScheduler()
								snake_case_           =				self.dummy_vqvae_and_unet
								snake_case_           =				AudioDiffusionPipeline(
								    vqvae=self.dummy_vqvae_and_unet[0]			,     unet=dummy_vqvae_and_unet[1]			,     mel=a			,     scheduler=a   )
								snake_case_           =				pipe.to(a   )
								pipe.set_progress_bar_config(disable=a   )

								np.random.seed(0   )
								snake_case_           =				np.random.uniform(-1			,     1			,     ((dummy_vqvae_and_unet[0].config.sample_size[1] - 1) * mel.hop_length,)   )
								snake_case_           =				torch.Generator(device=a   ).manual_seed(42   )
								snake_case_           =				pipe(raw_audio=a			,     generator=a			,     start_step=5			,     steps=10   )
								snake_case_           =				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]
								)
								snake_case_           =				np.frombuffer(image.tobytes()			,     dtype='uint8'   )[:10]
								snake_case_           =				np.array([1_20, 1_17, 1_10, 1_09, 1_38, 1_67, 1_38, 1_48, 1_32, 1_21]   )

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

								snake_case_           =				self.dummy_unet_condition
								snake_case_           =				AudioDiffusionPipeline(
								    vqvae=self.dummy_vqvae_and_unet[0]			,     unet=a			,     mel=a			,     scheduler=a   )
								snake_case_           =				pipe.to(a   )
								pipe.set_progress_bar_config(disable=a   )

								np.random.seed(0   )
								snake_case_           =				torch.rand((1, 1, 10)   )
								snake_case_           =				pipe(generator=a			,     encoding=a   )
								snake_case_           =				output.images[0]
								snake_case_           =				np.frombuffer(image.tobytes()			,     dtype='uint8'   )[:10]
								snake_case_           =				np.array([1_07, 1_03, 1_20, 1_27, 1_42, 1_22, 1_13, 1_22, 97, 1_11]   )

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






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



	'''simple docstring'''




	def 				_UpperCamelCase     (			self   )			->       Tuple:
								# clean up the VRAM after each test
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()






	def 				_UpperCamelCase     (			self   )			->       Union[str, Any]:
								snake_case_           =				torch_device

								snake_case_           =				DiffusionPipeline.from_pretrained('teticio/audio-diffusion-ddim-256'   )
								snake_case_           =				pipe.to(a   )
								pipe.set_progress_bar_config(disable=a   )

								snake_case_           =				torch.Generator(device=a   ).manual_seed(42   )
								snake_case_           =				pipe(generator=a   )
								snake_case_           =				output.audios[0]
								snake_case_           =				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]
								snake_case_           =				np.frombuffer(image.tobytes()			,     dtype='uint8'   )[:10]
								snake_case_           =				np.array([1_51, 1_67, 1_54, 1_44, 1_22, 1_34, 1_21, 1_05, 70, 26]   )

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













import unittest

import numpy as np

from transformers import MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING, TF_MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING
from transformers.pipelines import AudioClassificationPipeline, pipeline
from transformers.testing_utils import (
    is_pipeline_test,
    nested_simplify,
    require_tf,
    require_torch,
    require_torchaudio,
    slow,
)

from .test_pipelines_common import ANY




@is_pipeline_test
class 					UpperCamelCase_ (     unittest.TestCase							):



	'''simple docstring'''




	lowerCAmelCase						=					MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING
	lowerCAmelCase						=					TF_MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING




	def 				_UpperCamelCase     (			self			,     a			,     a			,     a   )			->       Union[str, Any]:
								snake_case_           =				AudioClassificationPipeline(model=a			,     feature_extractor=a   )

								# test with a raw waveform
								snake_case_           =				np.zeros((3_40_00,)   )
								snake_case_           =				np.zeros((1_40_00,)   )
								return audio_classifier, [audioa, audio]




	def 				_UpperCamelCase     (			self			,     a			,     a   )			->       Tuple:
								snake_case_     ,						snake_case_           =				examples
								snake_case_           =				audio_classifier(a   )
								# by default a model is initialized with num_labels=2
								self.assertEqual(
								    a			,     [
								        {'score': ANY(a   ), 'label': ANY(a   )},
								        {'score': ANY(a   ), 'label': ANY(a   )},
								    ]			,     )
								snake_case_           =				audio_classifier(a			,     top_k=1   )
								self.assertEqual(
								    a			,     [
								        {'score': ANY(a   ), 'label': ANY(a   )},
								    ]			,     )

								self.run_torchaudio(a   )




	@require_torchaudio
	def 				_UpperCamelCase     (			self			,     a   )			->       List[str]:
								import datasets

								# test with a local file
								snake_case_           =				datasets.load_dataset('hf-internal-testing/librispeech_asr_dummy'			,     'clean'			,     split='validation'   )
								snake_case_           =				dataset[0]['audio']['array']
								snake_case_           =				audio_classifier(a   )
								self.assertEqual(
								    a			,     [
								        {'score': ANY(a   ), 'label': ANY(a   )},
								        {'score': ANY(a   ), 'label': ANY(a   )},
								    ]			,     )




	@require_torch
	def 				_UpperCamelCase     (			self   )			->       Dict:
								snake_case_           =				'anton-l/wav2vec2-random-tiny-classifier'

								snake_case_           =				pipeline('audio-classification'			,     model=a   )

								snake_case_           =				np.ones((80_00,)   )
								snake_case_           =				audio_classifier(a			,     top_k=4   )

								snake_case_           =				[
								    {'score': 0.0_842, 'label': 'no'},
								    {'score': 0.0_838, 'label': 'up'},
								    {'score': 0.0_837, 'label': 'go'},
								    {'score': 0.0_834, 'label': 'right'},
								]
								snake_case_           =				[
								    {'score': 0.0_845, 'label': 'stop'},
								    {'score': 0.0_844, 'label': 'on'},
								    {'score': 0.0_841, 'label': 'right'},
								    {'score': 0.0_834, 'label': 'left'},
								]
								self.assertIn(nested_simplify(a			,     decimals=4   )			,     [EXPECTED_OUTPUT, EXPECTED_OUTPUT_PT_2]   )

								snake_case_           =				{'array': np.ones((80_00,)   ), 'sampling_rate': audio_classifier.feature_extractor.sampling_rate}
								snake_case_           =				audio_classifier(a			,     top_k=4   )
								self.assertIn(nested_simplify(a			,     decimals=4   )			,     [EXPECTED_OUTPUT, EXPECTED_OUTPUT_PT_2]   )




	@require_torch
	@slow
	def 				_UpperCamelCase     (			self   )			->       str:
								import datasets

								snake_case_           =				'superb/wav2vec2-base-superb-ks'

								snake_case_           =				pipeline('audio-classification'			,     model=a   )
								snake_case_           =				datasets.load_dataset('anton-l/superb_dummy'			,     'ks'			,     split='test'   )

								snake_case_           =				np.array(dataset[3]['speech']			,     dtype=np.floataa   )
								snake_case_           =				audio_classifier(a			,     top_k=4   )
								self.assertEqual(
								    nested_simplify(a			,     decimals=3   )			,     [
								        {'score': 0.981, 'label': 'go'},
								        {'score': 0.007, 'label': 'up'},
								        {'score': 0.006, 'label': '_unknown_'},
								        {'score': 0.001, 'label': 'down'},
								    ]			,     )






	@require_tf
	@unittest.skip('Audio classification is not implemented for TF'   )
	def 				_UpperCamelCase     (			self   )			->       Optional[int]:
								pass





