import timeit
from typing import Callable, Optional

from ..configuration_utils import PretrainedConfig
from ..models.auto.modeling_auto import MODEL_MAPPING, MODEL_WITH_LM_HEAD_MAPPING
from ..utils import is_pyanvml_available, is_torch_available, logging
from .benchmark_utils import (
    Benchmark,
    Memory,
    MemorySummary,
    measure_peak_memory_cpu,
    start_memory_tracing,
    stop_memory_tracing,
)


if is_torch_available():
    import torch

    from .benchmark_args import PyTorchBenchmarkArguments


if is_pyanvml_available():
    import pyanvml.pyanvml as nvml


lowerCAmelCase__      		=					logging.get_logger(__name__)
class 							a_					(		SCREAMING_SNAKE_CASE		):






 '''simple docstring'''






 UpperCAmelCase_      					=   42
 UpperCAmelCase_      					=   42
 UpperCAmelCase_      					=   "PyTorch"




 @property
 def 						__snake_case		(     self   :    int):







        '''simple docstring'''



        return torch.__version__




 def 						__snake_case		(     self   :    List[Any]     ,       lowercase__   :    str     ,       lowercase__   :    int     ,       lowercase__   :    int):







        '''simple docstring'''



        lowerCAmelCase__	      =   self._prepare_inference_func(lowercase__     ,       lowercase__     ,       lowercase__)
        return self._measure_speed(_inference)




 def 						__snake_case		(     self   :    int     ,       lowercase__   :    str     ,       lowercase__   :    int     ,       lowercase__   :    int):







        '''simple docstring'''



        lowerCAmelCase__	      =   self._prepare_inference_func(lowercase__     ,       lowercase__     ,       lowercase__)
        return self._measure_memory(_inference)




 def 						__snake_case		(     self   :    str     ,       lowercase__   :    str     ,       lowercase__   :    int     ,       lowercase__   :    int):







        '''simple docstring'''



        lowerCAmelCase__	      =   self._prepare_train_func(lowercase__     ,       lowercase__     ,       lowercase__)
        return self._measure_speed(_train)




 def 						__snake_case		(     self   :    Union[str, Any]     ,       lowercase__   :    str     ,       lowercase__   :    int     ,       lowercase__   :    int):







        '''simple docstring'''



        lowerCAmelCase__	      =   self._prepare_train_func(lowercase__     ,       lowercase__     ,       lowercase__)
        return self._measure_memory(_train)




 def 						__snake_case		(     self   :    Any     ,       lowercase__   :    str     ,       lowercase__   :    int     ,       lowercase__   :    int):







        '''simple docstring'''



        lowerCAmelCase__	      =   self.config_dict[model_name]

        if self.args.torchscript:
               lowerCAmelCase__	      =   True

        lowerCAmelCase__	      =   (
            hasattr(lowercase__     ,       'architectures')
            and isinstance(config.architectures     ,       lowercase__)
            and len(config.architectures) > 0
        )
        if not self.args.only_pretrain_model and has_model_class_in_config:
               try:
                      lowerCAmelCase__	      =   config.architectures[0]
                      lowerCAmelCase__	      =   __import__('transformers'     ,       fromlist=[model_class])
                      lowerCAmelCase__	      =   getattr(lowercase__     ,       lowercase__)
                      lowerCAmelCase__	      =   model_cls(lowercase__)
               except ImportError:
                      raise ImportError(
                          F"""{model_class} does not exist. If you just want to test the pretrained model, you might want to"""
                          ' set `--only_pretrain_model` or `args.only_pretrain_model=True`.')
        else:
               lowerCAmelCase__	      =   MODEL_MAPPING[config.__class__](lowercase__)

        model.eval()
        model.to(self.args.device)

        # encoder-decoder has vocab size saved differently
        lowerCAmelCase__	      =   config.vocab_size if hasattr(lowercase__     ,       'vocab_size') else config.encoder.vocab_size
        lowerCAmelCase__	      =   torch.randint(lowercase__     ,       (batch_size, sequence_length)     ,       dtype=torch.long     ,       device=self.args.device)

        if self.args.fpaa:
               logger.info('Running training in Mixed Precision...')
               if not self.args.is_gpu:
                      raise ValueError('Mixed precision is possible only for GPU.')
               # amp seems to have memory leaks so that memory usage
               # is measured using .half() for now https://github.com/NVIDIA/apex/issues/439
               model.half()

        if self.args.torchscript:
               with torch.no_grad():
                      lowerCAmelCase__	      =   torch.jit.trace(lowercase__     ,       lowercase__)
        else:
               lowerCAmelCase__	      =   model

        def encoder_decoder_forward():
               with torch.no_grad():
                      lowerCAmelCase__	      =   inference_model(lowercase__     ,       decoder_input_ids=lowercase__)
               return outputs

        def encoder_forward():
               with torch.no_grad():
                      lowerCAmelCase__	      =   inference_model(lowercase__)
               return outputs

        lowerCAmelCase__	      =   encoder_decoder_forward if config.is_encoder_decoder else encoder_forward
        return _forward




 def 						__snake_case		(     self   :    Any     ,       lowercase__   :    str     ,       lowercase__   :    int     ,       lowercase__   :    int):







        '''simple docstring'''



        lowerCAmelCase__	      =   self.config_dict[model_name]

        lowerCAmelCase__	      =   (
            hasattr(lowercase__     ,       'architectures')
            and isinstance(config.architectures     ,       lowercase__)
            and len(config.architectures) > 0
        )
        if not self.args.only_pretrain_model and has_model_class_in_config:
               try:
                      lowerCAmelCase__	      =   config.architectures[0]
                      lowerCAmelCase__	      =   __import__('transformers'     ,       fromlist=[model_class])
                      lowerCAmelCase__	      =   getattr(lowercase__     ,       lowercase__)
                      lowerCAmelCase__	      =   model_cls(lowercase__)
               except ImportError:
                      raise ImportError(
                          F"""{model_class} does not exist. If you just want to test the pretrained model, you might want to"""
                          ' set `--only_pretrain_model` or `args.only_pretrain_model=True`.')
        else:
               lowerCAmelCase__	      =   MODEL_WITH_LM_HEAD_MAPPING[config.__class__](lowercase__)

        if self.args.torchscript:
               raise NotImplementedError('Training for torchscript is currently not implemented')
        else:
               lowerCAmelCase__	      =   model

        model.train()
        model.to(self.args.device)

        # encoder-decoder has vocab size saved differently
        lowerCAmelCase__	      =   config.vocab_size if hasattr(lowercase__     ,       'vocab_size') else config.encoder.vocab_size
        lowerCAmelCase__	      =   torch.randint(lowercase__     ,       (batch_size, sequence_length)     ,       dtype=torch.long     ,       device=self.args.device)

        if self.args.fpaa:
               logger.info('Running training in Mixed Precision...')
               if not self.args.is_gpu:
                      raise ValueError('Mixed precision is possible only for GPU.')

               # amp seems to have memory leaks so that memory usage
               # is measured using .half() for now https://github.com/NVIDIA/apex/issues/439
               model.half()

        def compute_loss_and_backprob_encoder():
               lowerCAmelCase__	      =   train_model(lowercase__     ,       labels=lowercase__)[0]
               loss.backward()
               return loss

        def compute_loss_and_backprob_encoder_decoder():
               lowerCAmelCase__	      =   train_model(lowercase__     ,       decoder_input_ids=lowercase__     ,       labels=lowercase__)[0]
               loss.backward()
               return loss

        lowerCAmelCase__	      =   (
            compute_loss_and_backprob_encoder_decoder
            if config.is_encoder_decoder
            else compute_loss_and_backprob_encoder
        )
        return _train




 def 						__snake_case		(     self   :    Optional[Any]     ,       lowercase__   :    str):







        '''simple docstring'''



        try:
               if self.args.is_tpu or self.args.torchscript:
                      # run additional 10 times to stabilize compilation for tpu and torchscript
                      logger.info('Do inference on TPU or torchscript. Running model 5 times to stabilize compilation')
                      timeit.repeat(
                          lowercase__     ,       repeat=1     ,       number=5     ,       )

               # as written in https://docs.python.org/2/library/timeit.html#timeit.Timer.repeat, min should be taken rather than the average
               lowerCAmelCase__	      =   timeit.repeat(
                   lowercase__     ,       repeat=self.args.repeat     ,       number=10     ,       )

               if self.args.is_tpu and self.args.torch_xla_tpu_print_metrics:
                      import torch_xla.debug.metrics as met

                      self.print_fn(met.metrics_report())

               return min(lowercase__) / 10.0
        except RuntimeError as e:
               self.print_fn(F"""Doesn't fit on GPU. {e}""")
               return "N/A"




 def 						__snake_case		(     self   :    Union[str, Any]     ,       lowercase__   :    Callable[[], None]):







        '''simple docstring'''



        try:
               if self.args.trace_memory_line_by_line:
                      lowerCAmelCase__	      =   start_memory_tracing('transformers')

               if self.args.is_tpu:
                      # tpu
                      raise NotImplementedError(
                          'Memory Benchmarking is currently not implemented for TPU. Please disable memory benchmarking with'
                          ' `--no-memory` or `args.memory=False`')
               elif self.args.is_gpu:
                      if not is_pyanvml_available():
                             logger.warning(
                                 'py3nvml not installed, we won\'t log GPU memory usage. '
                                 'Install py3nvml (pip install py3nvml) to log information about GPU.')
                             lowerCAmelCase__	      =   'N/A'
                      else:
                             logger.info(
                                 'Measuring total GPU usage on GPU device. Make sure to not have additional processes running'
                                 ' on the same GPU.')
                             # init nvml
                             nvml.nvmlInit()
                             func()
                             lowerCAmelCase__	      =   nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx)
                             lowerCAmelCase__	      =   nvml.nvmlDeviceGetMemoryInfo(lowercase__)
                             lowerCAmelCase__	      =   meminfo.used
                             lowerCAmelCase__	      =   Memory(lowercase__)
                             # shutdown nvml
                             nvml.nvmlShutdown()
               else:
                      # cpu
                      lowerCAmelCase__	      =   measure_peak_memory_cpu(lowercase__)
                      lowerCAmelCase__	      =   Memory(lowercase__) if isinstance(lowercase__     ,       lowercase__) else memory_bytes

               if self.args.trace_memory_line_by_line:
                      lowerCAmelCase__	      =   stop_memory_tracing(lowercase__)
               else:
                      lowerCAmelCase__	      =   None

               return memory, summary
        except RuntimeError as e:
               self.print_fn(F"""Doesn't fit on GPU. {e}""")
               return "N/A", None

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__      		=					'./tests/fixtures/elise_format0.mid'
@unittest.skip('The note-seq package currently throws an error on import'		)
class 							a_					(		SCREAMING_SNAKE_CASE			,	unittest.TestCase		):






 '''simple docstring'''






 UpperCAmelCase_      					=   SpectrogramDiffusionPipeline
 UpperCAmelCase_      					=   PipelineTesterMixin.required_optional_params - {
     'callback',
     'latents',
     'callback_steps',
     'output_type',
     'num_images_per_prompt',
 }
 UpperCAmelCase_      					=   False
 UpperCAmelCase_      					=   False
 UpperCAmelCase_      					=   TOKENS_TO_AUDIO_GENERATION_PARAMS
 UpperCAmelCase_      					=   TOKENS_TO_AUDIO_GENERATION_BATCH_PARAMS




 def 						__snake_case		(     self   :    List[Any]):







        '''simple docstring'''



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

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

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

        lowerCAmelCase__	      =   DDPMScheduler()

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




 def 						__snake_case		(     self   :    Optional[int]     ,       lowercase__   :    int     ,       lowercase__   :    List[str]=0):







        '''simple docstring'''



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







        '''simple docstring'''



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

        lowerCAmelCase__	      =   self.get_dummy_inputs(lowercase__)
        lowerCAmelCase__	      =   pipe(**lowercase__)
        lowerCAmelCase__	      =   output.audios

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

        assert mel_slice.shape == (3, 3)
        lowerCAmelCase__	      =   np.array(
            [-11.512_925, -4.788_215, -0.46_172_905, -2.051_715, -10.539_147, -10.970_963, -9.091_634, 4.0, 4.0])
        assert np.abs(mel_slice.flatten() - expected_slice).max() < 1e-2




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







        '''simple docstring'''



        return super().test_save_load_local()




 @skip_mps
 def 						__snake_case		(     self   :    Optional[Any]):







        '''simple docstring'''



        return super().test_dict_tuple_outputs_equivalent()




 @skip_mps
 def 						__snake_case		(     self   :    Dict):







        '''simple docstring'''



        return super().test_save_load_optional_components()




 @skip_mps
 def 						__snake_case		(     self   :    str):







        '''simple docstring'''



        return super().test_attention_slicing_forward_pass()




 def 						__snake_case		(     self   :    Optional[Any]):







        '''simple docstring'''



        pass




 def 						__snake_case		(     self   :    List[Any]):







        '''simple docstring'''



        pass




 @skip_mps
 def 						__snake_case		(     self   :    Optional[Any]):







        '''simple docstring'''



        return super().test_progress_bar()






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






 '''simple docstring'''






 def 						__snake_case		(     self   :    Optional[int]):







        '''simple docstring'''



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




 def 						__snake_case		(     self   :    Optional[Any]):







        '''simple docstring'''



        lowerCAmelCase__	      =   torch_device

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

        lowerCAmelCase__	      =   pipe.to(lowercase__)
        pipe.set_progress_bar_config(disable=lowercase__)

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

        lowerCAmelCase__	      =   MidiProcessor()
        lowerCAmelCase__	      =   processor(lowercase__)

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




 def 						__snake_case		(     self   :    str):







        '''simple docstring'''



        lowerCAmelCase__	      =   torch_device

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

        lowerCAmelCase__	      =   processor(lowercase__)
        # just run two denoising loops
        lowerCAmelCase__	      =   input_tokens[:2]

        lowerCAmelCase__	      =   torch.manual_seed(0)
        lowerCAmelCase__	      =   pipe(lowercase__     ,       num_inference_steps=2     ,       generator=lowercase__)

        lowerCAmelCase__	      =   output.audios[0]

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




 def 						__snake_case		(     self   :    Dict):







        '''simple docstring'''



        lowerCAmelCase__	      =   torch_device

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

        lowerCAmelCase__	      =   MidiProcessor()

        lowerCAmelCase__	      =   processor(lowercase__)

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

        lowerCAmelCase__	      =   torch.manual_seed(0)
        lowerCAmelCase__	      =   pipe(lowercase__     ,       num_inference_steps=100     ,       generator=lowercase__)

        lowerCAmelCase__	      =   output.audios[0]
        assert abs(np.abs(lowercase__).sum() - 9_389.1_111) < 5e-2
