



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

from ...utils import logging


lowerCamelCase__      =			logging.get_logger(__name__)

def 	lowerCAmelCase__    (				_SCREAMING_SNAKE_CASE							:      List[Any]					,  _SCREAMING_SNAKE_CASE							:      Optional[int]     ):
				"""simple docstring"""







				__a	  =		nn.functional.normalize(_SCREAMING_SNAKE_CASE     )
				__a	  =		nn.functional.normalize(_SCREAMING_SNAKE_CASE     )
				return torch.mm(_SCREAMING_SNAKE_CASE					,  normalized_text_embeds.t()     )






class  SCREAMING_SNAKE_CASE      (				snake_case_       ):
							__lowerCamelCase						:						Optional[Any]     =CLIPConfig

							__lowerCamelCase						:						List[str]     =['CLIPEncoderLayer']





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

											'''simple docstring'''
											super().__init__(lowerCAmelCase__						)

											__a	  =		CLIPVisionModel(config.vision_config						)
											__a	  =		nn.Linear(config.vision_config.hidden_size   ,					config.projection_dim   ,					bias=lowerCAmelCase__						)

											__a	  =		nn.Parameter(torch.ones(17   ,					config.projection_dim						)   ,					requires_grad=lowerCAmelCase__						)
											__a	  =		nn.Parameter(torch.ones(3   ,					config.projection_dim						)   ,					requires_grad=lowerCAmelCase__						)

											__a	  =		nn.Parameter(torch.ones(17						)   ,					requires_grad=lowerCAmelCase__						)
											__a	  =		nn.Parameter(torch.ones(3						)   ,					requires_grad=lowerCAmelCase__						)





							@torch.no_grad()
							def     UpperCamelCase_       ( self		:		List[str]   ,					__lowercase		:		List[str]   ,					__lowercase		:		Optional[Any]						):

											'''simple docstring'''
											__a	  =		self.vision_model(lowerCAmelCase__						)[1]  # pooled_output
											__a	  =		self.visual_projection(lowerCAmelCase__						)

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

											__a	  =		[]
											__a	  =		image_embeds.shape[0]
											for i in range(lowerCAmelCase__						):
															__a	  =		{'''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
															__a	  =		0.0

															for concept_idx in range(len(special_cos_dist[0]						)						):
																			__a	  =		special_cos_dist[i][concept_idx]
																			__a	  =		self.special_care_embeds_weights[concept_idx].item()
																			__a	  =		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]}						)
																							__a	  =		0.01

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

															result.append(lowerCAmelCase__						)

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

											return images, has_nsfw_concepts







							@torch.no_grad()
							def     UpperCamelCase_       ( self		:		List[str]   ,					__lowercase		:		torch.FloatTensor   ,					__lowercase		:		torch.FloatTensor						):

											'''simple docstring'''
											__a	  =		self.vision_model(lowerCAmelCase__						)[1]  # pooled_output
											__a	  =		self.visual_projection(lowerCAmelCase__						)

											__a	  =		cosine_distance(lowerCAmelCase__   ,					self.special_care_embeds						)
											__a	  =		cosine_distance(lowerCAmelCase__   ,					self.concept_embeds						)

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

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

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

											return images, has_nsfw_concepts





'''simple docstring'''



from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_torch_available,
)


lowerCAmelCase_  :							Dict             =    {
    'configuration_encodec': [
        'ENCODEC_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'EncodecConfig',
    ],
    'feature_extraction_encodec': ['EncodecFeatureExtractor'],
}

try:
							if not is_torch_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							lowerCAmelCase_  :							List[Any]             =    [
							    'ENCODEC_PRETRAINED_MODEL_ARCHIVE_LIST',
							    'EncodecModel',
							    'EncodecPreTrainedModel',
							]

if TYPE_CHECKING:
							from .configuration_encodec import (
							    ENCODEC_PRETRAINED_CONFIG_ARCHIVE_MAP,
							    EncodecConfig,
							)
							from .feature_extraction_encodec import EncodecFeatureExtractor

							try:
														if not is_torch_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_encodec import (
														    ENCODEC_PRETRAINED_MODEL_ARCHIVE_LIST,
														    EncodecModel,
														    EncodecPreTrainedModel,
														)

else:
							import sys

							lowerCAmelCase_  :							List[str]             =    _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)




