





'''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_      (  SCREAMING_SNAKE_CASE__      :			str,			SCREAMING_SNAKE_CASE__      :			Any,			SCREAMING_SNAKE_CASE__      :			Union[str, Any]=1E-12   ) ->	Optional[Any]:
						UpperCAmelCase_			:      Any     						=  jnp.divide(emb_a.T,			jnp.clip(jnp.linalg.norm(SCREAMING_SNAKE_CASE__,			axis=1   ),			a_min=SCREAMING_SNAKE_CASE__   )   ).T
						UpperCAmelCase_			:      Optional[int]     						=  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 					__a      (nn.Module     ):
							__a					:   CLIPConfig
							__a					:   jnp.dtype									=       jnp.floataa

							def     UpperCAmelCase__	(							self							:     Union[str, Any]   )					->       Any:







													"""simple docstring"""
													UpperCAmelCase_			:      Union[str, Any]     						=  FlaxCLIPVisionModule(self.config.vision_config   )
													UpperCAmelCase_			:      Tuple     						=  nn.Dense(self.config.projection_dim       ,     use_bias=__magic_name__       ,     dtype=self.dtype   )

													UpperCAmelCase_			:      Optional[Any]     						=  self.param('''concept_embeds'''       ,     jax.nn.initializers.ones       ,     (17, self.config.projection_dim)   )
													UpperCAmelCase_			:      Union[str, Any]     						=  self.param(
													    '''special_care_embeds'''       ,     jax.nn.initializers.ones       ,     (3, self.config.projection_dim)   )

													UpperCAmelCase_			:      Tuple     						=  self.param('''concept_embeds_weights'''       ,     jax.nn.initializers.ones       ,     (17,)   )
													UpperCAmelCase_			:      Tuple     						=  self.param('''special_care_embeds_weights'''       ,     jax.nn.initializers.ones       ,     (3,)   )



							def __call__(							self							:     Any       ,     __magic_name__							:     Tuple   )					->       Optional[int]:







													"""simple docstring"""
													UpperCAmelCase_			:      Dict     						=  self.vision_model(__magic_name__   )[1]
													UpperCAmelCase_			:      Optional[Any]     						=  self.visual_projection(__magic_name__   )

													UpperCAmelCase_			:      Union[str, Any]     						=  jax_cosine_distance(__magic_name__       ,     self.special_care_embeds   )
													UpperCAmelCase_			:      Optional[int]     						=  jax_cosine_distance(__magic_name__       ,     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_			:      Any     						=  0.0

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

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

													return has_nsfw_concepts





class 					__a      (lowerCamelCase     ):
							__a					:   str									=       CLIPConfig
							__a					:   Optional[Any]									=       "clip_input"
							__a					:   List[str]									=       FlaxStableDiffusionSafetyCheckerModule

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







													"""simple docstring"""
													if input_shape is None:
																			UpperCAmelCase_			:      List[str]     						=  (1, 2_24, 2_24, 3)
													UpperCAmelCase_			:      Dict     						=  self.module_class(config=__magic_name__       ,     dtype=__magic_name__       ,     **__magic_name__   )
													super().__init__(__magic_name__       ,     __magic_name__       ,     input_shape=__magic_name__       ,     seed=__magic_name__       ,     dtype=__magic_name__       ,     _do_init=_do_init   )

							def     UpperCAmelCase__	(							self							:     Tuple       ,     __magic_name__							:     jax.random.KeyArray       ,     __magic_name__							:     Tuple       ,     __magic_name__							:     FrozenDict = None   )					->       FrozenDict:







													"""simple docstring"""
													# init input tensor
													UpperCAmelCase_			:      Dict     						=  jax.random.normal(__magic_name__       ,     __magic_name__   )

													UpperCAmelCase_						,       UpperCAmelCase_			:      str     						=  jax.random.split(__magic_name__   )
													UpperCAmelCase_			:      Union[str, Any]     						=  {'''params''': params_rng, '''dropout''': dropout_rng}

													UpperCAmelCase_			:      Optional[Any]     						=  self.module.init(__magic_name__       ,     __magic_name__   )['''params''']

													return random_params



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







													"""simple docstring"""
													UpperCAmelCase_			:      Optional[Any]     						=  jnp.transpose(__magic_name__       ,     (0, 2, 3, 1)   )

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












'''simple docstring'''







snake_case_     :       int						=  {
    "Pillow": "Pillow",
    "accelerate": "accelerate>=0.11.0",
    "compel": "compel==0.1.8",
    "black": "black~=23.1",
    "datasets": "datasets",
    "filelock": "filelock",
    "flax": "flax>=0.4.1",
    "hf-doc-builder": "hf-doc-builder>=0.3.0",
    "huggingface-hub": "huggingface-hub>=0.13.2",
    "requests-mock": "requests-mock==1.10.0",
    "importlib_metadata": "importlib_metadata",
    "invisible-watermark": "invisible-watermark",
    "isort": "isort>=5.5.4",
    "jax": "jax>=0.2.8,!=0.3.2",
    "jaxlib": "jaxlib>=0.1.65",
    "Jinja2": "Jinja2",
    "k-diffusion": "k-diffusion>=0.0.12",
    "torchsde": "torchsde",
    "note_seq": "note_seq",
    "librosa": "librosa",
    "numpy": "numpy",
    "omegaconf": "omegaconf",
    "parameterized": "parameterized",
    "protobuf": "protobuf>=3.20.3,<4",
    "pytest": "pytest",
    "pytest-timeout": "pytest-timeout",
    "pytest-xdist": "pytest-xdist",
    "ruff": "ruff>=0.0.241",
    "safetensors": "safetensors",
    "sentencepiece": "sentencepiece>=0.1.91,!=0.1.92",
    "scipy": "scipy",
    "onnx": "onnx",
    "regex": "regex!=2019.12.17",
    "requests": "requests",
    "tensorboard": "tensorboard",
    "torch": "torch>=1.4",
    "torchvision": "torchvision",
    "transformers": "transformers>=4.25.1",
    "urllib3": "urllib3<=2.0.0",
}





