

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 						a__							(       unittest.TestCase		):
			A__    :       Dict        =      MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING
			A__    :       Optional[int]        =      TF_MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING



			def __SCREAMING_SNAKE_CASE							(						self				,						UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase   )			->   List[Any]:
								__a		 =       AudioClassificationPipeline(model=UpperCAmelCase				,						feature_extractor=UpperCAmelCase   )

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



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

								self.run_torchaudio(UpperCAmelCase   )



			@require_torchaudio
			def __SCREAMING_SNAKE_CASE							(						self				,						UpperCAmelCase   )			->   List[Any]:
								import datasets

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



			@require_torch
			def __SCREAMING_SNAKE_CASE							(						self   )			->   Any:
								__a		 =       'anton-l/wav2vec2-random-tiny-classifier'

								__a		 =       pipeline('audio-classification'				,						model=UpperCAmelCase   )

								__a		 =       np.ones((8_0_0_0,)   )
								__a		 =       audio_classifier(UpperCAmelCase				,						top_k=4   )

								__a		 =       [
								    {'score': 0.0_842, 'label': 'no'},
								    {'score': 0.0_838, 'label': 'up'},
								    {'score': 0.0_837, 'label': 'go'},
								    {'score': 0.0_834, 'label': 'right'},
								]
								__a		 =       [
								    {'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(UpperCAmelCase				,						decimals=4   )				,						[EXPECTED_OUTPUT, EXPECTED_OUTPUT_PT_2]   )

								__a		 =       {'array': np.ones((8_0_0_0,)   ), 'sampling_rate': audio_classifier.feature_extractor.sampling_rate}
								__a		 =       audio_classifier(UpperCAmelCase				,						top_k=4   )
								self.assertIn(nested_simplify(UpperCAmelCase				,						decimals=4   )				,						[EXPECTED_OUTPUT, EXPECTED_OUTPUT_PT_2]   )



			@require_torch
			@slow
			def __SCREAMING_SNAKE_CASE							(						self   )			->   List[Any]:
								import datasets

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

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

								__a		 =       np.array(dataset[3]['speech']				,						dtype=np.floataa   )
								__a		 =       audio_classifier(UpperCAmelCase				,						top_k=4   )
								self.assertEqual(
								    nested_simplify(UpperCAmelCase				,						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 __SCREAMING_SNAKE_CASE							(						self   )			->   Any:
								pass





import tempfile

import numpy as np
import torch
from transformers import AutoTokenizer, TaEncoderModel

from diffusers import DDPMScheduler, UNetaDConditionModel
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.pipelines.deepfloyd_if import IFWatermarker
from diffusers.utils.testing_utils import torch_device

from ..test_pipelines_common import to_np


class 						a__							:



			def __SCREAMING_SNAKE_CASE							(						self   )			->   Optional[Any]:
								torch.manual_seed(0   )
								__a		 =       TaEncoderModel.from_pretrained('hf-internal-testing/tiny-random-t5'   )

								torch.manual_seed(0   )
								__a		 =       AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-t5'   )

								torch.manual_seed(0   )
								__a		 =       UNetaDConditionModel(
								    sample_size=3_2				,						layers_per_block=1				,						block_out_channels=[3_2, 6_4]				,						down_block_types=[
								        'ResnetDownsampleBlock2D',
								        'SimpleCrossAttnDownBlock2D',
								    ]				,						mid_block_type='UNetMidBlock2DSimpleCrossAttn'				,						up_block_types=['SimpleCrossAttnUpBlock2D', 'ResnetUpsampleBlock2D']				,						in_channels=3				,						out_channels=6				,						cross_attention_dim=3_2				,						encoder_hid_dim=3_2				,						attention_head_dim=8				,						addition_embed_type='text'				,						addition_embed_type_num_heads=2				,						cross_attention_norm='group_norm'				,						resnet_time_scale_shift='scale_shift'				,						act_fn='gelu'				,						)
								unet.set_attn_processor(AttnAddedKVProcessor()   )  # For reproducibility tests

								torch.manual_seed(0   )
								__a		 =       DDPMScheduler(
								    num_train_timesteps=1_0_0_0				,						beta_schedule='squaredcos_cap_v2'				,						beta_start=0.0_001				,						beta_end=0.02				,						thresholding=UpperCAmelCase				,						dynamic_thresholding_ratio=0.95				,						sample_max_value=1.0				,						prediction_type='epsilon'				,						variance_type='learned_range'				,						)

								torch.manual_seed(0   )
								__a		 =       IFWatermarker()

								return {
								    "text_encoder": text_encoder,
								    "tokenizer": tokenizer,
								    "unet": unet,
								    "scheduler": scheduler,
								    "watermarker": watermarker,
								    "safety_checker": None,
								    "feature_extractor": None,
								}



			def __SCREAMING_SNAKE_CASE							(						self   )			->   Optional[int]:
								torch.manual_seed(0   )
								__a		 =       TaEncoderModel.from_pretrained('hf-internal-testing/tiny-random-t5'   )

								torch.manual_seed(0   )
								__a		 =       AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-t5'   )

								torch.manual_seed(0   )
								__a		 =       UNetaDConditionModel(
								    sample_size=3_2				,						layers_per_block=[1, 2]				,						block_out_channels=[3_2, 6_4]				,						down_block_types=[
								        'ResnetDownsampleBlock2D',
								        'SimpleCrossAttnDownBlock2D',
								    ]				,						mid_block_type='UNetMidBlock2DSimpleCrossAttn'				,						up_block_types=['SimpleCrossAttnUpBlock2D', 'ResnetUpsampleBlock2D']				,						in_channels=6				,						out_channels=6				,						cross_attention_dim=3_2				,						encoder_hid_dim=3_2				,						attention_head_dim=8				,						addition_embed_type='text'				,						addition_embed_type_num_heads=2				,						cross_attention_norm='group_norm'				,						resnet_time_scale_shift='scale_shift'				,						act_fn='gelu'				,						class_embed_type='timestep'				,						mid_block_scale_factor=1.414				,						time_embedding_act_fn='gelu'				,						time_embedding_dim=3_2				,						)
								unet.set_attn_processor(AttnAddedKVProcessor()   )  # For reproducibility tests

								torch.manual_seed(0   )
								__a		 =       DDPMScheduler(
								    num_train_timesteps=1_0_0_0				,						beta_schedule='squaredcos_cap_v2'				,						beta_start=0.0_001				,						beta_end=0.02				,						thresholding=UpperCAmelCase				,						dynamic_thresholding_ratio=0.95				,						sample_max_value=1.0				,						prediction_type='epsilon'				,						variance_type='learned_range'				,						)

								torch.manual_seed(0   )
								__a		 =       DDPMScheduler(
								    num_train_timesteps=1_0_0_0				,						beta_schedule='squaredcos_cap_v2'				,						beta_start=0.0_001				,						beta_end=0.02				,						)

								torch.manual_seed(0   )
								__a		 =       IFWatermarker()

								return {
								    "text_encoder": text_encoder,
								    "tokenizer": tokenizer,
								    "unet": unet,
								    "scheduler": scheduler,
								    "image_noising_scheduler": image_noising_scheduler,
								    "watermarker": watermarker,
								    "safety_checker": None,
								    "feature_extractor": None,
								}



			def __SCREAMING_SNAKE_CASE							(						self   )			->   str:
								__a		 =       self.get_dummy_components()
								__a		 =       self.pipeline_class(**UpperCAmelCase   )
								pipe.to(UpperCAmelCase   )
								pipe.set_progress_bar_config(disable=UpperCAmelCase   )

								__a		 =       self.get_dummy_inputs(UpperCAmelCase   )

								__a		 =       inputs['prompt']
								__a		 =       inputs['generator']
								__a		 =       inputs['num_inference_steps']
								__a		 =       inputs['output_type']

								if "image" in inputs:
													__a		 =       inputs['image']
								else:
													__a		 =       None

								if "mask_image" in inputs:
													__a		 =       inputs['mask_image']
								else:
													__a		 =       None

								if "original_image" in inputs:
													__a		 =       inputs['original_image']
								else:
													__a		 =       None

								__a ,				__a		 =       pipe.encode_prompt(UpperCAmelCase   )

								# inputs with prompt converted to embeddings
								__a		 =       {
								    'prompt_embeds': prompt_embeds,
								    'negative_prompt_embeds': negative_prompt_embeds,
								    'generator': generator,
								    'num_inference_steps': num_inference_steps,
								    'output_type': output_type,
								}

								if image is not None:
													__a		 =       image

								if mask_image is not None:
													__a		 =       mask_image

								if original_image is not None:
													__a		 =       original_image

								# set all optional components to None
								for optional_component in pipe._optional_components:
													setattr(UpperCAmelCase				,						UpperCAmelCase				,						UpperCAmelCase   )

								__a		 =       pipe(**UpperCAmelCase   )[0]

								with tempfile.TemporaryDirectory() as tmpdir:
													pipe.save_pretrained(UpperCAmelCase   )
													__a		 =       self.pipeline_class.from_pretrained(UpperCAmelCase   )
													pipe_loaded.to(UpperCAmelCase   )
													pipe_loaded.set_progress_bar_config(disable=UpperCAmelCase   )

								pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor()   )  # For reproducibility tests

								for optional_component in pipe._optional_components:
													self.assertTrue(
													    getattr(UpperCAmelCase				,						UpperCAmelCase   ) is None				,						f'''`{optional_component}` did not stay set to None after loading.'''				,						)

								__a		 =       self.get_dummy_inputs(UpperCAmelCase   )

								__a		 =       inputs['generator']
								__a		 =       inputs['num_inference_steps']
								__a		 =       inputs['output_type']

								# inputs with prompt converted to embeddings
								__a		 =       {
								    'prompt_embeds': prompt_embeds,
								    'negative_prompt_embeds': negative_prompt_embeds,
								    'generator': generator,
								    'num_inference_steps': num_inference_steps,
								    'output_type': output_type,
								}

								if image is not None:
													__a		 =       image

								if mask_image is not None:
													__a		 =       mask_image

								if original_image is not None:
													__a		 =       original_image

								__a		 =       pipe_loaded(**UpperCAmelCase   )[0]

								__a		 =       np.abs(to_np(UpperCAmelCase   ) - to_np(UpperCAmelCase   )   ).max()
								self.assertLess(UpperCAmelCase				,						1e-4   )







			def __SCREAMING_SNAKE_CASE							(						self   )			->   int:
								__a		 =       self.get_dummy_components()
								__a		 =       self.pipeline_class(**UpperCAmelCase   )
								pipe.to(UpperCAmelCase   )
								pipe.set_progress_bar_config(disable=UpperCAmelCase   )

								__a		 =       self.get_dummy_inputs(UpperCAmelCase   )
								__a		 =       pipe(**UpperCAmelCase   )[0]

								with tempfile.TemporaryDirectory() as tmpdir:
													pipe.save_pretrained(UpperCAmelCase   )
													__a		 =       self.pipeline_class.from_pretrained(UpperCAmelCase   )
													pipe_loaded.to(UpperCAmelCase   )
													pipe_loaded.set_progress_bar_config(disable=UpperCAmelCase   )

								pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor()   )  # For reproducibility tests

								__a		 =       self.get_dummy_inputs(UpperCAmelCase   )
								__a		 =       pipe_loaded(**UpperCAmelCase   )[0]

								__a		 =       np.abs(to_np(UpperCAmelCase   ) - to_np(UpperCAmelCase   )   ).max()
								self.assertLess(UpperCAmelCase				,						1e-4   )




