






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     (__lowercase     ,   __lowercase     ,   __lowercase=1e-12):
   UpperCamelCase_						=  jnp.divide(emb_a.T     ,   jnp.clip(jnp.linalg.norm(__lowercase     ,   axis=1)     ,   a_min=__lowercase)).T
   UpperCamelCase_						=  jnp.divide(emb_a.T     ,   jnp.clip(jnp.linalg.norm(__lowercase     ,   axis=1)     ,   a_min=__lowercase)).T
   return jnp.matmul(__lowercase     ,   norm_emb_a.T)


class     _a							(	nn.Module     ):






  """simple docstring"""





  A_						       = 42
  A_						       = jnp.floataa




  def  _UpperCAmelCase  (  self			)      ->			Dict:
     UpperCamelCase_						=  FlaxCLIPVisionModule(self.config.vision_config			)
     UpperCamelCase_						=  nn.Dense(self.config.projection_dim		,		use_bias=_UpperCAmelCase		,		dtype=self.dtype			)

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

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




  def __call__(  self		,		_UpperCAmelCase			)      ->			List[str]:
     UpperCamelCase_						=  self.vision_model(_UpperCAmelCase			)[1]
     UpperCamelCase_						=  self.visual_projection(_UpperCAmelCase			)

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

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

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

     return has_nsfw_concepts




class     _a							(	UpperCAmelCase__     ):






  """simple docstring"""





  A_						       = CLIPConfig
  A_						       = """clip_input"""
  A_						       = FlaxStableDiffusionSafetyCheckerModule




  def __init__(  self		,		_UpperCAmelCase		,		_UpperCAmelCase = None		,		_UpperCAmelCase = 0		,		_UpperCAmelCase = jnp.floataa		,		_UpperCAmelCase = True		,		**_UpperCAmelCase		,		)      ->			Union[str, Any]:
     if input_shape is None:
        UpperCamelCase_						=  (1, 224, 224, 3)
     UpperCamelCase_						=  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			)      ->			FrozenDict:
     # init input tensor
     UpperCamelCase_						=  jax.random.normal(_UpperCAmelCase		,		_UpperCAmelCase			)

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

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

     return random_params




  def __call__(  self		,		_UpperCAmelCase		,		_UpperCAmelCase = None		,		)      ->			Optional[int]:
     UpperCamelCase_						=  jnp.transpose(_UpperCAmelCase		,		(0, 2, 3, 1)			)

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











import unittest
from typing import Tuple

import torch

from diffusers.utils import floats_tensor, randn_tensor, torch_all_close, torch_device
from diffusers.utils.testing_utils import require_torch





@require_torch
class 			lowerCamelCase_				:


 '''simple docstring'''







 @property
 def  lowerCAmelCase_    (				self    :						Optional[Any]		):
       return self.get_dummy_input()



 @property
 def  lowerCAmelCase_    (				self    :						Union[str, Any]		):
       if self.block_type == "down":
             return (4, 32, 16, 16)
       elif self.block_type == "mid":
             return (4, 32, 32, 32)
       elif self.block_type == "up":
             return (4, 32, 64, 64)

       raise ValueError(F"'{self.block_type}' is not a supported block_type. Set it to 'up', 'mid', or 'down'."		)



 def  lowerCAmelCase_    (				self    :						List[str]							,							_lowerCAmelCase    :						Optional[Any]=True							,							_lowerCAmelCase    :						str=False							,							_lowerCAmelCase    :						Optional[int]=False							,							_lowerCAmelCase    :						Dict=False							,							):
       SCREAMING_SNAKE_CASE_      		=     4
       SCREAMING_SNAKE_CASE_      		=     32
       SCREAMING_SNAKE_CASE_      		=     (32, 32)

       SCREAMING_SNAKE_CASE_      		=     torch.manual_seed(0		)
       SCREAMING_SNAKE_CASE_      		=     torch.device(_lowerCAmelCase		)
       SCREAMING_SNAKE_CASE_      		=     (batch_size, num_channels) + sizes
       SCREAMING_SNAKE_CASE_      		=     randn_tensor(_lowerCAmelCase							,							generator=_lowerCAmelCase							,							device=_lowerCAmelCase		)
       SCREAMING_SNAKE_CASE_      		=     {'hidden_states': hidden_states}

       if include_temb:
             SCREAMING_SNAKE_CASE_      		=     128
             SCREAMING_SNAKE_CASE_      		=     randn_tensor((batch_size, temb_channels)							,							generator=_lowerCAmelCase							,							device=_lowerCAmelCase		)

       if include_res_hidden_states_tuple:
             SCREAMING_SNAKE_CASE_      		=     torch.manual_seed(1		)
             SCREAMING_SNAKE_CASE_      		=     (randn_tensor(_lowerCAmelCase							,							generator=_lowerCAmelCase							,							device=_lowerCAmelCase		),)

       if include_encoder_hidden_states:
             SCREAMING_SNAKE_CASE_      		=     floats_tensor((batch_size, 32, 32)		).to(_lowerCAmelCase		)

       if include_skip_sample:
             SCREAMING_SNAKE_CASE_      		=     randn_tensor(((batch_size, 3) + sizes)							,							generator=_lowerCAmelCase							,							device=_lowerCAmelCase		)

       return dummy_input



 def  lowerCAmelCase_    (				self    :						Optional[int]		):
       SCREAMING_SNAKE_CASE_      		=     {
           'in_channels': 32,
           'out_channels': 32,
           'temb_channels': 128,
       }
       if self.block_type == "up":
             SCREAMING_SNAKE_CASE_      		=     32

       if self.block_type == "mid":
             init_dict.pop('out_channels'		)

       SCREAMING_SNAKE_CASE_      		=     self.dummy_input
       return init_dict, inputs_dict



 def  lowerCAmelCase_    (				self    :						Dict							,							_lowerCAmelCase    :						Optional[Any]		):
       SCREAMING_SNAKE_CASE_		,					SCREAMING_SNAKE_CASE_      		=     self.prepare_init_args_and_inputs_for_common()
       SCREAMING_SNAKE_CASE_      		=     self.block_class(**_lowerCAmelCase		)
       unet_block.to(_lowerCAmelCase		)
       unet_block.eval()

       with torch.no_grad():
             SCREAMING_SNAKE_CASE_      		=     unet_block(**_lowerCAmelCase		)

       if isinstance(_lowerCAmelCase							,							_lowerCAmelCase		):
             SCREAMING_SNAKE_CASE_      		=     output[0]

       self.assertEqual(output.shape							,							self.output_shape		)

       SCREAMING_SNAKE_CASE_      		=     output[0, -1, -3:, -3:]
       SCREAMING_SNAKE_CASE_      		=     torch.tensor(_lowerCAmelCase		).to(_lowerCAmelCase		)
       assert torch_all_close(output_slice.flatten()							,							_lowerCAmelCase							,							atol=5E-3		)




 @unittest.skipIf(torch_device == 'mps'							,							'Training is not supported in mps'		)
 def  lowerCAmelCase_    (				self    :						Any		):
       SCREAMING_SNAKE_CASE_		,					SCREAMING_SNAKE_CASE_      		=     self.prepare_init_args_and_inputs_for_common()
       SCREAMING_SNAKE_CASE_      		=     self.block_class(**_lowerCAmelCase		)
       model.to(_lowerCAmelCase		)
       model.train()
       SCREAMING_SNAKE_CASE_      		=     model(**_lowerCAmelCase		)

       if isinstance(_lowerCAmelCase							,							_lowerCAmelCase		):
             SCREAMING_SNAKE_CASE_      		=     output[0]

       SCREAMING_SNAKE_CASE_      		=     torch.device(_lowerCAmelCase		)
       SCREAMING_SNAKE_CASE_      		=     randn_tensor(output.shape							,							device=_lowerCAmelCase		)
       SCREAMING_SNAKE_CASE_      		=     torch.nn.functional.mse_loss(_lowerCAmelCase							,							_lowerCAmelCase		)
       loss.backward()