






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


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


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


   '''simple docstring'''

   A_		:      Any							=					SpectrogramDiffusionPipeline
   A_		:      List[str]							=					PipelineTesterMixin.required_optional_params - {
       "callback",
       "latents",
       "callback_steps",
       "output_type",
       "num_images_per_prompt",
   }
   A_		:      Dict							=					False
   A_		:      Union[str, Any]							=					False
   A_		:      List[str]							=					TOKENS_TO_AUDIO_GENERATION_PARAMS
   A_		:      str							=					TOKENS_TO_AUDIO_GENERATION_BATCH_PARAMS


   def 							_SCREAMING_SNAKE_CASE							(      self					:	List[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=2_00_00.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 							_SCREAMING_SNAKE_CASE							(      self					:	Tuple,			_lowerCamelCase					:	Optional[Any],			_lowerCamelCase					:	int=0    ):

     '''simple docstring'''
     if str(_lowerCamelCase    ).startswith('''mps'''    ):
       __A		     =					torch.manual_seed(_lowerCamelCase    )
     else:
       __A		     =					torch.Generator(device=_lowerCamelCase    ).manual_seed(_lowerCamelCase    )
     __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 							_SCREAMING_SNAKE_CASE							(      self					:	str    ):

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

     __A		     =					self.get_dummy_inputs(_lowerCamelCase    )
     __A		     =					pipe(**_lowerCamelCase    )
     __A		     =					output.audios

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

     assert mel_slice.shape == (3, 3)
     __A		     =					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]    ):

     '''simple docstring'''
     return super().test_save_load_local()


   @skip_mps
   def 							_SCREAMING_SNAKE_CASE							(      self					:	Dict    ):

     '''simple docstring'''
     return super().test_dict_tuple_outputs_equivalent()


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

     '''simple docstring'''
     return super().test_save_load_optional_components()


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

     '''simple docstring'''
     return super().test_attention_slicing_forward_pass()


   def 							_SCREAMING_SNAKE_CASE							(      self					:	List[Any]    ):

     '''simple docstring'''
     pass


   def 							_SCREAMING_SNAKE_CASE							(      self					:	List[Any]    ):

     '''simple docstring'''
     pass






   @skip_mps
   def 							_SCREAMING_SNAKE_CASE							(      self					:	Optional[int]    ):

     '''simple docstring'''
     return super().test_progress_bar()





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


   '''simple docstring'''

   def 							_SCREAMING_SNAKE_CASE							(      self					:	int    ):

     '''simple docstring'''
     # clean up the VRAM after each test
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()


   def 							_SCREAMING_SNAKE_CASE							(      self					:	str    ):

     '''simple docstring'''
     # TODO - test that pipeline can decode tokens in a callback
     # so that music can be played live
     __A		     =					torch_device

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

     __A		     =					pipe.to(_lowerCamelCase    )
     pipe.set_progress_bar_config(disable=_lowerCamelCase    )

     def callback(_lowerCamelCase					:	Tuple,			_lowerCamelCase					:	str    ):
       # 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(_lowerCamelCase    )

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


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

     '''simple docstring'''
     __A		     =					torch_device

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

     __A		     =					processor(_lowerCamelCase    )
     # just run two denoising loops
     __A		     =					input_tokens[:2]

     __A		     =					torch.manual_seed(0    )
     __A		     =					pipe(_lowerCamelCase,			num_inference_steps=2,			generator=_lowerCamelCase    )

     __A		     =					output.audios[0]

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






   def 							_SCREAMING_SNAKE_CASE							(      self					:	int    ):

     '''simple docstring'''
     __A		     =					torch_device

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

     __A		     =					MidiProcessor()

     __A		     =					processor(_lowerCamelCase    )

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

     __A		     =					torch.manual_seed(0    )
     __A		     =					pipe(_lowerCamelCase,			num_inference_steps=1_00,			generator=_lowerCamelCase    )

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












"""simple docstring"""


from typing import TYPE_CHECKING

from ..utils import OptionalDependencyNotAvailable, _LazyModule, is_flax_available, is_tf_available, is_torch_available


lowercase_								=	{
    'configuration_utils': ['GenerationConfig'],
    'streamers': ['TextIteratorStreamer', 'TextStreamer'],
}

try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       lowercase_								=	[
           'Constraint',
           'ConstraintListState',
           'DisjunctiveConstraint',
           'PhrasalConstraint',
       ]
       lowercase_								=	[
           'BeamHypotheses',
           'BeamScorer',
           'BeamSearchScorer',
           'ConstrainedBeamSearchScorer',
       ]
       lowercase_								=	[
           'EpsilonLogitsWarper',
           'EtaLogitsWarper',
           'ForcedBOSTokenLogitsProcessor',
           'ForcedEOSTokenLogitsProcessor',
           'HammingDiversityLogitsProcessor',
           'InfNanRemoveLogitsProcessor',
           'LogitsProcessor',
           'LogitsProcessorList',
           'LogitsWarper',
           'MinLengthLogitsProcessor',
           'MinNewTokensLengthLogitsProcessor',
           'NoBadWordsLogitsProcessor',
           'NoRepeatNGramLogitsProcessor',
           'PrefixConstrainedLogitsProcessor',
           'RepetitionPenaltyLogitsProcessor',
           'SequenceBiasLogitsProcessor',
           'EncoderRepetitionPenaltyLogitsProcessor',
           'TemperatureLogitsWarper',
           'TopKLogitsWarper',
           'TopPLogitsWarper',
           'TypicalLogitsWarper',
           'EncoderNoRepeatNGramLogitsProcessor',
           'ExponentialDecayLengthPenalty',
           'LogitNormalization',
       ]
       lowercase_								=	[
           'MaxNewTokensCriteria',
           'MaxLengthCriteria',
           'MaxTimeCriteria',
           'StoppingCriteria',
           'StoppingCriteriaList',
           'validate_stopping_criteria',
       ]
       lowercase_								=	[
           'GenerationMixin',
           'top_k_top_p_filtering',
           'GreedySearchEncoderDecoderOutput',
           'GreedySearchDecoderOnlyOutput',
           'SampleEncoderDecoderOutput',
           'SampleDecoderOnlyOutput',
           'BeamSearchEncoderDecoderOutput',
           'BeamSearchDecoderOnlyOutput',
           'BeamSampleEncoderDecoderOutput',
           'BeamSampleDecoderOnlyOutput',
           'ContrastiveSearchEncoderDecoderOutput',
           'ContrastiveSearchDecoderOnlyOutput',
       ]

try:
       if not is_tf_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       lowercase_								=	[
           'TFForcedBOSTokenLogitsProcessor',
           'TFForcedEOSTokenLogitsProcessor',
           'TFLogitsProcessor',
           'TFLogitsProcessorList',
           'TFLogitsWarper',
           'TFMinLengthLogitsProcessor',
           'TFNoBadWordsLogitsProcessor',
           'TFNoRepeatNGramLogitsProcessor',
           'TFRepetitionPenaltyLogitsProcessor',
           'TFTemperatureLogitsWarper',
           'TFTopKLogitsWarper',
           'TFTopPLogitsWarper',
           'TFForceTokensLogitsProcessor',
           'TFSuppressTokensAtBeginLogitsProcessor',
           'TFSuppressTokensLogitsProcessor',
       ]
       lowercase_								=	[
           'TFGenerationMixin',
           'tf_top_k_top_p_filtering',
           'TFGreedySearchDecoderOnlyOutput',
           'TFGreedySearchEncoderDecoderOutput',
           'TFSampleEncoderDecoderOutput',
           'TFSampleDecoderOnlyOutput',
           'TFBeamSearchEncoderDecoderOutput',
           'TFBeamSearchDecoderOnlyOutput',
           'TFBeamSampleEncoderDecoderOutput',
           'TFBeamSampleDecoderOnlyOutput',
           'TFContrastiveSearchEncoderDecoderOutput',
           'TFContrastiveSearchDecoderOnlyOutput',
       ]

try:
       if not is_flax_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       lowercase_								=	[
           'FlaxForcedBOSTokenLogitsProcessor',
           'FlaxForcedEOSTokenLogitsProcessor',
           'FlaxLogitsProcessor',
           'FlaxLogitsProcessorList',
           'FlaxLogitsWarper',
           'FlaxMinLengthLogitsProcessor',
           'FlaxTemperatureLogitsWarper',
           'FlaxTopKLogitsWarper',
           'FlaxTopPLogitsWarper',
       ]
       lowercase_								=	[
           'FlaxGenerationMixin',
           'FlaxGreedySearchOutput',
           'FlaxSampleOutput',
           'FlaxBeamSearchOutput',
       ]

if TYPE_CHECKING:
       from .configuration_utils import GenerationConfig
       from .streamers import TextIteratorStreamer, TextStreamer

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .beam_constraints import Constraint, ConstraintListState, DisjunctiveConstraint, PhrasalConstraint
              from .beam_search import BeamHypotheses, BeamScorer, BeamSearchScorer, ConstrainedBeamSearchScorer
              from .logits_process import (
                  EncoderNoRepeatNGramLogitsProcessor,
                  EncoderRepetitionPenaltyLogitsProcessor,
                  EpsilonLogitsWarper,
                  EtaLogitsWarper,
                  ExponentialDecayLengthPenalty,
                  ForcedBOSTokenLogitsProcessor,
                  ForcedEOSTokenLogitsProcessor,
                  HammingDiversityLogitsProcessor,
                  InfNanRemoveLogitsProcessor,
                  LogitNormalization,
                  LogitsProcessor,
                  LogitsProcessorList,
                  LogitsWarper,
                  MinLengthLogitsProcessor,
                  MinNewTokensLengthLogitsProcessor,
                  NoBadWordsLogitsProcessor,
                  NoRepeatNGramLogitsProcessor,
                  PrefixConstrainedLogitsProcessor,
                  RepetitionPenaltyLogitsProcessor,
                  SequenceBiasLogitsProcessor,
                  TemperatureLogitsWarper,
                  TopKLogitsWarper,
                  TopPLogitsWarper,
                  TypicalLogitsWarper,
              )
              from .stopping_criteria import (
                  MaxLengthCriteria,
                  MaxNewTokensCriteria,
                  MaxTimeCriteria,
                  StoppingCriteria,
                  StoppingCriteriaList,
                  validate_stopping_criteria,
              )
              from .utils import (
                  BeamSampleDecoderOnlyOutput,
                  BeamSampleEncoderDecoderOutput,
                  BeamSearchDecoderOnlyOutput,
                  BeamSearchEncoderDecoderOutput,
                  ContrastiveSearchDecoderOnlyOutput,
                  ContrastiveSearchEncoderDecoderOutput,
                  GenerationMixin,
                  GreedySearchDecoderOnlyOutput,
                  GreedySearchEncoderDecoderOutput,
                  SampleDecoderOnlyOutput,
                  SampleEncoderDecoderOutput,
                  top_k_top_p_filtering,
              )

       try:
              if not is_tf_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .tf_logits_process import (
                  TFForcedBOSTokenLogitsProcessor,
                  TFForcedEOSTokenLogitsProcessor,
                  TFForceTokensLogitsProcessor,
                  TFLogitsProcessor,
                  TFLogitsProcessorList,
                  TFLogitsWarper,
                  TFMinLengthLogitsProcessor,
                  TFNoBadWordsLogitsProcessor,
                  TFNoRepeatNGramLogitsProcessor,
                  TFRepetitionPenaltyLogitsProcessor,
                  TFSuppressTokensAtBeginLogitsProcessor,
                  TFSuppressTokensLogitsProcessor,
                  TFTemperatureLogitsWarper,
                  TFTopKLogitsWarper,
                  TFTopPLogitsWarper,
              )
              from .tf_utils import (
                  TFBeamSampleDecoderOnlyOutput,
                  TFBeamSampleEncoderDecoderOutput,
                  TFBeamSearchDecoderOnlyOutput,
                  TFBeamSearchEncoderDecoderOutput,
                  TFContrastiveSearchDecoderOnlyOutput,
                  TFContrastiveSearchEncoderDecoderOutput,
                  TFGenerationMixin,
                  TFGreedySearchDecoderOnlyOutput,
                  TFGreedySearchEncoderDecoderOutput,
                  TFSampleDecoderOnlyOutput,
                  TFSampleEncoderDecoderOutput,
                  tf_top_k_top_p_filtering,
              )

       try:
              if not is_flax_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .flax_logits_process import (
                  FlaxForcedBOSTokenLogitsProcessor,
                  FlaxForcedEOSTokenLogitsProcessor,
                  FlaxLogitsProcessor,
                  FlaxLogitsProcessorList,
                  FlaxLogitsWarper,
                  FlaxMinLengthLogitsProcessor,
                  FlaxTemperatureLogitsWarper,
                  FlaxTopKLogitsWarper,
                  FlaxTopPLogitsWarper,
              )
              from .flax_utils import FlaxBeamSearchOutput, FlaxGenerationMixin, FlaxGreedySearchOutput, FlaxSampleOutput
else:
       import sys

       lowercase_								=	_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)




