



'''simple docstring'''


from __future__ import annotations

import inspect
import unittest
from math import floor

import numpy as np

from transformers import CvtConfig
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 TFCvtForImageClassification, TFCvtModel
   from transformers.models.cvt.modeling_tf_cvt import TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
   from PIL import Image

   from transformers import AutoImageProcessor

class  SCREAMING_SNAKE_CASE (  _a		):






    """simple docstring"""





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


       """simple docstring"""







       UpperCamelCase					=  self.config_class(**self.inputs_dict )
       self.parent.assertTrue(hasattr(UpperCamelCase__  ,						'embed_dim' ) )
       self.parent.assertTrue(hasattr(UpperCamelCase__  ,						'num_heads' ) )

class  SCREAMING_SNAKE_CASE :






    """simple docstring"""





    def __init__(  self      :					str  ,						UpperCamelCase__      :					str  ,						UpperCamelCase__      :					Dict=1_3  ,						UpperCamelCase__      :					List[Any]=6_4  ,						UpperCamelCase__      :					Dict=3  ,						UpperCamelCase__      :					Union[str, Any]=[1_6, 4_8, 9_6]  ,						UpperCamelCase__      :					Tuple=[1, 3, 6]  ,						UpperCamelCase__      :					int=[1, 2, 1_0]  ,						UpperCamelCase__      :					Union[str, Any]=[7, 3, 3]  ,						UpperCamelCase__      :					str=[4, 2, 2]  ,						UpperCamelCase__      :					Tuple=[2, 1, 1]  ,						UpperCamelCase__      :					Any=[2, 2, 2]  ,						UpperCamelCase__      :					Union[str, Any]=[False, False, True]  ,						UpperCamelCase__      :					Dict=[0.0, 0.0, 0.0]  ,						UpperCamelCase__      :					Tuple=0.0_2  ,						UpperCamelCase__      :					Tuple=1E-1_2  ,						UpperCamelCase__      :					List[Any]=True  ,						UpperCamelCase__      :					Any=True  ,						UpperCamelCase__      :					Union[str, Any]=2  ,						):


       """simple docstring"""







       UpperCamelCase					=  parent
       UpperCamelCase					=  batch_size
       UpperCamelCase					=  image_size
       UpperCamelCase					=  patch_sizes
       UpperCamelCase					=  patch_stride
       UpperCamelCase					=  patch_padding
       UpperCamelCase					=  is_training
       UpperCamelCase					=  use_labels
       UpperCamelCase					=  num_labels
       UpperCamelCase					=  num_channels
       UpperCamelCase					=  embed_dim
       UpperCamelCase					=  num_heads
       UpperCamelCase					=  stride_kv
       UpperCamelCase					=  depth
       UpperCamelCase					=  cls_token
       UpperCamelCase					=  attention_drop_rate
       UpperCamelCase					=  initializer_range
       UpperCamelCase					=  layer_norm_eps






    def  A       (  self      :					Dict ):


       """simple docstring"""







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

       UpperCamelCase					=  None
       if self.use_labels:
          # create a random int32 tensor of given shape
          UpperCamelCase					=  ids_tensor([self.batch_size]  ,						self.num_labels )

       UpperCamelCase					=  self.get_config()
       return config, pixel_values, labels






    def  A       (  self      :					List[str] ):


       """simple docstring"""







       return CvtConfig(
           image_size=self.image_size  ,						num_labels=self.num_labels  ,						num_channels=self.num_channels  ,						embed_dim=self.embed_dim  ,						num_heads=self.num_heads  ,						patch_sizes=self.patch_sizes  ,						patch_padding=self.patch_padding  ,						patch_stride=self.patch_stride  ,						stride_kv=self.stride_kv  ,						depth=self.depth  ,						cls_token=self.cls_token  ,						attention_drop_rate=self.attention_drop_rate  ,						initializer_range=self.initializer_range  ,						)






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


       """simple docstring"""







       UpperCamelCase					=  TFCvtModel(config=UpperCamelCase__ )
       UpperCamelCase					=  model(UpperCamelCase__  ,						training=UpperCamelCase__ )
       UpperCamelCase					=  (self.image_size, self.image_size)
       UpperCamelCase				,      UpperCamelCase					=  image_size[0], image_size[1]
       for i in range(len(self.depth ) ):
          UpperCamelCase					=  floor(((height + 2 * self.patch_padding[i] - self.patch_sizes[i]) / self.patch_stride[i]) + 1 )
          UpperCamelCase					=  floor(((width + 2 * self.patch_padding[i] - self.patch_sizes[i]) / self.patch_stride[i]) + 1 )
       self.parent.assertEqual(result.last_hidden_state.shape  ,						(self.batch_size, self.embed_dim[-1], height, width) )






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


       """simple docstring"""







       UpperCamelCase					=  self.num_labels
       UpperCamelCase					=  TFCvtForImageClassification(UpperCamelCase__ )
       UpperCamelCase					=  model(UpperCamelCase__  ,						labels=UpperCamelCase__  ,						training=UpperCamelCase__ )
       self.parent.assertEqual(result.logits.shape  ,						(self.batch_size, self.num_labels) )






    def  A       (  self      :					List[Any] ):


       """simple docstring"""







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

@require_tf
class  SCREAMING_SNAKE_CASE (  _a   ,						_a   ,						unittest.TestCase		):






    """simple docstring"""





    _SCREAMING_SNAKE_CASE	  =  (TFCvtModel, TFCvtForImageClassification) if is_tf_available() else ()
    _SCREAMING_SNAKE_CASE	  =  (
        {"""feature-extraction""": TFCvtModel, """image-classification""": TFCvtForImageClassification}
        if is_tf_available()
        else {}
    )
    _SCREAMING_SNAKE_CASE	  =  False
    _SCREAMING_SNAKE_CASE	  =  False
    _SCREAMING_SNAKE_CASE	  =  False
    _SCREAMING_SNAKE_CASE	  =  False
    _SCREAMING_SNAKE_CASE	  =  False






    def  A       (  self      :					Optional[int] ):


       """simple docstring"""







       UpperCamelCase					=  TFCvtModelTester(self )
       UpperCamelCase					=  TFCvtConfigTester(self  ,						config_class=UpperCamelCase__  ,						has_text_modality=UpperCamelCase__  ,						hidden_size=3_7 )






    def  A       (  self      :					List[str] ):


       """simple docstring"""







       self.config_tester.create_and_test_config_common_properties()
       self.config_tester.create_and_test_config_to_json_string()
       self.config_tester.create_and_test_config_to_json_file()
       self.config_tester.create_and_test_config_from_and_save_pretrained()
       self.config_tester.create_and_test_config_with_num_labels()
       self.config_tester.check_config_can_be_init_without_params()
       self.config_tester.check_config_arguments_init()






    @unittest.skip(reason='Cvt does not output attentions' )
    def  A       (  self      :					List[Any] ):


       """simple docstring"""







       pass






    @unittest.skip(reason='Cvt does not use inputs_embeds' )
    def  A       (  self      :					str ):


       """simple docstring"""







       pass






    @unittest.skip(reason='Cvt does not support input and output embeddings' )
    def  A       (  self      :					List[str] ):


       """simple docstring"""







       pass






    @unittest.skipIf(
        not is_tf_available() or len(tf.config.list_physical_devices('GPU' ) ) == 0  ,						reason='TF does not support backprop for grouped convolutions on CPU.'  ,						)
    def  A       (  self      :					Dict ):


       """simple docstring"""







       super().test_dataset_conversion()






    @unittest.skipIf(
        not is_tf_available() or len(tf.config.list_physical_devices('GPU' ) ) == 0  ,						reason='TF does not support backprop for grouped convolutions on CPU.'  ,						)
    @slow
    def  A       (  self      :					str ):


       """simple docstring"""







       super().test_keras_fit()






    @unittest.skip(reason='Get `Failed to determine best cudnn convolution algo.` error after using TF 2.12+cuda 11.8' )
    def  A       (  self      :					List[str] ):


       """simple docstring"""







       UpperCamelCase					=  tf.keras.mixed_precision.Policy('mixed_float16' )
       tf.keras.mixed_precision.set_global_policy(UpperCamelCase__ )
       super().test_keras_fit()
       tf.keras.mixed_precision.set_global_policy('float32' )






    def  A       (  self      :					Optional[int] ):


       """simple docstring"""







       UpperCamelCase				,      UpperCamelCase					=  self.model_tester.prepare_config_and_inputs_for_common()

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

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






    def  A       (  self      :					Optional[Any] ):


       """simple docstring"""







       def check_hidden_states_output(UpperCamelCase__      :					List[str]  ,						UpperCamelCase__      :					Any  ,						UpperCamelCase__      :					Any ):
          UpperCamelCase					=  model_class(UpperCamelCase__ )

          UpperCamelCase					=  model(**self._prepare_for_class(UpperCamelCase__  ,						UpperCamelCase__ ) )
          UpperCamelCase					=  outputs.hidden_states

          UpperCamelCase					=  len(self.model_tester.depth )
          self.assertEqual(len(UpperCamelCase__ )  ,						UpperCamelCase__ )

          # verify the first hidden states (first block)
          self.assertListEqual(
              list(hidden_states[0].shape[-3:] )  ,						[
                  self.model_tester.embed_dim[0],
                  self.model_tester.image_size // 4,
                  self.model_tester.image_size // 4,
              ]  ,						)

       UpperCamelCase				,      UpperCamelCase					=  self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
          UpperCamelCase					=  True
          check_hidden_states_output(UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__ )

          # check that output_hidden_states also work using config
          del inputs_dict["output_hidden_states"]
          UpperCamelCase					=  True

          check_hidden_states_output(UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__ )






    def  A       (  self      :					Optional[Any] ):


       """simple docstring"""







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






    def  A       (  self      :					Dict ):


       """simple docstring"""







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






    @slow
    def  A       (  self      :					List[str] ):


       """simple docstring"""







       for model_name in TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          UpperCamelCase					=  TFCvtModel.from_pretrained(UpperCamelCase__ )
          self.assertIsNotNone(UpperCamelCase__ )







def 				__lowerCamelCase				(   )       ->	Tuple:






   """simple docstring"""
   UpperCamelCase					=  Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'    )
   return image

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






    """simple docstring"""





    @cached_property
    def  A       (  self      :					List[Any] ):


       """simple docstring"""







       return AutoImageProcessor.from_pretrained(TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST[0] )






    @slow
    def  A       (  self      :					List[str] ):


       """simple docstring"""







       UpperCamelCase					=  TFCvtForImageClassification.from_pretrained(TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST[0] )

       UpperCamelCase					=  self.default_image_processor
       UpperCamelCase					=  prepare_img()
       UpperCamelCase					=  image_processor(images=UpperCamelCase__  ,						return_tensors='tf' )

       # forward pass
       UpperCamelCase					=  model(**UpperCamelCase__ )

       # verify the logits
       UpperCamelCase					=  tf.TensorShape((1, 1_0_0_0) )
       self.assertEqual(outputs.logits.shape  ,						UpperCamelCase__ )

       UpperCamelCase					=  tf.constant([0.9_2_8_5, 0.9_0_1_5, -0.3_1_5_0] )
       self.assertTrue(np.allclose(outputs.logits[0, :3].numpy()  ,						UpperCamelCase__  ,						atol=1E-4 ) )






'''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 				__lowerCamelCase				(   A__    ,						A__    ,						A__=1e-1_2    )       ->	Dict:






   """simple docstring"""
   UpperCamelCase					=  jnp.divide(emb_a.T    ,						jnp.clip(jnp.linalg.norm(A__    ,						axis=1    )    ,						a_min=A__    )    ).T
   UpperCamelCase					=  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  SCREAMING_SNAKE_CASE (  nn.Module		):






    """simple docstring"""





    _SCREAMING_SNAKE_CASE	  =  42
    _SCREAMING_SNAKE_CASE	  =  jnp.floataa






    def  A       (  self      :					List[Any] ):


       """simple docstring"""







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

       UpperCamelCase					=  self.param('concept_embeds'  ,						jax.nn.initializers.ones  ,						(1_7, self.config.projection_dim) )
       UpperCamelCase					=  self.param(
           'special_care_embeds'  ,						jax.nn.initializers.ones  ,						(3, self.config.projection_dim) )

       UpperCamelCase					=  self.param('concept_embeds_weights'  ,						jax.nn.initializers.ones  ,						(1_7,) )
       UpperCamelCase					=  self.param('special_care_embeds_weights'  ,						jax.nn.initializers.ones  ,						(3,) )






    def __call__(  self      :					str  ,						UpperCamelCase__      :					List[str] ):


       """simple docstring"""







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

       UpperCamelCase					=  jax_cosine_distance(UpperCamelCase__  ,						self.special_care_embeds )
       UpperCamelCase					=  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
       UpperCamelCase					=  0.0

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

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

       return has_nsfw_concepts

class  SCREAMING_SNAKE_CASE (  _a		):






    """simple docstring"""





    _SCREAMING_SNAKE_CASE	  =  CLIPConfig
    _SCREAMING_SNAKE_CASE	  =  """clip_input"""
    _SCREAMING_SNAKE_CASE	  =  FlaxStableDiffusionSafetyCheckerModule






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


       """simple docstring"""







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






    def  A       (  self      :					int  ,						UpperCamelCase__      :					jax.random.KeyArray  ,						UpperCamelCase__      :					Tuple  ,						UpperCamelCase__      :					FrozenDict = None ):


       """simple docstring"""







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

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

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

       return random_params






    def __call__(  self      :					List[Any]  ,						UpperCamelCase__      :					Dict  ,						UpperCamelCase__      :					dict = None  ,						):


       """simple docstring"""







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

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

