





from __future__ import annotations

import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_sentencepiece, require_tf, require_tokenizers, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
      import numpy as np
      import tensorflow as tf

      from transformers import (
          TF_FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
          FlaubertConfig,
          TFFlaubertForMultipleChoice,
          TFFlaubertForQuestionAnsweringSimple,
          TFFlaubertForSequenceClassification,
          TFFlaubertForTokenClassification,
          TFFlaubertModel,
          TFFlaubertWithLMHeadModel,
      )
class        SCREAMING_SNAKE_CASE     :


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





     """simple docstring"""







     lowercase__			    =						parent
     lowercase__			    =						13
     lowercase__			    =						7
     lowercase__			    =						True
     lowercase__			    =						True
     lowercase__			    =						True
     lowercase__			    =						True
     lowercase__			    =						True
     lowercase__			    =						False
     lowercase__			    =						False
     lowercase__			    =						False
     lowercase__			    =						2
     lowercase__			    =						99
     lowercase__			    =						0
     lowercase__			    =						32
     lowercase__			    =						2
     lowercase__			    =						4
     lowercase__			    =						0.1
     lowercase__			    =						0.1
     lowercase__			    =						512
     lowercase__			    =						16
     lowercase__			    =						2
     lowercase__			    =						0.02
     lowercase__			    =						3
     lowercase__			    =						4
     lowercase__			    =						'last'
     lowercase__			    =						True
     lowercase__			    =						None
     lowercase__			    =						0


   def  SCREAMING_SNAKE_CASE_							(					self    :       Tuple )->     List[Any]:





     """simple docstring"""







     lowercase__			    =						ids_tensor([self.batch_size, self.seq_length]    ,  self.vocab_size )
     lowercase__			    =						random_attention_mask([self.batch_size, self.seq_length]    ,  dtype=tf.floataa )

     lowercase__			    =						None
     if self.use_input_lengths:
       lowercase__			    =						(
           ids_tensor([self.batch_size]    ,  vocab_size=2 ) + self.seq_length - 2
       )  # small variation of seq_length

     lowercase__			    =						None
     if self.use_token_type_ids:
       lowercase__			    =						ids_tensor([self.batch_size, self.seq_length]    ,  self.n_langs )

     lowercase__			    =						None
     lowercase__			    =						None
     lowercase__			    =						None
     if self.use_labels:
       lowercase__			    =						ids_tensor([self.batch_size]    ,  self.type_sequence_label_size )
       lowercase__			    =						ids_tensor([self.batch_size, self.seq_length]    ,  self.num_labels )
       lowercase__			    =						ids_tensor([self.batch_size]    ,  2    ,  dtype=tf.floataa )
       lowercase__			    =						ids_tensor([self.batch_size]    ,  self.num_choices )

     lowercase__			    =						FlaubertConfig(
         vocab_size=self.vocab_size    ,  n_special=self.n_special    ,  emb_dim=self.hidden_size    ,  n_layers=self.num_hidden_layers    ,  n_heads=self.num_attention_heads    ,  dropout=self.hidden_dropout_prob    ,  attention_dropout=self.attention_probs_dropout_prob    ,  gelu_activation=self.gelu_activation    ,  sinusoidal_embeddings=self.sinusoidal_embeddings    ,  asm=self.asm    ,  causal=self.causal    ,  n_langs=self.n_langs    ,  max_position_embeddings=self.max_position_embeddings    ,  initializer_range=self.initializer_range    ,  summary_type=self.summary_type    ,  use_proj=self.use_proj    ,  bos_token_id=self.bos_token_id    ,  )

     return (
         config,
         input_ids,
         token_type_ids,
         input_lengths,
         sequence_labels,
         token_labels,
         is_impossible_labels,
         choice_labels,
         input_mask,
     )


   def  SCREAMING_SNAKE_CASE_							(					self    :       int    ,  a    :       str    ,  a    :       Tuple    ,  a    :       Optional[Any]    ,  a    :       Dict    ,  a    :       Union[str, Any]    ,  a    :       Any    ,  a    :       Union[str, Any]    ,  a    :       int    ,  a    :       Optional[int]    ,  )->     Optional[int]:





     """simple docstring"""







     lowercase__			    =						TFFlaubertModel(config=a )
     lowercase__			    =						{'input_ids': input_ids, 'lengths': input_lengths, 'langs': token_type_ids}
     lowercase__			    =						model(a )

     lowercase__			    =						[input_ids, input_mask]
     lowercase__			    =						model(a )
     self.parent.assertEqual(result.last_hidden_state.shape    ,  (self.batch_size, self.seq_length, self.hidden_size) )


   def  SCREAMING_SNAKE_CASE_							(					self    :       str    ,  a    :       str    ,  a    :       Tuple    ,  a    :       List[str]    ,  a    :       List[Any]    ,  a    :       Optional[Any]    ,  a    :       Dict    ,  a    :       Optional[int]    ,  a    :       Optional[Any]    ,  a    :       Union[str, Any]    ,  )->     int:





     """simple docstring"""







     lowercase__			    =						TFFlaubertWithLMHeadModel(a )

     lowercase__			    =						{'input_ids': input_ids, 'lengths': input_lengths, 'langs': token_type_ids}
     lowercase__			    =						model(a )

     self.parent.assertEqual(result.logits.shape    ,  (self.batch_size, self.seq_length, self.vocab_size) )


   def  SCREAMING_SNAKE_CASE_							(					self    :       str    ,  a    :       Optional[int]    ,  a    :       Tuple    ,  a    :       int    ,  a    :       Tuple    ,  a    :       Optional[Any]    ,  a    :       List[str]    ,  a    :       List[Any]    ,  a    :       Optional[int]    ,  a    :       int    ,  )->     List[str]:





     """simple docstring"""







     lowercase__			    =						TFFlaubertForQuestionAnsweringSimple(a )

     lowercase__			    =						{'input_ids': input_ids, 'lengths': input_lengths}

     lowercase__			    =						model(a )

     self.parent.assertEqual(result.start_logits.shape    ,  (self.batch_size, self.seq_length) )
     self.parent.assertEqual(result.end_logits.shape    ,  (self.batch_size, self.seq_length) )


   def  SCREAMING_SNAKE_CASE_							(					self    :       List[Any]    ,  a    :       int    ,  a    :       Tuple    ,  a    :       List[Any]    ,  a    :       List[str]    ,  a    :       str    ,  a    :       int    ,  a    :       str    ,  a    :       Tuple    ,  a    :       str    ,  )->     Dict:





     """simple docstring"""







     lowercase__			    =						TFFlaubertForSequenceClassification(a )

     lowercase__			    =						{'input_ids': input_ids, 'lengths': input_lengths}

     lowercase__			    =						model(a )

     self.parent.assertEqual(result.logits.shape    ,  (self.batch_size, self.type_sequence_label_size) )


   def  SCREAMING_SNAKE_CASE_							(					self    :       Union[str, Any]    ,  a    :       Tuple    ,  a    :       Dict    ,  a    :       Any    ,  a    :       List[str]    ,  a    :       Union[str, Any]    ,  a    :       int    ,  a    :       Optional[Any]    ,  a    :       Dict    ,  a    :       Tuple    ,  )->     Optional[Any]:





     """simple docstring"""







     lowercase__			    =						self.num_labels
     lowercase__			    =						TFFlaubertForTokenClassification(config=a )
     lowercase__			    =						{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
     lowercase__			    =						model(a )
     self.parent.assertEqual(result.logits.shape    ,  (self.batch_size, self.seq_length, self.num_labels) )


   def  SCREAMING_SNAKE_CASE_							(					self    :       Union[str, Any]    ,  a    :       Optional[int]    ,  a    :       Optional[Any]    ,  a    :       Tuple    ,  a    :       Tuple    ,  a    :       Any    ,  a    :       Union[str, Any]    ,  a    :       int    ,  a    :       Tuple    ,  a    :       str    ,  )->     List[str]:





     """simple docstring"""







     lowercase__			    =						self.num_choices
     lowercase__			    =						TFFlaubertForMultipleChoice(config=a )
     lowercase__			    =						tf.tile(tf.expand_dims(a    ,  1 )    ,  (1, self.num_choices, 1) )
     lowercase__			    =						tf.tile(tf.expand_dims(a    ,  1 )    ,  (1, self.num_choices, 1) )
     lowercase__			    =						tf.tile(tf.expand_dims(a    ,  1 )    ,  (1, self.num_choices, 1) )
     lowercase__			    =						{
         'input_ids': multiple_choice_inputs_ids,
         'attention_mask': multiple_choice_input_mask,
         'token_type_ids': multiple_choice_token_type_ids,
     }
     lowercase__			    =						model(a )
     self.parent.assertEqual(result.logits.shape    ,  (self.batch_size, self.num_choices) )


   def  SCREAMING_SNAKE_CASE_							(					self    :       Any )->     Dict:





     """simple docstring"""







     lowercase__			    =						self.prepare_config_and_inputs()
     (
         (
         lowercase__
     )    ,					(
         lowercase__
     )    ,					(
         lowercase__
     )    ,					(
         lowercase__
     )    ,					(
         lowercase__
     )    ,					(
         lowercase__
     )    ,					(
         lowercase__
     )    ,					(
         lowercase__
     )    ,					(
         lowercase__
     )    ,					
     )			    =						config_and_inputs
     lowercase__			    =						{
         'input_ids': input_ids,
         'token_type_ids': token_type_ids,
         'langs': token_type_ids,
         'lengths': input_lengths,
     }
     return config, inputs_dict






@require_tf
class        SCREAMING_SNAKE_CASE     (UpperCAmelCase    ,			UpperCAmelCase    ,			unittest.TestCase   ):
   _UpperCamelCase					:	List[Any]				      = (
       (
           TFFlaubertModel,
           TFFlaubertWithLMHeadModel,
           TFFlaubertForSequenceClassification,
           TFFlaubertForQuestionAnsweringSimple,
           TFFlaubertForTokenClassification,
           TFFlaubertForMultipleChoice,
       )
       if is_tf_available()
       else ()
   )
   _UpperCamelCase					:	List[str]				      = (
       (TFFlaubertWithLMHeadModel,) if is_tf_available() else ()
   )  # TODO (PVP): Check other models whether language generation is also applicable
   _UpperCamelCase					:	List[str]				      = (
       {
           'feature-extraction': TFFlaubertModel,
           'fill-mask': TFFlaubertWithLMHeadModel,
           'question-answering': TFFlaubertForQuestionAnsweringSimple,
           'text-classification': TFFlaubertForSequenceClassification,
           'token-classification': TFFlaubertForTokenClassification,
           'zero-shot': TFFlaubertForSequenceClassification,
       }
       if is_tf_available()
       else {}
   )
   _UpperCamelCase					:	Optional[Any]				      = False
   _UpperCamelCase					:	Union[str, Any]				      = False


   def  SCREAMING_SNAKE_CASE_							(					self    :       int    ,  a    :       List[str]    ,  a    :       Dict    ,  a    :       Optional[int]    ,  a    :       Union[str, Any]    ,  a    :       List[Any] )->     Optional[Any]:





     """simple docstring"""







     if (
         pipeline_test_casse_name == "QAPipelineTests"
         and tokenizer_name is not None
         and not tokenizer_name.endswith('Fast' )
     ):
       # `QAPipelineTests` fails for a few models when the slower tokenizer are used.
       # (The slower tokenizers were never used for pipeline tests before the pipeline testing rework)
       # TODO: check (and possibly fix) the `QAPipelineTests` with slower tokenizer
       return True

     return False


   def  SCREAMING_SNAKE_CASE_							(					self    :       Optional[Any] )->     int:





     """simple docstring"""







     lowercase__			    =						TFFlaubertModelTester(self )
     lowercase__			    =						ConfigTester(self    ,  config_class=a    ,  emb_dim=37 )


   def  SCREAMING_SNAKE_CASE_							(					self    :       List[str] )->     Any:





     """simple docstring"""







     self.config_tester.run_common_tests()


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





     """simple docstring"""







     lowercase__			    =						self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_flaubert_model(*a )


   def  SCREAMING_SNAKE_CASE_							(					self    :       str )->     Optional[int]:





     """simple docstring"""







     lowercase__			    =						self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_flaubert_lm_head(*a )


   def  SCREAMING_SNAKE_CASE_							(					self    :       Tuple )->     Dict:





     """simple docstring"""







     lowercase__			    =						self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_flaubert_qa(*a )


   def  SCREAMING_SNAKE_CASE_							(					self    :       Optional[int] )->     Any:





     """simple docstring"""







     lowercase__			    =						self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_flaubert_sequence_classif(*a )


   def  SCREAMING_SNAKE_CASE_							(					self    :       List[Any] )->     Optional[int]:





     """simple docstring"""







     lowercase__			    =						self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_flaubert_for_token_classification(*a )


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





     """simple docstring"""







     lowercase__			    =						self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_flaubert_for_multiple_choice(*a )


   @slow
   def  SCREAMING_SNAKE_CASE_							(					self    :       int )->     Optional[Any]:





     """simple docstring"""







     for model_name in TF_FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       lowercase__			    =						TFFlaubertModel.from_pretrained(a )
       self.assertIsNotNone(a )






@require_tf
@require_sentencepiece
@require_tokenizers
class        SCREAMING_SNAKE_CASE     (unittest.TestCase   ):


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





     """simple docstring"""







     lowercase__			    =						TFFlaubertModel.from_pretrained('jplu/tf-flaubert-small-cased' )

     lowercase__			    =						tf.convert_to_tensor(
         [[0, 158, 735, 2_592, 1_424, 6_727, 82, 1]]    ,  dtype=tf.intaa    ,  )  # "J'aime flaubert !"

     lowercase__			    =						model(a )[0]
     lowercase__			    =						tf.TensorShape((1, 8, 512) )
     self.assertEqual(output.shape    ,  a )
     # compare the actual values for a slice.
     lowercase__			    =						tf.convert_to_tensor(
         [
             [
                 [-1.8768773, -1.566555, 0.27072418],
                 [-1.6920038, -0.5873505, 1.9329599],
                 [-2.9563985, -1.6993835, 1.7972052],
             ]
         ]    ,  dtype=tf.floataa    ,  )

     self.assertTrue(np.allclose(output[:, :3, :3].numpy()    ,  expected_slice.numpy()    ,  atol=1E-4 ) )









import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    StableDiffusionSAGPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()
class        SCREAMING_SNAKE_CASE     (UpperCAmelCase    ,			UpperCAmelCase    ,			unittest.TestCase   ):
   _UpperCamelCase					:	List[Any]				      = StableDiffusionSAGPipeline
   _UpperCamelCase					:	str				      = TEXT_TO_IMAGE_PARAMS
   _UpperCamelCase					:	Dict				      = TEXT_TO_IMAGE_BATCH_PARAMS
   _UpperCamelCase					:	str				      = TEXT_TO_IMAGE_IMAGE_PARAMS
   _UpperCamelCase					:	Optional[int]				      = TEXT_TO_IMAGE_IMAGE_PARAMS
   _UpperCamelCase					:	Union[str, Any]				      = False


   def  SCREAMING_SNAKE_CASE_							(					self    :       List[str] )->     Dict:





     """simple docstring"""







     torch.manual_seed(0 )
     lowercase__			    =						UNetaDConditionModel(
         block_out_channels=(32, 64)    ,  layers_per_block=2    ,  sample_size=32    ,  in_channels=4    ,  out_channels=4    ,  down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')    ,  up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')    ,  cross_attention_dim=32    ,  )
     lowercase__			    =						DDIMScheduler(
         beta_start=0.00085    ,  beta_end=0.012    ,  beta_schedule='scaled_linear'    ,  clip_sample=a    ,  set_alpha_to_one=a    ,  )
     torch.manual_seed(0 )
     lowercase__			    =						AutoencoderKL(
         block_out_channels=[32, 64]    ,  in_channels=3    ,  out_channels=3    ,  down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']    ,  up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']    ,  latent_channels=4    ,  )
     torch.manual_seed(0 )
     lowercase__			    =						CLIPTextConfig(
         bos_token_id=0    ,  eos_token_id=2    ,  hidden_size=32    ,  intermediate_size=37    ,  layer_norm_eps=1E-0_5    ,  num_attention_heads=4    ,  num_hidden_layers=5    ,  pad_token_id=1    ,  vocab_size=1_000    ,  )
     lowercase__			    =						CLIPTextModel(a )
     lowercase__			    =						CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip' )

     lowercase__			    =						{
         'unet': unet,
         'scheduler': scheduler,
         'vae': vae,
         'text_encoder': text_encoder,
         'tokenizer': tokenizer,
         'safety_checker': None,
         'feature_extractor': None,
     }
     return components


   def  SCREAMING_SNAKE_CASE_							(					self    :       Tuple    ,  a    :       Optional[int]    ,  a    :       Any=0 )->     Union[str, Any]:





     """simple docstring"""







     if str(a ).startswith('mps' ):
       lowercase__			    =						torch.manual_seed(a )
     else:
       lowercase__			    =						torch.Generator(device=a ).manual_seed(a )
     lowercase__			    =						{
         'prompt': '.',
         'generator': generator,
         'num_inference_steps': 2,
         'guidance_scale': 1.0,
         'sag_scale': 1.0,
         'output_type': 'numpy',
     }
     return inputs


   def  SCREAMING_SNAKE_CASE_							(					self    :       Any )->     Any:





     """simple docstring"""







     super().test_inference_batch_single_identical(expected_max_diff=3E-3 )






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


   def  SCREAMING_SNAKE_CASE_							(					self    :       str )->     int:





     """simple docstring"""







     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()


   def  SCREAMING_SNAKE_CASE_							(					self    :       Tuple )->     Optional[Any]:





     """simple docstring"""







     lowercase__			    =						StableDiffusionSAGPipeline.from_pretrained('CompVis/stable-diffusion-v1-4' )
     lowercase__			    =						sag_pipe.to(a )
     sag_pipe.set_progress_bar_config(disable=a )

     lowercase__			    =						'.'
     lowercase__			    =						torch.manual_seed(0 )
     lowercase__			    =						sag_pipe(
         [prompt]    ,  generator=a    ,  guidance_scale=7.5    ,  sag_scale=1.0    ,  num_inference_steps=20    ,  output_type='np' )

     lowercase__			    =						output.images

     lowercase__			    =						image[0, -3:, -3:, -1]

     assert image.shape == (1, 512, 512, 3)
     lowercase__			    =						np.array([0.1568, 0.1738, 0.1695, 0.1693, 0.1507, 0.1705, 0.1547, 0.1751, 0.1949] )

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


   def  SCREAMING_SNAKE_CASE_							(					self    :       Tuple )->     Optional[Any]:





     """simple docstring"""







     lowercase__			    =						StableDiffusionSAGPipeline.from_pretrained('stabilityai/stable-diffusion-2-1-base' )
     lowercase__			    =						sag_pipe.to(a )
     sag_pipe.set_progress_bar_config(disable=a )

     lowercase__			    =						'.'
     lowercase__			    =						torch.manual_seed(0 )
     lowercase__			    =						sag_pipe(
         [prompt]    ,  generator=a    ,  guidance_scale=7.5    ,  sag_scale=1.0    ,  num_inference_steps=20    ,  output_type='np' )

     lowercase__			    =						output.images

     lowercase__			    =						image[0, -3:, -3:, -1]

     assert image.shape == (1, 512, 512, 3)
     lowercase__			    =						np.array([0.3459, 0.2876, 0.2537, 0.3002, 0.2671, 0.2160, 0.3026, 0.2262, 0.2371] )

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


   def  SCREAMING_SNAKE_CASE_							(					self    :       Optional[int] )->     List[Any]:





     """simple docstring"""







     lowercase__			    =						StableDiffusionSAGPipeline.from_pretrained('stabilityai/stable-diffusion-2-1-base' )
     lowercase__			    =						sag_pipe.to(a )
     sag_pipe.set_progress_bar_config(disable=a )

     lowercase__			    =						'.'
     lowercase__			    =						torch.manual_seed(0 )
     lowercase__			    =						sag_pipe(
         [prompt]    ,  width=768    ,  height=512    ,  generator=a    ,  guidance_scale=7.5    ,  sag_scale=1.0    ,  num_inference_steps=20    ,  output_type='np'    ,  )

     lowercase__			    =						output.images

     assert image.shape == (1, 512, 768, 3)


