

"""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  (  UpperCamelCase_     :       Optional[Any])	->						Union[str, Any]:
   '''simple docstring'''


   __lowercase								=						384
   __lowercase								=						7
   if "tiny" in model_name:
      __lowercase								=						96
      __lowercase								=						(2, 2, 6, 2)
      __lowercase								=						(3, 6, 12, 24)
   elif "small" in model_name:
      __lowercase								=						96
      __lowercase								=						(2, 2, 18, 2)
      __lowercase								=						(3, 6, 12, 24)
   elif "base" in model_name:
      __lowercase								=						128
      __lowercase								=						(2, 2, 18, 2)
      __lowercase								=						(4, 8, 16, 32)
      __lowercase								=						12
      __lowercase								=						512
   elif "large" in model_name:
      __lowercase								=						192
      __lowercase								=						(2, 2, 18, 2)
      __lowercase								=						(6, 12, 24, 48)
      __lowercase								=						12
      __lowercase								=						768

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

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

   return config




def   _A  (  UpperCamelCase_     :       List[str])	->						Any:
   '''simple docstring'''


   __lowercase								=						[]

   # 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  (  UpperCamelCase_     :       Optional[Any],			UpperCamelCase_     :       Optional[int],			UpperCamelCase_     :       Tuple)	->						Tuple:
   '''simple docstring'''


   __lowercase								=						dct.pop(__a)
   __lowercase								=						val




def   _A  (  UpperCamelCase_     :       Optional[int],			UpperCamelCase_     :       str)	->						str:
   '''simple docstring'''


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




def   _A  (  UpperCamelCase_     :       Union[str, Any])	->						Optional[Any]:
   '''simple docstring'''


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




def   _A  (  UpperCamelCase_     :       Optional[int])	->						Tuple:
   '''simple docstring'''


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

   return x




def   _A  (  UpperCamelCase_     :       List[str])	->						Tuple:
   '''simple docstring'''


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




def   _A  (  UpperCamelCase_     :       int)	->						Any:
   '''simple docstring'''


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




def   _A  (  UpperCamelCase_     :       Union[str, Any],			UpperCamelCase_     :       Dict,			UpperCamelCase_     :       Union[str, Any])	->						List[str]:
   '''simple docstring'''


   __lowercase								=						{
       '''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''',
   }
   __lowercase								=						model_name_to_url[model_name]
   __lowercase								=						torch.hub.load_state_dict_from_url(__a,			map_location="cpu",			file_name=__a)[
       '''state_dict'''
   ]

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

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

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

   # rename keys
   __lowercase								=						create_rename_keys(__a)
   for src, dest in rename_keys:
      rename_key(__a,			__a,			__a)
   read_in_q_k_v(__a,			config.backbone_config)

   # fix downsample parameters
   for key, value in state_dict.items():
      if "downsample" in key:
         if "reduction" in key:
            __lowercase								=						reverse_correct_unfold_reduction_order(__a)
         if "norm" in key:
            __lowercase								=						reverse_correct_unfold_norm_order(__a)

   model.load_state_dict(__a)

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

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

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

   print(logits.shape)
   print("First values of logits:",			logits[0, 0, :3, :3])
   # assert values
   if model_name == "upernet-swin-tiny":
      __lowercase								=						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":
      __lowercase								=						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":
      __lowercase								=						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":
      __lowercase								=						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],			__a,			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(__a)
      print(F"""Saving processor to {pytorch_dump_folder_path}""")
      processor.save_pretrained(__a)

   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)







import re

import jax.numpy as jnp
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.random import PRNGKey

from ..utils import logging


a__				=		logging.get_logger(__name__)



def    __UpperCAmelCase (	__a      :			Dict	)			->     Tuple:

					"""simple docstring"""

					_a	:    Any         =		R'''\w+[.]\d+'''
					_a	:    Union[str, Any]         =		re.findall(__a		,__a	)
					for pat in pats:
										_a	:    int         =		key.replace(__a		,'''_'''.join(pat.split('''.'''	)	)	)
					return key



def    __UpperCAmelCase (	__a      :			List[str]		,__a      :			Union[str, Any]		,__a      :			Optional[int]	)			->     Tuple:

					"""simple docstring"""

					_a	:    Dict         =		pt_tuple_key[:-1] + ('''scale''',)
					if (
					    any('''norm''' in str_ for str_ in pt_tuple_key	)
					    and (pt_tuple_key[-1] == "bias")
					    and (pt_tuple_key[:-1] + ("bias",) not in random_flax_state_dict)
					    and (pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict)
					):
										_a	:    Dict         =		pt_tuple_key[:-1] + ('''scale''',)
										return renamed_pt_tuple_key, pt_tensor
					elif pt_tuple_key[-1] in ["weight", "gamma"] and pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict:
										_a	:    Optional[int]         =		pt_tuple_key[:-1] + ('''scale''',)
										return renamed_pt_tuple_key, pt_tensor

					# embedding
					if pt_tuple_key[-1] == "weight" and pt_tuple_key[:-1] + ("embedding",) in random_flax_state_dict:
										_a	:    Union[str, Any]         =		pt_tuple_key[:-1] + ('''embedding''',)
										return renamed_pt_tuple_key, pt_tensor

					# conv layer
					_a	:    Tuple         =		pt_tuple_key[:-1] + ('''kernel''',)
					if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4:
										_a	:    List[str]         =		pt_tensor.transpose(2		,3		,1		,0	)
										return renamed_pt_tuple_key, pt_tensor

					# linear layer
					_a	:    Dict         =		pt_tuple_key[:-1] + ('''kernel''',)
					if pt_tuple_key[-1] == "weight":
										_a	:    Union[str, Any]         =		pt_tensor.T
										return renamed_pt_tuple_key, pt_tensor

					# old PyTorch layer norm weight
					_a	:    Dict         =		pt_tuple_key[:-1] + ('''weight''',)
					if pt_tuple_key[-1] == "gamma":
										return renamed_pt_tuple_key, pt_tensor

					# old PyTorch layer norm bias
					_a	:    Union[str, Any]         =		pt_tuple_key[:-1] + ('''bias''',)
					if pt_tuple_key[-1] == "beta":
										return renamed_pt_tuple_key, pt_tensor

					return pt_tuple_key, pt_tensor







def    __UpperCAmelCase (	__a      :			Dict		,__a      :			str		,__a      :			str=42	)			->     Optional[int]:

					"""simple docstring"""

					_a	:    Tuple         =		{k: v.numpy() for k, v in pt_state_dict.items()}

					# Step 2: Since the model is stateless, get random Flax params
					_a	:    List[Any]         =		flax_model.init_weights(PRNGKey(__a	)	)

					_a	:    Optional[int]         =		flatten_dict(__a	)
					_a	:    Tuple         =		{}

					# Need to change some parameters name to match Flax names
					for pt_key, pt_tensor in pt_state_dict.items():
										_a	:    List[str]         =		rename_key(__a	)
										_a	:    Optional[Any]         =		tuple(renamed_pt_key.split('''.'''	)	)

										# Correctly rename weight parameters
										_a      ,					_a	:    List[str]         =		rename_key_and_reshape_tensor(__a		,__a		,__a	)

										if flax_key in random_flax_state_dict:
															if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
																				raise ValueError(
																				    F"""PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape """
																				    F"""{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}."""	)

        # also add unexpected weight so that warning is thrown
										_a	:    Dict         =		jnp.asarray(__a	)

					return unflatten_dict(__a	)



