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


UpperCamelCase__ :						Any    							=					"./tests/fixtures/elise_format0.mid"

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





       """simple docstring"""




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






       def     UpperCamelCase			(							self )      ->     Dict:
          torch.manual_seed(0 )
          _SCREAMING_SNAKE_CASE       	=			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"""    ,							)

          _SCREAMING_SNAKE_CASE       	=			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"""    ,							)

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

          _SCREAMING_SNAKE_CASE       	=			DDPMScheduler()

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






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

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

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

          assert mel_slice.shape == (3, 3)
          _SCREAMING_SNAKE_CASE       	=			np.array(
              [-11.51_2925, -4.78_8215, -0.4617_2905, -2.05_1715, -10.53_9147, -10.97_0963, -9.09_1634, 4.0, 4.0] )
          assert np.abs(mel_slice.flatten() - expected_slice ).max() < 1E-2






       @skip_mps
       def     UpperCamelCase			(							self )      ->     Tuple:
          return super().test_save_load_local()






       @skip_mps
       def     UpperCamelCase			(							self )      ->     List[str]:
          return super().test_dict_tuple_outputs_equivalent()






       @skip_mps
       def     UpperCamelCase			(							self )      ->     List[str]:
          return super().test_save_load_optional_components()






       @skip_mps
       def     UpperCamelCase			(							self )      ->     Tuple:
          return super().test_attention_slicing_forward_pass()






       def     UpperCamelCase			(							self )      ->     List[Any]:
          pass






       def     UpperCamelCase			(							self )      ->     Dict:
          pass






       @skip_mps
       def     UpperCamelCase			(							self )      ->     str:
          return super().test_progress_bar()






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





       """simple docstring"""




       def     UpperCamelCase			(							self )      ->     int:
          # clean up the VRAM after each test
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()






       def     UpperCamelCase			(							self )      ->     str:
          # TODO - test that pipeline can decode tokens in a callback
          # so that music can be played live
          _SCREAMING_SNAKE_CASE       	=			torch_device

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

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

          def callback(A__    ,							A__ ):
             # decode mel to audio
             _SCREAMING_SNAKE_CASE       	=			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

          _SCREAMING_SNAKE_CASE       	=			MidiProcessor()
          _SCREAMING_SNAKE_CASE       	=			processor(A__ )

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






       def     UpperCamelCase			(							self )      ->     Tuple:
          _SCREAMING_SNAKE_CASE       	=			torch_device

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

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

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

          _SCREAMING_SNAKE_CASE       	=			output.audios[0]

          assert abs(np.abs(A__ ).sum() - 3612.841 ) < 1E-1






       def     UpperCamelCase			(							self )      ->     Tuple:
          _SCREAMING_SNAKE_CASE       	=			torch_device

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

          _SCREAMING_SNAKE_CASE       	=			MidiProcessor()

          _SCREAMING_SNAKE_CASE       	=			processor(A__ )

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

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

          _SCREAMING_SNAKE_CASE       	=			output.audios[0]
          assert abs(np.abs(A__ ).sum() - 9389.1111 ) < 5E-2








'''simple docstring'''



#
# Sort large text files in a minimum amount of memory
#
import argparse
import os

class 				_a			:





       """simple docstring"""




       SCREAMING_SNAKE_CASE =					'block_{0}.dat'






       def __init__(							self    ,							A__ )      ->     Union[str, Any]:
          _SCREAMING_SNAKE_CASE       	=			filename
          _SCREAMING_SNAKE_CASE       	=			[]






       def     UpperCamelCase			(							self    ,							A__    ,							A__ )      ->     Optional[Any]:
          _SCREAMING_SNAKE_CASE       	=			self.BLOCK_FILENAME_FORMAT.format(A__ )
          with open(A__    ,							"""w""" ) as file:
             file.write(A__ )
          self.block_filenames.append(A__ )






       def     UpperCamelCase			(							self )      ->     Tuple:
          return self.block_filenames






       def     UpperCamelCase			(							self    ,							A__    ,							A__=None )      ->     Any:
          _SCREAMING_SNAKE_CASE       	=			0
          with open(self.filename ) as file:
             while True:
                _SCREAMING_SNAKE_CASE       	=			file.readlines(A__ )

                if lines == []:
                   break

                if sort_key is None:
                   lines.sort()
                else:
                   lines.sort(key=A__ )

                self.write_block("""""".join(A__ )    ,							A__ )
                i += 1






       def     UpperCamelCase			(							self )      ->     List[Any]:
          map(os.remove    ,							self.block_filenames )






class 				_a			:





       """simple docstring"""




       def     UpperCamelCase			(							self    ,							A__ )      ->     str:
          _SCREAMING_SNAKE_CASE       	=			-1
          _SCREAMING_SNAKE_CASE       	=			None

          for i in range(len(A__ ) ):
             if min_str is None or choices[i] < min_str:
                _SCREAMING_SNAKE_CASE       	=			i

          return min_index






class 				_a			:





       """simple docstring"""




       def __init__(							self    ,							A__ )      ->     Dict:
          _SCREAMING_SNAKE_CASE       	=			files
          _SCREAMING_SNAKE_CASE       	=			set()
          _SCREAMING_SNAKE_CASE       	=			len(A__ )
          _SCREAMING_SNAKE_CASE       	=			{i: None for i in range(self.num_buffers )}






       def     UpperCamelCase			(							self )      ->     Dict:
          return {
              i: self.buffers[i] for i in range(self.num_buffers ) if i not in self.empty
          }






       def     UpperCamelCase			(							self )      ->     Tuple:
          for i in range(self.num_buffers ):
             if self.buffers[i] is None and i not in self.empty:
                _SCREAMING_SNAKE_CASE       	=			self.files[i].readline()

                if self.buffers[i] == "":
                   self.empty.add(A__ )
                   self.files[i].close()

          if len(self.empty ) == self.num_buffers:
             return False

          return True






       def     UpperCamelCase			(							self    ,							A__ )      ->     List[str]:
          _SCREAMING_SNAKE_CASE       	=			self.buffers[index]
          _SCREAMING_SNAKE_CASE       	=			None

          return value






class 				_a			:





       """simple docstring"""




       def __init__(							self    ,							A__ )      ->     str:
          _SCREAMING_SNAKE_CASE       	=			merge_strategy






       def     UpperCamelCase			(							self    ,							A__    ,							A__    ,							A__ )      ->     Union[str, Any]:
          _SCREAMING_SNAKE_CASE       	=			FilesArray(self.get_file_handles(A__    ,							A__ ) )
          with open(A__    ,							"""w"""    ,							A__ ) as outfile:
             while buffers.refresh():
                _SCREAMING_SNAKE_CASE       	=			self.merge_strategy.select(buffers.get_dict() )
                outfile.write(buffers.unshift(A__ ) )






       def     UpperCamelCase			(							self    ,							A__    ,							A__ )      ->     str:
          _SCREAMING_SNAKE_CASE       	=			{}

          for i in range(len(A__ ) ):
             _SCREAMING_SNAKE_CASE       	=			open(filenames[i]    ,							"""r"""    ,							A__ )  # noqa: UP015

          return files






class 				_a			:





       """simple docstring"""




       def __init__(							self    ,							A__ )      ->     Union[str, Any]:
          _SCREAMING_SNAKE_CASE       	=			block_size






       def     UpperCamelCase			(							self    ,							A__    ,							A__=None )      ->     Optional[Any]:
          _SCREAMING_SNAKE_CASE       	=			self.get_number_blocks(A__    ,							self.block_size )
          _SCREAMING_SNAKE_CASE       	=			FileSplitter(A__ )
          splitter.split(self.block_size    ,							A__ )

          _SCREAMING_SNAKE_CASE       	=			FileMerger(NWayMerge() )
          _SCREAMING_SNAKE_CASE       	=			self.block_size / (num_blocks + 1)
          merger.merge(splitter.get_block_filenames()    ,							filename + """.out"""    ,							A__ )

          splitter.cleanup()






       def     UpperCamelCase			(							self    ,							A__    ,							A__ )      ->     Any:
          return (os.stat(A__ ).st_size / block_size) + 1






def 				lowerCAmelCase_		(			SCREAMING_SNAKE_CASE_  )       ->       str:




   """simple docstring"""
   if string[-1].lower() == "k":
      return int(string[:-1]  ) * 10_24
   elif string[-1].lower() == "m":
      return int(string[:-1]  ) * 10_24 * 10_24
   elif string[-1].lower() == "g":
      return int(string[:-1]  ) * 10_24 * 10_24 * 10_24
   else:
      return int(SCREAMING_SNAKE_CASE_  )
def 				lowerCAmelCase_		(			)       ->       List[str]:




   """simple docstring"""
   _SCREAMING_SNAKE_CASE       	=			argparse.ArgumentParser()
   parser.add_argument(
       """-m"""      , """--mem"""      , help="""amount of memory to use for sorting"""      , default="""100M"""  )
   parser.add_argument(
       """filename"""      , metavar="""<filename>"""      , nargs=1      , help="""name of file to sort"""  )
   _SCREAMING_SNAKE_CASE       	=			parser.parse_args()

   _SCREAMING_SNAKE_CASE       	=			ExternalSort(parse_memory(args.mem  )  )
   sorter.sort(args.filename[0]  )


if __name__ == "__main__":
     main()






