





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     _lowerCAmelCase		(  _lowerCAmelCase     ):
   '''simple docstring'''


   A_     :				Dict			      =	3_8_4
   if "tiny" in model_name:
      A_     :				Optional[Any]			      =	[3, 3, 9, 3]
      A_     :				str			      =	[9_6, 1_9_2, 3_8_4, 7_6_8]
   if "small" in model_name:
      A_     :				List[Any]			      =	[3, 3, 2_7, 3]
      A_     :				Union[str, Any]			      =	[9_6, 1_9_2, 3_8_4, 7_6_8]
   if "base" in model_name:
      A_     :				Any			      =	[3, 3, 2_7, 3]
      A_     :				Dict			      =	[1_2_8, 2_5_6, 5_1_2, 1_0_2_4]
      A_     :				Any			      =	5_1_2
   if "large" in model_name:
      A_     :				Optional[int]			      =	[3, 3, 2_7, 3]
      A_     :				int			      =	[1_9_2, 3_8_4, 7_6_8, 1_5_3_6]
      A_     :				Optional[Any]			      =	7_6_8
   if "xlarge" in model_name:
      A_     :				Union[str, Any]			      =	[3, 3, 2_7, 3]
      A_     :				Optional[Any]			      =	[2_5_6, 5_1_2, 1_0_2_4, 2_0_4_8]
      A_     :				int			      =	1_0_2_4

   # set label information
   A_     :				Tuple			      =	1_5_0
   A_     :				Optional[int]			      =	"""huggingface/label-files"""
   A_     :				Any			      =	"""ade20k-id2label.json"""
   A_     :				Dict			      =	json.load(open(hf_hub_download(_lowerCAmelCase  ,_lowerCAmelCase  ,repo_type="""dataset"""     )  ,"""r"""     )     )
   A_     :				List[str]			      =	{int(_lowerCAmelCase     ): v for k, v in idalabel.items()}
   A_     :				Any			      =	{v: k for k, v in idalabel.items()}

   A_     :				List[Any]			      =	ConvNextConfig(
       depths=_lowerCAmelCase  ,hidden_sizes=_lowerCAmelCase  ,out_features=["""stage1""", """stage2""", """stage3""", """stage4"""]     )
   A_     :				Any			      =	UperNetConfig(
       backbone_config=_lowerCAmelCase  ,auxiliary_in_channels=_lowerCAmelCase  ,num_labels=_lowerCAmelCase  ,idalabel=_lowerCAmelCase  ,labelaid=_lowerCAmelCase  ,)

   return config



def     _lowerCAmelCase		(  _lowerCAmelCase     ):
   '''simple docstring'''


   A_     :				int			      =	[]

   # 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     _lowerCAmelCase		(  _lowerCAmelCase  ,_lowerCAmelCase  ,_lowerCAmelCase     ):
   '''simple docstring'''


   A_     :				str			      =	dct.pop(_lowerCAmelCase     )
   A_     :				Optional[int]			      =	val



def     _lowerCAmelCase		(  _lowerCAmelCase  ,_lowerCAmelCase  ,_lowerCAmelCase     ):
   '''simple docstring'''


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

   A_     :				Any			      =	get_upernet_config(_lowerCAmelCase     )
   A_     :				List[str]			      =	UperNetForSemanticSegmentation(_lowerCAmelCase     )
   model.eval()

   # replace "bn" => "batch_norm"
   for key in state_dict.copy().keys():
      A_     :				List[Any]			      =	state_dict.pop(_lowerCAmelCase     )
      if "bn" in key:
         A_     :				Any			      =	key.replace("""bn"""  ,"""batch_norm"""     )
      A_     :				List[Any]			      =	val

   # rename keys
   A_     :				int			      =	create_rename_keys(_lowerCAmelCase     )
   for src, dest in rename_keys:
      rename_key(_lowerCAmelCase  ,_lowerCAmelCase  ,_lowerCAmelCase     )

   model.load_state_dict(_lowerCAmelCase     )

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

   A_     :				Optional[int]			      =	SegformerImageProcessor()
   A_     :				Union[str, Any]			      =	processor(_lowerCAmelCase  ,return_tensors="""pt"""     ).pixel_values

   with torch.no_grad():
      A_     :				List[Any]			      =	model(_lowerCAmelCase     )

   if model_name == "upernet-convnext-tiny":
      A_     :				int			      =	torch.tensor(
          [[-8.8110, -8.8110, -8.6521], [-8.8110, -8.8110, -8.6521], [-8.7746, -8.7746, -8.6130]]     )
   elif model_name == "upernet-convnext-small":
      A_     :				Dict			      =	torch.tensor(
          [[-8.8236, -8.8236, -8.6771], [-8.8236, -8.8236, -8.6771], [-8.7638, -8.7638, -8.6240]]     )
   elif model_name == "upernet-convnext-base":
      A_     :				str			      =	torch.tensor(
          [[-8.8558, -8.8558, -8.6905], [-8.8558, -8.8558, -8.6905], [-8.7669, -8.7669, -8.6021]]     )
   elif model_name == "upernet-convnext-large":
      A_     :				str			      =	torch.tensor(
          [[-8.6660, -8.6660, -8.6210], [-8.6660, -8.6660, -8.6210], [-8.6310, -8.6310, -8.5964]]     )
   elif model_name == "upernet-convnext-xlarge":
      A_     :				str			      =	torch.tensor(
          [[-8.4980, -8.4980, -8.3977], [-8.4980, -8.4980, -8.3977], [-8.4379, -8.4379, -8.3412]]     )
   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__":
      _lowerCAmelCase 				=     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."""
      )

      _lowerCAmelCase 				=     parser.parse_args()
      convert_upernet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)













from __future__ import annotations

import numpy as np
from numpy import floataa
from numpy.typing import NDArray



def     _lowerCAmelCase		(  _lowerCAmelCase  ,_lowerCAmelCase  ,_lowerCAmelCase  ,_lowerCAmelCase  ,):
   '''simple docstring'''


   A_      ,				A_     :				int			      =	coefficient_matrix.shape
   A_      ,				A_     :				Tuple			      =	constant_matrix.shape

   if rowsa != colsa:
      A_     :				int			      =	f"""Coefficient matrix dimensions must be nxn but received {rowsa}x{colsa}"""
      raise ValueError(_lowerCAmelCase     )

   if colsa != 1:
      A_     :				List[Any]			      =	f"""Constant matrix must be nx1 but received {rowsa}x{colsa}"""
      raise ValueError(_lowerCAmelCase     )

   if rowsa != rowsa:
      A_     :				str			      =	(
          """Coefficient and constant matrices dimensions must be nxn and nx1 but """
          f"""received {rowsa}x{colsa} and {rowsa}x{colsa}"""
      )
      raise ValueError(_lowerCAmelCase     )

   if len(_lowerCAmelCase     ) != rowsa:
      A_     :				Any			      =	(
          """Number of initial values must be equal to number of rows in coefficient """
          f"""matrix but received {len(_lowerCAmelCase     )} and {rowsa}"""
      )
      raise ValueError(_lowerCAmelCase     )

   if iterations <= 0:
      raise ValueError("""Iterations must be at least 1"""     )

   A_     :				NDArray[floataa]			      =	np.concatenate(
       (coefficient_matrix, constant_matrix)  ,axis=1     )

   A_      ,				A_     :				str			      =	table.shape

   strictly_diagonally_dominant(_lowerCAmelCase     )

   # Iterates the whole matrix for given number of times
   for _ in range(_lowerCAmelCase     ):
      A_     :				Union[str, Any]			      =	[]
      for row in range(_lowerCAmelCase     ):
         A_     :				str			      =	0
         for col in range(_lowerCAmelCase     ):
            if col == row:
               A_     :				Optional[Any]			      =	table[row][col]
            elif col == cols - 1:
               A_     :				List[str]			      =	table[row][col]
            else:
               temp += (-1) * table[row][col] * init_val[col]
         A_     :				str			      =	(temp + val) / denom
         new_val.append(_lowerCAmelCase     )
      A_     :				List[str]			      =	new_val

   return [float(_lowerCAmelCase     ) for i in new_val]



def     _lowerCAmelCase		(  _lowerCAmelCase     ):
   '''simple docstring'''


   A_      ,				A_     :				str			      =	table.shape

   A_     :				Any			      =	True

   for i in range(0  ,_lowerCAmelCase     ):
      A_     :				Optional[Any]			      =	0
      for j in range(0  ,cols - 1     ):
         if i == j:
            continue
         else:
            total += table[i][j]

      if table[i][i] <= total:
         raise ValueError("""Coefficient matrix is not strictly diagonally dominant"""     )

   return is_diagonally_dominant


# Test Cases
if __name__ == "__main__":
      import doctest

      doctest.testmod()






