
import json
import os
from typing import Optional

import numpy as np

from ...feature_extraction_utils import BatchFeature
from ...processing_utils import ProcessorMixin
from ...utils import logging
from ...utils.hub import get_file_from_repo
from ..auto import AutoTokenizer


lowerCamelCase  =logging.get_logger(__name__)




class 					_lowerCamelCase			(   UpperCamelCase_			):





     """simple docstring"""







     SCREAMING_SNAKE_CASE_		 =	'''AutoTokenizer'''
     SCREAMING_SNAKE_CASE_		 =	['''tokenizer''']

     SCREAMING_SNAKE_CASE_		 =	{
         '''semantic_prompt''': 1,
         '''coarse_prompt''': 2,
         '''fine_prompt''': 2,
     }






     def __init__(				self    ,    __SCREAMING_SNAKE_CASE    ,    __SCREAMING_SNAKE_CASE=None	)	-> int:




           """simple docstring"""







           super().__init__(__SCREAMING_SNAKE_CASE	)

           UpperCamelCase__     : Dict											=					speaker_embeddings






     @classmethod
     def  __SCREAMING_SNAKE_CASE      (				cls    ,    __SCREAMING_SNAKE_CASE    ,    __SCREAMING_SNAKE_CASE="speaker_embeddings_path.json"    ,    **__SCREAMING_SNAKE_CASE	)	-> Dict:




           """simple docstring"""







           if speaker_embeddings_dict_path is not None:
                 UpperCamelCase__     : Optional[Any]											=					get_file_from_repo(
                     __SCREAMING_SNAKE_CASE    ,    __SCREAMING_SNAKE_CASE    ,    subfolder=kwargs.pop('''subfolder'''    ,    __SCREAMING_SNAKE_CASE	)    ,    cache_dir=kwargs.pop('''cache_dir'''    ,    __SCREAMING_SNAKE_CASE	)    ,    force_download=kwargs.pop('''force_download'''    ,    __SCREAMING_SNAKE_CASE	)    ,    proxies=kwargs.pop('''proxies'''    ,    __SCREAMING_SNAKE_CASE	)    ,    resume_download=kwargs.pop('''resume_download'''    ,    __SCREAMING_SNAKE_CASE	)    ,    local_files_only=kwargs.pop('''local_files_only'''    ,    __SCREAMING_SNAKE_CASE	)    ,    use_auth_token=kwargs.pop('''use_auth_token'''    ,    __SCREAMING_SNAKE_CASE	)    ,    revision=kwargs.pop('''revision'''    ,    __SCREAMING_SNAKE_CASE	)    ,    )
                 if speaker_embeddings_path is None:
                       logger.warning(
                           F'''`{os.path.join(__SCREAMING_SNAKE_CASE    ,    __SCREAMING_SNAKE_CASE	)}` does not exists
                    , no preloaded speaker embeddings will be used - Make sure to provide a correct path to the json
                    dictionnary if wanted, otherwise set `speaker_embeddings_dict_path=None`.'''	)
                       UpperCamelCase__     : Union[str, Any]											=					None
                 else:
                       with open(__SCREAMING_SNAKE_CASE	) as speaker_embeddings_json:
                             UpperCamelCase__     : Dict											=					json.load(__SCREAMING_SNAKE_CASE	)
           else:
                 UpperCamelCase__     : Any											=					None

           UpperCamelCase__     : Union[str, Any]											=					AutoTokenizer.from_pretrained(__SCREAMING_SNAKE_CASE    ,    **__SCREAMING_SNAKE_CASE	)

           return cls(tokenizer=__SCREAMING_SNAKE_CASE    ,    speaker_embeddings=__SCREAMING_SNAKE_CASE	)






     def  __SCREAMING_SNAKE_CASE      (				self    ,    __SCREAMING_SNAKE_CASE    ,    __SCREAMING_SNAKE_CASE="speaker_embeddings_path.json"    ,    __SCREAMING_SNAKE_CASE="speaker_embeddings"    ,    __SCREAMING_SNAKE_CASE = False    ,    **__SCREAMING_SNAKE_CASE    ,    )	-> Optional[Any]:




           """simple docstring"""







           if self.speaker_embeddings is not None:
                 os.makedirs(os.path.join(__SCREAMING_SNAKE_CASE    ,    __SCREAMING_SNAKE_CASE    ,    '''v2'''	)    ,    exist_ok=__SCREAMING_SNAKE_CASE	)

                 UpperCamelCase__     : Optional[int]											=					{}

                 UpperCamelCase__     : Union[str, Any]											=					save_directory

                 for prompt_key in self.speaker_embeddings:
                       if prompt_key != "repo_or_path":
                             UpperCamelCase__     : List[Any]											=					self._load_voice_preset(__SCREAMING_SNAKE_CASE	)

                             UpperCamelCase__     : Dict											=					{}
                             for key in self.speaker_embeddings[prompt_key]:
                                   np.save(
                                       os.path.join(
                                           embeddings_dict['''repo_or_path''']    ,    __SCREAMING_SNAKE_CASE    ,    F'''{prompt_key}_{key}'''	)    ,    voice_preset[key]    ,    allow_pickle=__SCREAMING_SNAKE_CASE    ,    )
                                   UpperCamelCase__     : Optional[Any]											=					os.path.join(__SCREAMING_SNAKE_CASE    ,    F'''{prompt_key}_{key}.npy'''	)

                             UpperCamelCase__     : Union[str, Any]											=					tmp_dict

                 with open(os.path.join(__SCREAMING_SNAKE_CASE    ,    __SCREAMING_SNAKE_CASE	)    ,    '''w'''	) as fp:
                       json.dump(__SCREAMING_SNAKE_CASE    ,    __SCREAMING_SNAKE_CASE	)

           super().save_pretrained(__SCREAMING_SNAKE_CASE    ,    __SCREAMING_SNAKE_CASE    ,    **__SCREAMING_SNAKE_CASE	)






     def  __SCREAMING_SNAKE_CASE      (				self    ,    __SCREAMING_SNAKE_CASE = None    ,    **__SCREAMING_SNAKE_CASE	)	-> int:




           """simple docstring"""







           UpperCamelCase__     : Optional[int]											=					self.speaker_embeddings[voice_preset]

           UpperCamelCase__     : Tuple											=					{}
           for key in ["semantic_prompt", "coarse_prompt", "fine_prompt"]:
                 if key not in voice_preset_paths:
                       raise ValueError(
                           F'''Voice preset unrecognized, missing {key} as a key in self.speaker_embeddings[{voice_preset}].'''	)

                 UpperCamelCase__     : Any											=					get_file_from_repo(
                     self.speaker_embeddings.get('''repo_or_path'''    ,    '''/'''	)    ,    voice_preset_paths[key]    ,    subfolder=kwargs.pop('''subfolder'''    ,    __SCREAMING_SNAKE_CASE	)    ,    cache_dir=kwargs.pop('''cache_dir'''    ,    __SCREAMING_SNAKE_CASE	)    ,    force_download=kwargs.pop('''force_download'''    ,    __SCREAMING_SNAKE_CASE	)    ,    proxies=kwargs.pop('''proxies'''    ,    __SCREAMING_SNAKE_CASE	)    ,    resume_download=kwargs.pop('''resume_download'''    ,    __SCREAMING_SNAKE_CASE	)    ,    local_files_only=kwargs.pop('''local_files_only'''    ,    __SCREAMING_SNAKE_CASE	)    ,    use_auth_token=kwargs.pop('''use_auth_token'''    ,    __SCREAMING_SNAKE_CASE	)    ,    revision=kwargs.pop('''revision'''    ,    __SCREAMING_SNAKE_CASE	)    ,    )
                 if path is None:
                       raise ValueError(
                           F'''`{os.path.join(self.speaker_embeddings.get("repo_or_path"    ,    "/"	)    ,    voice_preset_paths[key]	)}` does not exists
                    , no preloaded voice preset will be used - Make sure to provide correct paths to the {voice_preset}
                    embeddings.'''	)

                 UpperCamelCase__     : List[Any]											=					np.load(__SCREAMING_SNAKE_CASE	)

           return voice_preset_dict






     def  __SCREAMING_SNAKE_CASE      (				self    ,    __SCREAMING_SNAKE_CASE = None	)	-> Optional[int]:




           """simple docstring"""







           for key in ["semantic_prompt", "coarse_prompt", "fine_prompt"]:
                 if key not in voice_preset:
                       raise ValueError(F'''Voice preset unrecognized, missing {key} as a key.'''	)

                 if not isinstance(voice_preset[key]    ,    np.ndarray	):
                       raise ValueError(F'''{key} voice preset must be a {str(self.preset_shape[key]	)}D ndarray.'''	)

                 if len(voice_preset[key].shape	) != self.preset_shape[key]:
                       raise ValueError(F'''{key} voice preset must be a {str(self.preset_shape[key]	)}D ndarray.'''	)






     def __call__(				self    ,    __SCREAMING_SNAKE_CASE=None    ,    __SCREAMING_SNAKE_CASE=None    ,    __SCREAMING_SNAKE_CASE="pt"    ,    __SCREAMING_SNAKE_CASE=2_5_6    ,    __SCREAMING_SNAKE_CASE=False    ,    __SCREAMING_SNAKE_CASE=True    ,    __SCREAMING_SNAKE_CASE=False    ,    **__SCREAMING_SNAKE_CASE    ,    )	-> Union[str, Any]:




           """simple docstring"""







           if voice_preset is not None and not isinstance(__SCREAMING_SNAKE_CASE    ,    __SCREAMING_SNAKE_CASE	):
                 if (
                     isinstance(__SCREAMING_SNAKE_CASE    ,    __SCREAMING_SNAKE_CASE	)
                     and self.speaker_embeddings is not None
                     and voice_preset in self.speaker_embeddings
                 ):
                       UpperCamelCase__     : str											=					self._load_voice_preset(__SCREAMING_SNAKE_CASE	)

                 else:
                       if isinstance(__SCREAMING_SNAKE_CASE    ,    __SCREAMING_SNAKE_CASE	) and not voice_preset.endswith('''.npz'''	):
                             UpperCamelCase__     : int											=					voice_preset + '''.npz'''

                       UpperCamelCase__     : Union[str, Any]											=					np.load(__SCREAMING_SNAKE_CASE	)

           if voice_preset is not None:
                 self._validate_voice_preset_dict(__SCREAMING_SNAKE_CASE    ,    **__SCREAMING_SNAKE_CASE	)
                 UpperCamelCase__     : List[Any]											=					BatchFeature(data=__SCREAMING_SNAKE_CASE    ,    tensor_type=__SCREAMING_SNAKE_CASE	)

           UpperCamelCase__     : List[str]											=					self.tokenizer(
               __SCREAMING_SNAKE_CASE    ,    return_tensors=__SCREAMING_SNAKE_CASE    ,    padding='''max_length'''    ,    max_length=__SCREAMING_SNAKE_CASE    ,    return_attention_mask=__SCREAMING_SNAKE_CASE    ,    return_token_type_ids=__SCREAMING_SNAKE_CASE    ,    add_special_tokens=__SCREAMING_SNAKE_CASE    ,    **__SCREAMING_SNAKE_CASE    ,    )

           if voice_preset is not None:
                 UpperCamelCase__     : Optional[int]											=					voice_preset

           return encoded_text




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			):





     """simple docstring"""







     SCREAMING_SNAKE_CASE_		 =	DanceDiffusionPipeline
     SCREAMING_SNAKE_CASE_		 =	UNCONDITIONAL_AUDIO_GENERATION_PARAMS
     SCREAMING_SNAKE_CASE_		 =	PipelineTesterMixin.required_optional_params - {
         '''callback''',
         '''latents''',
         '''callback_steps''',
         '''output_type''',
         '''num_images_per_prompt''',
     }
     SCREAMING_SNAKE_CASE_		 =	UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS
     SCREAMING_SNAKE_CASE_		 =	False
     SCREAMING_SNAKE_CASE_		 =	False






     def  __SCREAMING_SNAKE_CASE      (				self	)	-> Any:




           """simple docstring"""







           torch.manual_seed(0	)
           UpperCamelCase__     : List[str]											=					UNetaDModel(
               block_out_channels=(3_2, 3_2, 6_4)    ,    extra_in_channels=1_6    ,    sample_size=5_1_2    ,    sample_rate=1_6_0_0_0    ,    in_channels=2    ,    out_channels=2    ,    flip_sin_to_cos=__SCREAMING_SNAKE_CASE    ,    use_timestep_embedding=__SCREAMING_SNAKE_CASE    ,    time_embedding_type='''fourier'''    ,    mid_block_type='''UNetMidBlock1D'''    ,    down_block_types=('''DownBlock1DNoSkip''', '''DownBlock1D''', '''AttnDownBlock1D''')    ,    up_block_types=('''AttnUpBlock1D''', '''UpBlock1D''', '''UpBlock1DNoSkip''')    ,    )
           UpperCamelCase__     : Union[str, Any]											=					IPNDMScheduler()

           UpperCamelCase__     : List[str]											=					{
               '''unet''': unet,
               '''scheduler''': scheduler,
           }
           return components






     def  __SCREAMING_SNAKE_CASE      (				self    ,    __SCREAMING_SNAKE_CASE    ,    __SCREAMING_SNAKE_CASE=0	)	-> Any:




           """simple docstring"""







           if str(__SCREAMING_SNAKE_CASE	).startswith('''mps'''	):
                 UpperCamelCase__     : Optional[Any]											=					torch.manual_seed(__SCREAMING_SNAKE_CASE	)
           else:
                 UpperCamelCase__     : Optional[Any]											=					torch.Generator(device=__SCREAMING_SNAKE_CASE	).manual_seed(__SCREAMING_SNAKE_CASE	)
           UpperCamelCase__     : int											=					{
               '''batch_size''': 1,
               '''generator''': generator,
               '''num_inference_steps''': 4,
           }
           return inputs






     def  __SCREAMING_SNAKE_CASE      (				self	)	-> List[Any]:




           """simple docstring"""







           UpperCamelCase__     : str											=					'''cpu'''  # ensure determinism for the device-dependent torch.Generator
           UpperCamelCase__     : Dict											=					self.get_dummy_components()
           UpperCamelCase__     : str											=					DanceDiffusionPipeline(**__SCREAMING_SNAKE_CASE	)
           UpperCamelCase__     : List[Any]											=					pipe.to(__SCREAMING_SNAKE_CASE	)
           pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE	)

           UpperCamelCase__     : int											=					self.get_dummy_inputs(__SCREAMING_SNAKE_CASE	)
           UpperCamelCase__     : Dict											=					pipe(**__SCREAMING_SNAKE_CASE	)
           UpperCamelCase__     : List[str]											=					output.audios

           UpperCamelCase__     : Dict											=					audio[0, -3:, -3:]

           assert audio.shape == (1, 2, components["unet"].sample_size)
           UpperCamelCase__     : List[Any]											=					np.array([-0.7265, 1.0000, -0.8388, 0.1175, 0.9498, -1.0000]	)
           assert np.abs(audio_slice.flatten() - expected_slice	).max() < 1e-2






     @skip_mps
     def  __SCREAMING_SNAKE_CASE      (				self	)	-> List[Any]:




           """simple docstring"""







           return super().test_save_load_local()






     @skip_mps
     def  __SCREAMING_SNAKE_CASE      (				self	)	-> Tuple:




           """simple docstring"""







           return super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3	)






     @skip_mps
     def  __SCREAMING_SNAKE_CASE      (				self	)	-> str:




           """simple docstring"""







           return super().test_save_load_optional_components()






     @skip_mps
     def  __SCREAMING_SNAKE_CASE      (				self	)	-> Optional[int]:




           """simple docstring"""







           return super().test_attention_slicing_forward_pass()






     def  __SCREAMING_SNAKE_CASE      (				self	)	-> Tuple:




           """simple docstring"""







           super().test_inference_batch_single_identical(expected_max_diff=3e-3	)




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





     """simple docstring"""







     def  __SCREAMING_SNAKE_CASE      (				self	)	-> int:




           """simple docstring"""







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






     def  __SCREAMING_SNAKE_CASE      (				self	)	-> Union[str, Any]:




           """simple docstring"""







           UpperCamelCase__     : Any											=					torch_device

           UpperCamelCase__     : Any											=					DanceDiffusionPipeline.from_pretrained('''harmonai/maestro-150k'''	)
           UpperCamelCase__     : int											=					pipe.to(__SCREAMING_SNAKE_CASE	)
           pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE	)

           UpperCamelCase__     : int											=					torch.manual_seed(0	)
           UpperCamelCase__     : Optional[Any]											=					pipe(generator=__SCREAMING_SNAKE_CASE    ,    num_inference_steps=1_0_0    ,    audio_length_in_s=4.096	)
           UpperCamelCase__     : str											=					output.audios

           UpperCamelCase__     : List[Any]											=					audio[0, -3:, -3:]

           assert audio.shape == (1, 2, pipe.unet.sample_size)
           UpperCamelCase__     : Tuple											=					np.array([-0.0192, -0.0231, -0.0318, -0.0059, 0.0002, -0.0020]	)

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






     def  __SCREAMING_SNAKE_CASE      (				self	)	-> Union[str, Any]:




           """simple docstring"""







           UpperCamelCase__     : Any											=					torch_device

           UpperCamelCase__     : Union[str, Any]											=					DanceDiffusionPipeline.from_pretrained('''harmonai/maestro-150k'''    ,    torch_dtype=torch.floataa	)
           UpperCamelCase__     : Tuple											=					pipe.to(__SCREAMING_SNAKE_CASE	)
           pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE	)

           UpperCamelCase__     : int											=					torch.manual_seed(0	)
           UpperCamelCase__     : Union[str, Any]											=					pipe(generator=__SCREAMING_SNAKE_CASE    ,    num_inference_steps=1_0_0    ,    audio_length_in_s=4.096	)
           UpperCamelCase__     : List[Any]											=					output.audios

           UpperCamelCase__     : List[Any]											=					audio[0, -3:, -3:]

           assert audio.shape == (1, 2, pipe.unet.sample_size)
           UpperCamelCase__     : Optional[Any]											=					np.array([-0.0367, -0.0488, -0.0771, -0.0525, -0.0444, -0.0341]	)

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


