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


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



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






   def  __UpperCamelCase	( self							:	Optional[int]    )   ->		List[Any]:
       torch.manual_seed(0    )
       lowercase__					:  Any          =							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"      ,		)

       lowercase__					:  str          =							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"      ,		)

       lowercase__					:  int          =							TaFilmDecoder(
           input_dims=1_28      ,		targets_length=2_56      ,		max_decoder_noise_time=2_00_00.0      ,		d_model=7_68      ,		num_layers=1      ,		num_heads=1      ,		d_kv=4      ,		d_ff=20_48      ,		dropout_rate=0.1      ,		)

       lowercase__					:  int          =							DDPMScheduler()

       lowercase__					:  Optional[int]          =							{
           "notes_encoder": notes_encoder.eval(),
           "continuous_encoder": continuous_encoder.eval(),
           "decoder": decoder.eval(),
           "scheduler": scheduler,
           "melgan": None,
       }
       return components






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

       lowercase__					:  Optional[Any]          =							self.get_dummy_inputs(lowercase_    )
       lowercase__					:  Union[str, Any]          =							pipe(**lowercase_    )
       lowercase__					:  Dict          =							output.audios

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

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






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






   @skip_mps
   def  __UpperCamelCase	( self							:	Union[str, Any]    )   ->		Any:
       return super().test_save_load_optional_components()






   @skip_mps
   def  __UpperCamelCase	( self							:	Any    )   ->		Any:
       return super().test_attention_slicing_forward_pass()






   def  __UpperCamelCase	( self							:	List[Any]    )   ->		Union[str, Any]:
       pass






   def  __UpperCamelCase	( self							:	Tuple    )   ->		Dict:
       pass







   @skip_mps
   def  __UpperCamelCase	( self							:	str    )   ->		int:
       return super().test_progress_bar()



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






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






   def  __UpperCamelCase	( self							:	List[str]    )   ->		int:
       # TODO - test that pipeline can decode tokens in a callback
       # so that music can be played live
       lowercase__					:  Dict          =							torch_device

       lowercase__					:  Dict          =							SpectrogramDiffusionPipeline.from_pretrained("google/music-spectrogram-diffusion"    )
       lowercase__					:  List[str]          =							pipe.melgan
       lowercase__					:  Dict          =							None

       lowercase__					:  Any          =							pipe.to(lowercase_    )
       pipe.set_progress_bar_config(disable=lowercase_    )

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

       lowercase__					:  str          =							MidiProcessor()
       lowercase__					:  Optional[int]          =							processor(lowercase_    )

       lowercase__					:  Dict          =							input_tokens[:3]
       lowercase__					:  int          =							torch.manual_seed(0    )
       pipe(lowercase_      ,		num_inference_steps=5      ,		generator=lowercase_      ,		callback=lowercase_      ,		output_type="mel"    )






   def  __UpperCamelCase	( self							:	Any    )   ->		List[str]:
       lowercase__					:  Optional[Any]          =							torch_device

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

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

       lowercase__					:  int          =							torch.manual_seed(0    )
       lowercase__					:  Dict          =							pipe(lowercase_      ,		num_inference_steps=2      ,		generator=lowercase_    )

       lowercase__					:  Dict          =							output.audios[0]

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







   def  __UpperCamelCase	( self							:	Tuple    )   ->		Optional[Any]:
       lowercase__					:  Union[str, Any]          =							torch_device

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

       lowercase__					:  str          =							MidiProcessor()

       lowercase__					:  Dict          =							processor(lowercase_    )

       # just run 4 denoising loops
       lowercase__					:  Tuple          =							input_tokens[:4]

       lowercase__					:  int          =							torch.manual_seed(0    )
       lowercase__					:  str          =							pipe(lowercase_      ,		num_inference_steps=1_00      ,		generator=lowercase_    )

       lowercase__					:  int          =							output.audios[0]
       assert abs(np.abs(lowercase_    ).sum() - 93_89.11_11    ) < 5E-2




from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available


UpperCamelCase    				=		{'''configuration_vit_hybrid''': ['''VIT_HYBRID_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''ViTHybridConfig''']}

try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      UpperCamelCase    				=		[
          '''VIT_HYBRID_PRETRAINED_MODEL_ARCHIVE_LIST''',
          '''ViTHybridForImageClassification''',
          '''ViTHybridModel''',
          '''ViTHybridPreTrainedModel''',
      ]

try:
      if not is_vision_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      UpperCamelCase    				=		['''ViTHybridImageProcessor''']


if TYPE_CHECKING:
      from .configuration_vit_hybrid import VIT_HYBRID_PRETRAINED_CONFIG_ARCHIVE_MAP, ViTHybridConfig

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_vit_hybrid import (
                VIT_HYBRID_PRETRAINED_MODEL_ARCHIVE_LIST,
                ViTHybridForImageClassification,
                ViTHybridModel,
                ViTHybridPreTrainedModel,
            )

      try:
            if not is_vision_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .image_processing_vit_hybrid import ViTHybridImageProcessor


else:
      import sys

      UpperCamelCase    				=		_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)



