





"""simple docstring"""






from math import acos, sin
from typing import List, Tuple, Union

import numpy as np
import torch
from PIL import Image

from ...models import AutoencoderKL, UNetaDConditionModel
from ...schedulers import DDIMScheduler, DDPMScheduler
from ...utils import randn_tensor
from ..pipeline_utils import AudioPipelineOutput, BaseOutput, DiffusionPipeline, ImagePipelineOutput
from .mel import Mel



class 	lowerCAmelCase					(      UpperCamelCase_							):
       lowerCAmelCase__						 =  ["""vqvae"""]






       def __init__(   self    ,  a__    ,  a__    ,  a__    ,  a__    ,  ):
              super().__init__()
              self.register_modules(unet=UpperCamelCase__    ,  scheduler=UpperCamelCase__    ,  mel=UpperCamelCase__    ,  vqvae=UpperCamelCase__     )






       def 							__A (   self     ):
              return 50 if isinstance(self.scheduler    ,  UpperCamelCase__     ) else 10_00






       @torch.no_grad()
       def __call__(   self    ,  a__ = 1    ,  a__ = None    ,  a__ = None    ,  a__ = 0    ,  a__ = 0    ,  a__ = None    ,  a__ = None    ,  a__ = 0    ,  a__ = 0    ,  a__ = None    ,  a__ = 0    ,  a__ = None    ,  a__ = None    ,  a__=True    ,  ):
              _UpperCAmelCase     =     steps or self.get_default_steps()
              self.scheduler.set_timesteps(UpperCamelCase__     )
              _UpperCAmelCase     =     step_generator or generator
              # For backwards compatibility
              if type(self.unet.config.sample_size     ) == int:
                     _UpperCAmelCase     =     (self.unet.config.sample_size, self.unet.config.sample_size)
              if noise is None:
                     _UpperCAmelCase     =     randn_tensor(
                         (
                             batch_size,
                             self.unet.config.in_channels,
                             self.unet.config.sample_size[0],
                             self.unet.config.sample_size[1],
                         )    ,  generator=UpperCamelCase__    ,  device=self.device    ,  )
              _UpperCAmelCase     =     noise
              _UpperCAmelCase     =     None

              if audio_file is not None or raw_audio is not None:
                     self.mel.load_audio(UpperCamelCase__    ,  UpperCamelCase__     )
                     _UpperCAmelCase     =     self.mel.audio_slice_to_image(UpperCamelCase__     )
                     _UpperCAmelCase     =     np.frombuffer(input_image.tobytes()    ,  dtype='uint8'     ).reshape(
                         (input_image.height, input_image.width)     )
                     _UpperCAmelCase     =     (input_image / 2_55) * 2 - 1
                     _UpperCAmelCase     =     torch.tensor(input_image[np.newaxis, :, :]    ,  dtype=torch.float     ).to(self.device     )

                     if self.vqvae is not None:
                            _UpperCAmelCase     =     self.vqvae.encode(torch.unsqueeze(UpperCamelCase__    ,  0     )     ).latent_dist.sample(
                                generator=UpperCamelCase__     )[0]
                            _UpperCAmelCase     =     self.vqvae.config.scaling_factor * input_images

                     if start_step > 0:
                            _UpperCAmelCase     =     self.scheduler.add_noise(UpperCamelCase__    ,  UpperCamelCase__    ,  self.scheduler.timesteps[start_step - 1]     )

                     _UpperCAmelCase     =     (
                         self.unet.config.sample_size[1] * self.mel.get_sample_rate() / self.mel.x_res / self.mel.hop_length
                     )
                     _UpperCAmelCase     =     int(mask_start_secs * pixels_per_second     )
                     _UpperCAmelCase     =     int(mask_end_secs * pixels_per_second     )
                     _UpperCAmelCase     =     self.scheduler.add_noise(UpperCamelCase__    ,  UpperCamelCase__    ,  torch.tensor(self.scheduler.timesteps[start_step:]     )     )

              for step, t in enumerate(self.progress_bar(self.scheduler.timesteps[start_step:]     )     ):
                     if isinstance(self.unet    ,  UpperCamelCase__     ):
                            _UpperCAmelCase     =     self.unet(UpperCamelCase__    ,  UpperCamelCase__    ,  UpperCamelCase__     )['''sample''']
                     else:
                            _UpperCAmelCase     =     self.unet(UpperCamelCase__    ,  UpperCamelCase__     )['''sample''']

                     if isinstance(self.scheduler    ,  UpperCamelCase__     ):
                            _UpperCAmelCase     =     self.scheduler.step(
                                model_output=UpperCamelCase__    ,  timestep=UpperCamelCase__    ,  sample=UpperCamelCase__    ,  eta=UpperCamelCase__    ,  generator=UpperCamelCase__    ,  )['''prev_sample''']
                     else:
                            _UpperCAmelCase     =     self.scheduler.step(
                                model_output=UpperCamelCase__    ,  timestep=UpperCamelCase__    ,  sample=UpperCamelCase__    ,  generator=UpperCamelCase__    ,  )['''prev_sample''']

                     if mask is not None:
                            if mask_start > 0:
                                   _UpperCAmelCase     =     mask[:, step, :, :mask_start]
                            if mask_end > 0:
                                   _UpperCAmelCase     =     mask[:, step, :, -mask_end:]

              if self.vqvae is not None:
                     # 0.18215 was scaling factor used in training to ensure unit variance
                     _UpperCAmelCase     =     1 / self.vqvae.config.scaling_factor * images
                     _UpperCAmelCase     =     self.vqvae.decode(UpperCamelCase__     )['''sample''']

              _UpperCAmelCase     =     (images / 2 + 0.5).clamp(0    ,  1     )
              _UpperCAmelCase     =     images.cpu().permute(0    ,  2    ,  3    ,  1     ).numpy()
              _UpperCAmelCase     =     (images * 2_55).round().astype('uint8'     )
              _UpperCAmelCase     =     list(
                  (Image.fromarray(_[:, :, 0]     ) for _ in images)
                  if images.shape[3] == 1
                  else (Image.fromarray(UpperCamelCase__    ,  mode='RGB'     ).convert('L'     ) for _ in images)     )

              _UpperCAmelCase     =     [self.mel.image_to_audio(UpperCamelCase__     ) for _ in images]
              if not return_dict:
                     return images, (self.mel.get_sample_rate(), audios)

              return BaseOutput(**AudioPipelineOutput(np.array(UpperCamelCase__     )[:, np.newaxis, :]     )    ,  **ImagePipelineOutput(UpperCamelCase__     )     )






       @torch.no_grad()
       def 							__A (   self    ,  a__    ,  a__ = 50     ):
              assert isinstance(self.scheduler    ,  UpperCamelCase__     )
              self.scheduler.set_timesteps(UpperCamelCase__     )
              _UpperCAmelCase     =     np.array(
                  [np.frombuffer(image.tobytes()    ,  dtype='uint8'     ).reshape((1, image.height, image.width)     ) for image in images]     )
              _UpperCAmelCase     =     (sample / 2_55) * 2 - 1
              _UpperCAmelCase     =     torch.Tensor(UpperCamelCase__     ).to(self.device     )

              for t in self.progress_bar(torch.flip(self.scheduler.timesteps    ,  (0,)     )     ):
                     _UpperCAmelCase     =     t - self.scheduler.config.num_train_timesteps // self.scheduler.num_inference_steps
                     _UpperCAmelCase     =     self.scheduler.alphas_cumprod[t]
                     _UpperCAmelCase     =     (
                         self.scheduler.alphas_cumprod[prev_timestep]
                         if prev_timestep >= 0
                         else self.scheduler.final_alpha_cumprod
                     )
                     _UpperCAmelCase     =     1 - alpha_prod_t
                     _UpperCAmelCase     =     self.unet(UpperCamelCase__    ,  UpperCamelCase__     )['''sample''']
                     _UpperCAmelCase     =     (1 - alpha_prod_t_prev) ** 0.5 * model_output
                     _UpperCAmelCase     =     (sample - pred_sample_direction) * alpha_prod_t_prev ** (-0.5)
                     _UpperCAmelCase     =     sample * alpha_prod_t ** 0.5 + beta_prod_t ** 0.5 * model_output

              return sample






       @staticmethod
       def 							__A (   a__    ,  a__    ,  a__     ):
              _UpperCAmelCase     =     acos(torch.dot(torch.flatten(UpperCamelCase__     )    ,  torch.flatten(UpperCamelCase__     )     ) / torch.norm(UpperCamelCase__     ) / torch.norm(UpperCamelCase__     )     )
              return sin((1 - alpha) * theta     ) * xa / sin(UpperCamelCase__     ) + sin(alpha * theta     ) * xa / sin(UpperCamelCase__     )






"""simple docstring"""






from __future__ import annotations

import os
import tempfile
import unittest

import numpy as np
from huggingface_hub import hf_hub_download

from transformers import is_tensorflow_text_available, is_tf_available
from transformers.testing_utils import require_tensorflow_text, require_tf, slow

from ..test_modeling_tf_common import floats_tensor
from .test_framework_agnostic import GenerationIntegrationTestsMixin


if is_tf_available():
    import tensorflow as tf

    from transformers import (
        AutoTokenizer,
        TFAutoModelForCausalLM,
        TFAutoModelForSeqaSeqLM,
        TFAutoModelForSpeechSeqaSeq,
        TFAutoModelForVisionaSeq,
        TFBartForConditionalGeneration,
        TFLogitsProcessorList,
        TFMinLengthLogitsProcessor,
        tf_top_k_top_p_filtering,
    )

if is_tensorflow_text_available():
    import tensorflow_text as text



@require_tf
class 	lowerCAmelCase					(      unittest.TestCase							):






       def 							__A (   self     ):
              _UpperCAmelCase     =     tf.convert_to_tensor(
                  [
                      [
                          8.2_220_991,  # 3rd highest value; idx. 0
                          -0.5_620_044,
                          5.23_229_752,
                          4.0_386_393,
                          -6.8_798_378,
                          -0.54_785_802,
                          -3.2_012_153,
                          2.92_777_176,
                          1.88_171_953,
                          7.35_341_276,  # 5th highest value; idx. 9
                          8.43_207_833,  # 2nd highest value; idx. 10
                          -9.85_711_836,
                          -5.96_209_236,
                          -1.13_039_161,
                          -7.1_115_294,
                          -0.8_369_633,
                          -5.3_186_408,
                          7.06_427_407,
                          0.81_369_344,
                          -0.82_023_817,
                          -5.9_179_796,
                          0.58_813_443,
                          -6.99_778_438,
                          4.71_551_189,
                          -0.18_771_637,
                          7.44_020_759,  # 4th highest value; idx. 25
                          9.38_450_987,  # 1st highest value; idx. 26
                          2.12_662_941,
                          -9.32_562_038,
                          2.35_652_522,
                      ],  # cummulative prob of 5 highest values <= 0.6
                      [
                          0.58_425_518,
                          4.53_139_238,
                          -5.57_510_464,
                          -6.28_030_699,
                          -7.19_529_503,
                          -4.02_122_551,
                          1.39_337_037,
                          -6.06_707_057,
                          1.59_480_517,
                          -9.643_119,
                          0.03_907_799,
                          0.67_231_762,
                          -8.88_206_726,
                          6.27_115_922,  # 4th highest value; idx. 13
                          2.28_520_723,
                          4.82_767_506,
                          4.30_421_368,
                          8.8_275_313,  # 2nd highest value; idx. 17
                          5.44_029_958,  # 5th highest value; idx. 18
                          -4.4_735_794,
                          7.38_579_536,  # 3rd highest value; idx. 20
                          -2.91_051_663,
                          2.61_946_077,
                          -2.5_674_762,
                          -9.48_959_302,
                          -4.02_922_645,
                          -1.35_416_918,
                          9.67_702_323,  # 1st highest value; idx. 27
                          -5.89_478_553,
                          1.85_370_467,
                      ],  # cummulative prob of 5 highest values <= 0.6
                  ]    ,  dtype=tf.floataa    ,  )

              _UpperCAmelCase     =     tf.convert_to_tensor(
                  [[0, 0], [0, 9], [0, 10], [0, 25], [0, 26], [1, 13], [1, 17], [1, 18], [1, 20], [1, 27]]    ,  dtype=tf.intaa    ,  )  # expected non filtered idx as noted above

              _UpperCAmelCase     =     tf.convert_to_tensor(
                  [8.222_099, 7.3_534_126, 8.432_078, 7.4_402_075, 9.38_451, 6.271_159, 8.827_531, 5.4_402_995, 7.3_857_956, 9.677_023]    ,  dtype=tf.floataa    ,  )  # expected non filtered values as noted above

              _UpperCAmelCase     =     tf_top_k_top_p_filtering(a__    ,  top_k=10    ,  top_p=0.6    ,  min_tokens_to_keep=4     )

              _UpperCAmelCase     =     output[output != -float('inf'     )]
              _UpperCAmelCase     =     tf.cast(
                  tf.where(tf.not_equal(a__    ,  tf.constant(-float('inf'     )    ,  dtype=tf.floataa     )     )     )    ,  dtype=tf.intaa    ,  )

              tf.debugging.assert_near(a__    ,  a__    ,  rtol=1E-12     )
              tf.debugging.assert_equal(a__    ,  a__     )



@require_tf
class 	lowerCAmelCase					(      unittest.TestCase			,	snake_case							):
       # setting framework_dependent_parameters needs to be gated, just like its contents' imports
       if is_tf_available():
              lowerCAmelCase__						 =  {
                  """AutoModelForCausalLM""": TFAutoModelForCausalLM,
                  """AutoModelForSpeechSeq2Seq""": TFAutoModelForSpeechSeqaSeq,
                  """AutoModelForSeq2SeqLM""": TFAutoModelForSeqaSeqLM,
                  """AutoModelForVision2Seq""": TFAutoModelForVisionaSeq,
                  """LogitsProcessorList""": TFLogitsProcessorList,
                  """MinLengthLogitsProcessor""": TFMinLengthLogitsProcessor,
                  """create_tensor_fn""": tf.convert_to_tensor,
                  """floats_tensor""": floats_tensor,
                  """return_tensors""": """tf""",
              }






       @slow
       def 							__A (   self     ):
              # TF-only test: tf.saved_model export
              _UpperCAmelCase     =     TFAutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2'     )
              _UpperCAmelCase     =     2
              _UpperCAmelCase     =     2



              class 	lowerCAmelCase					(      tf.Module							):






                     def __init__(   self    ,  a__     ):
                            super(a__    ,  self     ).__init__()
                            _UpperCAmelCase     =     model






                     @tf.function(
                         input_signature=(
                             tf.TensorSpec((None, input_length)    ,  tf.intaa    ,  name='input_ids'     ),
                             tf.TensorSpec((None, input_length)    ,  tf.intaa    ,  name='attention_mask'     ),
                         )    ,  jit_compile=a__    ,  )
                     def 							__A (   self    ,  a__    ,  a__     ):
                            _UpperCAmelCase     =     self.model.generate(
                                input_ids=a__    ,  attention_mask=a__    ,  max_new_tokens=a__    ,  return_dict_in_generate=a__    ,  )
                            return {"sequences": outputs["sequences"]}


              _UpperCAmelCase     =     [[2, 0], [1_02, 1_03]]
              _UpperCAmelCase     =     [[1, 0], [1, 1]]
              _UpperCAmelCase     =     DummyModel(model=a__     )
              with tempfile.TemporaryDirectory() as tmp_dir:
                     tf.saved_model.save(a__    ,  a__    ,  signatures={'serving_default': dummy_model.serving}     )
                     _UpperCAmelCase     =     tf.saved_model.load(a__     ).signatures['serving_default']
                     for batch_size in range(1    ,  len(a__     ) + 1     ):
                            _UpperCAmelCase     =     {
                                'input_ids': tf.constant(dummy_input_ids[:batch_size]     ),
                                'attention_mask': tf.constant(dummy_attention_masks[:batch_size]     ),
                            }
                            _UpperCAmelCase     =     serving_func(**a__     )['sequences']
                            _UpperCAmelCase     =     test_model.generate(**a__    ,  max_new_tokens=a__     )
                            tf.debugging.assert_equal(a__    ,  a__     )






       @slow
       def 							__A (   self     ):
              # TF-only test: tf.saved_model export
              _UpperCAmelCase     =     TFAutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2'     )
              _UpperCAmelCase     =     1
              _UpperCAmelCase     =     2



              class 	lowerCAmelCase					(      tf.Module							):






                     def __init__(   self    ,  a__     ):
                            super(a__    ,  self     ).__init__()
                            _UpperCAmelCase     =     model






                     @tf.function(
                         input_signature=(
                             tf.TensorSpec((batch_size, None)    ,  tf.intaa    ,  name='input_ids'     ),
                             tf.TensorSpec((batch_size, None)    ,  tf.intaa    ,  name='attention_mask'     ),
                         )    ,  jit_compile=a__    ,  )
                     def 							__A (   self    ,  a__    ,  a__     ):
                            _UpperCAmelCase     =     self.model.generate(
                                input_ids=a__    ,  attention_mask=a__    ,  max_new_tokens=a__    ,  return_dict_in_generate=a__    ,  )
                            return {"sequences": outputs["sequences"]}


              _UpperCAmelCase     =     [[2], [1_02, 1_03]]
              _UpperCAmelCase     =     [[1], [1, 1]]
              _UpperCAmelCase     =     DummyModel(model=a__     )
              with tempfile.TemporaryDirectory() as tmp_dir:
                     tf.saved_model.save(a__    ,  a__    ,  signatures={'serving_default': dummy_model.serving}     )
                     _UpperCAmelCase     =     tf.saved_model.load(a__     ).signatures['serving_default']
                     for input_row in range(len(a__     )     ):
                            _UpperCAmelCase     =     {
                                'input_ids': tf.constant([dummy_input_ids[input_row]]     ),
                                'attention_mask': tf.constant([dummy_attention_masks[input_row]]     ),
                            }
                            _UpperCAmelCase     =     serving_func(**a__     )['sequences']
                            _UpperCAmelCase     =     test_model.generate(**a__    ,  max_new_tokens=a__     )
                            tf.debugging.assert_equal(a__    ,  a__     )






       @slow
       @require_tensorflow_text
       def 							__A (   self     ):
              # TF-only test: tf.saved_model export
              with tempfile.TemporaryDirectory() as tmp_dir:
                     # file needed to load the TF tokenizer
                     hf_hub_download(repo_id='google/flan-t5-small'    ,  filename='spiece.model'    ,  local_dir=a__     )



                     class 	lowerCAmelCase					(      tf.keras.layers.Layer							):






                            def __init__(   self     ):
                                   super().__init__()
                                   _UpperCAmelCase     =     text.SentencepieceTokenizer(
                                       model=tf.io.gfile.GFile(os.path.join(a__    ,  'spiece.model'     )    ,  'rb'     ).read()     )
                                   _UpperCAmelCase     =     TFAutoModelForSeqaSeqLM.from_pretrained('hf-internal-testing/tiny-random-t5'     )






                            def 							__A (   self    ,  a__    ,  *a__    ,  **a__     ):
                                   _UpperCAmelCase     =     self.tokenizer.tokenize(a__     )
                                   _UpperCAmelCase    ,							_UpperCAmelCase     =     text.pad_model_inputs(
                                       a__    ,  max_seq_length=64    ,  pad_value=self.model.config.pad_token_id     )
                                   _UpperCAmelCase     =     self.model.generate(input_ids=a__    ,  attention_mask=a__     )
                                   return self.tokenizer.detokenize(a__     )


                     _UpperCAmelCase     =     CompleteSentenceTransformer()
                     _UpperCAmelCase     =     tf.keras.layers.Input(shape=(1,)    ,  dtype=tf.string    ,  name='inputs'     )
                     _UpperCAmelCase     =     complete_model(a__     )
                     _UpperCAmelCase     =     tf.keras.Model(a__    ,  a__     )
                     keras_model.save(a__     )






       def 							__A (   self     ):
              # Has PT equivalent: this test relies on random sampling
              _UpperCAmelCase     =     {
                  'do_sample': True,
                  'num_beams': 1,
                  'top_p': 0.7,
                  'top_k': 10,
                  'temperature': 0.7,
              }
              _UpperCAmelCase     =     14

              _UpperCAmelCase     =     AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-gpt2'     )
              _UpperCAmelCase     =     'Hello, my dog is cute and'
              _UpperCAmelCase     =     tokenizer(a__    ,  return_tensors='tf'     )
              _UpperCAmelCase     =     TFAutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2'     )

              _UpperCAmelCase     =     6_38
              # forces the generation to happen on CPU, to avoid GPU-related quirks
              with tf.device(':/CPU:0'     ):
                     tf.random.set_seed(0     )
                     _UpperCAmelCase     =     model.generate(**a__    ,  eos_token_id=a__    ,  **a__     )
              self.assertTrue(expectation == len(generated_tokens[0]     )     )

              _UpperCAmelCase     =     [6_38, 1_98]
              with tf.device(':/CPU:0'     ):
                     tf.random.set_seed(0     )
                     _UpperCAmelCase     =     model.generate(**a__    ,  eos_token_id=a__    ,  **a__     )
              self.assertTrue(expectation == len(generated_tokens[0]     )     )






       def 							__A (   self     ):
              # Has PT equivalent: ample use of framework-specific code
              _UpperCAmelCase     =     AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-bart'     )
              _UpperCAmelCase     =     'Hugging Face is a technology company based in New York and Paris.'
              _UpperCAmelCase     =     bart_tokenizer(a__    ,  return_tensors='tf'     ).input_ids
              _UpperCAmelCase     =     TFBartForConditionalGeneration.from_pretrained('hf-internal-testing/tiny-random-bart'     )
              _UpperCAmelCase     =     bart_model.generate(a__     ).numpy()



              class 	lowerCAmelCase					(      snake_case							):






                     def 							__A (   self    ,  a__    ,  a__=None    ,  **a__     ):
                            return super().call(a__    ,  **a__     )


              _UpperCAmelCase     =     FakeBart.from_pretrained('hf-internal-testing/tiny-random-bart'     )
              _UpperCAmelCase     =     bart_model.generate(a__    ,  foo='bar'     ).numpy()
              self.assertTrue(np.array_equal(a__    ,  a__     )     )



              class 	lowerCAmelCase					(      bart_model.model.encoder.__class__							):






                     def 							__A (   self    ,  a__    ,  **a__     ):
                            return super().call(a__    ,  **a__     )


              _UpperCAmelCase     =     FakeEncoder(bart_model.config    ,  bart_model.model.shared     )
              _UpperCAmelCase     =     fake_encoder

              # Normal generation still works (the output will be different because the encoder weights are different)
              _UpperCAmelCase     =     bart_model.generate(a__     ).numpy()
              with self.assertRaises(a__     ):
                     # FakeEncoder.call() accepts **kwargs -> no filtering -> value error due to unexpected input "foo"
                     bart_model.generate(a__    ,  foo='bar'     )






