
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					( _lowercase						,     _lowercase						,     _lowercase=1e-12			)       ->      List[Any]:
     UpperCamelCase						    =					jnp.divide(emb_a.T						,     jnp.clip(jnp.linalg.norm(_lowercase						,     axis=1			)						,     a_min=_lowercase			)			).T
     UpperCamelCase						    =					jnp.divide(emb_a.T						,     jnp.clip(jnp.linalg.norm(_lowercase						,     axis=1			)						,     a_min=_lowercase			)			).T
     return jnp.matmul(_lowercase						,     norm_emb_a.T			)
class       _lowerCAmelCase  (		nn.Module							):

 """simple docstring"""
 SCREAMING_SNAKE_CASE_				:					CLIPConfig
 SCREAMING_SNAKE_CASE_				:					jnp.dtype       						=jnp.floataa

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






      """simple docstring"""



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

      UpperCamelCase						    =					self.param('concept_embeds'					, jax.nn.initializers.ones					, (17, 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					, (17,)   )
      UpperCamelCase						    =					self.param('special_care_embeds_weights'					, jax.nn.initializers.ones					, (3,)   )






 def __call__(						self			:			Any					, SCREAMING_SNAKE_CASE__			:			Tuple   ):






      """simple docstring"""



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

      UpperCamelCase						    =					jax_cosine_distance(SCREAMING_SNAKE_CASE__					, self.special_care_embeds   )
      UpperCamelCase						    =					jax_cosine_distance(SCREAMING_SNAKE_CASE__					, 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(SCREAMING_SNAKE_CASE__					, 3   )
      UpperCamelCase						    =					jnp.any(special_scores > 0					, axis=1					, keepdims=SCREAMING_SNAKE_CASE__   )
      # Use a lower threshold if an image has any special care concept
      UpperCamelCase						    =					is_special_care * 0.01

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

      return has_nsfw_concepts




class       _lowerCAmelCase  (		__magic_name__							):

 """simple docstring"""
 SCREAMING_SNAKE_CASE_				:					List[str]       						=CLIPConfig
 SCREAMING_SNAKE_CASE_				:					Any       						="clip_input"
 SCREAMING_SNAKE_CASE_				:					List[Any]       						=FlaxStableDiffusionSafetyCheckerModule

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






      """simple docstring"""



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

 def      __lowerCAmelCase  (						self			:			Any					, SCREAMING_SNAKE_CASE__			:			jax.random.KeyArray					, SCREAMING_SNAKE_CASE__			:			Tuple					, SCREAMING_SNAKE_CASE__			:			FrozenDict = None   ):






      """simple docstring"""



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

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

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

      return random_params






 def __call__(						self			:			Tuple					, SCREAMING_SNAKE_CASE__			:			Union[str, Any]					, SCREAMING_SNAKE_CASE__			:			dict = None					, ):






      """simple docstring"""



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

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




import json
import os
import tempfile

from transformers.testing_utils import check_json_file_has_correct_format
class       _lowerCAmelCase  :

 """simple docstring"""
 SCREAMING_SNAKE_CASE_				:					Optional[Any]       						=None

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






      """simple docstring"""



      UpperCamelCase						    =					self.feature_extraction_class(**self.feat_extract_dict   )
      UpperCamelCase						    =					json.loads(feat_extract.to_json_string()   )
      for key, value in self.feat_extract_dict.items():
           self.assertEqual(obj[key]					, SCREAMING_SNAKE_CASE__   )

 def      __lowerCAmelCase  (						self			:			Tuple   ):






      """simple docstring"""



      UpperCamelCase						    =					self.feature_extraction_class(**self.feat_extract_dict   )

      with tempfile.TemporaryDirectory() as tmpdirname:
           UpperCamelCase						    =					os.path.join(SCREAMING_SNAKE_CASE__					, 'feat_extract.json'   )
           feat_extract_first.to_json_file(SCREAMING_SNAKE_CASE__   )
           UpperCamelCase						    =					self.feature_extraction_class.from_json_file(SCREAMING_SNAKE_CASE__   )

      self.assertEqual(feat_extract_second.to_dict()					, feat_extract_first.to_dict()   )

 def      __lowerCAmelCase  (						self			:			Tuple   ):






      """simple docstring"""



      UpperCamelCase						    =					self.feature_extraction_class(**self.feat_extract_dict   )

      with tempfile.TemporaryDirectory() as tmpdirname:
           UpperCamelCase						    =					feat_extract_first.save_pretrained(SCREAMING_SNAKE_CASE__   )[0]
           check_json_file_has_correct_format(SCREAMING_SNAKE_CASE__   )
           UpperCamelCase						    =					self.feature_extraction_class.from_pretrained(SCREAMING_SNAKE_CASE__   )

      self.assertEqual(feat_extract_second.to_dict()					, feat_extract_first.to_dict()   )






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






      """simple docstring"""



      UpperCamelCase						    =					self.feature_extraction_class()
      self.assertIsNotNone(SCREAMING_SNAKE_CASE__   )


