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


__lowercase							:				List[str]					=       '''./tests/fixtures/elise_format0.mid'''

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

			def     UpperCAmelCase__	(self ):
								torch.manual_seed(0 )
								lowerCamelCase_   : Optional[Any]       =						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'''     ,				)

								lowerCamelCase_   : Tuple       =						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'''     ,				)

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

								lowerCamelCase_   : List[Any]       =						DDPMScheduler()

								lowerCamelCase_   : Optional[Any]       =						{
								    '''notes_encoder''': notes_encoder.eval(),
								    '''continuous_encoder''': continuous_encoder.eval(),
								    '''decoder''': decoder.eval(),
								    '''scheduler''': scheduler,
								    '''melgan''': None,
								}
								return components

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

								lowerCamelCase_   : Optional[int]       =						self.get_dummy_inputs(A )
								lowerCamelCase_   : int       =						pipe(**A )
								lowerCamelCase_   : Optional[int]       =						output.audios

								lowerCamelCase_   : str       =						mel[0, -3:, -3:]

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

			@skip_mps
			def     UpperCAmelCase__	(self ):
								return super().test_dict_tuple_outputs_equivalent()

			@skip_mps
			def     UpperCAmelCase__	(self ):
								return super().test_save_load_optional_components()

			@skip_mps
			def     UpperCAmelCase__	(self ):
								return super().test_attention_slicing_forward_pass()

			def     UpperCAmelCase__	(self ):
								pass

			def     UpperCAmelCase__	(self ):
								pass






			@skip_mps
			def     UpperCAmelCase__	(self ):
								return super().test_progress_bar()

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

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

			def     UpperCAmelCase__	(self ):
								# TODO - test that pipeline can decode tokens in a callback
								# so that music can be played live
								lowerCamelCase_   : int       =						torch_device

								lowerCamelCase_   : Any       =						SpectrogramDiffusionPipeline.from_pretrained('''google/music-spectrogram-diffusion''' )
								lowerCamelCase_   : Dict       =						pipe.melgan
								lowerCamelCase_   : Any       =						None

								lowerCamelCase_   : List[Any]       =						pipe.to(A )
								pipe.set_progress_bar_config(disable=A )

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

								lowerCamelCase_   : List[str]       =						MidiProcessor()
								lowerCamelCase_   : List[Any]       =						processor(A )

								lowerCamelCase_   : Tuple       =						input_tokens[:3]
								lowerCamelCase_   : int       =						torch.manual_seed(0 )
								pipe(A     ,				num_inference_steps=5     ,				generator=A     ,				callback=A     ,				output_type='''mel''' )

			def     UpperCAmelCase__	(self ):
								lowerCamelCase_   : List[str]       =						torch_device

								lowerCamelCase_   : List[str]       =						SpectrogramDiffusionPipeline.from_pretrained('''google/music-spectrogram-diffusion''' )
								lowerCamelCase_   : List[str]       =						pipe.to(A )
								pipe.set_progress_bar_config(disable=A )
								lowerCamelCase_   : List[str]       =						MidiProcessor()

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

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

								lowerCamelCase_   : int       =						output.audios[0]

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






			def     UpperCAmelCase__	(self ):
								lowerCamelCase_   : int       =						torch_device

								lowerCamelCase_   : Optional[Any]       =						SpectrogramDiffusionPipeline.from_pretrained('''google/music-spectrogram-diffusion''' )
								lowerCamelCase_   : List[Any]       =						pipe.to(A )
								pipe.set_progress_bar_config(disable=A )

								lowerCamelCase_   : List[Any]       =						MidiProcessor()

								lowerCamelCase_   : Tuple       =						processor(A )

								# just run 4 denoising loops
								lowerCamelCase_   : Optional[Any]       =						input_tokens[:4]

								lowerCamelCase_   : Dict       =						torch.manual_seed(0 )
								lowerCamelCase_   : Optional[int]       =						pipe(A     ,				num_inference_steps=1_0_0     ,				generator=A )

								lowerCamelCase_   : List[str]       =						output.audios[0]
								assert abs(np.abs(A ).sum() - 93_89.11_11 ) < 5E-2




'''simple docstring'''


import argparse
import logging
import math
import time

import torch

from transformers import TransfoXLCorpus, TransfoXLLMHeadModel


logging.basicConfig(
    format='''%(asctime)s - %(levelname)s - %(name)s - %(message)s''', datefmt='''%m/%d/%Y %H:%M:%S''', level=logging.INFO
)
__lowercase							:				Optional[Any]					=       logging.getLogger(__name__)




def lowercase_		(				)       ->							Union[str, Any]:
					'''simple docstring'''

					lowerCamelCase_   : Optional[int]       =						argparse.ArgumentParser(description='''PyTorch Transformer Language Model'''			)
					parser.add_argument('''--model_name'''     ,    type=_lowercase     ,    default='''transfo-xl-wt103'''     ,    help='''pretrained model name'''			)
					parser.add_argument(
					    '''--split'''     ,    type=_lowercase     ,    default='''test'''     ,    choices=['''all''', '''valid''', '''test''']     ,    help='''which split to evaluate'''			)
					parser.add_argument('''--batch_size'''     ,    type=_lowercase     ,    default=10     ,    help='''batch size'''			)
					parser.add_argument('''--tgt_len'''     ,    type=_lowercase     ,    default=128     ,    help='''number of tokens to predict'''			)
					parser.add_argument('''--ext_len'''     ,    type=_lowercase     ,    default=0     ,    help='''length of the extended context'''			)
					parser.add_argument('''--mem_len'''     ,    type=_lowercase     ,    default=1_600     ,    help='''length of the retained previous heads'''			)
					parser.add_argument('''--clamp_len'''     ,    type=_lowercase     ,    default=1_000     ,    help='''max positional embedding index'''			)
					parser.add_argument('''--no_cuda'''     ,    action='''store_true'''     ,    help='''Do not use CUDA even though CUA is available'''			)
					parser.add_argument('''--work_dir'''     ,    type=_lowercase     ,    required=_lowercase     ,    help='''path to the work_dir'''			)
					parser.add_argument('''--no_log'''     ,    action='''store_true'''     ,    help='''do not log the eval result'''			)
					parser.add_argument('''--same_length'''     ,    action='''store_true'''     ,    help='''set same length attention with masking'''			)
					parser.add_argument('''--server_ip'''     ,    type=_lowercase     ,    default=''''''     ,    help='''Can be used for distant debugging.'''			)
					parser.add_argument('''--server_port'''     ,    type=_lowercase     ,    default=''''''     ,    help='''Can be used for distant debugging.'''			)
					lowerCamelCase_   : Dict       =						parser.parse_args()
					assert args.ext_len >= 0, "extended context length must be non-negative"

					if args.server_ip and args.server_port:
										# Distant debugging - see https://code.visualstudio.com/docs/python/debugging#_attach-to-a-local-script
										import ptvsd

										print('''Waiting for debugger attach'''			)
										ptvsd.enable_attach(address=(args.server_ip, args.server_port)     ,    redirect_output=_lowercase			)
										ptvsd.wait_for_attach()

					lowerCamelCase_   : List[str]       =						torch.device('''cuda''' if torch.cuda.is_available() and not args.no_cuda else '''cpu'''			)
					logger.info('''device: {}'''.format(_lowercase			)			)

					# Load a pre-processed dataset
					# You can also build the corpus yourself using TransfoXLCorpus methods
					# The pre-processing involve computing word frequencies to prepare the Adaptive input and SoftMax
					# and tokenizing the dataset
					# The pre-processed corpus is a convertion (using the conversion script )
					lowerCamelCase_   : str       =						TransfoXLCorpus.from_pretrained(args.model_name			)

					lowerCamelCase_   : Optional[int]       =						corpus.get_iterator('''valid'''     ,    args.batch_size     ,    args.tgt_len     ,    device=_lowercase     ,    ext_len=args.ext_len			)
					lowerCamelCase_   : str       =						corpus.get_iterator('''test'''     ,    args.batch_size     ,    args.tgt_len     ,    device=_lowercase     ,    ext_len=args.ext_len			)

					# Load a pre-trained model
					lowerCamelCase_   : str       =						TransfoXLLMHeadModel.from_pretrained(args.model_name			)
					model.to(_lowercase			)

					logger.info(
					    '''Evaluating with bsz {} tgt_len {} ext_len {} mem_len {} clamp_len {}'''.format(
					        args.batch_size     ,    args.tgt_len     ,    args.ext_len     ,    args.mem_len     ,    args.clamp_len			)			)

					model.reset_memory_length(args.mem_len			)
					if args.clamp_len > 0:
										lowerCamelCase_   : Union[str, Any]       =						args.clamp_len
					if args.same_length:
										lowerCamelCase_   : int       =						True

					###############################################################################
					# Evaluation code
					###############################################################################
					def evaluate(_lowercase			):
										# Turn on evaluation mode which disables dropout.
										model.eval()
										lowerCamelCase_, lowerCamelCase_   : Union[str, Any]       =						0, 0.0
										lowerCamelCase_   : Optional[int]       =						time.time()
										with torch.no_grad():
															lowerCamelCase_   : Dict       =						None
															for idx, (data, target, seq_len) in enumerate(_lowercase			):
																				lowerCamelCase_   : str       =						model(_lowercase     ,    lm_labels=_lowercase     ,    mems=_lowercase			)
																				lowerCamelCase_, lowerCamelCase_, lowerCamelCase_   : int       =						ret
																				lowerCamelCase_   : List[str]       =						loss.mean()
																				total_loss += seq_len * loss.item()
																				total_len += seq_len
															lowerCamelCase_   : List[str]       =						time.time() - start_time
										logger.info('''Time : {:.2f}s, {:.2f}ms/segment'''.format(_lowercase     ,    1_000 * total_time / (idx + 1)			)			)
										return total_loss / total_len

					# Run on test data.
					if args.split == "all":
										lowerCamelCase_   : Tuple       =						evaluate(_lowercase			)
										lowerCamelCase_   : List[Any]       =						evaluate(_lowercase			)
					elif args.split == "valid":
										lowerCamelCase_   : Union[str, Any]       =						evaluate(_lowercase			)
										lowerCamelCase_   : str       =						None
					elif args.split == "test":
										lowerCamelCase_   : Any       =						evaluate(_lowercase			)
										lowerCamelCase_   : int       =						None

					def format_log(_lowercase     ,    _lowercase			):
										lowerCamelCase_   : Optional[Any]       =						'''| {0} loss {1:5.2f} | {0} ppl {2:9.3f} '''.format(_lowercase     ,    _lowercase     ,    math.exp(_lowercase			)			)
										return log_str

					lowerCamelCase_   : List[Any]       =						''''''
					if valid_loss is not None:
										log_str += format_log(_lowercase     ,    '''valid'''			)
					if test_loss is not None:
										log_str += format_log(_lowercase     ,    '''test'''			)

					logger.info('''=''' * 100			)
					logger.info(_lowercase			)
					logger.info('''=''' * 100			)


if __name__ == "__main__":
					main()


