
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_  :       Dict					      =     """./tests/fixtures/elise_format0.mid"""

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






			'''simple docstring'''

			__UpperCamelCase					:  int    					=       SpectrogramDiffusionPipeline
			__UpperCamelCase					:  Any    					=       PipelineTesterMixin.required_optional_params - {
			    "callback",
			    "latents",
			    "callback_steps",
			    "output_type",
			    "num_images_per_prompt",
			}
			__UpperCamelCase					:  Optional[int]    					=       False
			__UpperCamelCase					:  Optional[int]    					=       False
			__UpperCamelCase					:  Optional[Any]    					=       TOKENS_TO_AUDIO_GENERATION_PARAMS
			__UpperCamelCase					:  Any    					=       TOKENS_TO_AUDIO_GENERATION_BATCH_PARAMS

			def 					lowerCAmelCase__		(				self       :				Union[str, Any]					):
							torch.manual_seed(0					)
							UpperCamelCase_:							str		=  SpectrogramNotesEncoder(
							    max_length=2048  ,		vocab_size=1536  ,		d_model=768  ,		dropout_rate=0.1  ,		num_layers=1  ,		num_heads=1  ,		d_kv=4  ,		d_ff=2048  ,		feed_forward_proj="""gated-gelu"""  ,		)

							UpperCamelCase_:							Optional[Any]		=  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=2048  ,		feed_forward_proj="""gated-gelu"""  ,		)

							UpperCamelCase_:							Optional[int]		=  TaFilmDecoder(
							    input_dims=128  ,		targets_length=256  ,		max_decoder_noise_time=2_0000.0  ,		d_model=768  ,		num_layers=1  ,		num_heads=1  ,		d_kv=4  ,		d_ff=2048  ,		dropout_rate=0.1  ,		)

							UpperCamelCase_:							List[str]		=  DDPMScheduler()

							UpperCamelCase_:							List[Any]		=  {
							    """notes_encoder""": notes_encoder.eval(),
							    """continuous_encoder""": continuous_encoder.eval(),
							    """decoder""": decoder.eval(),
							    """scheduler""": scheduler,
							    """melgan""": None,
							}
							return components

			def 					lowerCAmelCase__		(				self       :				Any  ,		snake_case_       :				Optional[int]  ,		snake_case_       :				Union[str, Any]=0					):
							if str(lowerCamelCase_					).startswith("""mps"""					):
											UpperCamelCase_:							Optional[Any]		=  torch.manual_seed(lowerCamelCase_					)
							else:
											UpperCamelCase_:							Tuple		=  torch.Generator(device=lowerCamelCase_					).manual_seed(lowerCamelCase_					)
							UpperCamelCase_:							Union[str, Any]		=  {
							    """input_tokens""": [
							        [1134, 90, 1135, 1133, 1080, 112, 1132, 1080, 1133, 1079, 133, 1132, 1079, 1133, 1] + [0] * 2033
							    ],
							    """generator""": generator,
							    """num_inference_steps""": 4,
							    """output_type""": """mel""",
							}
							return inputs

			def 					lowerCAmelCase__		(				self       :				Tuple					):
							UpperCamelCase_:							List[Any]		=  """cpu"""  # ensure determinism for the device-dependent torch.Generator
							UpperCamelCase_:							int		=  self.get_dummy_components()
							UpperCamelCase_:							Optional[int]		=  SpectrogramDiffusionPipeline(**lowerCamelCase_					)
							UpperCamelCase_:							Tuple		=  pipe.to(lowerCamelCase_					)
							pipe.set_progress_bar_config(disable=lowerCamelCase_					)

							UpperCamelCase_:							Optional[Any]		=  self.get_dummy_inputs(lowerCamelCase_					)
							UpperCamelCase_:							List[Any]		=  pipe(**lowerCamelCase_					)
							UpperCamelCase_:							Dict		=  output.audios

							UpperCamelCase_:							Any		=  mel[0, -3:, -3:]

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

			@skip_mps
			def 					lowerCAmelCase__		(				self       :				Union[str, Any]					):
							return super().test_dict_tuple_outputs_equivalent()

			@skip_mps
			def 					lowerCAmelCase__		(				self       :				Optional[Any]					):
							return super().test_save_load_optional_components()

			@skip_mps
			def 					lowerCAmelCase__		(				self       :				List[str]					):
							return super().test_attention_slicing_forward_pass()

			def 					lowerCAmelCase__		(				self       :				Any					):
							pass

			def 					lowerCAmelCase__		(				self       :				Optional[Any]					):
							pass




			@skip_mps
			def 					lowerCAmelCase__		(				self       :				Optional[int]					):
							return super().test_progress_bar()




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






			'''simple docstring'''

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

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

							UpperCamelCase_:							Optional[Any]		=  SpectrogramDiffusionPipeline.from_pretrained("""google/music-spectrogram-diffusion"""					)
							UpperCamelCase_:							Optional[Any]		=  pipe.melgan
							UpperCamelCase_:							Any		=  None

							UpperCamelCase_:							Union[str, Any]		=  pipe.to(lowerCamelCase_					)
							pipe.set_progress_bar_config(disable=lowerCamelCase_					)

							def callback(snake_case_       :				Any  ,		snake_case_       :				str					):
											# decode mel to audio
											UpperCamelCase_:							Dict		=  melgan(input_features=mel_output.astype(np.floataa					)					)[0]
											assert len(audio[0]					) == 8_1920 * (step + 1)
											# simulate that audio is played
											return audio

							UpperCamelCase_:							Any		=  MidiProcessor()
							UpperCamelCase_:							Any		=  processor(lowerCamelCase_					)

							UpperCamelCase_:							Optional[int]		=  input_tokens[:3]
							UpperCamelCase_:							Union[str, Any]		=  torch.manual_seed(0					)
							pipe(lowerCamelCase_  ,		num_inference_steps=5  ,		generator=lowerCamelCase_  ,		callback=lowerCamelCase_  ,		output_type="""mel"""					)

			def 					lowerCAmelCase__		(				self       :				List[Any]					):
							UpperCamelCase_:							str		=  torch_device

							UpperCamelCase_:							Dict		=  SpectrogramDiffusionPipeline.from_pretrained("""google/music-spectrogram-diffusion"""					)
							UpperCamelCase_:							int		=  pipe.to(lowerCamelCase_					)
							pipe.set_progress_bar_config(disable=lowerCamelCase_					)
							UpperCamelCase_:							Union[str, Any]		=  MidiProcessor()

							UpperCamelCase_:							Optional[Any]		=  processor(lowerCamelCase_					)
							# just run two denoising loops
							UpperCamelCase_:							Any		=  input_tokens[:2]

							UpperCamelCase_:							Dict		=  torch.manual_seed(0					)
							UpperCamelCase_:							Optional[Any]		=  pipe(lowerCamelCase_  ,		num_inference_steps=2  ,		generator=lowerCamelCase_					)

							UpperCamelCase_:							Optional[int]		=  output.audios[0]

							assert abs(np.abs(lowerCamelCase_					).sum() - 3612.841					) < 1e-1




			def 					lowerCAmelCase__		(				self       :				Dict					):
							UpperCamelCase_:							Optional[int]		=  torch_device

							UpperCamelCase_:							Any		=  SpectrogramDiffusionPipeline.from_pretrained("""google/music-spectrogram-diffusion"""					)
							UpperCamelCase_:							Tuple		=  pipe.to(lowerCamelCase_					)
							pipe.set_progress_bar_config(disable=lowerCamelCase_					)

							UpperCamelCase_:							Dict		=  MidiProcessor()

							UpperCamelCase_:							str		=  processor(lowerCamelCase_					)

							# just run 4 denoising loops
							UpperCamelCase_:							List[Any]		=  input_tokens[:4]

							UpperCamelCase_:							Union[str, Any]		=  torch.manual_seed(0					)
							UpperCamelCase_:							Union[str, Any]		=  pipe(lowerCamelCase_  ,		num_inference_steps=100  ,		generator=lowerCamelCase_					)

							UpperCamelCase_:							Optional[Any]		=  output.audios[0]
							assert abs(np.abs(lowerCamelCase_					).sum() - 9389.1111					) < 5e-2










'''simple docstring'''






from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_sentencepiece_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


__UpperCAmelCase    				=							{
    '''configuration_switch_transformers''': [
        '''SWITCH_TRANSFORMERS_PRETRAINED_CONFIG_ARCHIVE_MAP''',
        '''SwitchTransformersConfig''',
        '''SwitchTransformersOnnxConfig''',
    ]
}

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		__UpperCAmelCase    				=							[
		    '''SWITCH_TRANSFORMERS_PRETRAINED_MODEL_ARCHIVE_LIST''',
		    '''SwitchTransformersEncoderModel''',
		    '''SwitchTransformersForConditionalGeneration''',
		    '''SwitchTransformersModel''',
		    '''SwitchTransformersPreTrainedModel''',
		    '''SwitchTransformersTop1Router''',
		    '''SwitchTransformersSparseMLP''',
		]


if TYPE_CHECKING:
		from .configuration_switch_transformers import (
		    SWITCH_TRANSFORMERS_PRETRAINED_CONFIG_ARCHIVE_MAP,
		    SwitchTransformersConfig,
		    SwitchTransformersOnnxConfig,
		)

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_switch_transformers import (
				    SWITCH_TRANSFORMERS_PRETRAINED_MODEL_ARCHIVE_LIST,
				    SwitchTransformersEncoderModel,
				    SwitchTransformersForConditionalGeneration,
				    SwitchTransformersModel,
				    SwitchTransformersPreTrainedModel,
				    SwitchTransformersSparseMLP,
				    SwitchTransformersTopaRouter,
				)


else:
		import sys

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