



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					(   lowerCAmelCase_       : Optional[int]    )				->			List[str]:






						'''simple docstring'''



						UpperCAmelCase_=       3_84
						UpperCAmelCase_=       7
						if "tiny" in model_name:
												UpperCAmelCase_=       96
												UpperCAmelCase_=       (2, 2, 6, 2)
												UpperCAmelCase_=       (3, 6, 12, 24)
						elif "small" in model_name:
												UpperCAmelCase_=       96
												UpperCAmelCase_=       (2, 2, 18, 2)
												UpperCAmelCase_=       (3, 6, 12, 24)
						elif "base" in model_name:
												UpperCAmelCase_=       1_28
												UpperCAmelCase_=       (2, 2, 18, 2)
												UpperCAmelCase_=       (4, 8, 16, 32)
												UpperCAmelCase_=       12
												UpperCAmelCase_=       5_12
						elif "large" in model_name:
												UpperCAmelCase_=       1_92
												UpperCAmelCase_=       (2, 2, 18, 2)
												UpperCAmelCase_=       (6, 12, 24, 48)
												UpperCAmelCase_=       12
												UpperCAmelCase_=       7_68

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

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

						return config



def 	__a					(   lowerCAmelCase_       : Optional[Any]    )				->			List[Any]:






						'''simple docstring'''



						UpperCAmelCase_=       []

						# 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					(   lowerCAmelCase_       : List[str]      ,lowerCAmelCase_       : List[str]      ,lowerCAmelCase_       : str    )				->			Tuple:






						'''simple docstring'''



						UpperCAmelCase_=       dct.pop(_UpperCAmelCase    )
						UpperCAmelCase_=       val



def 	__a					(   lowerCAmelCase_       : Any      ,lowerCAmelCase_       : Optional[int]    )				->			str:






						'''simple docstring'''



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



def 	__a					(   lowerCAmelCase_       : Any    )				->			Optional[Any]:






						'''simple docstring'''



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



def 	__a					(   lowerCAmelCase_       : str    )				->			List[Any]:






						'''simple docstring'''



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

						return x



def 	__a					(   lowerCAmelCase_       : List[str]    )				->			List[str]:






						'''simple docstring'''



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



def 	__a					(   lowerCAmelCase_       : Optional[int]    )				->			Any:






						'''simple docstring'''



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



def 	__a					(   lowerCAmelCase_       : Union[str, Any]      ,lowerCAmelCase_       : Optional[Any]      ,lowerCAmelCase_       : Dict    )				->			str:






						'''simple docstring'''



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

						UpperCAmelCase_=       get_upernet_config(_UpperCAmelCase    )
						UpperCAmelCase_=       UperNetForSemanticSegmentation(_UpperCAmelCase    )
						model.eval()

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

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

						model.load_state_dict(_UpperCAmelCase    )

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

						UpperCAmelCase_=       SegformerImageProcessor()
						UpperCAmelCase_=       processor(_UpperCAmelCase      ,return_tensors="""pt"""    ).pixel_values

						with torch.no_grad():
												UpperCAmelCase_=       model(_UpperCAmelCase    )
												UpperCAmelCase_=       outputs.logits

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











import pytest

import datasets


# Import fixture modules as plugins
__A       =						['''tests.fixtures.files''', '''tests.fixtures.hub''', '''tests.fixtures.fsspec''']



def 	__a					(   lowerCAmelCase_       : Optional[Any]      ,lowerCAmelCase_       : Any    )				->			Tuple:






						'''simple docstring'''



						for item in items:
												if any(marker in item.keywords for marker in ["""integration""", """unit"""]    ):
																		continue
												item.add_marker(pytest.mark.unit    )



def 	__a					(   lowerCAmelCase_       : Tuple    )				->			Optional[Any]:






						'''simple docstring'''



						config.addinivalue_line("""markers"""      ,"""torchaudio_latest: mark test to run with torchaudio>=0.12"""    )



@pytest.fixture(autouse=lowerCAmelCase_    )
def 	__a					(   lowerCAmelCase_       : List[str]      ,lowerCAmelCase_       : List[str]    )				->			int:






						'''simple docstring'''



						UpperCAmelCase_=       tmp_path_factory.getbasetemp() / """cache"""
						UpperCAmelCase_=       test_hf_cache_home / """datasets"""
						UpperCAmelCase_=       test_hf_cache_home / """metrics"""
						UpperCAmelCase_=       test_hf_cache_home / """modules"""
						monkeypatch.setattr("""datasets.config.HF_DATASETS_CACHE"""      ,str(lowerCAmelCase_    )    )
						monkeypatch.setattr("""datasets.config.HF_METRICS_CACHE"""      ,str(lowerCAmelCase_    )    )
						monkeypatch.setattr("""datasets.config.HF_MODULES_CACHE"""      ,str(lowerCAmelCase_    )    )
						UpperCAmelCase_=       test_hf_datasets_cache / """downloads"""
						monkeypatch.setattr("""datasets.config.DOWNLOADED_DATASETS_PATH"""      ,str(lowerCAmelCase_    )    )
						UpperCAmelCase_=       test_hf_datasets_cache / """downloads""" / """extracted"""
						monkeypatch.setattr("""datasets.config.EXTRACTED_DATASETS_PATH"""      ,str(lowerCAmelCase_    )    )



@pytest.fixture(autouse=lowerCAmelCase_      ,scope="""session"""    )
def 	__a					(   )				->			Optional[int]:






						'''simple docstring'''



						datasets.disable_progress_bar()



@pytest.fixture(autouse=lowerCAmelCase_    )
def 	__a					(   lowerCAmelCase_       : int    )				->			str:






						'''simple docstring'''



						monkeypatch.setattr("""datasets.config.HF_UPDATE_DOWNLOAD_COUNTS"""      ,lowerCAmelCase_    )



@pytest.fixture
def 	__a					(   lowerCAmelCase_       : List[str]    )				->			Union[str, Any]:






						'''simple docstring'''



						monkeypatch.setattr("""sqlalchemy.util.deprecations.SILENCE_UBER_WARNING"""      ,lowerCAmelCase_    )




