




from ...configuration_utils import PretrainedConfig
from ...utils import logging


_A     =			logging.get_logger(__name__)

_A     =			{
    '''google/pegasus-large''': '''https://huggingface.co/google/pegasus-large/resolve/main/config.json''',
    # See all PEGASUS models at https://huggingface.co/models?filter=pegasus
}


class  A		(  __UpperCAmelCase     ):
		__snake_case									=   'pegasus'
		__snake_case									=   ['past_key_values']
		__snake_case									=   {'num_attention_heads': 'encoder_attention_heads', 'hidden_size': 'd_model'}



		def __init__(     self,	UpperCamelCase__=5_0265,	UpperCamelCase__=1024,	UpperCamelCase__=12,	UpperCamelCase__=4096,	UpperCamelCase__=16,	UpperCamelCase__=12,	UpperCamelCase__=4096,	UpperCamelCase__=16,	UpperCamelCase__=0.0,	UpperCamelCase__=0.0,	UpperCamelCase__=True,	UpperCamelCase__=True,	UpperCamelCase__="gelu",	UpperCamelCase__=1024,	UpperCamelCase__=0.1,	UpperCamelCase__=0.0,	UpperCamelCase__=0.0,	UpperCamelCase__=0.02,	UpperCamelCase__=0,	UpperCamelCase__=False,	UpperCamelCase__=0,	UpperCamelCase__=1,	UpperCamelCase__=1,	**UpperCamelCase__,	):







								"""simple docstring"""



								lowerCAmelCase_      			=  vocab_size
								lowerCAmelCase_      			=  max_position_embeddings
								lowerCAmelCase_      			=  d_model
								lowerCAmelCase_      			=  encoder_ffn_dim
								lowerCAmelCase_      			=  encoder_layers
								lowerCAmelCase_      			=  encoder_attention_heads
								lowerCAmelCase_      			=  decoder_ffn_dim
								lowerCAmelCase_      			=  decoder_layers
								lowerCAmelCase_      			=  decoder_attention_heads
								lowerCAmelCase_      			=  dropout
								lowerCAmelCase_      			=  attention_dropout
								lowerCAmelCase_      			=  activation_dropout
								lowerCAmelCase_      			=  activation_function
								lowerCAmelCase_      			=  init_std
								lowerCAmelCase_      			=  encoder_layerdrop
								lowerCAmelCase_      			=  decoder_layerdrop
								lowerCAmelCase_      			=  use_cache
								lowerCAmelCase_      			=  encoder_layers
								lowerCAmelCase_      			=  scale_embedding  # scale factor will be sqrt(d_model) if True
								super().__init__(
								    pad_token_id=UpperCamelCase__,	eos_token_id=UpperCamelCase__,	is_encoder_decoder=UpperCamelCase__,	decoder_start_token_id=UpperCamelCase__,	forced_eos_token_id=UpperCamelCase__,	**UpperCamelCase__,	)



		@property
		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								return self.encoder_attention_heads







		@property
		def        SCREAMING_SNAKE_CASE__			(     self   ):







								"""simple docstring"""



								return self.d_model










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			(       _A  ):
						lowerCAmelCase_      			=  384
						lowerCAmelCase_      			=  7
						if "tiny" in model_name:
												lowerCAmelCase_      			=  96
												lowerCAmelCase_      			=  (2, 2, 6, 2)
												lowerCAmelCase_      			=  (3, 6, 12, 24)
						elif "small" in model_name:
												lowerCAmelCase_      			=  96
												lowerCAmelCase_      			=  (2, 2, 18, 2)
												lowerCAmelCase_      			=  (3, 6, 12, 24)
						elif "base" in model_name:
												lowerCAmelCase_      			=  128
												lowerCAmelCase_      			=  (2, 2, 18, 2)
												lowerCAmelCase_      			=  (4, 8, 16, 32)
												lowerCAmelCase_      			=  12
												lowerCAmelCase_      			=  512
						elif "large" in model_name:
												lowerCAmelCase_      			=  192
												lowerCAmelCase_      			=  (2, 2, 18, 2)
												lowerCAmelCase_      			=  (6, 12, 24, 48)
												lowerCAmelCase_      			=  12
												lowerCAmelCase_      			=  768

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

						lowerCAmelCase_      			=  SwinConfig(
						    embed_dim=_A ,  depths=_A ,  num_heads=_A ,  window_size=_A ,  out_features=['''stage1''', '''stage2''', '''stage3''', '''stage4'''] ,  )
						lowerCAmelCase_      			=  UperNetConfig(
						    backbone_config=_A ,  auxiliary_in_channels=_A ,  num_labels=_A ,  idalabel=_A ,  labelaid=_A ,  )

						return config





def __UpperCamelCase			(       _A  ):
						lowerCAmelCase_      			=  []

						# 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			(       _A ,  _A ,  _A  ):
						lowerCAmelCase_      			=  dct.pop(_A  )
						lowerCAmelCase_      			=  val





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





def __UpperCamelCase			(       _A  ):
						lowerCAmelCase_			,					lowerCAmelCase_      			=  x.shape
						lowerCAmelCase_      			=  x.reshape(_A ,  4 ,  in_channel // 4  )
						lowerCAmelCase_      			=  x[:, [0, 2, 1, 3], :].transpose(1 ,  2  ).reshape(_A ,  _A  )
						return x





def __UpperCamelCase			(       _A  ):
						lowerCAmelCase_			,					lowerCAmelCase_      			=  x.shape
						lowerCAmelCase_      			=  x.reshape(_A ,  in_channel // 4 ,  4  )
						lowerCAmelCase_      			=  x[:, :, [0, 2, 1, 3]].transpose(1 ,  2  ).reshape(_A ,  _A  )

						return x





def __UpperCamelCase			(       _A  ):
						lowerCAmelCase_      			=  x.shape[0]
						lowerCAmelCase_      			=  x.reshape(4 ,  in_channel // 4  )
						lowerCAmelCase_      			=  x[[0, 2, 1, 3], :].transpose(0 ,  1  ).reshape(_A  )
						return x





def __UpperCamelCase			(       _A  ):
						lowerCAmelCase_      			=  x.shape[0]
						lowerCAmelCase_      			=  x.reshape(in_channel // 4 ,  4  )
						lowerCAmelCase_      			=  x[:, [0, 2, 1, 3]].transpose(0 ,  1  ).reshape(_A  )
						return x





def __UpperCamelCase			(       _A ,  _A ,  _A  ):
						lowerCAmelCase_      			=  {
						    '''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''',
						}
						lowerCAmelCase_      			=  model_name_to_url[model_name]
						lowerCAmelCase_      			=  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  )

						lowerCAmelCase_      			=  get_upernet_config(_A  )
						lowerCAmelCase_      			=  UperNetForSemanticSegmentation(_A  )
						model.eval()

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

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

						model.load_state_dict(_A  )

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

						lowerCAmelCase_      			=  SegformerImageProcessor()
						lowerCAmelCase_      			=  processor(_A ,  return_tensors='''pt'''  ).pixel_values

						with torch.no_grad():
												lowerCAmelCase_      			=  model(_A  )
												lowerCAmelCase_      			=  outputs.logits

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




