

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

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

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

						return config






def 							_lowerCAmelCase (      __magic_name__	:    Tuple	)							->     Optional[Any]:
						lowercase   :       Optional[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 							_lowerCAmelCase (      __magic_name__	:    List[Any]      ,   __magic_name__	:    List[Any]      ,   __magic_name__	:    Union[str, Any]	)							->     int:
						lowercase   :       Union[str, Any]								=dct.pop(__magic_name__	)
						lowercase   :       Tuple								=val






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






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






def 							_lowerCAmelCase (      __magic_name__	:    List[str]	)							->     int:
						lowercase		,  lowercase   :       List[str]								=x.shape
						lowercase   :       Tuple								=x.reshape(__magic_name__      ,   in_channel // 4      ,   4	)
						lowercase   :       str								=x[:, :, [0, 2, 1, 3]].transpose(1      ,   2	).reshape(__magic_name__      ,   __magic_name__	)

						return x






def 							_lowerCAmelCase (      __magic_name__	:    Tuple	)							->     str:
						lowercase   :       Dict								=x.shape[0]
						lowercase   :       Any								=x.reshape(4      ,   in_channel // 4	)
						lowercase   :       Union[str, Any]								=x[[0, 2, 1, 3], :].transpose(0      ,   1	).reshape(__magic_name__	)
						return x






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






def 							_lowerCAmelCase (      __magic_name__	:    Union[str, Any]      ,   __magic_name__	:    Any      ,   __magic_name__	:    int	)							->     Union[str, Any]:
						lowercase   :       Union[str, 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''',
						}
						lowercase   :       int								=model_name_to_url[model_name]
						lowercase   :       Any								=torch.hub.load_state_dict_from_url(__magic_name__      ,   map_location='''cpu'''      ,   file_name=__magic_name__	)[
						    '''state_dict'''
						]

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

						lowercase   :       str								=get_upernet_config(__magic_name__	)
						lowercase   :       Optional[Any]								=UperNetForSemanticSegmentation(__magic_name__	)
						model.eval()

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

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

						model.load_state_dict(__magic_name__	)

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

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

						with torch.no_grad():
												lowercase   :       Tuple								=model(__magic_name__	)
												lowercase   :       Union[str, Any]								=outputs.logits

						print(logits.shape	)
						print('''First values of logits:'''      ,   logits[0, 0, :3, :3]	)
						# assert values
						if model_name == "upernet-swin-tiny":
												lowercase   :       Union[str, Any]								=torch.tensor(
												    [[-7.5_9_5_8, -7.5_9_5_8, -7.4_3_0_2], [-7.5_9_5_8, -7.5_9_5_8, -7.4_3_0_2], [-7.4_7_9_7, -7.4_7_9_7, -7.3_0_6_8]]	)
						elif model_name == "upernet-swin-small":
												lowercase   :       str								=torch.tensor(
												    [[-7.1_9_2_1, -7.1_9_2_1, -6.9_5_3_2], [-7.1_9_2_1, -7.1_9_2_1, -6.9_5_3_2], [-7.0_9_0_8, -7.0_9_0_8, -6.8_5_3_4]]	)
						elif model_name == "upernet-swin-base":
												lowercase   :       List[str]								=torch.tensor(
												    [[-6.5_8_5_1, -6.5_8_5_1, -6.4_3_3_0], [-6.5_8_5_1, -6.5_8_5_1, -6.4_3_3_0], [-6.4_7_6_3, -6.4_7_6_3, -6.3_2_5_4]]	)
						elif model_name == "upernet-swin-large":
												lowercase   :       Optional[int]								=torch.tensor(
												    [[-7.5_2_9_7, -7.5_2_9_7, -7.3_8_0_2], [-7.5_2_9_7, -7.5_2_9_7, -7.3_8_0_2], [-7.4_0_4_4, -7.4_0_4_4, -7.2_5_8_6]]	)
						print('''Logits:'''      ,   outputs.logits[0, 0, :3, :3]	)
						assert torch.allclose(outputs.logits[0, 0, :3, :3]      ,   __magic_name__      ,   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(__magic_name__	)
												print(f'''Saving processor to {pytorch_dump_folder_path}'''	)
												processor.save_pretrained(__magic_name__	)

						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__":
							UpperCamelCase_							=					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."""
							)

							UpperCamelCase_							=					parser.parse_args()
							convert_upernet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)








'''simple docstring'''





from transformers import HfArgumentParser, TensorFlowBenchmark, TensorFlowBenchmarkArguments






def 							_lowerCAmelCase (      )							->     List[Any]:
						lowercase   :       Tuple								=HfArgumentParser(__magic_name__	)
						lowercase   :       Union[str, Any]								=parser.parse_args_into_dataclasses()[0]
						lowercase   :       Any								=TensorFlowBenchmark(args=__magic_name__	)
						try:
												lowercase   :       List[Any]								=parser.parse_args_into_dataclasses()[0]
						except ValueError as e:
												lowercase   :       List[Any]								='''Arg --no_{0} is no longer used, please use --no-{0} instead.'''
												lowercase   :       Any								=''' '''.join(str(__magic_name__	).split(''' '''	)[:-1]	)
												lowercase   :       Optional[Any]								=''''''
												lowercase   :       List[str]								=eval(str(__magic_name__	).split(''' '''	)[-1]	)
												lowercase   :       Optional[Any]								=[]
												for arg in depreciated_args:
																		# arg[2:] removes '--'
																		if arg[2:] in TensorFlowBenchmark.deprecated_args:
																								# arg[5:] removes '--no_'
																								full_error_msg += arg_error_msg.format(arg[5:]	)
																		else:
																								wrong_args.append(__magic_name__	)
												if len(__magic_name__	) > 0:
																		lowercase   :       int								=full_error_msg + begin_error_msg + str(__magic_name__	)
												raise ValueError(__magic_name__	)
						benchmark.run()


if __name__ == "__main__":
							main()





