"""simple docstring"""
import argparse
import json

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

from transformers import ConvNextConfig, SegformerImageProcessor, UperNetConfig, UperNetForSemanticSegmentation







def 						__a       (     _SCREAMING_SNAKE_CASE       )     ->int:
  a__:					Optional[int]      							=     384
  if "tiny" in model_name:
    a__:					Any      							=     [3, 3, 9, 3]
    a__:					int      							=     [96, 192, 384, 768]
  if "small" in model_name:
    a__:					List[Any]      							=     [3, 3, 27, 3]
    a__:					List[Any]      							=     [96, 192, 384, 768]
  if "base" in model_name:
    a__:					Optional[int]      							=     [3, 3, 27, 3]
    a__:					List[str]      							=     [128, 256, 512, 1024]
    a__:					str      							=     512
  if "large" in model_name:
    a__:					Dict      							=     [3, 3, 27, 3]
    a__:					int      							=     [192, 384, 768, 1536]
    a__:					Optional[int]      							=     768
  if "xlarge" in model_name:
    a__:					Dict      							=     [3, 3, 27, 3]
    a__:					Dict      							=     [256, 512, 1024, 2048]
    a__:					Union[str, Any]      							=     1024

  # set label information
  a__:					Dict      							=     150
  a__:					Any      							=     'huggingface/label-files'
  a__:					Dict      							=     'ade20k-id2label.json'
  a__:					Tuple      							=     json.load(open(hf_hub_download(_SCREAMING_SNAKE_CASE		,					_SCREAMING_SNAKE_CASE		,					repo_type='dataset'       )		,					'r'       )       )
  a__:					Tuple      							=     {int(_SCREAMING_SNAKE_CASE       ): v for k, v in idalabel.items()}
  a__:					Optional[int]      							=     {v: k for k, v in idalabel.items()}

  a__:					Tuple      							=     ConvNextConfig(
      depths=_SCREAMING_SNAKE_CASE		,					hidden_sizes=_SCREAMING_SNAKE_CASE		,					out_features=['stage1', 'stage2', 'stage3', 'stage4']       )
  a__:					Union[str, Any]      							=     UperNetConfig(
      backbone_config=_SCREAMING_SNAKE_CASE		,					auxiliary_in_channels=_SCREAMING_SNAKE_CASE		,					num_labels=_SCREAMING_SNAKE_CASE		,					idalabel=_SCREAMING_SNAKE_CASE		,					labelaid=_SCREAMING_SNAKE_CASE		,					)

  return config







def 						__a       (     _SCREAMING_SNAKE_CASE       )     ->int:
  a__:					Tuple      							=     []

  # fmt: off
  # stem
  rename_keys.append(('backbone.downsample_layers.0.0.weight', 'backbone.embeddings.patch_embeddings.weight')       )
  rename_keys.append(('backbone.downsample_layers.0.0.bias', 'backbone.embeddings.patch_embeddings.bias')       )
  rename_keys.append(('backbone.downsample_layers.0.1.weight', 'backbone.embeddings.layernorm.weight')       )
  rename_keys.append(('backbone.downsample_layers.0.1.bias', 'backbone.embeddings.layernorm.bias')       )
  # stages
  for i in range(len(config.backbone_config.depths       )       ):
    for j in range(config.backbone_config.depths[i]       ):
      rename_keys.append((F'backbone.stages.{i}.{j}.gamma', F'backbone.encoder.stages.{i}.layers.{j}.layer_scale_parameter')       )
      rename_keys.append((F'backbone.stages.{i}.{j}.depthwise_conv.weight', F'backbone.encoder.stages.{i}.layers.{j}.dwconv.weight')       )
      rename_keys.append((F'backbone.stages.{i}.{j}.depthwise_conv.bias', F'backbone.encoder.stages.{i}.layers.{j}.dwconv.bias')       )
      rename_keys.append((F'backbone.stages.{i}.{j}.norm.weight', F'backbone.encoder.stages.{i}.layers.{j}.layernorm.weight')       )
      rename_keys.append((F'backbone.stages.{i}.{j}.norm.bias', F'backbone.encoder.stages.{i}.layers.{j}.layernorm.bias')       )
      rename_keys.append((F'backbone.stages.{i}.{j}.pointwise_conv1.weight', F'backbone.encoder.stages.{i}.layers.{j}.pwconv1.weight')       )
      rename_keys.append((F'backbone.stages.{i}.{j}.pointwise_conv1.bias', F'backbone.encoder.stages.{i}.layers.{j}.pwconv1.bias')       )
      rename_keys.append((F'backbone.stages.{i}.{j}.pointwise_conv2.weight', F'backbone.encoder.stages.{i}.layers.{j}.pwconv2.weight')       )
      rename_keys.append((F'backbone.stages.{i}.{j}.pointwise_conv2.bias', F'backbone.encoder.stages.{i}.layers.{j}.pwconv2.bias')       )
    if i > 0:
      rename_keys.append((F'backbone.downsample_layers.{i}.0.weight', F'backbone.encoder.stages.{i}.downsampling_layer.0.weight')       )
      rename_keys.append((F'backbone.downsample_layers.{i}.0.bias', F'backbone.encoder.stages.{i}.downsampling_layer.0.bias')       )
      rename_keys.append((F'backbone.downsample_layers.{i}.1.weight', F'backbone.encoder.stages.{i}.downsampling_layer.1.weight')       )
      rename_keys.append((F'backbone.downsample_layers.{i}.1.bias', F'backbone.encoder.stages.{i}.downsampling_layer.1.bias')       )

    rename_keys.append((F'backbone.norm{i}.weight', F'backbone.hidden_states_norms.stage{i+1}.weight')       )
    rename_keys.append((F'backbone.norm{i}.bias', F'backbone.hidden_states_norms.stage{i+1}.bias')       )

  # decode head
  rename_keys.extend(
      [
          ('decode_head.conv_seg.weight', 'decode_head.classifier.weight'),
          ('decode_head.conv_seg.bias', 'decode_head.classifier.bias'),
          ('auxiliary_head.conv_seg.weight', 'auxiliary_head.classifier.weight'),
          ('auxiliary_head.conv_seg.bias', 'auxiliary_head.classifier.bias'),
      ]       )
  # fmt: on

  return rename_keys







def 						__a       (     _SCREAMING_SNAKE_CASE		,					_SCREAMING_SNAKE_CASE		,					_SCREAMING_SNAKE_CASE       )     ->Dict:
  a__:					Union[str, Any]      							=     dct.pop(_SCREAMING_SNAKE_CASE       )
  a__:					Union[str, Any]      							=     val







def 						__a       (     _SCREAMING_SNAKE_CASE		,					_SCREAMING_SNAKE_CASE		,					_SCREAMING_SNAKE_CASE       )     ->int:
  a__:					Optional[Any]      							=     {
      'upernet-convnext-tiny': 'https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_tiny_fp16_512x512_160k_ade20k/upernet_convnext_tiny_fp16_512x512_160k_ade20k_20220227_124553-cad485de.pth',
      'upernet-convnext-small': 'https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_small_fp16_512x512_160k_ade20k/upernet_convnext_small_fp16_512x512_160k_ade20k_20220227_131208-1b1e394f.pth',
      'upernet-convnext-base': 'https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_base_fp16_512x512_160k_ade20k/upernet_convnext_base_fp16_512x512_160k_ade20k_20220227_181227-02a24fc6.pth',
      'upernet-convnext-large': 'https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_large_fp16_640x640_160k_ade20k/upernet_convnext_large_fp16_640x640_160k_ade20k_20220226_040532-e57aa54d.pth',
      'upernet-convnext-xlarge': 'https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_xlarge_fp16_640x640_160k_ade20k/upernet_convnext_xlarge_fp16_640x640_160k_ade20k_20220226_080344-95fc38c2.pth',
  }
  a__:					Union[str, Any]      							=     model_name_to_url[model_name]
  a__:					Tuple      							=     torch.hub.load_state_dict_from_url(_SCREAMING_SNAKE_CASE		,					map_location='cpu'       )['state_dict']

  a__:					List[Any]      							=     get_upernet_config(_SCREAMING_SNAKE_CASE       )
  a__:					Union[str, Any]      							=     UperNetForSemanticSegmentation(_SCREAMING_SNAKE_CASE       )
  model.eval()

  # replace "bn" => "batch_norm"
  for key in state_dict.copy().keys():
    a__:					Optional[Any]      							=     state_dict.pop(_SCREAMING_SNAKE_CASE       )
    if "bn" in key:
      a__:					List[Any]      							=     key.replace('bn'		,					'batch_norm'       )
    a__:					List[Any]      							=     val

  # rename keys
  a__:					List[Any]      							=     create_rename_keys(_SCREAMING_SNAKE_CASE       )
  for src, dest in rename_keys:
    rename_key(_SCREAMING_SNAKE_CASE		,					_SCREAMING_SNAKE_CASE		,					_SCREAMING_SNAKE_CASE       )

  model.load_state_dict(_SCREAMING_SNAKE_CASE       )

  # verify on image
  a__:					Optional[int]      							=     'https://huggingface.co/datasets/hf-internal-testing/fixtures_ade20k/resolve/main/ADE_val_00000001.jpg'
  a__:					str      							=     Image.open(requests.get(_SCREAMING_SNAKE_CASE		,					stream=_SCREAMING_SNAKE_CASE       ).raw       ).convert('RGB'       )

  a__:					Tuple      							=     SegformerImageProcessor()
  a__:					List[str]      							=     processor(_SCREAMING_SNAKE_CASE		,					return_tensors='pt'       ).pixel_values

  with torch.no_grad():
    a__:					Dict      							=     model(_SCREAMING_SNAKE_CASE       )

  if model_name == "upernet-convnext-tiny":
    a__:					Optional[int]      							=     torch.tensor(
        [[-8.8_110, -8.8_110, -8.6_521], [-8.8_110, -8.8_110, -8.6_521], [-8.7_746, -8.7_746, -8.6_130]]       )
  elif model_name == "upernet-convnext-small":
    a__:					str      							=     torch.tensor(
        [[-8.8_236, -8.8_236, -8.6_771], [-8.8_236, -8.8_236, -8.6_771], [-8.7_638, -8.7_638, -8.6_240]]       )
  elif model_name == "upernet-convnext-base":
    a__:					str      							=     torch.tensor(
        [[-8.8_558, -8.8_558, -8.6_905], [-8.8_558, -8.8_558, -8.6_905], [-8.7_669, -8.7_669, -8.6_021]]       )
  elif model_name == "upernet-convnext-large":
    a__:					Any      							=     torch.tensor(
        [[-8.6_660, -8.6_660, -8.6_210], [-8.6_660, -8.6_660, -8.6_210], [-8.6_310, -8.6_310, -8.5_964]]       )
  elif model_name == "upernet-convnext-xlarge":
    a__:					Union[str, Any]      							=     torch.tensor(
        [[-8.4_980, -8.4_980, -8.3_977], [-8.4_980, -8.4_980, -8.3_977], [-8.4_379, -8.4_379, -8.3_412]]       )
  print('Logits:'		,					outputs.logits[0, 0, :3, :3]       )
  assert torch.allclose(outputs.logits[0, 0, :3, :3]		,					_SCREAMING_SNAKE_CASE		,					atol=1e-4       )
  print('Looks ok!'       )

  if pytorch_dump_folder_path is not None:
    print(F'Saving model {model_name} to {pytorch_dump_folder_path}'       )
    model.save_pretrained(_SCREAMING_SNAKE_CASE       )
    print(F'Saving processor to {pytorch_dump_folder_path}'       )
    processor.save_pretrained(_SCREAMING_SNAKE_CASE       )

  if push_to_hub:
    print(F'Pushing model and processor for {model_name} to hub'       )
    model.push_to_hub(F'openmmlab/{model_name}'       )
    processor.push_to_hub(F'openmmlab/{model_name}'       )


if __name__ == "__main__":
      lowercase__ 	=     argparse.ArgumentParser()
      # Required parameters
      parser.add_argument(
          '--model_name',
          default='upernet-convnext-tiny',
          type=str,
          choices=[f"upernet-convnext-{size}" for size in ['tiny', 'small', 'base', 'large', 'xlarge']],
          help='Name of the ConvNext UperNet model you\'d like to convert.',
      )
      parser.add_argument(
          '--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model directory.'
      )
      parser.add_argument(
          '--push_to_hub', action='store_true', help='Whether or not to push the converted model to the 🤗 hub.'
      )

      lowercase__ 	=     parser.parse_args()
      convert_upernet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)



"""simple docstring"""
import inspect
import unittest

from transformers import RegNetConfig
from transformers.file_utils import cached_property, is_torch_available, is_vision_available
from transformers.testing_utils import require_torch, require_vision, slow, torch_device

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


if is_torch_available():
 import torch
 from torch import nn

 from transformers import RegNetForImageClassification, RegNetModel
 from transformers.models.regnet.modeling_regnet import REGNET_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
 from PIL import Image

 from transformers import AutoImageProcessor




class        __snake_case	:



      def __init__(						self   ,      lowercase   ,      lowercase=3   ,      lowercase=32   ,      lowercase=3   ,      lowercase=10   ,      lowercase=[10, 20, 30, 40]   ,      lowercase=[1, 1, 2, 1]   ,      lowercase=True   ,      lowercase=True   ,      lowercase="relu"   ,      lowercase=3   ,      lowercase=None   ,      )     ->      Any:





        '''simple docstring'''





        a__:					Union[str, Any]      							=     parent
        a__:					int      							=     batch_size
        a__:					List[Any]      							=     image_size
        a__:					Any      							=     num_channels
        a__:					Dict      							=     embeddings_size
        a__:					str      							=     hidden_sizes
        a__:					List[Any]      							=     depths
        a__:					Optional[int]      							=     is_training
        a__:					Optional[int]      							=     use_labels
        a__:					Tuple      							=     hidden_act
        a__:					Any      							=     num_labels
        a__:					Union[str, Any]      							=     scope
        a__:					Tuple      							=     len(lowercase)



      def   lowerCamelCase_       (						self)     ->      int:





        '''simple docstring'''





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

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

        a__:					int      							=     self.get_config()

        return config, pixel_values, labels



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





        '''simple docstring'''





        return RegNetConfig(
            num_channels=self.num_channels   ,      embeddings_size=self.embeddings_size   ,      hidden_sizes=self.hidden_sizes   ,      depths=self.depths   ,      hidden_act=self.hidden_act   ,      num_labels=self.num_labels   ,      )



      def   lowerCamelCase_       (						self   ,      lowercase   ,      lowercase   ,      lowercase)     ->      Any:





        '''simple docstring'''





        a__:					Dict      							=     RegNetModel(config=lowercase)
        model.to(lowercase)
        model.eval()
        a__:					List[str]      							=     model(lowercase)
        # expected last hidden states: B, C, H // 32, W // 32
        self.parent.assertEqual(
            result.last_hidden_state.shape   ,      (self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)   ,      )



      def   lowerCamelCase_       (						self   ,      lowercase   ,      lowercase   ,      lowercase)     ->      List[Any]:





        '''simple docstring'''





        a__:					Tuple      							=     self.num_labels
        a__:					Tuple      							=     RegNetForImageClassification(lowercase)
        model.to(lowercase)
        model.eval()
        a__:					Any      							=     model(lowercase   ,      labels=lowercase)
        self.parent.assertEqual(result.logits.shape   ,      (self.batch_size, self.num_labels))



      def   lowerCamelCase_       (						self)     ->      str:





        '''simple docstring'''





        a__:					int      							=     self.prepare_config_and_inputs()
        a__							,   a__							,   a__:					List[Any]      							=     config_and_inputs
        a__:					Union[str, Any]      							=     {'pixel_values': pixel_values}
        return config, inputs_dict







@require_torch
class        __snake_case	(					__lowerCAmelCase	,				__lowerCAmelCase	,				unittest.TestCase    ):
      a__				=		(RegNetModel, RegNetForImageClassification) if is_torch_available() else ()
      a__				=		(
          {"""feature-extraction""": RegNetModel, """image-classification""": RegNetForImageClassification}
          if is_torch_available()
          else {}
      )

      a__				=		False
      a__				=		False
      a__				=		False
      a__				=		False



      def   lowerCamelCase_       (						self)     ->      Dict:





        '''simple docstring'''





        a__:					Dict      							=     RegNetModelTester(self)
        a__:					Dict      							=     ConfigTester(self   ,      config_class=lowercase   ,      has_text_modality=lowercase)



      def   lowerCamelCase_       (						self)     ->      Dict:





        '''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   lowerCamelCase_       (						self)     ->      Optional[Any]:





        '''simple docstring'''





        return



      @unittest.skip(reason='RegNet does not use inputs_embeds')
      def   lowerCamelCase_       (						self)     ->      List[str]:





        '''simple docstring'''





        pass



      @unittest.skip(reason='RegNet does not support input and output embeddings')
      def   lowerCamelCase_       (						self)     ->      Optional[Any]:





        '''simple docstring'''





        pass



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





        '''simple docstring'''





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

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

          a__:					List[str]      							=     ['pixel_values']
          self.assertListEqual(arg_names[:1]   ,      lowercase)



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





        '''simple docstring'''





        a__:					Any      							=     self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*lowercase)



      def   lowerCamelCase_       (						self)     ->      Dict:





        '''simple docstring'''





        a__							,   a__:					int      							=     self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
          a__:					int      							=     model_class(config=lowercase)
          for name, module in model.named_modules():
            if isinstance(lowercase   ,      (nn.BatchNormad, nn.GroupNorm)):
              self.assertTrue(
                  torch.all(module.weight == 1)   ,      msg=f'Parameter {name} of model {model_class} seems not properly initialized'   ,      )
              self.assertTrue(
                  torch.all(module.bias == 0)   ,      msg=f'Parameter {name} of model {model_class} seems not properly initialized'   ,      )



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





        '''simple docstring'''





        def check_hidden_states_output(lowercase   ,      lowercase   ,      lowercase):
          a__:					int      							=     model_class(lowercase)
          model.to(lowercase)
          model.eval()

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

          a__:					Any      							=     outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

          a__:					List[str]      							=     self.model_tester.num_stages
          self.assertEqual(len(lowercase)   ,      expected_num_stages + 1)

          # RegNet's feature maps are of shape (batch_size, num_channels, height, width)
          self.assertListEqual(
              list(hidden_states[0].shape[-2:])   ,      [self.model_tester.image_size // 2, self.model_tester.image_size // 2]   ,      )

        a__							,   a__:					Optional[Any]      							=     self.model_tester.prepare_config_and_inputs_for_common()
        a__:					Dict      							=     ['basic', 'bottleneck']
        for model_class in self.all_model_classes:
          for layer_type in layers_type:
            a__:					Dict      							=     layer_type
            a__:					Dict      							=     True
            check_hidden_states_output(lowercase   ,      lowercase   ,      lowercase)

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

            check_hidden_states_output(lowercase   ,      lowercase   ,      lowercase)



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





        '''simple docstring'''





        a__:					Optional[Any]      							=     self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_image_classification(*lowercase)



      @slow
      def   lowerCamelCase_       (						self)     ->      str:





        '''simple docstring'''





        for model_name in REGNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          a__:					Tuple      							=     RegNetModel.from_pretrained(lowercase)
          self.assertIsNotNone(lowercase)







def 						__a       (     )     ->Dict:
  a__:					str      							=     Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'       )
  return image




@require_torch
@require_vision
class        __snake_case	(					unittest.TestCase    ):



      @cached_property
      def   lowerCamelCase_       (						self)     ->      Any:





        '''simple docstring'''





        return (
            AutoImageProcessor.from_pretrained(REGNET_PRETRAINED_MODEL_ARCHIVE_LIST[0])
            if is_vision_available()
            else None
        )



      @slow
      def   lowerCamelCase_       (						self)     ->      Optional[int]:





        '''simple docstring'''





        a__:					Tuple      							=     RegNetForImageClassification.from_pretrained(REGNET_PRETRAINED_MODEL_ARCHIVE_LIST[0]).to(lowercase)

        a__:					Tuple      							=     self.default_image_processor
        a__:					str      							=     prepare_img()
        a__:					Optional[int]      							=     image_processor(images=lowercase   ,      return_tensors='pt').to(lowercase)

        # forward pass
        with torch.no_grad():
          a__:					List[Any]      							=     model(**lowercase)

        # verify the logits
        a__:					Tuple      							=     torch.Size((1, 10_00))
        self.assertEqual(outputs.logits.shape   ,      lowercase)

        a__:					Optional[Any]      							=     torch.tensor([-0.4180, -1.5051, -3.4836]).to(lowercase)

        self.assertTrue(torch.allclose(outputs.logits[0, :3]   ,      lowercase   ,      atol=1e-4))


