

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							:     Optional[int] 						=	'''./tests/fixtures/elise_format0.mid'''



@unittest.skip('''The note-seq package currently throws an error on import'''    )
class    A__(snake_case__,       unittest.TestCase    ):
						"""simple docstring"""






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





						def 	UpperCamelCase__			(  self  )      ->				str:
									torch.manual_seed(0  )
									a_		:		Dict				    =				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"""       ,     )

									a_		:		Optional[int]				    =				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"""       ,     )

									a_		:		Optional[Any]				    =				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       ,     )

									a_		:		Dict				    =				DDPMScheduler()

									a_		:		List[str]				    =				{
									    'notes_encoder': notes_encoder.eval(),
									    'continuous_encoder': continuous_encoder.eval(),
									    'decoder': decoder.eval(),
									    'scheduler': scheduler,
									    'melgan': None,
									}
									return components





						def 	UpperCamelCase__			(  self       ,     _lowercase       ,     _lowercase=0  )      ->				List[str]:
									if str(_A  ).startswith("""mps"""  ):
												a_		:		Optional[int]				    =				torch.manual_seed(_A  )
									else:
												a_		:		Optional[int]				    =				torch.Generator(device=_A  ).manual_seed(_A  )
									a_		:		Dict				    =				{
									    '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 	UpperCamelCase__			(  self  )      ->				List[str]:
									a_		:		List[Any]				    =				'cpu'  # ensure determinism for the device-dependent torch.Generator
									a_		:		Dict				    =				self.get_dummy_components()
									a_		:		List[Any]				    =				SpectrogramDiffusionPipeline(**_A  )
									a_		:		List[Any]				    =				pipe.to(_A  )
									pipe.set_progress_bar_config(disable=_A  )

									a_		:		List[str]				    =				self.get_dummy_inputs(_A  )
									a_		:		Any				    =				pipe(**_A  )
									a_		:		Any				    =				output.audios

									a_		:		Dict				    =				mel[0, -3:, -3:]

									assert mel_slice.shape == (3, 3)
									a_		:		str				    =				np.array(
									    [-11.512_925, -4.7_8_8_2_1_5, -0.4_6_1_7_2_9_0_5, -2.0_5_1_7_1_5, -10.539_147, -10.970_963, -9.0_9_1_6_3_4, 4.0, 4.0]  )
									assert np.abs(mel_slice.flatten() - expected_slice  ).max() < 1e-2





						@skip_mps
						def 	UpperCamelCase__			(  self  )      ->				int:
									return super().test_save_load_local()





						@skip_mps
						def 	UpperCamelCase__			(  self  )      ->				List[str]:
									return super().test_dict_tuple_outputs_equivalent()





						@skip_mps
						def 	UpperCamelCase__			(  self  )      ->				str:
									return super().test_save_load_optional_components()





						@skip_mps
						def 	UpperCamelCase__			(  self  )      ->				int:
									return super().test_attention_slicing_forward_pass()





						def 	UpperCamelCase__			(  self  )      ->				Optional[int]:
									pass





						def 	UpperCamelCase__			(  self  )      ->				Optional[Any]:
									pass





						@skip_mps
						def 	UpperCamelCase__			(  self  )      ->				List[str]:
									return super().test_progress_bar()







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






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





						def 	UpperCamelCase__			(  self  )      ->				Any:
									# TODO - test that pipeline can decode tokens in a callback
									# so that music can be played live
									a_		:		List[Any]				    =				torch_device

									a_		:		Union[str, Any]				    =				SpectrogramDiffusionPipeline.from_pretrained("""google/music-spectrogram-diffusion"""  )
									a_		:		str				    =				pipe.melgan
									a_		:		Optional[int]				    =				None

									a_		:		Optional[int]				    =				pipe.to(_A  )
									pipe.set_progress_bar_config(disable=_A  )

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

									a_		:		List[str]				    =				MidiProcessor()
									a_		:		Any				    =				processor(_A  )

									a_		:		List[Any]				    =				input_tokens[:3]
									a_		:		Optional[int]				    =				torch.manual_seed(0  )
									pipe(_A       ,     num_inference_steps=5       ,     generator=_A       ,     callback=_A       ,     output_type="""mel"""  )





						def 	UpperCamelCase__			(  self  )      ->				int:
									a_		:		str				    =				torch_device

									a_		:		Union[str, Any]				    =				SpectrogramDiffusionPipeline.from_pretrained("""google/music-spectrogram-diffusion"""  )
									a_		:		Optional[int]				    =				pipe.to(_A  )
									pipe.set_progress_bar_config(disable=_A  )
									a_		:		Tuple				    =				MidiProcessor()

									a_		:		int				    =				processor(_A  )
									# just run two denoising loops
									a_		:		Union[str, Any]				    =				input_tokens[:2]

									a_		:		Dict				    =				torch.manual_seed(0  )
									a_		:		Any				    =				pipe(_A       ,     num_inference_steps=2       ,     generator=_A  )

									a_		:		Any				    =				output.audios[0]

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





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

									a_		:		Tuple				    =				SpectrogramDiffusionPipeline.from_pretrained("""google/music-spectrogram-diffusion"""  )
									a_		:		Union[str, Any]				    =				pipe.to(_A  )
									pipe.set_progress_bar_config(disable=_A  )

									a_		:		Union[str, Any]				    =				MidiProcessor()

									a_		:		List[str]				    =				processor(_A  )

									# just run 4 denoising loops
									a_		:		Optional[int]				    =				input_tokens[:4]

									a_		:		List[Any]				    =				torch.manual_seed(0  )
									a_		:		List[Any]				    =				pipe(_A       ,     num_inference_steps=100       ,     generator=_A  )

									a_		:		Union[str, Any]				    =				output.audios[0]
									assert abs(np.abs(_A  ).sum() - 9_389.1_111  ) < 5e-2





from ...configuration_utils import PretrainedConfig
from ...utils import logging


__snake_case							:     Optional[Any] 						=	logging.get_logger(__name__)

__snake_case							:     str 						=	{
    """ZinengTang/tvlt-base""": """https://huggingface.co/ZinengTang/tvlt-base/blob/main/config.json""",
}



class    A__(a_    ):
						"""simple docstring"""






						_A  :							Any   			=	'''tvlt'''





						def __init__(  self       ,     _lowercase=224       ,     _lowercase=2_048       ,     _lowercase=128       ,     _lowercase=[16, 16]       ,     _lowercase=[16, 16]       ,     _lowercase=3       ,     _lowercase=1       ,     _lowercase=8       ,     _lowercase=768       ,     _lowercase=12       ,     _lowercase=12       ,     _lowercase=3_072       ,     _lowercase="gelu"       ,     _lowercase=0.0       ,     _lowercase=0.0       ,     _lowercase=0.0_2       ,     _lowercase=1e-6       ,     _lowercase=True       ,     _lowercase=False       ,     _lowercase=16       ,     _lowercase=512       ,     _lowercase=8       ,     _lowercase=2_048       ,     _lowercase=0.7_5       ,     _lowercase=0.1_5       ,     _lowercase="frame-level"       ,     _lowercase=True       ,     _lowercase=True       ,     _lowercase="classification"       ,     **_lowercase       ,     )      ->				Tuple:
									super().__init__(**_lowercase  )

									if audio_mask_type not in ("frame-level", "patch_level"):
												raise ValueError(
												    """audio_mask_type must be one of two acceptable strategies - {'frame_level', 'patch-level') """
												    F'''got {audio_mask_type}'''  )

									a_		:		int				    =				image_size
									a_		:		str				    =				spectrogram_length
									a_		:		int				    =				frequency_length
									a_		:		str				    =				image_patch_size
									a_		:		int				    =				audio_patch_size
									a_		:		Any				    =				num_image_channels
									a_		:		str				    =				num_audio_channels
									a_		:		int				    =				num_frames

									a_		:		Optional[int]				    =				hidden_size
									a_		:		Tuple				    =				num_hidden_layers
									a_		:		Any				    =				num_attention_heads
									a_		:		Union[str, Any]				    =				intermediate_size
									a_		:		Tuple				    =				hidden_act
									a_		:		Union[str, Any]				    =				hidden_dropout_prob
									a_		:		int				    =				attention_probs_dropout_prob
									a_		:		Any				    =				initializer_range
									a_		:		int				    =				layer_norm_eps
									a_		:		Union[str, Any]				    =				qkv_bias
									a_		:		str				    =				use_mean_pooling

									a_		:		List[Any]				    =				decoder_num_attention_heads
									a_		:		Optional[Any]				    =				decoder_hidden_size
									a_		:		Dict				    =				decoder_num_hidden_layers
									a_		:		Optional[Any]				    =				decoder_intermediate_size
									a_		:		Dict				    =				pixel_mask_ratio
									a_		:		int				    =				audio_mask_ratio
									a_		:		Optional[Any]				    =				audio_mask_type

									a_		:		Union[str, Any]				    =				task_matching
									a_		:		Dict				    =				task_mae
									a_		:		Union[str, Any]				    =				loss_type


