

'''simple docstring'''

from __future__ import annotations

import inspect
import unittest

import numpy as np

from transformers import DeiTConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

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


if is_tf_available():
    import tensorflow as tf

    from transformers import (
        TFDeiTForImageClassification,
        TFDeiTForImageClassificationWithTeacher,
        TFDeiTForMaskedImageModeling,
        TFDeiTModel,
    )
    from transformers.models.deit.modeling_tf_deit import TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
    from PIL import Image

    from transformers import DeiTImageProcessor







class __a :





   def __init__(					self				:			Union[str, Any]				,lowerCamelCase				:			Dict				,lowerCamelCase				:			Optional[int]=13				,lowerCamelCase				:			List[Any]=30				,lowerCamelCase				:			Tuple=2				,lowerCamelCase				:			Tuple=3				,lowerCamelCase				:			str=True				,lowerCamelCase				:			str=True				,lowerCamelCase				:			Tuple=32				,lowerCamelCase				:			Dict=2				,lowerCamelCase				:			Dict=4				,lowerCamelCase				:			Tuple=37				,lowerCamelCase				:			Dict="gelu"				,lowerCamelCase				:			List[Any]=0.1				,lowerCamelCase				:			Optional[Any]=0.1				,lowerCamelCase				:			List[Any]=10				,lowerCamelCase				:			str=0.02				,lowerCamelCase				:			List[Any]=3				,lowerCamelCase				:			List[str]=None				,lowerCamelCase				:			Tuple=2				,):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE      							=       parent
      __SCREAMING_SNAKE_CASE      							=       batch_size
      __SCREAMING_SNAKE_CASE      							=       image_size
      __SCREAMING_SNAKE_CASE      							=       patch_size
      __SCREAMING_SNAKE_CASE      							=       num_channels
      __SCREAMING_SNAKE_CASE      							=       is_training
      __SCREAMING_SNAKE_CASE      							=       use_labels
      __SCREAMING_SNAKE_CASE      							=       hidden_size
      __SCREAMING_SNAKE_CASE      							=       num_hidden_layers
      __SCREAMING_SNAKE_CASE      							=       num_attention_heads
      __SCREAMING_SNAKE_CASE      							=       intermediate_size
      __SCREAMING_SNAKE_CASE      							=       hidden_act
      __SCREAMING_SNAKE_CASE      							=       hidden_dropout_prob
      __SCREAMING_SNAKE_CASE      							=       attention_probs_dropout_prob
      __SCREAMING_SNAKE_CASE      							=       type_sequence_label_size
      __SCREAMING_SNAKE_CASE      							=       initializer_range
      __SCREAMING_SNAKE_CASE      							=       scope
      __SCREAMING_SNAKE_CASE      							=       encoder_stride

      # in DeiT, the seq length equals the number of patches + 2 (we add 2 for the [CLS] and distilation tokens)
      __SCREAMING_SNAKE_CASE      							=       (image_size // patch_size) ** 2
      __SCREAMING_SNAKE_CASE      							=       num_patches + 2





   def      UpperCAmelCase__     (					self				:			Optional[int]    ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE      							=       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]    )

      __SCREAMING_SNAKE_CASE      							=       None
      if self.use_labels:
         __SCREAMING_SNAKE_CASE      							=       ids_tensor([self.batch_size]				,self.type_sequence_label_size    )

      __SCREAMING_SNAKE_CASE      							=       self.get_config()

      return config, pixel_values, labels





   def      UpperCAmelCase__     (					self				:			Optional[int]    ):

      '''simple docstring'''

      return DeiTConfig(
          image_size=self.image_size				,patch_size=self.patch_size				,num_channels=self.num_channels				,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				,is_decoder=_SCREAMING_SNAKE_CASE				,initializer_range=self.initializer_range				,encoder_stride=self.encoder_stride				,)





   def      UpperCAmelCase__     (					self				:			Dict				,lowerCamelCase				:			str				,lowerCamelCase				:			str				,lowerCamelCase				:			List[Any]    ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE      							=       TFDeiTModel(config=_SCREAMING_SNAKE_CASE    )
      __SCREAMING_SNAKE_CASE      							=       model(_SCREAMING_SNAKE_CASE    )
      self.parent.assertEqual(result.last_hidden_state.shape				,(self.batch_size, self.seq_length, self.hidden_size)    )





   def      UpperCAmelCase__     (					self				:			Any				,lowerCamelCase				:			Any				,lowerCamelCase				:			Dict				,lowerCamelCase				:			Dict    ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE      							=       TFDeiTForMaskedImageModeling(config=_SCREAMING_SNAKE_CASE    )
      __SCREAMING_SNAKE_CASE      							=       model(_SCREAMING_SNAKE_CASE    )
      self.parent.assertEqual(
          result.reconstruction.shape				,(self.batch_size, self.num_channels, self.image_size, self.image_size)    )

      # test greyscale images
      __SCREAMING_SNAKE_CASE      							=       1
      __SCREAMING_SNAKE_CASE      							=       TFDeiTForMaskedImageModeling(_SCREAMING_SNAKE_CASE    )

      __SCREAMING_SNAKE_CASE      							=       floats_tensor([self.batch_size, 1, self.image_size, self.image_size]    )
      __SCREAMING_SNAKE_CASE      							=       model(_SCREAMING_SNAKE_CASE    )
      self.parent.assertEqual(result.reconstruction.shape				,(self.batch_size, 1, self.image_size, self.image_size)    )





   def      UpperCAmelCase__     (					self				:			Optional[Any]				,lowerCamelCase				:			List[Any]				,lowerCamelCase				:			Any				,lowerCamelCase				:			Union[str, Any]    ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE      							=       self.type_sequence_label_size
      __SCREAMING_SNAKE_CASE      							=       TFDeiTForImageClassification(_SCREAMING_SNAKE_CASE    )
      __SCREAMING_SNAKE_CASE      							=       model(_SCREAMING_SNAKE_CASE				,labels=_SCREAMING_SNAKE_CASE    )
      self.parent.assertEqual(result.logits.shape				,(self.batch_size, self.type_sequence_label_size)    )

      # test greyscale images
      __SCREAMING_SNAKE_CASE      							=       1
      __SCREAMING_SNAKE_CASE      							=       TFDeiTForImageClassification(_SCREAMING_SNAKE_CASE    )

      __SCREAMING_SNAKE_CASE      							=       floats_tensor([self.batch_size, 1, self.image_size, self.image_size]    )
      __SCREAMING_SNAKE_CASE      							=       model(_SCREAMING_SNAKE_CASE				,labels=_SCREAMING_SNAKE_CASE    )
      self.parent.assertEqual(result.logits.shape				,(self.batch_size, self.type_sequence_label_size)    )





   def      UpperCAmelCase__     (					self				:			Dict    ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE      							=       self.prepare_config_and_inputs()
      __SCREAMING_SNAKE_CASE							,     __SCREAMING_SNAKE_CASE							,     __SCREAMING_SNAKE_CASE      							=       config_and_inputs
      __SCREAMING_SNAKE_CASE      							=       {"""pixel_values""": pixel_values}
      return config, inputs_dict







@require_tf
class __a (   _snake_case,   _snake_case,   unittest.TestCase			):
   __UpperCamelCase   : Union[str, Any]	      =	(
       (
           TFDeiTModel,
           TFDeiTForImageClassification,
           TFDeiTForImageClassificationWithTeacher,
           TFDeiTForMaskedImageModeling,
       )
       if is_tf_available()
       else ()
   )
   __UpperCamelCase   : Optional[int]	      =	(
       {
           'feature-extraction': TFDeiTModel,
           'image-classification': (TFDeiTForImageClassification, TFDeiTForImageClassificationWithTeacher),
       }
       if is_tf_available()
       else {}
   )

   __UpperCamelCase   : Union[str, Any]	      =	False
   __UpperCamelCase   : List[str]	      =	False
   __UpperCamelCase   : str	      =	False
   __UpperCamelCase   : Tuple	      =	False





   def      UpperCAmelCase__     (					self				:			Any    ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE      							=       TFDeiTModelTester(self    )
      __SCREAMING_SNAKE_CASE      							=       ConfigTester(self				,config_class=_SCREAMING_SNAKE_CASE				,has_text_modality=_SCREAMING_SNAKE_CASE				,hidden_size=37    )





   def      UpperCAmelCase__     (					self				:			Tuple    ):

      '''simple docstring'''

      self.config_tester.run_common_tests()





   @unittest.skip(reason="""DeiT does not use inputs_embeds"""    )
   def      UpperCAmelCase__     (					self				:			List[Any]    ):

      '''simple docstring'''

      pass





   def      UpperCAmelCase__     (					self				:			int    ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE							,     __SCREAMING_SNAKE_CASE      							=       self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
         __SCREAMING_SNAKE_CASE      							=       model_class(_SCREAMING_SNAKE_CASE    )
         self.assertIsInstance(model.get_input_embeddings()				,(tf.keras.layers.Layer)    )
         __SCREAMING_SNAKE_CASE      							=       model.get_output_embeddings()
         self.assertTrue(x is None or isinstance(_SCREAMING_SNAKE_CASE				,tf.keras.layers.Dense    )    )





   def      UpperCAmelCase__     (					self				:			Optional[int]    ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE							,     __SCREAMING_SNAKE_CASE      							=       self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
         __SCREAMING_SNAKE_CASE      							=       model_class(_SCREAMING_SNAKE_CASE    )
         __SCREAMING_SNAKE_CASE      							=       inspect.signature(model.call    )
         # signature.parameters is an OrderedDict => so arg_names order is deterministic
         __SCREAMING_SNAKE_CASE      							=       [*signature.parameters.keys()]

         __SCREAMING_SNAKE_CASE      							=       ["""pixel_values"""]
         self.assertListEqual(arg_names[:1]				,_SCREAMING_SNAKE_CASE    )





   def      UpperCAmelCase__     (					self				:			List[Any]    ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE      							=       self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*_SCREAMING_SNAKE_CASE    )





   def      UpperCAmelCase__     (					self				:			Optional[Any]    ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE      							=       self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_masked_image_modeling(*_SCREAMING_SNAKE_CASE    )





   def      UpperCAmelCase__     (					self				:			Optional[int]    ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE      							=       self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_image_classification(*_SCREAMING_SNAKE_CASE    )





   def      UpperCAmelCase__     (					self				:			Optional[Any]				,lowerCamelCase				:			List[str]				,lowerCamelCase				:			int				,lowerCamelCase				:			Dict=False    ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE      							=       super()._prepare_for_class(_SCREAMING_SNAKE_CASE				,_SCREAMING_SNAKE_CASE				,return_labels=_SCREAMING_SNAKE_CASE    )

      if return_labels:
         if "labels" in inputs_dict and "labels" not in inspect.signature(model_class.call    ).parameters:
            del inputs_dict["labels"]

      return inputs_dict





   @slow
   def      UpperCAmelCase__     (					self				:			Optional[int]    ):

      '''simple docstring'''

      for model_name in TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         __SCREAMING_SNAKE_CASE      							=       TFDeiTModel.from_pretrained(_SCREAMING_SNAKE_CASE    )
         self.assertIsNotNone(_SCREAMING_SNAKE_CASE    )





def 						__magic_name__       (							)						->					Dict:

   '''simple docstring'''







   __SCREAMING_SNAKE_CASE      							=       Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""      )
   return image







@require_tf
@require_vision
class __a (   unittest.TestCase			):





   @cached_property
   def      UpperCAmelCase__     (					self				:			List[str]    ):

      '''simple docstring'''

      return (
          DeiTImageProcessor.from_pretrained("""facebook/deit-base-distilled-patch16-224"""    )
          if is_vision_available()
          else None
      )





   @slow
   def      UpperCAmelCase__     (					self				:			List[Any]    ):

      '''simple docstring'''

      __SCREAMING_SNAKE_CASE      							=       TFDeiTForImageClassificationWithTeacher.from_pretrained("""facebook/deit-base-distilled-patch16-224"""    )

      __SCREAMING_SNAKE_CASE      							=       self.default_image_processor
      __SCREAMING_SNAKE_CASE      							=       prepare_img()
      __SCREAMING_SNAKE_CASE      							=       image_processor(images=_SCREAMING_SNAKE_CASE				,return_tensors="""tf"""    )

      # forward pass
      __SCREAMING_SNAKE_CASE      							=       model(**_SCREAMING_SNAKE_CASE    )

      # verify the logits
      __SCREAMING_SNAKE_CASE      							=       tf.TensorShape((1, 1000)    )
      self.assertEqual(outputs.logits.shape				,_SCREAMING_SNAKE_CASE    )

      __SCREAMING_SNAKE_CASE      							=       tf.constant([-1.0_266, 0.1_912, -1.2_861]    )

      self.assertTrue(np.allclose(outputs.logits[0, :3]				,_SCREAMING_SNAKE_CASE				,atol=1E-4    )    )












'''simple docstring'''

from ...utils import (
    OptionalDependencyNotAvailable,
    is_flax_available,
    is_torch_available,
    is_transformers_available,
)


try:
    if not (is_transformers_available() and is_torch_available()):
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    from ...utils.dummy_torch_and_transformers_objects import *  # noqa F403
else:
    from .multicontrolnet import MultiControlNetModel
    from .pipeline_controlnet import StableDiffusionControlNetPipeline
    from .pipeline_controlnet_imgaimg import StableDiffusionControlNetImgaImgPipeline
    from .pipeline_controlnet_inpaint import StableDiffusionControlNetInpaintPipeline


if is_transformers_available() and is_flax_available():
    from .pipeline_flax_controlnet import FlaxStableDiffusionControlNetPipeline





