


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       (    _lowerCAmelCase )						->	int:


						'''simple docstring'''






						__snake_case				=		384
						__snake_case				=		7
						if "tiny" in model_name:
												__snake_case				=		96
												__snake_case				=		(2, 2, 6, 2)
												__snake_case				=		(3, 6, 12, 24)
						elif "small" in model_name:
												__snake_case				=		96
												__snake_case				=		(2, 2, 18, 2)
												__snake_case				=		(3, 6, 12, 24)
						elif "base" in model_name:
												__snake_case				=		128
												__snake_case				=		(2, 2, 18, 2)
												__snake_case				=		(4, 8, 16, 32)
												__snake_case				=		12
												__snake_case				=		512
						elif "large" in model_name:
												__snake_case				=		192
												__snake_case				=		(2, 2, 18, 2)
												__snake_case				=		(6, 12, 24, 48)
												__snake_case				=		12
												__snake_case				=		768

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

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

						return config





def 	_lowerCAmelCase       (    _lowerCAmelCase )						->	Tuple:


						'''simple docstring'''






						__snake_case				=		[]

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


						'''simple docstring'''






						__snake_case				=		dct.pop(_lowerCAmelCase )
						__snake_case				=		val





def 	_lowerCAmelCase       (    _lowerCAmelCase      ,	_lowerCAmelCase )						->	int:


						'''simple docstring'''






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





def 	_lowerCAmelCase       (    _lowerCAmelCase )						->	Optional[int]:


						'''simple docstring'''






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





def 	_lowerCAmelCase       (    _lowerCAmelCase )						->	str:


						'''simple docstring'''






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

						return x





def 	_lowerCAmelCase       (    _lowerCAmelCase )						->	Optional[int]:


						'''simple docstring'''






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





def 	_lowerCAmelCase       (    _lowerCAmelCase )						->	Union[str, Any]:


						'''simple docstring'''






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





def 	_lowerCAmelCase       (    _lowerCAmelCase      ,	_lowerCAmelCase      ,	_lowerCAmelCase )						->	Union[str, Any]:


						'''simple docstring'''






						__snake_case				=		{
						    "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",
						}
						__snake_case				=		model_name_to_url[model_name]
						__snake_case				=		torch.hub.load_state_dict_from_url(_lowerCAmelCase      ,	map_location="cpu"      ,	file_name=_lowerCAmelCase )[
						    "state_dict"
						]

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

						__snake_case				=		get_upernet_config(_lowerCAmelCase )
						__snake_case				=		UperNetForSemanticSegmentation(_lowerCAmelCase )
						model.eval()

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

						# rename keys
						__snake_case				=		create_rename_keys(_lowerCAmelCase )
						for src, dest in rename_keys:
												rename_key(_lowerCAmelCase      ,	_lowerCAmelCase      ,	_lowerCAmelCase )
						read_in_q_k_v(_lowerCAmelCase      ,	config.backbone_config )

						# fix downsample parameters
						for key, value in state_dict.items():
												if "downsample" in key:
																		if "reduction" in key:
																								__snake_case				=		reverse_correct_unfold_reduction_order(_lowerCAmelCase )
																		if "norm" in key:
																								__snake_case				=		reverse_correct_unfold_norm_order(_lowerCAmelCase )

						model.load_state_dict(_lowerCAmelCase )

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

						__snake_case				=		SegformerImageProcessor()
						__snake_case				=		processor(_lowerCAmelCase      ,	return_tensors="pt" ).pixel_values

						with torch.no_grad():
												__snake_case				=		model(_lowerCAmelCase )
												__snake_case				=		outputs.logits

						print(logits.shape )
						print("First values of logits:"      ,	logits[0, 0, :3, :3] )
						# assert values
						if model_name == "upernet-swin-tiny":
												__snake_case				=		torch.tensor(
												    [[-7.59_58, -7.59_58, -7.43_02], [-7.59_58, -7.59_58, -7.43_02], [-7.47_97, -7.47_97, -7.30_68]] )
						elif model_name == "upernet-swin-small":
												__snake_case				=		torch.tensor(
												    [[-7.19_21, -7.19_21, -6.95_32], [-7.19_21, -7.19_21, -6.95_32], [-7.09_08, -7.09_08, -6.85_34]] )
						elif model_name == "upernet-swin-base":
												__snake_case				=		torch.tensor(
												    [[-6.58_51, -6.58_51, -6.43_30], [-6.58_51, -6.58_51, -6.43_30], [-6.47_63, -6.47_63, -6.32_54]] )
						elif model_name == "upernet-swin-large":
												__snake_case				=		torch.tensor(
												    [[-7.52_97, -7.52_97, -7.38_02], [-7.52_97, -7.52_97, -7.38_02], [-7.40_44, -7.40_44, -7.25_86]] )
						print("Logits:"      ,	outputs.logits[0, 0, :3, :3] )
						assert torch.allclose(outputs.logits[0, 0, :3, :3]      ,	_lowerCAmelCase      ,	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(_lowerCAmelCase )
												print(F'''Saving processor to {pytorch_dump_folder_path}''' )
												processor.save_pretrained(_lowerCAmelCase )

						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							:					Optional[Any]     						=     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							:					List[Any]     						=     parser.parse_args()
	convert_upernet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)









import os
from datetime import datetime as dt

from github import Github


A							:					Union[str, Any]     						=     [
    'good first issue',
    'good second issue',
    'good difficult issue',
    'enhancement',
    'new pipeline/model',
    'new scheduler',
    'wip',
]





def 	_lowerCAmelCase       (    )						->	str:


						'''simple docstring'''






						__snake_case				=		Github(os.environ["GITHUB_TOKEN"] )
						__snake_case				=		g.get_repo("huggingface/diffusers" )
						__snake_case				=		repo.get_issues(state="open" )

						for issue in open_issues:
												__snake_case				=		sorted(issue.get_comments()      ,	key=lambda _lowerCAmelCase : i.created_at      ,	reverse=_lowerCAmelCase )
												__snake_case				=		comments[0] if len(_lowerCAmelCase ) > 0 else None
												if (
												    last_comment is not None
												    and last_comment.user.login == "github-actions[bot]"
												    and (dt.utcnow() - issue.updated_at).days > 7
												    and (dt.utcnow() - issue.created_at).days >= 30
												    and not any(label.name.lower() in LABELS_TO_EXEMPT for label in issue.get_labels() )
												):
																		# Closes the issue after 7 days of inactivity since the Stalebot notification.
																		issue.edit(state="closed" )
												elif (
												    "stale" in issue.get_labels()
												    and last_comment is not None
												    and last_comment.user.login != "github-actions[bot]"
												):
																		# Opens the issue if someone other than Stalebot commented.
																		issue.edit(state="open" )
																		issue.remove_from_labels("stale" )
												elif (
												    (dt.utcnow() - issue.updated_at).days > 23
												    and (dt.utcnow() - issue.created_at).days >= 30
												    and not any(label.name.lower() in LABELS_TO_EXEMPT for label in issue.get_labels() )
												):
																		# Post a Stalebot notification after 23 days of inactivity.
																		issue.create_comment(
																		    "This issue has been automatically marked as stale because it has not had "
																		    "recent activity. If you think this still needs to be addressed "
																		    "please comment on this thread.\n\nPlease note that issues that do not follow the "
																		    "[contributing guidelines](https://github.com/huggingface/diffusers/blob/main/CONTRIBUTING.md) "
																		    "are likely to be ignored." )
																		issue.add_to_labels("stale" )


if __name__ == "__main__":
	main()





