






'''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        _SCREAMING_SNAKE_CASE  (				A  :       Tuple  )    -> List[Any]:







   """simple docstring"""
   __snake_case				:      Optional[int]     							= 3_84
   if "tiny" in model_name:
      __snake_case				:      Any     							= [3, 3, 9, 3]
      __snake_case				:      Union[str, Any]     							= [96, 1_92, 3_84, 7_68]
   if "small" in model_name:
      __snake_case				:      int     							= [3, 3, 27, 3]
      __snake_case				:      Dict     							= [96, 1_92, 3_84, 7_68]
   if "base" in model_name:
      __snake_case				:      Tuple     							= [3, 3, 27, 3]
      __snake_case				:      List[Any]     							= [1_28, 2_56, 5_12, 10_24]
      __snake_case				:      Union[str, Any]     							= 5_12
   if "large" in model_name:
      __snake_case				:      int     							= [3, 3, 27, 3]
      __snake_case				:      Union[str, Any]     							= [1_92, 3_84, 7_68, 15_36]
      __snake_case				:      Optional[Any]     							= 7_68
   if "xlarge" in model_name:
      __snake_case				:      Dict     							= [3, 3, 27, 3]
      __snake_case				:      str     							= [2_56, 5_12, 10_24, 20_48]
      __snake_case				:      List[str]     							= 10_24

   # set label information
   __snake_case				:      Any     							= 1_50
   __snake_case				:      Optional[int]     							= '''huggingface/label-files'''
   __snake_case				:      Union[str, Any]     							= '''ade20k-id2label.json'''
   __snake_case				:      List[Any]     							= json.load(open(hf_hub_download(__UpperCamelCase  ,		__UpperCamelCase  ,		repo_type='dataset'  )  ,		'r'  )  )
   __snake_case				:      Any     							= {int(__UpperCamelCase  ): v for k, v in idalabel.items()}
   __snake_case				:      Any     							= {v: k for k, v in idalabel.items()}

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

   return config




def        _SCREAMING_SNAKE_CASE  (				A  :       Dict  )    -> Dict:







   """simple docstring"""
   __snake_case				:      Optional[Any]     							= []

   # 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        _SCREAMING_SNAKE_CASE  (				A  :       Any  ,		A  :       List[Any]  ,		A  :       Optional[int]  )    -> Dict:







   """simple docstring"""
   __snake_case				:      Optional[Any]     							= dct.pop(__UpperCamelCase  )
   __snake_case				:      Dict     							= val




def        _SCREAMING_SNAKE_CASE  (				A  :       List[Any]  ,		A  :       str  ,		A  :       Optional[Any]  )    -> Tuple:







   """simple docstring"""
   __snake_case				:      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''',
   }
   __snake_case				:      Dict     							= model_name_to_url[model_name]
   __snake_case				:      str     							= torch.hub.load_state_dict_from_url(__UpperCamelCase  ,		map_location='cpu'  )['''state_dict''']

   __snake_case				:      Dict     							= get_upernet_config(__UpperCamelCase  )
   __snake_case				:      Union[str, Any]     							= UperNetForSemanticSegmentation(__UpperCamelCase  )
   model.eval()

   # replace "bn" => "batch_norm"
   for key in state_dict.copy().keys():
      __snake_case				:      int     							= state_dict.pop(__UpperCamelCase  )
      if "bn" in key:
         __snake_case				:      Union[str, Any]     							= key.replace('bn'  ,		'batch_norm'  )
      __snake_case				:      Union[str, Any]     							= val

   # rename keys
   __snake_case				:      List[str]     							= create_rename_keys(__UpperCamelCase  )
   for src, dest in rename_keys:
      rename_key(__UpperCamelCase  ,		__UpperCamelCase  ,		__UpperCamelCase  )

   model.load_state_dict(__UpperCamelCase  )

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

   __snake_case				:      Union[str, Any]     							= SegformerImageProcessor()
   __snake_case				:      Union[str, Any]     							= processor(__UpperCamelCase  ,		return_tensors='pt'  ).pixel_values

   with torch.no_grad():
      __snake_case				:      Dict     							= model(__UpperCamelCase  )

   if model_name == "upernet-convnext-tiny":
      __snake_case				:      Union[str, Any]     							= 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":
      __snake_case				:      List[str]     							= 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":
      __snake_case				:      Any     							= 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":
      __snake_case				:      Union[str, Any]     							= 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":
      __snake_case				:      List[Any]     							= 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]  ,		__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__":
      __A    	=   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.'''
      )

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







'''simple docstring'''






import unittest

import torch

from diffusers import VQModel
from diffusers.utils import floats_tensor, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()




class a_    (    UpperCamelCase_		,      UpperCamelCase_		,      unittest.TestCase							):
   _snake_case        =							VQModel
   _snake_case        =							"""sample"""







   @property
   def     SCREAMING_SNAKE_CASE__					(self					,   __a=(3_2, 3_2))			->      str:
      """simple docstring"""






      __snake_case				:      Dict     							= 4
      __snake_case				:      Optional[int]     							= 3

      __snake_case				:      str     							= floats_tensor((batch_size, num_channels) + sizes).to(__a)

      return {"sample": image}







   @property
   def     SCREAMING_SNAKE_CASE__					(self)			->      Union[str, Any]:
      """simple docstring"""






      return (3, 3_2, 3_2)







   @property
   def     SCREAMING_SNAKE_CASE__					(self)			->      List[Any]:
      """simple docstring"""






      return (3, 3_2, 3_2)







   def     SCREAMING_SNAKE_CASE__					(self)			->      Optional[int]:
      """simple docstring"""






      __snake_case				:      Optional[Any]     							= {
          'block_out_channels': [3_2, 6_4],
          'in_channels': 3,
          'out_channels': 3,
          'down_block_types': ['DownEncoderBlock2D', 'DownEncoderBlock2D'],
          'up_block_types': ['UpDecoderBlock2D', 'UpDecoderBlock2D'],
          'latent_channels': 3,
      }
      __snake_case				:      List[Any]     							= self.dummy_input
      return init_dict, inputs_dict







   def     SCREAMING_SNAKE_CASE__					(self)			->      Any:
      """simple docstring"""






      pass







   def     SCREAMING_SNAKE_CASE__					(self)			->      int:
      """simple docstring"""






      pass







   def     SCREAMING_SNAKE_CASE__					(self)			->      Union[str, Any]:
      """simple docstring"""






      __snake_case							,__snake_case				:      List[Any]     							= VQModel.from_pretrained('fusing/vqgan-dummy'					,   output_loading_info=__a)
      self.assertIsNotNone(__a)
      self.assertEqual(len(loading_info['missing_keys'])					,   0)

      model.to(__a)
      __snake_case				:      Any     							= model(**self.dummy_input)

      assert image is not None, "Make sure output is not None"







   def     SCREAMING_SNAKE_CASE__					(self)			->      Optional[int]:
      """simple docstring"""






      __snake_case				:      Union[str, Any]     							= VQModel.from_pretrained('fusing/vqgan-dummy')
      model.to(__a).eval()

      torch.manual_seed(0)
      if torch.cuda.is_available():
         torch.cuda.manual_seed_all(0)

      __snake_case				:      Tuple     							= torch.randn(1					,   model.config.in_channels					,   model.config.sample_size					,   model.config.sample_size)
      __snake_case				:      Optional[int]     							= image.to(__a)
      with torch.no_grad():
         __snake_case				:      List[Any]     							= model(__a).sample

      __snake_case				:      int     							= output[0, -1, -3:, -3:].flatten().cpu()
      # fmt: off
      __snake_case				:      int     							= torch.tensor([-0.0_153, -0.4_044, -0.1_880, -0.5_161, -0.2_418, -0.4_072, -0.1_612, -0.0_633, -0.0_143])
      # fmt: on
      self.assertTrue(torch.allclose(__a					,   __a					,   atol=1E-3))