
from typing import Optional

import torch
import torch.utils.checkpoint
from torch import Tensor, nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from ...activations import ACTaFN
from ...modeling_outputs import (
    BackboneOutput,
    BaseModelOutputWithNoAttention,
    BaseModelOutputWithPoolingAndNoAttention,
    ImageClassifierOutputWithNoAttention,
)
from ...modeling_utils import PreTrainedModel
from ...utils import (
    add_code_sample_docstrings,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    logging,
    replace_return_docstrings,
)
from ...utils.backbone_utils import BackboneMixin
from .configuration_resnet import ResNetConfig


__lowerCAmelCase    :	str						=						logging.get_logger(__name__)

# General docstring
__lowerCAmelCase    :	List[str]						=						'''ResNetConfig'''

# Base docstring
__lowerCAmelCase    :	Tuple						=						'''microsoft/resnet-50'''
__lowerCAmelCase    :	List[str]						=						[1, 2_048, 7, 7]

# Image classification docstring
__lowerCAmelCase    :	str						=						'''microsoft/resnet-50'''
__lowerCAmelCase    :	List[str]						=						'''tiger cat'''

__lowerCAmelCase    :	Dict						=						[
    '''microsoft/resnet-50''',
    # See all resnet models at https://huggingface.co/models?filter=resnet
]


class   __lowerCAmelCase		(		nn.Module					):



      """simple docstring"""







      def __init__(					self						:				int ,				_snake_case						:				int ,				_snake_case						:				int ,				_snake_case						:				int = 3 ,				_snake_case						:				int = 1 ,				_snake_case						:				str = "relu" ):
          super().__init__()
          __lowercase						:	Optional[Any]					=				nn.Convad(
              A_ ,				A_ ,				kernel_size=A_ ,				stride=A_ ,				padding=kernel_size // 2 ,				bias=A_ )
          __lowercase						:	Any					=				nn.BatchNormad(A_ )
          __lowercase						:	Optional[int]					=				ACTaFN[activation] if activation is not None else nn.Identity()




      def 						snake_case_	(					self						:				Union[str, Any] ,				_snake_case						:				Tensor ):
          __lowercase						:	Tuple					=				self.convolution(A_ )
          __lowercase						:	str					=				self.normalization(A_ )
          __lowercase						:	Tuple					=				self.activation(A_ )
          return hidden_state






class   __lowerCAmelCase		(		nn.Module					):



      """simple docstring"""







      def __init__(					self						:				Optional[Any] ,				_snake_case						:				ResNetConfig ):
          super().__init__()
          __lowercase						:	Union[str, Any]					=				ResNetConvLayer(
              config.num_channels ,				config.embedding_size ,				kernel_size=7 ,				stride=2 ,				activation=config.hidden_act )
          __lowercase						:	Any					=				nn.MaxPoolad(kernel_size=3 ,				stride=2 ,				padding=1 )
          __lowercase						:	Dict					=				config.num_channels




      def 						snake_case_	(					self						:				Any ,				_snake_case						:				Tensor ):
          __lowercase						:	Tuple					=				pixel_values.shape[1]
          if num_channels != self.num_channels:
              raise ValueError(
                  '''Make sure that the channel dimension of the pixel values match with the one set in the configuration.''' )
          __lowercase						:	int					=				self.embedder(A_ )
          __lowercase						:	List[str]					=				self.pooler(A_ )
          return embedding






class   __lowerCAmelCase		(		nn.Module					):



      """simple docstring"""







      def __init__(					self						:				List[str] ,				_snake_case						:				int ,				_snake_case						:				int ,				_snake_case						:				int = 2 ):
          super().__init__()
          __lowercase						:	Optional[Any]					=				nn.Convad(A_ ,				A_ ,				kernel_size=1 ,				stride=A_ ,				bias=A_ )
          __lowercase						:	int					=				nn.BatchNormad(A_ )




      def 						snake_case_	(					self						:				Any ,				_snake_case						:				Tensor ):
          __lowercase						:	str					=				self.convolution(A_ )
          __lowercase						:	Optional[int]					=				self.normalization(A_ )
          return hidden_state






class   __lowerCAmelCase		(		nn.Module					):



      """simple docstring"""







      def __init__(					self						:				int ,				_snake_case						:				int ,				_snake_case						:				int ,				_snake_case						:				int = 1 ,				_snake_case						:				str = "relu" ):
          super().__init__()
          __lowercase						:	Tuple					=				in_channels != out_channels or stride != 1
          __lowercase						:	str					=				(
              ResNetShortCut(A_ ,				A_ ,				stride=A_ ) if should_apply_shortcut else nn.Identity()
          )
          __lowercase						:	str					=				nn.Sequential(
              ResNetConvLayer(A_ ,				A_ ,				stride=A_ ) ,				ResNetConvLayer(A_ ,				A_ ,				activation=A_ ) ,				)
          __lowercase						:	List[Any]					=				ACTaFN[activation]




      def 						snake_case_	(					self						:				Union[str, Any] ,				_snake_case						:				Dict ):
          __lowercase						:	Optional[int]					=				hidden_state
          __lowercase						:	Any					=				self.layer(A_ )
          __lowercase						:	Optional[int]					=				self.shortcut(A_ )
          hidden_state += residual
          __lowercase						:	Tuple					=				self.activation(A_ )
          return hidden_state






class   __lowerCAmelCase		(		nn.Module					):



      """simple docstring"""







      def __init__(					self						:				Optional[int] ,				_snake_case						:				int ,				_snake_case						:				int ,				_snake_case						:				int = 1 ,				_snake_case						:				str = "relu" ,				_snake_case						:				int = 4 ):
          super().__init__()
          __lowercase						:	Union[str, Any]					=				in_channels != out_channels or stride != 1
          __lowercase						:	List[str]					=				out_channels // reduction
          __lowercase						:	Union[str, Any]					=				(
              ResNetShortCut(A_ ,				A_ ,				stride=A_ ) if should_apply_shortcut else nn.Identity()
          )
          __lowercase						:	List[Any]					=				nn.Sequential(
              ResNetConvLayer(A_ ,				A_ ,				kernel_size=1 ) ,				ResNetConvLayer(A_ ,				A_ ,				stride=A_ ) ,				ResNetConvLayer(A_ ,				A_ ,				kernel_size=1 ,				activation=A_ ) ,				)
          __lowercase						:	int					=				ACTaFN[activation]




      def 						snake_case_	(					self						:				Union[str, Any] ,				_snake_case						:				Dict ):
          __lowercase						:	List[str]					=				hidden_state
          __lowercase						:	Optional[Any]					=				self.layer(A_ )
          __lowercase						:	Optional[Any]					=				self.shortcut(A_ )
          hidden_state += residual
          __lowercase						:	Dict					=				self.activation(A_ )
          return hidden_state






class   __lowerCAmelCase		(		nn.Module					):



      """simple docstring"""







      def __init__(					self						:				int ,				_snake_case						:				ResNetConfig ,				_snake_case						:				int ,				_snake_case						:				int ,				_snake_case						:				int = 2 ,				_snake_case						:				int = 2 ,				):
          super().__init__()

          __lowercase						:	List[str]					=				ResNetBottleNeckLayer if config.layer_type == '''bottleneck''' else ResNetBasicLayer

          __lowercase						:	int					=				nn.Sequential(
              # downsampling is done in the first layer with stride of 2
              layer(A_ ,				A_ ,				stride=A_ ,				activation=config.hidden_act ) ,				*[layer(A_ ,				A_ ,				activation=config.hidden_act ) for _ in range(depth - 1 )] ,				)




      def 						snake_case_	(					self						:				Tuple ,				_snake_case						:				Tensor ):
          __lowercase						:	Union[str, Any]					=				input
          for layer in self.layers:
              __lowercase						:	List[str]					=				layer(A_ )
          return hidden_state






class   __lowerCAmelCase		(		nn.Module					):



      """simple docstring"""







      def __init__(					self						:				List[str] ,				_snake_case						:				ResNetConfig ):
          super().__init__()
          __lowercase						:	Union[str, Any]					=				nn.ModuleList([] )
          # based on `downsample_in_first_stage` the first layer of the first stage may or may not downsample the input
          self.stages.append(
              ResNetStage(
                  A_ ,				config.embedding_size ,				config.hidden_sizes[0] ,				stride=2 if config.downsample_in_first_stage else 1 ,				depth=config.depths[0] ,				) )
          __lowercase						:	Tuple					=				zip(config.hidden_sizes ,				config.hidden_sizes[1:] )
          for (in_channels, out_channels), depth in zip(A_ ,				config.depths[1:] ):
              self.stages.append(ResNetStage(A_ ,				A_ ,				A_ ,				depth=A_ ) )




      def 						snake_case_	(					self						:				Union[str, Any] ,				_snake_case						:				Tensor ,				_snake_case						:				bool = False ,				_snake_case						:				bool = True ):
          __lowercase						:	Union[str, Any]					=				() if output_hidden_states else None

          for stage_module in self.stages:
              if output_hidden_states:
                  __lowercase						:	int					=				hidden_states + (hidden_state,)

              __lowercase						:	int					=				stage_module(A_ )

          if output_hidden_states:
              __lowercase						:	List[Any]					=				hidden_states + (hidden_state,)

          if not return_dict:
              return tuple(v for v in [hidden_state, hidden_states] if v is not None )

          return BaseModelOutputWithNoAttention(
              last_hidden_state=A_ ,				hidden_states=A_ ,				)






class   __lowerCAmelCase		(		UpperCamelCase_					):



      """simple docstring"""







      A__		:	Union[str, Any]         =			ResNetConfig
      A__		:	Optional[int]         =			'''resnet'''
      A__		:	int         =			'''pixel_values'''
      A__		:	str         =			True




      def 						snake_case_	(					self						:				Optional[Any] ,				_snake_case						:				Any ):
          if isinstance(A_ ,				nn.Convad ):
              nn.init.kaiming_normal_(module.weight ,				mode='''fan_out''' ,				nonlinearity='''relu''' )
          elif isinstance(A_ ,				(nn.BatchNormad, nn.GroupNorm) ):
              nn.init.constant_(module.weight ,				1 )
              nn.init.constant_(module.bias ,				0 )




      def 						snake_case_	(					self						:				str ,				_snake_case						:				int ,				_snake_case						:				Union[str, Any]=False ):
          if isinstance(A_ ,				A_ ):
              __lowercase						:	Union[str, Any]					=				value






__lowerCAmelCase    :	Any						=						R'''
    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it
    as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
    behavior.

    Parameters:
        config ([`ResNetConfig`]): Model configuration class with all the parameters of the model.
            Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
'''

__lowerCAmelCase    :	Union[str, Any]						=						R'''
    Args:
        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
            [`ConvNextImageProcessor.__call__`] for details.

        output_hidden_states (`bool`, *optional*):
            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
            more detail.
        return_dict (`bool`, *optional*):
            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
'''


@add_start_docstrings(
    '''The bare ResNet model outputting raw features without any specific head on top.'''       ,       UpperCamelCase_       ,       )
class   __lowerCAmelCase		(		UpperCamelCase_					):



      """simple docstring"""







      def __init__(					self						:				Any ,				_snake_case						:				Any ):
          super().__init__(A_ )
          __lowercase						:	str					=				config
          __lowercase						:	Union[str, Any]					=				ResNetEmbeddings(A_ )
          __lowercase						:	Dict					=				ResNetEncoder(A_ )
          __lowercase						:	Any					=				nn.AdaptiveAvgPoolad((1, 1) )
          # Initialize weights and apply final processing
          self.post_init()




      @add_start_docstrings_to_model_forward(A_ )
      @add_code_sample_docstrings(
          checkpoint=_CHECKPOINT_FOR_DOC ,				output_type=A_ ,				config_class=_CONFIG_FOR_DOC ,				modality='''vision''' ,				expected_output=_EXPECTED_OUTPUT_SHAPE ,				)
      def 						snake_case_	(					self						:				Optional[int] ,				_snake_case						:				Tensor ,				_snake_case						:				Optional[bool] = None ,				_snake_case						:				Optional[bool] = None ):
          __lowercase						:	Optional[int]					=				(
              output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
          )
          __lowercase						:	Optional[int]					=				return_dict if return_dict is not None else self.config.use_return_dict

          __lowercase						:	Tuple					=				self.embedder(A_ )

          __lowercase						:	Tuple					=				self.encoder(
              A_ ,				output_hidden_states=A_ ,				return_dict=A_ )

          __lowercase						:	List[Any]					=				encoder_outputs[0]

          __lowercase						:	str					=				self.pooler(A_ )

          if not return_dict:
              return (last_hidden_state, pooled_output) + encoder_outputs[1:]

          return BaseModelOutputWithPoolingAndNoAttention(
              last_hidden_state=A_ ,				pooler_output=A_ ,				hidden_states=encoder_outputs.hidden_states ,				)






@add_start_docstrings(
    '''
    ResNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
    ImageNet.
    '''       ,       UpperCamelCase_       ,       )
class   __lowerCAmelCase		(		UpperCamelCase_					):



      """simple docstring"""







      def __init__(					self						:				Dict ,				_snake_case						:				List[str] ):
          super().__init__(A_ )
          __lowercase						:	Optional[Any]					=				config.num_labels
          __lowercase						:	Union[str, Any]					=				ResNetModel(A_ )
          # classification head
          __lowercase						:	List[str]					=				nn.Sequential(
              nn.Flatten() ,				nn.Linear(config.hidden_sizes[-1] ,				config.num_labels ) if config.num_labels > 0 else nn.Identity() ,				)
          # initialize weights and apply final processing
          self.post_init()




      @add_start_docstrings_to_model_forward(A_ )
      @add_code_sample_docstrings(
          checkpoint=_IMAGE_CLASS_CHECKPOINT ,				output_type=A_ ,				config_class=_CONFIG_FOR_DOC ,				expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT ,				)
      def 						snake_case_	(					self						:				int ,				_snake_case						:				Optional[torch.FloatTensor] = None ,				_snake_case						:				Optional[torch.LongTensor] = None ,				_snake_case						:				Optional[bool] = None ,				_snake_case						:				Optional[bool] = None ,				):
          __lowercase						:	Dict					=				return_dict if return_dict is not None else self.config.use_return_dict

          __lowercase						:	Dict					=				self.resnet(A_ ,				output_hidden_states=A_ ,				return_dict=A_ )

          __lowercase						:	Optional[Any]					=				outputs.pooler_output if return_dict else outputs[1]

          __lowercase						:	List[str]					=				self.classifier(A_ )

          __lowercase						:	List[str]					=				None

          if labels is not None:
              if self.config.problem_type is None:
                  if self.num_labels == 1:
                      __lowercase						:	List[Any]					=				'''regression'''
                  elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
                      __lowercase						:	Optional[Any]					=				'''single_label_classification'''
                  else:
                      __lowercase						:	str					=				'''multi_label_classification'''
              if self.config.problem_type == "regression":
                  __lowercase						:	List[str]					=				MSELoss()
                  if self.num_labels == 1:
                      __lowercase						:	List[Any]					=				loss_fct(logits.squeeze() ,				labels.squeeze() )
                  else:
                      __lowercase						:	Dict					=				loss_fct(A_ ,				A_ )
              elif self.config.problem_type == "single_label_classification":
                  __lowercase						:	int					=				CrossEntropyLoss()
                  __lowercase						:	Optional[int]					=				loss_fct(logits.view(-1 ,				self.num_labels ) ,				labels.view(-1 ) )
              elif self.config.problem_type == "multi_label_classification":
                  __lowercase						:	List[str]					=				BCEWithLogitsLoss()
                  __lowercase						:	Tuple					=				loss_fct(A_ ,				A_ )

          if not return_dict:
              __lowercase						:	int					=				(logits,) + outputs[2:]
              return (loss,) + output if loss is not None else output

          return ImageClassifierOutputWithNoAttention(loss=A_ ,				logits=A_ ,				hidden_states=outputs.hidden_states )






@add_start_docstrings(
    '''
    ResNet backbone, to be used with frameworks like DETR and MaskFormer.
    '''       ,       UpperCamelCase_       ,       )
class   __lowerCAmelCase		(		UpperCamelCase_       ,       UpperCamelCase_					):



      """simple docstring"""







      def __init__(					self						:				Dict ,				_snake_case						:				int ):
          super().__init__(A_ )
          super()._init_backbone(A_ )

          __lowercase						:	List[Any]					=				[config.embedding_size] + config.hidden_sizes
          __lowercase						:	Dict					=				ResNetEmbeddings(A_ )
          __lowercase						:	Any					=				ResNetEncoder(A_ )

          # initialize weights and apply final processing
          self.post_init()




      @add_start_docstrings_to_model_forward(A_ )
      @replace_return_docstrings(output_type=A_ ,				config_class=_CONFIG_FOR_DOC )
      def 						snake_case_	(					self						:				List[str] ,				_snake_case						:				Tensor ,				_snake_case						:				Optional[bool] = None ,				_snake_case						:				Optional[bool] = None ):
          __lowercase						:	Optional[Any]					=				return_dict if return_dict is not None else self.config.use_return_dict
          __lowercase						:	Optional[int]					=				(
              output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
          )

          __lowercase						:	Any					=				self.embedder(A_ )

          __lowercase						:	str					=				self.encoder(A_ ,				output_hidden_states=A_ ,				return_dict=A_ )

          __lowercase						:	Union[str, Any]					=				outputs.hidden_states

          __lowercase						:	List[Any]					=				()
          for idx, stage in enumerate(self.stage_names ):
              if stage in self.out_features:
                  feature_maps += (hidden_states[idx],)

          if not return_dict:
              __lowercase						:	Any					=				(feature_maps,)
              if output_hidden_states:
                  output += (outputs.hidden_states,)
              return output

          return BackboneOutput(
              feature_maps=A_ ,				hidden_states=outputs.hidden_states if output_hidden_states else None ,				attentions=A_ ,				)







import copy
from typing import Dict, List, Optional

from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ..auto import CONFIG_MAPPING


A__  :  Optional[Any]								=					{
    '''facebook/mask2former-swin-small-coco-instance''': (
        '''https://huggingface.co/facebook/mask2former-swin-small-coco-instance/blob/main/config.json'''
    )
    # See all Mask2Former models at https://huggingface.co/models?filter=mask2former
}

A__  :  Dict								=					logging.get_logger(__name__)



class 				__snake_case						(     UpperCamelCase_						):
							_a			      = '''mask2former'''
							_a			      = ['''swin''']
							_a			      = {'''hidden_size''': '''hidden_dim'''}




							def __init__(       self     :       Any				,				A_     :       Optional[Dict] = None				,				A_     :       int = 2_5_6				,				A_     :       int = 2_5_6				,				A_     :       int = 2_5_6				,				A_     :       int = 1_0_2_4				,				A_     :       str = "relu"				,				A_     :       int = 6				,				A_     :       int = 1_0				,				A_     :       int = 8				,				A_     :       float = 0.0				,				A_     :       int = 2_0_4_8				,				A_     :       bool = False				,				A_     :       bool = False				,				A_     :       int = 4				,				A_     :       int = 2_5_5				,				A_     :       int = 1_0_0				,				A_     :       float = 0.1				,				A_     :       float = 2.0				,				A_     :       float = 5.0				,				A_     :       float = 5.0				,				A_     :       int = 1_2_5_4_4				,				A_     :       float = 3.0				,				A_     :       float = 0.75				,				A_     :       float = 0.02				,				A_     :       float = 1.0				,				A_     :       bool = True				,				A_     :       List[int] = [4, 8, 1_6, 3_2]				,				A_     :       bool = None				,				**A_     :       Dict				,				):
									if backbone_config is None:
											logger.info('''`backbone_config` is `None`. Initializing the config with the default `Swin` backbone.''')
											lowerCAmelCase_			:     int									= CONFIG_MAPPING['''swin'''](
											    image_size=2_2_4				,				in_channels=3				,				patch_size=4				,				embed_dim=9_6				,				depths=[2, 2, 1_8, 2]				,				num_heads=[3, 6, 1_2, 2_4]				,				window_size=7				,				drop_path_rate=0.3				,				use_absolute_embeddings=A_				,				out_features=['''stage1''', '''stage2''', '''stage3''', '''stage4''']				,				)

									if isinstance(A_				,				A_):
											lowerCAmelCase_			:     List[Any]									= backbone_config.pop('''model_type''')
											lowerCAmelCase_			:     Optional[Any]									= CONFIG_MAPPING[backbone_model_type]
											lowerCAmelCase_			:     List[Any]									= config_class.from_dict(A_)

									# verify that the backbone is supported
									if backbone_config.model_type not in self.backbones_supported:
											logger.warning_once(
											    F"""Backbone {backbone_config.model_type} is not a supported model and may not be compatible with Mask2Former. """
											    F"""Supported model types: {",".join(self.backbones_supported)}""")

									lowerCAmelCase_			:     List[Any]									= backbone_config
									lowerCAmelCase_			:     str									= feature_size
									lowerCAmelCase_			:     Optional[Any]									= mask_feature_size
									lowerCAmelCase_			:     int									= hidden_dim
									lowerCAmelCase_			:     int									= encoder_feedforward_dim
									lowerCAmelCase_			:     Optional[int]									= activation_function
									lowerCAmelCase_			:     Any									= encoder_layers
									lowerCAmelCase_			:     Optional[Any]									= decoder_layers
									lowerCAmelCase_			:     Optional[Any]									= num_attention_heads
									lowerCAmelCase_			:     Optional[int]									= dropout
									lowerCAmelCase_			:     List[str]									= dim_feedforward
									lowerCAmelCase_			:     Optional[Any]									= pre_norm
									lowerCAmelCase_			:     List[str]									= enforce_input_projection
									lowerCAmelCase_			:     Tuple									= common_stride
									lowerCAmelCase_			:     Optional[Any]									= ignore_value
									lowerCAmelCase_			:     Optional[Any]									= num_queries
									lowerCAmelCase_			:     int									= no_object_weight
									lowerCAmelCase_			:     Tuple									= class_weight
									lowerCAmelCase_			:     int									= mask_weight
									lowerCAmelCase_			:     Dict									= dice_weight
									lowerCAmelCase_			:     str									= train_num_points
									lowerCAmelCase_			:     Dict									= oversample_ratio
									lowerCAmelCase_			:     Tuple									= importance_sample_ratio
									lowerCAmelCase_			:     List[str]									= init_std
									lowerCAmelCase_			:     List[str]									= init_xavier_std
									lowerCAmelCase_			:     Optional[Any]									= use_auxiliary_loss
									lowerCAmelCase_			:     List[Any]									= feature_strides
									lowerCAmelCase_			:     int									= output_auxiliary_logits
									lowerCAmelCase_			:     Optional[Any]									= decoder_layers

									super().__init__(**A_)




							@classmethod
							def    UpperCAmelCase__			(       cls     :       List[str]				,				A_     :       PretrainedConfig				,				**A_     :       List[Any]):
									return cls(
									    backbone_config=A_				,				**A_				,				)







							def    UpperCAmelCase__			(       self     :       List[Any]):
									lowerCAmelCase_			:     str									= copy.deepcopy(self.__dict__)
									lowerCAmelCase_			:     Dict									= self.backbone_config.to_dict()
									lowerCAmelCase_			:     Optional[int]									= self.__class__.model_type
									return output





