




"""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      __UpperCAmelCase     ( __UpperCamelCase	):
			__lowercase			:  Optional[Any]								=       3_84
			__lowercase			:  List[Any]								=       7
			if "tiny" in model_name:
						__lowercase			:  List[str]								=       96
						__lowercase			:  Dict								=       (2, 2, 6, 2)
						__lowercase			:  Tuple								=       (3, 6, 12, 24)
			elif "small" in model_name:
						__lowercase			:  Optional[Any]								=       96
						__lowercase			:  int								=       (2, 2, 18, 2)
						__lowercase			:  List[str]								=       (3, 6, 12, 24)
			elif "base" in model_name:
						__lowercase			:  List[Any]								=       1_28
						__lowercase			:  List[str]								=       (2, 2, 18, 2)
						__lowercase			:  Tuple								=       (4, 8, 16, 32)
						__lowercase			:  Union[str, Any]								=       12
						__lowercase			:  Union[str, Any]								=       5_12
			elif "large" in model_name:
						__lowercase			:  List[Any]								=       1_92
						__lowercase			:  Union[str, Any]								=       (2, 2, 18, 2)
						__lowercase			:  Optional[int]								=       (6, 12, 24, 48)
						__lowercase			:  Union[str, Any]								=       12
						__lowercase			:  Optional[Any]								=       7_68

			# set label information
			__lowercase			:  Any								=       1_50
			__lowercase			:  List[str]								=       '''huggingface/label-files'''
			__lowercase			:  int								=       '''ade20k-id2label.json'''
			__lowercase			:  str								=       json.load(open(hf_hub_download(__UpperCamelCase	,				__UpperCamelCase	,				repo_type='''dataset'''	)	,				'''r'''	)	)
			__lowercase			:  List[str]								=       {int(__UpperCamelCase	): v for k, v in idalabel.items()}
			__lowercase			:  Any								=       {v: k for k, v in idalabel.items()}

			__lowercase			:  Any								=       SwinConfig(
			    embed_dim=__UpperCamelCase	,				depths=__UpperCamelCase	,				num_heads=__UpperCamelCase	,				window_size=__UpperCamelCase	,				out_features=['''stage1''', '''stage2''', '''stage3''', '''stage4''']	,				)
			__lowercase			:  List[str]								=       UperNetConfig(
			    backbone_config=__UpperCamelCase	,				auxiliary_in_channels=__UpperCamelCase	,				num_labels=__UpperCamelCase	,				idalabel=__UpperCamelCase	,				labelaid=__UpperCamelCase	,				)

			return config







def      __UpperCAmelCase     ( __UpperCamelCase	):
			__lowercase			:  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      __UpperCAmelCase     ( __UpperCamelCase	,				__UpperCamelCase	,				__UpperCamelCase	):
			__lowercase			:  List[Any]								=       dct.pop(__UpperCamelCase	)
			__lowercase			:  Optional[Any]								=       val







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







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







def      __UpperCAmelCase     ( __UpperCamelCase	):
			__lowercase			,__lowercase			:  List[Any]								=       x.shape
			__lowercase			:  Tuple								=       x.reshape(__UpperCamelCase	,				in_channel // 4	,				4	)
			__lowercase			:  int								=       x[:, :, [0, 2, 1, 3]].transpose(1	,				2	).reshape(__UpperCamelCase	,				__UpperCamelCase	)

			return x







def      __UpperCAmelCase     ( __UpperCamelCase	):
			__lowercase			:  Optional[int]								=       x.shape[0]
			__lowercase			:  Dict								=       x.reshape(4	,				in_channel // 4	)
			__lowercase			:  int								=       x[[0, 2, 1, 3], :].transpose(0	,				1	).reshape(__UpperCamelCase	)
			return x







def      __UpperCAmelCase     ( __UpperCamelCase	):
			__lowercase			:  Dict								=       x.shape[0]
			__lowercase			:  Any								=       x.reshape(in_channel // 4	,				4	)
			__lowercase			:  Optional[Any]								=       x[:, [0, 2, 1, 3]].transpose(0	,				1	).reshape(__UpperCamelCase	)
			return x







def      __UpperCAmelCase     ( __UpperCamelCase	,				__UpperCamelCase	,				__UpperCamelCase	):
			__lowercase			:  str								=       {
			    '''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			:  str								=       model_name_to_url[model_name]
			__lowercase			:  List[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	)

			__lowercase			:  Optional[int]								=       get_upernet_config(__UpperCamelCase	)
			__lowercase			:  Any								=       UperNetForSemanticSegmentation(__UpperCamelCase	)
			model.eval()

			# replace "bn" => "batch_norm"
			for key in state_dict.copy().keys():
						__lowercase			:  Dict								=       state_dict.pop(__UpperCamelCase	)
						if "bn" in key:
									__lowercase			:  Optional[Any]								=       key.replace('''bn'''	,				'''batch_norm'''	)
						__lowercase			:  Union[str, Any]								=       val

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

			model.load_state_dict(__UpperCamelCase	)

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

			__lowercase			:  str								=       SegformerImageProcessor()
			__lowercase			:  List[str]								=       processor(__UpperCamelCase	,				return_tensors='''pt'''	).pixel_values

			with torch.no_grad():
						__lowercase			:  Optional[int]								=       model(__UpperCamelCase	)
						__lowercase			:  Optional[int]								=       outputs.logits

			print(logits.shape	)
			print('''First values of logits:'''	,				logits[0, 0, :3, :3]	)
			# assert values
			if model_name == "upernet-swin-tiny":
						__lowercase			:  Any								=       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			:  List[Any]								=       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			:  int								=       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			:  str								=       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]	,				__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 argparse
import os
import re

import packaging.version


__UpperCamelCase     :						Union[str, Any]						       =  '''examples/'''
__UpperCamelCase     :						str						       =  {
    '''examples''': (re.compile(R'''^check_min_version\("[^"]+"\)\s*$''', re.MULTILINE), '''check_min_version("VERSION")\n'''),
    '''init''': (re.compile(R'''^__version__\s+=\s+"([^"]+)"\s*$''', re.MULTILINE), '''__version__ = "VERSION"\n'''),
    '''setup''': (re.compile(R'''^(\s*)version\s*=\s*"[^"]+",''', re.MULTILINE), R'''\1version="VERSION",'''),
    '''doc''': (re.compile(R'''^(\s*)release\s*=\s*"[^"]+"$''', re.MULTILINE), '''release = "VERSION"\n'''),
}
__UpperCamelCase     :						List[str]						       =  {
    '''init''': '''src/transformers/__init__.py''',
    '''setup''': '''setup.py''',
}
__UpperCamelCase     :						Optional[int]						       =  '''README.md'''

def 		_SCREAMING_SNAKE_CASE			(_UpperCAmelCase					:    Any  ,     _UpperCAmelCase					:    Dict  ,     _UpperCAmelCase					:    Any			):
 with open(_UpperCAmelCase  ,     'r'  ,     encoding='utf-8'  ,     newline='\n'			) as f:
  lowerCAmelCase    				= f.read()
 lowerCAmelCase				,lowerCAmelCase    				= REPLACE_PATTERNS[pattern]
 lowerCAmelCase    				= replace.replace('VERSION'  ,     _UpperCAmelCase			)
 lowerCAmelCase    				= re_pattern.sub(_UpperCAmelCase  ,     _UpperCAmelCase			)
 with open(_UpperCAmelCase  ,     'w'  ,     encoding='utf-8'  ,     newline='\n'			) as f:
  f.write(_UpperCAmelCase			)

def 		_SCREAMING_SNAKE_CASE			(_UpperCAmelCase					:    Union[str, Any]			):
 for folder, directories, fnames in os.walk(_UpperCAmelCase			):
  # Removing some of the folders with non-actively maintained examples from the walk
  if "research_projects" in directories:
   directories.remove('research_projects'			)
  if "legacy" in directories:
   directories.remove('legacy'			)
  for fname in fnames:
   if fname.endswith('.py'			):
    update_version_in_file(os.path.join(_UpperCAmelCase  ,     _UpperCAmelCase			)  ,     _UpperCAmelCase  ,     pattern='examples'			)

def 		_SCREAMING_SNAKE_CASE			(_UpperCAmelCase					:    List[str]  ,     _UpperCAmelCase					:    Dict=False			):
 for pattern, fname in REPLACE_FILES.items():
  update_version_in_file(_UpperCAmelCase  ,     _UpperCAmelCase  ,     _UpperCAmelCase			)
 if not patch:
  update_version_in_examples(_UpperCAmelCase			)

def 		_SCREAMING_SNAKE_CASE			():
 lowerCAmelCase    				= '🤗 Transformers currently provides the following architectures'
 lowerCAmelCase    				= '1. Want to contribute a new model?'
 with open(_UpperCAmelCase  ,     'r'  ,     encoding='utf-8'  ,     newline='\n'			) as f:
  lowerCAmelCase    				= f.readlines()

 # Find the start of the list.
 lowerCAmelCase    				= 0
 while not lines[start_index].startswith(_start_prompt			):
  start_index += 1
 start_index += 1

 lowerCAmelCase    				= start_index
 # Update the lines in the model list.
 while not lines[index].startswith(_end_prompt			):
  if lines[index].startswith('1.'			):
   lowerCAmelCase    				= lines[index].replace(
       'https://huggingface.co/docs/transformers/main/model_doc'  ,     'https://huggingface.co/docs/transformers/model_doc'  ,     )
  index += 1

 with open(_UpperCAmelCase  ,     'w'  ,     encoding='utf-8'  ,     newline='\n'			) as f:
  f.writelines(_UpperCAmelCase			)

def 		_SCREAMING_SNAKE_CASE			():
 with open(REPLACE_FILES['init']  ,     'r'			) as f:
  lowerCAmelCase    				= f.read()
 lowerCAmelCase    				= REPLACE_PATTERNS['init'][0].search(_UpperCAmelCase			).groups()[0]
 return packaging.version.parse(_UpperCAmelCase			)

def 		_SCREAMING_SNAKE_CASE			(_UpperCAmelCase					:    Tuple=False			):
 lowerCAmelCase    				= get_version()
 if patch and default_version.is_devrelease:
  raise ValueError('Can\'t create a patch version from the dev branch, checkout a released version!'			)
 if default_version.is_devrelease:
  lowerCAmelCase    				= default_version.base_version
 elif patch:
  lowerCAmelCase    				= F'{default_version.major}.{default_version.minor}.{default_version.micro + 1}'
 else:
  lowerCAmelCase    				= F'{default_version.major}.{default_version.minor + 1}.0'

 # Now let's ask nicely if that's the right one.
 lowerCAmelCase    				= input(F'Which version are you releasing? [{default_version}]'			)
 if len(_UpperCAmelCase			) == 0:
  lowerCAmelCase    				= default_version

 print(F'Updating version to {version}.'			)
 global_version_update(_UpperCAmelCase  ,     patch=_UpperCAmelCase			)
 if not patch:
  print('Cleaning main README, don\'t forget to run `make fix-copies`.'			)
  clean_main_ref_in_model_list()

def 		_SCREAMING_SNAKE_CASE			():
 lowerCAmelCase    				= get_version()
 lowerCAmelCase    				= F'{current_version.major}.{current_version.minor + 1}.0.dev0'
 lowerCAmelCase    				= current_version.base_version

 # Check with the user we got that right.
 lowerCAmelCase    				= input(F'Which version are we developing now? [{dev_version}]'			)
 if len(_UpperCAmelCase			) == 0:
  lowerCAmelCase    				= dev_version

 print(F'Updating version to {version}.'			)
 global_version_update(_UpperCAmelCase			)
 print('Cleaning main README, don\'t forget to run `make fix-copies`.'			)
 clean_main_ref_in_model_list()


if __name__ == "__main__":
       __UpperCamelCase     :						Optional[int]						       =  argparse.ArgumentParser()
       parser.add_argument('''--post_release''', action='''store_true''', help='''Whether this is pre or post release.''')
       parser.add_argument('''--patch''', action='''store_true''', help='''Whether or not this is a patch release.''')
       __UpperCamelCase     :						Optional[int]						       =  parser.parse_args()
       if not args.post_release:
              pre_release_work(patch=args.patch)
       elif args.patch:
              print('''Nothing to do after a patch :-)''')
       else:
              post_release_work()





