
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  __snake_case							(	__magic_name__		,    __magic_name__		,    __magic_name__=1e-12						):







				'''simple docstring'''


				lowercase         =      jnp.divide(emb_a.T		,    jnp.clip(jnp.linalg.norm(__magic_name__		,    axis=1						)		,    a_min=__magic_name__						)						).T
				lowercase         =      jnp.divide(emb_a.T		,    jnp.clip(jnp.linalg.norm(__magic_name__		,    axis=1						)		,    a_min=__magic_name__						)						).T
				return jnp.matmul(__magic_name__		,    norm_emb_a.T						)





class   UpperCamelCase_      (      nn.Module		):







						'''simple docstring'''



						UpperCamelCase      :  CLIPConfig
						UpperCamelCase      :  jnp.dtype				   =					jnp.floataa





						def        SCREAMING_SNAKE_CASE(			self							:List[str]     )  ->Optional[int]:
										lowercase         =      FlaxCLIPVisionModule(self.config.vision_config     )
										lowercase         =      nn.Dense(self.config.projection_dim		,    use_bias=lowerCAmelCase__		,    dtype=self.dtype     )

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

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





						def __call__(			self							:Optional[Any]		,    lowerCAmelCase__							:Optional[int]     )  ->int:
										lowercase         =      self.vision_model(lowerCAmelCase__     )[1]
										lowercase         =      self.visual_projection(lowerCAmelCase__     )

										lowercase         =      jax_cosine_distance(lowerCAmelCase__		,    self.special_care_embeds     )
										lowercase         =      jax_cosine_distance(lowerCAmelCase__		,    self.concept_embeds     )

										# increase this value to create a stronger `nfsw` filter
										# at the cost of increasing the possibility of filtering benign image inputs
										lowercase         =      0.0

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

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

										return has_nsfw_concepts





class   UpperCamelCase_      (      __a		):







						'''simple docstring'''



						UpperCamelCase      :  List[Any]				   =					CLIPConfig
						UpperCamelCase      :  Tuple				   =					'''clip_input'''
						UpperCamelCase      :  List[Any]				   =					FlaxStableDiffusionSafetyCheckerModule





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





						def        SCREAMING_SNAKE_CASE(			self							:Optional[int]		,    lowerCAmelCase__							:jax.random.KeyArray		,    lowerCAmelCase__							:Tuple		,    lowerCAmelCase__							:FrozenDict = None     )  ->FrozenDict:
										# init input tensor
										lowercase         =      jax.random.normal(lowerCAmelCase__		,    lowerCAmelCase__     )

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

										lowercase         =      self.module.init(lowerCAmelCase__		,    lowerCAmelCase__     )["params"]

										return random_params





						def __call__(			self							:Optional[Any]		,    lowerCAmelCase__							:Tuple		,    lowerCAmelCase__							:dict = None		,    )  ->str:
										lowercase         =      jnp.transpose(lowerCAmelCase__		,    (0, 2, 3, 1)     )

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


from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


_snake_case      :							Union[str, Any]     	=    {"configuration_wavlm": ["WAVLM_PRETRAINED_CONFIG_ARCHIVE_MAP", "WavLMConfig"]}

try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						_snake_case      :							Union[str, Any]     	=    [
						    "WAVLM_PRETRAINED_MODEL_ARCHIVE_LIST",
						    "WavLMForAudioFrameClassification",
						    "WavLMForCTC",
						    "WavLMForSequenceClassification",
						    "WavLMForXVector",
						    "WavLMModel",
						    "WavLMPreTrainedModel",
						]

if TYPE_CHECKING:
						from .configuration_wavlm import WAVLM_PRETRAINED_CONFIG_ARCHIVE_MAP, WavLMConfig

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_wavlm import (
												    WAVLM_PRETRAINED_MODEL_ARCHIVE_LIST,
												    WavLMForAudioFrameClassification,
												    WavLMForCTC,
												    WavLMForSequenceClassification,
												    WavLMForXVector,
												    WavLMModel,
												    WavLMPreTrainedModel,
												)

else:
						import sys

						_snake_case      :							str     	=    _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
