
"""simple docstring"""







import argparse
import collections

import torch
from flax import traverse_util
from tax import checkpoints

from transformers import TaConfig, TaEncoderModel, TaForConditionalGeneration
from transformers.utils import logging


logging.set_verbosity_info()
def   __UpperCamelCase   (					SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE="attention" )  ->     Optional[Any]:







							"""simple docstring"""


							__snake_case     			=       params[F'''{prefix}/layers_{i}/{layer_name}/key/kernel''']
							__snake_case     			=       params[F'''{prefix}/layers_{i}/{layer_name}/out/kernel''']
							__snake_case     			=       params[F'''{prefix}/layers_{i}/{layer_name}/query/kernel''']
							__snake_case     			=       params[F'''{prefix}/layers_{i}/{layer_name}/value/kernel''']
							return k, o, q, v
def   __UpperCamelCase   (					SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE=False )  ->     Union[str, Any]:







							"""simple docstring"""


							if split_mlp_wi:
														__snake_case     			=       params[F'''{prefix}/layers_{i}/mlp/wi_0/kernel''']
														__snake_case     			=       params[F'''{prefix}/layers_{i}/mlp/wi_1/kernel''']
														__snake_case     			=       (wi_a, wi_a)
							else:
														__snake_case     			=       params[F'''{prefix}/layers_{i}/mlp/wi/kernel''']

							__snake_case     			=       params[F'''{prefix}/layers_{i}/mlp/wo/kernel''']
							return wi, wo
def   __UpperCamelCase   (					SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE )  ->     str:







							"""simple docstring"""


							return params[F'''{prefix}/layers_{i}/{layer_name}/scale''']
def   __UpperCamelCase   (					SCREAMING_SNAKE_CASE							,    *, SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE )  ->     Tuple:







							"""simple docstring"""


							__snake_case     			=       traverse_util.flatten_dict(variables["target"] )
							__snake_case     			=       {"/".join(SCREAMING_SNAKE_CASE ): v for k, v in old.items()}

							# v1.1 models have a gated GeLU with wi_0 and wi_1 instead of wi
							__snake_case     			=       "encoder/layers_0/mlp/wi_0/kernel" in old
							print("Split MLP:"							,    SCREAMING_SNAKE_CASE )

							__snake_case     			=       collections.OrderedDict()

							# Shared embeddings.
							__snake_case     			=       old["token_embedder/embedding"]

							# Encoder.
							for i in range(SCREAMING_SNAKE_CASE ):
														# Block i, layer 0 (Self Attention).
														__snake_case     			=       tax_layer_norm_lookup(SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE							,    "encoder"							,    "pre_attention_layer_norm" )
														__snake_case			,   __snake_case			,   __snake_case			,   __snake_case     			=       tax_attention_lookup(SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE							,    "encoder"							,    "attention" )
														__snake_case     			=       layer_norm
														__snake_case     			=       k.T
														__snake_case     			=       o.T
														__snake_case     			=       q.T
														__snake_case     			=       v.T

														# Block i, layer 1 (MLP).
														__snake_case     			=       tax_layer_norm_lookup(SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE							,    "encoder"							,    "pre_mlp_layer_norm" )
														__snake_case			,   __snake_case     			=       tax_mlp_lookup(SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE							,    "encoder"							,    SCREAMING_SNAKE_CASE )
														__snake_case     			=       layer_norm
														if split_mlp_wi:
																					__snake_case     			=       wi[0].T
																					__snake_case     			=       wi[1].T
														else:
																					__snake_case     			=       wi.T
														__snake_case     			=       wo.T

							__snake_case     			=       old[
							    "encoder/relpos_bias/rel_embedding"
							].T
							__snake_case     			=       old["encoder/encoder_norm/scale"]

							if not is_encoder_only:
														# Decoder.
														for i in range(SCREAMING_SNAKE_CASE ):
																					# Block i, layer 0 (Self Attention).
																					__snake_case     			=       tax_layer_norm_lookup(SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE							,    "decoder"							,    "pre_self_attention_layer_norm" )
																					__snake_case			,   __snake_case			,   __snake_case			,   __snake_case     			=       tax_attention_lookup(SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE							,    "decoder"							,    "self_attention" )
																					__snake_case     			=       layer_norm
																					__snake_case     			=       k.T
																					__snake_case     			=       o.T
																					__snake_case     			=       q.T
																					__snake_case     			=       v.T

																					# Block i, layer 1 (Cross Attention).
																					__snake_case     			=       tax_layer_norm_lookup(SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE							,    "decoder"							,    "pre_cross_attention_layer_norm" )
																					__snake_case			,   __snake_case			,   __snake_case			,   __snake_case     			=       tax_attention_lookup(SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE							,    "decoder"							,    "encoder_decoder_attention" )
																					__snake_case     			=       layer_norm
																					__snake_case     			=       k.T
																					__snake_case     			=       o.T
																					__snake_case     			=       q.T
																					__snake_case     			=       v.T

																					# Block i, layer 2 (MLP).
																					__snake_case     			=       tax_layer_norm_lookup(SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE							,    "decoder"							,    "pre_mlp_layer_norm" )
																					__snake_case			,   __snake_case     			=       tax_mlp_lookup(SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE							,    "decoder"							,    SCREAMING_SNAKE_CASE )
																					__snake_case     			=       layer_norm
																					if split_mlp_wi:
																												__snake_case     			=       wi[0].T
																												__snake_case     			=       wi[1].T
																					else:
																												__snake_case     			=       wi.T
																					__snake_case     			=       wo.T

														__snake_case     			=       old["decoder/decoder_norm/scale"]
														__snake_case     			=       old[
														    "decoder/relpos_bias/rel_embedding"
														].T

														# LM Head (only in v1.1 checkpoints, in v1.0 embeddings are used instead)
														if "decoder/logits_dense/kernel" in old:
																					__snake_case     			=       old["decoder/logits_dense/kernel"].T

							return new
def   __UpperCamelCase   (					SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE )  ->     Union[str, Any]:







							"""simple docstring"""


							__snake_case     			=       collections.OrderedDict([(k, torch.from_numpy(v.copy() )) for (k, v) in converted_params.items()] )

							# Add what is missing.
							if "encoder.embed_tokens.weight" not in state_dict:
														__snake_case     			=       state_dict["shared.weight"]

							if not is_encoder_only:
														if "decoder.embed_tokens.weight" not in state_dict:
																					__snake_case     			=       state_dict["shared.weight"]

														if "lm_head.weight" not in state_dict:  # For old 1.0 models.
																					print("Using shared word embeddings as lm_head." )
																					__snake_case     			=       state_dict["shared.weight"]

							return state_dict
def   __UpperCamelCase   (					SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE )  ->     Any:







							"""simple docstring"""


							__snake_case     			=       checkpoints.load_tax_checkpoint(SCREAMING_SNAKE_CASE )
							__snake_case     			=       convert_tax_to_pytorch(SCREAMING_SNAKE_CASE							,    num_layers=config.num_layers							,    is_encoder_only=SCREAMING_SNAKE_CASE )
							__snake_case     			=       make_state_dict(SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE )
							model.load_state_dict(SCREAMING_SNAKE_CASE							,    strict=SCREAMING_SNAKE_CASE )
def   __UpperCamelCase   (					SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE = False )  ->     List[str]:







							"""simple docstring"""


							__snake_case     			=       TaConfig.from_json_file(SCREAMING_SNAKE_CASE )
							print(F'''Building PyTorch model from configuration: {config}''' )
							# Non-v1.1 checkpoints could also use T5Model, but this works for all.
							# The v1.0 checkpoints will simply have an LM head that is the word embeddings.
							if is_encoder_only:
														__snake_case     			=       TaEncoderModel(SCREAMING_SNAKE_CASE )
							else:
														__snake_case     			=       TaForConditionalGeneration(SCREAMING_SNAKE_CASE )

							# Load weights from tf checkpoint
							load_tax_weights_in_ta(SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE )

							# Save pytorch-model
							print(F'''Save PyTorch model to {pytorch_dump_path}''' )
							model.save_pretrained(SCREAMING_SNAKE_CASE )

							# Verify that we can load the checkpoint.
							model.from_pretrained(SCREAMING_SNAKE_CASE )
							print("Done" )


if __name__ == "__main__":
		_SCREAMING_SNAKE_CASE	       =      argparse.ArgumentParser(description="""Converts a native T5X checkpoint into a PyTorch checkpoint.""")
		# Required parameters
		parser.add_argument(
		    """--t5x_checkpoint_path""", default=None, type=str, required=True, help="""Path to the T5X checkpoint."""
		)
		parser.add_argument(
		    """--config_file""",
		    default=None,
		    type=str,
		    required=True,
		    help="""The config json file corresponding to the pre-trained T5 model.\nThis specifies the model architecture.""",
		)
		parser.add_argument(
		    """--pytorch_dump_path""", default=None, type=str, required=True, help="""Path to the output PyTorch model."""
		)
		parser.add_argument(
		    """--is_encoder_only""", action="""store_true""", help="""Check if the model is encoder-decoder model""", default=False
		)
		_SCREAMING_SNAKE_CASE	       =      parser.parse_args()
		convert_tax_checkpoint_to_pytorch(
		    args.tax_checkpoint_path, args.config_file, args.pytorch_dump_path, args.is_encoder_only
		)


"""simple docstring"""







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   (					SCREAMING_SNAKE_CASE )  ->     List[str]:







							"""simple docstring"""


							__snake_case     			=       3_84
							if "tiny" in model_name:
														__snake_case     			=       [3, 3, 9, 3]
														__snake_case     			=       [96, 1_92, 3_84, 7_68]
							if "small" in model_name:
														__snake_case     			=       [3, 3, 27, 3]
														__snake_case     			=       [96, 1_92, 3_84, 7_68]
							if "base" in model_name:
														__snake_case     			=       [3, 3, 27, 3]
														__snake_case     			=       [1_28, 2_56, 5_12, 10_24]
														__snake_case     			=       5_12
							if "large" in model_name:
														__snake_case     			=       [3, 3, 27, 3]
														__snake_case     			=       [1_92, 3_84, 7_68, 15_36]
														__snake_case     			=       7_68
							if "xlarge" in model_name:
														__snake_case     			=       [3, 3, 27, 3]
														__snake_case     			=       [2_56, 5_12, 10_24, 20_48]
														__snake_case     			=       10_24

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

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

							return config
def   __UpperCamelCase   (					SCREAMING_SNAKE_CASE )  ->     Optional[Any]:







							"""simple docstring"""


							__snake_case     			=       []

							# 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   (					SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE )  ->     str:







							"""simple docstring"""


							__snake_case     			=       dct.pop(SCREAMING_SNAKE_CASE )
							__snake_case     			=       val
def   __UpperCamelCase   (					SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE )  ->     Optional[int]:







							"""simple docstring"""


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

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

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

							# rename keys
							__snake_case     			=       create_rename_keys(SCREAMING_SNAKE_CASE )
							for src, dest in rename_keys:
														rename_key(SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE )

							model.load_state_dict(SCREAMING_SNAKE_CASE )

							# 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(SCREAMING_SNAKE_CASE							,    stream=SCREAMING_SNAKE_CASE ).raw ).convert("RGB" )

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

							with torch.no_grad():
														__snake_case     			=       model(SCREAMING_SNAKE_CASE )

							if model_name == "upernet-convnext-tiny":
														__snake_case     			=       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":
														__snake_case     			=       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":
														__snake_case     			=       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":
														__snake_case     			=       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":
														__snake_case     			=       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]							,    SCREAMING_SNAKE_CASE							,    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(SCREAMING_SNAKE_CASE )
														print(F'''Saving processor to {pytorch_dump_folder_path}''' )
														processor.save_pretrained(SCREAMING_SNAKE_CASE )

							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	       =      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	       =      parser.parse_args()
		convert_upernet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)
