

"""simple docstring"""



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


SCREAMING_SNAKE_CASE_       :						str					=    './tests/fixtures/elise_format0.mid'



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



    """simple docstring"""



    UpperCAmelCase    			=       SpectrogramDiffusionPipeline
    UpperCAmelCase    			=       PipelineTesterMixin.required_optional_params - {
        "callback",
        "latents",
        "callback_steps",
        "output_type",
        "num_images_per_prompt",
    }
    UpperCAmelCase    			=       False
    UpperCAmelCase    			=       False
    UpperCAmelCase    			=       TOKENS_TO_AUDIO_GENERATION_PARAMS
    UpperCAmelCase    			=       TOKENS_TO_AUDIO_GENERATION_BATCH_PARAMS

    def 					UpperCamelCase						(  self:			Union[str, Any]							):






         """simple docstring"""


         torch.manual_seed(0							)
         A__				      =	SpectrogramNotesEncoder(
             max_length=20_48     ,    vocab_size=15_36     ,    d_model=7_68     ,    dropout_rate=0.1     ,    num_layers=1     ,    num_heads=1     ,    d_kv=4     ,    d_ff=20_48     ,    feed_forward_proj="""gated-gelu"""     ,    )

         A__				      =	SpectrogramContEncoder(
             input_dims=1_28     ,    targets_context_length=2_56     ,    d_model=7_68     ,    dropout_rate=0.1     ,    num_layers=1     ,    num_heads=1     ,    d_kv=4     ,    d_ff=20_48     ,    feed_forward_proj="""gated-gelu"""     ,    )

         A__				      =	TaFilmDecoder(
             input_dims=1_28     ,    targets_length=2_56     ,    max_decoder_noise_time=20_000.0     ,    d_model=7_68     ,    num_layers=1     ,    num_heads=1     ,    d_kv=4     ,    d_ff=20_48     ,    dropout_rate=0.1     ,    )

         A__				      =	DDPMScheduler()

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

    def 					UpperCamelCase						(  self:			Union[str, Any]     ,    UpperCamelCase:			Union[str, Any]     ,    UpperCamelCase:			Any=0							):






         """simple docstring"""


         if str(UpperCamelCase							).startswith("""mps"""							):
              A__				      =	torch.manual_seed(UpperCamelCase							)
         else:
              A__				      =	torch.Generator(device=UpperCamelCase							).manual_seed(UpperCamelCase							)
         A__				      =	{
             """input_tokens""": [
                 [11_34, 90, 11_35, 11_33, 10_80, 1_12, 11_32, 10_80, 11_33, 10_79, 1_33, 11_32, 10_79, 11_33, 1] + [0] * 20_33
             ],
             """generator""": generator,
             """num_inference_steps""": 4,
             """output_type""": """mel""",
         }
         return inputs

    def 					UpperCamelCase						(  self:			List[Any]							):






         """simple docstring"""


         A__				      =	"""cpu"""  # ensure determinism for the device-dependent torch.Generator
         A__				      =	self.get_dummy_components()
         A__				      =	SpectrogramDiffusionPipeline(**UpperCamelCase							)
         A__				      =	pipe.to(UpperCamelCase							)
         pipe.set_progress_bar_config(disable=UpperCamelCase							)

         A__				      =	self.get_dummy_inputs(UpperCamelCase							)
         A__				      =	pipe(**UpperCamelCase							)
         A__				      =	output.audios

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

         assert mel_slice.shape == (3, 3)
         A__				      =	np.array(
             [-11.512_925, -4.788_215, -0.46_172_905, -2.051_715, -10.539_147, -10.970_963, -9.091_634, 4.0, 4.0]							)
         assert np.abs(mel_slice.flatten() - expected_slice							).max() < 1e-2

    @skip_mps
    def 					UpperCamelCase						(  self:			int							):






         """simple docstring"""


         return super().test_save_load_local()

    @skip_mps
    def 					UpperCamelCase						(  self:			List[Any]							):






         """simple docstring"""


         return super().test_dict_tuple_outputs_equivalent()

    @skip_mps
    def 					UpperCamelCase						(  self:			Union[str, Any]							):






         """simple docstring"""


         return super().test_save_load_optional_components()

    @skip_mps
    def 					UpperCamelCase						(  self:			Dict							):






         """simple docstring"""


         return super().test_attention_slicing_forward_pass()

    def 					UpperCamelCase						(  self:			Any							):






         """simple docstring"""


         pass

    def 					UpperCamelCase						(  self:			List[str]							):






         """simple docstring"""


         pass





    @skip_mps
    def 					UpperCamelCase						(  self:			Tuple							):






         """simple docstring"""


         return super().test_progress_bar()



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



    """simple docstring"""



    def 					UpperCamelCase						(  self:			int							):






         """simple docstring"""


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

    def 					UpperCamelCase						(  self:			Tuple							):






         """simple docstring"""


         A__				      =	torch_device

         A__				      =	SpectrogramDiffusionPipeline.from_pretrained("""google/music-spectrogram-diffusion"""							)
         A__				      =	pipe.melgan
         A__				      =	None

         A__				      =	pipe.to(UpperCamelCase							)
         pipe.set_progress_bar_config(disable=UpperCamelCase							)

         def callback(UpperCamelCase:			Optional[Any]     ,    UpperCamelCase:			Optional[int]							):
              # decode mel to audio
              A__				      =	melgan(input_features=mel_output.astype(np.floataa							)							)[0]
              assert len(audio[0]							) == 8_19_20 * (step + 1)
              # simulate that audio is played
              return audio

         A__				      =	MidiProcessor()
         A__				      =	processor(UpperCamelCase							)

         A__				      =	input_tokens[:3]
         A__				      =	torch.manual_seed(0							)
         pipe(UpperCamelCase     ,    num_inference_steps=5     ,    generator=UpperCamelCase     ,    callback=UpperCamelCase     ,    output_type="""mel"""							)

    def 					UpperCamelCase						(  self:			List[Any]							):






         """simple docstring"""


         A__				      =	torch_device

         A__				      =	SpectrogramDiffusionPipeline.from_pretrained("""google/music-spectrogram-diffusion"""							)
         A__				      =	pipe.to(UpperCamelCase							)
         pipe.set_progress_bar_config(disable=UpperCamelCase							)
         A__				      =	MidiProcessor()

         A__				      =	processor(UpperCamelCase							)
         # just run two denoising loops
         A__				      =	input_tokens[:2]

         A__				      =	torch.manual_seed(0							)
         A__				      =	pipe(UpperCamelCase     ,    num_inference_steps=2     ,    generator=UpperCamelCase							)

         A__				      =	output.audios[0]

         assert abs(np.abs(UpperCamelCase							).sum() - 3_612.841							) < 1e-1





    def 					UpperCamelCase						(  self:			List[str]							):






         """simple docstring"""


         A__				      =	torch_device

         A__				      =	SpectrogramDiffusionPipeline.from_pretrained("""google/music-spectrogram-diffusion"""							)
         A__				      =	pipe.to(UpperCamelCase							)
         pipe.set_progress_bar_config(disable=UpperCamelCase							)

         A__				      =	MidiProcessor()

         A__				      =	processor(UpperCamelCase							)

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

         A__				      =	torch.manual_seed(0							)
         A__				      =	pipe(UpperCamelCase     ,    num_inference_steps=1_00     ,    generator=UpperCamelCase							)

         A__				      =	output.audios[0]
         assert abs(np.abs(UpperCamelCase							).sum() - 9_389.1_111							) < 5e-2









"""simple docstring"""



from typing import List

from tokenizers import NormalizedString, PreTokenizedString, normalizers



class        a    :



    """simple docstring"""



    def __init__(  self:			List[Any]     ,    UpperCamelCase:			str							):






         """simple docstring"""


         A__				      =	vocab
         A__				      =	normalizers.BertNormalizer(
             clean_text=UpperCamelCase     ,    handle_chinese_chars=UpperCamelCase     ,    strip_accents=UpperCamelCase     ,    lowercase=UpperCamelCase     ,    )
         try:
              import rjieba
         except ImportError:
              raise ImportError(
                  """You need to install rjieba to use RoFormerTokenizer. """
                  """See https://pypi.org/project/rjieba/ for installation."""							)
         A__				      =	rjieba

    def 					UpperCamelCase						(  self:			Optional[Any]     ,    UpperCamelCase:			int     ,    UpperCamelCase:			NormalizedString							):






         """simple docstring"""


         A__				      =	[]

         # this code slice normalized_string is too slow (6s) but test_alignement_methods can pass
         for token, start, end in self.jieba.tokenize(str(UpperCamelCase							)     ,    hmm=UpperCamelCase							):
              if token in self.vocab:
                   splits.append(normalized_string[start:end]							)
              else:
                   A__				      =	self.normalizers.normalize_str(UpperCamelCase							).split()
                   for token in token_list:
                        if token:
                             A__				      =	start + len(UpperCamelCase							)
                             splits.append(normalized_string[start:end]							)
                             A__				      =	end

        # this code test_alignement_methods can't pass but fast (300ms)
        # for token in self.jieba.cut(str(normalized_string), False):
        #     if token in self.vocab:
        #         splits.append(NormalizedString(token))
        #     else:
        #         token_list = self.normalizers.normalize_str(token).split()
        #         for token in token_list:
        #             if token:
        #                 splits.append(NormalizedString(token))

         return splits





    def 					UpperCamelCase						(  self:			Tuple     ,    UpperCamelCase:			PreTokenizedString							):






         """simple docstring"""


         pretok.split(self.jieba_split							)






