
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     A__   (  lowerCamelCase       ,      lowerCamelCase       ,      lowerCamelCase=1E-1_2  )					->		Tuple:
				UpperCamelCase_:							Dict		=  jnp.divide(emb_a.T       ,      jnp.clip(jnp.linalg.norm(lowerCamelCase       ,      axis=1  )       ,      a_min=lowerCamelCase  )  ).T
				UpperCamelCase_:							Optional[Any]		=  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     _UpperCamelCase			(   nn.Module					):






			'''simple docstring'''

			__UpperCamelCase :  CLIPConfig
			__UpperCamelCase :  jnp.dtype    =	jnp.floataa

			def  lowerCAmelCase__		(    self	:    List[str] ):
							UpperCamelCase_:							Tuple		=  FlaxCLIPVisionModule(self.config.vision_config )
							UpperCamelCase_:							List[str]		=  nn.Dense(self.config.projection_dim  ,		use_bias=snake_case_  ,		dtype=self.dtype )

							UpperCamelCase_:							List[str]		=  self.param("""concept_embeds"""  ,		jax.nn.initializers.ones  ,		(17, self.config.projection_dim) )
							UpperCamelCase_:							Tuple		=  self.param(
							    """special_care_embeds"""  ,		jax.nn.initializers.ones  ,		(3, self.config.projection_dim) )

							UpperCamelCase_:							List[str]		=  self.param("""concept_embeds_weights"""  ,		jax.nn.initializers.ones  ,		(17,) )
							UpperCamelCase_:							Dict		=  self.param("""special_care_embeds_weights"""  ,		jax.nn.initializers.ones  ,		(3,) )




			def __call__(    self	:    Any  ,		snake_case_	:    str ):
							UpperCamelCase_:							Any		=  self.vision_model(snake_case_ )[1]
							UpperCamelCase_:							Union[str, Any]		=  self.visual_projection(snake_case_ )

							UpperCamelCase_:							Dict		=  jax_cosine_distance(snake_case_  ,		self.special_care_embeds )
							UpperCamelCase_:							str		=  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
							UpperCamelCase_:							Optional[Any]		=  0.0

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

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

							return has_nsfw_concepts




class     _UpperCamelCase			(   _A					):






			'''simple docstring'''

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

			def __init__(    self	:    Dict  ,		snake_case_	:    CLIPConfig  ,		snake_case_	:    Optional[Tuple] = None  ,		snake_case_	:    int = 0  ,		snake_case_	:    jnp.dtype = jnp.floataa  ,		snake_case_	:    bool = True  ,		**snake_case_	:    Optional[Any]  ,		):
							if input_shape is None:
											UpperCamelCase_:							Union[str, Any]		=  (1, 224, 224, 3)
							UpperCamelCase_:							Optional[Any]		=  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  lowerCAmelCase__		(    self	:    str  ,		snake_case_	:    jax.random.KeyArray  ,		snake_case_	:    Tuple  ,		snake_case_	:    FrozenDict = None ):
							# init input tensor
							UpperCamelCase_:							List[Any]		=  jax.random.normal(snake_case_  ,		snake_case_ )

							UpperCamelCase_,				UpperCamelCase_:							Any		=  jax.random.split(snake_case_ )
							UpperCamelCase_:							Optional[int]		=  {"""params""": params_rng, """dropout""": dropout_rng}

							UpperCamelCase_:							Tuple		=  self.module.init(snake_case_  ,		snake_case_ )["""params"""]

							return random_params




			def __call__(    self	:    str  ,		snake_case_	:    Dict  ,		snake_case_	:    dict = None  ,		):
							UpperCamelCase_:							Optional[int]		=  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={}  ,		)








from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCamelCase_  :	Optional[int] 				=					logging.get_logger(__name__)

lowerCamelCase_  :	Optional[int] 				=					{"""ctrl""": """https://huggingface.co/ctrl/resolve/main/config.json"""}

class     _UpperCamelCase			(   _A					):






			'''simple docstring'''

			__UpperCamelCase :  int    =	"""ctrl"""
			__UpperCamelCase :  Dict    =	["""past_key_values"""]
			__UpperCamelCase :  List[str]    =	{
			    """max_position_embeddings""": """n_positions""",
			    """hidden_size""": """n_embd""",
			    """num_attention_heads""": """n_head""",
			    """num_hidden_layers""": """n_layer""",
			}

			def __init__(    self	:    Dict  ,		snake_case_	:    Any=24_6534  ,		snake_case_	:    Dict=256  ,		snake_case_	:    str=1280  ,		snake_case_	:    Optional[int]=8192  ,		snake_case_	:    Union[str, Any]=48  ,		snake_case_	:    Any=16  ,		snake_case_	:    Optional[int]=0.1  ,		snake_case_	:    Any=0.1  ,		snake_case_	:    Any=1e-6  ,		snake_case_	:    Optional[Any]=0.02  ,		snake_case_	:    Optional[int]=True  ,		**snake_case_	:    Union[str, Any]  ,		):
							UpperCamelCase_:							Union[str, Any]		=  vocab_size
							UpperCamelCase_:							Union[str, Any]		=  n_positions
							UpperCamelCase_:							Optional[int]		=  n_embd
							UpperCamelCase_:							int		=  n_layer
							UpperCamelCase_:							str		=  n_head
							UpperCamelCase_:							Optional[int]		=  dff
							UpperCamelCase_:							Optional[Any]		=  resid_pdrop
							UpperCamelCase_:							Union[str, Any]		=  embd_pdrop
							UpperCamelCase_:							List[str]		=  layer_norm_epsilon
							UpperCamelCase_:							Optional[Any]		=  initializer_range

							UpperCamelCase_:							Optional[Any]		=  use_cache

							super().__init__(**snake_case_ )






