






import gc
import unittest

import numpy as np
import torch

from diffusers import DDPMScheduler, MidiProcessor, SpectrogramDiffusionPipeline
from diffusers.pipelines.spectrogram_diffusion import SpectrogramContEncoder, SpectrogramNotesEncoder, TaFilmDecoder
from diffusers.utils import require_torch_gpu, skip_mps, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_note_seq, require_onnxruntime

from ..pipeline_params import TOKENS_TO_AUDIO_GENERATION_BATCH_PARAMS, TOKENS_TO_AUDIO_GENERATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()


lowerCamelCase_						     =  """./tests/fixtures/elise_format0.mid"""




@unittest.skip('''The note-seq package currently throws an error on import'''		)
class 							a_ (						a_	,	unittest.TestCase		):


							'''simple docstring'''






							__a:	Tuple     =    SpectrogramDiffusionPipeline
							__a:	Dict     =    PipelineTesterMixin.required_optional_params - {
							    '''callback''',
							    '''latents''',
							    '''callback_steps''',
							    '''output_type''',
							    '''num_images_per_prompt''',
							}
							__a:	Optional[int]     =    False
							__a:	Dict     =    False
							__a:	Dict     =    TOKENS_TO_AUDIO_GENERATION_PARAMS
							__a:	List[str]     =    TOKENS_TO_AUDIO_GENERATION_BATCH_PARAMS

							def      _lowercase       (		self   )					->					int:
														'''simple docstring'''

														torch.manual_seed(0   )
														lowerCAmelCase_										=     SpectrogramNotesEncoder(
														    max_length=2_0_4_8			,			vocab_size=1_5_3_6			,			d_model=7_6_8			,			dropout_rate=0.1			,			num_layers=1			,			num_heads=1			,			d_kv=4			,			d_ff=2_0_4_8			,			feed_forward_proj='gated-gelu'			,			)

														lowerCAmelCase_										=     SpectrogramContEncoder(
														    input_dims=1_2_8			,			targets_context_length=2_5_6			,			d_model=7_6_8			,			dropout_rate=0.1			,			num_layers=1			,			num_heads=1			,			d_kv=4			,			d_ff=2_0_4_8			,			feed_forward_proj='gated-gelu'			,			)

														lowerCAmelCase_										=     TaFilmDecoder(
														    input_dims=1_2_8			,			targets_length=2_5_6			,			max_decoder_noise_time=2_00_00.0			,			d_model=7_6_8			,			num_layers=1			,			num_heads=1			,			d_kv=4			,			d_ff=2_0_4_8			,			dropout_rate=0.1			,			)

														lowerCAmelCase_										=     DDPMScheduler()

														lowerCAmelCase_										=     {
														    'notes_encoder': notes_encoder.eval(),
														    'continuous_encoder': continuous_encoder.eval(),
														    'decoder': decoder.eval(),
														    'scheduler': scheduler,
														    'melgan': None,
														}
														return components

							def      _lowercase       (		self			,			lowercase_			,			lowercase_=0   )					->					List[str]:
														'''simple docstring'''

														if str(lowercase_   ).startswith('mps'   ):
																					lowerCAmelCase_										=     torch.manual_seed(lowercase_   )
														else:
																					lowerCAmelCase_										=     torch.Generator(device=lowercase_   ).manual_seed(lowercase_   )
														lowerCAmelCase_										=     {
														    'input_tokens': [
														        [1_1_3_4, 9_0, 1_1_3_5, 1_1_3_3, 1_0_8_0, 1_1_2, 1_1_3_2, 1_0_8_0, 1_1_3_3, 1_0_7_9, 1_3_3, 1_1_3_2, 1_0_7_9, 1_1_3_3, 1] + [0] * 2_0_3_3
														    ],
														    'generator': generator,
														    'num_inference_steps': 4,
														    'output_type': 'mel',
														}
														return inputs

							def      _lowercase       (		self   )					->					Union[str, Any]:
														'''simple docstring'''

														lowerCAmelCase_										=     'cpu'  # ensure determinism for the device-dependent torch.Generator
														lowerCAmelCase_										=     self.get_dummy_components()
														lowerCAmelCase_										=     SpectrogramDiffusionPipeline(**lowercase_   )
														lowerCAmelCase_										=     pipe.to(lowercase_   )
														pipe.set_progress_bar_config(disable=lowercase_   )

														lowerCAmelCase_										=     self.get_dummy_inputs(lowercase_   )
														lowerCAmelCase_										=     pipe(**lowercase_   )
														lowerCAmelCase_										=     output.audios

														lowerCAmelCase_										=     mel[0, -3:, -3:]

														assert mel_slice.shape == (3, 3)
														lowerCAmelCase_										=     np.array(
														    [-11.51_29_25, -4.78_82_15, -0.46_17_29_05, -2.05_17_15, -10.53_91_47, -10.97_09_63, -9.09_16_34, 4.0, 4.0]   )
														assert np.abs(mel_slice.flatten() - expected_slice   ).max() < 1e-2

							@skip_mps
							def      _lowercase       (		self   )					->					Tuple:
														'''simple docstring'''

														return super().test_save_load_local()

							@skip_mps
							def      _lowercase       (		self   )					->					str:
														'''simple docstring'''

														return super().test_dict_tuple_outputs_equivalent()

							@skip_mps
							def      _lowercase       (		self   )					->					int:
														'''simple docstring'''

														return super().test_save_load_optional_components()

							@skip_mps
							def      _lowercase       (		self   )					->					Dict:
														'''simple docstring'''

														return super().test_attention_slicing_forward_pass()

							def      _lowercase       (		self   )					->					List[Any]:
														'''simple docstring'''

														pass

							def      _lowercase       (		self   )					->					Optional[Any]:
														'''simple docstring'''

														pass






							@skip_mps
							def      _lowercase       (		self   )					->					Optional[int]:
														'''simple docstring'''

														return super().test_progress_bar()




@slow
@require_torch_gpu
@require_onnxruntime
@require_note_seq
class 							a_ (						unittest.TestCase		):


							'''simple docstring'''






							def      _lowercase       (		self   )					->					int:
														'''simple docstring'''

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

							def      _lowercase       (		self   )					->					Any:
														'''simple docstring'''

														lowerCAmelCase_										=     torch_device

														lowerCAmelCase_										=     SpectrogramDiffusionPipeline.from_pretrained('google/music-spectrogram-diffusion'   )
														lowerCAmelCase_										=     pipe.melgan
														lowerCAmelCase_										=     None

														lowerCAmelCase_										=     pipe.to(lowercase_   )
														pipe.set_progress_bar_config(disable=lowercase_   )

														def callback(lowercase_			,			lowercase_   ):
																					# decode mel to audio
																					lowerCAmelCase_										=     melgan(input_features=mel_output.astype(np.floataa   )   )[0]
																					assert len(audio[0]   ) == 8_1_9_2_0 * (step + 1)
																					# simulate that audio is played
																					return audio

														lowerCAmelCase_										=     MidiProcessor()
														lowerCAmelCase_										=     processor(lowercase_   )

														lowerCAmelCase_										=     input_tokens[:3]
														lowerCAmelCase_										=     torch.manual_seed(0   )
														pipe(lowercase_			,			num_inference_steps=5			,			generator=lowercase_			,			callback=lowercase_			,			output_type='mel'   )

							def      _lowercase       (		self   )					->					List[str]:
														'''simple docstring'''

														lowerCAmelCase_										=     torch_device

														lowerCAmelCase_										=     SpectrogramDiffusionPipeline.from_pretrained('google/music-spectrogram-diffusion'   )
														lowerCAmelCase_										=     pipe.to(lowercase_   )
														pipe.set_progress_bar_config(disable=lowercase_   )
														lowerCAmelCase_										=     MidiProcessor()

														lowerCAmelCase_										=     processor(lowercase_   )
														# just run two denoising loops
														lowerCAmelCase_										=     input_tokens[:2]

														lowerCAmelCase_										=     torch.manual_seed(0   )
														lowerCAmelCase_										=     pipe(lowercase_			,			num_inference_steps=2			,			generator=lowercase_   )

														lowerCAmelCase_										=     output.audios[0]

														assert abs(np.abs(lowercase_   ).sum() - 36_12.8_41   ) < 1e-1






							def      _lowercase       (		self   )					->					int:
														'''simple docstring'''

														lowerCAmelCase_										=     torch_device

														lowerCAmelCase_										=     SpectrogramDiffusionPipeline.from_pretrained('google/music-spectrogram-diffusion'   )
														lowerCAmelCase_										=     pipe.to(lowercase_   )
														pipe.set_progress_bar_config(disable=lowercase_   )

														lowerCAmelCase_										=     MidiProcessor()

														lowerCAmelCase_										=     processor(lowercase_   )

														# just run 4 denoising loops
														lowerCAmelCase_										=     input_tokens[:4]

														lowerCAmelCase_										=     torch.manual_seed(0   )
														lowerCAmelCase_										=     pipe(lowercase_			,			num_inference_steps=1_0_0			,			generator=lowercase_   )

														lowerCAmelCase_										=     output.audios[0]
														assert abs(np.abs(lowercase_   ).sum() - 93_89.11_11   ) < 5e-2











import os
import re
import unicodedata
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import AddedToken, PreTrainedTokenizer
from ...utils import logging


lowerCamelCase_						     =  logging.get_logger(__name__)
lowerCamelCase_						     =  {"""vocab_file""": """spiece.model"""}

lowerCamelCase_						     =  {
    """vocab_file""": {
        """google/fnet-base""": """https://huggingface.co/google/fnet-base/resolve/main/spiece.model""",
        """google/fnet-large""": """https://huggingface.co/google/fnet-large/resolve/main/spiece.model""",
    },
}

lowerCamelCase_						     =  {
    """google/fnet-base""": 5_1_2,
    """google/fnet-large""": 5_1_2,
}

lowerCamelCase_						     =  """▁"""




class 							a_ (						a_		):


							'''simple docstring'''






							__a:	Optional[int]     =    VOCAB_FILES_NAMES
							__a:	Union[str, Any]     =    PRETRAINED_VOCAB_FILES_MAP
							__a:	List[Any]     =    PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
							__a:	Any     =    ['''input_ids''', '''token_type_ids''']

							def __init__(		self			,			lowercase_			,			lowercase_=False			,			lowercase_=True			,			lowercase_=True			,			lowercase_="<unk>"			,			lowercase_="[SEP]"			,			lowercase_="<pad>"			,			lowercase_="[CLS]"			,			lowercase_="[MASK]"			,			lowercase_ = None			,			**lowercase_			,			)					->					None:
														'''simple docstring'''

														lowerCAmelCase_										=     (
														    AddedToken(lowercase_			,			lstrip=lowercase_			,			rstrip=lowercase_			,			normalized=lowercase_   )
														    if isinstance(lowercase_			,			lowercase_   )
														    else mask_token
														)

														lowerCAmelCase_										=     {} if sp_model_kwargs is None else sp_model_kwargs

														super().__init__(
														    do_lower_case=lowercase_			,			remove_space=lowercase_			,			keep_accents=lowercase_			,			unk_token=lowercase_			,			sep_token=lowercase_			,			pad_token=lowercase_			,			cls_token=lowercase_			,			mask_token=lowercase_			,			sp_model_kwargs=self.sp_model_kwargs			,			**lowercase_			,			)

														lowerCAmelCase_										=     do_lower_case
														lowerCAmelCase_										=     remove_space
														lowerCAmelCase_										=     keep_accents
														lowerCAmelCase_										=     vocab_file

														lowerCAmelCase_										=     spm.SentencePieceProcessor(**self.sp_model_kwargs   )
														self.sp_model.Load(lowercase_   )

							@property
							def      _lowercase       (		self   )					->					Optional[Any]:
														'''simple docstring'''

														return len(self.sp_model   )

							def      _lowercase       (		self   )					->					Dict:
														'''simple docstring'''

														lowerCAmelCase_										=     {self.convert_ids_to_tokens(lowercase_   ): i for i in range(self.vocab_size   )}
														vocab.update(self.added_tokens_encoder   )
														return vocab

							def __getstate__(		self   )					->					Dict:
														'''simple docstring'''

														lowerCAmelCase_										=     self.__dict__.copy()
														lowerCAmelCase_										=     None
														return state

							def __setstate__(		self			,			lowercase_   )					->					Optional[int]:
														'''simple docstring'''

														lowerCAmelCase_										=     d

														# for backward compatibility
														if not hasattr(self			,			'sp_model_kwargs'   ):
																					lowerCAmelCase_										=     {}

														lowerCAmelCase_										=     spm.SentencePieceProcessor(**self.sp_model_kwargs   )
														self.sp_model.Load(self.vocab_file   )

							def      _lowercase       (		self			,			lowercase_   )					->					str:
														'''simple docstring'''

														if self.remove_space:
																					lowerCAmelCase_										=     ' '.join(inputs.strip().split()   )
														else:
																					lowerCAmelCase_										=     inputs
														lowerCAmelCase_										=     outputs.replace('``'			,			'"'   ).replace('\'\''			,			'"'   )

														if not self.keep_accents:
																					lowerCAmelCase_										=     unicodedata.normalize('NFKD'			,			lowercase_   )
																					lowerCAmelCase_										=     ''.join([c for c in outputs if not unicodedata.combining(lowercase_   )]   )
														if self.do_lower_case:
																					lowerCAmelCase_										=     outputs.lower()

														return outputs

							def      _lowercase       (		self			,			lowercase_   )					->					List[str]:
														'''simple docstring'''

														lowerCAmelCase_										=     self.preprocess_text(lowercase_   )
														lowerCAmelCase_										=     self.sp_model.encode(lowercase_			,			out_type=lowercase_   )
														lowerCAmelCase_										=     []
														for piece in pieces:
																					if len(lowercase_   ) > 1 and piece[-1] == str(','   ) and piece[-2].isdigit():
																												lowerCAmelCase_										=     self.sp_model.EncodeAsPieces(piece[:-1].replace(lowercase_			,			''   )   )
																												if piece[0] != SPIECE_UNDERLINE and cur_pieces[0][0] == SPIECE_UNDERLINE:
																																			if len(cur_pieces[0]   ) == 1:
																																										lowerCAmelCase_										=     cur_pieces[1:]
																																			else:
																																										lowerCAmelCase_										=     cur_pieces[0][1:]
																												cur_pieces.append(piece[-1]   )
																												new_pieces.extend(lowercase_   )
																					else:
																												new_pieces.append(lowercase_   )

														return new_pieces

							def      _lowercase       (		self			,			lowercase_   )					->					Tuple:
														'''simple docstring'''

														return self.sp_model.PieceToId(lowercase_   )

							def      _lowercase       (		self			,			lowercase_   )					->					Union[str, Any]:
														'''simple docstring'''

														return self.sp_model.IdToPiece(lowercase_   )

							def      _lowercase       (		self			,			lowercase_   )					->					Union[str, Any]:
														'''simple docstring'''

														lowerCAmelCase_										=     []
														lowerCAmelCase_										=     ''
														lowerCAmelCase_										=     False
														for token in tokens:
																					# make sure that special tokens are not decoded using sentencepiece model
																					if token in self.all_special_tokens:
																												if not prev_is_special:
																																			out_string += " "
																												out_string += self.sp_model.decode(lowercase_   ) + token
																												lowerCAmelCase_										=     True
																												lowerCAmelCase_										=     []
																					else:
																												current_sub_tokens.append(lowercase_   )
																												lowerCAmelCase_										=     False
														out_string += self.sp_model.decode(lowercase_   )
														return out_string.strip()

							def      _lowercase       (		self			,			lowercase_			,			lowercase_ = False			,			lowercase_ = None			,			lowercase_ = True			,			**lowercase_			,			)					->					str:
														'''simple docstring'''

														lowerCAmelCase_										=     kwargs.pop('use_source_tokenizer'			,			lowercase_   )

														lowerCAmelCase_										=     self.convert_ids_to_tokens(lowercase_			,			skip_special_tokens=lowercase_   )

														# To avoid mixing byte-level and unicode for byte-level BPT
														# we need to build string separately for added tokens and byte-level tokens
														# cf. https://github.com/huggingface/transformers/issues/1133
														lowerCAmelCase_										=     []
														lowerCAmelCase_										=     []
														for token in filtered_tokens:
																					if skip_special_tokens and token in self.all_special_ids:
																												continue
																					if token in self.added_tokens_encoder:
																												if current_sub_text:
																																			sub_texts.append(self.convert_tokens_to_string(lowercase_   )   )
																																			lowerCAmelCase_										=     []
																												sub_texts.append(lowercase_   )
																					else:
																												current_sub_text.append(lowercase_   )
														if current_sub_text:
																					sub_texts.append(self.convert_tokens_to_string(lowercase_   )   )

														# Mimic the behavior of the Rust tokenizer:
														# No space after <unk>
														if spaces_between_special_tokens:
																					lowerCAmelCase_										=     re.sub(R'(<unk>) '			,			R'\1'			,			' '.join(lowercase_   )   )
														else:
																					lowerCAmelCase_										=     ''.join(lowercase_   )

														lowerCAmelCase_										=     (
														    clean_up_tokenization_spaces
														    if clean_up_tokenization_spaces is not None
														    else self.clean_up_tokenization_spaces
														)
														if clean_up_tokenization_spaces:
																					lowerCAmelCase_										=     self.clean_up_tokenization(lowercase_   )
																					return clean_text
														else:
																					return text

							def      _lowercase       (		self			,			lowercase_			,			lowercase_ = None   )					->					List[int]:
														'''simple docstring'''

														lowerCAmelCase_										=     [self.sep_token_id]
														lowerCAmelCase_										=     [self.cls_token_id]
														if token_ids_a is None:
																					return cls + token_ids_a + sep
														return cls + token_ids_a + sep + token_ids_a + sep

							def      _lowercase       (		self			,			lowercase_			,			lowercase_ = None			,			lowercase_ = False   )					->					List[int]:
														'''simple docstring'''

														if already_has_special_tokens:
																					return super().get_special_tokens_mask(
																					    token_ids_a=lowercase_			,			token_ids_a=lowercase_			,			already_has_special_tokens=lowercase_   )

														if token_ids_a is not None:
																					return [1] + ([0] * len(lowercase_   )) + [1] + ([0] * len(lowercase_   )) + [1]
														return [1] + ([0] * len(lowercase_   )) + [1]

							def      _lowercase       (		self			,			lowercase_			,			lowercase_ = None   )					->					List[int]:
														'''simple docstring'''

														lowerCAmelCase_										=     [self.sep_token_id]
														lowerCAmelCase_										=     [self.cls_token_id]

														if token_ids_a is None:
																					return len(cls + token_ids_a + sep   ) * [0]
														return len(cls + token_ids_a + sep   ) * [0] + len(token_ids_a + sep   ) * [1]






							def      _lowercase       (		self			,			lowercase_			,			lowercase_ = None   )					->					Tuple[str]:
														'''simple docstring'''

														if not os.path.isdir(lowercase_   ):
																					logger.error(f'''Vocabulary path ({save_directory}) should be a directory'''   )
																					return
														lowerCAmelCase_										=     os.path.join(
														    lowercase_			,			(filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['vocab_file']   )

														if os.path.abspath(self.vocab_file   ) != os.path.abspath(lowercase_   ) and os.path.isfile(self.vocab_file   ):
																					copyfile(self.vocab_file			,			lowercase_   )
														elif not os.path.isfile(self.vocab_file   ):
																					with open(lowercase_			,			'wb'   ) as fi:
																												lowerCAmelCase_										=     self.sp_model.serialized_model_proto()
																												fi.write(lowercase_   )

														return (out_vocab_file,)



