





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        __lowercase       (				_SCREAMING_SNAKE_CASE     )			->       Dict:


      '''simple docstring'''


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

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

      SCREAMING_SNAKE_CASE		=						ConvNextConfig(
          depths=_SCREAMING_SNAKE_CASE			,						hidden_sizes=_SCREAMING_SNAKE_CASE			,						out_features=["""stage1""", """stage2""", """stage3""", """stage4"""]     )
      SCREAMING_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        __lowercase       (				_SCREAMING_SNAKE_CASE     )			->       List[str]:


      '''simple docstring'''


      SCREAMING_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        __lowercase       (				_SCREAMING_SNAKE_CASE			,						_SCREAMING_SNAKE_CASE			,						_SCREAMING_SNAKE_CASE     )			->       Optional[int]:


      '''simple docstring'''


      SCREAMING_SNAKE_CASE		=						dct.pop(_SCREAMING_SNAKE_CASE     )
      SCREAMING_SNAKE_CASE		=						val



def        __lowercase       (				_SCREAMING_SNAKE_CASE			,						_SCREAMING_SNAKE_CASE			,						_SCREAMING_SNAKE_CASE     )			->       Optional[Any]:


      '''simple docstring'''


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

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

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

      # rename keys
      SCREAMING_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
      SCREAMING_SNAKE_CASE		=						"""https://huggingface.co/datasets/hf-internal-testing/fixtures_ade20k/resolve/main/ADE_val_00000001.jpg"""
      SCREAMING_SNAKE_CASE		=						Image.open(requests.get(_SCREAMING_SNAKE_CASE			,						stream=_SCREAMING_SNAKE_CASE     ).raw     ).convert("""RGB"""     )

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

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

      if model_name == "upernet-convnext-tiny":
            SCREAMING_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":
            SCREAMING_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":
            SCREAMING_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":
            SCREAMING_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":
            SCREAMING_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)









from itertools import permutations



def        __lowercase       (				_SCREAMING_SNAKE_CASE     )			->       bool:


      '''simple docstring'''


      if num[3] % 2 != 0:
            return False

      if (num[2] + num[3] + num[4]) % 3 != 0:
            return False

      if num[5] % 5 != 0:
            return False

      SCREAMING_SNAKE_CASE		=						[7, 11, 13, 17]
      for i, test in enumerate(_SCREAMING_SNAKE_CASE     ):
            if (num[i + 4] * 1_00 + num[i + 5] * 10 + num[i + 6]) % test != 0:
                  return False
      return True



def        __lowercase       (				_SCREAMING_SNAKE_CASE = 10     )			->       int:


      '''simple docstring'''


      return sum(
          int("""""".join(map(_SCREAMING_SNAKE_CASE			,						_SCREAMING_SNAKE_CASE     )     )     )
          for num in permutations(range(_SCREAMING_SNAKE_CASE     )     )
          if is_substring_divisible(_SCREAMING_SNAKE_CASE     )     )


if __name__ == "__main__":
      print(F'''{solution() = }''')


