





import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import BlenderbotSmallConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

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
      a__      :     Dict         =		'platform'

      import jax
      import jax.numpy as jnp

      from transformers.models.blenderbot_small.modeling_flax_blenderbot_small import (
          FlaxBlenderbotSmallForConditionalGeneration,
          FlaxBlenderbotSmallModel,
          shift_tokens_right,
      )

def      UpperCAmelCase_    (    _UpperCAmelCase							:str  ,						_UpperCAmelCase							:int  ,						_UpperCAmelCase							:Optional[int]=None  ,						_UpperCAmelCase							:Tuple=None  ,						_UpperCAmelCase							:str=None  ,						_UpperCAmelCase							:Optional[int]=None  ,						_UpperCAmelCase							:Union[str, Any]=None  ,						_UpperCAmelCase							:Optional[int]=None  ,						)	->				List[Any]:
     '''simple docstring'''


     if attention_mask is None:
          A_          =       np.where(input_ids != config.pad_token_id  ,						1  ,						0				)
     if decoder_attention_mask is None:
          A_          =       np.where(decoder_input_ids != config.pad_token_id  ,						1  ,						0				)
     if head_mask is None:
          A_          =       np.ones((config.encoder_layers, config.encoder_attention_heads)				)
     if decoder_head_mask is None:
          A_          =       np.ones((config.decoder_layers, config.decoder_attention_heads)				)
     if cross_attn_head_mask is None:
          A_          =       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 ,__snake_case ,__snake_case=1_3 ,__snake_case=7 ,__snake_case=True ,__snake_case=False ,__snake_case=9_9 ,__snake_case=1_6 ,__snake_case=2 ,__snake_case=4 ,__snake_case=4 ,__snake_case="gelu" ,__snake_case=0.1 ,__snake_case=0.1 ,__snake_case=3_2 ,__snake_case=2 ,__snake_case=1 ,__snake_case=0 ,__snake_case=0.02 ,):



       """simple docstring"""





       A_          =       parent
       A_          =       batch_size
       A_          =       seq_length
       A_          =       is_training
       A_          =       use_labels
       A_          =       vocab_size
       A_          =       hidden_size
       A_          =       num_hidden_layers
       A_          =       num_attention_heads
       A_          =       intermediate_size
       A_          =       hidden_act
       A_          =       hidden_dropout_prob
       A_          =       attention_probs_dropout_prob
       A_          =       max_position_embeddings
       A_          =       eos_token_id
       A_          =       pad_token_id
       A_          =       bos_token_id
       A_          =       initializer_range


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_          =       np.clip(ids_tensor([self.batch_size, self.seq_length - 1] ,self.vocab_size ) ,3 ,self.vocab_size )
       A_          =       np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1) ,dtype=np.intaa )) ,-1 )

       A_          =       shift_tokens_right(__snake_case ,1 ,2 )

       A_          =       BlenderbotSmallConfig(
           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=__snake_case ,)
       A_          =       prepare_blenderbot_inputs_dict(__snake_case ,__snake_case ,__snake_case )
       return config, inputs_dict


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_   ,    A_          =       self.prepare_config_and_inputs()
       return config, inputs_dict


  def 						__UpperCAmelCase     (							self ,__snake_case ,__snake_case ,__snake_case ):



       """simple docstring"""





       A_          =       2_0
       A_          =       model_class_name(__snake_case )

       A_          =       model.encode(inputs_dict['''input_ids'''] )

       A_   ,    A_          =       (
           inputs_dict['''decoder_input_ids'''],
           inputs_dict['''decoder_attention_mask'''],
       )

       A_          =       model.init_cache(decoder_input_ids.shape[0] ,__snake_case ,__snake_case )
       A_          =       jnp.ones((decoder_input_ids.shape[0], max_decoder_length) ,dtype='''i4''' )

       A_          =       jnp.broadcast_to(
           jnp.arange(decoder_input_ids.shape[-1] - 1 )[None, :] ,(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1) ,)
       A_          =       model.decode(
           decoder_input_ids[:, :-1] ,__snake_case ,decoder_attention_mask=__snake_case ,past_key_values=__snake_case ,decoder_position_ids=__snake_case ,)

       A_          =       jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]] ,dtype='''i4''' )
       A_          =       model.decode(
           decoder_input_ids[:, -1:] ,__snake_case ,decoder_attention_mask=__snake_case ,past_key_values=outputs_cache.past_key_values ,decoder_position_ids=__snake_case ,)

       A_          =       model.decode(__snake_case ,__snake_case )

       A_          =       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 						__UpperCAmelCase     (							self ,__snake_case ,__snake_case ,__snake_case ):



       """simple docstring"""





       A_          =       2_0
       A_          =       model_class_name(__snake_case )

       A_          =       model.encode(inputs_dict['''input_ids'''] )

       A_   ,    A_          =       (
           inputs_dict['''decoder_input_ids'''],
           inputs_dict['''decoder_attention_mask'''],
       )

       A_          =       jnp.concatenate(
           [
               decoder_attention_mask,
               jnp.zeros((decoder_attention_mask.shape[0], max_decoder_length - decoder_attention_mask.shape[1]) ),
           ] ,axis=-1 ,)

       A_          =       model.init_cache(decoder_input_ids.shape[0] ,__snake_case ,__snake_case )
       A_          =       jnp.broadcast_to(
           jnp.arange(decoder_input_ids.shape[-1] - 1 )[None, :] ,(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1) ,)

       A_          =       model.decode(
           decoder_input_ids[:, :-1] ,__snake_case ,decoder_attention_mask=__snake_case ,past_key_values=__snake_case ,decoder_position_ids=__snake_case ,)
       A_          =       jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]] ,dtype='''i4''' )
       A_          =       model.decode(
           decoder_input_ids[:, -1:] ,__snake_case ,past_key_values=outputs_cache.past_key_values ,decoder_attention_mask=__snake_case ,decoder_position_ids=__snake_case ,)

       A_          =       model.decode(__snake_case ,__snake_case ,decoder_attention_mask=__snake_case )

       A_          =       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_				( unittest.TestCase ):
  __lowerCAmelCase   :						Optional[int]							=   99


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_          =       np.array(
           [
               [7_1, 8_2, 1_8, 3_3, 4_6, 9_1, 2],
               [6_8, 3_4, 2_6, 5_8, 3_0, 8_2, 2],
               [5, 9_7, 1_7, 3_9, 9_4, 4_0, 2],
               [7_6, 8_3, 9_4, 2_5, 7_0, 7_8, 2],
               [8_7, 5_9, 4_1, 3_5, 4_8, 6_6, 2],
               [5_5, 1_3, 1_6, 5_8, 5, 2, 1],  # note padding
               [6_4, 2_7, 3_1, 5_1, 1_2, 7_5, 2],
               [5_2, 6_4, 8_6, 1_7, 8_3, 3_9, 2],
               [4_8, 6_1, 9, 2_4, 7_1, 8_2, 2],
               [2_6, 1, 6_0, 4_8, 2_2, 1_3, 2],
               [2_1, 5, 6_2, 2_8, 1_4, 7_6, 2],
               [4_5, 9_8, 3_7, 8_6, 5_9, 4_8, 2],
               [7_0, 7_0, 5_0, 9, 2_8, 0, 2],
           ] ,dtype=np.intaa ,)

       A_          =       input_ids.shape[0]
       A_          =       BlenderbotSmallConfig(
           vocab_size=self.vocab_size ,d_model=2_4 ,encoder_layers=2 ,decoder_layers=2 ,encoder_attention_heads=2 ,decoder_attention_heads=2 ,encoder_ffn_dim=3_2 ,decoder_ffn_dim=3_2 ,max_position_embeddings=4_8 ,eos_token_id=2 ,pad_token_id=1 ,bos_token_id=0 ,)
       return config, input_ids, batch_size


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_   ,    A_   ,    A_          =       self._get_config_and_data()
       A_          =       FlaxBlenderbotSmallForConditionalGeneration(__snake_case )
       A_          =       lm_model(input_ids=__snake_case )
       A_          =       (batch_size, input_ids.shape[1], config.vocab_size)
       self.assertEqual(outputs['''logits'''].shape ,__snake_case )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_          =       BlenderbotSmallConfig(
           vocab_size=self.vocab_size ,d_model=1_4 ,encoder_layers=2 ,decoder_layers=2 ,encoder_attention_heads=2 ,decoder_attention_heads=2 ,encoder_ffn_dim=8 ,decoder_ffn_dim=8 ,max_position_embeddings=4_8 ,)
       A_          =       FlaxBlenderbotSmallForConditionalGeneration(__snake_case )
       A_          =       np.array([[7_1, 8_2, 1_8, 3_3, 4_6, 9_1, 2], [6_8, 3_4, 2_6, 5_8, 3_0, 2, 1]] ,dtype=np.intaa )
       A_          =       np.array([[8_2, 7_1, 8_2, 1_8, 2], [5_8, 6_8, 2, 1, 1]] ,dtype=np.intaa )
       A_          =       lm_model(input_ids=__snake_case ,decoder_input_ids=__snake_case )
       A_          =       (*summary.shape, config.vocab_size)
       self.assertEqual(outputs['''logits'''].shape ,__snake_case )




  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_          =       np.array([[7_1, 8_2, 1_8, 3_3, 2, 1, 1], [6_8, 3_4, 2_6, 5_8, 3_0, 8_2, 2]] ,dtype=np.intaa )
       A_          =       shift_tokens_right(__snake_case ,1 ,2 )
       A_          =       np.equal(__snake_case ,1 ).astype(np.floataa ).sum()
       A_          =       np.equal(__snake_case ,1 ).astype(np.floataa ).sum()
       self.assertEqual(shifted.shape ,input_ids.shape )
       self.assertEqual(__snake_case ,n_pad_before - 1 )
       self.assertTrue(np.equal(shifted[:, 0] ,2 ).all() )





@require_flax
class UpperCAmelCase_				( _SCREAMING_SNAKE_CASE     ,		unittest.TestCase     ,		_SCREAMING_SNAKE_CASE ):
  __lowerCAmelCase   :						Tuple							=   True
  __lowerCAmelCase   :						Union[str, Any]							=   (
      (
          FlaxBlenderbotSmallModel,
          FlaxBlenderbotSmallForConditionalGeneration,
      )
      if is_flax_available()
      else ()
  )
  __lowerCAmelCase   :						int							=   (FlaxBlenderbotSmallForConditionalGeneration,) if is_flax_available() else ()


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_          =       FlaxBlenderbotSmallModelTester(self )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_   ,    A_          =       self.model_tester.prepare_config_and_inputs()
       for model_class in self.all_model_classes:
            self.model_tester.check_use_cache_forward(__snake_case ,__snake_case ,__snake_case )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_   ,    A_          =       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(__snake_case ,__snake_case ,__snake_case )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_   ,    A_          =       self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
            with self.subTest(model_class.__name__ ):
                 A_          =       self._prepare_for_class(__snake_case ,__snake_case )
                 A_          =       model_class(__snake_case )

                 @jax.jit
                 def encode_jitted(__snake_case ,__snake_case=None ,**__snake_case ):
                      return model.encode(input_ids=__snake_case ,attention_mask=__snake_case )

                 with self.subTest('''JIT Enabled''' ):
                      A_          =       encode_jitted(**__snake_case ).to_tuple()

                 with self.subTest('''JIT Disabled''' ):
                      with jax.disable_jit():
                           A_          =       encode_jitted(**__snake_case ).to_tuple()

                 self.assertEqual(len(__snake_case ) ,len(__snake_case ) )
                 for jitted_output, output in zip(__snake_case ,__snake_case ):
                      self.assertEqual(jitted_output.shape ,output.shape )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_   ,    A_          =       self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
            with self.subTest(model_class.__name__ ):
                 A_          =       model_class(__snake_case )
                 A_          =       model.encode(inputs_dict['''input_ids'''] ,inputs_dict['''attention_mask'''] )

                 A_          =       {
                     '''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(__snake_case ,__snake_case ,__snake_case ):
                      return model.decode(
                          decoder_input_ids=__snake_case ,decoder_attention_mask=__snake_case ,encoder_outputs=__snake_case ,)

                 with self.subTest('''JIT Enabled''' ):
                      A_          =       decode_jitted(**__snake_case ).to_tuple()

                 with self.subTest('''JIT Disabled''' ):
                      with jax.disable_jit():
                           A_          =       decode_jitted(**__snake_case ).to_tuple()

                 self.assertEqual(len(__snake_case ) ,len(__snake_case ) )
                 for jitted_output, output in zip(__snake_case ,__snake_case ):
                      self.assertEqual(jitted_output.shape ,output.shape )




  @slow
  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       for model_class_name in self.all_model_classes:
            A_          =       model_class_name.from_pretrained('''facebook/blenderbot_small-90M''' )
            # FlaxBlenderbotForSequenceClassification expects eos token in input_ids
            A_          =       np.ones((1, 1) ) * model.config.eos_token_id
            A_          =       model(__snake_case )
            self.assertIsNotNone(__snake_case )










import unittest

import numpy as np
import torch

from diffusers import VersatileDiffusionImageVariationPipeline
from diffusers.utils.testing_utils import load_image, require_torch_gpu, slow, torch_device


a__      :     Optional[int]         =		False





class UpperCAmelCase_				( unittest.TestCase ):
  pass





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


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_          =       VersatileDiffusionImageVariationPipeline.from_pretrained('''shi-labs/versatile-diffusion''' )
       pipe.to(__snake_case )
       pipe.set_progress_bar_config(disable=__snake_case )

       A_          =       load_image(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/versatile_diffusion/benz.jpg''' )
       A_          =       torch.manual_seed(0 )
       A_          =       pipe(
           image=__snake_case ,generator=__snake_case ,guidance_scale=7.5 ,num_inference_steps=5_0 ,output_type='''numpy''' ,).images

       A_          =       image[0, 2_5_3:2_5_6, 2_5_3:2_5_6, -1]

       assert image.shape == (1, 5_1_2, 5_1_2, 3)
       A_          =       np.array([0.0441, 0.0469, 0.0507, 0.0575, 0.0632, 0.0650, 0.0865, 0.0909, 0.0945] )

       assert np.abs(image_slice.flatten() - expected_slice ).max() < 1E-2



