




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

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

   snake_case				:    Optional[int]           =	SwinConfig(
       embed_dim=lowercase						,depths=lowercase						,num_heads=lowercase						,window_size=lowercase						,out_features=["""stage1""", """stage2""", """stage3""", """stage4"""]						,)
   snake_case				:    Dict           =	UperNetConfig(
       backbone_config=lowercase						,auxiliary_in_channels=lowercase						,num_labels=lowercase						,idalabel=lowercase						,labelaid=lowercase						,)

   return config




def 						SCREAMING_SNAKE_CASE__ (   lowercase					)  ->       Tuple:
   snake_case				:    Dict           =	[]

   # 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 						SCREAMING_SNAKE_CASE__ (   lowercase						,lowercase						,lowercase					)  ->       List[Any]:
   snake_case				:    List[str]           =	dct.pop(lowercase					)
   snake_case				:    str           =	val




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




def 						SCREAMING_SNAKE_CASE__ (   lowercase					)  ->       int:
   snake_case     , snake_case				:    Optional[Any]           =	x.shape
   snake_case				:    Dict           =	x.reshape(lowercase						,4						,in_channel // 4					)
   snake_case				:    Optional[int]           =	x[:, [0, 2, 1, 3], :].transpose(1						,2					).reshape(lowercase						,lowercase					)
   return x




def 						SCREAMING_SNAKE_CASE__ (   lowercase					)  ->       Dict:
   snake_case     , snake_case				:    int           =	x.shape
   snake_case				:    Union[str, Any]           =	x.reshape(lowercase						,in_channel // 4						,4					)
   snake_case				:    List[Any]           =	x[:, :, [0, 2, 1, 3]].transpose(1						,2					).reshape(lowercase						,lowercase					)

   return x




def 						SCREAMING_SNAKE_CASE__ (   lowercase					)  ->       List[str]:
   snake_case				:    Dict           =	x.shape[0]
   snake_case				:    Optional[int]           =	x.reshape(4						,in_channel // 4					)
   snake_case				:    int           =	x[[0, 2, 1, 3], :].transpose(0						,1					).reshape(lowercase					)
   return x




def 						SCREAMING_SNAKE_CASE__ (   lowercase					)  ->       Any:
   snake_case				:    Optional[Any]           =	x.shape[0]
   snake_case				:    Dict           =	x.reshape(in_channel // 4						,4					)
   snake_case				:    str           =	x[:, [0, 2, 1, 3]].transpose(0						,1					).reshape(lowercase					)
   return x




def 						SCREAMING_SNAKE_CASE__ (   lowercase						,lowercase						,lowercase					)  ->       Any:
   snake_case				:    List[Any]           =	{
       """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""",
   }
   snake_case				:    List[str]           =	model_name_to_url[model_name]
   snake_case				:    Dict           =	torch.hub.load_state_dict_from_url(lowercase						,map_location="""cpu"""						,file_name=lowercase					)[
       """state_dict"""
   ]

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

   snake_case				:    Tuple           =	get_upernet_config(lowercase					)
   snake_case				:    Tuple           =	UperNetForSemanticSegmentation(lowercase					)
   model.eval()

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

   # rename keys
   snake_case				:    Dict           =	create_rename_keys(lowercase					)
   for src, dest in rename_keys:
      rename_key(lowercase						,lowercase						,lowercase					)
   read_in_q_k_v(lowercase						,config.backbone_config					)

   # fix downsample parameters
   for key, value in state_dict.items():
      if "downsample" in key:
         if "reduction" in key:
            snake_case				:    Tuple           =	reverse_correct_unfold_reduction_order(lowercase					)
         if "norm" in key:
            snake_case				:    str           =	reverse_correct_unfold_norm_order(lowercase					)

   model.load_state_dict(lowercase					)

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

   snake_case				:    int           =	SegformerImageProcessor()
   snake_case				:    str           =	processor(lowercase						,return_tensors="""pt"""					).pixel_values

   with torch.no_grad():
      snake_case				:    Tuple           =	model(lowercase					)
      snake_case				:    str           =	outputs.logits

   print(logits.shape					)
   print("""First values of logits:"""						,logits[0, 0, :3, :3]					)
   # assert values
   if model_name == "upernet-swin-tiny":
      snake_case				:    Tuple           =	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":
      snake_case				:    List[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":
      snake_case				:    Tuple           =	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":
      snake_case				:    Optional[int]           =	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]						,lowercase						,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(lowercase					)
      print(f"""Saving processor to {pytorch_dump_folder_path}"""					)
      processor.save_pretrained(lowercase					)

   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 warnings
from functools import wraps
from typing import Callable




def 						SCREAMING_SNAKE_CASE__ (   lowercase					)  ->       Callable:

   @wraps(lowercase					)
   def _inner_fn(*lowercase						,**lowercase					):
      warnings.warn(
          (f"""'{fn.__name__}' is experimental and might be subject to breaking changes in the future.""")						,lowercase						,)
      return fn(*lowercase						,**lowercase					)

   return _inner_fn





