


'''simple docstring'''

import inspect
import unittest

from transformers import ViTConfig
from transformers.testing_utils import (
    require_accelerate,
    require_torch,
    require_torch_gpu,
    require_vision,
    slow,
    torch_device,
)
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
 import torch
 from torch import nn

 from transformers import ViTForImageClassification, ViTForMaskedImageModeling, ViTModel
 from transformers.models.vit.modeling_vit import VIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
 from PIL import Image

 from transformers import ViTImageProcessor
class     __lowerCAmelCase		:





      '''simple docstring'''






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

             '''simple docstring'''







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

             # in ViT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
             lowercase__      =					(image_size // patch_size) ** 2
             lowercase__      =					num_patches + 1



      def  UpperCamelCase__						(self       :       Any					):

             '''simple docstring'''







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

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

             lowercase__      =					self.get_config()

             return config, pixel_values, labels



      def  UpperCamelCase__						(self       :       int					):

             '''simple docstring'''







             return ViTConfig(
                 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=UpperCamelCase  ,		initializer_range=self.initializer_range  ,		encoder_stride=self.encoder_stride  ,		)



      def  UpperCamelCase__						(self       :       List[str]  ,		UpperCamelCase       :       Optional[int]  ,		UpperCamelCase       :       str  ,		UpperCamelCase       :       List[str]					):

             '''simple docstring'''







             lowercase__      =					ViTModel(config=UpperCamelCase					)
             model.to(UpperCamelCase					)
             model.eval()
             lowercase__      =					model(UpperCamelCase					)
             self.parent.assertEqual(result.last_hidden_state.shape  ,		(self.batch_size, self.seq_length, self.hidden_size)					)



      def  UpperCamelCase__						(self       :       Union[str, Any]  ,		UpperCamelCase       :       str  ,		UpperCamelCase       :       List[str]  ,		UpperCamelCase       :       int					):

             '''simple docstring'''







             lowercase__      =					ViTForMaskedImageModeling(config=UpperCamelCase					)
             model.to(UpperCamelCase					)
             model.eval()
             lowercase__      =					model(UpperCamelCase					)
             self.parent.assertEqual(
                 result.reconstruction.shape  ,		(self.batch_size, self.num_channels, self.image_size, self.image_size)					)

             # test greyscale images
             lowercase__      =					1
             lowercase__      =					ViTForMaskedImageModeling(UpperCamelCase					)
             model.to(UpperCamelCase					)
             model.eval()

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



      def  UpperCamelCase__						(self       :       List[str]  ,		UpperCamelCase       :       Tuple  ,		UpperCamelCase       :       Tuple  ,		UpperCamelCase       :       List[str]					):

             '''simple docstring'''







             lowercase__      =					self.type_sequence_label_size
             lowercase__      =					ViTForImageClassification(UpperCamelCase					)
             model.to(UpperCamelCase					)
             model.eval()
             lowercase__      =					model(UpperCamelCase  ,		labels=UpperCamelCase					)
             self.parent.assertEqual(result.logits.shape  ,		(self.batch_size, self.type_sequence_label_size)					)

             # test greyscale images
             lowercase__      =					1
             lowercase__      =					ViTForImageClassification(UpperCamelCase					)
             model.to(UpperCamelCase					)
             model.eval()

             lowercase__      =					floats_tensor([self.batch_size, 1, self.image_size, self.image_size]					)
             lowercase__      =					model(UpperCamelCase					)
             self.parent.assertEqual(result.logits.shape  ,		(self.batch_size, self.type_sequence_label_size)					)






      def  UpperCamelCase__						(self       :       int					):

             '''simple docstring'''







             lowercase__      =					self.prepare_config_and_inputs()
             (
                 (
                 lowercase__
             )     ,(
                 lowercase__
             )     ,(
                 lowercase__
             )     ,
             )      =					config_and_inputs
             lowercase__      =					{'''pixel_values''': pixel_values}
             return config, inputs_dict






@require_torch
class     __lowerCAmelCase		(lowercase_    ,	lowercase_    ,	unittest.TestCase	):





      '''simple docstring'''






      lowerCAmelCase__   :	List[str]       =				(
          (
              ViTModel,
              ViTForImageClassification,
              ViTForMaskedImageModeling,
          )
          if is_torch_available()
          else ()
      )
      lowerCAmelCase__   :	str       =				(
          {"""feature-extraction""": ViTModel, """image-classification""": ViTForImageClassification}
          if is_torch_available()
          else {}
      )
      lowerCAmelCase__   :	Tuple       =				True

      lowerCAmelCase__   :	List[str]       =				False
      lowerCAmelCase__   :	Dict       =				False
      lowerCAmelCase__   :	List[str]       =				False



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

             '''simple docstring'''







             lowercase__      =					ViTModelTester(self					)
             lowercase__      =					ConfigTester(self  ,		config_class=UpperCamelCase  ,		has_text_modality=UpperCamelCase  ,		hidden_size=37					)



      def  UpperCamelCase__						(self       :       Dict					):

             '''simple docstring'''







             self.config_tester.run_common_tests()



      @unittest.skip(reason='''ViT does not use inputs_embeds'''					)
      def  UpperCamelCase__						(self       :       Optional[Any]					):

             '''simple docstring'''







             pass



      def  UpperCamelCase__						(self       :       Optional[int]					):

             '''simple docstring'''







             lowercase__     ,lowercase__      =					self.model_tester.prepare_config_and_inputs_for_common()

             for model_class in self.all_model_classes:
                    lowercase__      =					model_class(UpperCamelCase					)
                    self.assertIsInstance(model.get_input_embeddings()  ,		(nn.Module)					)
                    lowercase__      =					model.get_output_embeddings()
                    self.assertTrue(x is None or isinstance(UpperCamelCase  ,		nn.Linear					)					)



      def  UpperCamelCase__						(self       :       List[Any]					):

             '''simple docstring'''







             lowercase__     ,lowercase__      =					self.model_tester.prepare_config_and_inputs_for_common()

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

                    lowercase__      =					['''pixel_values''']
                    self.assertListEqual(arg_names[:1]  ,		UpperCamelCase					)



      def  UpperCamelCase__						(self       :       List[str]					):

             '''simple docstring'''







             lowercase__      =					self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_model(*UpperCamelCase					)



      def  UpperCamelCase__						(self       :       List[str]					):

             '''simple docstring'''







             lowercase__      =					self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_masked_image_modeling(*UpperCamelCase					)



      def  UpperCamelCase__						(self       :       Dict					):

             '''simple docstring'''







             lowercase__      =					self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_image_classification(*UpperCamelCase					)






      @slow
      def  UpperCamelCase__						(self       :       Dict					):

             '''simple docstring'''







             for model_name in VIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                    lowercase__      =					ViTModel.from_pretrained(UpperCamelCase					)
                    self.assertIsNotNone(UpperCamelCase					)






def _SCREAMING_SNAKE_CASE ()			->			Optional[Any]:






       """simple docstring"""







       lowercase__      =					Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''			)
       return image

@require_torch
@require_vision
class     __lowerCAmelCase		(unittest.TestCase	):





      '''simple docstring'''






      @cached_property
      def  UpperCamelCase__						(self       :       Optional[int]					):

             '''simple docstring'''







             return ViTImageProcessor.from_pretrained('''google/vit-base-patch16-224'''					) if is_vision_available() else None



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

             '''simple docstring'''







             lowercase__      =					ViTForImageClassification.from_pretrained('''google/vit-base-patch16-224'''					).to(UpperCamelCase					)

             lowercase__      =					self.default_image_processor
             lowercase__      =					prepare_img()
             lowercase__      =					image_processor(images=UpperCamelCase  ,		return_tensors='''pt'''					).to(UpperCamelCase					)

             # forward pass
             with torch.no_grad():
                    lowercase__      =					model(**UpperCamelCase					)

             # verify the logits
             lowercase__      =					torch.Size((1, 1000)					)
             self.assertEqual(outputs.logits.shape  ,		UpperCamelCase					)

             lowercase__      =					torch.tensor([-0.27_44, 0.82_15, -0.08_36]					).to(UpperCamelCase					)

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



      @slow
      def  UpperCamelCase__						(self       :       str					):

             '''simple docstring'''







             lowercase__      =					ViTModel.from_pretrained('''facebook/dino-vits8'''					).to(UpperCamelCase					)

             lowercase__      =					ViTImageProcessor.from_pretrained('''facebook/dino-vits8'''  ,		size=480					)
             lowercase__      =					prepare_img()
             lowercase__      =					image_processor(images=UpperCamelCase  ,		return_tensors='''pt'''					)
             lowercase__      =					inputs.pixel_values.to(UpperCamelCase					)

             # forward pass
             with torch.no_grad():
                    lowercase__      =					model(UpperCamelCase  ,		interpolate_pos_encoding=UpperCamelCase					)

             # verify the logits
             lowercase__      =					torch.Size((1, 3601, 384)					)
             self.assertEqual(outputs.last_hidden_state.shape  ,		UpperCamelCase					)

             lowercase__      =					torch.tensor(
                 [[4.23_40, 4.39_06, -6.66_92], [4.54_63, 1.89_28, -6.72_57], [4.44_29, 0.84_96, -5.85_85]]					).to(UpperCamelCase					)

             self.assertTrue(torch.allclose(outputs.last_hidden_state[0, :3, :3]  ,		UpperCamelCase  ,		atol=1E-4					)					)






      @slow
      @require_accelerate
      @require_torch_gpu
      def  UpperCamelCase__						(self       :       int					):

             '''simple docstring'''







             lowercase__      =					ViTModel.from_pretrained('''facebook/dino-vits8'''  ,		torch_dtype=torch.floataa  ,		device_map='''auto'''					)
             lowercase__      =					self.default_image_processor

             lowercase__      =					prepare_img()
             lowercase__      =					image_processor(images=UpperCamelCase  ,		return_tensors='''pt'''					)
             lowercase__      =					inputs.pixel_values.to(UpperCamelCase					)

             # forward pass to make sure inference works in fp16
             with torch.no_grad():
                    lowercase__      =					model(UpperCamelCase					)








'''simple docstring'''

from typing import Optional, Tuple

import jax
import jax.numpy as jnp
from flax import linen as nn
from flax.core.frozen_dict import FrozenDict
from transformers import CLIPConfig, FlaxPreTrainedModel
from transformers.models.clip.modeling_flax_clip import FlaxCLIPVisionModule

def _SCREAMING_SNAKE_CASE (A    ,					A    ,					A=1E-12			)			->			str:






       """simple docstring"""







       lowercase__      =					jnp.divide(emb_a.T    ,					jnp.clip(jnp.linalg.norm(A    ,					axis=1			)    ,					a_min=A			)			).T
       lowercase__      =					jnp.divide(emb_a.T    ,					jnp.clip(jnp.linalg.norm(A    ,					axis=1			)    ,					a_min=A			)			).T
       return jnp.matmul(A    ,					norm_emb_a.T			)
class     __lowerCAmelCase		(nn.Module	):





      '''simple docstring'''






      lowerCAmelCase__   :	CLIPConfig
      lowerCAmelCase__   :	jnp.dtype       =				jnp.floataa



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

             '''simple docstring'''







             lowercase__      =					FlaxCLIPVisionModule(self.config.vision_config					)
             lowercase__      =					nn.Dense(self.config.projection_dim  ,		use_bias=UpperCamelCase  ,		dtype=self.dtype					)

             lowercase__      =					self.param('''concept_embeds'''  ,		jax.nn.initializers.ones  ,		(17, self.config.projection_dim)					)
             lowercase__      =					self.param(
                 '''special_care_embeds'''  ,		jax.nn.initializers.ones  ,		(3, self.config.projection_dim)					)

             lowercase__      =					self.param('''concept_embeds_weights'''  ,		jax.nn.initializers.ones  ,		(17,)					)
             lowercase__      =					self.param('''special_care_embeds_weights'''  ,		jax.nn.initializers.ones  ,		(3,)					)






      def __call__(self       :       Union[str, Any]  ,		UpperCamelCase       :       Tuple					):

             '''simple docstring'''







             lowercase__      =					self.vision_model(UpperCamelCase					)[1]
             lowercase__      =					self.visual_projection(UpperCamelCase					)

             lowercase__      =					jax_cosine_distance(UpperCamelCase  ,		self.special_care_embeds					)
             lowercase__      =					jax_cosine_distance(UpperCamelCase  ,		self.concept_embeds					)

             # increase this value to create a stronger `nfsw` filter
             # at the cost of increasing the possibility of filtering benign image inputs
             lowercase__      =					0.0

             lowercase__      =					special_cos_dist - self.special_care_embeds_weights[None, :] + adjustment
             lowercase__      =					jnp.round(UpperCamelCase  ,		3					)
             lowercase__      =					jnp.any(special_scores > 0  ,		axis=1  ,		keepdims=UpperCamelCase					)
             # Use a lower threshold if an image has any special care concept
             lowercase__      =					is_special_care * 0.01

             lowercase__      =					cos_dist - self.concept_embeds_weights[None, :] + special_adjustment
             lowercase__      =					jnp.round(UpperCamelCase  ,		3					)
             lowercase__      =					jnp.any(concept_scores > 0  ,		axis=1					)

             return has_nsfw_concepts






class     __lowerCAmelCase		(lowercase_	):





      '''simple docstring'''






      lowerCAmelCase__   :	Union[str, Any]       =				CLIPConfig
      lowerCAmelCase__   :	Optional[int]       =				"""clip_input"""
      lowerCAmelCase__   :	Any       =				FlaxStableDiffusionSafetyCheckerModule



      def __init__(self       :       Optional[int]  ,		UpperCamelCase       :       CLIPConfig  ,		UpperCamelCase       :       Optional[Tuple] = None  ,		UpperCamelCase       :       int = 0  ,		UpperCamelCase       :       jnp.dtype = jnp.floataa  ,		UpperCamelCase       :       bool = True  ,		**UpperCamelCase       :       List[Any]  ,		):

             '''simple docstring'''







             if input_shape is None:
                    lowercase__      =					(1, 224, 224, 3)
             lowercase__      =					self.module_class(config=UpperCamelCase  ,		dtype=UpperCamelCase  ,		**UpperCamelCase					)
             super().__init__(UpperCamelCase  ,		UpperCamelCase  ,		input_shape=UpperCamelCase  ,		seed=UpperCamelCase  ,		dtype=UpperCamelCase  ,		_do_init=_do_init					)



      def  UpperCamelCase__						(self       :       Any  ,		UpperCamelCase       :       jax.random.KeyArray  ,		UpperCamelCase       :       Tuple  ,		UpperCamelCase       :       FrozenDict = None					):

             '''simple docstring'''







             lowercase__      =					jax.random.normal(UpperCamelCase  ,		UpperCamelCase					)

             lowercase__     ,lowercase__      =					jax.random.split(UpperCamelCase					)
             lowercase__      =					{'''params''': params_rng, '''dropout''': dropout_rng}

             lowercase__      =					self.module.init(UpperCamelCase  ,		UpperCamelCase					)['''params''']

             return random_params






      def __call__(self       :       Optional[Any]  ,		UpperCamelCase       :       Tuple  ,		UpperCamelCase       :       dict = None  ,		):

             '''simple docstring'''







             lowercase__      =					jnp.transpose(UpperCamelCase  ,		(0, 2, 3, 1)					)

             return self.module.apply(
                 {'''params''': params or self.params}  ,		jnp.array(UpperCamelCase  ,		dtype=jnp.floataa					)  ,		rngs={}  ,		)




