



from collections import OrderedDict
from typing import Mapping

from packaging import version

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


__UpperCAmelCase      = logging.get_logger(__name__)

__UpperCAmelCase      = {
    '''apple/mobilevitv2-1.0''': '''https://huggingface.co/apple/mobilevitv2-1.0/resolve/main/config.json''',
}



class        lowerCAmelCase_	(						a__							):
       UpperCAmelCase__       :      Optional[Any]       =						"mobilevitv2"
       def __init__(    self,    SCREAMING_SNAKE_CASE_=3,    SCREAMING_SNAKE_CASE_=256,    SCREAMING_SNAKE_CASE_=2,    SCREAMING_SNAKE_CASE_=2.0,    SCREAMING_SNAKE_CASE_="swish",    SCREAMING_SNAKE_CASE_=3,    SCREAMING_SNAKE_CASE_=32,    SCREAMING_SNAKE_CASE_=0.1,    SCREAMING_SNAKE_CASE_=0.02,    SCREAMING_SNAKE_CASE_=1e-5,    SCREAMING_SNAKE_CASE_=512,    SCREAMING_SNAKE_CASE_=[6, 12, 18],    SCREAMING_SNAKE_CASE_=0.1,    SCREAMING_SNAKE_CASE_=255,    SCREAMING_SNAKE_CASE_=[2, 4, 3],    SCREAMING_SNAKE_CASE_=[128, 192, 256],    SCREAMING_SNAKE_CASE_=1.0,    SCREAMING_SNAKE_CASE_=2,    SCREAMING_SNAKE_CASE_=0.0,    SCREAMING_SNAKE_CASE_=0.0,    **SCREAMING_SNAKE_CASE_,    )  ->	int:
          super().__init__(**SCREAMING_SNAKE_CASE_ )

          UpperCamelCase					:						Union[str, Any]					=			num_channels
          UpperCamelCase					:						List[Any]					=			image_size
          UpperCamelCase					:						str					=			patch_size
          UpperCamelCase					:						List[str]					=			expand_ratio
          UpperCamelCase					:						List[str]					=			hidden_act
          UpperCamelCase					:						str					=			conv_kernel_size
          UpperCamelCase					:						int					=			output_stride
          UpperCamelCase					:						Union[str, Any]					=			initializer_range
          UpperCamelCase					:						Optional[Any]					=			layer_norm_eps
          UpperCamelCase					:						Dict					=			n_attn_blocks
          UpperCamelCase					:						Dict					=			base_attn_unit_dims
          UpperCamelCase					:						List[str]					=			width_multiplier
          UpperCamelCase					:						Optional[int]					=			ffn_multiplier
          UpperCamelCase					:						str					=			ffn_dropout
          UpperCamelCase					:						Dict					=			attn_dropout
          UpperCamelCase					:						List[Any]					=			classifier_dropout_prob

          # decode head attributes for semantic segmentation
          UpperCamelCase					:						str					=			aspp_out_channels
          UpperCamelCase					:						Optional[Any]					=			atrous_rates
          UpperCamelCase					:						Optional[int]					=			aspp_dropout_prob
          UpperCamelCase					:						Dict					=			semantic_loss_ignore_index






class        lowerCAmelCase_	(						a__							):
       UpperCAmelCase__       :      Dict       =						version.parse("1.11"							)
       @property
       def       snake_case_ (    self )  ->	Mapping[str, Mapping[int, str]]:
          return OrderedDict([('pixel_values', {0: 'batch', 1: 'num_channels', 2: 'height', 3: 'width'})] )
       @property
       def       snake_case_ (    self )  ->	Mapping[str, Mapping[int, str]]:
          if self.task == "image-classification":
             return OrderedDict([('logits', {0: 'batch'})] )
          else:
             return OrderedDict([('last_hidden_state', {0: 'batch'}), ('pooler_output', {0: 'batch'})] )






       @property
       def       snake_case_ (    self )  ->	float:
          return 1e-4








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

from ...utils import logging


__UpperCAmelCase      = logging.get_logger(__name__)
def 	UpperCamelCase			(     snake_case__ :					List[str] ,  snake_case__ :					str			)			->  Dict:
   UpperCamelCase					:						int					=			nn.functional.normalize(snake_case__			)
   UpperCamelCase					:						Union[str, Any]					=			nn.functional.normalize(snake_case__			)
   return torch.mm(snake_case__ ,  normalized_text_embeds.t()			)



class        lowerCAmelCase_	(						a__							):
       UpperCAmelCase__       :      int       =						CLIPConfig

       UpperCAmelCase__       :      Union[str, Any]       =						["CLIPEncoderLayer"]
       def __init__(    self,    SCREAMING_SNAKE_CASE_ )  ->	List[str]:
          super().__init__(SCREAMING_SNAKE_CASE_ )

          UpperCamelCase					:						int					=			CLIPVisionModel(config.vision_config )
          UpperCamelCase					:						int					=			nn.Linear(config.vision_config.hidden_size,    config.projection_dim,    bias=SCREAMING_SNAKE_CASE_ )

          UpperCamelCase					:						List[Any]					=			nn.Parameter(torch.ones(17,    config.projection_dim ),    requires_grad=SCREAMING_SNAKE_CASE_ )
          UpperCamelCase					:						str					=			nn.Parameter(torch.ones(3,    config.projection_dim ),    requires_grad=SCREAMING_SNAKE_CASE_ )

          UpperCamelCase					:						Dict					=			nn.Parameter(torch.ones(17 ),    requires_grad=SCREAMING_SNAKE_CASE_ )
          UpperCamelCase					:						Optional[Any]					=			nn.Parameter(torch.ones(3 ),    requires_grad=SCREAMING_SNAKE_CASE_ )
       @torch.no_grad()
       def       snake_case_ (    self,    SCREAMING_SNAKE_CASE_,    SCREAMING_SNAKE_CASE_ )  ->	Union[str, Any]:
          UpperCamelCase					:						Optional[int]					=			self.vision_model(SCREAMING_SNAKE_CASE_ )[1]  # pooled_output
          UpperCamelCase					:						int					=			self.visual_projection(SCREAMING_SNAKE_CASE_ )

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

          UpperCamelCase					:						str					=			[]
          UpperCamelCase					:						int					=			image_embeds.shape[0]
          for i in range(SCREAMING_SNAKE_CASE_ ):
             UpperCamelCase					:						Optional[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
             UpperCamelCase					:						Optional[Any]					=			0.0

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

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

             result.append(SCREAMING_SNAKE_CASE_ )

          UpperCamelCase					:						int					=			[len(res['bad_concepts'] ) > 0 for res in result]

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

          if any(SCREAMING_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,    SCREAMING_SNAKE_CASE_,    SCREAMING_SNAKE_CASE_ )  ->	Optional[int]:
          UpperCamelCase					:						Union[str, Any]					=			self.vision_model(SCREAMING_SNAKE_CASE_ )[1]  # pooled_output
          UpperCamelCase					:						str					=			self.visual_projection(SCREAMING_SNAKE_CASE_ )

          UpperCamelCase					:						str					=			cosine_distance(SCREAMING_SNAKE_CASE_,    self.special_care_embeds )
          UpperCamelCase					:						Any					=			cosine_distance(SCREAMING_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					:						Dict					=			0.0

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

          UpperCamelCase					:						List[str]					=			(cos_dist - self.concept_embeds_weights) + special_adjustment
          # concept_scores = concept_scores.round(decimals=3)
          UpperCamelCase					:						Tuple					=			torch.any(concept_scores > 0,    dim=1 )

          UpperCamelCase					:						Any					=			0.0  # black image

          return images, has_nsfw_concepts



