






'''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 							__SCREAMING_SNAKE_CASE       (							unittest.TestCase     ):




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



         '''simple docstring'''
         __a				     =    tf.convert_to_tensor(
             [
                 [
                     8.2220991,  # 3rd highest value; idx. 0
                     -0.5620044,
                     5.23229752,
                     4.0386393,
                     -6.8798378,
                     -0.54785802,
                     -3.2012153,
                     2.92777176,
                     1.88171953,
                     7.35341276,  # 5th highest value; idx. 9
                     8.43207833,  # 2nd highest value; idx. 10
                     -9.85711836,
                     -5.96209236,
                     -1.13039161,
                     -7.1115294,
                     -0.8369633,
                     -5.3186408,
                     7.06427407,
                     0.81369344,
                     -0.82023817,
                     -5.9179796,
                     0.58813443,
                     -6.99778438,
                     4.71551189,
                     -0.18771637,
                     7.44020759,  # 4th highest value; idx. 25
                     9.38450987,  # 1st highest value; idx. 26
                     2.12662941,
                     -9.32562038,
                     2.35652522,
                 ],  # cummulative prob of 5 highest values <= 0.6
                 [
                     0.58425518,
                     4.53139238,
                     -5.57510464,
                     -6.28030699,
                     -7.19529503,
                     -4.02122551,
                     1.39337037,
                     -6.06707057,
                     1.59480517,
                     -9.643119,
                     0.03907799,
                     0.67231762,
                     -8.88206726,
                     6.27115922,  # 4th highest value; idx. 13
                     2.28520723,
                     4.82767506,
                     4.30421368,
                     8.8275313,  # 2nd highest value; idx. 17
                     5.44029958,  # 5th highest value; idx. 18
                     -4.4735794,
                     7.38579536,  # 3rd highest value; idx. 20
                     -2.91051663,
                     2.61946077,
                     -2.5674762,
                     -9.48959302,
                     -4.02922645,
                     -1.35416918,
                     9.67702323,  # 1st highest value; idx. 27
                     -5.89478553,
                     1.85370467,
                 ],  # cummulative prob of 5 highest values <= 0.6
             ] ,   dtype=tf.floataa ,   )

         __a				     =    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

         __a				     =    tf.convert_to_tensor(
             [8.222099, 7.3534126, 8.432078, 7.4402075, 9.38451, 6.271159, 8.827531, 5.4402995, 7.3857956, 9.677023] ,   dtype=tf.floataa ,   )  # expected non filtered values as noted above

         __a				     =    tf_top_k_top_p_filtering(lowerCamelCase ,   top_k=10 ,   top_p=0.6 ,   min_tokens_to_keep=4      )

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

         tf.debugging.assert_near(lowerCamelCase ,   lowerCamelCase ,   rtol=1e-12      )
         tf.debugging.assert_equal(lowerCamelCase ,   lowerCamelCase      )







@require_tf
class 							__SCREAMING_SNAKE_CASE       (							unittest.TestCase      ,      _lowerCAmelCase     ):
      # setting framework_dependent_parameters needs to be gated, just like its contents' imports
      if is_tf_available():
            __a  					={
                "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  __UpperCamelCase							(  self      )							->Tuple:



         '''simple docstring'''
         # TF-only test: tf.saved_model export
         __a				     =    TFAutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2'      )
         __a				     =    2
         __a				     =    2
         class 							__SCREAMING_SNAKE_CASE       (							tf.Module     ):




               def __init__(  self ,   lowerCamelCase      )							->Union[str, Any]:



                  '''simple docstring'''
                  super(lowerCamelCase ,   self      ).__init__()
                  __a				     =    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=lowerCamelCase ,   )
               def  __UpperCamelCase							(  self ,   lowerCamelCase ,   lowerCamelCase      )							->Any:



                  '''simple docstring'''
                  __a				     =    self.model.generate(
                      input_ids=lowerCamelCase ,   attention_mask=lowerCamelCase ,   max_new_tokens=lowerCamelCase ,   return_dict_in_generate=lowerCamelCase ,   )
                  return {"sequences": outputs["sequences"]}







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




      @slow
      def  __UpperCamelCase							(  self      )							->List[str]:



         '''simple docstring'''
         # TF-only test: tf.saved_model export
         __a				     =    TFAutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2'      )
         __a				     =    1
         __a				     =    2
         class 							__SCREAMING_SNAKE_CASE       (							tf.Module     ):




               def __init__(  self ,   lowerCamelCase      )							->Dict:



                  '''simple docstring'''
                  super(lowerCamelCase ,   self      ).__init__()
                  __a				     =    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=lowerCamelCase ,   )
               def  __UpperCamelCase							(  self ,   lowerCamelCase ,   lowerCamelCase      )							->int:



                  '''simple docstring'''
                  __a				     =    self.model.generate(
                      input_ids=lowerCamelCase ,   attention_mask=lowerCamelCase ,   max_new_tokens=lowerCamelCase ,   return_dict_in_generate=lowerCamelCase ,   )
                  return {"sequences": outputs["sequences"]}







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




      @slow
      @require_tensorflow_text
      def  __UpperCamelCase							(  self      )							->List[str]:



         '''simple docstring'''
         # 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=lowerCamelCase      )
            class 							__SCREAMING_SNAKE_CASE       (							tf.keras.layers.Layer     ):




                  def __init__(  self      )							->Optional[int]:



                     '''simple docstring'''
                     super().__init__()
                     __a				     =    text.SentencepieceTokenizer(
                         model=tf.io.gfile.GFile(os.path.join(lowerCamelCase ,   'spiece.model'      ) ,   'rb'      ).read()      )
                     __a				     =    TFAutoModelForSeqaSeqLM.from_pretrained('hf-internal-testing/tiny-random-t5'      )






                  def  __UpperCamelCase							(  self ,   lowerCamelCase ,   *lowerCamelCase ,   **lowerCamelCase      )							->Union[str, Any]:



                     '''simple docstring'''
                     __a				     =    self.tokenizer.tokenize(lowerCamelCase      )
                     __a     ,			__a				     =    text.pad_model_inputs(
                         lowerCamelCase ,   max_seq_length=64 ,   pad_value=self.model.config.pad_token_id      )
                     __a				     =    self.model.generate(input_ids=lowerCamelCase ,   attention_mask=lowerCamelCase      )
                     return self.tokenizer.detokenize(lowerCamelCase      )







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




      def  __UpperCamelCase							(  self      )							->List[Any]:



         '''simple docstring'''
         # Has PT equivalent: this test relies on random sampling
         __a				     =    {
             'do_sample': True,
             'num_beams': 1,
             'top_p': 0.7,
             'top_k': 10,
             'temperature': 0.7,
         }
         __a				     =    14

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

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

         __a				     =    [638, 198]
         with tf.device(':/CPU:0'      ):
            tf.random.set_seed(0      )
            __a				     =    model.generate(**lowerCamelCase ,   eos_token_id=lowerCamelCase ,   **lowerCamelCase      )
         self.assertTrue(expectation == len(generated_tokens[0]      )      )






      def  __UpperCamelCase							(  self      )							->Optional[Any]:



         '''simple docstring'''
         # Has PT equivalent: ample use of framework-specific code
         __a				     =    AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-bart'      )
         __a				     =    'Hugging Face is a technology company based in New York and Paris.'
         __a				     =    bart_tokenizer(lowerCamelCase ,   return_tensors='tf'      ).input_ids
         __a				     =    TFBartForConditionalGeneration.from_pretrained('hf-internal-testing/tiny-random-bart'      )
         __a				     =    bart_model.generate(lowerCamelCase      ).numpy()
         class 							__SCREAMING_SNAKE_CASE       (							_lowerCAmelCase     ):




               def  __UpperCamelCase							(  self ,   lowerCamelCase ,   lowerCamelCase=None ,   **lowerCamelCase      )							->Dict:



                  '''simple docstring'''
                  return super().call(lowerCamelCase ,   **lowerCamelCase      )







         __a				     =    FakeBart.from_pretrained('hf-internal-testing/tiny-random-bart'      )
         __a				     =    bart_model.generate(lowerCamelCase ,   foo='bar'      ).numpy()
         self.assertTrue(np.array_equal(lowerCamelCase ,   lowerCamelCase      )      )
         class 							__SCREAMING_SNAKE_CASE       (							bart_model.model.encoder.__class__     ):




               def  __UpperCamelCase							(  self ,   lowerCamelCase ,   **lowerCamelCase      )							->Any:



                  '''simple docstring'''
                  return super().call(lowerCamelCase ,   **lowerCamelCase      )







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

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







'''simple docstring'''



import inspect
from typing import List, Optional, Tuple, Union

import numpy as np
import PIL
import torch
import torch.utils.checkpoint

from ...models import UNetaDModel, VQModel
from ...schedulers import (
    DDIMScheduler,
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
)
from ...utils import PIL_INTERPOLATION, randn_tensor
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput



def    __UpperCAmelCase	(				SCREAMING_SNAKE_CASE__:	List[Any]    )  ->			Dict:

   """simple docstring"""
   __a     ,			__a				     =    image.size
   __a     ,			__a				     =    (x - x % 32 for x in (w, h))  # resize to integer multiple of 32
   __a				     =    image.resize((w, h),					resample=PIL_INTERPOLATION['lanczos']    )
   __a				     =    np.array(SCREAMING_SNAKE_CASE__    ).astype(np.floataa    ) / 2_5_5.0
   __a				     =    image[None].transpose(0,					3,					1,					2    )
   __a				     =    torch.from_numpy(SCREAMING_SNAKE_CASE__    )
   return 2.0 * image - 1.0



class 							__SCREAMING_SNAKE_CASE       (							_lowerCAmelCase     ):




      def __init__(  self ,   lowerCamelCase ,   lowerCamelCase ,   lowerCamelCase ,   )							->Dict:



         '''simple docstring'''
         super().__init__()
         self.register_modules(vqvae=lowerCamelCase ,   unet=lowerCamelCase ,   scheduler=lowerCamelCase      )






      @torch.no_grad()
      def __call__(  self ,   lowerCamelCase = None ,   lowerCamelCase = 1 ,   lowerCamelCase = 100 ,   lowerCamelCase = 0.0 ,   lowerCamelCase = None ,   lowerCamelCase = "pil" ,   lowerCamelCase = True ,   )							->Union[Tuple, ImagePipelineOutput]:



         '''simple docstring'''
         if isinstance(lowerCamelCase ,   PIL.Image.Image      ):
            __a				     =    1
         elif isinstance(lowerCamelCase ,   torch.Tensor      ):
            __a				     =    image.shape[0]
         else:
            raise ValueError(F"""`image` has to be of type `PIL.Image.Image` or `torch.Tensor` but is {type(lowerCamelCase      )}"""      )

         if isinstance(lowerCamelCase ,   PIL.Image.Image      ):
            __a				     =    preprocess(lowerCamelCase      )

         __a     ,			__a				     =    image.shape[-2:]

         # in_channels should be 6: 3 for latents, 3 for low resolution image
         __a				     =    (batch_size, self.unet.config.in_channels // 2, height, width)
         __a				     =    next(self.unet.parameters()      ).dtype

         __a				     =    randn_tensor(lowerCamelCase ,   generator=lowerCamelCase ,   device=self.device ,   dtype=lowerCamelCase      )

         __a				     =    image.to(device=self.device ,   dtype=lowerCamelCase      )

         # set timesteps and move to the correct device
         self.scheduler.set_timesteps(lowerCamelCase ,   device=self.device      )
         __a				     =    self.scheduler.timesteps

         # scale the initial noise by the standard deviation required by the scheduler
         __a				     =    latents * self.scheduler.init_noise_sigma

         # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature.
         # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
         # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
         # and should be between [0, 1]
         __a				     =    'eta' in set(inspect.signature(self.scheduler.step      ).parameters.keys()      )
         __a				     =    {}
         if accepts_eta:
            __a				     =    eta

         for t in self.progress_bar(lowerCamelCase      ):
            # concat latents and low resolution image in the channel dimension.
            __a				     =    torch.cat([latents, image] ,   dim=1      )
            __a				     =    self.scheduler.scale_model_input(lowerCamelCase ,   lowerCamelCase      )
            # predict the noise residual
            __a				     =    self.unet(lowerCamelCase ,   lowerCamelCase      ).sample
            # compute the previous noisy sample x_t -> x_t-1
            __a				     =    self.scheduler.step(lowerCamelCase ,   lowerCamelCase ,   lowerCamelCase ,   **lowerCamelCase      ).prev_sample

         # decode the image latents with the VQVAE
         __a				     =    self.vqvae.decode(lowerCamelCase      ).sample
         __a				     =    torch.clamp(lowerCamelCase ,   -1.0 ,   1.0      )
         __a				     =    image / 2 + 0.5
         __a				     =    image.cpu().permute(0 ,   2 ,   3 ,   1      ).numpy()

         if output_type == "pil":
            __a				     =    self.numpy_to_pil(lowerCamelCase      )

         if not return_dict:
            return (image,)

         return ImagePipelineOutput(images=lowerCamelCase      )