


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      ,      SCREAMING_SNAKE_CASE      ,      SCREAMING_SNAKE_CASE=1E-12      )->						List[str]:




      """simple docstring"""

      snake_case_ =   jnp.divide(emb_a.T      ,      jnp.clip(jnp.linalg.norm(SCREAMING_SNAKE_CASE      ,      axis=1      )      ,      a_min=SCREAMING_SNAKE_CASE      )      ).T
      snake_case_ =   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'''

 __snake_case				=						42
 __snake_case				=						jnp.floataa



 def 		UpperCamelCase__  (			self      ):
       snake_case_ =   FlaxCLIPVisionModule(self.config.vision_config      )
       snake_case_ =   nn.Dense(self.config.projection_dim	,   use_bias=_UpperCAmelCase	,   dtype=self.dtype      )

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

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





 def __call__(			self	,   _UpperCAmelCase      ):
       snake_case_ =   self.vision_model(_UpperCAmelCase      )[1]
       snake_case_ =   self.visual_projection(_UpperCAmelCase      )

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

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

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

       return has_nsfw_concepts







class lowerCAmelCase_		(				lowerCamelCase__ ):

 '''simple docstring'''

 __snake_case				=						CLIPConfig
 __snake_case				=						"clip_input"
 __snake_case				=						FlaxStableDiffusionSafetyCheckerModule



 def __init__(			self	,   _UpperCAmelCase	,   _UpperCAmelCase = None	,   _UpperCAmelCase = 0	,   _UpperCAmelCase = jnp.floataa	,   _UpperCAmelCase = True	,   **_UpperCAmelCase	,   ):
       if input_shape is None:
             snake_case_ =   (1, 2_24, 2_24, 3)
       snake_case_ =   self.module_class(config=_UpperCAmelCase	,   dtype=_UpperCAmelCase	,   **_UpperCAmelCase      )
       super().__init__(_UpperCAmelCase	,   _UpperCAmelCase	,   input_shape=_UpperCAmelCase	,   seed=_UpperCAmelCase	,   dtype=_UpperCAmelCase	,   _do_init=_do_init      )



 def 		UpperCamelCase__  (			self	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase = None      ):
       # init input tensor
       snake_case_ =   jax.random.normal(_UpperCAmelCase	,   _UpperCAmelCase      )

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

       snake_case_ =   self.module.init(_UpperCAmelCase	,   _UpperCAmelCase      )['''params''']

       return random_params





 def __call__(			self	,   _UpperCAmelCase	,   _UpperCAmelCase = None	,   ):
       snake_case_ =   jnp.transpose(_UpperCAmelCase	,   (0, 2, 3, 1)      )

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



from ...utils import is_note_seq_available, is_transformers_available, is_torch_available
from ...utils import OptionalDependencyNotAvailable


try:
   if not (is_transformers_available() and is_torch_available()):
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   from ...utils.dummy_torch_and_transformers_objects import *  # noqa F403
else:
   from .notes_encoder import SpectrogramNotesEncoder
   from .continous_encoder import SpectrogramContEncoder
   from .pipeline_spectrogram_diffusion import (
       SpectrogramContEncoder,
       SpectrogramDiffusionPipeline,
       TaFilmDecoder,
   )

try:
   if not (is_transformers_available() and is_torch_available() and is_note_seq_available()):
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   from ...utils.dummy_transformers_and_torch_and_note_seq_objects import *  # noqa F403
else:
   from .midi_utils import MidiProcessor