



"""simple docstring"""






from __future__ import annotations

import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_tf, 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
     import tensorflow as tf

     from transformers import (
         TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
         TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
         TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
         BertConfig,
         DPRConfig,
         TFDPRContextEncoder,
         TFDPRQuestionEncoder,
         TFDPRReader,
     )



class     UpperCAmelCase    :



   def __init__(       self    :	List[Any]							,       __lowerCamelCase    :	Optional[int]							,       __lowerCamelCase    :	Union[str, Any]=1_3							,       __lowerCamelCase    :	str=7							,       __lowerCamelCase    :	Optional[int]=True							,       __lowerCamelCase    :	str=True							,       __lowerCamelCase    :	int=True							,       __lowerCamelCase    :	List[Any]=True							,       __lowerCamelCase    :	Optional[int]=9_9							,       __lowerCamelCase    :	str=3_2							,       __lowerCamelCase    :	Optional[int]=2							,       __lowerCamelCase    :	Union[str, Any]=4							,       __lowerCamelCase    :	Optional[Any]=3_7							,       __lowerCamelCase    :	List[str]="gelu"							,       __lowerCamelCase    :	List[Any]=0.1							,       __lowerCamelCase    :	int=0.1							,       __lowerCamelCase    :	Any=5_1_2							,       __lowerCamelCase    :	Any=1_6							,       __lowerCamelCase    :	Dict=2							,       __lowerCamelCase    :	Optional[int]=0.0_2							,       __lowerCamelCase    :	Any=3							,       __lowerCamelCase    :	Optional[int]=4							,       __lowerCamelCase    :	Optional[Any]=None							,       __lowerCamelCase    :	List[str]=0							,       ):



     """simple docstring"""
     _snake_case   	=      parent
     _snake_case   	=      batch_size
     _snake_case   	=      seq_length
     _snake_case   	=      is_training
     _snake_case   	=      use_input_mask
     _snake_case   	=      use_token_type_ids
     _snake_case   	=      use_labels
     _snake_case   	=      vocab_size
     _snake_case   	=      hidden_size
     _snake_case   	=      num_hidden_layers
     _snake_case   	=      num_attention_heads
     _snake_case   	=      intermediate_size
     _snake_case   	=      hidden_act
     _snake_case   	=      hidden_dropout_prob
     _snake_case   	=      attention_probs_dropout_prob
     _snake_case   	=      max_position_embeddings
     _snake_case   	=      type_vocab_size
     _snake_case   	=      type_sequence_label_size
     _snake_case   	=      initializer_range
     _snake_case   	=      num_labels
     _snake_case   	=      num_choices
     _snake_case   	=      scope
     _snake_case   	=      projection_dim



   def    __UpperCAmelCase    (       self    :	List[Any]				):



     """simple docstring"""
     _snake_case   	=      ids_tensor([self.batch_size, self.seq_length]							,       self.vocab_size				)

     _snake_case   	=      None
     if self.use_input_mask:
       # follow test_modeling_tf_ctrl.py
       _snake_case   	=      random_attention_mask([self.batch_size, self.seq_length]				)

     _snake_case   	=      None
     if self.use_token_type_ids:
       _snake_case   	=      ids_tensor([self.batch_size, self.seq_length]							,       self.type_vocab_size				)

     _snake_case   	=      None
     _snake_case   	=      None
     _snake_case   	=      None
     if self.use_labels:
       _snake_case   	=      ids_tensor([self.batch_size]							,       self.type_sequence_label_size				)
       _snake_case   	=      ids_tensor([self.batch_size, self.seq_length]							,       self.num_labels				)
       _snake_case   	=      ids_tensor([self.batch_size]							,       self.num_choices				)

     _snake_case   	=      BertConfig(
         vocab_size=self.vocab_size							,       hidden_size=self.hidden_size							,       num_hidden_layers=self.num_hidden_layers							,       num_attention_heads=self.num_attention_heads							,       intermediate_size=self.intermediate_size							,       hidden_act=self.hidden_act							,       hidden_dropout_prob=self.hidden_dropout_prob							,       attention_probs_dropout_prob=self.attention_probs_dropout_prob							,       max_position_embeddings=self.max_position_embeddings							,       type_vocab_size=self.type_vocab_size							,       is_decoder=__lowerCamelCase							,       initializer_range=self.initializer_range							,       )
     _snake_case   	=      DPRConfig(projection_dim=self.projection_dim							,       **config.to_dict()				)

     return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels



   def    __UpperCAmelCase    (       self    :	Optional[int]							,       __lowerCamelCase    :	Union[str, Any]							,       __lowerCamelCase    :	List[Any]							,       __lowerCamelCase    :	int							,       __lowerCamelCase    :	List[Any]							,       __lowerCamelCase    :	Dict							,       __lowerCamelCase    :	Any							,       __lowerCamelCase    :	Tuple				):



     """simple docstring"""
     _snake_case   	=      TFDPRContextEncoder(config=__lowerCamelCase				)
     _snake_case   	=      model(__lowerCamelCase							,       attention_mask=__lowerCamelCase							,       token_type_ids=__lowerCamelCase				)
     _snake_case   	=      model(__lowerCamelCase							,       token_type_ids=__lowerCamelCase				)
     _snake_case   	=      model(__lowerCamelCase				)
     self.parent.assertEqual(result.pooler_output.shape							,       (self.batch_size, self.projection_dim or self.hidden_size)				)



   def    __UpperCAmelCase    (       self    :	List[Any]							,       __lowerCamelCase    :	List[str]							,       __lowerCamelCase    :	List[Any]							,       __lowerCamelCase    :	Dict							,       __lowerCamelCase    :	Any							,       __lowerCamelCase    :	str							,       __lowerCamelCase    :	str							,       __lowerCamelCase    :	List[str]				):



     """simple docstring"""
     _snake_case   	=      TFDPRQuestionEncoder(config=__lowerCamelCase				)
     _snake_case   	=      model(__lowerCamelCase							,       attention_mask=__lowerCamelCase							,       token_type_ids=__lowerCamelCase				)
     _snake_case   	=      model(__lowerCamelCase							,       token_type_ids=__lowerCamelCase				)
     _snake_case   	=      model(__lowerCamelCase				)
     self.parent.assertEqual(result.pooler_output.shape							,       (self.batch_size, self.projection_dim or self.hidden_size)				)



   def    __UpperCAmelCase    (       self    :	List[Any]							,       __lowerCamelCase    :	str							,       __lowerCamelCase    :	Tuple							,       __lowerCamelCase    :	List[str]							,       __lowerCamelCase    :	Dict							,       __lowerCamelCase    :	Tuple							,       __lowerCamelCase    :	Tuple							,       __lowerCamelCase    :	Optional[int]				):



     """simple docstring"""
     _snake_case   	=      TFDPRReader(config=__lowerCamelCase				)
     _snake_case   	=      model(__lowerCamelCase							,       attention_mask=__lowerCamelCase				)

     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)				)
     self.parent.assertEqual(result.relevance_logits.shape							,       (self.batch_size,)				)





   def    __UpperCAmelCase    (       self    :	Optional[Any]				):



     """simple docstring"""
     _snake_case   	=      self.prepare_config_and_inputs()
     (
         (
         _snake_case
     )	,     (
         _snake_case
     )	,     (
         _snake_case
     )	,     (
         _snake_case
     )	,     (
         _snake_case
     )	,     (
         _snake_case
     )	,     (
         _snake_case
     )	,     
     )   	=      config_and_inputs
     _snake_case   	=      {'''input_ids''': input_ids}
     return config, inputs_dict



@require_tf
class     UpperCAmelCase    (	__SCREAMING_SNAKE_CASE,__SCREAMING_SNAKE_CASE,unittest.TestCase  ):
   A__		:	Any           =  (
       (
           TFDPRContextEncoder,
           TFDPRQuestionEncoder,
           TFDPRReader,
       )
       if is_tf_available()
       else ()
   )
   A__		:	Optional[int]           =  {'''feature-extraction''': TFDPRQuestionEncoder} if is_tf_available() else {}

   A__		:	List[Any]           =  False
   A__		:	Any           =  False
   A__		:	Optional[int]           =  False
   A__		:	Optional[Any]           =  False
   A__		:	str           =  False



   def    __UpperCAmelCase    (       self    :	Dict				):



     """simple docstring"""
     _snake_case   	=      TFDPRModelTester(self				)
     _snake_case   	=      ConfigTester(self							,       config_class=__lowerCamelCase							,       hidden_size=3_7				)



   def    __UpperCAmelCase    (       self    :	Tuple				):



     """simple docstring"""
     self.config_tester.run_common_tests()



   def    __UpperCAmelCase    (       self    :	Union[str, Any]				):



     """simple docstring"""
     _snake_case   	=      self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_dpr_context_encoder(*__lowerCamelCase				)



   def    __UpperCAmelCase    (       self    :	Union[str, Any]				):



     """simple docstring"""
     _snake_case   	=      self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_dpr_question_encoder(*__lowerCamelCase				)



   def    __UpperCAmelCase    (       self    :	Any				):



     """simple docstring"""
     _snake_case   	=      self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_dpr_reader(*__lowerCamelCase				)





   @slow
   def    __UpperCAmelCase    (       self    :	Optional[Any]				):



     """simple docstring"""
     for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       _snake_case   	=      TFDPRContextEncoder.from_pretrained(__lowerCamelCase				)
       self.assertIsNotNone(__lowerCamelCase				)

     for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       _snake_case   	=      TFDPRContextEncoder.from_pretrained(__lowerCamelCase				)
       self.assertIsNotNone(__lowerCamelCase				)

     for model_name in TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       _snake_case   	=      TFDPRQuestionEncoder.from_pretrained(__lowerCamelCase				)
       self.assertIsNotNone(__lowerCamelCase				)

     for model_name in TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       _snake_case   	=      TFDPRReader.from_pretrained(__lowerCamelCase				)
       self.assertIsNotNone(__lowerCamelCase				)



@require_tf
class     UpperCAmelCase    (	unittest.TestCase  ):



   @slow
   def    __UpperCAmelCase    (       self    :	str				):



     """simple docstring"""
     _snake_case   	=      TFDPRQuestionEncoder.from_pretrained('''facebook/dpr-question_encoder-single-nq-base'''				)

     _snake_case   	=      tf.constant(
         [[1_0_1, 7_5_9_2, 1_0_1_0, 2_0_0_3, 2_0_2_6, 3_8_9_9, 1_0_1_4_0, 1_0_2_9, 1_0_2]]				)  # [CLS] hello, is my dog cute? [SEP]
     _snake_case   	=      model(__lowerCamelCase				)[0]  # embedding shape = (1, 768)
     # compare the actual values for a slice.
     _snake_case   	=      tf.constant(
         [
             [
                 0.0_3_2_3_6_2_5_3,
                 0.1_2_7_5_3_3_3_5,
                 0.1_6_8_1_8_5_0_9,
                 0.0_0_2_7_9_7_8_6,
                 0.3_8_9_6_9_3_3,
                 0.2_4_2_6_4_9_4_5,
                 0.2_1_7_8_9_7_1,
                 -0.0_2_3_3_5_2_2_7,
                 -0.0_8_4_8_1_9_5_9,
                 -0.1_4_3_2_4_1_1_7,
             ]
         ]				)
     self.assertTrue(numpy.allclose(output[:, :1_0].numpy()							,       expected_slice.numpy()							,       atol=1E-4				)				)









"""simple docstring"""
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 	lowercase__				(					A_ ,A_ ,unittest.TestCase				):
   __UpperCAmelCase							       =					StableDiffusionSAGPipeline
   __UpperCAmelCase							       =					TEXT_TO_IMAGE_PARAMS
   __UpperCAmelCase							       =					TEXT_TO_IMAGE_BATCH_PARAMS
   __UpperCAmelCase							       =					TEXT_TO_IMAGE_IMAGE_PARAMS
   __UpperCAmelCase							       =					TEXT_TO_IMAGE_IMAGE_PARAMS
   __UpperCAmelCase							       =					False





   def 		UpperCamelCase_  (							self)  ->  Optional[Any]:
      torch.manual_seed(0)
      _lowerCamelCase	:					Dict   =		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      ,   )
      _lowerCamelCase	:					int   =		DDIMScheduler(
          beta_start=0.0_00_85      ,   beta_end=0.0_12      ,   beta_schedule="""scaled_linear"""      ,   clip_sample=SCREAMING_SNAKE_CASE      ,   set_alpha_to_one=SCREAMING_SNAKE_CASE      ,   )
      torch.manual_seed(0)
      _lowerCamelCase	:					Tuple   =		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)
      _lowerCamelCase	:					Tuple   =		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=1000      ,   )
      _lowerCamelCase	:					List[Any]   =		CLIPTextModel(SCREAMING_SNAKE_CASE)
      _lowerCamelCase	:					Tuple   =		CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip""")

      _lowerCamelCase	:					List[Any]   =		{
          """unet""": unet,
          """scheduler""": scheduler,
          """vae""": vae,
          """text_encoder""": text_encoder,
          """tokenizer""": tokenizer,
          """safety_checker""": None,
          """feature_extractor""": None,
      }
      return components





   def 		UpperCamelCase_  (							self      ,   SCREAMING_SNAKE_CASE      ,   SCREAMING_SNAKE_CASE=0)  ->  List[Any]:
      if str(SCREAMING_SNAKE_CASE).startswith("""mps"""):
         _lowerCamelCase	:					List[str]   =		torch.manual_seed(SCREAMING_SNAKE_CASE)
      else:
         _lowerCamelCase	:					List[str]   =		torch.Generator(device=SCREAMING_SNAKE_CASE).manual_seed(SCREAMING_SNAKE_CASE)
      _lowerCamelCase	:					List[Any]   =		{
          """prompt""": """.""",
          """generator""": generator,
          """num_inference_steps""": 2,
          """guidance_scale""": 1.0,
          """sag_scale""": 1.0,
          """output_type""": """numpy""",
      }
      return inputs





   def 		UpperCamelCase_  (							self)  ->  Tuple:
      super().test_inference_batch_single_identical(expected_max_diff=3e-3)







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





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





   def 		UpperCamelCase_  (							self)  ->  Optional[Any]:
      _lowerCamelCase	:					Any   =		StableDiffusionSAGPipeline.from_pretrained("""CompVis/stable-diffusion-v1-4""")
      _lowerCamelCase	:					Union[str, Any]   =		sag_pipe.to(SCREAMING_SNAKE_CASE)
      sag_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE)

      _lowerCamelCase	:					Optional[int]   =		"""."""
      _lowerCamelCase	:					int   =		torch.manual_seed(0)
      _lowerCamelCase	:					Tuple   =		sag_pipe(
          [prompt]      ,   generator=SCREAMING_SNAKE_CASE      ,   guidance_scale=7.5      ,   sag_scale=1.0      ,   num_inference_steps=20      ,   output_type="""np""")

      _lowerCamelCase	:					Dict   =		output.images

      _lowerCamelCase	:					Tuple   =		image[0, -3:, -3:, -1]

      assert image.shape == (1, 512, 512, 3)
      _lowerCamelCase	:					Optional[Any]   =		np.array([0.15_68, 0.17_38, 0.16_95, 0.16_93, 0.15_07, 0.17_05, 0.15_47, 0.17_51, 0.19_49])

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





   def 		UpperCamelCase_  (							self)  ->  List[str]:
      _lowerCamelCase	:					Optional[Any]   =		StableDiffusionSAGPipeline.from_pretrained("""stabilityai/stable-diffusion-2-1-base""")
      _lowerCamelCase	:					Dict   =		sag_pipe.to(SCREAMING_SNAKE_CASE)
      sag_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE)

      _lowerCamelCase	:					Union[str, Any]   =		"""."""
      _lowerCamelCase	:					List[str]   =		torch.manual_seed(0)
      _lowerCamelCase	:					int   =		sag_pipe(
          [prompt]      ,   generator=SCREAMING_SNAKE_CASE      ,   guidance_scale=7.5      ,   sag_scale=1.0      ,   num_inference_steps=20      ,   output_type="""np""")

      _lowerCamelCase	:					Any   =		output.images

      _lowerCamelCase	:					List[str]   =		image[0, -3:, -3:, -1]

      assert image.shape == (1, 512, 512, 3)
      _lowerCamelCase	:					Any   =		np.array([0.34_59, 0.28_76, 0.25_37, 0.30_02, 0.26_71, 0.21_60, 0.30_26, 0.22_62, 0.23_71])

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





   def 		UpperCamelCase_  (							self)  ->  List[str]:
      _lowerCamelCase	:					int   =		StableDiffusionSAGPipeline.from_pretrained("""stabilityai/stable-diffusion-2-1-base""")
      _lowerCamelCase	:					Optional[Any]   =		sag_pipe.to(SCREAMING_SNAKE_CASE)
      sag_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE)

      _lowerCamelCase	:					Dict   =		"""."""
      _lowerCamelCase	:					Union[str, Any]   =		torch.manual_seed(0)
      _lowerCamelCase	:					Optional[int]   =		sag_pipe(
          [prompt]      ,   width=768      ,   height=512      ,   generator=SCREAMING_SNAKE_CASE      ,   guidance_scale=7.5      ,   sag_scale=1.0      ,   num_inference_steps=20      ,   output_type="""np"""      ,   )

      _lowerCamelCase	:					Union[str, Any]   =		output.images

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

