




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

from ...utils import logging


a_ 		=  logging.get_logger(__name__)




def        __SCREAMING_SNAKE_CASE (	lowercase_       ,  lowercase_   )							->						Optional[int]:

    """simple docstring"""


    __UpperCamelCase			=    nn.functional.normalize(lowercase_   )
    __UpperCamelCase			=    nn.functional.normalize(lowercase_   )
    return torch.mm(lowercase_       ,  normalized_text_embeds.t()   )


class 			_lowerCamelCase    (      _SCREAMING_SNAKE_CASE ):

      """simple docstring"""

      lowerCAmelCase__      :			Tuple     =       CLIPConfig

      lowerCAmelCase__      :			Dict     =       ["CLIPEncoderLayer"]







      def __init__(  self       :			List[str]			,							snake_case       :			CLIPConfig					):
          super().__init__(snake_case					)

          __UpperCamelCase			=    CLIPVisionModel(config.vision_config					)
          __UpperCamelCase			=    nn.Linear(config.vision_config.hidden_size			,							config.projection_dim			,							bias=snake_case					)

          __UpperCamelCase			=    nn.Parameter(torch.ones(17			,							config.projection_dim					)			,							requires_grad=snake_case					)
          __UpperCamelCase			=    nn.Parameter(torch.ones(3			,							config.projection_dim					)			,							requires_grad=snake_case					)

          __UpperCamelCase			=    nn.Parameter(torch.ones(17					)			,							requires_grad=snake_case					)
          __UpperCamelCase			=    nn.Parameter(torch.ones(3					)			,							requires_grad=snake_case					)







      @torch.no_grad()
      def    snake_case						(  self       :			List[Any]			,							snake_case       :			Tuple			,							snake_case       :			List[str]					):
          __UpperCamelCase			=    self.vision_model(snake_case					)[1]  # pooled_output
          __UpperCamelCase			=    self.visual_projection(snake_case					)

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

          __UpperCamelCase			=    []
          __UpperCamelCase			=    image_embeds.shape[0]
          for i in range(snake_case					):
              __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(snake_case					)

              result.append(snake_case					)

          __UpperCamelCase			=    [len(res['''bad_concepts''']					) > 0 for res in result]

          for idx, has_nsfw_concept in enumerate(snake_case					):
              if has_nsfw_concept:
                  if torch.is_tensor(snake_case					) 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(snake_case					):
              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    snake_case						(  self       :			int			,							snake_case       :			torch.FloatTensor			,							snake_case       :			torch.FloatTensor					):
          __UpperCamelCase			=    self.vision_model(snake_case					)[1]  # pooled_output
          __UpperCamelCase			=    self.visual_projection(snake_case					)

          __UpperCamelCase			=    cosine_distance(snake_case			,							self.special_care_embeds					)
          __UpperCamelCase			=    cosine_distance(snake_case			,							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'''







from functools import partial
from typing import Any, Callable, Dict, List, Type, TypeVar, Union, overload

import torch
import torch.nn as nn
import torch.types
def 			_lowerCAmelCase(      UpperCAmelCase_      :	torch.Tensor					, UpperCAmelCase_      :	torch.Tensor					, UpperCAmelCase_      :	bool						) ->							torch.Tensor:
							# The first operation in a checkpoint can't be in-place, but it's
							# nice to have in-place addition during inference. Thus...
							if not inplace:
														lowerCAmelCase__  =   ma + ma
							else:
														ma += ma

							return ma
def 			_lowerCAmelCase(      UpperCAmelCase_      :	torch.Tensor					, UpperCAmelCase_      :	List[int]						) ->							torch.Tensor:
							lowerCAmelCase__  =   -1 * len(UpperCAmelCase_						)
							lowerCAmelCase__  =   list(range(len(tensor.shape[:zero_index]						)						)						)
							return tensor.permute(first_inds + [zero_index + i for i in inds]						)
def 			_lowerCAmelCase(      UpperCAmelCase_      :	torch.Tensor					, UpperCAmelCase_      :	int						) ->							torch.Tensor:
							return t.reshape(t.shape[:-no_dims] + (-1,)						)
def 			_lowerCAmelCase(      UpperCAmelCase_      :	torch.Tensor					, UpperCAmelCase_      :	torch.Tensor					, UpperCAmelCase_      :	int					, UpperCAmelCase_      :	float = 1E-4						) ->							torch.Tensor:
							lowerCAmelCase__  =   mask.expand(*value.shape						)
							return torch.sum(mask * value					, dim=UpperCAmelCase_						) / (eps + torch.sum(UpperCAmelCase_					, dim=UpperCAmelCase_						))
def 			_lowerCAmelCase(      UpperCAmelCase_      :	torch.Tensor					, UpperCAmelCase_      :	torch.types.Number = 2.31_25					, UpperCAmelCase_      :	torch.types.Number = 21.68_75					, UpperCAmelCase_      :	int = 64						) ->							torch.Tensor:
							lowerCAmelCase__  =   torch.linspace(UpperCAmelCase_					, UpperCAmelCase_					, no_bins - 1					, device=pts.device						)
							lowerCAmelCase__  =   torch.sqrt(torch.sum((pts.unsqueeze(-2						) - pts.unsqueeze(-3						)) ** 2					, dim=-1						)						)
							return torch.bucketize(UpperCAmelCase_					, UpperCAmelCase_						)
def 			_lowerCAmelCase(      UpperCAmelCase_      :	Callable[[list], Any]					, UpperCAmelCase_      :	List[dict]						) ->							dict:
							lowerCAmelCase__  =   dicts[0]
							lowerCAmelCase__  =   {}
							for k, v in first.items():
														lowerCAmelCase__  =   [d[k] for d in dicts]
														if isinstance(UpperCAmelCase_					, UpperCAmelCase_						):
																					lowerCAmelCase__  =   dict_multimap(UpperCAmelCase_					, UpperCAmelCase_						)
														else:
																					lowerCAmelCase__  =   fn(UpperCAmelCase_						)

							return new_dict
def 			_lowerCAmelCase(      UpperCAmelCase_      :	torch.Tensor					, UpperCAmelCase_      :	torch.Tensor						) ->							torch.Tensor:
							lowerCAmelCase__  =   v_bins.view(((1,) * len(x.shape						)) + (len(UpperCAmelCase_						),)						)
							lowerCAmelCase__  =   x[..., None] - reshaped_bins
							lowerCAmelCase__  =   torch.argmin(torch.abs(UpperCAmelCase_						)					, dim=-1						)
							return nn.functional.one_hot(UpperCAmelCase_					, num_classes=len(UpperCAmelCase_						)						).float()
def 			_lowerCAmelCase(      UpperCAmelCase_      :	torch.Tensor					, UpperCAmelCase_      :	torch.Tensor					, UpperCAmelCase_      :	int = 0					, UpperCAmelCase_      :	int = 0						) ->							torch.Tensor:
							lowerCAmelCase__  =   []
							for i, s in enumerate(data.shape[:no_batch_dims]						):
														lowerCAmelCase__  =   torch.arange(UpperCAmelCase_						)
														lowerCAmelCase__  =   r.view(*(*((1,) * i), -1, *((1,) * (len(inds.shape						) - i - 1)))						)
														ranges.append(UpperCAmelCase_						)

							lowerCAmelCase__  =   [slice(UpperCAmelCase_						) for _ in range(len(data.shape						) - no_batch_dims						)]
							lowerCAmelCase__  =   inds
							ranges.extend(UpperCAmelCase_						)
							# Matt note: Editing this to get around the behaviour of using a list as an array index changing
							# in recent Numpy versions
							return data[tuple(UpperCAmelCase_						)]


_UpperCamelCase  						=			TypeVar("""T""")
def 			_lowerCAmelCase(      UpperCAmelCase_      :	Callable[[T], Any]					, UpperCAmelCase_      :	Dict[Any, Union[dict, list, tuple, T]]					, UpperCAmelCase_      :	Type[T]						) ->							Dict[Any, Union[dict, list, tuple, Any]]:
							lowerCAmelCase__  =   {}
							for k, v in dic.items():
														if isinstance(UpperCAmelCase_					, UpperCAmelCase_						):
																					lowerCAmelCase__  =   dict_map(UpperCAmelCase_					, UpperCAmelCase_					, UpperCAmelCase_						)
														else:
																					lowerCAmelCase__  =   tree_map(UpperCAmelCase_					, UpperCAmelCase_					, UpperCAmelCase_						)

							return new_dict
@overload
def 			_lowerCAmelCase(      UpperCAmelCase_      :	Callable[[T], Any]					, UpperCAmelCase_      :	T					, UpperCAmelCase_      :	Type[T]						) ->							Any:
							...
@overload
def 			_lowerCAmelCase(      UpperCAmelCase_      :	Callable[[T], Any]					, UpperCAmelCase_      :	dict					, UpperCAmelCase_      :	Type[T]						) ->							dict:
							...
@overload
def 			_lowerCAmelCase(      UpperCAmelCase_      :	Callable[[T], Any]					, UpperCAmelCase_      :	list					, UpperCAmelCase_      :	Type[T]						) ->							list:
							...
@overload
def 			_lowerCAmelCase(      UpperCAmelCase_      :	Callable[[T], Any]					, UpperCAmelCase_      :	tuple					, UpperCAmelCase_      :	Type[T]						) ->							tuple:
							...
def 			_lowerCAmelCase(      UpperCAmelCase_      :	str					, UpperCAmelCase_      :	Dict					, UpperCAmelCase_      :	Union[str, Any]						) ->							List[str]:
							if isinstance(UpperCAmelCase_					, UpperCAmelCase_						):
														return dict_map(UpperCAmelCase_					, UpperCAmelCase_					, UpperCAmelCase_						)
							elif isinstance(UpperCAmelCase_					, UpperCAmelCase_						):
														return [tree_map(UpperCAmelCase_					, UpperCAmelCase_					, UpperCAmelCase_						) for x in tree]
							elif isinstance(UpperCAmelCase_					, UpperCAmelCase_						):
														return tuple(tree_map(UpperCAmelCase_					, UpperCAmelCase_					, UpperCAmelCase_						) for x in tree						)
							elif isinstance(UpperCAmelCase_					, UpperCAmelCase_						):
														return fn(UpperCAmelCase_						)
							else:
														print(type(UpperCAmelCase_						)						)
														raise ValueError("""Not supported"""						)






_UpperCamelCase  						=			partial(tree_map, leaf_type=torch.Tensor)





