





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 (  SCREAMING_SNAKE_CASE__      :  List[str]			,				SCREAMING_SNAKE_CASE__      :  Dict			,				SCREAMING_SNAKE_CASE__      :  List[Any]=1E-1_2       )	->    Union[str, Any]:

							'''simple docstring'''






							SCREAMING_SNAKE_CASE__       :      Tuple     =						jnp.divide(emb_a.T			,				jnp.clip(jnp.linalg.norm(SCREAMING_SNAKE_CASE__			,				axis=1       )			,				a_min=SCREAMING_SNAKE_CASE__       )       ).T
							SCREAMING_SNAKE_CASE__       :      Tuple     =						jnp.divide(emb_a.T			,				jnp.clip(jnp.linalg.norm(SCREAMING_SNAKE_CASE__			,				axis=1       )			,				a_min=SCREAMING_SNAKE_CASE__       )       ).T
							return jnp.matmul(SCREAMING_SNAKE_CASE__			,				norm_emb_a.T       )


class 						lowerCamelCase       (nn.Module      ):







							"""simple docstring"""




							UpperCAmelCase_      						=   42
							UpperCAmelCase_      						=   jnp.floataa






							def  A_     ( self     :					List[Any]						)				->			Optional[int]:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      int     =						FlaxCLIPVisionModule(self.config.vision_config						)
														SCREAMING_SNAKE_CASE__       :      Tuple     =						nn.Dense(self.config.projection_dim,			use_bias=_UpperCAmelCase,			dtype=self.dtype						)

														SCREAMING_SNAKE_CASE__       :      str     =						self.param("concept_embeds",			jax.nn.initializers.ones,			(1_7, self.config.projection_dim)						)
														SCREAMING_SNAKE_CASE__       :      Any     =						self.param(
														    "special_care_embeds",			jax.nn.initializers.ones,			(3, self.config.projection_dim)						)

														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						self.param("concept_embeds_weights",			jax.nn.initializers.ones,			(1_7,)						)
														SCREAMING_SNAKE_CASE__       :      str     =						self.param("special_care_embeds_weights",			jax.nn.initializers.ones,			(3,)						)






							def __call__( self     :					int,			_UpperCAmelCase     :					str						)				->			int:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						self.vision_model(_UpperCAmelCase						)[1]
														SCREAMING_SNAKE_CASE__       :      Tuple     =						self.visual_projection(_UpperCAmelCase						)

														SCREAMING_SNAKE_CASE__       :      Any     =						jax_cosine_distance(_UpperCAmelCase,			self.special_care_embeds						)
														SCREAMING_SNAKE_CASE__       :      Any     =						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
														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						0.0

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

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

														return has_nsfw_concepts


class 						lowerCamelCase       (__lowerCamelCase      ):







							"""simple docstring"""




							UpperCAmelCase_      						=   CLIPConfig
							UpperCAmelCase_      						=   "clip_input"
							UpperCAmelCase_      						=   FlaxStableDiffusionSafetyCheckerModule






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





														"""simple docstring"""
														if input_shape is None:
																					SCREAMING_SNAKE_CASE__       :      int     =						(1, 2_2_4, 2_2_4, 3)
														SCREAMING_SNAKE_CASE__       :      List[Any]     =						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     :					Tuple,			_UpperCAmelCase     :					jax.random.KeyArray,			_UpperCAmelCase     :					Tuple,			_UpperCAmelCase     :					FrozenDict = None						)				->			FrozenDict:





														"""simple docstring"""
														# init input tensor
														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						jax.random.normal(_UpperCAmelCase,			_UpperCAmelCase						)

														SCREAMING_SNAKE_CASE__   ,SCREAMING_SNAKE_CASE__       :      Tuple     =						jax.random.split(_UpperCAmelCase						)
														SCREAMING_SNAKE_CASE__       :      Any     =						{"params": params_rng, "dropout": dropout_rng}

														SCREAMING_SNAKE_CASE__       :      Any     =						self.module.init(_UpperCAmelCase,			_UpperCAmelCase						)["params"]

														return random_params






							def __call__( self     :					List[Any],			_UpperCAmelCase     :					Optional[Any],			_UpperCAmelCase     :					dict = None,			)				->			Any:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Tuple     =						jnp.transpose(_UpperCAmelCase,			(0, 2, 3, 1)						)

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













def _a (  SCREAMING_SNAKE_CASE__      :  str       )	->    str:

							'''simple docstring'''






							if not all(char in "01" for char in bin_string       ):
														raise ValueError("Non-binary value was passed to the function"       )
							if not bin_string:
														raise ValueError("Empty string was passed to the function"       )
							SCREAMING_SNAKE_CASE__       :      List[Any]     =						""
							while len(SCREAMING_SNAKE_CASE__       ) % 3 != 0:
														SCREAMING_SNAKE_CASE__       :      str     =						"0" + bin_string
							SCREAMING_SNAKE_CASE__       :      List[Any]     =						[
							    bin_string[index : index + 3]
							    for index in range(len(SCREAMING_SNAKE_CASE__       )       )
							    if index % 3 == 0
							]
							for bin_group in bin_string_in_3_list:
														SCREAMING_SNAKE_CASE__       :      List[Any]     =						0
														for index, val in enumerate(SCREAMING_SNAKE_CASE__       ):
																					oct_val += int(2 ** (2 - index) * int(SCREAMING_SNAKE_CASE__       )       )
														oct_string += str(SCREAMING_SNAKE_CASE__       )
							return oct_string


if __name__ == "__main__":
		from doctest import testmod

		testmod()






