






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


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




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






						def 		lowerCAmelCase (       self)      ->	str:




											"""simple docstring"""





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

											lowerCamelCase__					=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'     ,    )

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

											lowerCamelCase__					=DDPMScheduler()

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






						def 		lowerCAmelCase (       self     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE=0)      ->	Optional[Any]:




											"""simple docstring"""





											if str(_SCREAMING_SNAKE_CASE).startswith('mps'):
																lowerCamelCase__					=torch.manual_seed(_SCREAMING_SNAKE_CASE)
											else:
																lowerCamelCase__					=torch.Generator(device=_SCREAMING_SNAKE_CASE).manual_seed(_SCREAMING_SNAKE_CASE)
											lowerCamelCase__					={
											    '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 		lowerCAmelCase (       self)      ->	List[Any]:




											"""simple docstring"""





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

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

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

											assert mel_slice.shape == (3, 3)
											lowerCamelCase__					=np.array(
											    [-11.512925, -4.788215, -0.46172905, -2.051715, -10.539147, -10.970963, -9.091634, 4.0, 4.0])
											assert np.abs(mel_slice.flatten() - expected_slice).max() < 1E-2






						@skip_mps
						def 		lowerCAmelCase (       self)      ->	Any:




											"""simple docstring"""





											return super().test_save_load_local()






						@skip_mps
						def 		lowerCAmelCase (       self)      ->	List[str]:




											"""simple docstring"""





											return super().test_dict_tuple_outputs_equivalent()






						@skip_mps
						def 		lowerCAmelCase (       self)      ->	int:




											"""simple docstring"""





											return super().test_save_load_optional_components()






						@skip_mps
						def 		lowerCAmelCase (       self)      ->	int:




											"""simple docstring"""





											return super().test_attention_slicing_forward_pass()






						def 		lowerCAmelCase (       self)      ->	str:




											"""simple docstring"""





											pass






						def 		lowerCAmelCase (       self)      ->	List[str]:




											"""simple docstring"""





											pass






						@skip_mps
						def 		lowerCAmelCase (       self)      ->	Any:




											"""simple docstring"""





											return super().test_progress_bar()




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






						def 		lowerCAmelCase (       self)      ->	Optional[Any]:




											"""simple docstring"""





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






						def 		lowerCAmelCase (       self)      ->	List[str]:




											"""simple docstring"""





											lowerCamelCase__					=torch_device

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

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

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

											lowerCamelCase__					=MidiProcessor()
											lowerCamelCase__					=processor(_SCREAMING_SNAKE_CASE)

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






						def 		lowerCAmelCase (       self)      ->	Optional[int]:




											"""simple docstring"""





											lowerCamelCase__					=torch_device

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

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

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

											lowerCamelCase__					=output.audios[0]

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






						def 		lowerCAmelCase (       self)      ->	Dict:




											"""simple docstring"""





											lowerCamelCase__					=torch_device

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

											lowerCamelCase__					=MidiProcessor()

											lowerCamelCase__					=processor(_SCREAMING_SNAKE_CASE)

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

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

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







"""simple docstring"""



from typing import Any, Dict, List, Mapping

import numpy as np
import torch

from ...utils import is_cython_available, requires_backends


if is_cython_available():
						import pyximport

						pyximport.install(setup_args={"include_dirs": np.get_include()})
						from . import algos_graphormer  # noqa E402




def     lowerCamelCase_					(	__A	:		List[str]      ,						__A	:		int = 512      ):
					lowerCamelCase__					=x.shape[1] if len(x.shape      ) > 1 else 1
					lowerCamelCase__					=1 + np.arange(0      ,						feature_num * offset      ,						__A      ,						dtype=np.intaa      )
					lowerCamelCase__					=x + feature_offset
					return x




def     lowerCamelCase_					(	__A	:		str      ,						__A	:		Optional[Any]=True      ):
					requires_backends(__A      ,						['cython']      )

					if keep_features and "edge_attr" in item.keys():  # edge_attr
										lowerCamelCase__					=np.asarray(item['edge_attr']      ,						dtype=np.intaa      )
					else:
										lowerCamelCase__					=np.ones((len(item['edge_index'][0]      ), 1)      ,						dtype=np.intaa      )  # same embedding for all

					if keep_features and "node_feat" in item.keys():  # input_nodes
										lowerCamelCase__					=np.asarray(item['node_feat']      ,						dtype=np.intaa      )
					else:
										lowerCamelCase__					=np.ones((item['num_nodes'], 1)      ,						dtype=np.intaa      )  # same embedding for all

					lowerCamelCase__					=np.asarray(item['edge_index']      ,						dtype=np.intaa      )

					lowerCamelCase__					=convert_to_single_emb(__A      ) + 1
					lowerCamelCase__					=item['num_nodes']

					if len(edge_attr.shape      ) == 1:
										lowerCamelCase__					=edge_attr[:, None]
					lowerCamelCase__					=np.zeros([num_nodes, num_nodes, edge_attr.shape[-1]]      ,						dtype=np.intaa      )
					lowerCamelCase__					=convert_to_single_emb(__A      ) + 1

					# node adj matrix [num_nodes, num_nodes] bool
					lowerCamelCase__					=np.zeros([num_nodes, num_nodes]      ,						dtype=__A      )
					lowerCamelCase__					=True

					lowerCamelCase__					,  lowerCamelCase__					=algos_graphormer.floyd_warshall(__A      )
					lowerCamelCase__					=np.amax(__A      )

					lowerCamelCase__					=algos_graphormer.gen_edge_input(__A      ,						__A      ,						__A      )
					lowerCamelCase__					=np.zeros([num_nodes + 1, num_nodes + 1]      ,						dtype=np.single      )  # with graph token

					# combine
					lowerCamelCase__					=input_nodes + 1  # we shift all indices by one for padding
					lowerCamelCase__					=attn_bias
					lowerCamelCase__					=attn_edge_type
					lowerCamelCase__					=shortest_path_result.astype(np.intaa      ) + 1  # we shift all indices by one for padding
					lowerCamelCase__					=np.sum(__A      ,						axis=1      ).reshape(-1      ) + 1  # we shift all indices by one for padding
					lowerCamelCase__					=item['in_degree']  # for undirected graph
					lowerCamelCase__					=input_edges + 1  # we shift all indices by one for padding
					if "labels" not in item:
										lowerCamelCase__					=item['y']

					return item




class 							__lowercase:






						def __init__(       self     ,    _SCREAMING_SNAKE_CASE=20     ,    _SCREAMING_SNAKE_CASE=False)      ->	Any:




											"""simple docstring"""





											if not is_cython_available():
																raise ImportError('Graphormer preprocessing needs Cython (pyximport)')

											lowerCamelCase__					=spatial_pos_max
											lowerCamelCase__					=on_the_fly_processing






						def __call__(       self     ,    _SCREAMING_SNAKE_CASE)      ->	Dict[str, Any]:




											"""simple docstring"""





											if self.on_the_fly_processing:
																lowerCamelCase__					=[preprocess_item(_SCREAMING_SNAKE_CASE) for i in features]

											if not isinstance(features[0]     ,    _SCREAMING_SNAKE_CASE):
																lowerCamelCase__					=[vars(_SCREAMING_SNAKE_CASE) for f in features]
											lowerCamelCase__					={}

											lowerCamelCase__					=max(len(i['input_nodes']) for i in features)
											lowerCamelCase__					=len(features[0]['input_nodes'][0])
											lowerCamelCase__					=len(features[0]['attn_edge_type'][0][0])
											lowerCamelCase__					=max(len(i['input_edges'][0][0]) for i in features)
											lowerCamelCase__					=len(features[0]['input_edges'][0][0][0])
											lowerCamelCase__					=len(_SCREAMING_SNAKE_CASE)

											lowerCamelCase__					=torch.zeros(_SCREAMING_SNAKE_CASE     ,    max_node_num + 1     ,    max_node_num + 1     ,    dtype=torch.float)
											lowerCamelCase__					=torch.zeros(_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    dtype=torch.long)
											lowerCamelCase__					=torch.zeros(_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    dtype=torch.long)
											lowerCamelCase__					=torch.zeros(_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    dtype=torch.long)
											lowerCamelCase__					=torch.zeros(_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    dtype=torch.long)
											lowerCamelCase__					=torch.zeros(
											    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    dtype=torch.long)

											for ix, f in enumerate(_SCREAMING_SNAKE_CASE):
																for k in ["attn_bias", "attn_edge_type", "spatial_pos", "in_degree", "input_nodes", "input_edges"]:
																					lowerCamelCase__					=torch.tensor(f[k])

																if len(f['attn_bias'][1:, 1:][f['spatial_pos'] >= self.spatial_pos_max]) > 0:
																					lowerCamelCase__					=float('-inf')

																lowerCamelCase__					=f['attn_bias']
																lowerCamelCase__					=f[
																    'attn_edge_type'
																]
																lowerCamelCase__					=f['spatial_pos']
																lowerCamelCase__					=f['in_degree']
																lowerCamelCase__					=f['input_nodes']
																lowerCamelCase__					=f['input_edges']

											lowerCamelCase__					=batch['in_degree']

											lowerCamelCase__					=features[0]['labels']
											if len(_SCREAMING_SNAKE_CASE) == 1:  # one task
																if isinstance(sample[0]     ,    _SCREAMING_SNAKE_CASE):  # regression
																					lowerCamelCase__					=torch.from_numpy(np.concatenate([i['labels'] for i in features]))
																else:  # binary classification
																					lowerCamelCase__					=torch.from_numpy(np.concatenate([i['labels'] for i in features]))
											else:  # multi task classification, left to float to keep the NaNs
																lowerCamelCase__					=torch.from_numpy(np.stack([i['labels'] for i in features]     ,    axis=0))

											return batch