




"""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__			:     str			       =				"./tests/fixtures/elise_format0.mid"
@unittest.skip("""The note-seq package currently throws an error on import"""  )
class      A__   (   _lowerCAmelCase  ,						unittest.TestCase  ):
	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 				a_		( self	:						List[Any]						)       ->			Dict:
						torch.manual_seed(0						)
						__snake_case        =						SpectrogramNotesEncoder(
						    max_length=2_048   ,     vocab_size=1_536   ,     d_model=768   ,     dropout_rate=0.1   ,     num_layers=1   ,     num_heads=1   ,     d_kv=4   ,     d_ff=2_048   ,     feed_forward_proj="""gated-gelu"""   ,     )

						__snake_case        =						SpectrogramContEncoder(
						    input_dims=128   ,     targets_context_length=256   ,     d_model=768   ,     dropout_rate=0.1   ,     num_layers=1   ,     num_heads=1   ,     d_kv=4   ,     d_ff=2_048   ,     feed_forward_proj="""gated-gelu"""   ,     )

						__snake_case        =						TaFilmDecoder(
						    input_dims=128   ,     targets_length=256   ,     max_decoder_noise_time=20_000.0   ,     d_model=768   ,     num_layers=1   ,     num_heads=1   ,     d_kv=4   ,     d_ff=2_048   ,     dropout_rate=0.1   ,     )

						__snake_case        =						DDPMScheduler()

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




	def 				a_		( self	:						Tuple   ,     lowercase	:						Optional[Any]   ,     lowercase	:						int=0						)       ->			Dict:
						if str(_lowerCamelCase						).startswith("""mps"""						):
											__snake_case        =						torch.manual_seed(_lowerCamelCase						)
						else:
											__snake_case        =						torch.Generator(device=_lowerCamelCase						).manual_seed(_lowerCamelCase						)
						__snake_case        =						{
						    """input_tokens""": [
						        [1_134, 90, 1_135, 1_133, 1_080, 112, 1_132, 1_080, 1_133, 1_079, 133, 1_132, 1_079, 1_133, 1] + [0] * 2_033
						    ],
						    """generator""": generator,
						    """num_inference_steps""": 4,
						    """output_type""": """mel""",
						}
						return inputs




	def 				a_		( self	:						str						)       ->			Optional[Any]:
						__snake_case        =						"""cpu"""  # ensure determinism for the device-dependent torch.Generator
						__snake_case        =						self.get_dummy_components()
						__snake_case        =						SpectrogramDiffusionPipeline(**_lowerCamelCase						)
						__snake_case        =						pipe.to(_lowerCamelCase						)
						pipe.set_progress_bar_config(disable=_lowerCamelCase						)

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

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

						assert mel_slice.shape == (3, 3)
						__snake_case        =						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 				a_		( self	:						List[str]						)       ->			Optional[int]:
						return super().test_save_load_local()




	@skip_mps
	def 				a_		( self	:						Dict						)       ->			int:
						return super().test_dict_tuple_outputs_equivalent()




	@skip_mps
	def 				a_		( self	:						Union[str, Any]						)       ->			Any:
						return super().test_save_load_optional_components()




	@skip_mps
	def 				a_		( self	:						Union[str, Any]						)       ->			Any:
						return super().test_attention_slicing_forward_pass()




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




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




	@skip_mps
	def 				a_		( self	:						Optional[int]						)       ->			Optional[int]:
						return super().test_progress_bar()

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




	def 				a_		( self	:						int						)       ->			int:
						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()




	def 				a_		( self	:						str						)       ->			List[Any]:
						__snake_case        =						torch_device

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

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

						def callback(lowercase	:						Tuple   ,     lowercase	:						str						):
											# decode mel to audio
											__snake_case        =						melgan(input_features=mel_output.astype(np.floataa						)						)[0]
											assert len(audio[0]						) == 81_920 * (step + 1)
											# simulate that audio is played
											return audio

						__snake_case        =						MidiProcessor()
						__snake_case        =						processor(_lowerCamelCase						)

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




	def 				a_		( self	:						Union[str, Any]						)       ->			Tuple:
						__snake_case        =						torch_device

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

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

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

						__snake_case        =						output.audios[0]

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




	def 				a_		( self	:						int						)       ->			List[str]:
						__snake_case        =						torch_device

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

						__snake_case        =						MidiProcessor()

						__snake_case        =						processor(_lowerCamelCase						)

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

						__snake_case        =						torch.manual_seed(0						)
						__snake_case        =						pipe(_lowerCamelCase   ,     num_inference_steps=100   ,     generator=_lowerCamelCase						)

						__snake_case        =						output.audios[0]
						assert abs(np.abs(_lowerCamelCase						).sum() - 9_389.1_111						) < 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__)




