



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__	:   Tuple	,					snake_case__	:   Any     )     ->					Optional[Any]:
   UpperCamelCase   :  Any						     =     nn.functional.normalize(snake_case__     )
   UpperCamelCase   :  Optional[Any]						     =     nn.functional.normalize(snake_case__     )
   return torch.mm(snake_case__	,					normalized_text_embeds.t()     )



class     lowerCAmelCase_		(  a__							):
       UpperCAmelCase__    : Tuple       					=  CLIPConfig

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

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

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

          UpperCamelCase   :  Any						     =     nn.Parameter(torch.ones(17	),      requires_grad=SCREAMING_SNAKE_CASE_	)
          UpperCamelCase   :  str						     =     nn.Parameter(torch.ones(3	),      requires_grad=SCREAMING_SNAKE_CASE_	)
       @torch.no_grad()
       def  snake_case_	(      self,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_	)					->		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   :  Union[str, Any]						     =     cosine_distance(SCREAMING_SNAKE_CASE_,      self.special_care_embeds	).cpu().float().numpy()
          UpperCamelCase   :  List[str]						     =     cosine_distance(SCREAMING_SNAKE_CASE_,      self.concept_embeds	).cpu().float().numpy()

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

             for concept_idx in range(len(special_cos_dist[0]	)	):
                UpperCamelCase   :  Optional[Any]						     =     special_cos_dist[i][concept_idx]
                UpperCamelCase   :  int						     =     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   :  int						     =     cos_dist[i][concept_idx]
                UpperCamelCase   :  Optional[Any]						     =     self.concept_embeds_weights[concept_idx].item()
                UpperCamelCase   :  Optional[int]						     =     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   :  Optional[Any]						     =     [len(res['bad_concepts']	) > 0 for res in result]

          return images, has_nsfw_concepts






       @torch.no_grad()
       def  snake_case_	(      self,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_	)					->		Tuple:
          UpperCamelCase   :  Union[str, Any]						     =     self.vision_model(SCREAMING_SNAKE_CASE_	)[1]  # pooled_output
          UpperCamelCase   :  str						     =     self.visual_projection(SCREAMING_SNAKE_CASE_	)

          UpperCamelCase   :  Dict						     =     cosine_distance(SCREAMING_SNAKE_CASE_,      self.special_care_embeds	)
          UpperCamelCase   :  Union[str, 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   :  Optional[Any]						     =     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   :  Union[str, 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   :  Union[str, Any]						     =     torch.any(concept_scores > 0,      dim=1	)

          return images, has_nsfw_concepts








import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
    import torch

if is_vision_available():
    from PIL import Image

    from transformers import ViTImageProcessor



class     lowerCAmelCase_		(  unittest.TestCase							):
       def __init__(      self,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_=13,      SCREAMING_SNAKE_CASE_=3,      SCREAMING_SNAKE_CASE_=224,      SCREAMING_SNAKE_CASE_=30,      SCREAMING_SNAKE_CASE_=400,      SCREAMING_SNAKE_CASE_=True,      SCREAMING_SNAKE_CASE_=None,      SCREAMING_SNAKE_CASE_=True,      SCREAMING_SNAKE_CASE_=[0.5, 0.5, 0.5],      SCREAMING_SNAKE_CASE_=[0.5, 0.5, 0.5],      )					->		List[str]:
          UpperCamelCase   :  Optional[int]						     =     size if size is not None else {'height': 18, 'width': 18}
          UpperCamelCase   :  List[Any]						     =     parent
          UpperCamelCase   :  List[Any]						     =     batch_size
          UpperCamelCase   :  int						     =     num_channels
          UpperCamelCase   :  int						     =     image_size
          UpperCamelCase   :  List[Any]						     =     min_resolution
          UpperCamelCase   :  int						     =     max_resolution
          UpperCamelCase   :  Any						     =     do_resize
          UpperCamelCase   :  Optional[int]						     =     size
          UpperCamelCase   :  List[str]						     =     do_normalize
          UpperCamelCase   :  Optional[Any]						     =     image_mean
          UpperCamelCase   :  Tuple						     =     image_std






       def  snake_case_	(      self	)					->		List[Any]:
          return {
              "image_mean": self.image_mean,
              "image_std": self.image_std,
              "do_normalize": self.do_normalize,
              "do_resize": self.do_resize,
              "size": self.size,
          }






@require_torch
@require_vision
class     lowerCAmelCase_		(  a__	,   unittest.TestCase							):
       UpperCAmelCase__    : Optional[Any]       					=  ViTImageProcessor if is_vision_available() else None
       def  snake_case_	(      self	)					->		Any:
          UpperCamelCase   :  Dict						     =     EfficientFormerImageProcessorTester(self	)
       @property
       def  snake_case_	(      self	)					->		List[Any]:
          return self.image_proc_tester.prepare_image_processor_dict()
       def  snake_case_	(      self	)					->		Optional[int]:
          UpperCamelCase   :  Optional[int]						     =     self.image_processing_class(**self.image_processor_dict	)
          self.assertTrue(hasattr(SCREAMING_SNAKE_CASE_,      'image_mean'	)	)
          self.assertTrue(hasattr(SCREAMING_SNAKE_CASE_,      'image_std'	)	)
          self.assertTrue(hasattr(SCREAMING_SNAKE_CASE_,      'do_normalize'	)	)
          self.assertTrue(hasattr(SCREAMING_SNAKE_CASE_,      'do_resize'	)	)
          self.assertTrue(hasattr(SCREAMING_SNAKE_CASE_,      'size'	)	)
       def  snake_case_	(      self	)					->		Any:
          pass
       def  snake_case_	(      self	)					->		int:
          # Initialize image_processor
          UpperCamelCase   :  Union[str, Any]						     =     self.image_processing_class(**self.image_processor_dict	)
          # create random PIL images
          UpperCamelCase   :  List[str]						     =     prepare_image_inputs(self.image_proc_tester,      equal_resolution=SCREAMING_SNAKE_CASE_	)
          for image in image_inputs:
             self.assertIsInstance(SCREAMING_SNAKE_CASE_,      Image.Image	)

          # Test not batched input
          UpperCamelCase   :  str						     =     image_processor(image_inputs[0],      return_tensors='pt'	).pixel_values
          self.assertEqual(
              encoded_images.shape,      (
                  1,
                  self.image_proc_tester.num_channels,
                  self.image_proc_tester.size['height'],
                  self.image_proc_tester.size['width'],
              ),      )

          # Test batched
          UpperCamelCase   :  Optional[Any]						     =     image_processor(SCREAMING_SNAKE_CASE_,      return_tensors='pt'	).pixel_values
          self.assertEqual(
              encoded_images.shape,      (
                  self.image_proc_tester.batch_size,
                  self.image_proc_tester.num_channels,
                  self.image_proc_tester.size['height'],
                  self.image_proc_tester.size['width'],
              ),      )
       def  snake_case_	(      self	)					->		str:
          # Initialize image_processor
          UpperCamelCase   :  List[str]						     =     self.image_processing_class(**self.image_processor_dict	)
          # create random numpy tensors
          UpperCamelCase   :  Union[str, Any]						     =     prepare_image_inputs(self.image_proc_tester,      equal_resolution=SCREAMING_SNAKE_CASE_,      numpify=SCREAMING_SNAKE_CASE_	)
          for image in image_inputs:
             self.assertIsInstance(SCREAMING_SNAKE_CASE_,      np.ndarray	)

          # Test not batched input
          UpperCamelCase   :  Dict						     =     image_processor(image_inputs[0],      return_tensors='pt'	).pixel_values
          self.assertEqual(
              encoded_images.shape,      (
                  1,
                  self.image_proc_tester.num_channels,
                  self.image_proc_tester.size['height'],
                  self.image_proc_tester.size['width'],
              ),      )

          # Test batched
          UpperCamelCase   :  Dict						     =     image_processor(SCREAMING_SNAKE_CASE_,      return_tensors='pt'	).pixel_values
          self.assertEqual(
              encoded_images.shape,      (
                  self.image_proc_tester.batch_size,
                  self.image_proc_tester.num_channels,
                  self.image_proc_tester.size['height'],
                  self.image_proc_tester.size['width'],
              ),      )






       def  snake_case_	(      self	)					->		Tuple:
          # Initialize image_processor
          UpperCamelCase   :  Any						     =     self.image_processing_class(**self.image_processor_dict	)
          # create random PyTorch tensors
          UpperCamelCase   :  int						     =     prepare_image_inputs(self.image_proc_tester,      equal_resolution=SCREAMING_SNAKE_CASE_,      torchify=SCREAMING_SNAKE_CASE_	)
          for image in image_inputs:
             self.assertIsInstance(SCREAMING_SNAKE_CASE_,      torch.Tensor	)

          # Test not batched input
          UpperCamelCase   :  Optional[int]						     =     image_processor(image_inputs[0],      return_tensors='pt'	).pixel_values
          self.assertEqual(
              encoded_images.shape,      (
                  1,
                  self.image_proc_tester.num_channels,
                  self.image_proc_tester.size['height'],
                  self.image_proc_tester.size['width'],
              ),      )

          # Test batched
          UpperCamelCase   :  int						     =     image_processor(SCREAMING_SNAKE_CASE_,      return_tensors='pt'	).pixel_values
          self.assertEqual(
              encoded_images.shape,      (
                  self.image_proc_tester.batch_size,
                  self.image_proc_tester.num_channels,
                  self.image_proc_tester.size['height'],
                  self.image_proc_tester.size['width'],
              ),      )



