


import inspect
import unittest

import numpy as np

from tests.test_modeling_common import floats_tensor
from transformers import MaskaFormerConfig, is_torch_available, is_vision_available
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import cached_property

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
   import torch

   from transformers import MaskaFormerForUniversalSegmentation, MaskaFormerModel

   if is_vision_available():
      from transformers import MaskaFormerImageProcessor

if is_vision_available():
   from PIL import Image
class 		__magic_name__						:





     '''simple docstring'''



     def __init__(					self:int     ,		_a:Optional[Any]     ,		_a:Any=2     ,		_a:Dict=True     ,		_a:List[Any]=False     ,		_a:List[str]=10     ,		_a:Union[str, Any]=3     ,		_a:Tuple=32 * 8     ,		_a:Dict=32 * 8     ,		_a:List[str]=4     ,		_a:Union[str, Any]=64     ,		):
       snake_case__						= parent
       snake_case__						= batch_size
       snake_case__						= is_training
       snake_case__						= use_auxiliary_loss
       snake_case__						= num_queries
       snake_case__						= num_channels
       snake_case__						= min_size
       snake_case__						= max_size
       snake_case__						= num_labels
       snake_case__						= hidden_dim
       snake_case__						= hidden_dim

     def 	SCREAMING_SNAKE_CASE__						(					self:Tuple   ):
       snake_case__						= floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]   ).to(
           _a   )

       snake_case__						= torch.ones([self.batch_size, self.min_size, self.max_size]     ,		device=_a   )

       snake_case__						= (
           torch.rand([self.batch_size, self.num_labels, self.min_size, self.max_size]     ,		device=_a   ) > 0.5
       ).float()
       snake_case__						= (torch.rand((self.batch_size, self.num_labels)     ,		device=_a   ) > 0.5).long()

       snake_case__						= self.get_config()
       return config, pixel_values, pixel_mask, mask_labels, class_labels

     def 	SCREAMING_SNAKE_CASE__						(					self:List[str]   ):
       snake_case__						= MaskaFormerConfig(
           hidden_size=self.hidden_dim     ,		)
       snake_case__						= self.num_queries
       snake_case__						= self.num_labels

       snake_case__						= [1, 1, 1, 1]
       snake_case__						= self.num_channels

       snake_case__						= 64
       snake_case__						= 1_28
       snake_case__						= self.hidden_dim
       snake_case__						= self.hidden_dim
       snake_case__						= self.hidden_dim
       return config

     def 	SCREAMING_SNAKE_CASE__						(					self:int   ):
       snake_case__				,				snake_case__				,				snake_case__				,				snake_case__				,				snake_case__						= self.prepare_config_and_inputs()
       snake_case__						= {'''pixel_values''': pixel_values, '''pixel_mask''': pixel_mask}
       return config, inputs_dict

     def 	SCREAMING_SNAKE_CASE__						(					self:List[Any]     ,		_a:int     ,		_a:str   ):
       snake_case__						= output.encoder_hidden_states
       snake_case__						= output.pixel_decoder_hidden_states
       snake_case__						= output.transformer_decoder_hidden_states

       self.parent.assertTrue(len(_a   )     ,		len(config.backbone_config.depths   )   )
       self.parent.assertTrue(len(_a   )     ,		len(config.backbone_config.depths   )   )
       self.parent.assertTrue(len(_a   )     ,		config.decoder_layers   )

     def 	SCREAMING_SNAKE_CASE__						(					self:List[str]     ,		_a:Optional[Any]     ,		_a:Optional[int]     ,		_a:List[str]     ,		_a:Tuple=False   ):
       with torch.no_grad():
         snake_case__						= MaskaFormerModel(config=_a   )
         model.to(_a   )
         model.eval()

         snake_case__						= model(pixel_values=_a     ,		pixel_mask=_a   )
         snake_case__						= model(_a     ,		output_hidden_states=_a   )

       self.parent.assertEqual(
           output.transformer_decoder_last_hidden_state.shape     ,		(self.batch_size, self.num_queries, self.hidden_dim)     ,		)
       # let's ensure the other two hidden state exists
       self.parent.assertTrue(output.pixel_decoder_last_hidden_state is not None   )
       self.parent.assertTrue(output.encoder_last_hidden_state is not None   )

       if output_hidden_states:
         self.check_output_hidden_state(_a     ,		_a   )







     def 	SCREAMING_SNAKE_CASE__						(					self:Optional[int]     ,		_a:str     ,		_a:Dict     ,		_a:Optional[Any]     ,		_a:str     ,		_a:str   ):
       snake_case__						= MaskaFormerForUniversalSegmentation(config=_a   )
       model.to(_a   )
       model.eval()

       def comm_check_on_output(_a:Any   ):
         # let's still check that all the required stuff is there
         self.parent.assertTrue(result.transformer_decoder_last_hidden_state is not None   )
         self.parent.assertTrue(result.pixel_decoder_last_hidden_state is not None   )
         self.parent.assertTrue(result.encoder_last_hidden_state is not None   )
         # okay, now we need to check the logits shape
         # due to the encoder compression, masks have a //4 spatial size
         self.parent.assertEqual(
             result.masks_queries_logits.shape     ,		(self.batch_size, self.num_queries, self.min_size // 4, self.max_size // 4)     ,		)
         # + 1 for null class
         self.parent.assertEqual(
             result.class_queries_logits.shape     ,		(self.batch_size, self.num_queries, self.num_labels + 1)   )

       with torch.no_grad():
         snake_case__						= model(pixel_values=_a     ,		pixel_mask=_a   )
         snake_case__						= model(_a   )

         comm_check_on_output(_a   )

         snake_case__						= model(
             pixel_values=_a     ,		pixel_mask=_a     ,		mask_labels=_a     ,		class_labels=_a   )

       comm_check_on_output(_a   )

       self.parent.assertTrue(result.loss is not None   )
       self.parent.assertEqual(result.loss.shape     ,		torch.Size([1]   )   )




@require_torch
class 		__magic_name__						(snake_case_   ,snake_case_   ,unittest.TestCase				):





     '''simple docstring'''



     __lowercase				:							Optional[Any]	   =       (MaskaFormerModel, MaskaFormerForUniversalSegmentation) if is_torch_available() else ()
     __lowercase				:							int	   =       {'feature-extraction': MaskaFormerModel} if is_torch_available() else {}

     __lowercase				:							List[Any]	   =       False
     __lowercase				:							str	   =       False
     __lowercase				:							Union[str, Any]	   =       False
     __lowercase				:							int	   =       False

     def 	SCREAMING_SNAKE_CASE__						(					self:Optional[int]   ):
       snake_case__						= MaskaFormerModelTester(self   )
       snake_case__						= ConfigTester(self     ,		config_class=_a     ,		has_text_modality=_a   )

     def 	SCREAMING_SNAKE_CASE__						(					self:List[Any]   ):
       self.config_tester.run_common_tests()

     def 	SCREAMING_SNAKE_CASE__						(					self:Tuple   ):
       snake_case__				,				snake_case__						= self.model_tester.prepare_config_and_inputs_for_common()
       self.model_tester.create_and_check_maskaformer_model(_a     ,		**_a     ,		output_hidden_states=_a   )

     def 	SCREAMING_SNAKE_CASE__						(					self:List[Any]   ):
       snake_case__						= self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_maskaformer_instance_segmentation_head_model(*_a   )

     @unittest.skip(reason='''Mask2Former does not use inputs_embeds'''   )
     def 	SCREAMING_SNAKE_CASE__						(					self:Tuple   ):
       pass

     @unittest.skip(reason='''Mask2Former does not have a get_input_embeddings method'''   )
     def 	SCREAMING_SNAKE_CASE__						(					self:Dict   ):
       pass

     @unittest.skip(reason='''Mask2Former is not a generative model'''   )
     def 	SCREAMING_SNAKE_CASE__						(					self:List[str]   ):
       pass

     @unittest.skip(reason='''Mask2Former does not use token embeddings'''   )
     def 	SCREAMING_SNAKE_CASE__						(					self:int   ):
       pass

     @require_torch_multi_gpu
     @unittest.skip(
         reason='''Mask2Former has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`'''   )
     def 	SCREAMING_SNAKE_CASE__						(					self:Any   ):
       pass

     @unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''   )
     def 	SCREAMING_SNAKE_CASE__						(					self:List[Any]   ):
       pass

     def 	SCREAMING_SNAKE_CASE__						(					self:Any   ):
       snake_case__				,				snake_case__						= self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
         snake_case__						= model_class(_a   )
         snake_case__						= inspect.signature(model.forward   )
         # signature.parameters is an OrderedDict => so arg_names order is deterministic
         snake_case__						= [*signature.parameters.keys()]

         snake_case__						= ['''pixel_values''']
         self.assertListEqual(arg_names[:1]     ,		_a   )

     @slow
     def 	SCREAMING_SNAKE_CASE__						(					self:int   ):
       for model_name in ["facebook/mask2former-swin-small-coco-instance"]:
         snake_case__						= MaskaFormerModel.from_pretrained(_a   )
         self.assertIsNotNone(_a   )

     def 	SCREAMING_SNAKE_CASE__						(					self:Optional[Any]   ):
       snake_case__						= (self.model_tester.min_size,) * 2
       snake_case__						= {
           '''pixel_values''': torch.randn((2, 3, *size)     ,		device=_a   ),
           '''mask_labels''': torch.randn((2, 10, *size)     ,		device=_a   ),
           '''class_labels''': torch.zeros(2     ,		10     ,		device=_a   ).long(),
       }
       snake_case__						= self.model_tester.get_config()

       snake_case__						= MaskaFormerForUniversalSegmentation(_a   ).to(_a   )
       snake_case__						= model(**_a   )
       self.assertTrue(outputs.loss is not None   )

     def 	SCREAMING_SNAKE_CASE__						(					self:Tuple   ):
       snake_case__				,				snake_case__						= self.model_tester.prepare_config_and_inputs_for_common()
       self.model_tester.create_and_check_maskaformer_model(_a     ,		**_a     ,		output_hidden_states=_a   )

     def 	SCREAMING_SNAKE_CASE__						(					self:Dict   ):
       snake_case__				,				snake_case__						= self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
         snake_case__						= model_class(_a   ).to(_a   )
         snake_case__						= model(**_a     ,		output_attentions=_a   )
         self.assertTrue(outputs.attentions is not None   )

     def 	SCREAMING_SNAKE_CASE__						(					self:List[Any]   ):
       if not self.model_tester.is_training:
         return

       snake_case__						= self.all_model_classes[1]
       snake_case__				,				snake_case__				,				snake_case__				,				snake_case__				,				snake_case__						= self.model_tester.prepare_config_and_inputs()

       snake_case__						= model_class(_a   )
       model.to(_a   )
       model.train()

       snake_case__						= model(_a     ,		mask_labels=_a     ,		class_labels=_a   ).loss
       loss.backward()







     def 	SCREAMING_SNAKE_CASE__						(					self:Dict   ):
       snake_case__						= self.all_model_classes[1]
       snake_case__				,				snake_case__				,				snake_case__				,				snake_case__				,				snake_case__						= self.model_tester.prepare_config_and_inputs()
       snake_case__						= True
       snake_case__						= True

       snake_case__						= model_class(_a   ).to(_a   )
       model.train()

       snake_case__						= model(_a     ,		mask_labels=_a     ,		class_labels=_a   )

       snake_case__						= outputs.encoder_hidden_states[0]
       encoder_hidden_states.retain_grad()

       snake_case__						= outputs.pixel_decoder_hidden_states[0]
       pixel_decoder_hidden_states.retain_grad()

       snake_case__						= outputs.transformer_decoder_hidden_states[0]
       transformer_decoder_hidden_states.retain_grad()

       snake_case__						= outputs.attentions[0]
       attentions.retain_grad()

       outputs.loss.backward(retain_graph=_a   )

       self.assertIsNotNone(encoder_hidden_states.grad   )
       self.assertIsNotNone(pixel_decoder_hidden_states.grad   )
       self.assertIsNotNone(transformer_decoder_hidden_states.grad   )
       self.assertIsNotNone(attentions.grad   )




lowerCamelCase__							:						Tuple							     =   1E-4







def SCREAMING_SNAKE_CASE	(			)    ->				Tuple:
  snake_case__						= Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''				)
  return image


@require_vision
@slow
class 		__magic_name__						(unittest.TestCase				):





     '''simple docstring'''



     @cached_property
     def 	SCREAMING_SNAKE_CASE__						(					self:Any   ):
       return "facebook/mask2former-swin-small-coco-instance"

     @cached_property
     def 	SCREAMING_SNAKE_CASE__						(					self:int   ):
       return MaskaFormerImageProcessor.from_pretrained(self.model_checkpoints   ) if is_vision_available() else None

     def 	SCREAMING_SNAKE_CASE__						(					self:Optional[Any]   ):
       snake_case__						= MaskaFormerModel.from_pretrained(self.model_checkpoints   ).to(_a   )
       snake_case__						= self.default_image_processor
       snake_case__						= prepare_img()
       snake_case__						= image_processor(_a     ,		return_tensors='''pt'''   ).to(_a   )
       snake_case__						= inputs['''pixel_values'''].shape
       # check size is divisible by 32
       self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0   )
       # check size
       self.assertEqual(_a     ,		(1, 3, 3_84, 3_84)   )

       with torch.no_grad():
         snake_case__						= model(**_a   )

       snake_case__						= torch.tensor(
           [[-0.2790, -1.0717, -1.1668], [-0.5128, -0.3128, -0.4987], [-0.5832, 0.1971, -0.0197]]   ).to(_a   )
       self.assertTrue(
           torch.allclose(
               outputs.encoder_last_hidden_state[0, 0, :3, :3]     ,		_a     ,		atol=_a   )   )

       snake_case__						= torch.tensor(
           [[0.8973, 1.1847, 1.1776], [1.1934, 1.5040, 1.5128], [1.1153, 1.4486, 1.4951]]   ).to(_a   )
       self.assertTrue(
           torch.allclose(
               outputs.pixel_decoder_last_hidden_state[0, 0, :3, :3]     ,		_a     ,		atol=_a   )   )

       snake_case__						= torch.tensor(
           [[2.1152, 1.7000, -0.8603], [1.5808, 1.8004, -0.9353], [1.6043, 1.7495, -0.5999]]   ).to(_a   )
       self.assertTrue(
           torch.allclose(
               outputs.transformer_decoder_last_hidden_state[0, :3, :3]     ,		_a     ,		atol=_a   )   )

     def 	SCREAMING_SNAKE_CASE__						(					self:Tuple   ):
       snake_case__						= MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints   ).to(_a   ).eval()
       snake_case__						= self.default_image_processor
       snake_case__						= prepare_img()
       snake_case__						= image_processor(_a     ,		return_tensors='''pt'''   ).to(_a   )
       snake_case__						= inputs['''pixel_values'''].shape
       # check size is divisible by 32
       self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0   )
       # check size
       self.assertEqual(_a     ,		(1, 3, 3_84, 3_84)   )

       with torch.no_grad():
         snake_case__						= model(**_a   )
       # masks_queries_logits
       snake_case__						= outputs.masks_queries_logits
       self.assertEqual(
           masks_queries_logits.shape     ,		(1, model.config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)   )
       snake_case__						= [
           [-8.7839, -9.0056, -8.8121],
           [-7.4104, -7.0313, -6.5401],
           [-6.6105, -6.3427, -6.4675],
       ]
       snake_case__						= torch.tensor(_a   ).to(_a   )
       self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]     ,		_a     ,		atol=_a   )   )
       # class_queries_logits
       snake_case__						= outputs.class_queries_logits
       self.assertEqual(class_queries_logits.shape     ,		(1, model.config.num_queries, model.config.num_labels + 1)   )
       snake_case__						= torch.tensor(
           [
               [1.8324, -8.0835, -4.1922],
               [0.8450, -9.0050, -3.6053],
               [0.3045, -7.7293, -3.0275],
           ]   ).to(_a   )
       self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]     ,		_a     ,		atol=_a   )   )







     def 	SCREAMING_SNAKE_CASE__						(					self:Union[str, Any]   ):
       snake_case__						= MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints   ).to(_a   ).eval()
       snake_case__						= self.default_image_processor

       snake_case__						= image_processor(
           [np.zeros((3, 8_00, 13_33)   ), np.zeros((3, 8_00, 13_33)   )]     ,		segmentation_maps=[np.zeros((3_84, 3_84)   ).astype(np.floataa   ), np.zeros((3_84, 3_84)   ).astype(np.floataa   )]     ,		return_tensors='''pt'''     ,		)

       snake_case__						= inputs['''pixel_values'''].to(_a   )
       snake_case__						= [el.to(_a   ) for el in inputs['''mask_labels''']]
       snake_case__						= [el.to(_a   ) for el in inputs['''class_labels''']]

       with torch.no_grad():
         snake_case__						= model(**_a   )

       self.assertTrue(outputs.loss is not None   )





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


lowerCamelCase__							:						List[Any]							     =   logging.get_logger(__name__)

lowerCamelCase__							:						Optional[int]							     =   {
    """facebook/data2vec-vision-base-ft""": (
        """https://huggingface.co/facebook/data2vec-vision-base-ft/resolve/main/config.json"""
    ),
}
class 		__magic_name__						(snake_case_				):





     '''simple docstring'''



     __lowercase				:							Optional[int]	   =       'data2vec-vision'

     def __init__(					self:int     ,		_a:Tuple=7_68     ,		_a:int=12     ,		_a:Any=12     ,		_a:Optional[int]=30_72     ,		_a:Optional[int]="gelu"     ,		_a:Any=0.0     ,		_a:Any=0.0     ,		_a:List[str]=0.02     ,		_a:Dict=1e-12     ,		_a:Tuple=2_24     ,		_a:Any=16     ,		_a:str=3     ,		_a:str=False     ,		_a:Union[str, Any]=False     ,		_a:Optional[int]=False     ,		_a:Any=False     ,		_a:Dict=0.1     ,		_a:Dict=0.1     ,		_a:str=True     ,		_a:str=[3, 5, 7, 11]     ,		_a:List[str]=[1, 2, 3, 6]     ,		_a:List[str]=True     ,		_a:Any=0.4     ,		_a:str=2_56     ,		_a:Union[str, Any]=1     ,		_a:int=False     ,		_a:Optional[int]=2_55     ,		**_a:Dict     ,		):
       super().__init__(**_a   )

       snake_case__						= hidden_size
       snake_case__						= num_hidden_layers
       snake_case__						= num_attention_heads
       snake_case__						= intermediate_size
       snake_case__						= hidden_act
       snake_case__						= hidden_dropout_prob
       snake_case__						= attention_probs_dropout_prob
       snake_case__						= initializer_range
       snake_case__						= layer_norm_eps

       snake_case__						= image_size
       snake_case__						= patch_size
       snake_case__						= num_channels
       snake_case__						= use_mask_token
       snake_case__						= use_absolute_position_embeddings
       snake_case__						= use_relative_position_bias
       snake_case__						= use_shared_relative_position_bias
       snake_case__						= layer_scale_init_value
       snake_case__						= drop_path_rate
       snake_case__						= use_mean_pooling
       # decode head attributes (semantic segmentation)
       snake_case__						= out_indices
       snake_case__						= pool_scales
       # auxiliary head attributes (semantic segmentation)
       snake_case__						= use_auxiliary_head
       snake_case__						= auxiliary_loss_weight
       snake_case__						= auxiliary_channels
       snake_case__						= auxiliary_num_convs
       snake_case__						= auxiliary_concat_input
       snake_case__						= semantic_loss_ignore_index




class 		__magic_name__						(snake_case_				):





     '''simple docstring'''



     __lowercase				:							Any	   =       version.parse('1.11'				)

     @property
     def 	SCREAMING_SNAKE_CASE__						(					self:List[str]   ):
       return OrderedDict(
           [
               ('''pixel_values''', {0: '''batch''', 1: '''num_channels''', 2: '''height''', 3: '''width'''}),
           ]   )







     @property
     def 	SCREAMING_SNAKE_CASE__						(					self:Tuple   ):
       return 1e-4

