





import unittest
from pathlib import Path
from tempfile import TemporaryDirectory

from transformers import AutoConfig, TFGPTaLMHeadModel, is_keras_nlp_available, is_tf_available
from transformers.models.gpta.tokenization_gpta import GPTaTokenizer
from transformers.testing_utils import require_keras_nlp, require_tf, slow


if is_tf_available():
     import tensorflow as tf

if is_keras_nlp_available():
     from transformers.models.gpta import TFGPTaTokenizer


_snake_case  :   List[Any]						   =  ['gpt2']
_snake_case  :   Union[str, Any]						   =  'gpt2'

if is_tf_available():



     class     _UpperCAmelCase   (   tf.Module   ):

       """simple docstring"""


       def __init__(    self			:		Union[str, Any] ,  lowerCAmelCase_			:		List[Any]					)							->						List[str]:
          super().__init__()
          __lowerCAmelCase							     =     tokenizer
          __lowerCAmelCase							     =     AutoConfig.from_pretrained(lowerCAmelCase_					)
          __lowerCAmelCase							     =     TFGPTaLMHeadModel.from_config(lowerCAmelCase_					)






       @tf.function(input_signature=(tf.TensorSpec((None,) ,  tf.string ,  name='text'					),)					)
       def     lowercase		(    self			:		Union[str, Any] ,  lowerCAmelCase_			:		Optional[Any]					)							->						Dict:
          __lowerCAmelCase							     =     self.tokenizer(lowerCAmelCase_					)
          __lowerCAmelCase							     =     tokenized['input_ids'].to_tensor()

          __lowerCAmelCase							     =     tf.cast(input_ids_dense > 0 ,  tf.intaa					)
          # input_mask = tf.reshape(input_mask, [-1, MAX_SEQ_LEN])

          __lowerCAmelCase							     =     self.model(input_ids=lowerCAmelCase_ ,  attention_mask=lowerCAmelCase_					)['logits']

          return outputs



@require_tf
@require_keras_nlp
class     _UpperCAmelCase   (   unittest.TestCase   ):

  """simple docstring"""


  def     lowercase		(    self			:		Optional[int]					)							->						str:
     super().setUp()

     __lowerCAmelCase							     =     [GPTaTokenizer.from_pretrained(lowerCAmelCase_					) for checkpoint in (TOKENIZER_CHECKPOINTS)]
     __lowerCAmelCase							     =     [TFGPTaTokenizer.from_pretrained(lowerCAmelCase_					) for checkpoint in TOKENIZER_CHECKPOINTS]
     assert len(self.tokenizers					) == len(self.tf_tokenizers					)

     __lowerCAmelCase							     =     [
         'This is a straightforward English test sentence.',
         'This one has some weird characters\rto\nsee\r\nif  those\u00E9break things.',
         'Now we\'re going to add some Chinese: 一 二 三 一二三',
         'And some much more rare Chinese: 齉 堃 齉堃',
         'Je vais aussi écrire en français pour tester les accents',
         'Classical Irish also has some unusual characters, so in they go: Gaelaċ, ꝼ',
     ]
     __lowerCAmelCase							     =     list(zip(self.test_sentences ,  self.test_sentences[::-1]					)					)






  def     lowercase		(    self			:		List[str]					)							->						Optional[int]:
     for tokenizer, tf_tokenizer in zip(self.tokenizers ,  self.tf_tokenizers					):
        for test_inputs in self.test_sentences:
           __lowerCAmelCase							     =     tokenizer([test_inputs] ,  return_tensors='tf'					)
           __lowerCAmelCase							     =     tf_tokenizer([test_inputs]					)

           for key in python_outputs.keys():
              # convert them to numpy to avoid messing with ragged tensors
              __lowerCAmelCase							     =     python_outputs[key].numpy()
              __lowerCAmelCase							     =     tf_outputs[key].numpy()

              self.assertTrue(tf.reduce_all(python_outputs_values.shape == tf_outputs_values.shape					)					)
              self.assertTrue(tf.reduce_all(tf.cast(lowerCAmelCase_ ,  tf.intaa					) == tf_outputs_values					)					)






  @slow
  def     lowercase		(    self			:		Dict					)							->						Dict:
     for tf_tokenizer in self.tf_tokenizers:
        __lowerCAmelCase							     =     tf.function(lowerCAmelCase_					)
        for test_inputs in self.test_sentences:
           __lowerCAmelCase							     =     tf.constant(lowerCAmelCase_					)
           __lowerCAmelCase							     =     compiled_tokenizer(lowerCAmelCase_					)
           __lowerCAmelCase							     =     tf_tokenizer(lowerCAmelCase_					)

           for key in eager_outputs.keys():
              self.assertTrue(tf.reduce_all(eager_outputs[key] == compiled_outputs[key]					)					)






  @slow
  def     lowercase		(    self			:		Dict					)							->						Dict:
     for tf_tokenizer in self.tf_tokenizers:
        __lowerCAmelCase							     =     ModelToSave(tokenizer=lowerCAmelCase_					)
        __lowerCAmelCase							     =     tf.convert_to_tensor([self.test_sentences[0]]					)
        __lowerCAmelCase							     =     model.serving(lowerCAmelCase_					)  # Build model with some sample inputs
        with TemporaryDirectory() as tempdir:
           __lowerCAmelCase							     =     Path(lowerCAmelCase_					) / 'saved.model'
           tf.saved_model.save(lowerCAmelCase_ ,  lowerCAmelCase_ ,  signatures={'serving_default': model.serving}					)
           __lowerCAmelCase							     =     tf.saved_model.load(lowerCAmelCase_					)
        __lowerCAmelCase							     =     loaded_model.signatures['serving_default'](lowerCAmelCase_					)['output_0']
        # We may see small differences because the loaded model is compiled, so we need an epsilon for the test
        self.assertTrue(tf.reduce_all(out == loaded_output					)					)






  @slow
  def     lowercase		(    self			:		Dict					)							->						Dict:
     for tf_tokenizer in self.tf_tokenizers:
        __lowerCAmelCase							     =     tf.convert_to_tensor([self.test_sentences[0]]					)
        __lowerCAmelCase							     =     tf_tokenizer(lowerCAmelCase_					)  # Build model with some sample inputs

        __lowerCAmelCase							     =     tf_tokenizer.get_config()
        __lowerCAmelCase							     =     TFGPTaTokenizer.from_config(lowerCAmelCase_					)
        __lowerCAmelCase							     =     model_from_config(lowerCAmelCase_					)

        for key in from_config_output.keys():
           self.assertTrue(tf.reduce_all(from_config_output[key] == out[key]					)					)






  @slow
  def     lowercase		(    self			:		Dict					)							->						Dict:
     for tf_tokenizer in self.tf_tokenizers:
        # for the test to run
        __lowerCAmelCase							     =     1_2_3_1_2_3

        for max_length in [3, 5, 1_0_2_4]:
           __lowerCAmelCase							     =     tf.convert_to_tensor([self.test_sentences[0]]					)
           __lowerCAmelCase							     =     tf_tokenizer(lowerCAmelCase_ ,  max_length=lowerCAmelCase_					)

           __lowerCAmelCase							     =     out['input_ids'].numpy().shape[1]

           assert out_length == max_length








import gc
import unittest

import numpy as np
import torch

from diffusers import DanceDiffusionPipeline, IPNDMScheduler, UNetaDModel
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS, UNCONDITIONAL_AUDIO_GENERATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()



class     _UpperCAmelCase   (   _UpperCamelCase		,   unittest.TestCase   ):

  """simple docstring"""


  a_      =			DanceDiffusionPipeline
  a_      =			UNCONDITIONAL_AUDIO_GENERATION_PARAMS
  a_      =			PipelineTesterMixin.required_optional_params - {
      """callback""",
      """latents""",
      """callback_steps""",
      """output_type""",
      """num_images_per_prompt""",
  }
  a_      =			UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS
  a_      =			False
  a_      =			False






  def     lowercase		(    self			:		List[Any]					)							->						Dict:
     torch.manual_seed(0					)
     __lowerCAmelCase							     =     UNetaDModel(
         block_out_channels=(3_2, 3_2, 6_4) ,  extra_in_channels=1_6 ,  sample_size=5_1_2 ,  sample_rate=1_6_0_0_0 ,  in_channels=2 ,  out_channels=2 ,  flip_sin_to_cos=lowerCAmelCase_ ,  use_timestep_embedding=lowerCAmelCase_ ,  time_embedding_type='fourier' ,  mid_block_type='UNetMidBlock1D' ,  down_block_types=('DownBlock1DNoSkip', 'DownBlock1D', 'AttnDownBlock1D') ,  up_block_types=('AttnUpBlock1D', 'UpBlock1D', 'UpBlock1DNoSkip') ,  )
     __lowerCAmelCase							     =     IPNDMScheduler()

     __lowerCAmelCase							     =     {
         'unet': unet,
         'scheduler': scheduler,
     }
     return components






  def     lowercase		(    self			:		Tuple ,  lowerCAmelCase_			:		Tuple ,  lowerCAmelCase_			:		List[str]=0					)							->						Any:
     if str(lowerCAmelCase_					).startswith('mps'					):
        __lowerCAmelCase							     =     torch.manual_seed(lowerCAmelCase_					)
     else:
        __lowerCAmelCase							     =     torch.Generator(device=lowerCAmelCase_					).manual_seed(lowerCAmelCase_					)
     __lowerCAmelCase							     =     {
         'batch_size': 1,
         'generator': generator,
         'num_inference_steps': 4,
     }
     return inputs






  def     lowercase		(    self			:		Union[str, Any]					)							->						int:
     __lowerCAmelCase							     =     'cpu'  # ensure determinism for the device-dependent torch.Generator
     __lowerCAmelCase							     =     self.get_dummy_components()
     __lowerCAmelCase							     =     DanceDiffusionPipeline(**lowerCAmelCase_					)
     __lowerCAmelCase							     =     pipe.to(lowerCAmelCase_					)
     pipe.set_progress_bar_config(disable=lowerCAmelCase_					)

     __lowerCAmelCase							     =     self.get_dummy_inputs(lowerCAmelCase_					)
     __lowerCAmelCase							     =     pipe(**lowerCAmelCase_					)
     __lowerCAmelCase							     =     output.audios

     __lowerCAmelCase							     =     audio[0, -3:, -3:]

     assert audio.shape == (1, 2, components["unet"].sample_size)
     __lowerCAmelCase							     =     np.array([-0.72_65, 1.00_00, -0.83_88, 0.11_75, 0.94_98, -1.00_00]					)
     assert np.abs(audio_slice.flatten() - expected_slice					).max() < 1e-2






  @skip_mps
  def     lowercase		(    self			:		Union[str, Any]					)							->						Tuple:
     return super().test_save_load_local()






  @skip_mps
  def     lowercase		(    self			:		List[str]					)							->						Dict:
     return super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3					)






  @skip_mps
  def     lowercase		(    self			:		str					)							->						List[str]:
     return super().test_save_load_optional_components()






  @skip_mps
  def     lowercase		(    self			:		List[Any]					)							->						List[str]:
     return super().test_attention_slicing_forward_pass()






  def     lowercase		(    self			:		str					)							->						int:
     super().test_inference_batch_single_identical(expected_max_diff=3e-3					)





@slow
@require_torch_gpu
class     _UpperCAmelCase   (   unittest.TestCase   ):

  """simple docstring"""


  def     lowercase		(    self			:		Any					)							->						Union[str, Any]:
     # clean up the VRAM after each test
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()






  def     lowercase		(    self			:		List[str]					)							->						List[str]:
     __lowerCAmelCase							     =     torch_device

     __lowerCAmelCase							     =     DanceDiffusionPipeline.from_pretrained('harmonai/maestro-150k'					)
     __lowerCAmelCase							     =     pipe.to(lowerCAmelCase_					)
     pipe.set_progress_bar_config(disable=lowerCAmelCase_					)

     __lowerCAmelCase							     =     torch.manual_seed(0					)
     __lowerCAmelCase							     =     pipe(generator=lowerCAmelCase_ ,  num_inference_steps=1_0_0 ,  audio_length_in_s=4.0_96					)
     __lowerCAmelCase							     =     output.audios

     __lowerCAmelCase							     =     audio[0, -3:, -3:]

     assert audio.shape == (1, 2, pipe.unet.sample_size)
     __lowerCAmelCase							     =     np.array([-0.01_92, -0.02_31, -0.03_18, -0.00_59, 0.00_02, -0.00_20]					)

     assert np.abs(audio_slice.flatten() - expected_slice					).max() < 1e-2






  def     lowercase		(    self			:		Tuple					)							->						Dict:
     __lowerCAmelCase							     =     torch_device

     __lowerCAmelCase							     =     DanceDiffusionPipeline.from_pretrained('harmonai/maestro-150k' ,  torch_dtype=torch.floataa					)
     __lowerCAmelCase							     =     pipe.to(lowerCAmelCase_					)
     pipe.set_progress_bar_config(disable=lowerCAmelCase_					)

     __lowerCAmelCase							     =     torch.manual_seed(0					)
     __lowerCAmelCase							     =     pipe(generator=lowerCAmelCase_ ,  num_inference_steps=1_0_0 ,  audio_length_in_s=4.0_96					)
     __lowerCAmelCase							     =     output.audios

     __lowerCAmelCase							     =     audio[0, -3:, -3:]

     assert audio.shape == (1, 2, pipe.unet.sample_size)
     __lowerCAmelCase							     =     np.array([-0.03_67, -0.04_88, -0.07_71, -0.05_25, -0.04_44, -0.03_41]					)

     assert np.abs(audio_slice.flatten() - expected_slice					).max() < 1e-2

