'''simple docstring'''

import argparse
import collections
import json
from pathlib import Path

import requests
import torch
import yaml
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import (
    MobileViTImageProcessor,
    MobileViTVaConfig,
    MobileViTVaForImageClassification,
    MobileViTVaForSemanticSegmentation,
)
from transformers.utils import logging


logging.set_verbosity_info()
_A			:      List[str]					      =				logging.get_logger(__name__)







def      UpperCamelCase_						(					snake_case_       :       Optional[Any]      )					->     int:






 '''simple docstring'''

 print("""Loading config file..."""      )

 def flatten_yaml_as_dict(snake_case_       :       Union[str, Any]						,		snake_case_       :       List[str]=""						,		snake_case_       :       Optional[Any]="."      ):
  __lowerCAmelCase											=						[]
  for k, v in d.items():
   __lowerCAmelCase											=						parent_key + sep + k if parent_key else k
   if isinstance(snake_case_						,		collections.abc.MutableMapping      ):
    items.extend(flatten_yaml_as_dict(snake_case_						,		snake_case_						,		sep=snake_case_      ).items()      )
   else:
    items.append((new_key, v)      )
  return dict(snake_case_      )

 __lowerCAmelCase											=						argparse.Namespace()
 with open(snake_case_						,		"""r"""      ) as yaml_file:
  try:
   __lowerCAmelCase											=						yaml.load(snake_case_						,		Loader=yaml.FullLoader      )

   __lowerCAmelCase											=						flatten_yaml_as_dict(snake_case_      )
   for k, v in flat_cfg.items():
    setattr(snake_case_						,		snake_case_						,		snake_case_      )
  except yaml.YAMLError as exc:
   logger.error("""Error while loading config file: {}. Error message: {}""".format(snake_case_						,		str(snake_case_      )      )      )
 return config







def      UpperCamelCase_						(					snake_case_       :       Any						,		snake_case_       :       Union[str, Any]      )					->     int:






 '''simple docstring'''

 __lowerCAmelCase											=						MobileViTVaConfig()

 __lowerCAmelCase											=						False

 # dataset
 if task_name.startswith("""imagenet1k_"""      ):
  __lowerCAmelCase											=						10_00
  if int(task_name.strip().split("""_"""      )[-1]      ) == 3_84:
   __lowerCAmelCase											=						3_84
  else:
   __lowerCAmelCase											=						2_56
  __lowerCAmelCase											=						"""imagenet-1k-id2label.json"""
 elif task_name.startswith("""imagenet21k_to_1k_"""      ):
  __lowerCAmelCase											=						2_10_00
  if int(task_name.strip().split("""_"""      )[-1]      ) == 3_84:
   __lowerCAmelCase											=						3_84
  else:
   __lowerCAmelCase											=						2_56
  __lowerCAmelCase											=						"""imagenet-22k-id2label.json"""
 elif task_name.startswith("""ade20k_"""      ):
  __lowerCAmelCase											=						1_51
  __lowerCAmelCase											=						5_12
  __lowerCAmelCase											=						"""ade20k-id2label.json"""
  __lowerCAmelCase											=						True
 elif task_name.startswith("""voc_"""      ):
  __lowerCAmelCase											=						21
  __lowerCAmelCase											=						5_12
  __lowerCAmelCase											=						"""pascal-voc-id2label.json"""
  __lowerCAmelCase											=						True

 # orig_config
 __lowerCAmelCase											=						load_orig_config_file(snake_case_      )
 assert getattr(snake_case_						,		"""model.classification.name"""						,		-1      ) == "mobilevit_v2", "Invalid model"
 __lowerCAmelCase											=						getattr(snake_case_						,		"""model.classification.mitv2.width_multiplier"""						,		1.0      )
 assert (
     getattr(snake_case_						,		"""model.classification.mitv2.attn_norm_layer"""						,		-1      ) == "layer_norm_2d"
 ), "Norm layers other than layer_norm_2d is not supported"
 __lowerCAmelCase											=						getattr(snake_case_						,		"""model.classification.activation.name"""						,		"""swish"""      )
 # config.image_size == getattr(orig_config,  'sampler.bs.crop_size_width', 256)

 if is_segmentation_model:
  __lowerCAmelCase											=						getattr(snake_case_						,		"""model.segmentation.output_stride"""						,		16      )
  if "_deeplabv3" in task_name:
   __lowerCAmelCase											=						getattr(snake_case_						,		"""model.segmentation.deeplabv3.aspp_rates"""						,		[12, 24, 36]      )
   __lowerCAmelCase											=						getattr(snake_case_						,		"""model.segmentation.deeplabv3.aspp_out_channels"""						,		5_12      )
   __lowerCAmelCase											=						getattr(snake_case_						,		"""model.segmentation.deeplabv3.aspp_dropout"""						,		0.1      )

    # id2label
 __lowerCAmelCase											=						"""huggingface/label-files"""
 __lowerCAmelCase											=						json.load(open(hf_hub_download(snake_case_						,		snake_case_						,		repo_type="""dataset"""      )						,		"""r"""      )      )
 __lowerCAmelCase											=						{int(snake_case_      ): v for k, v in idalabel.items()}
 __lowerCAmelCase											=						idalabel
 __lowerCAmelCase											=						{v: k for k, v in idalabel.items()}

 return config







def      UpperCamelCase_						(					snake_case_       :       Union[str, Any]						,		snake_case_       :       Union[str, Any]						,		snake_case_       :       Union[str, Any]      )					->     str:






 '''simple docstring'''

 __lowerCAmelCase											=						dct.pop(snake_case_      )
 __lowerCAmelCase											=						val







def      UpperCamelCase_						(					snake_case_       :       str						,		snake_case_       :       Any=False      )					->     Any:






 '''simple docstring'''

 if base_model:
  __lowerCAmelCase											=						""""""
 else:
  __lowerCAmelCase											=						"""mobilevitv2."""

 __lowerCAmelCase											=						[]
 for k in state_dict.keys():
  if k[:8] == "encoder.":
   __lowerCAmelCase											=						k[8:]
  else:
   __lowerCAmelCase											=						k

  if ".block." in k:
   __lowerCAmelCase											=						k_new.replace(""".block."""						,		"""."""      )
  if ".conv." in k:
   __lowerCAmelCase											=						k_new.replace(""".conv."""						,		""".convolution."""      )
  if ".norm." in k:
   __lowerCAmelCase											=						k_new.replace(""".norm."""						,		""".normalization."""      )

  if "conv_1." in k:
   __lowerCAmelCase											=						k_new.replace("""conv_1."""						,		f"""{model_prefix}conv_stem."""      )
  for i in [1, 2]:
   if f"""layer_{i}.""" in k:
    __lowerCAmelCase											=						k_new.replace(f"""layer_{i}."""						,		f"""{model_prefix}encoder.layer.{i-1}.layer."""      )
  if ".exp_1x1." in k:
   __lowerCAmelCase											=						k_new.replace(""".exp_1x1."""						,		""".expand_1x1."""      )
  if ".red_1x1." in k:
   __lowerCAmelCase											=						k_new.replace(""".red_1x1."""						,		""".reduce_1x1."""      )

  for i in [3, 4, 5]:
   if f"""layer_{i}.0.""" in k:
    __lowerCAmelCase											=						k_new.replace(f"""layer_{i}.0."""						,		f"""{model_prefix}encoder.layer.{i-1}.downsampling_layer."""      )
   if f"""layer_{i}.1.local_rep.0.""" in k:
    __lowerCAmelCase											=						k_new.replace(f"""layer_{i}.1.local_rep.0."""						,		f"""{model_prefix}encoder.layer.{i-1}.conv_kxk."""      )
   if f"""layer_{i}.1.local_rep.1.""" in k:
    __lowerCAmelCase											=						k_new.replace(f"""layer_{i}.1.local_rep.1."""						,		f"""{model_prefix}encoder.layer.{i-1}.conv_1x1."""      )

  for i in [3, 4, 5]:
   if i == 3:
    __lowerCAmelCase											=						[0, 1]
   elif i == 4:
    __lowerCAmelCase											=						[0, 1, 2, 3]
   elif i == 5:
    __lowerCAmelCase											=						[0, 1, 2]

   for j in j_in:
    if f"""layer_{i}.1.global_rep.{j}.""" in k:
     __lowerCAmelCase											=						k_new.replace(
         f"""layer_{i}.1.global_rep.{j}."""						,		f"""{model_prefix}encoder.layer.{i-1}.transformer.layer.{j}."""      )
   if f"""layer_{i}.1.global_rep.{j+1}.""" in k:
    __lowerCAmelCase											=						k_new.replace(
        f"""layer_{i}.1.global_rep.{j+1}."""						,		f"""{model_prefix}encoder.layer.{i-1}.layernorm."""      )

   if f"""layer_{i}.1.conv_proj.""" in k:
    __lowerCAmelCase											=						k_new.replace(f"""layer_{i}.1.conv_proj."""						,		f"""{model_prefix}encoder.layer.{i-1}.conv_projection."""      )

  if "pre_norm_attn.0." in k:
   __lowerCAmelCase											=						k_new.replace("""pre_norm_attn.0."""						,		"""layernorm_before."""      )
  if "pre_norm_attn.1." in k:
   __lowerCAmelCase											=						k_new.replace("""pre_norm_attn.1."""						,		"""attention."""      )
  if "pre_norm_ffn.0." in k:
   __lowerCAmelCase											=						k_new.replace("""pre_norm_ffn.0."""						,		"""layernorm_after."""      )
  if "pre_norm_ffn.1." in k:
   __lowerCAmelCase											=						k_new.replace("""pre_norm_ffn.1."""						,		"""ffn.conv1."""      )
  if "pre_norm_ffn.3." in k:
   __lowerCAmelCase											=						k_new.replace("""pre_norm_ffn.3."""						,		"""ffn.conv2."""      )

  if "classifier.1." in k:
   __lowerCAmelCase											=						k_new.replace("""classifier.1."""						,		"""classifier."""      )

  if "seg_head." in k:
   __lowerCAmelCase											=						k_new.replace("""seg_head."""						,		"""segmentation_head."""      )
  if ".aspp_layer." in k:
   __lowerCAmelCase											=						k_new.replace(""".aspp_layer."""						,		"""."""      )
  if ".aspp_pool." in k:
   __lowerCAmelCase											=						k_new.replace(""".aspp_pool."""						,		"""."""      )

  rename_keys.append((k, k_new)      )
 return rename_keys







def      UpperCamelCase_						(					snake_case_       :       Optional[int]      )					->     Optional[Any]:






 '''simple docstring'''

 __lowerCAmelCase											=						[]
 for k in state_dict.keys():
  if k.startswith("""seg_head.aux_head."""      ):
   keys_to_ignore.append(snake_case_      )
 for k in keys_to_ignore:
  state_dict.pop(snake_case_						,		snake_case_      )







def      UpperCamelCase_						(					)					->     List[str]:






 '''simple docstring'''

 __lowerCAmelCase											=						"""http://images.cocodataset.org/val2017/000000039769.jpg"""
 # url = "https://cdn.britannica.com/86/141086-050-9D7C75EE/Gulfstream-G450-business-jet-passengers.jpg"
 __lowerCAmelCase											=						Image.open(requests.get(snake_case_						,		stream=snake_case_      ).raw      )
 return im







@torch.no_grad()
def      UpperCamelCase_						(					snake_case_       :       int						,		snake_case_       :       Tuple						,		snake_case_       :       Tuple						,		snake_case_       :       str      )					->     Union[str, Any]:






 '''simple docstring'''

 __lowerCAmelCase											=						get_mobilevitva_config(snake_case_						,		snake_case_      )

 # load original state_dict
 __lowerCAmelCase											=						torch.load(snake_case_						,		map_location="""cpu"""      )

 # load huggingface model
 if task_name.startswith("""ade20k_"""      ) or task_name.startswith("""voc_"""      ):
  __lowerCAmelCase											=						MobileViTVaForSemanticSegmentation(snake_case_      ).eval()
  __lowerCAmelCase											=						False
 else:
  __lowerCAmelCase											=						MobileViTVaForImageClassification(snake_case_      ).eval()
  __lowerCAmelCase											=						False

 # remove and rename some keys of load the original model
 __lowerCAmelCase											=						checkpoint
 remove_unused_keys(snake_case_      )
 __lowerCAmelCase											=						create_rename_keys(snake_case_						,		base_model=snake_case_      )
 for rename_key_src, rename_key_dest in rename_keys:
  rename_key(snake_case_						,		snake_case_						,		snake_case_      )

 # load modified state_dict
 model.load_state_dict(snake_case_      )

 # Check outputs on an image, prepared by MobileViTImageProcessor
 __lowerCAmelCase											=						MobileViTImageProcessor(crop_size=config.image_size						,		size=config.image_size + 32      )
 __lowerCAmelCase											=						image_processor(images=prepare_img()						,		return_tensors="""pt"""      )
 __lowerCAmelCase											=						model(**snake_case_      )

 # verify classification model
 if task_name.startswith("""imagenet"""      ):
  __lowerCAmelCase											=						outputs.logits
  __lowerCAmelCase											=						logits.argmax(-1      ).item()
  print("""Predicted class:"""						,		model.config.idalabel[predicted_class_idx]      )
  if task_name.startswith("""imagenet1k_256"""      ) and config.width_multiplier == 1.0:
   # expected_logits for base variant
   __lowerCAmelCase											=						torch.tensor([-1.6336E00, -7.3204E-02, -5.1883E-01]      )
   assert torch.allclose(logits[0, :3]						,		snake_case_						,		atol=1E-4      )

 Path(snake_case_      ).mkdir(exist_ok=snake_case_      )
 print(f"""Saving model {task_name} to {pytorch_dump_folder_path}"""      )
 model.save_pretrained(snake_case_      )
 print(f"""Saving image processor to {pytorch_dump_folder_path}"""      )
 image_processor.save_pretrained(snake_case_      )


if __name__ == "__main__":
 _A			:      int					      =				argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     '''--task''',
     default='''imagenet1k_256''',
     type=str,
     help=(
         '''Name of the task for which the MobileViTV2 model you\'d like to convert is trained on . '''
         '''
                Classification (ImageNet-1k)
                    - MobileViTV2 (256x256) : imagenet1k_256
                    - MobileViTV2 (Trained on 256x256 and Finetuned on 384x384) : imagenet1k_384
                    - MobileViTV2 (Trained on ImageNet-21k and Finetuned on ImageNet-1k 256x256) :
                      imagenet21k_to_1k_256
                    - MobileViTV2 (Trained on ImageNet-21k, Finetuned on ImageNet-1k 256x256, and Finetuned on
                      ImageNet-1k 384x384) : imagenet21k_to_1k_384
                Segmentation
                    - ADE20K Dataset : ade20k_deeplabv3
                    - Pascal VOC 2012 Dataset: voc_deeplabv3
            '''
     ),
     choices=[
         '''imagenet1k_256''',
         '''imagenet1k_384''',
         '''imagenet21k_to_1k_256''',
         '''imagenet21k_to_1k_384''',
         '''ade20k_deeplabv3''',
         '''voc_deeplabv3''',
     ],
 )

 parser.add_argument(
     '''--orig_checkpoint_path''', required=True, type=str, help='''Path to the original state dict (.pt file).'''
 )
 parser.add_argument('''--orig_config_path''', required=True, type=str, help='''Path to the original config file.''')
 parser.add_argument(
     '''--pytorch_dump_folder_path''', required=True, type=str, help='''Path to the output PyTorch model directory.'''
 )

 _A			:      Optional[int]					      =				parser.parse_args()
 convert_mobilevitva_checkpoint(
     args.task, args.orig_checkpoint_path, args.orig_config_path, args.pytorch_dump_folder_path
 )







'''simple docstring'''

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


_A			:      Optional[Any]					      =				logging.get_logger(__name__)

# General docstring
_A			:      Optional[Any]					      =				'''ResNetConfig'''

# Base docstring
_A			:      Tuple					      =				'''microsoft/resnet-50'''
_A			:      List[str]					      =				[1, 2048, 7, 7]

# Image classification docstring
_A			:      str					      =				'''microsoft/resnet-50'''
_A			:      Dict					      =				'''tiger cat'''

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







class 				_lowercase     (				nn.Module			):


 '''simple docstring'''






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





 def a					(      self   :	List[Any]      ,       SCREAMING_SNAKE_CASE__   :	Tensor							)  ->    Tensor:
  __lowerCAmelCase											=						self.convolution(SCREAMING_SNAKE_CASE__							)
  __lowerCAmelCase											=						self.normalization(SCREAMING_SNAKE_CASE__							)
  __lowerCAmelCase											=						self.activation(SCREAMING_SNAKE_CASE__							)
  return hidden_state







class 				_lowercase     (				nn.Module			):


 '''simple docstring'''






 def __init__(      self   :	Optional[int]      ,       SCREAMING_SNAKE_CASE__   :	ResNetConfig							)  ->    List[str]:
  super().__init__()
  __lowerCAmelCase											=						ResNetConvLayer(
      config.num_channels      ,       config.embedding_size      ,       kernel_size=7      ,       stride=2      ,       activation=config.hidden_act							)
  __lowerCAmelCase											=						nn.MaxPoolad(kernel_size=3      ,       stride=2      ,       padding=1							)
  __lowerCAmelCase											=						config.num_channels





 def a					(      self   :	List[Any]      ,       SCREAMING_SNAKE_CASE__   :	Tensor							)  ->    Tensor:
  __lowerCAmelCase											=						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."""							)
  __lowerCAmelCase											=						self.embedder(SCREAMING_SNAKE_CASE__							)
  __lowerCAmelCase											=						self.pooler(SCREAMING_SNAKE_CASE__							)
  return embedding







class 				_lowercase     (				nn.Module			):


 '''simple docstring'''






 def __init__(      self   :	List[Any]      ,       SCREAMING_SNAKE_CASE__   :	int      ,       SCREAMING_SNAKE_CASE__   :	int      ,       SCREAMING_SNAKE_CASE__   :	int = 2							)  ->    Dict:
  super().__init__()
  __lowerCAmelCase											=						nn.Convad(SCREAMING_SNAKE_CASE__      ,       SCREAMING_SNAKE_CASE__      ,       kernel_size=1      ,       stride=SCREAMING_SNAKE_CASE__      ,       bias=SCREAMING_SNAKE_CASE__							)
  __lowerCAmelCase											=						nn.BatchNormad(SCREAMING_SNAKE_CASE__							)





 def a					(      self   :	Optional[Any]      ,       SCREAMING_SNAKE_CASE__   :	Tensor							)  ->    Tensor:
  __lowerCAmelCase											=						self.convolution(SCREAMING_SNAKE_CASE__							)
  __lowerCAmelCase											=						self.normalization(SCREAMING_SNAKE_CASE__							)
  return hidden_state







class 				_lowercase     (				nn.Module			):


 '''simple docstring'''






 def __init__(      self   :	int      ,       SCREAMING_SNAKE_CASE__   :	int      ,       SCREAMING_SNAKE_CASE__   :	int      ,       SCREAMING_SNAKE_CASE__   :	int = 1      ,       SCREAMING_SNAKE_CASE__   :	str = "relu"							)  ->    Dict:
  super().__init__()
  __lowerCAmelCase											=						in_channels != out_channels or stride != 1
  __lowerCAmelCase											=						(
      ResNetShortCut(SCREAMING_SNAKE_CASE__      ,       SCREAMING_SNAKE_CASE__      ,       stride=SCREAMING_SNAKE_CASE__							) if should_apply_shortcut else nn.Identity()
  )
  __lowerCAmelCase											=						nn.Sequential(
      ResNetConvLayer(SCREAMING_SNAKE_CASE__      ,       SCREAMING_SNAKE_CASE__      ,       stride=SCREAMING_SNAKE_CASE__							)      ,       ResNetConvLayer(SCREAMING_SNAKE_CASE__      ,       SCREAMING_SNAKE_CASE__      ,       activation=SCREAMING_SNAKE_CASE__							)      ,       )
  __lowerCAmelCase											=						ACTaFN[activation]





 def a					(      self   :	Optional[Any]      ,       SCREAMING_SNAKE_CASE__   :	List[Any]							)  ->    int:
  __lowerCAmelCase											=						hidden_state
  __lowerCAmelCase											=						self.layer(SCREAMING_SNAKE_CASE__							)
  __lowerCAmelCase											=						self.shortcut(SCREAMING_SNAKE_CASE__							)
  hidden_state += residual
  __lowerCAmelCase											=						self.activation(SCREAMING_SNAKE_CASE__							)
  return hidden_state







class 				_lowercase     (				nn.Module			):


 '''simple docstring'''






 def __init__(      self   :	List[str]      ,       SCREAMING_SNAKE_CASE__   :	int      ,       SCREAMING_SNAKE_CASE__   :	int      ,       SCREAMING_SNAKE_CASE__   :	int = 1      ,       SCREAMING_SNAKE_CASE__   :	str = "relu"      ,       SCREAMING_SNAKE_CASE__   :	int = 4							)  ->    int:
  super().__init__()
  __lowerCAmelCase											=						in_channels != out_channels or stride != 1
  __lowerCAmelCase											=						out_channels // reduction
  __lowerCAmelCase											=						(
      ResNetShortCut(SCREAMING_SNAKE_CASE__      ,       SCREAMING_SNAKE_CASE__      ,       stride=SCREAMING_SNAKE_CASE__							) if should_apply_shortcut else nn.Identity()
  )
  __lowerCAmelCase											=						nn.Sequential(
      ResNetConvLayer(SCREAMING_SNAKE_CASE__      ,       SCREAMING_SNAKE_CASE__      ,       kernel_size=1							)      ,       ResNetConvLayer(SCREAMING_SNAKE_CASE__      ,       SCREAMING_SNAKE_CASE__      ,       stride=SCREAMING_SNAKE_CASE__							)      ,       ResNetConvLayer(SCREAMING_SNAKE_CASE__      ,       SCREAMING_SNAKE_CASE__      ,       kernel_size=1      ,       activation=SCREAMING_SNAKE_CASE__							)      ,       )
  __lowerCAmelCase											=						ACTaFN[activation]





 def a					(      self   :	Any      ,       SCREAMING_SNAKE_CASE__   :	Optional[int]							)  ->    Tuple:
  __lowerCAmelCase											=						hidden_state
  __lowerCAmelCase											=						self.layer(SCREAMING_SNAKE_CASE__							)
  __lowerCAmelCase											=						self.shortcut(SCREAMING_SNAKE_CASE__							)
  hidden_state += residual
  __lowerCAmelCase											=						self.activation(SCREAMING_SNAKE_CASE__							)
  return hidden_state







class 				_lowercase     (				nn.Module			):


 '''simple docstring'''






 def __init__(      self   :	Optional[int]      ,       SCREAMING_SNAKE_CASE__   :	ResNetConfig      ,       SCREAMING_SNAKE_CASE__   :	int      ,       SCREAMING_SNAKE_CASE__   :	int      ,       SCREAMING_SNAKE_CASE__   :	int = 2      ,       SCREAMING_SNAKE_CASE__   :	int = 2      ,       )  ->    int:
  super().__init__()

  __lowerCAmelCase											=						ResNetBottleNeckLayer if config.layer_type == """bottleneck""" else ResNetBasicLayer

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





 def a					(      self   :	List[Any]      ,       SCREAMING_SNAKE_CASE__   :	Tensor							)  ->    Tensor:
  __lowerCAmelCase											=						input
  for layer in self.layers:
   __lowerCAmelCase											=						layer(SCREAMING_SNAKE_CASE__							)
  return hidden_state







class 				_lowercase     (				nn.Module			):


 '''simple docstring'''






 def __init__(      self   :	Optional[int]      ,       SCREAMING_SNAKE_CASE__   :	ResNetConfig							)  ->    Optional[int]:
  super().__init__()
  __lowerCAmelCase											=						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(
          SCREAMING_SNAKE_CASE__      ,       config.embedding_size      ,       config.hidden_sizes[0]      ,       stride=2 if config.downsample_in_first_stage else 1      ,       depth=config.depths[0]      ,       )							)
  __lowerCAmelCase											=						zip(config.hidden_sizes      ,       config.hidden_sizes[1:]							)
  for (in_channels, out_channels), depth in zip(SCREAMING_SNAKE_CASE__      ,       config.depths[1:]							):
   self.stages.append(ResNetStage(SCREAMING_SNAKE_CASE__      ,       SCREAMING_SNAKE_CASE__      ,       SCREAMING_SNAKE_CASE__      ,       depth=SCREAMING_SNAKE_CASE__							)							)





 def a					(      self   :	List[str]      ,       SCREAMING_SNAKE_CASE__   :	Tensor      ,       SCREAMING_SNAKE_CASE__   :	bool = False      ,       SCREAMING_SNAKE_CASE__   :	bool = True							)  ->    BaseModelOutputWithNoAttention:
  __lowerCAmelCase											=						() if output_hidden_states else None

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

   __lowerCAmelCase											=						stage_module(SCREAMING_SNAKE_CASE__							)

  if output_hidden_states:
   __lowerCAmelCase											=						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=SCREAMING_SNAKE_CASE__      ,       hidden_states=SCREAMING_SNAKE_CASE__      ,       )







class 				_lowercase     (				UpperCAmelCase__			):


 '''simple docstring'''






 _SCREAMING_SNAKE_CASE     :       int      	=					ResNetConfig
 _SCREAMING_SNAKE_CASE     :       Union[str, Any]      	=					"""resnet"""
 _SCREAMING_SNAKE_CASE     :       Union[str, Any]      	=					"""pixel_values"""
 _SCREAMING_SNAKE_CASE     :       Union[str, Any]      	=					True
 def a					(      self   :	Optional[int]      ,       SCREAMING_SNAKE_CASE__   :	Optional[int]							)  ->    str:
  if isinstance(SCREAMING_SNAKE_CASE__      ,       nn.Convad							):
   nn.init.kaiming_normal_(module.weight      ,       mode="""fan_out"""      ,       nonlinearity="""relu"""							)
  elif isinstance(SCREAMING_SNAKE_CASE__      ,       (nn.BatchNormad, nn.GroupNorm)							):
   nn.init.constant_(module.weight      ,       1							)
   nn.init.constant_(module.bias      ,       0							)





 def a					(      self   :	List[Any]      ,       SCREAMING_SNAKE_CASE__   :	Optional[int]      ,       SCREAMING_SNAKE_CASE__   :	Union[str, Any]=False							)  ->    int:
  if isinstance(SCREAMING_SNAKE_CASE__      ,       SCREAMING_SNAKE_CASE__							):
   __lowerCAmelCase											=						value







_A			:      Dict					      =				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.
'''

_A			:      Optional[int]					      =				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 				_lowercase     (				UpperCAmelCase__			):


 '''simple docstring'''






 def __init__(      self   :	Union[str, Any]      ,       SCREAMING_SNAKE_CASE__   :	Optional[int]							)  ->    List[str]:
  super().__init__(SCREAMING_SNAKE_CASE__							)
  __lowerCAmelCase											=						config
  __lowerCAmelCase											=						ResNetEmbeddings(SCREAMING_SNAKE_CASE__							)
  __lowerCAmelCase											=						ResNetEncoder(SCREAMING_SNAKE_CASE__							)
  __lowerCAmelCase											=						nn.AdaptiveAvgPoolad((1, 1)							)
  # Initialize weights and apply final processing
  self.post_init()





 @add_start_docstrings_to_model_forward(SCREAMING_SNAKE_CASE__							)
 @add_code_sample_docstrings(
     checkpoint=_CHECKPOINT_FOR_DOC      ,       output_type=SCREAMING_SNAKE_CASE__      ,       config_class=_CONFIG_FOR_DOC      ,       modality="""vision"""      ,       expected_output=_EXPECTED_OUTPUT_SHAPE      ,       )
 def a					(      self   :	List[Any]      ,       SCREAMING_SNAKE_CASE__   :	Tensor      ,       SCREAMING_SNAKE_CASE__   :	Optional[bool] = None      ,       SCREAMING_SNAKE_CASE__   :	Optional[bool] = None							)  ->    BaseModelOutputWithPoolingAndNoAttention:
  __lowerCAmelCase											=						(
      output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
  )
  __lowerCAmelCase											=						return_dict if return_dict is not None else self.config.use_return_dict

  __lowerCAmelCase											=						self.embedder(SCREAMING_SNAKE_CASE__							)

  __lowerCAmelCase											=						self.encoder(
      SCREAMING_SNAKE_CASE__      ,       output_hidden_states=SCREAMING_SNAKE_CASE__      ,       return_dict=SCREAMING_SNAKE_CASE__							)

  __lowerCAmelCase											=						encoder_outputs[0]

  __lowerCAmelCase											=						self.pooler(SCREAMING_SNAKE_CASE__							)

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

  return BaseModelOutputWithPoolingAndNoAttention(
      last_hidden_state=SCREAMING_SNAKE_CASE__      ,       pooler_output=SCREAMING_SNAKE_CASE__      ,       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 				_lowercase     (				UpperCAmelCase__			):


 '''simple docstring'''






 def __init__(      self   :	List[str]      ,       SCREAMING_SNAKE_CASE__   :	Tuple							)  ->    Any:
  super().__init__(SCREAMING_SNAKE_CASE__							)
  __lowerCAmelCase											=						config.num_labels
  __lowerCAmelCase											=						ResNetModel(SCREAMING_SNAKE_CASE__							)
  # classification head
  __lowerCAmelCase											=						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(SCREAMING_SNAKE_CASE__							)
 @add_code_sample_docstrings(
     checkpoint=_IMAGE_CLASS_CHECKPOINT      ,       output_type=SCREAMING_SNAKE_CASE__      ,       config_class=_CONFIG_FOR_DOC      ,       expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT      ,       )
 def a					(      self   :	int      ,       SCREAMING_SNAKE_CASE__   :	Optional[torch.FloatTensor] = None      ,       SCREAMING_SNAKE_CASE__   :	Optional[torch.LongTensor] = None      ,       SCREAMING_SNAKE_CASE__   :	Optional[bool] = None      ,       SCREAMING_SNAKE_CASE__   :	Optional[bool] = None      ,       )  ->    ImageClassifierOutputWithNoAttention:
  __lowerCAmelCase											=						return_dict if return_dict is not None else self.config.use_return_dict

  __lowerCAmelCase											=						self.resnet(SCREAMING_SNAKE_CASE__      ,       output_hidden_states=SCREAMING_SNAKE_CASE__      ,       return_dict=SCREAMING_SNAKE_CASE__							)

  __lowerCAmelCase											=						outputs.pooler_output if return_dict else outputs[1]

  __lowerCAmelCase											=						self.classifier(SCREAMING_SNAKE_CASE__							)

  __lowerCAmelCase											=						None

  if labels is not None:
   if self.config.problem_type is None:
    if self.num_labels == 1:
     __lowerCAmelCase											=						"""regression"""
    elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
     __lowerCAmelCase											=						"""single_label_classification"""
    else:
     __lowerCAmelCase											=						"""multi_label_classification"""
   if self.config.problem_type == "regression":
    __lowerCAmelCase											=						MSELoss()
    if self.num_labels == 1:
     __lowerCAmelCase											=						loss_fct(logits.squeeze()      ,       labels.squeeze()							)
    else:
     __lowerCAmelCase											=						loss_fct(SCREAMING_SNAKE_CASE__      ,       SCREAMING_SNAKE_CASE__							)
   elif self.config.problem_type == "single_label_classification":
    __lowerCAmelCase											=						CrossEntropyLoss()
    __lowerCAmelCase											=						loss_fct(logits.view(-1      ,       self.num_labels							)      ,       labels.view(-1							)							)
   elif self.config.problem_type == "multi_label_classification":
    __lowerCAmelCase											=						BCEWithLogitsLoss()
    __lowerCAmelCase											=						loss_fct(SCREAMING_SNAKE_CASE__      ,       SCREAMING_SNAKE_CASE__							)

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

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







@add_start_docstrings(
    """
    ResNet backbone, to be used with frameworks like DETR and MaskFormer.
    """  ,							UpperCAmelCase__  ,							)
class 				_lowercase     (				UpperCAmelCase__  ,							UpperCAmelCase__			):


 '''simple docstring'''






 def __init__(      self   :	List[str]      ,       SCREAMING_SNAKE_CASE__   :	List[Any]							)  ->    Tuple:
  super().__init__(SCREAMING_SNAKE_CASE__							)
  super()._init_backbone(SCREAMING_SNAKE_CASE__							)

  __lowerCAmelCase											=						[config.embedding_size] + config.hidden_sizes
  __lowerCAmelCase											=						ResNetEmbeddings(SCREAMING_SNAKE_CASE__							)
  __lowerCAmelCase											=						ResNetEncoder(SCREAMING_SNAKE_CASE__							)

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





 @add_start_docstrings_to_model_forward(SCREAMING_SNAKE_CASE__							)
 @replace_return_docstrings(output_type=SCREAMING_SNAKE_CASE__      ,       config_class=_CONFIG_FOR_DOC							)
 def a					(      self   :	List[Any]      ,       SCREAMING_SNAKE_CASE__   :	Tensor      ,       SCREAMING_SNAKE_CASE__   :	Optional[bool] = None      ,       SCREAMING_SNAKE_CASE__   :	Optional[bool] = None							)  ->    BackboneOutput:
  __lowerCAmelCase											=						return_dict if return_dict is not None else self.config.use_return_dict
  __lowerCAmelCase											=						(
      output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
  )

  __lowerCAmelCase											=						self.embedder(SCREAMING_SNAKE_CASE__							)

  __lowerCAmelCase											=						self.encoder(SCREAMING_SNAKE_CASE__      ,       output_hidden_states=SCREAMING_SNAKE_CASE__      ,       return_dict=SCREAMING_SNAKE_CASE__							)

  __lowerCAmelCase											=						outputs.hidden_states

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

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

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






