






def 			_SCREAMING_SNAKE_CASE						( __lowercase   :			Tuple    )						->      bool:





    """simple docstring"""

    if not isinstance(_SCREAMING_SNAKE_CASE			,				_SCREAMING_SNAKE_CASE    ):
        raise ValueError("""Input series is not valid, valid series - [2, 4, 6]"""    )
    if len(_SCREAMING_SNAKE_CASE    ) == 0:
        raise ValueError("""Input list must be a non empty list"""    )
    if len(_SCREAMING_SNAKE_CASE    ) == 1:
        return True
    __A         =	series[1] - series[0]
    for index in range(len(_SCREAMING_SNAKE_CASE    ) - 1    ):
        if series[index + 1] - series[index] != common_diff:
            return False
    return True



def 			_SCREAMING_SNAKE_CASE						( __lowercase   :			Dict    )						->      float:





    """simple docstring"""

    if not isinstance(_SCREAMING_SNAKE_CASE			,				_SCREAMING_SNAKE_CASE    ):
        raise ValueError("""Input series is not valid, valid series - [2, 4, 6]"""    )
    if len(_SCREAMING_SNAKE_CASE    ) == 0:
        raise ValueError("""Input list must be a non empty list"""    )
    __A         =	0
    for val in series:
        answer += val
    return answer / len(_SCREAMING_SNAKE_CASE    )


if __name__ == "__main__":
  import doctest

  doctest.testmod()












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   :			Optional[Any]    )						->      Dict:





    """simple docstring"""

    __A         =	3_8_4
    __A         =	7
    if "tiny" in model_name:
        __A         =	9_6
        __A         =	(2, 2, 6, 2)
        __A         =	(3, 6, 1_2, 2_4)
    elif "small" in model_name:
        __A         =	9_6
        __A         =	(2, 2, 1_8, 2)
        __A         =	(3, 6, 1_2, 2_4)
    elif "base" in model_name:
        __A         =	1_2_8
        __A         =	(2, 2, 1_8, 2)
        __A         =	(4, 8, 1_6, 3_2)
        __A         =	1_2
        __A         =	5_1_2
    elif "large" in model_name:
        __A         =	1_9_2
        __A         =	(2, 2, 1_8, 2)
        __A         =	(6, 1_2, 2_4, 4_8)
        __A         =	1_2
        __A         =	7_6_8

    # set label information
    __A         =	1_5_0
    __A         =	"""huggingface/label-files"""
    __A         =	"""ade20k-id2label.json"""
    __A         =	json.load(open(hf_hub_download(__lowercase			,				__lowercase			,				repo_type="""dataset"""    )			,				"""r"""    )    )
    __A         =	{int(__lowercase    ): v for k, v in idalabel.items()}
    __A         =	{v: k for k, v in idalabel.items()}

    __A         =	SwinConfig(
        embed_dim=__lowercase			,				depths=__lowercase			,				num_heads=__lowercase			,				window_size=__lowercase			,				out_features=["""stage1""", """stage2""", """stage3""", """stage4"""]			,				)
    __A         =	UperNetConfig(
        backbone_config=__lowercase			,				auxiliary_in_channels=__lowercase			,				num_labels=__lowercase			,				idalabel=__lowercase			,				labelaid=__lowercase			,				)

    return config



def 			_SCREAMING_SNAKE_CASE						( __lowercase   :			Optional[Any]    )						->      Dict:





    """simple docstring"""

    __A         =	[]

    # 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   :			Optional[Any]			,				__lowercase   :			int			,				__lowercase   :			Optional[int]    )						->      Any:





    """simple docstring"""

    __A         =	dct.pop(__lowercase    )
    __A         =	val



def 			_SCREAMING_SNAKE_CASE						( __lowercase   :			List[str]			,				__lowercase   :			int    )						->      Any:





    """simple docstring"""

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



def 			_SCREAMING_SNAKE_CASE						( __lowercase   :			Any    )						->      List[str]:





    """simple docstring"""

    __A    ,	__A         =	x.shape
    __A         =	x.reshape(__lowercase			,				4			,				in_channel // 4    )
    __A         =	x[:, [0, 2, 1, 3], :].transpose(1			,				2    ).reshape(__lowercase			,				__lowercase    )
    return x



def 			_SCREAMING_SNAKE_CASE						( __lowercase   :			Optional[Any]    )						->      Dict:





    """simple docstring"""

    __A    ,	__A         =	x.shape
    __A         =	x.reshape(__lowercase			,				in_channel // 4			,				4    )
    __A         =	x[:, :, [0, 2, 1, 3]].transpose(1			,				2    ).reshape(__lowercase			,				__lowercase    )

    return x



def 			_SCREAMING_SNAKE_CASE						( __lowercase   :			Tuple    )						->      Tuple:





    """simple docstring"""

    __A         =	x.shape[0]
    __A         =	x.reshape(4			,				in_channel // 4    )
    __A         =	x[[0, 2, 1, 3], :].transpose(0			,				1    ).reshape(__lowercase    )
    return x



def 			_SCREAMING_SNAKE_CASE						( __lowercase   :			Optional[int]    )						->      Optional[Any]:





    """simple docstring"""

    __A         =	x.shape[0]
    __A         =	x.reshape(in_channel // 4			,				4    )
    __A         =	x[:, [0, 2, 1, 3]].transpose(0			,				1    ).reshape(__lowercase    )
    return x



def 			_SCREAMING_SNAKE_CASE						( __lowercase   :			Dict			,				__lowercase   :			List[Any]			,				__lowercase   :			str    )						->      Tuple:





    """simple docstring"""

    __A         =	{
        """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""",
    }
    __A         =	model_name_to_url[model_name]
    __A         =	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    )

    __A         =	get_upernet_config(__lowercase    )
    __A         =	UperNetForSemanticSegmentation(__lowercase    )
    model.eval()

    # replace "bn" => "batch_norm"
    for key in state_dict.copy().keys():
        __A         =	state_dict.pop(__lowercase    )
        if "bn" in key:
            __A         =	key.replace("""bn"""			,				"""batch_norm"""    )
        __A         =	val

    # rename keys
    __A         =	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:
                __A         =	reverse_correct_unfold_reduction_order(__lowercase    )
            if "norm" in key:
                __A         =	reverse_correct_unfold_norm_order(__lowercase    )

    model.load_state_dict(__lowercase    )

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

    __A         =	SegformerImageProcessor()
    __A         =	processor(__lowercase			,				return_tensors="""pt"""    ).pixel_values

    with torch.no_grad():
        __A         =	model(__lowercase    )
        __A         =	outputs.logits

    print(logits.shape    )
    print("""First values of logits:"""			,				logits[0, 0, :3, :3]    )
    # assert values
    if model_name == "upernet-swin-tiny":
        __A         =	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":
        __A         =	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":
        __A         =	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":
        __A         =	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]			,				__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__":
  __a      :       Union[str, 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."
  )

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




