code: 
  
  
  
  
  
  import unittest
  
  import numpy as np
  import timeout_decorator  # noqa
  
  from transformers import MarianConfig, is_flax_available
  from transformers.testing_utils import require_flax, require_sentencepiece, require_tokenizers, slow
  from transformers.utils import cached_property
  
  from ...generation.test_flax_utils import FlaxGenerationTesterMixin
  from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor
  
  
  if is_flax_available():
  		import os
  
  		# The slow tests are often failing with OOM error on GPU
  		# This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
  		# but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
  		_lowercase		:     Any 					='''platform'''
  
  		import jax
  		import jax.numpy as jnp
  
  		from transformers import MarianTokenizer
  		from transformers.models.marian.modeling_flax_marian import FlaxMarianModel, FlaxMarianMTModel, shift_tokens_right
  
  
  def 		_SCREAMING_SNAKE_CASE    (   lowerCAmelCase__       ,lowerCAmelCase__       ,lowerCAmelCase__=None       ,lowerCAmelCase__=None       ,lowerCAmelCase__=None       ,lowerCAmelCase__=None       ,lowerCAmelCase__=None       ,lowerCAmelCase__=None       ,):
  		if attention_mask is None:
  				lowerCamelCase_  :    List[Any]   	=			np.where(input_ids != config.pad_token_id       ,1       ,0	)
  		if decoder_attention_mask is None:
  				lowerCamelCase_  :    List[Any]   	=			np.where(decoder_input_ids != config.pad_token_id       ,1       ,0	)
  		if head_mask is None:
  				lowerCamelCase_  :    Tuple   	=			np.ones((config.encoder_layers, config.encoder_attention_heads)	)
  		if decoder_head_mask is None:
  				lowerCamelCase_  :    str   	=			np.ones((config.decoder_layers, config.decoder_attention_heads)	)
  		if cross_attn_head_mask is None:
  				lowerCamelCase_  :    List[Any]   	=			np.ones((config.decoder_layers, config.decoder_attention_heads)	)
  		return {
  		    "input_ids": input_ids,
  		    "decoder_input_ids": decoder_input_ids,
  		    "attention_mask": attention_mask,
  		    "decoder_attention_mask": attention_mask,
  		}
  
  
  
  
  
  
  class 	UpperCamelCase_     :
  
  
  
  					def __init__(		self    :    int					,			lowerCamelCase    :    Optional[Any]					,			lowerCamelCase    :    Tuple=13					,			lowerCamelCase    :    Tuple=7					,			lowerCamelCase    :    Union[str, Any]=True					,			lowerCamelCase    :    List[Any]=False					,			lowerCamelCase    :    Dict=99					,			lowerCamelCase    :    Union[str, Any]=16					,			lowerCamelCase    :    str=2					,			lowerCamelCase    :    Optional[Any]=4					,			lowerCamelCase    :    str=4					,			lowerCamelCase    :    Optional[int]="gelu"					,			lowerCamelCase    :    Optional[Any]=0.1					,			lowerCamelCase    :    Union[str, Any]=0.1					,			lowerCamelCase    :    Dict=32					,			lowerCamelCase    :    Any=2					,			lowerCamelCase    :    Union[str, Any]=1					,			lowerCamelCase    :    Optional[Any]=0					,			lowerCamelCase    :    Tuple=0.02					,			):
  							lowerCamelCase_  :    str   	=			parent
  							lowerCamelCase_  :    str   	=			batch_size
  							lowerCamelCase_  :    Any   	=			seq_length
  							lowerCamelCase_  :    int   	=			is_training
  							lowerCamelCase_  :    Union[str, Any]   	=			use_labels
  							lowerCamelCase_  :    Optional[int]   	=			vocab_size
  							lowerCamelCase_  :    Optional[Any]   	=			hidden_size
  							lowerCamelCase_  :    Tuple   	=			num_hidden_layers
  							lowerCamelCase_  :    Optional[int]   	=			num_attention_heads
  							lowerCamelCase_  :    Optional[Any]   	=			intermediate_size
  							lowerCamelCase_  :    int   	=			hidden_act
  							lowerCamelCase_  :    int   	=			hidden_dropout_prob
  							lowerCamelCase_  :    List[Any]   	=			attention_probs_dropout_prob
  							lowerCamelCase_  :    Union[str, Any]   	=			max_position_embeddings
  							lowerCamelCase_  :    Optional[Any]   	=			eos_token_id
  							lowerCamelCase_  :    Optional[int]   	=			pad_token_id
  							lowerCamelCase_  :    Dict   	=			bos_token_id
  							lowerCamelCase_  :    List[Any]   	=			initializer_range
  
  
  
  					def      __a    (		self    :    List[str]	):
  							lowerCamelCase_  :    Tuple   	=			np.clip(ids_tensor([self.batch_size, self.seq_length - 1]					,			self.vocab_size	)					,			3					,			self.vocab_size	)
  							lowerCamelCase_  :    List[Any]   	=			np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)					,			dtype=np.intaa	))					,			-1	)
  
  							lowerCamelCase_  :    List[str]   	=			shift_tokens_right(lowerCamelCase					,			1					,			2	)
  
  							lowerCamelCase_  :    Optional[Any]   	=			MarianConfig(
  							    vocab_size=self.vocab_size					,			d_model=self.hidden_size					,			encoder_layers=self.num_hidden_layers					,			decoder_layers=self.num_hidden_layers					,			encoder_attention_heads=self.num_attention_heads					,			decoder_attention_heads=self.num_attention_heads					,			encoder_ffn_dim=self.intermediate_size					,			decoder_ffn_dim=self.intermediate_size					,			dropout=self.hidden_dropout_prob					,			attention_dropout=self.attention_probs_dropout_prob					,			max_position_embeddings=self.max_position_embeddings					,			eos_token_id=self.eos_token_id					,			bos_token_id=self.bos_token_id					,			pad_token_id=self.pad_token_id					,			initializer_range=self.initializer_range					,			use_cache=lowerCamelCase					,			)
  							lowerCamelCase_  :    Tuple   	=			prepare_marian_inputs_dict(lowerCamelCase					,			lowerCamelCase					,			lowerCamelCase	)
  							return config, inputs_dict
  
  
  
  					def      __a    (		self    :    Any	):
  							lowerCamelCase_  :    List[Any]   	=			self.prepare_config_and_inputs()
  							return config, inputs_dict
  
  
  
  					def      __a    (		self    :    int					,			lowerCamelCase    :    List[str]					,			lowerCamelCase    :    Union[str, Any]					,			lowerCamelCase    :    List[Any]	):
  							lowerCamelCase_  :    List[str]   	=			20
  							lowerCamelCase_  :    str   	=			model_class_name(lowerCamelCase	)
  
  							lowerCamelCase_  :    int   	=			model.encode(inputs_dict['input_ids']	)
  
  							lowerCamelCase_  :    Tuple   	=			(
  							    inputs_dict["decoder_input_ids"],
  							    inputs_dict["decoder_attention_mask"],
  							)
  
  							lowerCamelCase_  :    Optional[Any]   	=			model.init_cache(decoder_input_ids.shape[0]					,			lowerCamelCase					,			lowerCamelCase	)
  							lowerCamelCase_  :    Optional[int]   	=			jnp.ones((decoder_input_ids.shape[0], max_decoder_length)					,			dtype='i4'	)
  
  							lowerCamelCase_  :    List[str]   	=			jnp.broadcast_to(
  							    jnp.arange(decoder_input_ids.shape[-1] - 1	)[None, :]					,			(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)					,			)
  							lowerCamelCase_  :    int   	=			model.decode(
  							    decoder_input_ids[:, :-1]					,			lowerCamelCase					,			decoder_attention_mask=lowerCamelCase					,			past_key_values=lowerCamelCase					,			decoder_position_ids=lowerCamelCase					,			)
  
  							lowerCamelCase_  :    List[Any]   	=			jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]					,			dtype='i4'	)
  							lowerCamelCase_  :    Tuple   	=			model.decode(
  							    decoder_input_ids[:, -1:]					,			lowerCamelCase					,			decoder_attention_mask=lowerCamelCase					,			past_key_values=outputs_cache.past_key_values					,			decoder_position_ids=lowerCamelCase					,			)
  
  							lowerCamelCase_  :    Dict   	=			model.decode(lowerCamelCase					,			lowerCamelCase	)
  
  							lowerCamelCase_  :    Optional[int]   	=			np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])	)	)
  							self.parent.assertTrue(diff < 1E-3					,			msg=F"Max diff is {diff}"	)
  
  
  
  
  
  					def      __a    (		self    :    Dict					,			lowerCamelCase    :    List[str]					,			lowerCamelCase    :    Optional[Any]					,			lowerCamelCase    :    Tuple	):
  							lowerCamelCase_  :    Tuple   	=			20
  							lowerCamelCase_  :    List[Any]   	=			model_class_name(lowerCamelCase	)
  
  							lowerCamelCase_  :    str   	=			model.encode(inputs_dict['input_ids']	)
  
  							lowerCamelCase_  :    Optional[Any]   	=			(
  							    inputs_dict["decoder_input_ids"],
  							    inputs_dict["decoder_attention_mask"],
  							)
  
  							lowerCamelCase_  :    str   	=			jnp.concatenate(
  							    [
  							        decoder_attention_mask,
  							        jnp.zeros((decoder_attention_mask.shape[0], max_decoder_length - decoder_attention_mask.shape[1])	),
  							    ]					,			axis=-1					,			)
  
  							lowerCamelCase_  :    Optional[Any]   	=			model.init_cache(decoder_input_ids.shape[0]					,			lowerCamelCase					,			lowerCamelCase	)
  							lowerCamelCase_  :    Optional[int]   	=			jnp.broadcast_to(
  							    jnp.arange(decoder_input_ids.shape[-1] - 1	)[None, :]					,			(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)					,			)
  
  							lowerCamelCase_  :    Union[str, Any]   	=			model.decode(
  							    decoder_input_ids[:, :-1]					,			lowerCamelCase					,			decoder_attention_mask=lowerCamelCase					,			past_key_values=lowerCamelCase					,			decoder_position_ids=lowerCamelCase					,			)
  							lowerCamelCase_  :    List[Any]   	=			jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]					,			dtype='i4'	)
  							lowerCamelCase_  :    Any   	=			model.decode(
  							    decoder_input_ids[:, -1:]					,			lowerCamelCase					,			past_key_values=outputs_cache.past_key_values					,			decoder_attention_mask=lowerCamelCase					,			decoder_position_ids=lowerCamelCase					,			)
  
  							lowerCamelCase_  :    List[Any]   	=			model.decode(lowerCamelCase					,			lowerCamelCase					,			decoder_attention_mask=lowerCamelCase	)
  
  							lowerCamelCase_  :    Tuple   	=			np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])	)	)
  							self.parent.assertTrue(diff < 1E-3					,			msg=F"Max diff is {diff}"	)
  
  
  
  
  
  
  @require_flax
  class 	UpperCamelCase_     (     A__     ,    unittest.TestCase     ,    A__       ):
  					_a			:	Union[str, Any]							=				True
  					_a			:	int							=				(FlaxMarianModel, FlaxMarianMTModel) if is_flax_available() else ()
  					_a			:	List[str]							=				(FlaxMarianMTModel,) if is_flax_available() else ()
  
  
  
  					def      __a    (		self    :    Optional[Any]	):
  							lowerCamelCase_  :    List[str]   	=			FlaxMarianModelTester(self	)
  
  
  
  					def      __a    (		self    :    Dict	):
  							lowerCamelCase_  :    List[Any]   	=			self.model_tester.prepare_config_and_inputs()
  							for model_class in self.all_model_classes:
  									self.model_tester.check_use_cache_forward(lowerCamelCase					,			lowerCamelCase					,			lowerCamelCase	)
  
  
  
  					def      __a    (		self    :    Optional[int]	):
  							lowerCamelCase_  :    Dict   	=			self.model_tester.prepare_config_and_inputs()
  							for model_class in self.all_model_classes:
  									self.model_tester.check_use_cache_forward_with_attn_mask(lowerCamelCase					,			lowerCamelCase					,			lowerCamelCase	)
  
  
  
  					def      __a    (		self    :    Any	):
  							lowerCamelCase_  :    Union[str, Any]   	=			self.model_tester.prepare_config_and_inputs_for_common()
  
  							for model_class in self.all_model_classes:
  									with self.subTest(model_class.__name__	):
  											lowerCamelCase_  :    Optional[int]   	=			self._prepare_for_class(lowerCamelCase					,			lowerCamelCase	)
  											lowerCamelCase_  :    Tuple   	=			model_class(lowerCamelCase	)
  
  											@jax.jit
  											def encode_jitted(lowerCamelCase    :    Tuple					,			lowerCamelCase    :    List[Any]=None					,			**lowerCamelCase    :    Optional[int]	):
  													return model.encode(input_ids=lowerCamelCase					,			attention_mask=lowerCamelCase	)
  
  											with self.subTest('JIT Enabled'	):
  													lowerCamelCase_  :    str   	=			encode_jitted(**lowerCamelCase	).to_tuple()
  
  											with self.subTest('JIT Disabled'	):
  													with jax.disable_jit():
  															lowerCamelCase_  :    List[str]   	=			encode_jitted(**lowerCamelCase	).to_tuple()
  
  											self.assertEqual(len(lowerCamelCase	)					,			len(lowerCamelCase	)	)
  											for jitted_output, output in zip(lowerCamelCase					,			lowerCamelCase	):
  													self.assertEqual(jitted_output.shape					,			output.shape	)
  
  
  
  					def      __a    (		self    :    List[str]	):
  							lowerCamelCase_  :    Dict   	=			self.model_tester.prepare_config_and_inputs_for_common()
  
  							for model_class in self.all_model_classes:
  									with self.subTest(model_class.__name__	):
  											lowerCamelCase_  :    List[Any]   	=			model_class(lowerCamelCase	)
  											lowerCamelCase_  :    Any   	=			model.encode(inputs_dict['input_ids']					,			inputs_dict['attention_mask']	)
  
  											lowerCamelCase_  :    Any   	=			{
  											    "decoder_input_ids": inputs_dict["decoder_input_ids"],
  											    "decoder_attention_mask": inputs_dict["decoder_attention_mask"],
  											    "encoder_outputs": encoder_outputs,
  											}
  
  											@jax.jit
  											def decode_jitted(lowerCamelCase    :    Dict					,			lowerCamelCase    :    int					,			lowerCamelCase    :    List[str]	):
  													return model.decode(
  													    decoder_input_ids=lowerCamelCase					,			decoder_attention_mask=lowerCamelCase					,			encoder_outputs=lowerCamelCase					,			)
  
  											with self.subTest('JIT Enabled'	):
  													lowerCamelCase_  :    Union[str, Any]   	=			decode_jitted(**lowerCamelCase	).to_tuple()
  
  											with self.subTest('JIT Disabled'	):
  													with jax.disable_jit():
  															lowerCamelCase_  :    Optional[int]   	=			decode_jitted(**lowerCamelCase	).to_tuple()
  
  											self.assertEqual(len(lowerCamelCase	)					,			len(lowerCamelCase	)	)
  											for jitted_output, output in zip(lowerCamelCase					,			lowerCamelCase	):
  													self.assertEqual(jitted_output.shape					,			output.shape	)
  
  
  
  
  
  					@slow
  					def      __a    (		self    :    Optional[Any]	):
  
  
  
  
  
  
  							for model_class_name in self.all_model_classes:
  									lowerCamelCase_  :    Union[str, Any]   	=			model_class_name.from_pretrained('Helsinki-NLP/opus-mt-en-de'	)
  									# FlaxMarianForSequenceClassification expects eos token in input_ids
  									lowerCamelCase_  :    Optional[int]   	=			np.ones((1, 1)	) * model.config.eos_token_id
  									lowerCamelCase_  :    Optional[int]   	=			model(lowerCamelCase	)
  									self.assertIsNotNone(lowerCamelCase	)
  
  
  
  
  
  
  @require_flax
  @require_sentencepiece
  @require_tokenizers
  class 	UpperCamelCase_     (     unittest.TestCase       ):
  					_a			:	Any							=				None
  					_a			:	Tuple							=				None
  
  
  
  					@classmethod
  					def      __a    (		cls    :    Dict	):
  							lowerCamelCase_  :    List[Any]   	=			F"Helsinki-NLP/opus-mt-{cls.src}-{cls.tgt}"
  							return cls
  
  
  
  					@cached_property
  					def      __a    (		self    :    Any	):
  							return MarianTokenizer.from_pretrained(self.model_name	)
  
  
  
  					@property
  					def      __a    (		self    :    List[Any]	):
  							return self.tokenizer.eos_token_id
  
  
  
  					@cached_property
  					def      __a    (		self    :    Union[str, Any]	):
  							lowerCamelCase_  :    FlaxMarianMTModel   	=			FlaxMarianMTModel.from_pretrained(self.model_name	)
  							self.assertEqual(model.config.decoder_start_token_id					,			model.config.pad_token_id	)
  							return model
  
  
  
  					def      __a    (		self    :    Any					,			**lowerCamelCase    :    List[Any]	):
  							lowerCamelCase_  :    List[Any]   	=			self.translate_src_text(**lowerCamelCase	)
  							self.assertListEqual(self.expected_text					,			lowerCamelCase	)
  
  
  
  
  
  					def      __a    (		self    :    Union[str, Any]					,			**lowerCamelCase    :    str	):
  							lowerCamelCase_  :    List[Any]   	=			self.tokenizer(self.src_text					,			padding=lowerCamelCase					,			return_tensors='np'					,			**lowerCamelCase	)
  							lowerCamelCase_  :    Optional[int]   	=			self.model.generate(
  							    model_inputs.input_ids					,			attention_mask=model_inputs.attention_mask					,			num_beams=2					,			max_length=1_28					,			).sequences
  							lowerCamelCase_  :    List[Any]   	=			self.tokenizer.batch_decode(lowerCamelCase					,			skip_special_tokens=lowerCamelCase	)
  							return generated_words
  
  
  
  
  
  
  @require_flax
  @require_sentencepiece
  @require_tokenizers
  class 	UpperCamelCase_     (     A__       ):
  					_a			:	int							=				'''en'''
  					_a			:	List[Any]							=				'''fr'''
  					_a			:	int							=				[
  					    '''I am a small frog.''',
  					    '''Now I can forget the 100 words of german that I know.''',
  					]
  					_a			:	List[str]							=				[
  					    '''Je suis une petite grenouille.''',
  					    '''Maintenant, je peux oublier les 100 mots d\'allemand que je connais.''',
  					]
  
  
  
  					@slow
  					def      __a    (		self    :    int	):
  							self._assert_generated_batch_equal_expected()
  
  
  
  
  
  
  @require_flax
  @require_sentencepiece
  @require_tokenizers
  class 	UpperCamelCase_     (     A__       ):
  					_a			:	int							=				'''fr'''
  					_a			:	List[Any]							=				'''en'''
  					_a			:	Optional[Any]							=				[
  					    '''Donnez moi le micro.''',
  					    '''Tom et Mary étaient assis à une table.''',  # Accents
  					]
  					_a			:	str							=				[
  					    '''Give me the microphone.''',
  					    '''Tom and Mary were sitting at a table.''',
  					]
  
  
  
  					@slow
  					def      __a    (		self    :    Optional[Any]	):
  							self._assert_generated_batch_equal_expected()
  
  
  
  
  
  
  @require_flax
  @require_sentencepiece
  @require_tokenizers
  class 	UpperCamelCase_     (     A__       ):
  					_a			:	Optional[int]							=				'''mt'''
  					_a			:	Tuple							=				'''en'''
  					_a			:	Optional[int]							=				['''Billi messu b\'mod ġentili, Ġesù fejjaq raġel li kien milqut bil - marda kerha tal - ġdiem.''']
  					_a			:	int							=				['''Touching gently, Jesus healed a man who was affected by the sad disease of leprosy.''']
  
  
  
  					@slow
  					def      __a    (		self    :    Union[str, Any]	):
  							self._assert_generated_batch_equal_expected()
  
  
  
  
  
  
  @require_flax
  @require_sentencepiece
  @require_tokenizers
  class 	UpperCamelCase_     (     A__       ):
  					_a			:	List[str]							=				'''en'''
  					_a			:	Any							=				'''de'''
  					_a			:	str							=				[
  					    '''I am a small frog.''',
  					    '''Now I can forget the 100 words of german that I know.''',
  					    '''Tom asked his teacher for advice.''',
  					    '''That\'s how I would do it.''',
  					    '''Tom really admired Mary\'s courage.''',
  					    '''Turn around and close your eyes.''',
  					]
  					_a			:	Dict							=				[
  					    '''Ich bin ein kleiner Frosch.''',
  					    '''Jetzt kann ich die 100 Wörter des Deutschen vergessen, die ich kenne.''',
  					    '''Tom bat seinen Lehrer um Rat.''',
  					    '''So würde ich das machen.''',
  					    '''Tom bewunderte Marias Mut wirklich.''',
  					    '''Drehen Sie sich um und schließen Sie die Augen.''',
  					]
  
  
  
  					@slow
  					def      __a    (		self    :    Dict	):
  							self._assert_generated_batch_equal_expected()
  
  
  
  
  
  
  @require_flax
  @require_sentencepiece
  @require_tokenizers
  class 	UpperCamelCase_     (     A__       ):
  					_a			:	List[Any]							=				'''en'''
  					_a			:	List[Any]							=				'''zh'''
  					_a			:	Optional[int]							=				['''My name is Wolfgang and I live in Berlin''']
  					_a			:	int							=				['''我叫沃尔夫冈 我住在柏林''']
  
  
  
  					@slow
  					def      __a    (		self    :    int	):
  							self._assert_generated_batch_equal_expected()
  
  
  
  
  
  
  @require_flax
  @require_sentencepiece
  @require_tokenizers
  class 	UpperCamelCase_     (     A__       ):
  					_a			:	int							=				'''ru'''
  					_a			:	List[Any]							=				'''fr'''
  					_a			:	List[Any]							=				['''Он показал мне рукопись своей новой пьесы.''']
  					_a			:	List[str]							=				['''Il m\'a montré le manuscrit de sa nouvelle pièce.''']
  
  
  
  					@slow
  					def      __a    (		self    :    Tuple	):
  							self._assert_generated_batch_equal_expected()
  
  
  
  
  
  
  @require_flax
  @require_sentencepiece
  @require_tokenizers
  class 	UpperCamelCase_     (     A__       ):
  					_a			:	str							=				'''en'''
  					_a			:	List[str]							=				'''ROMANCE'''
  					_a			:	Optional[Any]							=				[
  					    '''>>fr<< Don\'t spend so much time watching TV.''',
  					    '''>>pt<< Your message has been sent.''',
  					    '''>>es<< He\'s two years older than me.''',
  					]
  					_a			:	Any							=				[
  					    '''Ne passez pas autant de temps à regarder la télé.''',
  					    '''A sua mensagem foi enviada.''',
  					    '''Es dos años más viejo que yo.''',
  					]
  
  
  
  					@slow
  					def      __a    (		self    :    Union[str, Any]	):
  							self._assert_generated_batch_equal_expected()
  
style_context: 
  
  
  
  
  
  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		:     Optional[Any] 					="""./tests/fixtures/elise_format0.mid"""
  
  
  
  
  
  
  @unittest.skip('The note-seq package currently throws an error on import'       )
  class 	UpperCamelCase_     (     snake_case__     ,    unittest.TestCase       ):
  					_a			:	int							=				SpectrogramDiffusionPipeline
  					_a			:	Dict							=				PipelineTesterMixin.required_optional_params - {
  					    'callback',
  					    'latents',
  					    'callback_steps',
  					    'output_type',
  					    'num_images_per_prompt',
  					}
  					_a			:	str							=				False
  					_a			:	Union[str, Any]							=				False
  					_a			:	Optional[int]							=				TOKENS_TO_AUDIO_GENERATION_PARAMS
  					_a			:	Optional[Any]							=				TOKENS_TO_AUDIO_GENERATION_BATCH_PARAMS
  
  
  
  					def      __a    (		self    :    Dict	):
  							torch.manual_seed(0	)
  							lowerCamelCase_  :    List[str]   	=			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_  :    Union[str, Any]   	=			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_  :    Union[str, Any]   	=			TaFilmDecoder(
  							    input_dims=1_28					,			targets_length=2_56					,			max_decoder_noise_time=20_000.0					,			d_model=7_68					,			num_layers=1					,			num_heads=1					,			d_kv=4					,			d_ff=20_48					,			dropout_rate=0.1					,			)
  
  							lowerCamelCase_  :    Any   	=			DDPMScheduler()
  
  							lowerCamelCase_  :    str   	=			{
  							    'notes_encoder': notes_encoder.eval(),
  							    'continuous_encoder': continuous_encoder.eval(),
  							    'decoder': decoder.eval(),
  							    'scheduler': scheduler,
  							    'melgan': None,
  							}
  							return components
  
  
  
  					def      __a    (		self    :    Any					,			lowerCamelCase    :    Any					,			lowerCamelCase    :    int=0	):
  							if str(lowerCamelCase	).startswith('mps'	):
  									lowerCamelCase_  :    Optional[int]   	=			torch.manual_seed(lowerCamelCase	)
  							else:
  									lowerCamelCase_  :    Dict   	=			torch.Generator(device=lowerCamelCase	).manual_seed(lowerCamelCase	)
  							lowerCamelCase_  :    str   	=			{
  							    '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      __a    (		self    :    Dict	):
  							lowerCamelCase_  :    str   	=			'cpu'  # ensure determinism for the device-dependent torch.Generator
  							lowerCamelCase_  :    Dict   	=			self.get_dummy_components()
  							lowerCamelCase_  :    Tuple   	=			SpectrogramDiffusionPipeline(**lowerCamelCase	)
  							lowerCamelCase_  :    List[str]   	=			pipe.to(lowerCamelCase	)
  							pipe.set_progress_bar_config(disable=lowerCamelCase	)
  
  							lowerCamelCase_  :    Optional[Any]   	=			self.get_dummy_inputs(lowerCamelCase	)
  							lowerCamelCase_  :    str   	=			pipe(**lowerCamelCase	)
  							lowerCamelCase_  :    Tuple   	=			output.audios
  
  							lowerCamelCase_  :    Tuple   	=			mel[0, -3:, -3:]
  
  							assert mel_slice.shape == (3, 3)
  							lowerCamelCase_  :    List[str]   	=			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      __a    (		self    :    Any	):
  							return super().test_save_load_local()
  
  
  
  					@skip_mps
  					def      __a    (		self    :    Tuple	):
  							return super().test_dict_tuple_outputs_equivalent()
  
  
  
  					@skip_mps
  					def      __a    (		self    :    Dict	):
  							return super().test_save_load_optional_components()
  
  
  
  					@skip_mps
  					def      __a    (		self    :    str	):
  							return super().test_attention_slicing_forward_pass()
  
  
  
  					def      __a    (		self    :    Any	):
  							pass
  
  
  
  					def      __a    (		self    :    Optional[int]	):
  							pass
  
  
  
  
  
  					@skip_mps
  					def      __a    (		self    :    Optional[Any]	):
  							return super().test_progress_bar()
  
  
  
  
  
  
  @slow
  @require_torch_gpu
  @require_onnxruntime
  @require_note_seq
  class 	UpperCamelCase_     (     unittest.TestCase       ):
  
  
  
  					def      __a    (		self    :    Optional[int]	):
  							# clean up the VRAM after each test
  							super().tearDown()
  							gc.collect()
  							torch.cuda.empty_cache()
  
  
  
  					def      __a    (		self    :    str	):
  							# TODO - test that pipeline can decode tokens in a callback
  							# so that music can be played live
  							lowerCamelCase_  :    Optional[Any]   	=			torch_device
  
  							lowerCamelCase_  :    Union[str, Any]   	=			SpectrogramDiffusionPipeline.from_pretrained('google/music-spectrogram-diffusion'	)
  							lowerCamelCase_  :    str   	=			pipe.melgan
  							lowerCamelCase_  :    Optional[int]   	=			None
  
  							lowerCamelCase_  :    str   	=			pipe.to(lowerCamelCase	)
  							pipe.set_progress_bar_config(disable=lowerCamelCase	)
  
  							def callback(lowerCamelCase    :    int					,			lowerCamelCase    :    int	):
  									# decode mel to audio
  									lowerCamelCase_  :    Union[str, Any]   	=			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_  :    Any   	=			MidiProcessor()
  							lowerCamelCase_  :    str   	=			processor(lowerCamelCase	)
  
  							lowerCamelCase_  :    str   	=			input_tokens[:3]
  							lowerCamelCase_  :    Dict   	=			torch.manual_seed(0	)
  							pipe(lowerCamelCase					,			num_inference_steps=5					,			generator=lowerCamelCase					,			callback=lowerCamelCase					,			output_type='mel'	)
  
  
  
  					def      __a    (		self    :    str	):
  							lowerCamelCase_  :    List[str]   	=			torch_device
  
  							lowerCamelCase_  :    List[str]   	=			SpectrogramDiffusionPipeline.from_pretrained('google/music-spectrogram-diffusion'	)
  							lowerCamelCase_  :    Union[str, Any]   	=			pipe.to(lowerCamelCase	)
  							pipe.set_progress_bar_config(disable=lowerCamelCase	)
  							lowerCamelCase_  :    int   	=			MidiProcessor()
  
  							lowerCamelCase_  :    Dict   	=			processor(lowerCamelCase	)
  							# just run two denoising loops
  							lowerCamelCase_  :    Optional[Any]   	=			input_tokens[:2]
  
  							lowerCamelCase_  :    List[Any]   	=			torch.manual_seed(0	)
  							lowerCamelCase_  :    str   	=			pipe(lowerCamelCase					,			num_inference_steps=2					,			generator=lowerCamelCase	)
  
  							lowerCamelCase_  :    Dict   	=			output.audios[0]
  
  							assert abs(np.abs(lowerCamelCase	).sum() - 3_612.841	) < 1E-1
  
  
  
  
  
  					def      __a    (		self    :    List[Any]	):
  							lowerCamelCase_  :    int   	=			torch_device
  
  							lowerCamelCase_  :    str   	=			SpectrogramDiffusionPipeline.from_pretrained('google/music-spectrogram-diffusion'	)
  							lowerCamelCase_  :    Optional[int]   	=			pipe.to(lowerCamelCase	)
  							pipe.set_progress_bar_config(disable=lowerCamelCase	)
  
  							lowerCamelCase_  :    Tuple   	=			MidiProcessor()
  
  							lowerCamelCase_  :    int   	=			processor(lowerCamelCase	)
  
  							# just run 4 denoising loops
  							lowerCamelCase_  :    Optional[Any]   	=			input_tokens[:4]
  
  							lowerCamelCase_  :    List[Any]   	=			torch.manual_seed(0	)
  							lowerCamelCase_  :    int   	=			pipe(lowerCamelCase					,			num_inference_steps=1_00					,			generator=lowerCamelCase	)
  
  							lowerCamelCase_  :    List[Any]   	=			output.audios[0]
  							assert abs(np.abs(lowerCamelCase	).sum() - 9_389.1_111	) < 5E-2
  