



"""simple docstring"""







import math
from collections import OrderedDict

import torch
from packaging import version
from torch import Tensor, nn

from .utils import logging


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





class    UpperCamelCase     (			nn.Module    ):




	"""simple docstring"""






	def __init__(			self					):
			super().__init__()
			if version.parse(torch.__version__					) < version.parse("""1.12.0"""					):
					raise ImportError(
					    f"""You are using torch=={torch.__version__}, but torch>=1.12.0 is required to use """
					    """PytorchGELUTanh. Please upgrade torch."""					)






	def 					lowerCamelCase__ (			self							,UpperCAmelCase_					):
			return nn.functional.gelu(UpperCAmelCase_							,approximate="""tanh"""					)







class    UpperCamelCase     (			nn.Module    ):




	"""simple docstring"""






	def 					lowerCamelCase__ (			self							,UpperCAmelCase_					):
			return 0.5 * input * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi					) * (input + 0.044715 * torch.pow(UpperCAmelCase_							,3.0					))					))







class    UpperCamelCase     (			nn.Module    ):




	"""simple docstring"""






	def __init__(			self							,UpperCAmelCase_ = False					):
			super().__init__()
			if use_gelu_python:
					_lowercase							:      str			=					self._gelu_python
			else:
					_lowercase							:      List[str]			=					nn.functional.gelu






	def 					lowerCamelCase__ (			self							,UpperCAmelCase_					):
			return input * 0.5 * (1.0 + torch.erf(input / math.sqrt(2.0					)					))






	def 					lowerCamelCase__ (			self							,UpperCAmelCase_					):
			return self.act(UpperCAmelCase_					)







class    UpperCamelCase     (			nn.Module    ):




	"""simple docstring"""






	def 					lowerCamelCase__ (			self							,UpperCAmelCase_					):
			return 0.5 * input * (1.0 + torch.tanh(input * 0.7978845608 * (1.0 + 0.044715 * input * input)					))







class    UpperCamelCase     (			nn.Module    ):




	"""simple docstring"""






	def 					lowerCamelCase__ (			self							,UpperCAmelCase_					):
			return input * torch.sigmoid(1.702 * input					)







class    UpperCamelCase     (			nn.Module    ):




	"""simple docstring"""






	def __init__(			self							,UpperCAmelCase_							,UpperCAmelCase_					):
			if min > max:
					raise ValueError(f"""min should be < max (got min: {min}, max: {max})"""					)

			super().__init__()
			_lowercase							:      Tuple			=					min
			_lowercase							:      List[Any]			=					max






	def 					lowerCamelCase__ (			self							,UpperCAmelCase_					):
			return torch.clip(gelu(UpperCAmelCase_					)							,self.min							,self.max					)







class    UpperCamelCase     (			nn.Module    ):




	"""simple docstring"""






	def __init__(			self					):
			super().__init__()
			_lowercase							:      Dict			=					math.sqrt(2 / math.pi					)






	def 					lowerCamelCase__ (			self							,UpperCAmelCase_					):
			return 0.5 * input * (1 + torch.tanh(self.precomputed_constant * (input + 0.044715 * torch.pow(UpperCAmelCase_							,3					))					))







class    UpperCamelCase     (			nn.Module    ):




	"""simple docstring"""






	def 					lowerCamelCase__ (			self							,UpperCAmelCase_					):
			return nn.functional.silu(UpperCAmelCase_					)







class    UpperCamelCase     (			nn.Module    ):




	"""simple docstring"""






	def __init__(			self					):
			super().__init__()
			if version.parse(torch.__version__					) < version.parse("""1.9.0"""					):
					_lowercase							:      str			=					self._mish_python
			else:
					_lowercase							:      str			=					nn.functional.mish






	def 					lowerCamelCase__ (			self							,UpperCAmelCase_					):
			return input * torch.tanh(nn.functional.softplus(UpperCAmelCase_					)					)






	def 					lowerCamelCase__ (			self							,UpperCAmelCase_					):
			return self.act(UpperCAmelCase_					)







class    UpperCamelCase     (			nn.Module    ):




	"""simple docstring"""






	def 					lowerCamelCase__ (			self							,UpperCAmelCase_					):
			return input







class    UpperCamelCase     (			nn.Module    ):




	"""simple docstring"""






	def 					lowerCamelCase__ (			self							,UpperCAmelCase_							,UpperCAmelCase_=0.707107							,UpperCAmelCase_=0.282095					):
			_lowercase							:      List[Any]			=					(input - mu).div(sigma * math.sqrt(2.0					)					)
			return 0.5 * (1.0 + torch.erf(UpperCAmelCase_					))







class    UpperCamelCase     (			nn.Module    ):




	"""simple docstring"""






	def 					lowerCamelCase__ (			self							,UpperCAmelCase_					):
			_lowercase							:      Any			=					nn.functional.relu(UpperCAmelCase_					)
			_lowercase							:      Tuple			=					torch.square(UpperCAmelCase_					)
			return squared







class    UpperCamelCase     (			snake_case    ):




	"""simple docstring"""






	def __getitem__(			self							,UpperCAmelCase_					):
			_lowercase							:      Any			=					super().__getitem__(UpperCAmelCase_					)
			_lowercase			,   _lowercase							:      Optional[int]			=					content if isinstance(UpperCAmelCase_							,UpperCAmelCase_					) else (content, {})
			return cls(**UpperCAmelCase_					)







UpperCAmelCase:   List[str]       			=							{
    """gelu""": GELUActivation,
    """gelu_10""": (ClippedGELUActivation, {"""min""": -10, """max""": 10}),
    """gelu_fast""": FastGELUActivation,
    """gelu_new""": NewGELUActivation,
    """gelu_python""": (GELUActivation, {"""use_gelu_python""": True}),
    """gelu_pytorch_tanh""": PytorchGELUTanh,
    """gelu_accurate""": AccurateGELUActivation,
    """laplace""": LaplaceActivation,
    """linear""": LinearActivation,
    """mish""": MishActivation,
    """quick_gelu""": QuickGELUActivation,
    """relu""": nn.ReLU,
    """relu2""": ReLUSquaredActivation,
    """relu6""": nn.ReLUa,
    """sigmoid""": nn.Sigmoid,
    """silu""": SiLUActivation,
    """swish""": SiLUActivation,
    """tanh""": nn.Tanh,
}
UpperCAmelCase:   Any       			=							ClassInstantier(ACTaCLS)


def       __SCREAMING_SNAKE_CASE      (   __UpperCAmelCase	):
		if activation_string in ACTaFN:
				return ACTaFN[activation_string]
		else:
				raise KeyError(F"""function {activation_string} not found in ACT2FN mapping {list(ACTaFN.keys()	)}"""	)


# For backwards compatibility with: from activations import gelu_python
UpperCAmelCase:   Any       			=							get_activation("""gelu_python""")
UpperCAmelCase:   int       			=							get_activation("""gelu_new""")
UpperCAmelCase:   List[Any]       			=							get_activation("""gelu""")
UpperCAmelCase:   Union[str, Any]       			=							get_activation("""gelu_fast""")
UpperCAmelCase:   int       			=							get_activation("""quick_gelu""")
UpperCAmelCase:   List[str]       			=							get_activation("""silu""")
UpperCAmelCase:   Tuple       			=							get_activation("""mish""")
UpperCAmelCase:   Optional[int]       			=							get_activation("""linear""")








"""simple docstring"""







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

from ...utils import logging


UpperCAmelCase:   Dict       			=							logging.get_logger(__name__)


def       __SCREAMING_SNAKE_CASE      (   __UpperCAmelCase					,							__UpperCAmelCase	):
		_lowercase							:      List[str]			=					nn.functional.normalize(__UpperCAmelCase	)
		_lowercase							:      int			=					nn.functional.normalize(__UpperCAmelCase	)
		return torch.mm(__UpperCAmelCase					,							normalized_text_embeds.t()	)





class    UpperCamelCase     (			snake_case    ):




	"""simple docstring"""
	SCREAMING_SNAKE_CASE_				:    List[str]		    =   CLIPConfig

	SCREAMING_SNAKE_CASE_				:    List[Any]		    =   ["CLIPEncoderLayer"]






	def __init__(			self							,UpperCAmelCase_					):
			super().__init__(UpperCAmelCase_					)

			_lowercase							:      Any			=					CLIPVisionModel(config.vision_config					)
			_lowercase							:      Any			=					nn.Linear(config.vision_config.hidden_size							,config.projection_dim							,bias=UpperCAmelCase_					)

			_lowercase							:      Dict			=					nn.Parameter(torch.ones(17							,config.projection_dim					)							,requires_grad=UpperCAmelCase_					)
			_lowercase							:      Tuple			=					nn.Parameter(torch.ones(3							,config.projection_dim					)							,requires_grad=UpperCAmelCase_					)

			_lowercase							:      Optional[int]			=					nn.Parameter(torch.ones(17					)							,requires_grad=UpperCAmelCase_					)
			_lowercase							:      List[Any]			=					nn.Parameter(torch.ones(3					)							,requires_grad=UpperCAmelCase_					)






	@torch.no_grad()
	def 					lowerCamelCase__ (			self							,UpperCAmelCase_							,UpperCAmelCase_					):
			_lowercase							:      List[str]			=					self.vision_model(UpperCAmelCase_					)[1]  # pooled_output
			_lowercase							:      str			=					self.visual_projection(UpperCAmelCase_					)

			# we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
			_lowercase							:      Union[str, Any]			=					cosine_distance(UpperCAmelCase_							,self.special_care_embeds					).cpu().float().numpy()
			_lowercase							:      str			=					cosine_distance(UpperCAmelCase_							,self.concept_embeds					).cpu().float().numpy()

			_lowercase							:      List[str]			=					[]
			_lowercase							:      Optional[Any]			=					image_embeds.shape[0]
			for i in range(UpperCAmelCase_					):
					_lowercase							:      int			=					{"""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
					_lowercase							:      str			=					0.0

					for concept_idx in range(len(special_cos_dist[0]					)					):
							_lowercase							:      Optional[Any]			=					special_cos_dist[i][concept_idx]
							_lowercase							:      str			=					self.special_care_embeds_weights[concept_idx].item()
							_lowercase							:      Tuple			=					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]}					)
									_lowercase							:      Union[str, Any]			=					0.01

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

					result.append(UpperCAmelCase_					)

			_lowercase							:      List[str]			=					[len(res["""bad_concepts"""]					) > 0 for res in result]

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

			if any(UpperCAmelCase_					):
					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 					lowerCamelCase__ (			self							,UpperCAmelCase_							,UpperCAmelCase_					):
			_lowercase							:      List[Any]			=					self.vision_model(UpperCAmelCase_					)[1]  # pooled_output
			_lowercase							:      Optional[int]			=					self.visual_projection(UpperCAmelCase_					)

			_lowercase							:      Union[str, Any]			=					cosine_distance(UpperCAmelCase_							,self.special_care_embeds					)
			_lowercase							:      Tuple			=					cosine_distance(UpperCAmelCase_							,self.concept_embeds					)

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

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

			_lowercase							:      Optional[Any]			=					(cos_dist - self.concept_embeds_weights) + special_adjustment
			# concept_scores = concept_scores.round(decimals=3)
			_lowercase							:      Optional[int]			=					torch.any(concept_scores > 0							,dim=1					)

			_lowercase							:      Tuple			=					0.0  # black image

			return images, has_nsfw_concepts



