



'''simple docstring'''





import argparse
import json

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

from transformers import SegformerImageProcessor, SwinConfig, UperNetConfig, UperNetForSemanticSegmentation

def        a_     (							_lowerCAmelCase       )      ->      int:
  __lowerCamelCase       :       Optional[int]           = 384
  __lowerCamelCase       :       List[str]           = 7
  if "tiny" in model_name:
    __lowerCamelCase       :       List[str]           = 96
    __lowerCamelCase       :       Optional[int]           = (2, 2, 6, 2)
    __lowerCamelCase       :       List[str]           = (3, 6, 12, 24)
  elif "small" in model_name:
    __lowerCamelCase       :       str           = 96
    __lowerCamelCase       :       Tuple           = (2, 2, 18, 2)
    __lowerCamelCase       :       List[str]           = (3, 6, 12, 24)
  elif "base" in model_name:
    __lowerCamelCase       :       Tuple           = 128
    __lowerCamelCase       :       int           = (2, 2, 18, 2)
    __lowerCamelCase       :       Any           = (4, 8, 16, 32)
    __lowerCamelCase       :       List[str]           = 12
    __lowerCamelCase       :       Optional[Any]           = 512
  elif "large" in model_name:
    __lowerCamelCase       :       Union[str, Any]           = 192
    __lowerCamelCase       :       List[Any]           = (2, 2, 18, 2)
    __lowerCamelCase       :       Tuple           = (6, 12, 24, 48)
    __lowerCamelCase       :       Optional[Any]           = 12
    __lowerCamelCase       :       List[str]           = 768

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

  __lowerCamelCase       :       Optional[Any]           = SwinConfig(
      embed_dim=_lowerCAmelCase       ,depths=_lowerCAmelCase       ,num_heads=_lowerCAmelCase       ,window_size=_lowerCAmelCase       ,out_features=['stage1', 'stage2', 'stage3', 'stage4']       ,)
  __lowerCamelCase       :       Dict           = UperNetConfig(
      backbone_config=_lowerCAmelCase       ,auxiliary_in_channels=_lowerCAmelCase       ,num_labels=_lowerCAmelCase       ,idalabel=_lowerCAmelCase       ,labelaid=_lowerCAmelCase       ,)

  return config

def        a_     (							_lowerCAmelCase       )      ->      Union[str, Any]:
  __lowerCamelCase       :       List[str]           = []

  # fmt: off
  # stem
  rename_keys.append(('backbone.patch_embed.projection.weight', 'backbone.embeddings.patch_embeddings.projection.weight')       )
  rename_keys.append(('backbone.patch_embed.projection.bias', 'backbone.embeddings.patch_embeddings.projection.bias')       )
  rename_keys.append(('backbone.patch_embed.norm.weight', 'backbone.embeddings.norm.weight')       )
  rename_keys.append(('backbone.patch_embed.norm.bias', 'backbone.embeddings.norm.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}.blocks.{j}.norm1.weight', F'backbone.encoder.layers.{i}.blocks.{j}.layernorm_before.weight')       )
      rename_keys.append((F'backbone.stages.{i}.blocks.{j}.norm1.bias', F'backbone.encoder.layers.{i}.blocks.{j}.layernorm_before.bias')       )
      rename_keys.append((F'backbone.stages.{i}.blocks.{j}.attn.w_msa.relative_position_bias_table', F'backbone.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_bias_table')       )
      rename_keys.append((F'backbone.stages.{i}.blocks.{j}.attn.w_msa.relative_position_index', F'backbone.encoder.layers.{i}.blocks.{j}.attention.self.relative_position_index')       )
      rename_keys.append((F'backbone.stages.{i}.blocks.{j}.attn.w_msa.proj.weight', F'backbone.encoder.layers.{i}.blocks.{j}.attention.output.dense.weight')       )
      rename_keys.append((F'backbone.stages.{i}.blocks.{j}.attn.w_msa.proj.bias', F'backbone.encoder.layers.{i}.blocks.{j}.attention.output.dense.bias')       )
      rename_keys.append((F'backbone.stages.{i}.blocks.{j}.norm2.weight', F'backbone.encoder.layers.{i}.blocks.{j}.layernorm_after.weight')       )
      rename_keys.append((F'backbone.stages.{i}.blocks.{j}.norm2.bias', F'backbone.encoder.layers.{i}.blocks.{j}.layernorm_after.bias')       )
      rename_keys.append((F'backbone.stages.{i}.blocks.{j}.ffn.layers.0.0.weight', F'backbone.encoder.layers.{i}.blocks.{j}.intermediate.dense.weight')       )
      rename_keys.append((F'backbone.stages.{i}.blocks.{j}.ffn.layers.0.0.bias', F'backbone.encoder.layers.{i}.blocks.{j}.intermediate.dense.bias')       )
      rename_keys.append((F'backbone.stages.{i}.blocks.{j}.ffn.layers.1.weight', F'backbone.encoder.layers.{i}.blocks.{j}.output.dense.weight')       )
      rename_keys.append((F'backbone.stages.{i}.blocks.{j}.ffn.layers.1.bias', F'backbone.encoder.layers.{i}.blocks.{j}.output.dense.bias')       )

    if i < 3:
      rename_keys.append((F'backbone.stages.{i}.downsample.reduction.weight', F'backbone.encoder.layers.{i}.downsample.reduction.weight')       )
      rename_keys.append((F'backbone.stages.{i}.downsample.norm.weight', F'backbone.encoder.layers.{i}.downsample.norm.weight')       )
      rename_keys.append((F'backbone.stages.{i}.downsample.norm.bias', F'backbone.encoder.layers.{i}.downsample.norm.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_     (							_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase       )      ->      List[str]:
  __lowerCamelCase       :       Dict           = dct.pop(_lowerCAmelCase       )
  __lowerCamelCase       :       List[Any]           = val

def        a_     (							_lowerCAmelCase       ,_lowerCAmelCase       )      ->      Optional[Any]:
  __lowerCamelCase       :       int           = [int(backbone_config.embed_dim * 2**i       ) for i in range(len(backbone_config.depths       )       )]
  for i in range(len(backbone_config.depths       )       ):
    __lowerCamelCase       :       str           = num_features[i]
    for j in range(backbone_config.depths[i]       ):
      # fmt: off
      # read in weights + bias of input projection layer (in original implementation, this is a single matrix + bias)
      __lowerCamelCase       :       Optional[int]           = state_dict.pop(F'backbone.stages.{i}.blocks.{j}.attn.w_msa.qkv.weight'       )
      __lowerCamelCase       :       Union[str, Any]           = state_dict.pop(F'backbone.stages.{i}.blocks.{j}.attn.w_msa.qkv.bias'       )
      # next, add query, keys and values (in that order) to the state dict
      __lowerCamelCase       :       Union[str, Any]           = in_proj_weight[:dim, :]
      __lowerCamelCase       :       str           = in_proj_bias[: dim]
      __lowerCamelCase       :       Union[str, Any]           = in_proj_weight[
          dim : dim * 2, :
      ]
      __lowerCamelCase       :       List[str]           = in_proj_bias[
          dim : dim * 2
      ]
      __lowerCamelCase       :       Optional[Any]           = in_proj_weight[
          -dim :, :
      ]
      __lowerCamelCase       :       str           = in_proj_bias[-dim :]
      # fmt: on

def        a_     (							_lowerCAmelCase       )      ->      Dict:
  __lowerCamelCase       ,__lowerCamelCase       :       List[Any]           = x.shape
  __lowerCamelCase       :       Any           = x.reshape(_lowerCAmelCase       ,4       ,in_channel // 4       )
  __lowerCamelCase       :       str           = x[:, [0, 2, 1, 3], :].transpose(1       ,2       ).reshape(_lowerCAmelCase       ,_lowerCAmelCase       )
  return x

def        a_     (							_lowerCAmelCase       )      ->      int:
  __lowerCamelCase       ,__lowerCamelCase       :       str           = x.shape
  __lowerCamelCase       :       Optional[int]           = x.reshape(_lowerCAmelCase       ,in_channel // 4       ,4       )
  __lowerCamelCase       :       int           = x[:, :, [0, 2, 1, 3]].transpose(1       ,2       ).reshape(_lowerCAmelCase       ,_lowerCAmelCase       )

  return x

def        a_     (							_lowerCAmelCase       )      ->      int:
  __lowerCamelCase       :       int           = x.shape[0]
  __lowerCamelCase       :       int           = x.reshape(4       ,in_channel // 4       )
  __lowerCamelCase       :       Union[str, Any]           = x[[0, 2, 1, 3], :].transpose(0       ,1       ).reshape(_lowerCAmelCase       )
  return x

def        a_     (							_lowerCAmelCase       )      ->      List[str]:
  __lowerCamelCase       :       List[str]           = x.shape[0]
  __lowerCamelCase       :       Any           = x.reshape(in_channel // 4       ,4       )
  __lowerCamelCase       :       List[Any]           = x[:, [0, 2, 1, 3]].transpose(0       ,1       ).reshape(_lowerCAmelCase       )
  return x

def        a_     (							_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase       )      ->      Optional[int]:
  __lowerCamelCase       :       Tuple           = {
      'upernet-swin-tiny': 'https://download.openmmlab.com/mmsegmentation/v0.5/swin/upernet_swin_tiny_patch4_window7_512x512_160k_ade20k_pretrain_224x224_1K/upernet_swin_tiny_patch4_window7_512x512_160k_ade20k_pretrain_224x224_1K_20210531_112542-e380ad3e.pth',
      'upernet-swin-small': 'https://download.openmmlab.com/mmsegmentation/v0.5/swin/upernet_swin_small_patch4_window7_512x512_160k_ade20k_pretrain_224x224_1K/upernet_swin_small_patch4_window7_512x512_160k_ade20k_pretrain_224x224_1K_20210526_192015-ee2fff1c.pth',
      'upernet-swin-base': 'https://download.openmmlab.com/mmsegmentation/v0.5/swin/upernet_swin_base_patch4_window12_512x512_160k_ade20k_pretrain_384x384_22K/upernet_swin_base_patch4_window12_512x512_160k_ade20k_pretrain_384x384_22K_20210531_125459-429057bf.pth',
      'upernet-swin-large': 'https://download.openmmlab.com/mmsegmentation/v0.5/swin/upernet_swin_large_patch4_window12_512x512_pretrain_384x384_22K_160k_ade20k/upernet_swin_large_patch4_window12_512x512_pretrain_384x384_22K_160k_ade20k_20220318_091743-9ba68901.pth',
  }
  __lowerCamelCase       :       Optional[Any]           = model_name_to_url[model_name]
  __lowerCamelCase       :       str           = torch.hub.load_state_dict_from_url(_lowerCAmelCase       ,map_location='cpu'       ,file_name=_lowerCAmelCase       )[
      'state_dict'
  ]

  for name, param in state_dict.items():
    print(_lowerCAmelCase       ,param.shape       )

  __lowerCamelCase       :       str           = get_upernet_config(_lowerCAmelCase       )
  __lowerCamelCase       :       int           = UperNetForSemanticSegmentation(_lowerCAmelCase       )
  model.eval()

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

  # rename keys
  __lowerCamelCase       :       List[Any]           = create_rename_keys(_lowerCAmelCase       )
  for src, dest in rename_keys:
    rename_key(_lowerCAmelCase       ,_lowerCAmelCase       ,_lowerCAmelCase       )
  read_in_q_k_v(_lowerCAmelCase       ,config.backbone_config       )

  # fix downsample parameters
  for key, value in state_dict.items():
    if "downsample" in key:
      if "reduction" in key:
        __lowerCamelCase       :       int           = reverse_correct_unfold_reduction_order(_lowerCAmelCase       )
      if "norm" in key:
        __lowerCamelCase       :       Optional[Any]           = reverse_correct_unfold_norm_order(_lowerCAmelCase       )

  model.load_state_dict(_lowerCAmelCase       )

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

  __lowerCamelCase       :       str           = SegformerImageProcessor()
  __lowerCamelCase       :       List[Any]           = processor(_lowerCAmelCase       ,return_tensors='pt'       ).pixel_values

  with torch.no_grad():
    __lowerCamelCase       :       Dict           = model(_lowerCAmelCase       )
    __lowerCamelCase       :       Tuple           = outputs.logits

  print(logits.shape       )
  print('First values of logits:'       ,logits[0, 0, :3, :3]       )
  # assert values
  if model_name == "upernet-swin-tiny":
    __lowerCamelCase       :       Union[str, Any]           = torch.tensor(
        [[-7.5958, -7.5958, -7.4302], [-7.5958, -7.5958, -7.4302], [-7.4797, -7.4797, -7.3068]]       )
  elif model_name == "upernet-swin-small":
    __lowerCamelCase       :       Optional[Any]           = torch.tensor(
        [[-7.1921, -7.1921, -6.9532], [-7.1921, -7.1921, -6.9532], [-7.0908, -7.0908, -6.8534]]       )
  elif model_name == "upernet-swin-base":
    __lowerCamelCase       :       Dict           = torch.tensor(
        [[-6.5851, -6.5851, -6.4330], [-6.5851, -6.5851, -6.4330], [-6.4763, -6.4763, -6.3254]]       )
  elif model_name == "upernet-swin-large":
    __lowerCamelCase       :       List[Any]           = torch.tensor(
        [[-7.5297, -7.5297, -7.3802], [-7.5297, -7.5297, -7.3802], [-7.4044, -7.4044, -7.2586]]       )
  print('Logits:'       ,outputs.logits[0, 0, :3, :3]       )
  assert torch.allclose(outputs.logits[0, 0, :3, :3]       ,_lowerCAmelCase       ,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(_lowerCAmelCase       )
    print(F'Saving processor to {pytorch_dump_folder_path}'       )
    processor.save_pretrained(_lowerCAmelCase       )

  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__":
   _UpperCamelCase											=						argparse.ArgumentParser()
   # Required parameters
   parser.add_argument(
       '--model_name',
       default='upernet-swin-tiny',
       type=str,
       choices=[f'''upernet-swin-{size}''' for size in ['tiny', 'small', 'base', 'large']],
       help='Name of the Swin + 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.'
   )

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











'''simple docstring'''





import argparse

from torch import nn

# transformers_old should correspond to branch `save_old_prophetnet_model_structure` here
# original prophetnet_checkpoints are saved under `patrickvonplaten/..._old` respectively
from transformers_old.modeling_prophetnet import (
    ProphetNetForConditionalGeneration as ProphetNetForConditionalGenerationOld,
)
from transformers_old.modeling_xlm_prophetnet import (
    XLMProphetNetForConditionalGeneration as XLMProphetNetForConditionalGenerationOld,
)

from transformers import ProphetNetForConditionalGeneration, XLMProphetNetForConditionalGeneration, logging


_UpperCamelCase											=						logging.get_logger(__name__)
logging.set_verbosity_info()

def        a_     (							_lowerCAmelCase       ,_lowerCAmelCase       )      ->      Optional[int]:
  if "xprophetnet" in prophetnet_checkpoint_path:
    __lowerCamelCase       :       Union[str, Any]           = XLMProphetNetForConditionalGenerationOld.from_pretrained(_lowerCAmelCase       )
    __lowerCamelCase       ,__lowerCamelCase       :       List[str]           = XLMProphetNetForConditionalGeneration.from_pretrained(
        _lowerCAmelCase       ,output_loading_info=_lowerCAmelCase       )
  else:
    __lowerCamelCase       :       Optional[int]           = ProphetNetForConditionalGenerationOld.from_pretrained(_lowerCAmelCase       )
    __lowerCamelCase       ,__lowerCamelCase       :       List[str]           = ProphetNetForConditionalGeneration.from_pretrained(
        _lowerCAmelCase       ,output_loading_info=_lowerCAmelCase       )

  __lowerCamelCase       :       Union[str, Any]           = ['key_proj', 'value_proj', 'query_proj']

  __lowerCamelCase       :       Optional[Any]           = {
      'self_attn': 'ngram_self_attn',
      'cross_attn': 'encoder_attn',
      'cross_attn_layer_norm': 'encoder_attn_layer_norm',
      'feed_forward_layer_norm': 'final_layer_norm',
      'feed_forward': '',
      'intermediate': 'fc1',
      'output': 'fc2',
      'key_proj': 'k_proj',
      'query_proj': 'q_proj',
      'value_proj': 'v_proj',
      'word_embeddings': 'embed_tokens',
      'embeddings_layer_norm': 'emb_layer_norm',
      'relative_pos_embeddings': 'relative_linear',
      'ngram_embeddings': 'ngram_input_embed',
      'position_embeddings': 'embed_positions',
  }

  for key in loading_info["missing_keys"]:
    __lowerCamelCase       :       Optional[int]           = key.split('.'       )

    if attributes[0] == "lm_head":
      __lowerCamelCase       :       Dict           = prophet
      __lowerCamelCase       :       List[Any]           = prophet_old
    else:
      __lowerCamelCase       :       Any           = prophet.prophetnet
      __lowerCamelCase       :       Any           = prophet_old.model

    __lowerCamelCase       :       Optional[Any]           = False
    for attribute in attributes:
      if attribute in mapping:
        __lowerCamelCase       :       Any           = mapping[attribute]
        if not hasattr(_lowerCAmelCase       ,_lowerCAmelCase       ) and len(_lowerCAmelCase       ) > 0:
          __lowerCamelCase       :       int           = attribute
      elif hasattr(_lowerCAmelCase       ,_lowerCAmelCase       ):
        __lowerCamelCase       :       Optional[int]           = attribute

      if attribute == "weight":
        assert old_model.weight.shape == model.weight.shape, "Shapes have to match!"
        __lowerCamelCase       :       List[Any]           = old_model.weight
        logger.info(F'{attribute} is initialized.'       )
        __lowerCamelCase       :       List[Any]           = True
        break
      elif attribute == "bias":
        assert old_model.bias.shape == model.bias.shape, "Shapes have to match!"
        __lowerCamelCase       :       List[Any]           = old_model.bias
        logger.info(F'{attribute} is initialized'       )
        __lowerCamelCase       :       Dict           = True
        break
      elif attribute in special_keys and hasattr(_lowerCAmelCase       ,'in_proj_weight'       ):
        __lowerCamelCase       :       Optional[Any]           = old_model.in_proj_weight.shape[0] // 3
        __lowerCamelCase       :       Optional[Any]           = getattr(_lowerCAmelCase       ,_lowerCAmelCase       )
        param.weight.shape == old_model.in_proj_weight[:embed_dim, :].shape, "Shapes have to match"
        param.bias.shape == old_model.in_proj_bias[:embed_dim].shape, "Shapes have to match"
        if attribute == "query_proj":
          __lowerCamelCase       :       Optional[int]           = nn.Parameter(old_model.in_proj_weight[:embed_dim, :]       )
          __lowerCamelCase       :       Dict           = nn.Parameter(old_model.in_proj_bias[:embed_dim]       )

        elif attribute == "key_proj":
          __lowerCamelCase       :       List[str]           = nn.Parameter(old_model.in_proj_weight[embed_dim : 2 * embed_dim, :]       )
          __lowerCamelCase       :       Dict           = nn.Parameter(old_model.in_proj_bias[embed_dim : 2 * embed_dim]       )
        elif attribute == "value_proj":
          __lowerCamelCase       :       str           = nn.Parameter(old_model.in_proj_weight[2 * embed_dim :, :]       )
          __lowerCamelCase       :       Optional[int]           = nn.Parameter(old_model.in_proj_bias[2 * embed_dim :]       )
        __lowerCamelCase       :       Optional[int]           = True
        break
      elif attribute == "position_embeddings":
        assert (
            model.position_embeddings.weight.shape[-1] == old_model.embed_positions.weight.shape[-1]
        ), "Hidden size has to match"
        assert model.position_embeddings.weight.shape[0] == 512, "We want 512 position_embeddings."
        __lowerCamelCase       :       Optional[int]           = nn.Parameter(old_model.embed_positions.weight[:512, :]       )
        __lowerCamelCase       :       Dict           = True
        break

      if attribute.isdigit():
        __lowerCamelCase       :       List[str]           = model[int(_lowerCAmelCase       )]
        __lowerCamelCase       :       Union[str, Any]           = old_model[int(_lowerCAmelCase       )]
      else:
        __lowerCamelCase       :       Union[str, Any]           = getattr(_lowerCAmelCase       ,_lowerCAmelCase       )

        if old_attribute == "":
          __lowerCamelCase       :       str           = old_model
        else:
          if not hasattr(_lowerCAmelCase       ,_lowerCAmelCase       ):
            raise ValueError(F'{old_model} does not have {old_attribute}'       )
          __lowerCamelCase       :       str           = getattr(_lowerCAmelCase       ,_lowerCAmelCase       )

    if not is_key_init:
      raise ValueError(F'{key} was not correctly initialized!'       )

  print(F'Saving model to {pytorch_dump_folder_path}'       )
  prophet.save_pretrained(_lowerCAmelCase       )


if __name__ == "__main__":
   _UpperCamelCase											=						argparse.ArgumentParser()
   # Required parameters
   parser.add_argument(
       '--prophetnet_checkpoint_path', default=None, type=str, required=True, help='Path the official PyTorch dump.'
   )
   parser.add_argument(
       '--pytorch_dump_folder_path', default=None, type=str, required=True, help='Path to the output PyTorch model.'
   )
   _UpperCamelCase											=						parser.parse_args()
   convert_prophetnet_checkpoint_to_pytorch(args.prophetnet_checkpoint_path, args.pytorch_dump_folder_path)






