





import argparse
import json

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import ConvNextConfig, SegformerImageProcessor, UperNetConfig, UperNetForSemanticSegmentation


def       UpperCamelCase__    (		lowerCAmelCase__ ):
				lowercase							=     384
				if "tiny" in model_name:
								lowercase							=     [3, 3, 9, 3]
								lowercase							=     [96, 192, 384, 768]
				if "small" in model_name:
								lowercase							=     [3, 3, 27, 3]
								lowercase							=     [96, 192, 384, 768]
				if "base" in model_name:
								lowercase							=     [3, 3, 27, 3]
								lowercase							=     [128, 256, 512, 1_024]
								lowercase							=     512
				if "large" in model_name:
								lowercase							=     [3, 3, 27, 3]
								lowercase							=     [192, 384, 768, 1_536]
								lowercase							=     768
				if "xlarge" in model_name:
								lowercase							=     [3, 3, 27, 3]
								lowercase							=     [256, 512, 1_024, 2_048]
								lowercase							=     1_024

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

				lowercase							=     ConvNextConfig(
				    depths=__lowerCamelCase       ,hidden_sizes=__lowerCamelCase       ,out_features=["""stage1""", """stage2""", """stage3""", """stage4"""] )
				lowercase							=     UperNetConfig(
				    backbone_config=__lowerCamelCase       ,auxiliary_in_channels=__lowerCamelCase       ,num_labels=__lowerCamelCase       ,idalabel=__lowerCamelCase       ,labelaid=__lowerCamelCase       ,)

				return config


def       UpperCamelCase__    (		lowerCAmelCase__ ):
				lowercase							=     []

				# fmt: off
				# stem
				rename_keys.append(("""backbone.downsample_layers.0.0.weight""", """backbone.embeddings.patch_embeddings.weight""") )
				rename_keys.append(("""backbone.downsample_layers.0.0.bias""", """backbone.embeddings.patch_embeddings.bias""") )
				rename_keys.append(("""backbone.downsample_layers.0.1.weight""", """backbone.embeddings.layernorm.weight""") )
				rename_keys.append(("""backbone.downsample_layers.0.1.bias""", """backbone.embeddings.layernorm.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}.{j}.gamma""", f"""backbone.encoder.stages.{i}.layers.{j}.layer_scale_parameter""") )
												rename_keys.append((f"""backbone.stages.{i}.{j}.depthwise_conv.weight""", f"""backbone.encoder.stages.{i}.layers.{j}.dwconv.weight""") )
												rename_keys.append((f"""backbone.stages.{i}.{j}.depthwise_conv.bias""", f"""backbone.encoder.stages.{i}.layers.{j}.dwconv.bias""") )
												rename_keys.append((f"""backbone.stages.{i}.{j}.norm.weight""", f"""backbone.encoder.stages.{i}.layers.{j}.layernorm.weight""") )
												rename_keys.append((f"""backbone.stages.{i}.{j}.norm.bias""", f"""backbone.encoder.stages.{i}.layers.{j}.layernorm.bias""") )
												rename_keys.append((f"""backbone.stages.{i}.{j}.pointwise_conv1.weight""", f"""backbone.encoder.stages.{i}.layers.{j}.pwconv1.weight""") )
												rename_keys.append((f"""backbone.stages.{i}.{j}.pointwise_conv1.bias""", f"""backbone.encoder.stages.{i}.layers.{j}.pwconv1.bias""") )
												rename_keys.append((f"""backbone.stages.{i}.{j}.pointwise_conv2.weight""", f"""backbone.encoder.stages.{i}.layers.{j}.pwconv2.weight""") )
												rename_keys.append((f"""backbone.stages.{i}.{j}.pointwise_conv2.bias""", f"""backbone.encoder.stages.{i}.layers.{j}.pwconv2.bias""") )
								if i > 0:
												rename_keys.append((f"""backbone.downsample_layers.{i}.0.weight""", f"""backbone.encoder.stages.{i}.downsampling_layer.0.weight""") )
												rename_keys.append((f"""backbone.downsample_layers.{i}.0.bias""", f"""backbone.encoder.stages.{i}.downsampling_layer.0.bias""") )
												rename_keys.append((f"""backbone.downsample_layers.{i}.1.weight""", f"""backbone.encoder.stages.{i}.downsampling_layer.1.weight""") )
												rename_keys.append((f"""backbone.downsample_layers.{i}.1.bias""", f"""backbone.encoder.stages.{i}.downsampling_layer.1.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__    (		lowerCAmelCase__       ,lowerCAmelCase__       ,lowerCAmelCase__ ):
				lowercase							=     dct.pop(__lowerCamelCase )
				lowercase							=     val


def       UpperCamelCase__    (		lowerCAmelCase__       ,lowerCAmelCase__       ,lowerCAmelCase__ ):
				lowercase							=     {
				    """upernet-convnext-tiny""": """https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_tiny_fp16_512x512_160k_ade20k/upernet_convnext_tiny_fp16_512x512_160k_ade20k_20220227_124553-cad485de.pth""",
				    """upernet-convnext-small""": """https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_small_fp16_512x512_160k_ade20k/upernet_convnext_small_fp16_512x512_160k_ade20k_20220227_131208-1b1e394f.pth""",
				    """upernet-convnext-base""": """https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_base_fp16_512x512_160k_ade20k/upernet_convnext_base_fp16_512x512_160k_ade20k_20220227_181227-02a24fc6.pth""",
				    """upernet-convnext-large""": """https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_large_fp16_640x640_160k_ade20k/upernet_convnext_large_fp16_640x640_160k_ade20k_20220226_040532-e57aa54d.pth""",
				    """upernet-convnext-xlarge""": """https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_xlarge_fp16_640x640_160k_ade20k/upernet_convnext_xlarge_fp16_640x640_160k_ade20k_20220226_080344-95fc38c2.pth""",
				}
				lowercase							=     model_name_to_url[model_name]
				lowercase							=     torch.hub.load_state_dict_from_url(__lowerCamelCase       ,map_location="""cpu""" )["""state_dict"""]

				lowercase							=     get_upernet_config(__lowerCamelCase )
				lowercase							=     UperNetForSemanticSegmentation(__lowerCamelCase )
				model.eval()

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

				# rename keys
				lowercase							=     create_rename_keys(__lowerCamelCase )
				for src, dest in rename_keys:
								rename_key(__lowerCamelCase       ,__lowerCamelCase       ,__lowerCamelCase )

				model.load_state_dict(__lowerCamelCase )

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

				lowercase							=     SegformerImageProcessor()
				lowercase							=     processor(__lowerCamelCase       ,return_tensors="""pt""" ).pixel_values

				with torch.no_grad():
								lowercase							=     model(__lowerCamelCase )

				if model_name == "upernet-convnext-tiny":
								lowercase							=     torch.tensor(
								    [[-8.8_110, -8.8_110, -8.6_521], [-8.8_110, -8.8_110, -8.6_521], [-8.7_746, -8.7_746, -8.6_130]] )
				elif model_name == "upernet-convnext-small":
								lowercase							=     torch.tensor(
								    [[-8.8_236, -8.8_236, -8.6_771], [-8.8_236, -8.8_236, -8.6_771], [-8.7_638, -8.7_638, -8.6_240]] )
				elif model_name == "upernet-convnext-base":
								lowercase							=     torch.tensor(
								    [[-8.8_558, -8.8_558, -8.6_905], [-8.8_558, -8.8_558, -8.6_905], [-8.7_669, -8.7_669, -8.6_021]] )
				elif model_name == "upernet-convnext-large":
								lowercase							=     torch.tensor(
								    [[-8.6_660, -8.6_660, -8.6_210], [-8.6_660, -8.6_660, -8.6_210], [-8.6_310, -8.6_310, -8.5_964]] )
				elif model_name == "upernet-convnext-xlarge":
								lowercase							=     torch.tensor(
								    [[-8.4_980, -8.4_980, -8.3_977], [-8.4_980, -8.4_980, -8.3_977], [-8.4_379, -8.4_379, -8.3_412]] )
				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__":
					__SCREAMING_SNAKE_CASE							:    Any          =argparse.ArgumentParser()
					# Required parameters
					parser.add_argument(
					    '''--model_name''',
					    default='''upernet-convnext-tiny''',
					    type=str,
					    choices=[f'''upernet-convnext-{size}''' for size in ['''tiny''', '''small''', '''base''', '''large''', '''xlarge''']],
					    help='''Name of the ConvNext 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							:    Union[str, Any]          =parser.parse_args()
					convert_upernet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)







import pytest

from datasets import Dataset, DatasetDict, Features, NamedSplit, Value
from datasets.io.text import TextDatasetReader

from ..utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases






def        __lowerCAmelCase		(	__lowerCamelCase						:     List[Any]					, __lowerCamelCase						:     Any						)	->    Dict:
					assert isinstance(__lowerCamelCase					, __lowerCamelCase						)
					assert dataset.num_rows == 4
					assert dataset.num_columns == 1
					assert dataset.column_names == ["text"]
					for feature, expected_dtype in expected_features.items():
										assert dataset.features[feature].dtype == expected_dtype






@pytest.mark.parametrize("""keep_in_memory"""					, [False, True]						)
def        __lowerCAmelCase		(	__lowerCamelCase						:     Tuple					, __lowerCamelCase						:     str					, __lowerCamelCase						:     Optional[Any]						)	->    List[Any]:
					__lowerCAmelCase						 =tmp_path / """cache"""
					__lowerCAmelCase						 ={"""text""": """string"""}
					with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
										__lowerCAmelCase						 =TextDatasetReader(__lowerCamelCase					, cache_dir=__lowerCamelCase					, keep_in_memory=__lowerCamelCase						).read()
					_check_text_dataset(__lowerCamelCase					, __lowerCamelCase						)






@pytest.mark.parametrize(
    """features"""					, [
        None,
        {"""text""": """string"""},
        {"""text""": """int32"""},
        {"""text""": """float32"""},
    ]					, )
def        __lowerCAmelCase		(	__lowerCamelCase						:     Tuple					, __lowerCamelCase						:     Any					, __lowerCamelCase						:     Tuple						)	->    List[Any]:
					__lowerCAmelCase						 =tmp_path / """cache"""
					__lowerCAmelCase						 ={"""text""": """string"""}
					__lowerCAmelCase						 =features.copy() if features else default_expected_features
					__lowerCAmelCase						 =(
					    Features({feature: Value(__lowerCamelCase						) for feature, dtype in features.items()}						) if features is not None else None
					)
					__lowerCAmelCase						 =TextDatasetReader(__lowerCamelCase					, features=__lowerCamelCase					, cache_dir=__lowerCamelCase						).read()
					_check_text_dataset(__lowerCamelCase					, __lowerCamelCase						)






@pytest.mark.parametrize("""split"""					, [None, NamedSplit("""train"""						), """train""", """test"""]						)
def        __lowerCAmelCase		(	__lowerCamelCase						:     Union[str, Any]					, __lowerCamelCase						:     List[str]					, __lowerCamelCase						:     Union[str, Any]						)	->    Any:
					__lowerCAmelCase						 =tmp_path / """cache"""
					__lowerCAmelCase						 ={"""text""": """string"""}
					__lowerCAmelCase						 =TextDatasetReader(__lowerCamelCase					, cache_dir=__lowerCamelCase					, split=__lowerCamelCase						).read()
					_check_text_dataset(__lowerCamelCase					, __lowerCamelCase						)
					assert dataset.split == split if split else "train"






@pytest.mark.parametrize("""path_type"""					, [str, list]						)
def        __lowerCAmelCase		(	__lowerCamelCase						:     List[Any]					, __lowerCamelCase						:     int					, __lowerCamelCase						:     List[str]						)	->    Any:
					if issubclass(__lowerCamelCase					, __lowerCamelCase						):
										__lowerCAmelCase						 =text_path
					elif issubclass(__lowerCamelCase					, __lowerCamelCase						):
										__lowerCAmelCase						 =[text_path]
					__lowerCAmelCase						 =tmp_path / """cache"""
					__lowerCAmelCase						 ={"""text""": """string"""}
					__lowerCAmelCase						 =TextDatasetReader(__lowerCamelCase					, cache_dir=__lowerCamelCase						).read()
					_check_text_dataset(__lowerCamelCase					, __lowerCamelCase						)






def        __lowerCAmelCase		(	__lowerCamelCase						:     Tuple					, __lowerCamelCase						:     List[str]					, __lowerCamelCase						:     str=("train",)						)	->    Dict:
					assert isinstance(__lowerCamelCase					, __lowerCamelCase						)
					for split in splits:
										__lowerCAmelCase						 =dataset_dict[split]
										assert dataset.num_rows == 4
										assert dataset.num_columns == 1
										assert dataset.column_names == ["text"]
										for feature, expected_dtype in expected_features.items():
															assert dataset.features[feature].dtype == expected_dtype






@pytest.mark.parametrize("""keep_in_memory"""					, [False, True]						)
def        __lowerCAmelCase		(	__lowerCamelCase						:     Optional[int]					, __lowerCamelCase						:     List[Any]					, __lowerCamelCase						:     List[Any]						)	->    Dict:
					__lowerCAmelCase						 =tmp_path / """cache"""
					__lowerCAmelCase						 ={"""text""": """string"""}
					with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
										__lowerCAmelCase						 =TextDatasetReader({"""train""": text_path}					, cache_dir=__lowerCamelCase					, keep_in_memory=__lowerCamelCase						).read()
					_check_text_datasetdict(__lowerCamelCase					, __lowerCamelCase						)






@pytest.mark.parametrize(
    """features"""					, [
        None,
        {"""text""": """string"""},
        {"""text""": """int32"""},
        {"""text""": """float32"""},
    ]					, )
def        __lowerCAmelCase		(	__lowerCamelCase						:     Any					, __lowerCamelCase						:     Optional[int]					, __lowerCamelCase						:     Union[str, Any]						)	->    Tuple:
					__lowerCAmelCase						 =tmp_path / """cache"""
					# CSV file loses col_1 string dtype information: default now is "int64" instead of "string"
					__lowerCAmelCase						 ={"""text""": """string"""}
					__lowerCAmelCase						 =features.copy() if features else default_expected_features
					__lowerCAmelCase						 =(
					    Features({feature: Value(__lowerCamelCase						) for feature, dtype in features.items()}						) if features is not None else None
					)
					__lowerCAmelCase						 =TextDatasetReader({"""train""": text_path}					, features=__lowerCamelCase					, cache_dir=__lowerCamelCase						).read()
					_check_text_datasetdict(__lowerCamelCase					, __lowerCamelCase						)






@pytest.mark.parametrize("""split"""					, [None, NamedSplit("""train"""						), """train""", """test"""]						)
def        __lowerCAmelCase		(	__lowerCamelCase						:     int					, __lowerCamelCase						:     List[str]					, __lowerCamelCase						:     Any						)	->    int:
					if split:
										__lowerCAmelCase						 ={split: text_path}
					else:
										__lowerCAmelCase						 ="""train"""
										__lowerCAmelCase						 ={"""train""": text_path, """test""": text_path}
					__lowerCAmelCase						 =tmp_path / """cache"""
					__lowerCAmelCase						 ={"""text""": """string"""}
					__lowerCAmelCase						 =TextDatasetReader(__lowerCamelCase					, cache_dir=__lowerCamelCase						).read()
					_check_text_datasetdict(__lowerCamelCase					, __lowerCamelCase					, splits=list(path.keys()						)						)
					assert all(dataset[split].split == split for split in path.keys()						)




