


'''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_ ):
							A_			    =		3_8_4
							A_			    =		7
							if "tiny" in model_name:
														A_			    =		9_6
														A_			    =		(2, 2, 6, 2)
														A_			    =		(3, 6, 1_2, 2_4)
							elif "small" in model_name:
														A_			    =		9_6
														A_			    =		(2, 2, 1_8, 2)
														A_			    =		(3, 6, 1_2, 2_4)
							elif "base" in model_name:
														A_			    =		1_2_8
														A_			    =		(2, 2, 1_8, 2)
														A_			    =		(4, 8, 1_6, 3_2)
														A_			    =		1_2
														A_			    =		5_1_2
							elif "large" in model_name:
														A_			    =		1_9_2
														A_			    =		(2, 2, 1_8, 2)
														A_			    =		(6, 1_2, 2_4, 4_8)
														A_			    =		1_2
														A_			    =		7_6_8

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

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

							return config







def       a_				(	UpperCamelCase_ ):
							A_			    =		[]

							# 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_						,  UpperCamelCase_						,  UpperCamelCase_ ):
							A_			    =		dct.pop(UpperCamelCase_ )
							A_			    =		val







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







def       a_				(	UpperCamelCase_ ):
							A_    ,					A_			    =		x.shape
							A_			    =		x.reshape(UpperCamelCase_						,  4						,  in_channel // 4 )
							A_			    =		x[:, [0, 2, 1, 3], :].transpose(1						,  2 ).reshape(UpperCamelCase_						,  UpperCamelCase_ )
							return x







def       a_				(	UpperCamelCase_ ):
							A_    ,					A_			    =		x.shape
							A_			    =		x.reshape(UpperCamelCase_						,  in_channel // 4						,  4 )
							A_			    =		x[:, :, [0, 2, 1, 3]].transpose(1						,  2 ).reshape(UpperCamelCase_						,  UpperCamelCase_ )

							return x







def       a_				(	UpperCamelCase_ ):
							A_			    =		x.shape[0]
							A_			    =		x.reshape(4						,  in_channel // 4 )
							A_			    =		x[[0, 2, 1, 3], :].transpose(0						,  1 ).reshape(UpperCamelCase_ )
							return x







def       a_				(	UpperCamelCase_ ):
							A_			    =		x.shape[0]
							A_			    =		x.reshape(in_channel // 4						,  4 )
							A_			    =		x[:, [0, 2, 1, 3]].transpose(0						,  1 ).reshape(UpperCamelCase_ )
							return x







def       a_				(	UpperCamelCase_						,  UpperCamelCase_						,  UpperCamelCase_ ):
							A_			    =		{
							    "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",
							}
							A_			    =		model_name_to_url[model_name]
							A_			    =		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 )

							A_			    =		get_upernet_config(UpperCamelCase_ )
							A_			    =		UperNetForSemanticSegmentation(UpperCamelCase_ )
							model.eval()

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

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

							model.load_state_dict(UpperCamelCase_ )

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

							A_			    =		SegformerImageProcessor()
							A_			    =		processor(UpperCamelCase_						,  return_tensors="pt" ).pixel_values

							with torch.no_grad():
														A_			    =		model(UpperCamelCase_ )
														A_			    =		outputs.logits

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

							__SCREAMING_SNAKE_CASE					:				List[str]									=		parser.parse_args()
							convert_upernet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)









'''simple docstring'''
import unittest

from transformers import AutoConfig, AutoTokenizer, BertConfig, TensorType, is_flax_available
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, require_flax, slow


if is_flax_available():
							import jax

							from transformers.models.auto.modeling_flax_auto import FlaxAutoModel
							from transformers.models.bert.modeling_flax_bert import FlaxBertModel
							from transformers.models.roberta.modeling_flax_roberta import FlaxRobertaModel






@require_flax
class 					__lowerCAmelCase   (   unittest.TestCase						):

						"""simple docstring"""






						@slow
						def 				_UpperCAmelCase       (  self							:      Dict	):
													for model_name in ["bert-base-cased", "bert-large-uncased"]:
																				with self.subTest(lowerCAmelCase	):
																											A_			    =		AutoConfig.from_pretrained(lowerCAmelCase	)
																											self.assertIsNotNone(lowerCAmelCase	)
																											self.assertIsInstance(lowerCAmelCase ,  lowerCAmelCase	)

																											A_			    =		FlaxAutoModel.from_pretrained(lowerCAmelCase	)
																											self.assertIsNotNone(lowerCAmelCase	)
																											self.assertIsInstance(lowerCAmelCase ,  lowerCAmelCase	)






						@slow
						def 				_UpperCAmelCase       (  self							:      Union[str, Any]	):
													for model_name in ["roberta-base", "roberta-large"]:
																				with self.subTest(lowerCAmelCase	):
																											A_			    =		AutoConfig.from_pretrained(lowerCAmelCase	)
																											self.assertIsNotNone(lowerCAmelCase	)
																											self.assertIsInstance(lowerCAmelCase ,  lowerCAmelCase	)

																											A_			    =		FlaxAutoModel.from_pretrained(lowerCAmelCase	)
																											self.assertIsNotNone(lowerCAmelCase	)
																											self.assertIsInstance(lowerCAmelCase ,  lowerCAmelCase	)






						@slow
						def 				_UpperCAmelCase       (  self							:      List[Any]	):
													for model_name in ["bert-base-cased", "bert-large-uncased"]:
																				A_			    =		AutoTokenizer.from_pretrained(lowerCAmelCase	)
																				A_			    =		FlaxBertModel.from_pretrained(lowerCAmelCase	)
																				A_			    =		tokenizer("Do you support jax jitted function?" ,  return_tensors=TensorType.JAX	)

																				@jax.jit
																				def eval(**lowerCAmelCase							:      Tuple	):
																											return model(**lowerCAmelCase	)

																				eval(**lowerCAmelCase	).block_until_ready()






						@slow
						def 				_UpperCAmelCase       (  self							:      str	):
													for model_name in ["roberta-base", "roberta-large"]:
																				A_			    =		AutoTokenizer.from_pretrained(lowerCAmelCase	)
																				A_			    =		FlaxRobertaModel.from_pretrained(lowerCAmelCase	)
																				A_			    =		tokenizer("Do you support jax jitted function?" ,  return_tensors=TensorType.JAX	)

																				@jax.jit
																				def eval(**lowerCAmelCase							:      Any	):
																											return model(**lowerCAmelCase	)

																				eval(**lowerCAmelCase	).block_until_ready()






						def 				_UpperCAmelCase       (  self							:      List[Any]	):
													with self.assertRaisesRegex(
													    lowerCAmelCase ,  "bert-base is not a local folder and is not a valid model identifier"	):
																				A_			    =		FlaxAutoModel.from_pretrained("bert-base"	)






						def 				_UpperCAmelCase       (  self							:      Any	):
													with self.assertRaisesRegex(
													    lowerCAmelCase ,  r"aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)"	):
																				A_			    =		FlaxAutoModel.from_pretrained(lowerCAmelCase ,  revision="aaaaaa"	)






						def 				_UpperCAmelCase       (  self							:      Dict	):
													with self.assertRaisesRegex(
													    lowerCAmelCase ,  "hf-internal-testing/config-no-model does not appear to have a file named flax_model.msgpack" ,  ):
																				A_			    =		FlaxAutoModel.from_pretrained("hf-internal-testing/config-no-model"	)






						def 				_UpperCAmelCase       (  self							:      Any	):
													with self.assertRaisesRegex(lowerCAmelCase ,  "Use `from_pt=True` to load this model"	):
																				A_			    =		FlaxAutoModel.from_pretrained("hf-internal-testing/tiny-bert-pt-only"	)





