

import collections
import inspect
import unittest

from transformers import DinatConfig
from transformers.testing_utils import require_natten, require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

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


if is_torch_available():
       import torch
       from torch import nn

       from transformers import DinatBackbone, DinatForImageClassification, DinatModel
       from transformers.models.dinat.modeling_dinat import DINAT_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
       from PIL import Image

       from transformers import AutoImageProcessor







class 		__magic_name__     :

   '''simple docstring'''




   def __init__(							self,       lowercase_,       lowercase_=13,       lowercase_=64,       lowercase_=4,       lowercase_=3,       lowercase_=16,       lowercase_=[1, 2, 1],       lowercase_=[2, 4, 8],       lowercase_=3,       lowercase_=[[3], [1, 2], [1]],       lowercase_=2.0,       lowercase_=True,       lowercase_=0.0,       lowercase_=0.0,       lowercase_=0.1,       lowercase_="gelu",       lowercase_=True,       lowercase_=0.02,       lowercase_=1E-5,       lowercase_=True,       lowercase_=None,       lowercase_=True,       lowercase_=10,       lowercase_=["stage1", "stage2"],       lowercase_=[1, 2],       )				->       str:






     """simple docstring"""






     a__    =parent
     a__    =batch_size
     a__    =image_size
     a__    =patch_size
     a__    =num_channels
     a__    =embed_dim
     a__    =depths
     a__    =num_heads
     a__    =kernel_size
     a__    =dilations
     a__    =mlp_ratio
     a__    =qkv_bias
     a__    =hidden_dropout_prob
     a__    =attention_probs_dropout_prob
     a__    =drop_path_rate
     a__    =hidden_act
     a__    =patch_norm
     a__    =layer_norm_eps
     a__    =initializer_range
     a__    =is_training
     a__    =scope
     a__    =use_labels
     a__    =num_labels
     a__    =out_features
     a__    =out_indices



   def 					_UpperCAmelCase					(							self			)				->       Any:






     """simple docstring"""






     a__    =floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]			)

     a__    =None
     if self.use_labels:
       a__    =ids_tensor([self.batch_size],       self.num_labels			)

     a__    =self.get_config()

     return config, pixel_values, labels



   def 					_UpperCAmelCase					(							self			)				->       Dict:






     """simple docstring"""






     return DinatConfig(
         num_labels=self.num_labels,       image_size=self.image_size,       patch_size=self.patch_size,       num_channels=self.num_channels,       embed_dim=self.embed_dim,       depths=self.depths,       num_heads=self.num_heads,       kernel_size=self.kernel_size,       dilations=self.dilations,       mlp_ratio=self.mlp_ratio,       qkv_bias=self.qkv_bias,       hidden_dropout_prob=self.hidden_dropout_prob,       attention_probs_dropout_prob=self.attention_probs_dropout_prob,       drop_path_rate=self.drop_path_rate,       hidden_act=self.hidden_act,       patch_norm=self.patch_norm,       layer_norm_eps=self.layer_norm_eps,       initializer_range=self.initializer_range,       out_features=self.out_features,       out_indices=self.out_indices,       )



   def 					_UpperCAmelCase					(							self,       lowercase_,       lowercase_,       lowercase_			)				->       List[Any]:






     """simple docstring"""






     a__    =DinatModel(config=lowercase_			)
     model.to(lowercase_			)
     model.eval()
     a__    =model(lowercase_			)

     a__    =a__    =(config.image_size // config.patch_size) // (2 ** (len(config.depths			) - 1))
     a__    =int(config.embed_dim * 2 ** (len(config.depths			) - 1)			)

     self.parent.assertEqual(
         result.last_hidden_state.shape,       (self.batch_size, expected_height, expected_width, expected_dim)			)



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






     """simple docstring"""






     a__    =DinatForImageClassification(lowercase_			)
     model.to(lowercase_			)
     model.eval()
     a__    =model(lowercase_,       labels=lowercase_			)
     self.parent.assertEqual(result.logits.shape,       (self.batch_size, self.num_labels)			)

     # test greyscale images
     a__    =1
     a__    =DinatForImageClassification(lowercase_			)
     model.to(lowercase_			)
     model.eval()

     a__    =floats_tensor([self.batch_size, 1, self.image_size, self.image_size]			)
     a__    =model(lowercase_			)
     self.parent.assertEqual(result.logits.shape,       (self.batch_size, self.num_labels)			)



   def 					_UpperCAmelCase					(							self,       lowercase_,       lowercase_,       lowercase_			)				->       Any:






     """simple docstring"""






     a__    =DinatBackbone(config=lowercase_			)
     model.to(lowercase_			)
     model.eval()
     a__    =model(lowercase_			)

     # verify hidden states
     self.parent.assertEqual(len(result.feature_maps			),       len(config.out_features			)			)
     self.parent.assertListEqual(list(result.feature_maps[0].shape			),       [self.batch_size, model.channels[0], 16, 16]			)

     # verify channels
     self.parent.assertEqual(len(model.channels			),       len(config.out_features			)			)

     # verify backbone works with out_features=None
     a__    =None
     a__    =DinatBackbone(config=lowercase_			)
     model.to(lowercase_			)
     model.eval()
     a__    =model(lowercase_			)

     # verify feature maps
     self.parent.assertEqual(len(result.feature_maps			),       1			)
     self.parent.assertListEqual(list(result.feature_maps[0].shape			),       [self.batch_size, model.channels[-1], 4, 4]			)

     # verify channels
     self.parent.assertEqual(len(model.channels			),       1			)



   def 					_UpperCAmelCase					(							self			)				->       Optional[int]:






     """simple docstring"""






     a__    =self.prepare_config_and_inputs()
     a__, a__, a__    =config_and_inputs
     a__    ={'''pixel_values''': pixel_values}
     return config, inputs_dict







@require_natten
@require_torch
class 		__magic_name__     (				lowerCamelCase__						,			lowerCamelCase__						,			unittest.TestCase				):

   '''simple docstring'''




   lowerCamelCase__					:      Any        =  (
       (
           DinatModel,
           DinatForImageClassification,
           DinatBackbone,
       )
       if is_torch_available()
       else ()
   )
   lowerCamelCase__					:      List[str]        =  (
       {'feature-extraction': DinatModel, 'image-classification': DinatForImageClassification}
       if is_torch_available()
       else {}
   )
   lowerCamelCase__					:      Optional[Any]        =  False

   lowerCamelCase__					:      str        =  False
   lowerCamelCase__					:      List[Any]        =  False
   lowerCamelCase__					:      Dict        =  False
   lowerCamelCase__					:      Tuple        =  False



   def 					_UpperCAmelCase					(							self			)				->       int:






     """simple docstring"""






     a__    =DinatModelTester(self			)
     a__    =ConfigTester(self,       config_class=lowercase_,       embed_dim=37			)



   def 					_UpperCAmelCase					(							self			)				->       Optional[int]:






     """simple docstring"""






     self.create_and_test_config_common_properties()
     self.config_tester.create_and_test_config_to_json_string()
     self.config_tester.create_and_test_config_to_json_file()
     self.config_tester.create_and_test_config_from_and_save_pretrained()
     self.config_tester.create_and_test_config_with_num_labels()
     self.config_tester.check_config_can_be_init_without_params()
     self.config_tester.check_config_arguments_init()



   def 					_UpperCAmelCase					(							self			)				->       Optional[int]:






     """simple docstring"""






     return



   def 					_UpperCAmelCase					(							self			)				->       List[Any]:






     """simple docstring"""






     a__    =self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*lowercase_			)



   def 					_UpperCAmelCase					(							self			)				->       List[Any]:






     """simple docstring"""






     a__    =self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_image_classification(*lowercase_			)



   def 					_UpperCAmelCase					(							self			)				->       str:






     """simple docstring"""






     a__    =self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_backbone(*lowercase_			)



   @unittest.skip(reason='''Dinat does not use inputs_embeds'''			)
   def 					_UpperCAmelCase					(							self			)				->       Tuple:






     """simple docstring"""






     pass



   @unittest.skip(reason='''Dinat does not use feedforward chunking'''			)
   def 					_UpperCAmelCase					(							self			)				->       Tuple:






     """simple docstring"""






     pass



   def 					_UpperCAmelCase					(							self			)				->       Tuple:






     """simple docstring"""






     a__, a__    =self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
       a__    =model_class(lowercase_			)
       self.assertIsInstance(model.get_input_embeddings(),       (nn.Module)			)
       a__    =model.get_output_embeddings()
       self.assertTrue(x is None or isinstance(lowercase_,       nn.Linear			)			)



   def 					_UpperCAmelCase					(							self			)				->       Dict:






     """simple docstring"""






     a__, a__    =self.model_tester.prepare_config_and_inputs_for_common()

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

       a__    =['''pixel_values''']
       self.assertListEqual(arg_names[:1],       lowercase_			)



   def 					_UpperCAmelCase					(							self			)				->       Dict:






     """simple docstring"""






     self.skipTest('''Dinat\'s attention operation is handled entirely by NATTEN.'''			)



   def 					_UpperCAmelCase					(							self,       lowercase_,       lowercase_,       lowercase_,       lowercase_			)				->       Union[str, Any]:






     """simple docstring"""






     a__    =model_class(lowercase_			)
     model.to(lowercase_			)
     model.eval()

     with torch.no_grad():
       a__    =model(**self._prepare_for_class(lowercase_,       lowercase_			)			)

     a__    =outputs.hidden_states

     a__    =getattr(
         self.model_tester,       '''expected_num_hidden_layers''',       len(self.model_tester.depths			) + 1			)
     self.assertEqual(len(lowercase_			),       lowercase_			)

     # Dinat has a different seq_length
     a__    =(
         config.patch_size
         if isinstance(config.patch_size,       collections.abc.Iterable			)
         else (config.patch_size, config.patch_size)
     )

     a__    =image_size[0] // patch_size[0]
     a__    =image_size[1] // patch_size[1]

     self.assertListEqual(
         list(hidden_states[0].shape[-3:]			),       [height, width, self.model_tester.embed_dim],       )

     if model_class.__name__ != "DinatBackbone":
       a__    =outputs.reshaped_hidden_states
       self.assertEqual(len(lowercase_			),       lowercase_			)

       a__, a__, a__, a__    =reshaped_hidden_states[0].shape
       a__    =(
           reshaped_hidden_states[0].view(lowercase_,       lowercase_,       lowercase_,       lowercase_			).permute(0,       2,       3,       1			)
       )
       self.assertListEqual(
           list(reshaped_hidden_states.shape[-3:]			),       [height, width, self.model_tester.embed_dim],       )



   def 					_UpperCAmelCase					(							self			)				->       Optional[Any]:






     """simple docstring"""






     a__, a__    =self.model_tester.prepare_config_and_inputs_for_common()

     a__    =(
         self.model_tester.image_size
         if isinstance(self.model_tester.image_size,       collections.abc.Iterable			)
         else (self.model_tester.image_size, self.model_tester.image_size)
     )

     for model_class in self.all_model_classes:
       a__    =True
       self.check_hidden_states_output(lowercase_,       lowercase_,       lowercase_,       lowercase_			)

       # check that output_hidden_states also work using config
       del inputs_dict["output_hidden_states"]
       a__    =True

       self.check_hidden_states_output(lowercase_,       lowercase_,       lowercase_,       lowercase_			)



   @slow
   def 					_UpperCAmelCase					(							self			)				->       Any:






     """simple docstring"""






     for model_name in DINAT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       a__    =DinatModel.from_pretrained(lowercase_			)
       self.assertIsNotNone(lowercase_			)



   def 					_UpperCAmelCase					(							self			)				->       List[Any]:






     """simple docstring"""






     a__, a__    =self.model_tester.prepare_config_and_inputs_for_common()

     a__    =_config_zero_init(lowercase_			)
     for model_class in self.all_model_classes:
       a__    =model_class(config=lowercase_			)
       for name, param in model.named_parameters():
         if "embeddings" not in name and param.requires_grad:
           self.assertIn(
               ((param.data.mean() * 1E9).round() / 1E9).item(),       [0.0, 1.0],       msg=F"""Parameter {name} of model {model_class} seems not properly initialized""",       )







@require_natten
@require_vision
@require_torch
class 		__magic_name__     (				unittest.TestCase				):

   '''simple docstring'''




   @cached_property
   def 					_UpperCAmelCase					(							self			)				->       List[str]:






     """simple docstring"""






     return AutoImageProcessor.from_pretrained('''shi-labs/dinat-mini-in1k-224'''			) if is_vision_available() else None



   @slow
   def 					_UpperCAmelCase					(							self			)				->       int:






     """simple docstring"""






     a__    =DinatForImageClassification.from_pretrained('''shi-labs/dinat-mini-in1k-224'''			).to(lowercase_			)
     a__    =self.default_image_processor

     a__    =Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''			)
     a__    =image_processor(images=lowercase_,       return_tensors='''pt'''			).to(lowercase_			)

     # forward pass
     with torch.no_grad():
       a__    =model(**lowercase_			)

     # verify the logits
     a__    =torch.Size((1, 1000)			)
     self.assertEqual(outputs.logits.shape,       lowercase_			)
     a__    =torch.tensor([-0.1545, -0.7667, 0.4642]			).to(lowercase_			)
     self.assertTrue(torch.allclose(outputs.logits[0, :3],       lowercase_,       atol=1E-4			)			)







@require_torch
@require_natten
class 		__magic_name__     (				unittest.TestCase						,			lowerCamelCase__				):

   '''simple docstring'''




   lowerCamelCase__					:      Tuple        =  (DinatBackbone,) if is_torch_available() else ()
   lowerCamelCase__					:      Optional[int]        =  DinatConfig



   def 					_UpperCAmelCase					(							self			)				->       str:






     """simple docstring"""






     a__    =DinatModelTester(self			)



import copy

from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCamelCase 	=   logging.get_logger(__name__)

lowerCamelCase 	=   {
    '''facebook/sam-vit-huge''': '''https://huggingface.co/facebook/sam-vit-huge/resolve/main/config.json''',
    '''facebook/sam-vit-large''': '''https://huggingface.co/facebook/sam-vit-large/resolve/main/config.json''',
    '''facebook/sam-vit-base''': '''https://huggingface.co/facebook/sam-vit-base/resolve/main/config.json''',
}







class 		__magic_name__     (				lowerCamelCase__				):

   '''simple docstring'''




   def __init__(							self,       lowercase_=256,       lowercase_=1024,       lowercase_=16,       lowercase_=16,       lowercase_=4,       lowercase_="gelu",       lowercase_=1E-6,       **lowercase_,       )				->       int:






     """simple docstring"""






     super().__init__(**lowercase_			)
     a__    =hidden_size
     a__    =image_size
     a__    =patch_size
     a__    =image_size // patch_size
     a__    =mask_input_channels
     a__    =num_point_embeddings
     a__    =hidden_act
     a__    =layer_norm_eps







class 		__magic_name__     (				lowerCamelCase__				):

   '''simple docstring'''




   def __init__(							self,       lowercase_=256,       lowercase_="relu",       lowercase_=2048,       lowercase_=2,       lowercase_=8,       lowercase_=2,       lowercase_=3,       lowercase_=3,       lowercase_=256,       lowercase_=1E-6,       **lowercase_,       )				->       int:






     """simple docstring"""






     super().__init__(**lowercase_			)
     a__    =hidden_size
     a__    =hidden_act
     a__    =mlp_dim
     a__    =num_hidden_layers
     a__    =num_attention_heads
     a__    =attention_downsample_rate
     a__    =num_multimask_outputs
     a__    =iou_head_depth
     a__    =iou_head_hidden_dim
     a__    =layer_norm_eps







class 		__magic_name__     (				lowerCamelCase__				):

   '''simple docstring'''




   def __init__(							self,       lowercase_=768,       lowercase_=256,       lowercase_=12,       lowercase_=12,       lowercase_=3,       lowercase_=1024,       lowercase_=16,       lowercase_="gelu",       lowercase_=1E-06,       lowercase_=0.0,       lowercase_=1E-10,       lowercase_=True,       lowercase_=4.0,       lowercase_=True,       lowercase_=True,       lowercase_=14,       lowercase_=[2, 5, 8, 11],       lowercase_=128,       lowercase_=None,       **lowercase_,       )				->       Union[str, Any]:






     """simple docstring"""






     super().__init__(**lowercase_			)

     a__    =hidden_size
     a__    =output_channels
     a__    =num_hidden_layers
     a__    =num_attention_heads
     a__    =num_channels
     a__    =image_size
     a__    =patch_size
     a__    =hidden_act
     a__    =layer_norm_eps
     a__    =attention_dropout
     a__    =initializer_range
     a__    =qkv_bias
     a__    =mlp_ratio
     a__    =use_abs_pos
     a__    =use_rel_pos
     a__    =window_size
     a__    =global_attn_indexes
     a__    =num_pos_feats
     a__    =int(hidden_size * mlp_ratio			) if mlp_dim is None else mlp_dim







class 		__magic_name__     (				lowerCamelCase__				):

   '''simple docstring'''




   lowerCamelCase__					:      Optional[int]        =  'sam'
   lowerCamelCase__					:      str        =  True



   def __init__(							self,       lowercase_=None,       lowercase_=None,       lowercase_=None,       lowercase_=0.02,       **lowercase_,       )				->       Tuple:






     """simple docstring"""






     super().__init__(**lowercase_			)
     a__    =vision_config if vision_config is not None else {}
     a__    =prompt_encoder_config if prompt_encoder_config is not None else {}
     a__    =mask_decoder_config if mask_decoder_config is not None else {}

     if isinstance(lowercase_,       lowercase_			):
       a__    =vision_config.to_dict()
     if isinstance(lowercase_,       lowercase_			):
       a__    =prompt_encoder_config.to_dict()
     if isinstance(lowercase_,       lowercase_			):
       a__    =mask_decoder_config.to_dict()

     a__    =SamVisionConfig(**lowercase_			)
     a__    =SamPromptEncoderConfig(**lowercase_			)
     a__    =SamMaskDecoderConfig(**lowercase_			)
     a__    =initializer_range



   def 					_UpperCAmelCase					(							self			)				->       List[str]:






     """simple docstring"""






     a__    =copy.deepcopy(self.__dict__			)
     a__    =self.vision_config.to_dict()
     a__    =self.prompt_encoder_config.to_dict()
     a__    =self.mask_decoder_config.to_dict()
     a__    =self.__class__.model_type
     return output
