




'''simple docstring'''

from collections import OrderedDict
from typing import Any, Mapping, Optional

from ... import PreTrainedTokenizer
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig, OnnxSeqaSeqConfigWithPast
from ...onnx.utils import compute_effective_axis_dimension
from ...utils import TensorType, is_torch_available, logging


a__         =     logging.get_logger(__name__)

a__         =     {
    '''facebook/m2m100_418M''': '''https://huggingface.co/facebook/m2m100_418M/resolve/main/config.json''',
    # See all M2M100 models at https://huggingface.co/models?filter=m2m_100
}





class __magic_name__(						__lowerCAmelCase      ):
						UpperCAmelCase_		:					Optional[Any]    =  """m2m_100"""
						UpperCAmelCase_		:					Any    =  ["""past_key_values"""]
						UpperCAmelCase_		:					Optional[Any]    =  {"""num_attention_heads""": """encoder_attention_heads""", """hidden_size""": """d_model"""}






						def __init__(   self						:						int					,   __UpperCamelCase						:						List[str]=1_2_8_1_1_2					,   __UpperCamelCase						:						List[str]=1_0_2_4					,   __UpperCamelCase						:						List[str]=1_2					,   __UpperCamelCase						:						int=4_0_9_6					,   __UpperCamelCase						:						Optional[Any]=1_6					,   __UpperCamelCase						:						Optional[Any]=1_2					,   __UpperCamelCase						:						int=4_0_9_6					,   __UpperCamelCase						:						Dict=1_6					,   __UpperCamelCase						:						List[Any]=0.05					,   __UpperCamelCase						:						List[str]=0.05					,   __UpperCamelCase						:						List[str]=True					,   __UpperCamelCase						:						int=True					,   __UpperCamelCase						:						List[Any]="relu"					,   __UpperCamelCase						:						Tuple=1_0_2_4					,   __UpperCamelCase						:						Dict=0.1					,   __UpperCamelCase						:						List[str]=0.1					,   __UpperCamelCase						:						List[Any]=0.0					,   __UpperCamelCase						:						str=0.02					,   __UpperCamelCase						:						Any=2					,   __UpperCamelCase						:						Any=True					,   __UpperCamelCase						:						Dict=1					,   __UpperCamelCase						:						Tuple=0					,   __UpperCamelCase						:						List[Any]=2					,   **__UpperCamelCase						:						Union[str, Any]					,   ):



									'''simple docstring'''






									snake_case__          =						vocab_size
									snake_case__          =						max_position_embeddings
									snake_case__          =						d_model
									snake_case__          =						encoder_ffn_dim
									snake_case__          =						encoder_layers
									snake_case__          =						encoder_attention_heads
									snake_case__          =						decoder_ffn_dim
									snake_case__          =						decoder_layers
									snake_case__          =						decoder_attention_heads
									snake_case__          =						dropout
									snake_case__          =						attention_dropout
									snake_case__          =						activation_dropout
									snake_case__          =						activation_function
									snake_case__          =						init_std
									snake_case__          =						encoder_layerdrop
									snake_case__          =						decoder_layerdrop
									snake_case__          =						use_cache
									snake_case__          =						encoder_layers
									snake_case__          =						scale_embedding  # scale factor will be sqrt(d_model) if True

									super().__init__(
									    pad_token_id=__UpperCamelCase					,   bos_token_id=__UpperCamelCase					,   eos_token_id=__UpperCamelCase					,   is_encoder_decoder=__UpperCamelCase					,   decoder_start_token_id=__UpperCamelCase					,   **__UpperCamelCase					,   )







class __magic_name__(						__lowerCAmelCase      ):






						@property
						def       __lowerCAmelCase(   self						:						Any       ):



									'''simple docstring'''






									snake_case__          =						OrderedDict(
									    [
									        ("""input_ids""", {0: """batch""", 1: """encoder_sequence"""}),
									        ("""attention_mask""", {0: """batch""", 1: """encoder_sequence"""}),
									    ]       )

									if self.use_past:
												snake_case__          =						{0: """batch"""}
												snake_case__          =						{0: """batch""", 1: """past_decoder_sequence + sequence"""}
									else:
												snake_case__          =						{0: """batch""", 1: """decoder_sequence"""}
												snake_case__          =						{0: """batch""", 1: """decoder_sequence"""}

									if self.use_past:
												self.fill_with_past_key_values_(__UpperCamelCase					,   direction="""inputs"""       )
									return common_inputs






						def       __lowerCAmelCase(   self						:						int					,   __UpperCamelCase						:						PreTrainedTokenizer					,   __UpperCamelCase						:						int = -1					,   __UpperCamelCase						:						int = -1					,   __UpperCamelCase						:						bool = False					,   __UpperCamelCase						:						Optional[TensorType] = None					,   ):



									'''simple docstring'''






									snake_case__          =						compute_effective_axis_dimension(
									    __UpperCamelCase					,   fixed_dimension=OnnxConfig.default_fixed_batch					,   num_token_to_add=0       )

									# If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
									snake_case__          =						tokenizer.num_special_tokens_to_add(__UpperCamelCase       )
									snake_case__          =						compute_effective_axis_dimension(
									    __UpperCamelCase					,   fixed_dimension=OnnxConfig.default_fixed_sequence					,   num_token_to_add=__UpperCamelCase       )

									# Generate dummy inputs according to compute batch and sequence
									snake_case__          =						[""" """.join([tokenizer.unk_token]       ) * seq_length] * batch_size
									snake_case__          =						dict(tokenizer(__UpperCamelCase					,   return_tensors=__UpperCamelCase       )       )
									return common_inputs






						def       __lowerCAmelCase(   self						:						int					,   __UpperCamelCase						:						PreTrainedTokenizer					,   __UpperCamelCase						:						int = -1					,   __UpperCamelCase						:						int = -1					,   __UpperCamelCase						:						bool = False					,   __UpperCamelCase						:						Optional[TensorType] = None					,   ):



									'''simple docstring'''






									snake_case__          =						self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
									    __UpperCamelCase					,   __UpperCamelCase					,   __UpperCamelCase					,   __UpperCamelCase					,   __UpperCamelCase       )

									# Generate decoder inputs
									snake_case__          =						seq_length if not self.use_past else 1
									snake_case__          =						self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
									    __UpperCamelCase					,   __UpperCamelCase					,   __UpperCamelCase					,   __UpperCamelCase					,   __UpperCamelCase       )
									snake_case__          =						{f"""decoder_{name}""": tensor for name, tensor in decoder_inputs.items()}
									snake_case__          =						dict(**__UpperCamelCase					,   **__UpperCamelCase       )

									if self.use_past:
												if not is_torch_available():
															raise ValueError("""Cannot generate dummy past_keys inputs without PyTorch installed."""       )
												else:
															import torch
												snake_case__    ,			snake_case__          =						common_inputs["""input_ids"""].shape
												snake_case__          =						common_inputs["""decoder_input_ids"""].shape[1]
												snake_case__    ,			snake_case__          =						self.num_attention_heads
												snake_case__          =						(
												    batch,
												    num_encoder_attention_heads,
												    encoder_seq_length,
												    self._config.hidden_size // num_encoder_attention_heads,
												)
												snake_case__          =						decoder_seq_length + 3
												snake_case__          =						(
												    batch,
												    num_decoder_attention_heads,
												    decoder_past_length,
												    self._config.hidden_size // num_decoder_attention_heads,
												)

												snake_case__          =						torch.cat(
												    [common_inputs["""decoder_attention_mask"""], torch.ones(__UpperCamelCase					,   __UpperCamelCase       )]					,   dim=1       )

												snake_case__          =						[]
												# If the number of encoder and decoder layers are present in the model configuration, both are considered
												snake_case__    ,			snake_case__          =						self.num_layers
												snake_case__          =						min(__UpperCamelCase					,   __UpperCamelCase       )
												snake_case__          =						max(__UpperCamelCase					,   __UpperCamelCase       ) - min_num_layers
												snake_case__          =						"""encoder""" if num_encoder_layers > num_decoder_layers else """decoder"""

												for _ in range(__UpperCamelCase       ):
															common_inputs["past_key_values"].append(
															    (
															        torch.zeros(__UpperCamelCase       ),
															        torch.zeros(__UpperCamelCase       ),
															        torch.zeros(__UpperCamelCase       ),
															        torch.zeros(__UpperCamelCase       ),
															    )       )
												# TODO: test this.
												snake_case__          =						encoder_shape if remaining_side_name == """encoder""" else decoder_shape
												for _ in range(__UpperCamelCase					,   __UpperCamelCase       ):
															common_inputs["past_key_values"].append((torch.zeros(__UpperCamelCase       ), torch.zeros(__UpperCamelCase       ))       )
									return common_inputs







						UpperCAmelCase_		:					Optional[Any]    =  _generate_dummy_inputs_for_default_and_seqaseq_lm





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


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





@unittest.skip("""The note-seq package currently throws an error on import"""      )
class __magic_name__(						__lowerCAmelCase    ,		unittest.TestCase      ):
						UpperCAmelCase_		:					Tuple    =  SpectrogramDiffusionPipeline
						UpperCAmelCase_		:					List[str]    =  PipelineTesterMixin.required_optional_params - {
						    """callback""",
						    """latents""",
						    """callback_steps""",
						    """output_type""",
						    """num_images_per_prompt""",
						}
						UpperCAmelCase_		:					Dict    =  False
						UpperCAmelCase_		:					Dict    =  False
						UpperCAmelCase_		:					List[str]    =  TOKENS_TO_AUDIO_GENERATION_PARAMS
						UpperCAmelCase_		:					Union[str, Any]    =  TOKENS_TO_AUDIO_GENERATION_BATCH_PARAMS






						def       __lowerCAmelCase(   self						:						Tuple       ):



									'''simple docstring'''






									torch.manual_seed(0       )
									snake_case__          =						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"""					,   )

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

									snake_case__          =						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					,   )

									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       __lowerCAmelCase(   self						:						Tuple					,   __UpperCamelCase						:						Any					,   __UpperCamelCase						:						str=0       ):



									'''simple docstring'''






									if str(__UpperCamelCase       ).startswith("""mps"""       ):
												snake_case__          =						torch.manual_seed(__UpperCamelCase       )
									else:
												snake_case__          =						torch.Generator(device=__UpperCamelCase       ).manual_seed(__UpperCamelCase       )
									snake_case__          =						{
									    """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       __lowerCAmelCase(   self						:						Optional[Any]       ):



									'''simple docstring'''






									snake_case__          =						"""cpu"""  # ensure determinism for the device-dependent torch.Generator
									snake_case__          =						self.get_dummy_components()
									snake_case__          =						SpectrogramDiffusionPipeline(**__UpperCamelCase       )
									snake_case__          =						pipe.to(__UpperCamelCase       )
									pipe.set_progress_bar_config(disable=__UpperCamelCase       )

									snake_case__          =						self.get_dummy_inputs(__UpperCamelCase       )
									snake_case__          =						pipe(**__UpperCamelCase       )
									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       __lowerCAmelCase(   self						:						Tuple       ):



									'''simple docstring'''






									return super().test_save_load_local()






						@skip_mps
						def       __lowerCAmelCase(   self						:						List[Any]       ):



									'''simple docstring'''






									return super().test_dict_tuple_outputs_equivalent()






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



									'''simple docstring'''






									return super().test_save_load_optional_components()






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



									'''simple docstring'''






									return super().test_attention_slicing_forward_pass()






						def       __lowerCAmelCase(   self						:						Dict       ):



									'''simple docstring'''






									pass






						def       __lowerCAmelCase(   self						:						int       ):



									'''simple docstring'''






									pass







						@skip_mps
						def       __lowerCAmelCase(   self						:						Dict       ):



									'''simple docstring'''






									return super().test_progress_bar()







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






						def       __lowerCAmelCase(   self						:						List[Any]       ):



									'''simple docstring'''






									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()






						def       __lowerCAmelCase(   self						:						int       ):



									'''simple docstring'''






									snake_case__          =						torch_device

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

									snake_case__          =						pipe.to(__UpperCamelCase       )
									pipe.set_progress_bar_config(disable=__UpperCamelCase       )

									def callback(__UpperCamelCase						:						List[str]					,   __UpperCamelCase						:						List[str]       ):
												# decode mel to audio
												snake_case__          =						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

									snake_case__          =						MidiProcessor()
									snake_case__          =						processor(__UpperCamelCase       )

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






						def       __lowerCAmelCase(   self						:						List[Any]       ):



									'''simple docstring'''






									snake_case__          =						torch_device

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

									snake_case__          =						processor(__UpperCamelCase       )
									# just run two denoising loops
									snake_case__          =						input_tokens[:2]

									snake_case__          =						torch.manual_seed(0       )
									snake_case__          =						pipe(__UpperCamelCase					,   num_inference_steps=2					,   generator=__UpperCamelCase       )

									snake_case__          =						output.audios[0]

									assert abs(np.abs(__UpperCamelCase       ).sum() - 36_12.8_41       ) < 1E-1







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



									'''simple docstring'''






									snake_case__          =						torch_device

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

									snake_case__          =						MidiProcessor()

									snake_case__          =						processor(__UpperCamelCase       )

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

									snake_case__          =						torch.manual_seed(0       )
									snake_case__          =						pipe(__UpperCamelCase					,   num_inference_steps=1_0_0					,   generator=__UpperCamelCase       )

									snake_case__          =						output.audios[0]
									assert abs(np.abs(__UpperCamelCase       ).sum() - 93_89.11_11       ) < 5E-2