




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

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

      lowerCAmelCase__      :			Optional[Any]   =	SwinConfig(
          embed_dim=__UpperCAmelCase    ,		depths=__UpperCAmelCase    ,		num_heads=__UpperCAmelCase    ,		window_size=__UpperCAmelCase    ,		out_features=["""stage1""", """stage2""", """stage3""", """stage4"""]    ,		)
      lowerCAmelCase__      :			str   =	UperNetConfig(
          backbone_config=__UpperCAmelCase    ,		auxiliary_in_channels=__UpperCAmelCase    ,		num_labels=__UpperCAmelCase    ,		idalabel=__UpperCAmelCase    ,		labelaid=__UpperCAmelCase    ,		)

      return config



def 	lowercase_				(			__UpperCAmelCase				)						->    str:
      lowerCAmelCase__      :			int   =	[]

      # 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 	lowercase_				(			__UpperCAmelCase    ,		__UpperCAmelCase    ,		__UpperCAmelCase				)						->    Optional[int]:
      lowerCAmelCase__      :			List[Any]   =	dct.pop(__UpperCAmelCase				)
      lowerCAmelCase__      :			List[str]   =	val



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



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



def 	lowercase_				(			__UpperCAmelCase				)						->    Dict:
      lowerCAmelCase__ ,   lowerCAmelCase__      :			List[Any]   =	x.shape
      lowerCAmelCase__      :			List[str]   =	x.reshape(__UpperCAmelCase    ,		in_channel // 4    ,		4				)
      lowerCAmelCase__      :			Dict   =	x[:, :, [0, 2, 1, 3]].transpose(1    ,		2				).reshape(__UpperCAmelCase    ,		__UpperCAmelCase				)

      return x



def 	lowercase_				(			__UpperCAmelCase				)						->    int:
      lowerCAmelCase__      :			Tuple   =	x.shape[0]
      lowerCAmelCase__      :			Any   =	x.reshape(4    ,		in_channel // 4				)
      lowerCAmelCase__      :			List[Any]   =	x[[0, 2, 1, 3], :].transpose(0    ,		1				).reshape(__UpperCAmelCase				)
      return x



def 	lowercase_				(			__UpperCAmelCase				)						->    List[str]:
      lowerCAmelCase__      :			List[Any]   =	x.shape[0]
      lowerCAmelCase__      :			Tuple   =	x.reshape(in_channel // 4    ,		4				)
      lowerCAmelCase__      :			Union[str, Any]   =	x[:, [0, 2, 1, 3]].transpose(0    ,		1				).reshape(__UpperCAmelCase				)
      return x



def 	lowercase_				(			__UpperCAmelCase    ,		__UpperCAmelCase    ,		__UpperCAmelCase				)						->    List[str]:
      lowerCAmelCase__      :			int   =	{
          """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__      :			str   =	model_name_to_url[model_name]
      lowerCAmelCase__      :			Optional[Any]   =	torch.hub.load_state_dict_from_url(__UpperCAmelCase    ,		map_location="""cpu"""    ,		file_name=__UpperCAmelCase				)[
          """state_dict"""
      ]

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

      lowerCAmelCase__      :			List[Any]   =	get_upernet_config(__UpperCAmelCase				)
      lowerCAmelCase__      :			Any   =	UperNetForSemanticSegmentation(__UpperCAmelCase				)
      model.eval()

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

      # rename keys
      lowerCAmelCase__      :			List[str]   =	create_rename_keys(__UpperCAmelCase				)
      for src, dest in rename_keys:
            rename_key(__UpperCAmelCase    ,		__UpperCAmelCase    ,		__UpperCAmelCase				)
      read_in_q_k_v(__UpperCAmelCase    ,		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(__UpperCAmelCase				)
                  if "norm" in key:
                        lowerCAmelCase__      :			int   =	reverse_correct_unfold_norm_order(__UpperCAmelCase				)

      model.load_state_dict(__UpperCAmelCase				)

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

      lowerCAmelCase__      :			Tuple   =	SegformerImageProcessor()
      lowerCAmelCase__      :			List[Any]   =	processor(__UpperCAmelCase    ,		return_tensors="""pt"""				).pixel_values

      with torch.no_grad():
            lowerCAmelCase__      :			Tuple   =	model(__UpperCAmelCase				)
            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__      :			List[str]   =	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__      :			List[str]   =	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__      :			List[str]   =	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]    ,		__UpperCAmelCase    ,		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(__UpperCAmelCase				)
            print(f"""Saving processor to {pytorch_dump_folder_path}"""				)
            processor.save_pretrained(__UpperCAmelCase				)

      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													=      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													=      parser.parse_args()
  convert_upernet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)











"""simple docstring"""


import copy
from dataclasses import dataclass, field
from typing import ClassVar, Dict

from ..features import Audio, ClassLabel, Features
from .base import TaskTemplate


@dataclass(frozen=a_						)
class _lowerCamelCase    (					a_						):
   _lowerCamelCase    :str     				=     field(default="audio-classification"       ,	metadata={"include_in_asdict_even_if_is_default": True}						)
   _lowerCamelCase    :ClassVar[Features]     				=     Features({"audio": Audio()}						)
   _lowerCamelCase    :ClassVar[Features]     				=     Features({"labels": ClassLabel}						)
   _lowerCamelCase    :str     				=     "audio"
   _lowerCamelCase    :str     				=     "labels"






   def 	_lowerCAmelCase   (		self       :		str   ,     UpperCamelCase       :		Optional[int]   )			->						Union[str, Any]:

         """simple docstring"""
         if self.label_column not in features:
               raise ValueError(f"""Column {self.label_column} is not present in features."""   )
         if not isinstance(features[self.label_column]   ,     UpperCamelCase   ):
               raise ValueError(f"""Column {self.label_column} is not a ClassLabel."""   )
         lowerCAmelCase__      :			str   =	copy.deepcopy(self   )
         lowerCAmelCase__      :			Optional[int]   =	self.label_schema.copy()
         lowerCAmelCase__      :			List[Any]   =	features[self.label_column]
         lowerCAmelCase__      :			Optional[int]   =	label_schema
         return task_template






   @property
   def 	_lowerCAmelCase   (		self       :		List[Any]   )			->						Dict[str, str]:

         """simple docstring"""
         return {
             self.audio_column: "audio",
             self.label_column: "labels",
         }





