

'''simple docstring'''






import argparse
import json
from collections import OrderedDict
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import (
    SegformerConfig,
    SegformerForImageClassification,
    SegformerForSemanticSegmentation,
    SegformerImageProcessor,
)
from transformers.utils import logging


logging.set_verbosity_info()
lowerCAmelCase__   		=      logging.get_logger(__name__)


def 						_A	(       A__	,							A__=False ):






 """simple docstring"""



 __lowercase				    =					OrderedDict()
 for key, value in state_dict.items():
  if encoder_only and not key.startswith('''head''' ):
   __lowercase				    =					'''segformer.encoder.''' + key
  if key.startswith('''backbone''' ):
   __lowercase				    =					key.replace('''backbone'''	,							'''segformer.encoder''' )
  if "patch_embed" in key:
   # replace for example patch_embed1 by patch_embeddings.0
   __lowercase				    =					key[key.find('''patch_embed''' ) + len('''patch_embed''' )]
   __lowercase				    =					key.replace(F"patch_embed{idx}"	,							F"patch_embeddings.{int(A__ )-1}" )
  if "norm" in key:
   __lowercase				    =					key.replace('''norm'''	,							'''layer_norm''' )
  if "segformer.encoder.layer_norm" in key:
   # replace for example layer_norm1 by layer_norm.0
   __lowercase				    =					key[key.find('''segformer.encoder.layer_norm''' ) + len('''segformer.encoder.layer_norm''' )]
   __lowercase				    =					key.replace(F"layer_norm{idx}"	,							F"layer_norm.{int(A__ )-1}" )
  if "layer_norm1" in key:
   __lowercase				    =					key.replace('''layer_norm1'''	,							'''layer_norm_1''' )
  if "layer_norm2" in key:
   __lowercase				    =					key.replace('''layer_norm2'''	,							'''layer_norm_2''' )
  if "block" in key:
   # replace for example block1 by block.0
   __lowercase				    =					key[key.find('''block''' ) + len('''block''' )]
   __lowercase				    =					key.replace(F"block{idx}"	,							F"block.{int(A__ )-1}" )
  if "attn.q" in key:
   __lowercase				    =					key.replace('''attn.q'''	,							'''attention.self.query''' )
  if "attn.proj" in key:
   __lowercase				    =					key.replace('''attn.proj'''	,							'''attention.output.dense''' )
  if "attn" in key:
   __lowercase				    =					key.replace('''attn'''	,							'''attention.self''' )
  if "fc1" in key:
   __lowercase				    =					key.replace('''fc1'''	,							'''dense1''' )
  if "fc2" in key:
   __lowercase				    =					key.replace('''fc2'''	,							'''dense2''' )
  if "linear_pred" in key:
   __lowercase				    =					key.replace('''linear_pred'''	,							'''classifier''' )
  if "linear_fuse" in key:
   __lowercase				    =					key.replace('''linear_fuse.conv'''	,							'''linear_fuse''' )
   __lowercase				    =					key.replace('''linear_fuse.bn'''	,							'''batch_norm''' )
  if "linear_c" in key:
   # replace for example linear_c4 by linear_c.3
   __lowercase				    =					key[key.find('''linear_c''' ) + len('''linear_c''' )]
   __lowercase				    =					key.replace(F"linear_c{idx}"	,							F"linear_c.{int(A__ )-1}" )
  if key.startswith('''head''' ):
   __lowercase				    =					key.replace('''head'''	,							'''classifier''' )
  __lowercase				    =					value

 return new_state_dict


def 						_A	(       A__	,							A__ ):






 """simple docstring"""



 for i in range(config.num_encoder_blocks ):
  for j in range(config.depths[i] ):
   # read in weights + bias of keys and values (which is a single matrix in the original implementation)
   __lowercase				    =					state_dict.pop(F"segformer.encoder.block.{i}.{j}.attention.self.kv.weight" )
   __lowercase				    =					state_dict.pop(F"segformer.encoder.block.{i}.{j}.attention.self.kv.bias" )
   # next, add keys and values (in that order) to the state dict
   __lowercase				    =					kv_weight[
       : config.hidden_sizes[i], :
   ]
   __lowercase				    =					kv_bias[: config.hidden_sizes[i]]
   __lowercase				    =					kv_weight[
       config.hidden_sizes[i] :, :
   ]
   __lowercase				    =					kv_bias[
       config.hidden_sizes[i] :
   ]


def 						_A	(       ):






 """simple docstring"""



 __lowercase				    =					'''http://images.cocodataset.org/val2017/000000039769.jpg'''
 __lowercase				    =					Image.open(requests.get(A__	,							stream=A__ ).raw )

 return image


@torch.no_grad()
def 						_A	(       A__	,							A__	,							A__ ):






 """simple docstring"""



 __lowercase				    =					SegformerConfig()
 __lowercase				    =					False

 # set attributes based on model_name
 __lowercase				    =					'''huggingface/label-files'''
 if "segformer" in model_name:
  __lowercase				    =					model_name[len('''segformer.''' ) : len('''segformer.''' ) + 2]
  if "ade" in model_name:
   __lowercase				    =					150
   __lowercase				    =					'''ade20k-id2label.json'''
   __lowercase				    =					(1, 150, 128, 128)
  elif "city" in model_name:
   __lowercase				    =					19
   __lowercase				    =					'''cityscapes-id2label.json'''
   __lowercase				    =					(1, 19, 128, 128)
  else:
   raise ValueError(F"Model {model_name} not supported" )
 elif "mit" in model_name:
  __lowercase				    =					True
  __lowercase				    =					model_name[4:6]
  __lowercase				    =					1000
  __lowercase				    =					'''imagenet-1k-id2label.json'''
  __lowercase				    =					(1, 1000)
 else:
  raise ValueError(F"Model {model_name} not supported" )

 # set config attributes
 __lowercase				    =					json.load(open(hf_hub_download(A__	,							A__	,							repo_type='''dataset''' )	,							'''r''' ) )
 __lowercase				    =					{int(A__ ): v for k, v in idalabel.items()}
 __lowercase				    =					idalabel
 __lowercase				    =					{v: k for k, v in idalabel.items()}
 if size == "b0":
  pass
 elif size == "b1":
  __lowercase				    =					[64, 128, 320, 512]
  __lowercase				    =					256
 elif size == "b2":
  __lowercase				    =					[64, 128, 320, 512]
  __lowercase				    =					768
  __lowercase				    =					[3, 4, 6, 3]
 elif size == "b3":
  __lowercase				    =					[64, 128, 320, 512]
  __lowercase				    =					768
  __lowercase				    =					[3, 4, 18, 3]
 elif size == "b4":
  __lowercase				    =					[64, 128, 320, 512]
  __lowercase				    =					768
  __lowercase				    =					[3, 8, 27, 3]
 elif size == "b5":
  __lowercase				    =					[64, 128, 320, 512]
  __lowercase				    =					768
  __lowercase				    =					[3, 6, 40, 3]
 else:
  raise ValueError(F"Size {size} not supported" )

 # load image processor (only resize + normalize)
 __lowercase				    =					SegformerImageProcessor(
     image_scale=(512, 512)	,							keep_ratio=A__	,							align=A__	,							do_random_crop=A__ )

 # prepare image
 __lowercase				    =					prepare_img()
 __lowercase				    =					image_processor(images=A__	,							return_tensors='''pt''' ).pixel_values

 logger.info(F"Converting model {model_name}..." )

 # load original state dict
 if encoder_only:
  __lowercase				    =					torch.load(A__	,							map_location=torch.device('''cpu''' ) )
 else:
  __lowercase				    =					torch.load(A__	,							map_location=torch.device('''cpu''' ) )['''state_dict''']

 # rename keys
 __lowercase				    =					rename_keys(A__	,							encoder_only=A__ )
 if not encoder_only:
  del state_dict["decode_head.conv_seg.weight"]
  del state_dict["decode_head.conv_seg.bias"]

 # key and value matrices need special treatment
 read_in_k_v(A__	,							A__ )

 # create HuggingFace model and load state dict
 if encoder_only:
  __lowercase				    =					False
  __lowercase				    =					SegformerForImageClassification(A__ )
 else:
  __lowercase				    =					SegformerForSemanticSegmentation(A__ )
 model.load_state_dict(A__ )
 model.eval()

 # forward pass
 __lowercase				    =					model(A__ )
 __lowercase				    =					outputs.logits

 # set expected_slice based on model name
 # ADE20k checkpoints
 if model_name == "segformer.b0.512x512.ade.160k":
  __lowercase				    =					torch.tensor(
      [
          [[-4.6_3_1_0, -5.5_2_3_2, -6.2_3_5_6], [-5.1_9_2_1, -6.1_4_4_4, -6.5_9_9_6], [-5.4_4_2_4, -6.2_7_9_0, -6.7_5_7_4]],
          [[-1_2.1_3_9_1, -1_3.3_1_2_2, -1_3.9_5_5_4], [-1_2.8_7_3_2, -1_3.9_3_5_2, -1_4.3_5_6_3], [-1_2.9_4_3_8, -1_3.8_2_2_6, -1_4.2_5_1_3]],
          [[-1_2.5_1_3_4, -1_3.4_6_8_6, -1_4.4_9_1_5], [-1_2.8_6_6_9, -1_4.4_3_4_3, -1_4.7_7_5_8], [-1_3.2_5_2_3, -1_4.5_8_1_9, -1_5.0_6_9_4]],
      ] )
 elif model_name == "segformer.b1.512x512.ade.160k":
  __lowercase				    =					torch.tensor(
      [
          [[-7.5_8_2_0, -8.7_2_3_1, -8.3_2_1_5], [-8.0_6_0_0, -1_0.3_5_2_9, -1_0.0_3_0_4], [-7.5_2_0_8, -9.4_1_0_3, -9.6_2_3_9]],
          [[-1_2.6_9_1_8, -1_3.8_9_9_4, -1_3.7_1_3_7], [-1_3.3_1_9_6, -1_5.7_5_2_3, -1_5.4_7_8_9], [-1_2.9_3_4_3, -1_4.8_7_5_7, -1_4.9_6_8_9]],
          [[-1_1.1_9_1_1, -1_1.9_4_2_1, -1_1.3_2_4_3], [-1_1.3_3_4_2, -1_3.6_8_3_9, -1_3.3_5_8_1], [-1_0.3_9_0_9, -1_2.1_8_3_2, -1_2.4_8_5_8]],
      ] )
 elif model_name == "segformer.b2.512x512.ade.160k":
  __lowercase				    =					torch.tensor(
      [
          [[-1_1.8_1_7_3, -1_4.3_8_5_0, -1_6.3_1_2_8], [-1_4.5_6_4_8, -1_6.5_8_0_4, -1_8.6_5_6_8], [-1_4.7_2_2_3, -1_5.7_3_8_7, -1_8.4_2_1_8]],
          [[-1_5.7_2_9_0, -1_7.9_1_7_1, -1_9.4_4_2_3], [-1_8.3_1_0_5, -1_9.9_4_4_8, -2_1.4_6_6_1], [-1_7.9_2_9_6, -1_8.6_4_9_7, -2_0.7_9_1_0]],
          [[-1_5.0_7_8_3, -1_7.0_3_3_6, -1_8.2_7_8_9], [-1_6.8_7_7_1, -1_8.6_8_7_0, -2_0.1_6_1_2], [-1_6.2_4_5_4, -1_7.1_4_2_6, -1_9.5_0_5_5]],
      ] )
 elif model_name == "segformer.b3.512x512.ade.160k":
  __lowercase				    =					torch.tensor(
      [
          [[-9.0_8_7_8, -1_0.2_0_8_1, -1_0.1_8_9_1], [-9.3_1_4_4, -1_0.7_9_4_1, -1_0.9_8_4_3], [-9.2_2_9_4, -1_0.3_8_5_5, -1_0.5_7_0_4]],
          [[-1_2.2_3_1_6, -1_3.9_0_6_8, -1_3.6_1_0_2], [-1_2.9_1_6_1, -1_4.3_7_0_2, -1_4.3_2_3_5], [-1_2.5_2_3_3, -1_3.7_1_7_4, -1_3.7_9_3_2]],
          [[-1_4.6_2_7_5, -1_5.2_4_9_0, -1_4.9_7_2_7], [-1_4.3_4_0_0, -1_5.9_6_8_7, -1_6.2_8_2_7], [-1_4.1_4_8_4, -1_5.4_0_3_3, -1_5.8_9_3_7]],
      ] )
 elif model_name == "segformer.b4.512x512.ade.160k":
  __lowercase				    =					torch.tensor(
      [
          [[-1_2.3_1_4_4, -1_3.2_4_4_7, -1_4.0_8_0_2], [-1_3.3_6_1_4, -1_4.5_8_1_6, -1_5.6_1_1_7], [-1_3.3_3_4_0, -1_4.4_4_3_3, -1_6.2_2_1_9]],
          [[-1_9.2_7_8_1, -2_0.4_1_2_8, -2_0.7_5_0_6], [-2_0.6_1_5_3, -2_1.6_5_6_6, -2_2.0_9_9_8], [-1_9.9_8_0_0, -2_1.0_4_3_0, -2_2.1_4_9_4]],
          [[-1_8.8_7_3_9, -1_9.7_8_0_4, -2_1.1_8_3_4], [-2_0.1_2_3_3, -2_1.6_7_6_5, -2_3.2_9_4_4], [-2_0.0_3_1_5, -2_1.2_6_4_1, -2_3.6_9_4_4]],
      ] )
 elif model_name == "segformer.b5.640x640.ade.160k":
  __lowercase				    =					torch.tensor(
      [
          [[-9.5_5_2_4, -1_2.0_8_3_5, -1_1.7_3_4_8], [-1_0.5_2_2_9, -1_3.6_4_4_6, -1_4.5_6_6_2], [-9.5_8_4_2, -1_2.8_8_5_1, -1_3.9_4_1_4]],
          [[-1_5.3_4_3_2, -1_7.5_3_2_3, -1_7.0_8_1_8], [-1_6.3_3_3_0, -1_8.9_2_5_5, -1_9.2_1_0_1], [-1_5.1_3_4_0, -1_7.7_8_4_8, -1_8.3_9_7_1]],
          [[-1_2.6_0_7_2, -1_4.9_4_8_6, -1_4.6_6_3_1], [-1_3.7_6_2_9, -1_7.0_9_0_7, -1_7.7_7_4_5], [-1_2.7_8_9_9, -1_6.1_6_9_5, -1_7.1_6_7_1]],
      ] )
 # Cityscapes checkpoints
 elif model_name == "segformer.b0.1024x1024.city.160k":
  __lowercase				    =					torch.tensor(
      [
          [[-1_1.9_2_9_5, -1_3.4_0_5_7, -1_4.8_1_0_6], [-1_3.3_4_3_1, -1_4.8_1_7_9, -1_5.3_7_8_1], [-1_4.2_8_3_6, -1_5.5_9_4_2, -1_6.1_5_8_8]],
          [[-1_1.4_9_0_6, -1_2.8_0_6_7, -1_3.6_5_6_4], [-1_3.1_1_8_9, -1_4.0_5_0_0, -1_4.1_5_4_3], [-1_3.8_7_4_8, -1_4.5_1_3_6, -1_4.8_7_8_9]],
          [[0.5_3_7_4, 0.1_0_6_7, -0.4_7_4_2], [0.1_1_4_1, -0.2_2_5_5, -0.7_0_9_9], [-0.3_0_0_0, -0.5_9_2_4, -1.3_1_0_5]],
      ] )
 elif model_name == "segformer.b0.512x1024.city.160k":
  __lowercase				    =					torch.tensor(
      [
          [[-7.8_2_1_7, -9.8_7_6_7, -1_0.1_7_1_7], [-9.4_4_3_8, -1_0.9_0_5_8, -1_1.4_0_4_7], [-9.7_9_3_9, -1_2.3_4_9_5, -1_2.1_0_7_9]],
          [[-7.1_5_1_4, -9.5_3_3_6, -1_0.0_8_6_0], [-9.7_7_7_6, -1_1.6_8_2_2, -1_1.8_4_3_9], [-1_0.1_4_1_1, -1_2.7_6_5_5, -1_2.8_9_7_2]],
          [[0.3_0_2_1, 0.0_8_0_5, -0.2_3_1_0], [-0.0_3_2_8, -0.1_6_0_5, -0.2_7_1_4], [-0.1_4_0_8, -0.5_4_7_7, -0.6_9_7_6]],
      ] )
 elif model_name == "segformer.b0.640x1280.city.160k":
  __lowercase				    =					torch.tensor(
      [
          [
              [-1.13_72e01, -1.27_87e01, -1.34_77e01],
              [-1.25_36e01, -1.41_94e01, -1.44_09e01],
              [-1.32_17e01, -1.48_88e01, -1.53_27e01],
          ],
          [
              [-1.47_91e01, -1.71_22e01, -1.82_77e01],
              [-1.71_63e01, -1.91_92e01, -1.95_33e01],
              [-1.78_97e01, -1.99_91e01, -2.03_15e01],
          ],
          [
              [7.67_23e-01, 4.19_21e-01, -7.78_78e-02],
              [4.77_72e-01, 9.55_57e-03, -2.80_82e-01],
              [3.60_32e-01, -2.48_26e-01, -5.11_68e-01],
          ],
      ] )
 elif model_name == "segformer.b0.768x768.city.160k":
  __lowercase				    =					torch.tensor(
      [
          [[-9.4_9_5_9, -1_1.3_0_8_7, -1_1.7_4_7_9], [-1_1.0_0_2_5, -1_2.6_5_4_0, -1_2.3_3_1_9], [-1_1.4_0_6_4, -1_3.0_4_8_7, -1_2.9_9_0_5]],
          [[-9.8_9_0_5, -1_1.3_0_8_4, -1_2.0_8_5_4], [-1_1.1_7_2_6, -1_2.7_6_9_8, -1_2.9_5_8_3], [-1_1.5_9_8_5, -1_3.3_2_7_8, -1_4.1_7_7_4]],
          [[0.2_2_1_3, 0.0_1_9_2, -0.2_4_6_6], [-0.1_7_3_1, -0.4_2_1_3, -0.4_8_7_4], [-0.3_1_2_6, -0.6_5_4_1, -1.1_3_8_9]],
      ] )
 elif model_name == "segformer.b1.1024x1024.city.160k":
  __lowercase				    =					torch.tensor(
      [
          [[-1_3.5_7_4_8, -1_3.9_1_1_1, -1_2.6_5_0_0], [-1_4.3_5_0_0, -1_5.3_6_8_3, -1_4.2_3_2_8], [-1_4.7_5_3_2, -1_6.0_4_2_4, -1_5.6_0_8_7]],
          [[-1_7.1_6_5_1, -1_5.8_7_2_5, -1_2.9_6_5_3], [-1_7.2_5_8_0, -1_7.3_7_1_8, -1_4.8_2_2_3], [-1_6.6_0_5_8, -1_6.8_7_8_3, -1_6.7_4_5_2]],
          [[-3.6_4_5_6, -3.0_2_0_9, -1.4_2_0_3], [-3.0_7_9_7, -3.1_9_5_9, -2.0_0_0_0], [-1.8_7_5_7, -1.9_2_1_7, -1.6_9_9_7]],
      ] )
 elif model_name == "segformer.b2.1024x1024.city.160k":
  __lowercase				    =					torch.tensor(
      [
          [[-1_6.0_9_7_6, -1_6.4_8_5_6, -1_7.3_9_6_2], [-1_6.6_2_3_4, -1_9.0_3_4_2, -1_9.7_6_8_5], [-1_6.0_9_0_0, -1_8.0_6_6_1, -1_9.1_1_8_0]],
          [[-1_8.4_7_5_0, -1_8.8_4_8_8, -1_9.5_0_7_4], [-1_9.4_0_3_0, -2_2.1_5_7_0, -2_2.5_9_7_7], [-1_9.1_1_9_1, -2_0.8_4_8_6, -2_2.3_7_8_3]],
          [[-4.5_1_7_8, -5.5_0_3_7, -6.5_1_0_9], [-5.0_8_8_4, -7.2_1_7_4, -8.0_3_3_4], [-4.4_1_5_6, -5.8_1_1_7, -7.2_9_7_0]],
      ] )
 elif model_name == "segformer.b3.1024x1024.city.160k":
  __lowercase				    =					torch.tensor(
      [
          [[-1_4.2_0_8_1, -1_4.4_7_3_2, -1_4.1_9_7_7], [-1_4.5_8_6_7, -1_6.4_4_2_3, -1_6.6_3_5_6], [-1_3.4_4_4_1, -1_4.9_6_8_5, -1_6.8_6_9_6]],
          [[-1_4.4_5_7_6, -1_4.7_0_7_3, -1_5.0_4_5_1], [-1_5.0_8_1_6, -1_7.6_2_3_7, -1_7.9_8_7_3], [-1_4.4_2_1_3, -1_6.0_1_9_9, -1_8.5_9_9_2]],
          [[-4.7_3_4_9, -4.9_5_8_8, -5.0_9_6_6], [-4.3_2_1_0, -6.9_3_2_5, -7.2_5_9_1], [-3.4_3_1_2, -4.7_4_8_4, -7.1_9_1_7]],
      ] )
 elif model_name == "segformer.b4.1024x1024.city.160k":
  __lowercase				    =					torch.tensor(
      [
          [[-1_1.7_7_3_7, -1_1.9_5_2_6, -1_1.3_2_7_3], [-1_3.6_6_9_2, -1_4.4_5_7_4, -1_3.8_8_7_8], [-1_3.8_9_3_7, -1_4.6_9_2_4, -1_5.9_3_4_5]],
          [[-1_4.6_7_0_6, -1_4.5_3_3_0, -1_4.1_3_0_6], [-1_6.1_5_0_2, -1_6.8_1_8_0, -1_6.4_2_6_9], [-1_6.8_3_3_8, -1_7.8_9_3_9, -2_0.1_7_4_6]],
          [[1.0_4_9_1, 0.8_2_8_9, 1.0_3_1_0], [1.1_0_4_4, 0.5_2_1_9, 0.8_0_5_5], [1.0_8_9_9, 0.6_9_2_6, 0.5_5_9_0]],
      ] )
 elif model_name == "segformer.b5.1024x1024.city.160k":
  __lowercase				    =					torch.tensor(
      [
          [[-1_2.5_6_4_1, -1_3.4_7_7_7, -1_3.0_6_8_4], [-1_3.9_5_8_7, -1_5.8_9_8_3, -1_6.6_5_5_7], [-1_3.3_1_0_9, -1_5.7_3_5_0, -1_6.3_1_4_1]],
          [[-1_4.7_0_7_4, -1_5.4_3_5_2, -1_4.5_9_4_4], [-1_6.6_3_5_3, -1_8.1_6_6_3, -1_8.6_1_2_0], [-1_5.1_7_0_2, -1_8.0_3_2_9, -1_8.1_5_4_7]],
          [[-1.7_9_9_0, -2.0_9_5_1, -1.7_7_8_4], [-2.6_3_9_7, -3.8_2_4_5, -3.9_6_8_6], [-1.5_2_6_4, -2.8_1_2_6, -2.9_3_1_6]],
      ] )
 else:
  __lowercase				    =					logits.argmax(-1 ).item()
  print('''Predicted class:'''	,							model.config.idalabel[predicted_class_idx] )

 # verify logits
 if not encoder_only:
  assert logits.shape == expected_shape
  assert torch.allclose(logits[0, :3, :3, :3]	,							A__	,							atol=1e-2 )

 # finally, save model and image processor
 logger.info(F"Saving PyTorch model and image processor to {pytorch_dump_folder_path}..." )
 Path(A__ ).mkdir(exist_ok=A__ )
 model.save_pretrained(A__ )
 image_processor.save_pretrained(A__ )


if __name__ == "__main__":
 lowerCAmelCase__   		=      argparse.ArgumentParser()

 parser.add_argument(
     '''--model_name''',
     default='''segformer.b0.512x512.ade.160k''',
     type=str,
     help='''Name of the model you\'d like to convert.''',
 )
 parser.add_argument(
     '''--checkpoint_path''', default=None, type=str, help='''Path to the original PyTorch checkpoint (.pth file).'''
 )
 parser.add_argument(
     '''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the folder to output PyTorch model.'''
 )
 lowerCAmelCase__   		=      parser.parse_args()
 convert_segformer_checkpoint(args.model_name, args.checkpoint_path, args.pytorch_dump_folder_path)





"""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      UpperCAmelCase     (       _lowercase :   str			,   _lowercase :   int			,   _lowercase :   Optional[int]=1E-12	)						-> List[str]:


	"""simple docstring"""




	lowerCAmelCase_      							=			jnp.divide(emb_a.T			,   jnp.clip(jnp.linalg.norm(_lowercase			,   axis=1	)			,   a_min=_lowercase	)	).T
	lowerCAmelCase_      							=			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							):
				lowerCamelCase     :       CLIPConfig
				lowerCamelCase     :       jnp.dtype			       =jnp.floataa







				def      lowerCamelCase_     ( self   ):

					'''simple docstring'''






					lowerCAmelCase_      							=			FlaxCLIPVisionModule(self.config.vision_config   )
					lowerCAmelCase_      							=			nn.Dense(self.config.projection_dim      ,	use_bias=UpperCAmelCase      ,	dtype=self.dtype   )

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

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







				def __call__( self      ,	UpperCAmelCase   ):

					'''simple docstring'''






					lowerCAmelCase_      							=			self.vision_model(UpperCAmelCase   )[1]
					lowerCAmelCase_      							=			self.visual_projection(UpperCAmelCase   )

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

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

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

					return has_nsfw_concepts







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







				def __init__( self      ,	UpperCAmelCase      ,	UpperCAmelCase = None      ,	UpperCAmelCase = 0      ,	UpperCAmelCase = jnp.floataa      ,	UpperCAmelCase = True      ,	**UpperCAmelCase      ,	):

					'''simple docstring'''






					if input_shape is None:
						lowerCAmelCase_      							=			(1, 224, 224, 3)
					lowerCAmelCase_      							=			self.module_class(config=UpperCAmelCase      ,	dtype=UpperCAmelCase      ,	**UpperCAmelCase   )
					super().__init__(UpperCAmelCase      ,	UpperCAmelCase      ,	input_shape=UpperCAmelCase      ,	seed=UpperCAmelCase      ,	dtype=UpperCAmelCase      ,	_do_init=_do_init   )







				def      lowerCamelCase_     ( self      ,	UpperCAmelCase      ,	UpperCAmelCase      ,	UpperCAmelCase = None   ):

					'''simple docstring'''






					lowerCAmelCase_      							=			jax.random.normal(UpperCAmelCase      ,	UpperCAmelCase   )

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

					lowerCAmelCase_      							=			self.module.init(UpperCAmelCase      ,	UpperCAmelCase   )['''params''']

					return random_params







				def __call__( self      ,	UpperCAmelCase      ,	UpperCAmelCase = None      ,	):

					'''simple docstring'''






					lowerCAmelCase_      							=			jnp.transpose(UpperCAmelCase      ,	(0, 2, 3, 1)   )

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