





import numpy as np
import torch
import torch.nn as nn
from transformers import CLIPConfig, CLIPVisionModel, PreTrainedModel

from ...utils import logging


UpperCamelCase__ :			Optional[Any]       	=			logging.get_logger(__name__)




def  A_(      A					,   A ):
	UpperCAmelCase_		    =  nn.functional.normalize(UpperCamelCase_ )
	UpperCAmelCase_		    =  nn.functional.normalize(UpperCamelCase_ )
	return torch.mm(UpperCamelCase_					,   normalized_text_embeds.t() )







class 	_UpperCamelCase	(							snake_case_	):







					'''simple docstring'''

					lowerCamelCase  :  List[str]           =							CLIPConfig

					lowerCamelCase  :  Union[str, Any]           =							['CLIPEncoderLayer']





					def __init__(    self     :	Any			,       __lowercase     :	CLIPConfig		):






						'''simple docstring'''







						super().__init__(lowercase_		)

						UpperCAmelCase_		    =  CLIPVisionModel(config.vision_config		)
						UpperCAmelCase_		    =  nn.Linear(config.vision_config.hidden_size			,       config.projection_dim			,       bias=lowercase_		)

						UpperCAmelCase_		    =  nn.Parameter(torch.ones(17			,       config.projection_dim		)			,       requires_grad=lowercase_		)
						UpperCAmelCase_		    =  nn.Parameter(torch.ones(3			,       config.projection_dim		)			,       requires_grad=lowercase_		)

						UpperCAmelCase_		    =  nn.Parameter(torch.ones(17		)			,       requires_grad=lowercase_		)
						UpperCAmelCase_		    =  nn.Parameter(torch.ones(3		)			,       requires_grad=lowercase_		)





					@torch.no_grad()
					def 							SCREAMING_SNAKE_CASE				(    self     :	Optional[int]			,       __lowercase     :	Optional[int]			,       __lowercase     :	Tuple		):






						'''simple docstring'''







						UpperCAmelCase_		    =  self.vision_model(lowercase_		)[1]  # pooled_output
						UpperCAmelCase_		    =  self.visual_projection(lowercase_		)

						# we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
						UpperCAmelCase_		    =  cosine_distance(lowercase_			,       self.special_care_embeds		).cpu().float().numpy()
						UpperCAmelCase_		    =  cosine_distance(lowercase_			,       self.concept_embeds		).cpu().float().numpy()

						UpperCAmelCase_		    =  []
						UpperCAmelCase_		    =  image_embeds.shape[0]
						for i in range(lowercase_		):
							UpperCAmelCase_		    =  {'special_scores': {}, 'special_care': [], 'concept_scores': {}, 'bad_concepts': []}

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

							for concept_idx in range(len(special_cos_dist[0]		)		):
								UpperCAmelCase_		    =  special_cos_dist[i][concept_idx]
								UpperCAmelCase_		    =  self.special_care_embeds_weights[concept_idx].item()
								UpperCAmelCase_		    =  round(concept_cos - concept_threshold + adjustment			,       3		)
								if result_img["special_scores"][concept_idx] > 0:
									result_img["special_care"].append({concept_idx, result_img["""special_scores"""][concept_idx]}		)
									UpperCAmelCase_		    =  0.01

							for concept_idx in range(len(cos_dist[0]		)		):
								UpperCAmelCase_		    =  cos_dist[i][concept_idx]
								UpperCAmelCase_		    =  self.concept_embeds_weights[concept_idx].item()
								UpperCAmelCase_		    =  round(concept_cos - concept_threshold + adjustment			,       3		)
								if result_img["concept_scores"][concept_idx] > 0:
									result_img["bad_concepts"].append(lowercase_		)

							result.append(lowercase_		)

						UpperCAmelCase_		    =  [len(res["""bad_concepts"""]		) > 0 for res in result]

						for idx, has_nsfw_concept in enumerate(lowercase_		):
							if has_nsfw_concept:
								if torch.is_tensor(lowercase_		) or torch.is_tensor(images[0]		):
									UpperCAmelCase_		    =  torch.zeros_like(images[idx]		)  # black image
								else:
									UpperCAmelCase_		    =  np.zeros(images[idx].shape		)  # black image

						if any(lowercase_		):
							logger.warning(
							    """Potential NSFW content was detected in one or more images. A black image will be returned instead."""
							    """ Try again with a different prompt and/or seed."""		)

						return images, has_nsfw_concepts





					@torch.no_grad()
					def 							SCREAMING_SNAKE_CASE				(    self     :	Tuple			,       __lowercase     :	torch.FloatTensor			,       __lowercase     :	torch.FloatTensor		):






						'''simple docstring'''







						UpperCAmelCase_		    =  self.vision_model(lowercase_		)[1]  # pooled_output
						UpperCAmelCase_		    =  self.visual_projection(lowercase_		)

						UpperCAmelCase_		    =  cosine_distance(lowercase_			,       self.special_care_embeds		)
						UpperCAmelCase_		    =  cosine_distance(lowercase_			,       self.concept_embeds		)

						# increase this value to create a stronger `nsfw` filter
						# at the cost of increasing the possibility of filtering benign images
						UpperCAmelCase_		    =  0.0

						UpperCAmelCase_		    =  special_cos_dist - self.special_care_embeds_weights + adjustment
						# special_scores = special_scores.round(decimals=3)
						UpperCAmelCase_		    =  torch.any(special_scores > 0			,       dim=1		)
						UpperCAmelCase_		    =  special_care * 0.01
						UpperCAmelCase_		    =  special_adjustment.unsqueeze(1		).expand(-1			,       cos_dist.shape[1]		)

						UpperCAmelCase_		    =  (cos_dist - self.concept_embeds_weights) + special_adjustment
						# concept_scores = concept_scores.round(decimals=3)
						UpperCAmelCase_		    =  torch.any(concept_scores > 0			,       dim=1		)

						UpperCAmelCase_		    =  0.0  # black image

						return images, has_nsfw_concepts










"""simple docstring"""


import numpy as np

_UpperCamelCase :    int      =			[
    ["""a""", """b""", """c""", """d""", """e"""],
    ["""f""", """g""", """h""", """i""", """k"""],
    ["""l""", """m""", """n""", """o""", """p"""],
    ["""q""", """r""", """s""", """t""", """u"""],
    ["""v""", """w""", """x""", """y""", """z"""],
]






class __a :



  def __init__(self	:		Any )		->				None:
   __A						:							Any      =						np.array(lowercase_ )



  def     A_					(self	:		Dict					,      lowercase_	:		str )		->				np.ndarray:
   __A     ,       __A						:							Any      =						np.where(letter == self.SQUARE )
   __A						:							Tuple      =						np.concatenate([indexa + 1, indexa + 1] )
   return indexes



  def     A_					(self	:		Tuple					,      lowercase_	:		int					,      lowercase_	:		int )		->				str:
   return self.SQUARE[indexa - 1, indexa - 1]



  def     A_					(self	:		Union[str, Any]					,      lowercase_	:		str )		->				str:
   __A						:							Any      =						message.lower()
   __A						:							Optional[Any]      =						message.replace('j'					,      'i' )

   __A						:							Union[str, Any]      =						''
   for letter_index in range(len(lowercase_ ) ):
    if message[letter_index] != " ":
     __A						:							Union[str, Any]      =						self.letter_to_numbers(message[letter_index] )
     __A						:							Union[str, Any]      =						encoded_message + str(numbers[0] ) + str(numbers[1] )
    elif message[letter_index] == " ":
     __A						:							Optional[Any]      =						encoded_message + ' '

   return encoded_message



  def     A_					(self	:		Dict					,      lowercase_	:		str )		->				str:
   __A						:							List[str]      =						message.replace(' '					,      '  ' )
   __A						:							List[Any]      =						''
   for numbers_index in range(int(len(lowercase_ ) / 2 ) ):
    if message[numbers_index * 2] != " ":
     __A						:							Any      =						message[numbers_index * 2]
     __A						:							Tuple      =						message[numbers_index * 2 + 1]

     __A						:							Dict      =						self.numbers_to_letter(int(lowercase_ )					,      int(lowercase_ ) )
     __A						:							Union[str, Any]      =						decoded_message + letter
    elif message[numbers_index * 2] == " ":
     __A						:							int      =						decoded_message + ' '

   return decoded_message





