


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


__lowerCamelCase       				=	'''./tests/fixtures/elise_format0.mid'''




@unittest.skip("The note-seq package currently throws an error on import"	)
class 							UpperCAmelCase ( _snake_case	,   unittest.TestCase	):
 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 					__SCREAMING_SNAKE_CASE				(			self       :  Optional[int]   ):
    torch.manual_seed(0   )
    UpperCAmelCase__  :Optional[Any]            =			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''' ,					)

    UpperCAmelCase__  :List[str]            =			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''' ,					)

    UpperCAmelCase__  :Dict            =			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 ,					)

    UpperCAmelCase__  :List[str]            =			DDPMScheduler()

    UpperCAmelCase__  :List[str]            =			{
        '''notes_encoder''': notes_encoder.eval(),
        '''continuous_encoder''': continuous_encoder.eval(),
        '''decoder''': decoder.eval(),
        '''scheduler''': scheduler,
        '''melgan''': None,
    }
    return components


 def 					__SCREAMING_SNAKE_CASE				(			self       :  int ,					__lowerCamelCase       :  Any ,					__lowerCamelCase       :  str=0   ):
    if str(__lowerCamelCase   ).startswith('''mps'''   ):
       UpperCAmelCase__  :List[str]            =			torch.manual_seed(__lowerCamelCase   )
    else:
       UpperCAmelCase__  :Optional[int]            =			torch.Generator(device=__lowerCamelCase   ).manual_seed(__lowerCamelCase   )
    UpperCAmelCase__  :int            =			{
        '''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 					__SCREAMING_SNAKE_CASE				(			self       :  List[str]   ):
    UpperCAmelCase__  :Any            =			'''cpu'''  # ensure determinism for the device-dependent torch.Generator
    UpperCAmelCase__  :List[str]            =			self.get_dummy_components()
    UpperCAmelCase__  :Dict            =			SpectrogramDiffusionPipeline(**__lowerCamelCase   )
    UpperCAmelCase__  :Optional[Any]            =			pipe.to(__lowerCamelCase   )
    pipe.set_progress_bar_config(disable=__lowerCamelCase   )

    UpperCAmelCase__  :int            =			self.get_dummy_inputs(__lowerCamelCase   )
    UpperCAmelCase__  :List[Any]            =			pipe(**__lowerCamelCase   )
    UpperCAmelCase__  :int            =			output.audios

    UpperCAmelCase__  :Optional[int]            =			mel[0, -3:, -3:]

    assert mel_slice.shape == (3, 3)
    UpperCAmelCase__  :Union[str, Any]            =			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 					__SCREAMING_SNAKE_CASE				(			self       :  List[str]   ):
    return super().test_save_load_local()


 @skip_mps
 def 					__SCREAMING_SNAKE_CASE				(			self       :  List[Any]   ):
    return super().test_dict_tuple_outputs_equivalent()


 @skip_mps
 def 					__SCREAMING_SNAKE_CASE				(			self       :  List[Any]   ):
    return super().test_save_load_optional_components()


 @skip_mps
 def 					__SCREAMING_SNAKE_CASE				(			self       :  Optional[int]   ):
    return super().test_attention_slicing_forward_pass()


 def 					__SCREAMING_SNAKE_CASE				(			self       :  int   ):
    pass


 def 					__SCREAMING_SNAKE_CASE				(			self       :  List[str]   ):
    pass





 @skip_mps
 def 					__SCREAMING_SNAKE_CASE				(			self       :  Optional[Any]   ):
    return super().test_progress_bar()




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


 def 					__SCREAMING_SNAKE_CASE				(			self       :  str   ):
    # clean up the VRAM after each test
    super().tearDown()
    gc.collect()
    torch.cuda.empty_cache()


 def 					__SCREAMING_SNAKE_CASE				(			self       :  str   ):
    # TODO - test that pipeline can decode tokens in a callback
    # so that music can be played live
    UpperCAmelCase__  :Optional[Any]            =			torch_device

    UpperCAmelCase__  :List[Any]            =			SpectrogramDiffusionPipeline.from_pretrained('''google/music-spectrogram-diffusion'''   )
    UpperCAmelCase__  :str            =			pipe.melgan
    UpperCAmelCase__  :Optional[int]            =			None

    UpperCAmelCase__  :int            =			pipe.to(__lowerCamelCase   )
    pipe.set_progress_bar_config(disable=__lowerCamelCase   )

    def callback(__lowerCamelCase       :  Optional[Any] ,					__lowerCamelCase       :  List[Any]   ):
       # decode mel to audio
       UpperCAmelCase__  :List[Any]            =			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

    UpperCAmelCase__  :int            =			MidiProcessor()
    UpperCAmelCase__  :int            =			processor(__lowerCamelCase   )

    UpperCAmelCase__  :Dict            =			input_tokens[:3]
    UpperCAmelCase__  :Optional[Any]            =			torch.manual_seed(0   )
    pipe(__lowerCamelCase ,					num_inference_steps=5 ,					generator=__lowerCamelCase ,					callback=__lowerCamelCase ,					output_type='''mel'''   )


 def 					__SCREAMING_SNAKE_CASE				(			self       :  List[Any]   ):
    UpperCAmelCase__  :Tuple            =			torch_device

    UpperCAmelCase__  :Dict            =			SpectrogramDiffusionPipeline.from_pretrained('''google/music-spectrogram-diffusion'''   )
    UpperCAmelCase__  :str            =			pipe.to(__lowerCamelCase   )
    pipe.set_progress_bar_config(disable=__lowerCamelCase   )
    UpperCAmelCase__  :Dict            =			MidiProcessor()

    UpperCAmelCase__  :Any            =			processor(__lowerCamelCase   )
    # just run two denoising loops
    UpperCAmelCase__  :List[str]            =			input_tokens[:2]

    UpperCAmelCase__  :str            =			torch.manual_seed(0   )
    UpperCAmelCase__  :List[str]            =			pipe(__lowerCamelCase ,					num_inference_steps=2 ,					generator=__lowerCamelCase   )

    UpperCAmelCase__  :List[Any]            =			output.audios[0]

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





 def 					__SCREAMING_SNAKE_CASE				(			self       :  int   ):
    UpperCAmelCase__  :List[Any]            =			torch_device

    UpperCAmelCase__  :List[Any]            =			SpectrogramDiffusionPipeline.from_pretrained('''google/music-spectrogram-diffusion'''   )
    UpperCAmelCase__  :Dict            =			pipe.to(__lowerCamelCase   )
    pipe.set_progress_bar_config(disable=__lowerCamelCase   )

    UpperCAmelCase__  :str            =			MidiProcessor()

    UpperCAmelCase__  :List[Any]            =			processor(__lowerCamelCase   )

    # just run 4 denoising loops
    UpperCAmelCase__  :Union[str, Any]            =			input_tokens[:4]

    UpperCAmelCase__  :str            =			torch.manual_seed(0   )
    UpperCAmelCase__  :List[Any]            =			pipe(__lowerCamelCase ,					num_inference_steps=1_0_0 ,					generator=__lowerCamelCase   )

    UpperCAmelCase__  :List[Any]            =			output.audios[0]
    assert abs(np.abs(__lowerCamelCase   ).sum() - 93_89.11_11   ) < 5e-2







'''simple docstring'''





import json
import os
import unittest

from transformers.models.fsmt.tokenization_fsmt import VOCAB_FILES_NAMES, FSMTTokenizer
from transformers.testing_utils import slow
from transformers.utils import cached_property

from ...test_tokenization_common import TokenizerTesterMixin


# using a different tiny model than the one used for default params defined in init to ensure proper testing
__lowerCamelCase       				=	'''stas/tiny-wmt19-en-ru'''




class 							UpperCAmelCase ( _snake_case	,   unittest.TestCase	):
 UpperCAmelCase				=				FSMTTokenizer
 UpperCAmelCase				=				False


 def 					__SCREAMING_SNAKE_CASE				(			self       :  Union[str, Any]   ):
    super().setUp()

    # Adapted from Sennrich et al. 2015 and https://github.com/rsennrich/subword-nmt
    UpperCAmelCase__  :Any            =			[
        '''l''',
        '''o''',
        '''w''',
        '''e''',
        '''r''',
        '''s''',
        '''t''',
        '''i''',
        '''d''',
        '''n''',
        '''w</w>''',
        '''r</w>''',
        '''t</w>''',
        '''lo''',
        '''low''',
        '''er</w>''',
        '''low</w>''',
        '''lowest</w>''',
        '''newer</w>''',
        '''wider</w>''',
        '''<unk>''',
    ]
    UpperCAmelCase__  :Dict            =			dict(zip(__lowerCamelCase ,					range(len(__lowerCamelCase   )   )   )   )
    UpperCAmelCase__  :Dict            =			['''l o 123''', '''lo w 1456''', '''e r</w> 1789''', '''''']

    UpperCAmelCase__  :Dict            =			['''en''', '''ru''']
    UpperCAmelCase__  :Tuple            =			{
        '''langs''': self.langs,
        '''src_vocab_size''': 1_0,
        '''tgt_vocab_size''': 2_0,
    }

    UpperCAmelCase__  :Any            =			os.path.join(self.tmpdirname ,					VOCAB_FILES_NAMES['''src_vocab_file''']   )
    UpperCAmelCase__  :Any            =			os.path.join(self.tmpdirname ,					VOCAB_FILES_NAMES['''tgt_vocab_file''']   )
    UpperCAmelCase__  :Dict            =			os.path.join(self.tmpdirname ,					'''tokenizer_config.json'''   )
    UpperCAmelCase__  :Union[str, Any]            =			os.path.join(self.tmpdirname ,					VOCAB_FILES_NAMES['''merges_file''']   )
    with open(self.src_vocab_file ,					'''w'''   ) as fp:
       fp.write(json.dumps(__lowerCamelCase   )   )
    with open(self.tgt_vocab_file ,					'''w'''   ) as fp:
       fp.write(json.dumps(__lowerCamelCase   )   )
    with open(self.merges_file ,					'''w'''   ) as fp:
       fp.write('''\n'''.join(__lowerCamelCase   )   )
    with open(__lowerCamelCase ,					'''w'''   ) as fp:
       fp.write(json.dumps(__lowerCamelCase   )   )


 @cached_property
 def 					__SCREAMING_SNAKE_CASE				(			self       :  List[Any]   ):
    return FSMTTokenizer.from_pretrained('''facebook/wmt19-ru-en'''   )


 @cached_property
 def 					__SCREAMING_SNAKE_CASE				(			self       :  Tuple   ):
    return FSMTTokenizer.from_pretrained('''facebook/wmt19-en-ru'''   )


 def 					__SCREAMING_SNAKE_CASE				(			self       :  Optional[Any]   ):
    UpperCAmelCase__  :int            =			FSMTTokenizer.from_pretrained(__lowerCamelCase   )
    self.assertListEqual([tokenizer.src_lang, tokenizer.tgt_lang] ,					['''en''', '''ru''']   )
    self.assertEqual(tokenizer.src_vocab_size ,					2_1   )
    self.assertEqual(tokenizer.tgt_vocab_size ,					2_1   )


 def 					__SCREAMING_SNAKE_CASE				(			self       :  Optional[int]   ):
    UpperCAmelCase__  :Tuple            =			FSMTTokenizer(self.langs ,					self.src_vocab_file ,					self.tgt_vocab_file ,					self.merges_file   )

    UpperCAmelCase__  :Any            =			'''lower'''
    UpperCAmelCase__  :List[Any]            =			['''low''', '''er</w>''']
    UpperCAmelCase__  :Union[str, Any]            =			tokenizer.tokenize(__lowerCamelCase   )
    self.assertListEqual(__lowerCamelCase ,					__lowerCamelCase   )

    UpperCAmelCase__  :Dict            =			tokens + ['''<unk>''']
    UpperCAmelCase__  :Union[str, Any]            =			[1_4, 1_5, 2_0]
    self.assertListEqual(tokenizer.convert_tokens_to_ids(__lowerCamelCase   ) ,					__lowerCamelCase   )


 @slow
 def 					__SCREAMING_SNAKE_CASE				(			self       :  Any   ):
    UpperCAmelCase__  :str            =			self.tokenizer_ru_en

    UpperCAmelCase__  :Union[str, Any]            =			tokenizer.encode('''sequence builders''' ,					add_special_tokens=__lowerCamelCase   )
    UpperCAmelCase__  :Dict            =			tokenizer.encode('''multi-sequence build''' ,					add_special_tokens=__lowerCamelCase   )

    UpperCAmelCase__  :Any            =			tokenizer.build_inputs_with_special_tokens(__lowerCamelCase   )
    UpperCAmelCase__  :Union[str, Any]            =			tokenizer.build_inputs_with_special_tokens(__lowerCamelCase ,					__lowerCamelCase   )

    assert encoded_sentence == text + [2]
    assert encoded_pair == text + [2] + text_a + [2]


 @slow
 def 					__SCREAMING_SNAKE_CASE				(			self       :  Union[str, Any]   ):
    UpperCAmelCase__  :List[str]            =			self.tokenizer_en_ru
    UpperCAmelCase__  :Union[str, Any]            =			self.tokenizer_ru_en

    UpperCAmelCase__  :Any            =			[
        [
            '''Here\'s a little song I wrote. Don\'t worry, be happy.''',
            [2_4_7_0, 3_9, 1_1, 2_3_4_9, 7_2_2_2, 7_0, 5_9_7_9, 7, 8_4_5_0, 1_0_5_0, 1_3_1_6_0, 5, 2_6, 6_4_4_5, 7, 2],
        ],
        ['''This is it. No more. I\'m done!''', [1_3_2, 2_1, 3_7, 7, 1_4_3_4, 8_6, 7, 7_0, 6_4_7_6, 1_3_0_5, 4_2_7, 2]],
    ]

    # if data needs to be recreated or added, run:
    # import torch
    # model = torch.hub.load("pytorch/fairseq", "transformer.wmt19.en-ru", checkpoint_file="model4.pt", tokenizer="moses", bpe="fastbpe")
    # for src_text, _ in targets: print(f"""[\n"{src_text}",\n {model.encode(src_text).tolist()}\n],""")

    for src_text, tgt_input_ids in targets:
       UpperCAmelCase__  :List[Any]            =			tokenizer_enc.encode(__lowerCamelCase ,					return_tensors=__lowerCamelCase   )
       self.assertListEqual(__lowerCamelCase ,					__lowerCamelCase   )

       # and decode backward, using the reversed languages model
       UpperCAmelCase__  :str            =			tokenizer_dec.decode(__lowerCamelCase ,					skip_special_tokens=__lowerCamelCase   )
       self.assertEqual(__lowerCamelCase ,					__lowerCamelCase   )


 @slow
 def 					__SCREAMING_SNAKE_CASE				(			self       :  Union[str, Any]   ):
    UpperCAmelCase__  :Optional[Any]            =			FSMTTokenizer.from_pretrained('''facebook/wmt19-ru-en''' ,					do_lower_case=__lowerCamelCase   )
    UpperCAmelCase__  :str            =			tokenizer.tokenize('''USA is United States of America'''   )
    UpperCAmelCase__  :str            =			['''us''', '''a</w>''', '''is</w>''', '''un''', '''i''', '''ted</w>''', '''st''', '''ates</w>''', '''of</w>''', '''am''', '''er''', '''ica</w>''']
    self.assertListEqual(__lowerCamelCase ,					__lowerCamelCase   )


 @unittest.skip('''FSMTConfig.__init__  requires non-optional args'''   )
 def 					__SCREAMING_SNAKE_CASE				(			self       :  List[str]   ):
    pass





 @unittest.skip('''FSMTConfig.__init__  requires non-optional args'''   )
 def 					__SCREAMING_SNAKE_CASE				(			self       :  Dict   ):
    pass



