



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		(	__lowerCamelCase						:     Optional[int]					, __lowerCamelCase						:     Tuple					, __lowerCamelCase						:     str=1E-1_2						)	->    int:
					__lowerCAmelCase						 =jnp.divide(emb_a.T					, jnp.clip(jnp.linalg.norm(__lowerCamelCase					, axis=1						)					, a_min=__lowerCamelCase						)						).T
					__lowerCAmelCase						 =jnp.divide(emb_a.T					, jnp.clip(jnp.linalg.norm(__lowerCamelCase					, axis=1						)					, a_min=__lowerCamelCase						)						).T
					return jnp.matmul(__lowerCamelCase					, norm_emb_a.T						)






class     __a			(  nn.Module					):
	SCREAMING_SNAKE_CASE  	=							42
	SCREAMING_SNAKE_CASE  	=							jnp.floataa



	def 		UpperCamelCase				(    self				:     List[str])->				List[Any]:
						__lowerCAmelCase						 =FlaxCLIPVisionModule(self.config.vision_config)
						__lowerCAmelCase						 =nn.Dense(self.config.projection_dim					,						use_bias=snake_case_					,						dtype=self.dtype)

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

						__lowerCAmelCase						 =self.param("""concept_embeds_weights"""					,						jax.nn.initializers.ones					,						(17,))
						__lowerCAmelCase						 =self.param("""special_care_embeds_weights"""					,						jax.nn.initializers.ones					,						(3,))



	def __call__(    self				:     Tuple					,						snake_case_				:     int)->				int:
						__lowerCAmelCase						 =self.vision_model(snake_case_)[1]
						__lowerCAmelCase						 =self.visual_projection(snake_case_)

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

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

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

						return has_nsfw_concepts






class     __a			(  SCREAMING_SNAKE_CASE					):
	SCREAMING_SNAKE_CASE  	=							CLIPConfig
	SCREAMING_SNAKE_CASE  	=							"clip_input"
	SCREAMING_SNAKE_CASE  	=							FlaxStableDiffusionSafetyCheckerModule



	def __init__(    self				:     Optional[int]					,						snake_case_				:     CLIPConfig					,						snake_case_				:     Optional[Tuple] = None					,						snake_case_				:     int = 0					,						snake_case_				:     jnp.dtype = jnp.floataa					,						snake_case_				:     bool = True					,						**snake_case_				:     Any					,						)->				Any:
						if input_shape is None:
											__lowerCAmelCase						 =(1, 2_24, 2_24, 3)
						__lowerCAmelCase						 =self.module_class(config=snake_case_					,						dtype=snake_case_					,						**snake_case_)
						super().__init__(snake_case_					,						snake_case_					,						input_shape=snake_case_					,						seed=snake_case_					,						dtype=snake_case_					,						_do_init=_do_init)



	def 		UpperCamelCase				(    self				:     Optional[Any]					,						snake_case_				:     jax.random.KeyArray					,						snake_case_				:     Tuple					,						snake_case_				:     FrozenDict = None)->				FrozenDict:
						# init input tensor
						__lowerCAmelCase						 =jax.random.normal(snake_case_					,						snake_case_)

						__lowerCAmelCase ,						__lowerCAmelCase						 =jax.random.split(snake_case_)
						__lowerCAmelCase						 ={"""params""": params_rng, """dropout""": dropout_rng}

						__lowerCAmelCase						 =self.module.init(snake_case_					,						snake_case_)["""params"""]

						return random_params



	def __call__(    self				:     List[Any]					,						snake_case_				:     Union[str, Any]					,						snake_case_				:     dict = None					,						)->				Optional[Any]:
						__lowerCAmelCase						 =jnp.transpose(snake_case_					,						(0, 2, 3, 1))

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









def        __lowerCAmelCase		(	__lowerCamelCase						:     int					, __lowerCamelCase						:     int						)	->    bool:
					return numa ^ numa < 0


if __name__ == "__main__":
				import doctest

				doctest.testmod()




