






import gc
import unittest

import numpy as np
import torch

from diffusers import DanceDiffusionPipeline, IPNDMScheduler, UNetaDModel
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS, UNCONDITIONAL_AUDIO_GENERATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()
class  lowerCAmelCase_       (				UpperCAmelCase__     ,    unittest.TestCase					):
							UpperCamelCase_		:str		      =    DanceDiffusionPipeline
							UpperCamelCase_		:int		      =    UNCONDITIONAL_AUDIO_GENERATION_PARAMS
							UpperCamelCase_		:Optional[int]		      =    PipelineTesterMixin.required_optional_params - {
							    'callback',
							    'latents',
							    'callback_steps',
							    'output_type',
							    'num_images_per_prompt',
							}
							UpperCamelCase_		:Optional[int]		      =    UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS
							UpperCamelCase_		:Union[str, Any]		      =    False
							UpperCamelCase_		:List[str]		      =    False




							def 						__snake_case					(   self     :   Union[str, Any]	):
											torch.manual_seed(0	)
											lowerCAmelCase__      			=						UNetaDModel(
											    block_out_channels=(32, 32, 64) ,							extra_in_channels=16 ,							sample_size=512 ,							sample_rate=16_000 ,							in_channels=2 ,							out_channels=2 ,							flip_sin_to_cos=__lowerCAmelCase ,							use_timestep_embedding=__lowerCAmelCase ,							time_embedding_type='''fourier''' ,							mid_block_type='''UNetMidBlock1D''' ,							down_block_types=('''DownBlock1DNoSkip''', '''DownBlock1D''', '''AttnDownBlock1D''') ,							up_block_types=('''AttnUpBlock1D''', '''UpBlock1D''', '''UpBlock1DNoSkip''') ,							)
											lowerCAmelCase__      			=						IPNDMScheduler()

											lowerCAmelCase__      			=						{
											    '''unet''': unet,
											    '''scheduler''': scheduler,
											}
											return components




							def 						__snake_case					(   self     :   Dict ,							SCREAMING_SNAKE_CASE_     :   int ,							SCREAMING_SNAKE_CASE_     :   Optional[int]=0	):






											if str(__lowerCAmelCase	).startswith('''mps'''	):
															lowerCAmelCase__      			=						torch.manual_seed(__lowerCAmelCase	)
											else:
															lowerCAmelCase__      			=						torch.Generator(device=__lowerCAmelCase	).manual_seed(__lowerCAmelCase	)
											lowerCAmelCase__      			=						{
											    '''batch_size''': 1,
											    '''generator''': generator,
											    '''num_inference_steps''': 4,
											}
											return inputs




							def 						__snake_case					(   self     :   Tuple	):
											lowerCAmelCase__      			=						'''cpu'''  # ensure determinism for the device-dependent torch.Generator
											lowerCAmelCase__      			=						self.get_dummy_components()
											lowerCAmelCase__      			=						DanceDiffusionPipeline(**__lowerCAmelCase	)
											lowerCAmelCase__      			=						pipe.to(__lowerCAmelCase	)
											pipe.set_progress_bar_config(disable=__lowerCAmelCase	)

											lowerCAmelCase__      			=						self.get_dummy_inputs(__lowerCAmelCase	)
											lowerCAmelCase__      			=						pipe(**__lowerCAmelCase	)
											lowerCAmelCase__      			=						output.audios

											lowerCAmelCase__      			=						audio[0, -3:, -3:]

											assert audio.shape == (1, 2, components["unet"].sample_size)
											lowerCAmelCase__      			=						np.array([-0.7_265, 1.0_000, -0.8_388, 0.1_175, 0.9_498, -1.0_000]	)
											assert np.abs(audio_slice.flatten() - expected_slice	).max() < 1e-2




							@skip_mps
							def 						__snake_case					(   self     :   Dict	):
											return super().test_save_load_local()




							@skip_mps
							def 						__snake_case					(   self     :   int	):
											return super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3	)




							@skip_mps
							def 						__snake_case					(   self     :   Optional[Any]	):
											return super().test_save_load_optional_components()




							@skip_mps
							def 						__snake_case					(   self     :   List[Any]	):
											return super().test_attention_slicing_forward_pass()




							def 						__snake_case					(   self     :   Tuple	):
											super().test_inference_batch_single_identical(expected_max_diff=3e-3	)



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




							def 						__snake_case					(   self     :   Any	):
											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()




							def 						__snake_case					(   self     :   Tuple	):
											lowerCAmelCase__      			=						torch_device

											lowerCAmelCase__      			=						DanceDiffusionPipeline.from_pretrained('''harmonai/maestro-150k'''	)
											lowerCAmelCase__      			=						pipe.to(__lowerCAmelCase	)
											pipe.set_progress_bar_config(disable=__lowerCAmelCase	)

											lowerCAmelCase__      			=						torch.manual_seed(0	)
											lowerCAmelCase__      			=						pipe(generator=__lowerCAmelCase ,							num_inference_steps=100 ,							audio_length_in_s=4.096	)
											lowerCAmelCase__      			=						output.audios

											lowerCAmelCase__      			=						audio[0, -3:, -3:]

											assert audio.shape == (1, 2, pipe.unet.sample_size)
											lowerCAmelCase__      			=						np.array([-0.0_192, -0.0_231, -0.0_318, -0.0_059, 0.0_002, -0.0_020]	)

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




							def 						__snake_case					(   self     :   Optional[Any]	):
											lowerCAmelCase__      			=						torch_device

											lowerCAmelCase__      			=						DanceDiffusionPipeline.from_pretrained('''harmonai/maestro-150k''' ,							torch_dtype=torch.floataa	)
											lowerCAmelCase__      			=						pipe.to(__lowerCAmelCase	)
											pipe.set_progress_bar_config(disable=__lowerCAmelCase	)

											lowerCAmelCase__      			=						torch.manual_seed(0	)
											lowerCAmelCase__      			=						pipe(generator=__lowerCAmelCase ,							num_inference_steps=100 ,							audio_length_in_s=4.096	)
											lowerCAmelCase__      			=						output.audios

											lowerCAmelCase__      			=						audio[0, -3:, -3:]

											assert audio.shape == (1, 2, pipe.unet.sample_size)
											lowerCAmelCase__      			=						np.array([-0.0_367, -0.0_488, -0.0_771, -0.0_525, -0.0_444, -0.0_341]	)

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







import os
import shutil
import tempfile
import unittest

import numpy as np

from transformers import AutoTokenizer, BarkProcessor
from transformers.testing_utils import require_torch, slow
@require_torch
class  lowerCAmelCase_       (				unittest.TestCase					):




							def 						__snake_case					(   self     :   Optional[int]	):
											lowerCAmelCase__      			=						'''ylacombe/bark-small'''
											lowerCAmelCase__      			=						tempfile.mkdtemp()
											lowerCAmelCase__      			=						'''en_speaker_1'''
											lowerCAmelCase__      			=						'''This is a test string'''
											lowerCAmelCase__      			=						'''speaker_embeddings_path.json'''
											lowerCAmelCase__      			=						'''speaker_embeddings'''




							def 						__snake_case					(   self     :   Optional[Any] ,							**SCREAMING_SNAKE_CASE_     :   Optional[int]	):
											return AutoTokenizer.from_pretrained(self.checkpoint ,							**SCREAMING_SNAKE_CASE_	)




							def 						__snake_case					(   self     :   List[str]	):
											shutil.rmtree(self.tmpdirname	)




							def 						__snake_case					(   self     :   Optional[Any]	):
											lowerCAmelCase__      			=						self.get_tokenizer()

											lowerCAmelCase__      			=						BarkProcessor(tokenizer=SCREAMING_SNAKE_CASE_	)

											processor.save_pretrained(self.tmpdirname	)
											lowerCAmelCase__      			=						BarkProcessor.from_pretrained(self.tmpdirname	)

											self.assertEqual(processor.tokenizer.get_vocab() ,							tokenizer.get_vocab()	)




							@slow
							def 						__snake_case					(   self     :   Any	):
											lowerCAmelCase__      			=						BarkProcessor.from_pretrained(
											    pretrained_processor_name_or_path=self.checkpoint ,							speaker_embeddings_dict_path=self.speaker_embeddings_dict_path ,							)
											processor.save_pretrained(
											    self.tmpdirname ,							speaker_embeddings_dict_path=self.speaker_embeddings_dict_path ,							speaker_embeddings_directory=self.speaker_embeddings_directory ,							)

											lowerCAmelCase__      			=						self.get_tokenizer(bos_token='''(BOS)''' ,							eos_token='''(EOS)'''	)

											lowerCAmelCase__      			=						BarkProcessor.from_pretrained(
											    self.tmpdirname ,							self.speaker_embeddings_dict_path ,							bos_token='''(BOS)''' ,							eos_token='''(EOS)''' ,							)

											self.assertEqual(processor.tokenizer.get_vocab() ,							tokenizer_add_kwargs.get_vocab()	)




							def 						__snake_case					(   self     :   Dict	):
											lowerCAmelCase__      			=						BarkProcessor.from_pretrained(
											    pretrained_processor_name_or_path=self.checkpoint ,							speaker_embeddings_dict_path=self.speaker_embeddings_dict_path ,							)

											lowerCAmelCase__      			=						35
											lowerCAmelCase__      			=						2
											lowerCAmelCase__      			=						8

											lowerCAmelCase__      			=						{
											    '''semantic_prompt''': np.ones(SCREAMING_SNAKE_CASE_	),
											    '''coarse_prompt''': np.ones((nb_codebooks_coarse, seq_len)	),
											    '''fine_prompt''': np.ones((nb_codebooks_total, seq_len)	),
											}

											# test providing already loaded voice_preset
											lowerCAmelCase__      			=						processor(text=self.input_string ,							voice_preset=SCREAMING_SNAKE_CASE_	)

											lowerCAmelCase__      			=						inputs['''history_prompt''']
											for key in voice_preset:
															self.assertListEqual(voice_preset[key].tolist() ,							processed_voice_preset.get(SCREAMING_SNAKE_CASE_ ,							np.array([]	)	).tolist()	)

											# test loading voice preset from npz file
											lowerCAmelCase__      			=						os.path.join(self.tmpdirname ,							'''file.npz'''	)
											np.savez(SCREAMING_SNAKE_CASE_ ,							**SCREAMING_SNAKE_CASE_	)
											lowerCAmelCase__      			=						processor(text=self.input_string ,							voice_preset=SCREAMING_SNAKE_CASE_	)
											lowerCAmelCase__      			=						inputs['''history_prompt''']

											for key in voice_preset:
															self.assertListEqual(voice_preset[key].tolist() ,							processed_voice_preset.get(SCREAMING_SNAKE_CASE_ ,							np.array([]	)	).tolist()	)

											# test loading voice preset from the hub
											lowerCAmelCase__      			=						processor(text=self.input_string ,							voice_preset=self.voice_preset	)




							def 						__snake_case					(   self     :   Optional[int]	):
											lowerCAmelCase__      			=						self.get_tokenizer()

											lowerCAmelCase__      			=						BarkProcessor(tokenizer=SCREAMING_SNAKE_CASE_	)

											lowerCAmelCase__      			=						processor(text=self.input_string	)

											lowerCAmelCase__      			=						tokenizer(
											    self.input_string ,							padding='''max_length''' ,							max_length=256 ,							add_special_tokens=SCREAMING_SNAKE_CASE_ ,							return_attention_mask=SCREAMING_SNAKE_CASE_ ,							return_token_type_ids=SCREAMING_SNAKE_CASE_ ,							)

											for key in encoded_tok.keys():
															self.assertListEqual(encoded_tok[key] ,							encoded_processor[key].squeeze().tolist()	)





