

"""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 	lowercase      (   __UpperCamelCase						) ->      Optional[Any]:
 __magic_name__     =      384
 if "tiny" in model_name:
  __magic_name__     =      [3, 3, 9, 3]
  __magic_name__     =      [96, 192, 384, 768]
 if "small" in model_name:
  __magic_name__     =      [3, 3, 27, 3]
  __magic_name__     =      [96, 192, 384, 768]
 if "base" in model_name:
  __magic_name__     =      [3, 3, 27, 3]
  __magic_name__     =      [128, 256, 512, 1024]
  __magic_name__     =      512
 if "large" in model_name:
  __magic_name__     =      [3, 3, 27, 3]
  __magic_name__     =      [192, 384, 768, 1536]
  __magic_name__     =      768
 if "xlarge" in model_name:
  __magic_name__     =      [3, 3, 27, 3]
  __magic_name__     =      [256, 512, 1024, 2048]
  __magic_name__     =      1024

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

 __magic_name__     =      ConvNextConfig(
     depths=__UpperCamelCase      ,   hidden_sizes=__UpperCamelCase      ,   out_features=['''stage1''', '''stage2''', '''stage3''', '''stage4''']						)
 __magic_name__     =      UperNetConfig(
     backbone_config=__UpperCamelCase      ,   auxiliary_in_channels=__UpperCamelCase      ,   num_labels=__UpperCamelCase      ,   idalabel=__UpperCamelCase      ,   labelaid=__UpperCamelCase      ,   )

 return config





def 	lowercase      (   __UpperCamelCase						) ->      List[Any]:
 __magic_name__     =      []

 # 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      (   __UpperCamelCase      ,   __UpperCamelCase      ,   __UpperCamelCase						) ->      List[Any]:
 __magic_name__     =      dct.pop(__UpperCamelCase						)
 __magic_name__     =      val





def 	lowercase      (   __UpperCamelCase      ,   __UpperCamelCase      ,   __UpperCamelCase						) ->      int:
 __magic_name__     =      {
     '''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''',
 }
 __magic_name__     =      model_name_to_url[model_name]
 __magic_name__     =      torch.hub.load_state_dict_from_url(__UpperCamelCase      ,   map_location='''cpu'''						)['''state_dict''']

 __magic_name__     =      get_upernet_config(__UpperCamelCase						)
 __magic_name__     =      UperNetForSemanticSegmentation(__UpperCamelCase						)
 model.eval()

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

 # rename keys
 __magic_name__     =      create_rename_keys(__UpperCamelCase						)
 for src, dest in rename_keys:
  rename_key(__UpperCamelCase      ,   __UpperCamelCase      ,   __UpperCamelCase						)

 model.load_state_dict(__UpperCamelCase						)

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

 __magic_name__     =      SegformerImageProcessor()
 __magic_name__     =      processor(__UpperCamelCase      ,   return_tensors='''pt'''						).pixel_values

 with torch.no_grad():
  __magic_name__     =      model(__UpperCamelCase						)

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

 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)




"""simple docstring"""






import pytest

import datasets


# Import fixture modules as plugins
__lowerCamelCase										=		["tests.fixtures.files", "tests.fixtures.hub", "tests.fixtures.fsspec"]





def 	lowercase      (   __UpperCamelCase      ,   __UpperCamelCase						) ->      str:
 # Mark tests as "unit" by default if not marked as "integration" (or already marked as "unit")
 for item in items:
  if any(marker in item.keywords for marker in ['''integration''', '''unit''']						):
   continue
  item.add_marker(pytest.mark.unit						)





def 	lowercase      (   __UpperCamelCase						) ->      Union[str, Any]:
 config.addinivalue_line('''markers'''      ,   '''torchaudio_latest: mark test to run with torchaudio>=0.12'''						)





@pytest.fixture(autouse=__UpperCamelCase						)
def 	lowercase      (   __UpperCamelCase      ,   __UpperCamelCase						) ->      Dict:
 # test_hf_cache_home = tmp_path_factory.mktemp("cache")  # TODO: why a cache dir per test function does not work?
 __magic_name__     =      tmp_path_factory.getbasetemp() / '''cache'''
 __magic_name__     =      test_hf_cache_home / '''datasets'''
 __magic_name__     =      test_hf_cache_home / '''metrics'''
 __magic_name__     =      test_hf_cache_home / '''modules'''
 monkeypatch.setattr('''datasets.config.HF_DATASETS_CACHE'''      ,   str(__UpperCamelCase						)						)
 monkeypatch.setattr('''datasets.config.HF_METRICS_CACHE'''      ,   str(__UpperCamelCase						)						)
 monkeypatch.setattr('''datasets.config.HF_MODULES_CACHE'''      ,   str(__UpperCamelCase						)						)
 __magic_name__     =      test_hf_datasets_cache / '''downloads'''
 monkeypatch.setattr('''datasets.config.DOWNLOADED_DATASETS_PATH'''      ,   str(__UpperCamelCase						)						)
 __magic_name__     =      test_hf_datasets_cache / '''downloads''' / '''extracted'''
 monkeypatch.setattr('''datasets.config.EXTRACTED_DATASETS_PATH'''      ,   str(__UpperCamelCase						)						)





@pytest.fixture(autouse=__UpperCamelCase      ,   scope='''session'''						)
def 	lowercase      (   ) ->      Any:
 datasets.disable_progress_bar()





@pytest.fixture(autouse=__UpperCamelCase						)
def 	lowercase      (   __UpperCamelCase						) ->      str:
 # don't take tests into account when counting downloads
 monkeypatch.setattr('''datasets.config.HF_UPDATE_DOWNLOAD_COUNTS'''      ,   __UpperCamelCase						)





@pytest.fixture
def 	lowercase      (   __UpperCamelCase						) ->      int:
 # Required to suppress RemovedIn20Warning when feature(s) are not compatible with SQLAlchemy 2.0
 # To be removed once SQLAlchemy 2.0 supported
 monkeypatch.setattr('''sqlalchemy.util.deprecations.SILENCE_UBER_WARNING'''      ,   __UpperCamelCase						)

