



'''simple docstring'''



import inspect
import unittest

from transformers import SegformerConfig, is_torch_available, is_vision_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
							import torch

							from transformers import (
							    MODEL_MAPPING,
							    SegformerForImageClassification,
							    SegformerForSemanticSegmentation,
							    SegformerModel,
							)
							from transformers.models.segformer.modeling_segformer import SEGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
							from PIL import Image

							from transformers import SegformerImageProcessor
class 						lowerCamelCase_				(       snake_case_      ):







						def        __lowercase     (	self :   Union[str, Any]					):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   Tuple         =      self.config_class(**self.inputs_dict					)
													self.parent.assertTrue(hasattr(lowerCAmelCase__					,     '''hidden_sizes'''					)					)
													self.parent.assertTrue(hasattr(lowerCAmelCase__					,     '''num_attention_heads'''					)					)
													self.parent.assertTrue(hasattr(lowerCAmelCase__					,     '''num_encoder_blocks'''					)					)


class 						lowerCamelCase_				:







						def __init__(	self :   str					,     lowerCAmelCase__ :   str					,     lowerCAmelCase__ :   Union[str, Any]=13					,     lowerCAmelCase__ :   Dict=64					,     lowerCAmelCase__ :   Tuple=3					,     lowerCAmelCase__ :   Optional[int]=4					,     lowerCAmelCase__ :   str=[2, 2, 2, 2]					,     lowerCAmelCase__ :   Dict=[8, 4, 2, 1]					,     lowerCAmelCase__ :   Optional[Any]=[16, 32, 64, 1_28]					,     lowerCAmelCase__ :   Optional[Any]=[1, 4, 8, 16]					,     lowerCAmelCase__ :   List[str]=[1, 2, 4, 8]					,     lowerCAmelCase__ :   Any=True					,     lowerCAmelCase__ :   Union[str, Any]=True					,     lowerCAmelCase__ :   Optional[int]="gelu"					,     lowerCAmelCase__ :   str=0.1					,     lowerCAmelCase__ :   Any=0.1					,     lowerCAmelCase__ :   List[str]=0.02					,     lowerCAmelCase__ :   str=3					,     lowerCAmelCase__ :   Tuple=None					,     ):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   List[Any]         =      parent
													SCREAMING_SNAKE_CASE       :   Dict         =      batch_size
													SCREAMING_SNAKE_CASE       :   List[Any]         =      image_size
													SCREAMING_SNAKE_CASE       :   Optional[int]         =      num_channels
													SCREAMING_SNAKE_CASE       :   Dict         =      num_encoder_blocks
													SCREAMING_SNAKE_CASE       :   List[str]         =      sr_ratios
													SCREAMING_SNAKE_CASE       :   Any         =      depths
													SCREAMING_SNAKE_CASE       :   Any         =      hidden_sizes
													SCREAMING_SNAKE_CASE       :   List[Any]         =      downsampling_rates
													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      num_attention_heads
													SCREAMING_SNAKE_CASE       :   Optional[int]         =      is_training
													SCREAMING_SNAKE_CASE       :   List[str]         =      use_labels
													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      hidden_act
													SCREAMING_SNAKE_CASE       :   List[Any]         =      hidden_dropout_prob
													SCREAMING_SNAKE_CASE       :   Optional[Any]         =      attention_probs_dropout_prob
													SCREAMING_SNAKE_CASE       :   List[str]         =      initializer_range
													SCREAMING_SNAKE_CASE       :   int         =      num_labels
													SCREAMING_SNAKE_CASE       :   Dict         =      scope







						def        __lowercase     (	self :   List[Any]					):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   Optional[int]         =      floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]					)

													SCREAMING_SNAKE_CASE       :   List[str]         =      None
													if self.use_labels:
																				SCREAMING_SNAKE_CASE       :   Tuple         =      ids_tensor([self.batch_size, self.image_size, self.image_size]					,     self.num_labels					)

													SCREAMING_SNAKE_CASE       :   Optional[Any]         =      self.get_config()
													return config, pixel_values, labels







						def        __lowercase     (	self :   str					):

													"""simple docstring"""
													return SegformerConfig(
													    image_size=self.image_size					,     num_channels=self.num_channels					,     num_encoder_blocks=self.num_encoder_blocks					,     depths=self.depths					,     hidden_sizes=self.hidden_sizes					,     num_attention_heads=self.num_attention_heads					,     hidden_act=self.hidden_act					,     hidden_dropout_prob=self.hidden_dropout_prob					,     attention_probs_dropout_prob=self.attention_probs_dropout_prob					,     initializer_range=self.initializer_range					,     )







						def        __lowercase     (	self :   str					,     lowerCAmelCase__ :   Union[str, Any]					,     lowerCAmelCase__ :   Optional[Any]					,     lowerCAmelCase__ :   int					):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   Optional[Any]         =      SegformerModel(config=lowerCAmelCase__					)
													model.to(lowerCAmelCase__					)
													model.eval()
													SCREAMING_SNAKE_CASE       :   str         =      model(lowerCAmelCase__					)
													SCREAMING_SNAKE_CASE       :   int         =      self.image_size // (self.downsampling_rates[-1] * 2)
													self.parent.assertEqual(
													    result.last_hidden_state.shape					,     (self.batch_size, self.hidden_sizes[-1], expected_height, expected_width)					)







						def        __lowercase     (	self :   List[Any]					,     lowerCAmelCase__ :   int					,     lowerCAmelCase__ :   Union[str, Any]					,     lowerCAmelCase__ :   List[Any]					):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   str         =      self.num_labels
													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      SegformerForSemanticSegmentation(lowerCAmelCase__					)
													model.to(lowerCAmelCase__					)
													model.eval()
													SCREAMING_SNAKE_CASE       :   Dict         =      model(lowerCAmelCase__					)
													self.parent.assertEqual(
													    result.logits.shape					,     (self.batch_size, self.num_labels, self.image_size // 4, self.image_size // 4)					)
													SCREAMING_SNAKE_CASE       :   str         =      model(lowerCAmelCase__					,     labels=lowerCAmelCase__					)
													self.parent.assertEqual(
													    result.logits.shape					,     (self.batch_size, self.num_labels, self.image_size // 4, self.image_size // 4)					)
													self.parent.assertGreater(result.loss					,     0.0					)







						def        __lowercase     (	self :   Optional[int]					,     lowerCAmelCase__ :   Tuple					,     lowerCAmelCase__ :   List[str]					,     lowerCAmelCase__ :   Union[str, Any]					):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   str         =      1
													SCREAMING_SNAKE_CASE       :   str         =      SegformerForSemanticSegmentation(config=lowerCAmelCase__					)
													model.to(lowerCAmelCase__					)
													model.eval()
													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      torch.randint(0					,     1					,     (self.batch_size, self.image_size, self.image_size)					).to(lowerCAmelCase__					)
													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      model(lowerCAmelCase__					,     labels=lowerCAmelCase__					)
													self.parent.assertGreater(result.loss					,     0.0					)







						def        __lowercase     (	self :   Tuple					):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      self.prepare_config_and_inputs()
													SCREAMING_SNAKE_CASE       ,SCREAMING_SNAKE_CASE       ,SCREAMING_SNAKE_CASE       :   int         =      config_and_inputs
													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      {'''pixel_values''': pixel_values}
													return config, inputs_dict


@require_torch
class 						lowerCamelCase_				(       snake_case_     ,					snake_case_     ,					unittest.TestCase      ):
						_lowerCAmelCase  : Any						  =					(
						    (
						        SegformerModel,
						        SegformerForSemanticSegmentation,
						        SegformerForImageClassification,
						    )
						    if is_torch_available()
						    else ()
						)
						_lowerCAmelCase  : str						  =					(
						    {
						        'feature-extraction': SegformerModel,
						        'image-classification': SegformerForImageClassification,
						        'image-segmentation': SegformerForSemanticSegmentation,
						    }
						    if is_torch_available()
						    else {}
						)

						_lowerCAmelCase  : int						  =					True
						_lowerCAmelCase  : List[Any]						  =					False
						_lowerCAmelCase  : List[Any]						  =					False
						_lowerCAmelCase  : str						  =					False







						def        __lowercase     (	self :   Optional[int]					):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   List[str]         =      SegformerModelTester(self					)
													SCREAMING_SNAKE_CASE       :   Any         =      SegformerConfigTester(self					,     config_class=lowerCAmelCase__					)







						def        __lowercase     (	self :   Dict					):

													"""simple docstring"""
													self.config_tester.run_common_tests()







						def        __lowercase     (	self :   Optional[int]					):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_model(*lowerCAmelCase__					)







						def        __lowercase     (	self :   List[str]					):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   Dict         =      self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_binary_image_segmentation(*lowerCAmelCase__					)







						def        __lowercase     (	self :   List[str]					):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   List[str]         =      self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_image_segmentation(*lowerCAmelCase__					)







						@unittest.skip('''SegFormer does not use inputs_embeds'''					)
						def        __lowercase     (	self :   List[str]					):

													"""simple docstring"""
													pass







						@unittest.skip('''SegFormer does not have get_input_embeddings method and get_output_embeddings methods'''					)
						def        __lowercase     (	self :   Tuple					):

													"""simple docstring"""
													pass







						def        __lowercase     (	self :   Optional[int]					):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       ,SCREAMING_SNAKE_CASE       :   Optional[Any]         =      self.model_tester.prepare_config_and_inputs_for_common()

													for model_class in self.all_model_classes:
																				SCREAMING_SNAKE_CASE       :   List[str]         =      model_class(lowerCAmelCase__					)
																				SCREAMING_SNAKE_CASE       :   List[Any]         =      inspect.signature(model.forward					)
																				# signature.parameters is an OrderedDict => so arg_names order is deterministic
																				SCREAMING_SNAKE_CASE       :   int         =      [*signature.parameters.keys()]

																				SCREAMING_SNAKE_CASE       :   Optional[Any]         =      ['''pixel_values''']
																				self.assertListEqual(arg_names[:1]					,     lowerCAmelCase__					)







						def        __lowercase     (	self :   List[str]					):

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       ,SCREAMING_SNAKE_CASE       :   int         =      self.model_tester.prepare_config_and_inputs_for_common()
													SCREAMING_SNAKE_CASE       :   Dict         =      True

													for model_class in self.all_model_classes:
																				SCREAMING_SNAKE_CASE       :   int         =      True
																				SCREAMING_SNAKE_CASE       :   int         =      False
																				SCREAMING_SNAKE_CASE       :   int         =      True
																				SCREAMING_SNAKE_CASE       :   List[Any]         =      model_class(lowerCAmelCase__					)
																				model.to(lowerCAmelCase__					)
																				model.eval()
																				with torch.no_grad():
																											SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      model(**self._prepare_for_class(lowerCAmelCase__					,     lowerCAmelCase__					)					)
																				SCREAMING_SNAKE_CASE       :   Dict         =      outputs.attentions

																				SCREAMING_SNAKE_CASE       :   Optional[Any]         =      sum(self.model_tester.depths					)
																				self.assertEqual(len(lowerCAmelCase__					)					,     lowerCAmelCase__					)

																				# check that output_attentions also work using config
																				del inputs_dict["output_attentions"]
																				SCREAMING_SNAKE_CASE       :   Dict         =      True
																				SCREAMING_SNAKE_CASE       :   Dict         =      model_class(lowerCAmelCase__					)
																				model.to(lowerCAmelCase__					)
																				model.eval()
																				with torch.no_grad():
																											SCREAMING_SNAKE_CASE       :   Any         =      model(**self._prepare_for_class(lowerCAmelCase__					,     lowerCAmelCase__					)					)
																				SCREAMING_SNAKE_CASE       :   int         =      outputs.attentions

																				self.assertEqual(len(lowerCAmelCase__					)					,     lowerCAmelCase__					)

																				# verify the first attentions (first block, first layer)
																				SCREAMING_SNAKE_CASE       :   Any         =      (self.model_tester.image_size // 4) ** 2
																				SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      (self.model_tester.image_size // (4 * self.model_tester.sr_ratios[0])) ** 2
																				self.assertListEqual(
																				    list(attentions[0].shape[-3:]					)					,     [self.model_tester.num_attention_heads[0], expected_seq_len, expected_reduced_seq_len]					,     )

																				# verify the last attentions (last block, last layer)
																				SCREAMING_SNAKE_CASE       :   Optional[Any]         =      (self.model_tester.image_size // 32) ** 2
																				SCREAMING_SNAKE_CASE       :   Optional[int]         =      (self.model_tester.image_size // (32 * self.model_tester.sr_ratios[-1])) ** 2
																				self.assertListEqual(
																				    list(attentions[-1].shape[-3:]					)					,     [self.model_tester.num_attention_heads[-1], expected_seq_len, expected_reduced_seq_len]					,     )
																				SCREAMING_SNAKE_CASE       :   Any         =      len(lowerCAmelCase__					)

																				# Check attention is always last and order is fine
																				SCREAMING_SNAKE_CASE       :   Tuple         =      True
																				SCREAMING_SNAKE_CASE       :   int         =      True
																				SCREAMING_SNAKE_CASE       :   int         =      model_class(lowerCAmelCase__					)
																				model.to(lowerCAmelCase__					)
																				model.eval()
																				with torch.no_grad():
																											SCREAMING_SNAKE_CASE       :   str         =      model(**self._prepare_for_class(lowerCAmelCase__					,     lowerCAmelCase__					)					)

																				self.assertEqual(out_len + 1					,     len(lowerCAmelCase__					)					)

																				SCREAMING_SNAKE_CASE       :   int         =      outputs.attentions

																				self.assertEqual(len(lowerCAmelCase__					)					,     lowerCAmelCase__					)
																				# verify the first attentions (first block, first layer)
																				SCREAMING_SNAKE_CASE       :   Optional[int]         =      (self.model_tester.image_size // 4) ** 2
																				SCREAMING_SNAKE_CASE       :   List[str]         =      (self.model_tester.image_size // (4 * self.model_tester.sr_ratios[0])) ** 2
																				self.assertListEqual(
																				    list(self_attentions[0].shape[-3:]					)					,     [self.model_tester.num_attention_heads[0], expected_seq_len, expected_reduced_seq_len]					,     )







						def        __lowercase     (	self :   Union[str, Any]					):

													"""simple docstring"""
													def check_hidden_states_output(lowerCAmelCase__ :   str					,     lowerCAmelCase__ :   List[Any]					,     lowerCAmelCase__ :   Optional[Any]					):
																				SCREAMING_SNAKE_CASE       :   List[str]         =      model_class(lowerCAmelCase__					)
																				model.to(lowerCAmelCase__					)
																				model.eval()

																				with torch.no_grad():
																											SCREAMING_SNAKE_CASE       :   str         =      model(**self._prepare_for_class(lowerCAmelCase__					,     lowerCAmelCase__					)					)

																				SCREAMING_SNAKE_CASE       :   int         =      outputs.hidden_states

																				SCREAMING_SNAKE_CASE       :   Optional[int]         =      self.model_tester.num_encoder_blocks
																				self.assertEqual(len(lowerCAmelCase__					)					,     lowerCAmelCase__					)

																				# verify the first hidden states (first block)
																				self.assertListEqual(
																				    list(hidden_states[0].shape[-3:]					)					,     [
																				        self.model_tester.hidden_sizes[0],
																				        self.model_tester.image_size // 4,
																				        self.model_tester.image_size // 4,
																				    ]					,     )

													SCREAMING_SNAKE_CASE       ,SCREAMING_SNAKE_CASE       :   str         =      self.model_tester.prepare_config_and_inputs_for_common()

													for model_class in self.all_model_classes:
																				SCREAMING_SNAKE_CASE       :   int         =      True
																				check_hidden_states_output(lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__					)

																				# check that output_hidden_states also work using config
																				del inputs_dict["output_hidden_states"]
																				SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      True

																				check_hidden_states_output(lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__					)







						def        __lowercase     (	self :   Optional[int]					):

													"""simple docstring"""
													if not self.model_tester.is_training:
																				return

													SCREAMING_SNAKE_CASE       ,SCREAMING_SNAKE_CASE       :   Any         =      self.model_tester.prepare_config_and_inputs_for_common()
													SCREAMING_SNAKE_CASE       :   Tuple         =      True

													for model_class in self.all_model_classes:
																				if model_class in get_values(lowerCAmelCase__					):
																											continue

																				SCREAMING_SNAKE_CASE       :   Optional[Any]         =      model_class(lowerCAmelCase__					)
																				model.to(lowerCAmelCase__					)
																				model.train()
																				SCREAMING_SNAKE_CASE       :   Optional[int]         =      self._prepare_for_class(lowerCAmelCase__					,     lowerCAmelCase__					,     return_labels=lowerCAmelCase__					)
																				SCREAMING_SNAKE_CASE       :   Tuple         =      model(**lowerCAmelCase__					).loss
																				loss.backward()







						@unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''					)
						def        __lowercase     (	self :   Tuple					):

													"""simple docstring"""
													pass







						@slow
						def        __lowercase     (	self :   Optional[int]					):

													"""simple docstring"""
													for model_name in SEGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																				SCREAMING_SNAKE_CASE       :   Any         =      SegformerModel.from_pretrained(lowerCAmelCase__					)
																				self.assertIsNotNone(lowerCAmelCase__					)


def        UpperCAmelCase (    ):
							SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''					)
							return image






@require_torch
class 						lowerCamelCase_				(       unittest.TestCase      ):







						@slow
						def        __lowercase     (	self :   Optional[Any]					):

													"""simple docstring"""
													# only resize + normalize
													SCREAMING_SNAKE_CASE       :   List[Any]         =      SegformerImageProcessor(
													    image_scale=(5_12, 5_12)					,     keep_ratio=lowerCAmelCase__					,     align=lowerCAmelCase__					,     do_random_crop=lowerCAmelCase__					)
													SCREAMING_SNAKE_CASE       :   Optional[Any]         =      SegformerForSemanticSegmentation.from_pretrained('''nvidia/segformer-b0-finetuned-ade-512-512'''					).to(
													    lowerCAmelCase__					)

													SCREAMING_SNAKE_CASE       :   Optional[int]         =      prepare_img()
													SCREAMING_SNAKE_CASE       :   List[Any]         =      image_processor(images=lowerCAmelCase__					,     return_tensors='''pt'''					)
													SCREAMING_SNAKE_CASE       :   Optional[Any]         =      encoded_inputs.pixel_values.to(lowerCAmelCase__					)

													with torch.no_grad():
																				SCREAMING_SNAKE_CASE       :   List[Any]         =      model(lowerCAmelCase__					)

													SCREAMING_SNAKE_CASE       :   List[str]         =      torch.Size((1, model.config.num_labels, 1_28, 1_28)					)
													self.assertEqual(outputs.logits.shape					,     lowerCAmelCase__					)

													SCREAMING_SNAKE_CASE       :   List[str]         =      torch.tensor(
													    [
													        [[-4.6310, -5.5232, -6.2356], [-5.1921, -6.1444, -6.5996], [-5.4424, -6.2790, -6.7574]],
													        [[-12.1391, -13.3122, -13.9554], [-12.8732, -13.9352, -14.3563], [-12.9438, -13.8226, -14.2513]],
													        [[-12.5134, -13.4686, -14.4915], [-12.8669, -14.4343, -14.7758], [-13.2523, -14.5819, -15.0694]],
													    ]					).to(lowerCAmelCase__					)
													self.assertTrue(torch.allclose(outputs.logits[0, :3, :3, :3]					,     lowerCAmelCase__					,     atol=1e-4					)					)







						@slow
						def        __lowercase     (	self :   str					):

													"""simple docstring"""
													# only resize + normalize
													SCREAMING_SNAKE_CASE       :   Optional[Any]         =      SegformerImageProcessor(
													    image_scale=(5_12, 5_12)					,     keep_ratio=lowerCAmelCase__					,     align=lowerCAmelCase__					,     do_random_crop=lowerCAmelCase__					)
													SCREAMING_SNAKE_CASE       :   int         =      SegformerForSemanticSegmentation.from_pretrained(
													    '''nvidia/segformer-b1-finetuned-cityscapes-1024-1024'''					).to(lowerCAmelCase__					)

													SCREAMING_SNAKE_CASE       :   Any         =      prepare_img()
													SCREAMING_SNAKE_CASE       :   Optional[Any]         =      image_processor(images=lowerCAmelCase__					,     return_tensors='''pt'''					)
													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      encoded_inputs.pixel_values.to(lowerCAmelCase__					)

													with torch.no_grad():
																				SCREAMING_SNAKE_CASE       :   Dict         =      model(lowerCAmelCase__					)

													SCREAMING_SNAKE_CASE       :   Dict         =      torch.Size((1, model.config.num_labels, 1_28, 1_28)					)
													self.assertEqual(outputs.logits.shape					,     lowerCAmelCase__					)

													SCREAMING_SNAKE_CASE       :   Optional[int]         =      torch.tensor(
													    [
													        [[-13.5748, -13.9111, -12.6500], [-14.3500, -15.3683, -14.2328], [-14.7532, -16.0424, -15.6087]],
													        [[-17.1651, -15.8725, -12.9653], [-17.2580, -17.3718, -14.8223], [-16.6058, -16.8783, -16.7452]],
													        [[-3.6456, -3.0209, -1.4203], [-3.0797, -3.1959, -2.0000], [-1.8757, -1.9217, -1.6997]],
													    ]					).to(lowerCAmelCase__					)
													self.assertTrue(torch.allclose(outputs.logits[0, :3, :3, :3]					,     lowerCAmelCase__					,     atol=1e-1					)					)







						@slow
						def        __lowercase     (	self :   List[str]					):

													"""simple docstring"""
													# only resize + normalize
													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      SegformerImageProcessor(
													    image_scale=(5_12, 5_12)					,     keep_ratio=lowerCAmelCase__					,     align=lowerCAmelCase__					,     do_random_crop=lowerCAmelCase__					)
													SCREAMING_SNAKE_CASE       :   Optional[int]         =      SegformerForSemanticSegmentation.from_pretrained('''nvidia/segformer-b0-finetuned-ade-512-512'''					).to(
													    lowerCAmelCase__					)

													SCREAMING_SNAKE_CASE       :   Optional[int]         =      prepare_img()
													SCREAMING_SNAKE_CASE       :   List[Any]         =      image_processor(images=lowerCAmelCase__					,     return_tensors='''pt'''					)
													SCREAMING_SNAKE_CASE       :   List[str]         =      encoded_inputs.pixel_values.to(lowerCAmelCase__					)

													with torch.no_grad():
																				SCREAMING_SNAKE_CASE       :   List[Any]         =      model(lowerCAmelCase__					)

													SCREAMING_SNAKE_CASE       :   Dict         =      outputs.logits.detach().cpu()

													SCREAMING_SNAKE_CASE       :   Dict         =      image_processor.post_process_semantic_segmentation(outputs=lowerCAmelCase__					,     target_sizes=[(5_00, 3_00)]					)
													SCREAMING_SNAKE_CASE       :   Tuple         =      torch.Size((5_00, 3_00)					)
													self.assertEqual(segmentation[0].shape					,     lowerCAmelCase__					)

													SCREAMING_SNAKE_CASE       :   List[Any]         =      image_processor.post_process_semantic_segmentation(outputs=lowerCAmelCase__					)
													SCREAMING_SNAKE_CASE       :   Any         =      torch.Size((1_28, 1_28)					)
													self.assertEqual(segmentation[0].shape					,     lowerCAmelCase__					)









'''simple docstring'''



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

from ...utils import logging


lowerCAmelCase_  :							int             =    logging.get_logger(__name__)
def        UpperCAmelCase (    A	:					List[Any]       ,		A	:					Optional[int]					):
							SCREAMING_SNAKE_CASE       :   Any         =      nn.functional.normalize(A					)
							SCREAMING_SNAKE_CASE       :   List[str]         =      nn.functional.normalize(A					)
							return torch.mm(A       ,		normalized_text_embeds.t()					)






class 						lowerCamelCase_				(       snake_case_      ):
						_lowerCAmelCase  : Optional[Any]						  =					CLIPConfig

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







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

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

													SCREAMING_SNAKE_CASE       :   Any         =      CLIPVisionModel(config.vision_config					)
													SCREAMING_SNAKE_CASE       :   List[str]         =      nn.Linear(config.vision_config.hidden_size					,     config.projection_dim					,     bias=lowerCAmelCase__					)

													SCREAMING_SNAKE_CASE       :   Any         =      nn.Parameter(torch.ones(17					,     config.projection_dim					)					,     requires_grad=lowerCAmelCase__					)
													SCREAMING_SNAKE_CASE       :   int         =      nn.Parameter(torch.ones(3					,     config.projection_dim					)					,     requires_grad=lowerCAmelCase__					)

													SCREAMING_SNAKE_CASE       :   Dict         =      nn.Parameter(torch.ones(17					)					,     requires_grad=lowerCAmelCase__					)
													SCREAMING_SNAKE_CASE       :   List[str]         =      nn.Parameter(torch.ones(3					)					,     requires_grad=lowerCAmelCase__					)







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

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   Any         =      self.vision_model(lowerCAmelCase__					)[1]  # pooled_output
													SCREAMING_SNAKE_CASE       :   Any         =      self.visual_projection(lowerCAmelCase__					)

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

													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      []
													SCREAMING_SNAKE_CASE       :   int         =      image_embeds.shape[0]
													for i in range(lowerCAmelCase__					):
																				SCREAMING_SNAKE_CASE       :   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
																				SCREAMING_SNAKE_CASE       :   int         =      0.0

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

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

																				result.append(lowerCAmelCase__					)

													SCREAMING_SNAKE_CASE       :   Dict         =      [len(res['''bad_concepts''']					) > 0 for res in result]

													return images, has_nsfw_concepts







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

													"""simple docstring"""
													SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      self.vision_model(lowerCAmelCase__					)[1]  # pooled_output
													SCREAMING_SNAKE_CASE       :   List[Any]         =      self.visual_projection(lowerCAmelCase__					)

													SCREAMING_SNAKE_CASE       :   List[str]         =      cosine_distance(lowerCAmelCase__					,     self.special_care_embeds					)
													SCREAMING_SNAKE_CASE       :   List[str]         =      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
													SCREAMING_SNAKE_CASE       :   Optional[int]         =      0.0

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

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

													return images, has_nsfw_concepts




