





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     )						-> List[Any]:
 UpperCamelCase	:					int       	=   384
 UpperCamelCase	:					Optional[Any]       	=   7
 if "tiny" in model_name:
  UpperCamelCase	:					Optional[Any]       	=   96
  UpperCamelCase	:					Union[str, Any]       	=   (2, 2, 6, 2)
  UpperCamelCase	:					str       	=   (3, 6, 12, 24)
 elif "small" in model_name:
  UpperCamelCase	:					Any       	=   96
  UpperCamelCase	:					Dict       	=   (2, 2, 18, 2)
  UpperCamelCase	:					Tuple       	=   (3, 6, 12, 24)
 elif "base" in model_name:
  UpperCamelCase	:					Optional[int]       	=   128
  UpperCamelCase	:					Optional[Any]       	=   (2, 2, 18, 2)
  UpperCamelCase	:					Optional[Any]       	=   (4, 8, 16, 32)
  UpperCamelCase	:					Union[str, Any]       	=   12
  UpperCamelCase	:					Dict       	=   512
 elif "large" in model_name:
  UpperCamelCase	:					Dict       	=   192
  UpperCamelCase	:					List[Any]       	=   (2, 2, 18, 2)
  UpperCamelCase	:					Tuple       	=   (6, 12, 24, 48)
  UpperCamelCase	:					List[str]       	=   12
  UpperCamelCase	:					Optional[int]       	=   768

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

 UpperCamelCase	:					Tuple       	=   SwinConfig(
     embed_dim=_lowerCAmelCase							,   depths=_lowerCAmelCase							,   num_heads=_lowerCAmelCase							,   window_size=_lowerCAmelCase							,   out_features=["stage1", "stage2", "stage3", "stage4"]							,   )
 UpperCamelCase	:					List[Any]       	=   UperNetConfig(
     backbone_config=_lowerCAmelCase							,   auxiliary_in_channels=_lowerCAmelCase							,   num_labels=_lowerCAmelCase							,   idalabel=_lowerCAmelCase							,   labelaid=_lowerCAmelCase							,   )

 return config



def 		A_				(     _lowerCAmelCase     )						-> Optional[Any]:
 UpperCamelCase	:					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[Any]:
 UpperCamelCase	:					Optional[int]       	=   dct.pop(_lowerCAmelCase     )
 UpperCamelCase	:					Any       	=   val



def 		A_				(     _lowerCAmelCase							,   _lowerCAmelCase     )						-> Dict:
 UpperCamelCase	:					Tuple       	=   [int(backbone_config.embed_dim * 2**i     ) for i in range(len(backbone_config.depths     )     )]
 for i in range(len(backbone_config.depths     )     ):
  UpperCamelCase	:					Dict       	=   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)
   UpperCamelCase	:					Union[str, Any]       	=   state_dict.pop(F"""backbone.stages.{i}.blocks.{j}.attn.w_msa.qkv.weight"""     )
   UpperCamelCase	:					List[str]       	=   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
   UpperCamelCase	:					Any       	=   in_proj_weight[:dim, :]
   UpperCamelCase	:					Dict       	=   in_proj_bias[: dim]
   UpperCamelCase	:					str       	=   in_proj_weight[
       dim : dim * 2, :
   ]
   UpperCamelCase	:					Optional[Any]       	=   in_proj_bias[
       dim : dim * 2
   ]
   UpperCamelCase	:					Dict       	=   in_proj_weight[
       -dim :, :
   ]
   UpperCamelCase	:					Dict       	=   in_proj_bias[-dim :]
   # fmt: on



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



def 		A_				(     _lowerCAmelCase     )						-> Tuple:
 UpperCamelCase	,    UpperCamelCase	:					List[Any]       	=   x.shape
 UpperCamelCase	:					Optional[int]       	=   x.reshape(_lowerCAmelCase							,   in_channel // 4							,   4     )
 UpperCamelCase	:					List[Any]       	=   x[:, :, [0, 2, 1, 3]].transpose(1							,   2     ).reshape(_lowerCAmelCase							,   _lowerCAmelCase     )

 return x



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



def 		A_				(     _lowerCAmelCase     )						-> Optional[int]:
 UpperCamelCase	:					Tuple       	=   x.shape[0]
 UpperCamelCase	:					List[Any]       	=   x.reshape(in_channel // 4							,   4     )
 UpperCamelCase	:					Optional[Any]       	=   x[:, [0, 2, 1, 3]].transpose(0							,   1     ).reshape(_lowerCAmelCase     )
 return x



def 		A_				(     _lowerCAmelCase							,   _lowerCAmelCase							,   _lowerCAmelCase     )						-> Dict:
 UpperCamelCase	:					Dict       	=   {
     "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",
 }
 UpperCamelCase	:					Dict       	=   model_name_to_url[model_name]
 UpperCamelCase	:					Optional[int]       	=   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     )

 UpperCamelCase	:					Dict       	=   get_upernet_config(_lowerCAmelCase     )
 UpperCamelCase	:					Optional[int]       	=   UperNetForSemanticSegmentation(_lowerCAmelCase     )
 model.eval()

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

 # rename keys
 UpperCamelCase	:					Optional[int]       	=   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:
    UpperCamelCase	:					str       	=   reverse_correct_unfold_reduction_order(_lowerCAmelCase     )
   if "norm" in key:
    UpperCamelCase	:					List[Any]       	=   reverse_correct_unfold_norm_order(_lowerCAmelCase     )

 model.load_state_dict(_lowerCAmelCase     )

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

 UpperCamelCase	:					Dict       	=   SegformerImageProcessor()
 UpperCamelCase	:					str       	=   processor(_lowerCAmelCase							,   return_tensors="pt"     ).pixel_values

 with torch.no_grad():
  UpperCamelCase	:					Any       	=   model(_lowerCAmelCase     )
  UpperCamelCase	:					Union[str, Any]       	=   outputs.logits

 print(logits.shape     )
 print("First values of logits:"							,   logits[0, 0, :3, :3]     )
 # assert values
 if model_name == "upernet-swin-tiny":
  UpperCamelCase	:					List[str]       	=   torch.tensor(
      [[-7.5_958, -7.5_958, -7.4_302], [-7.5_958, -7.5_958, -7.4_302], [-7.4_797, -7.4_797, -7.3_068]]     )
 elif model_name == "upernet-swin-small":
  UpperCamelCase	:					int       	=   torch.tensor(
      [[-7.1_921, -7.1_921, -6.9_532], [-7.1_921, -7.1_921, -6.9_532], [-7.0_908, -7.0_908, -6.8_534]]     )
 elif model_name == "upernet-swin-base":
  UpperCamelCase	:					List[str]       	=   torch.tensor(
      [[-6.5_851, -6.5_851, -6.4_330], [-6.5_851, -6.5_851, -6.4_330], [-6.4_763, -6.4_763, -6.3_254]]     )
 elif model_name == "upernet-swin-large":
  UpperCamelCase	:					Optional[Any]       	=   torch.tensor(
      [[-7.5_297, -7.5_297, -7.3_802], [-7.5_297, -7.5_297, -7.3_802], [-7.4_044, -7.4_044, -7.2_586]]     )
 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__":
    __lowerCamelCase       :     List[Any]   =     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."""
    )

    __lowerCamelCase       :     Optional[Any]   =     parser.parse_args()
    convert_upernet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)







import inspect
import re

from transformers.utils import direct_transformers_import


# All paths are set with the intent you should run this script from the root of the repo with the command
# python utils/check_config_docstrings.py
__lowerCamelCase       :     str   =     """src/transformers"""


# This is to make sure the transformers module imported is the one in the repo.
__lowerCamelCase       :     Tuple   =     direct_transformers_import(PATH_TO_TRANSFORMERS)

__lowerCamelCase       :     List[str]   =     transformers.models.auto.configuration_auto.CONFIG_MAPPING

# Regex pattern used to find the checkpoint mentioned in the docstring of `config_class`.
# For example, `[bert-base-uncased](https://huggingface.co/bert-base-uncased)`
__lowerCamelCase       :     Optional[Any]   =     re.compile(r"""\[(.+?)\]\((https://huggingface\.co/.+?)\)""")


__lowerCamelCase       :     List[str]   =     {
    """DecisionTransformerConfig""",
    """EncoderDecoderConfig""",
    """MusicgenConfig""",
    """RagConfig""",
    """SpeechEncoderDecoderConfig""",
    """TimmBackboneConfig""",
    """VisionEncoderDecoderConfig""",
    """VisionTextDualEncoderConfig""",
    """LlamaConfig""",
}



def 		A_				(     _lowerCAmelCase     )						-> List[str]:
 UpperCamelCase	:					Optional[Any]       	=   None

 # source code of `config_class`
 UpperCamelCase	:					Tuple       	=   inspect.getsource(_lowerCAmelCase     )
 UpperCamelCase	:					Optional[Any]       	=   _re_checkpoint.findall(_lowerCAmelCase     )

 # Each `checkpoint` is a tuple of a checkpoint name and a checkpoint link.
 # For example, `('bert-base-uncased', 'https://huggingface.co/bert-base-uncased')`
 for ckpt_name, ckpt_link in checkpoints:
  # allow the link to end with `/`
  if ckpt_link.endswith("/"     ):
   UpperCamelCase	:					Dict       	=   ckpt_link[:-1]

  # verify the checkpoint name corresponds to the checkpoint link
  UpperCamelCase	:					Any       	=   F"""https://huggingface.co/{ckpt_name}"""
  if ckpt_link == ckpt_link_from_name:
   UpperCamelCase	:					List[Any]       	=   ckpt_name
   break

 return checkpoint



def 		A_				(     )						-> List[str]:
 UpperCamelCase	:					Optional[int]       	=   []

 for config_class in list(CONFIG_MAPPING.values()     ):
  # Skip deprecated models
  if "models.deprecated" in config_class.__module__:
   continue
  UpperCamelCase	:					Union[str, Any]       	=   get_checkpoint_from_config_class(_lowerCAmelCase     )

  UpperCamelCase	:					Optional[int]       	=   config_class.__name__
  if checkpoint is None and name not in CONFIG_CLASSES_TO_IGNORE_FOR_DOCSTRING_CHECKPOINT_CHECK:
   configs_without_checkpoint.append(_lowerCAmelCase     )

 if len(_lowerCAmelCase     ) > 0:
  UpperCamelCase	:					Any       	=   "\n".join(sorted(_lowerCAmelCase     )     )
  raise ValueError(F"""The following configurations don't contain any valid checkpoint:\n{message}"""     )


if __name__ == "__main__":
    check_config_docstrings_have_checkpoints()
