
'''simple docstring'''
import unittest

from transformers.testing_utils import is_torch_available, require_torch, tooslow

from ...generation.test_utils import torch_device
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import (
      CpmAntConfig,
      CpmAntForCausalLM,
      CpmAntModel,
      CpmAntTokenizer,
  )






@require_torch
class 				UpperCAmelCase_			:





 def __init__(      self	,    lowercase_	,    lowercase_=2	,    lowercase_=8	,    lowercase_=True	,    lowercase_=False	,    lowercase_=False	,    lowercase_=False	,    lowercase_=False	,    lowercase_=99	,    lowercase_=32	,    lowercase_=3	,    lowercase_=4	,    lowercase_=37	,    lowercase_=32	,    lowercase_=1_28	,    lowercase_=8	,    lowercase_=8	,    lowercase_=8	,    lowercase_=1.0	,    lowercase_=True	,    ):
     snake_case_	:      Union[str, Any]            =  parent
     snake_case_	:      Tuple            =  batch_size
     snake_case_	:      Dict            =  seq_length
     snake_case_	:      Optional[Any]            =  is_training
     snake_case_	:      Optional[int]            =  use_token_type_ids
     snake_case_	:      List[str]            =  use_input_mask
     snake_case_	:      str            =  use_labels
     snake_case_	:      Tuple            =  use_mc_token_ids
     snake_case_	:      Dict            =  vocab_size
     snake_case_	:      Any            =  hidden_size
     snake_case_	:      Union[str, Any]            =  num_hidden_layers
     snake_case_	:      List[Any]            =  num_attention_heads
     snake_case_	:      Union[str, Any]            =  intermediate_size
     snake_case_	:      Tuple            =  num_buckets
     snake_case_	:      Optional[Any]            =  max_distance
     snake_case_	:      List[Any]            =  prompt_length
     snake_case_	:      Optional[int]            =  prompt_types
     snake_case_	:      List[Any]            =  segment_types
     snake_case_	:      List[str]            =  init_std
     snake_case_	:      int            =  return_dict





 def     snake_case__						(      self):
     snake_case_	:      Optional[Any]            =  {}
     snake_case_	:      Optional[int]            =  ids_tensor([self.batch_size, self.seq_length]	,    self.vocab_size).type(torch.intaa)
     snake_case_	:      Any            =  False

     snake_case_	:      Any            =  self.get_config()

     return (config, input_ids)





 def     snake_case__						(      self):
     return CpmAntConfig(
         vocab_size=self.vocab_size	,    hidden_size=self.hidden_size	,    num_hidden_layers=self.num_hidden_layers	,    num_attention_heads=self.num_attention_heads	,    dim_ff=self.intermediate_size	,    position_bias_num_buckets=self.num_buckets	,    position_bias_max_distance=self.max_distance	,    prompt_types=self.prompt_types	,    prompt_length=self.prompt_length	,    segment_types=self.segment_types	,    use_cache=lowercase_	,    init_std=self.init_std	,    return_dict=self.return_dict	,    )





 def     snake_case__						(      self	,    lowercase_	,    lowercase_	,    *lowercase_):
     snake_case_	:      Union[str, Any]            =  CpmAntModel(config=lowercase_)
     model.to(lowercase_)
     model.eval()

     snake_case_	:      Any            =  model(**lowercase_).last_hidden_state

     self.parent.assertEqual(hidden_states.shape	,    (self.batch_size, self.seq_length, config.hidden_size))





 def     snake_case__						(      self	,    lowercase_	,    lowercase_	,    *lowercase_):
     snake_case_	:      Dict            =  CpmAntForCausalLM(lowercase_)
     model.to(lowercase_)
     snake_case_	:      List[Any]            =  input_ids['input_ids'].to(lowercase_)
     model.eval()

     snake_case_	:      Any            =  model(**lowercase_)
     self.parent.assertEqual(
         model_output.logits.shape	,    (self.batch_size, self.seq_length, config.vocab_size + config.prompt_types * config.prompt_length)	,    )





 def     snake_case__						(      self):
     snake_case_	:      Optional[Any]            =  self.prepare_config_and_inputs()
     return config, inputs_dict






@require_torch
class 				UpperCAmelCase_			(		snake_case__      ,  snake_case__      ,  unittest.TestCase   ):
 UpperCAmelCase_						      =  (CpmAntModel, CpmAntForCausalLM) if is_torch_available() else ()
 UpperCAmelCase_						      =  (
     {'feature-extraction': CpmAntModel, 'text-generation': CpmAntForCausalLM} if is_torch_available() else {}
 )

 UpperCAmelCase_						      =  False
 UpperCAmelCase_						      =  False
 UpperCAmelCase_						      =  False
 UpperCAmelCase_						      =  False
 UpperCAmelCase_						      =  False





 def     snake_case__						(      self):
     snake_case_	:      Dict            =  CpmAntModelTester(self)
     snake_case_	:      List[str]            =  ConfigTester(self	,    config_class=lowercase_)





 def     snake_case__						(      self):
     self.config_tester.create_and_test_config_common_properties()
     self.config_tester.create_and_test_config_to_json_string()
     self.config_tester.create_and_test_config_to_json_file()
     self.config_tester.create_and_test_config_from_and_save_pretrained()
     self.config_tester.check_config_can_be_init_without_params()
     self.config_tester.check_config_arguments_init()





 def     snake_case__						(      self):
     unittest.skip("CPMAnt doesn\'t support input_embeds.")(self.test_inputs_embeds)





 def     snake_case__						(      self):
     unittest.skip(
         "CPMAnt doesn\'t support retain grad in hidden_states or attentions, because prompt management will peel off the output.hidden_states from graph.\
                 So is attentions. We strongly recommand you use loss to tune model.")(self.test_retain_grad_hidden_states_attentions)





 def     snake_case__						(      self):
     snake_case_	:      Union[str, Any]            =  self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_cpmant_model(lowercase_	,    lowercase_)





 def     snake_case__						(      self):
     snake_case_	:      Any            =  self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_lm_head_model(lowercase_	,    lowercase_)






@require_torch
class 				UpperCAmelCase_			(		unittest.TestCase   ):





 @tooslow
 def     snake_case__						(      self):
     snake_case_	:      Optional[int]            =  '今天天气真好！'
     snake_case_	:      List[str]            =  'openbmb/cpm-ant-10b'
     snake_case_	:      Optional[Any]            =  CpmAntModel.from_pretrained(lowercase_)
     snake_case_	:      Optional[Any]            =  CpmAntTokenizer.from_pretrained(lowercase_)
     snake_case_	:      List[Any]            =  tokenizer(lowercase_	,    return_tensors="pt")
     snake_case_	:      str            =  model(**lowercase_).last_hidden_state

     snake_case_	:      Union[str, Any]            =  torch.tensor(
         [[[6.1_708, 5.9_244, 1.0_835], [6.5_207, 6.2_893, -11.3_324], [-1.0_107, -0.0_576, -5.9_577]]]	,    )
     self.assertTrue(torch.allclose(hidden_states[:, :3, :3]	,    lowercase_	,    atol=1E-2))






@require_torch
class 				UpperCAmelCase_			(		unittest.TestCase   ):





 @tooslow
 def     snake_case__						(      self):
     snake_case_	:      Optional[int]            =  '今天天气真好！'
     snake_case_	:      Any            =  'openbmb/cpm-ant-10b'
     snake_case_	:      int            =  CpmAntForCausalLM.from_pretrained(lowercase_)
     snake_case_	:      Optional[Any]            =  CpmAntTokenizer.from_pretrained(lowercase_)
     snake_case_	:      Dict            =  tokenizer(lowercase_	,    return_tensors="pt")
     snake_case_	:      Dict            =  model(**lowercase_).logits

     snake_case_	:      List[Any]            =  torch.tensor(
         [[[-6.4_267, -6.4_083, -6.3_958], [-5.8_802, -5.9_447, -5.7_811], [-5.3_896, -5.4_820, -5.4_295]]]	,    )
     self.assertTrue(torch.allclose(hidden_states[:, :3, :3]	,    lowercase_	,    atol=1E-2))





 @tooslow
 def     snake_case__						(      self):
     snake_case_	:      int            =  'openbmb/cpm-ant-10b'
     snake_case_	:      Dict            =  CpmAntForCausalLM.from_pretrained(lowercase_)
     snake_case_	:      Optional[int]            =  CpmAntTokenizer.from_pretrained(lowercase_)
     snake_case_	:      Dict            =  '今天天气不错，'
     snake_case_	:      int            =  '今天天气不错，阳光明媚，我和妈妈一起去超市买东西。\n在超市里，我看到了一个很好玩的玩具，它的名字叫“机器人”。它有一个圆圆的脑袋，两只圆圆的眼睛，还有一个圆圆的'
     snake_case_	:      List[str]            =  tokenizer(lowercase_	,    return_tensors="pt")
     snake_case_	:      List[str]            =  model.generate(**lowercase_)
     snake_case_	:      Tuple            =  tokenizer.batch_decode(lowercase_)
     self.assertEqual(lowercase_	,    lowercase_)





 @tooslow
 def     snake_case__						(      self):
     snake_case_	:      Tuple            =  'openbmb/cpm-ant-10b'
     snake_case_	:      Any            =  CpmAntForCausalLM.from_pretrained(lowercase_)
     snake_case_	:      Optional[Any]            =  CpmAntTokenizer.from_pretrained(lowercase_)
     snake_case_	:      Dict            =  ['今天天气不错，', '新年快乐，万事如意！']
     snake_case_	:      Union[str, Any]            =  [
         '今天天气不错，阳光明媚，我和妈妈一起去超市买东西。\n在超市里，我看到了一个很好玩的玩具，它的名字叫“机器人”。它有一个圆圆的脑袋，两只圆圆的眼睛，还有一个圆圆的',
         '新年快乐，万事如意！在这辞旧迎新的美好时刻，我谨代表《农村新技术》杂志社全体同仁，向一直以来关心、支持《农村新技术》杂志发展的各级领导、各界朋友和广大读者致以最诚挚的',
     ]
     snake_case_	:      Any            =  tokenizer(lowercase_	,    return_tensors="pt"	,    padding=lowercase_)
     snake_case_	:      str            =  model.generate(**lowercase_)
     snake_case_	:      str            =  tokenizer.batch_decode(lowercase_)
     self.assertEqual(lowercase_	,    lowercase_)






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


a_     =  "./tests/fixtures/elise_format0.mid"






@unittest.skip("""The note-seq package currently throws an error on import"""   )
class 				UpperCAmelCase_			(		snake_case__      ,  unittest.TestCase   ):
 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):
     torch.manual_seed(0)
     snake_case_	:      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"	,    )

     snake_case_	:      int            =  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"	,    )

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

     snake_case_	:      Union[str, Any]            =  DDPMScheduler()

     snake_case_	:      List[Any]            =  {
         "notes_encoder": notes_encoder.eval(),
         "continuous_encoder": continuous_encoder.eval(),
         "decoder": decoder.eval(),
         "scheduler": scheduler,
         "melgan": None,
     }
     return components





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

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

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

     assert mel_slice.shape == (3, 3)
     snake_case_	:      Any            =  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):
     return super().test_save_load_local()





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





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





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





 def     snake_case__						(      self):
     pass





 def     snake_case__						(      self):
     pass





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






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





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





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

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

     snake_case_	:      int            =  pipe.to(lowercase_)
     pipe.set_progress_bar_config(disable=lowercase_)

     def callback(lowercase_	,    lowercase_):
         # decode mel to audio
         snake_case_	:      Dict            =  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

     snake_case_	:      Tuple            =  MidiProcessor()
     snake_case_	:      int            =  processor(lowercase_)

     snake_case_	:      int            =  input_tokens[:3]
     snake_case_	:      Union[str, Any]            =  torch.manual_seed(0)
     pipe(lowercase_	,    num_inference_steps=5	,    generator=lowercase_	,    callback=lowercase_	,    output_type="mel")





 def     snake_case__						(      self):
     snake_case_	:      Any            =  torch_device

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

     snake_case_	:      List[str]            =  processor(lowercase_)
     # just run two denoising loops
     snake_case_	:      int            =  input_tokens[:2]

     snake_case_	:      Optional[Any]            =  torch.manual_seed(0)
     snake_case_	:      Tuple            =  pipe(lowercase_	,    num_inference_steps=2	,    generator=lowercase_)

     snake_case_	:      Optional[Any]            =  output.audios[0]

     assert abs(np.abs(lowercase_).sum() - 3_612.841) < 1E-1





 def     snake_case__						(      self):
     snake_case_	:      Union[str, Any]            =  torch_device

     snake_case_	:      Tuple            =  SpectrogramDiffusionPipeline.from_pretrained("google/music-spectrogram-diffusion")
     snake_case_	:      Optional[Any]            =  pipe.to(lowercase_)
     pipe.set_progress_bar_config(disable=lowercase_)

     snake_case_	:      Tuple            =  MidiProcessor()

     snake_case_	:      int            =  processor(lowercase_)

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

     snake_case_	:      Optional[Any]            =  torch.manual_seed(0)
     snake_case_	:      Tuple            =  pipe(lowercase_	,    num_inference_steps=1_00	,    generator=lowercase_)

     snake_case_	:      Tuple            =  output.audios[0]
     assert abs(np.abs(lowercase_).sum() - 9_389.1_111) < 5E-2




