





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__(       UpperCamelCase__			:       int  )->Tuple:
       A__											=						3_84
       if "tiny" in model_name:
              A__											=						[3, 3, 9, 3]
              A__											=						[96, 1_92, 3_84, 7_68]
       if "small" in model_name:
              A__											=						[3, 3, 27, 3]
              A__											=						[96, 1_92, 3_84, 7_68]
       if "base" in model_name:
              A__											=						[3, 3, 27, 3]
              A__											=						[1_28, 2_56, 5_12, 10_24]
              A__											=						5_12
       if "large" in model_name:
              A__											=						[3, 3, 27, 3]
              A__											=						[1_92, 3_84, 7_68, 15_36]
              A__											=						7_68
       if "xlarge" in model_name:
              A__											=						[3, 3, 27, 3]
              A__											=						[2_56, 5_12, 10_24, 20_48]
              A__											=						10_24

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

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

       return config


def 	UpperCamelCase__(       UpperCamelCase__			:       List[Any]  )->int:
       A__											=						[]

       # 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__(       UpperCamelCase__			:       Dict ,     UpperCamelCase__			:       Optional[Any] ,     UpperCamelCase__			:       Optional[int]  )->str:
       A__											=						dct.pop(UpperCamelCase__  )
       A__											=						val


def 	UpperCamelCase__(       UpperCamelCase__			:       List[str] ,     UpperCamelCase__			:       List[Any] ,     UpperCamelCase__			:       Any  )->str:
       A__											=						{
           '''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__											=						model_name_to_url[model_name]
       A__											=						torch.hub.load_state_dict_from_url(UpperCamelCase__ ,     map_location='''cpu'''  )['''state_dict''']

       A__											=						get_upernet_config(UpperCamelCase__  )
       A__											=						UperNetForSemanticSegmentation(UpperCamelCase__  )
       model.eval()

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

       # rename keys
       A__											=						create_rename_keys(UpperCamelCase__  )
       for src, dest in rename_keys:
              rename_key(UpperCamelCase__ ,     UpperCamelCase__ ,     UpperCamelCase__  )

       model.load_state_dict(UpperCamelCase__  )

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

       A__											=						SegformerImageProcessor()
       A__											=						processor(UpperCamelCase__ ,     return_tensors='''pt'''  ).pixel_values

       with torch.no_grad():
              A__											=						model(UpperCamelCase__  )

       if model_name == "upernet-convnext-tiny":
              A__											=						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__											=						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__											=						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__											=						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__											=						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__:      Dict									=		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__:      Any									=		parser.parse_args()
       convert_upernet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)












import unittest

from transformers import AutoConfig, AutoTokenizer, BertConfig, TensorType, is_flax_available
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, require_flax, slow


if is_flax_available():
       import jax

       from transformers.models.auto.modeling_flax_auto import FlaxAutoModel
       from transformers.models.bert.modeling_flax_bert import FlaxBertModel
       from transformers.models.roberta.modeling_flax_roberta import FlaxRobertaModel




@require_flax
class     SCREAMING_SNAKE_CASE__ (						unittest.TestCase      ):



 @slow
 def 				UpperCamelCase			(							self ):
        for model_name in ["bert-base-cased", "bert-large-uncased"]:
               with self.subTest(__lowerCamelCase ):
                      A__											=						AutoConfig.from_pretrained(__lowerCamelCase )
                      self.assertIsNotNone(__lowerCamelCase )
                      self.assertIsInstance(__lowerCamelCase,__lowerCamelCase )

                      A__											=						FlaxAutoModel.from_pretrained(__lowerCamelCase )
                      self.assertIsNotNone(__lowerCamelCase )
                      self.assertIsInstance(__lowerCamelCase,__lowerCamelCase )



 @slow
 def 				UpperCamelCase			(							self ):
        for model_name in ["roberta-base", "roberta-large"]:
               with self.subTest(__lowerCamelCase ):
                      A__											=						AutoConfig.from_pretrained(__lowerCamelCase )
                      self.assertIsNotNone(__lowerCamelCase )
                      self.assertIsInstance(__lowerCamelCase,__lowerCamelCase )

                      A__											=						FlaxAutoModel.from_pretrained(__lowerCamelCase )
                      self.assertIsNotNone(__lowerCamelCase )
                      self.assertIsInstance(__lowerCamelCase,__lowerCamelCase )



 @slow
 def 				UpperCamelCase			(							self ):
        for model_name in ["bert-base-cased", "bert-large-uncased"]:
               A__											=						AutoTokenizer.from_pretrained(__lowerCamelCase )
               A__											=						FlaxBertModel.from_pretrained(__lowerCamelCase )
               A__											=						tokenizer('''Do you support jax jitted function?''',return_tensors=TensorType.JAX )

               @jax.jit
               def eval(**__lowerCamelCase ):
                      return model(**__lowerCamelCase )

               eval(**__lowerCamelCase ).block_until_ready()



 @slow
 def 				UpperCamelCase			(							self ):
        for model_name in ["roberta-base", "roberta-large"]:
               A__											=						AutoTokenizer.from_pretrained(__lowerCamelCase )
               A__											=						FlaxRobertaModel.from_pretrained(__lowerCamelCase )
               A__											=						tokenizer('''Do you support jax jitted function?''',return_tensors=TensorType.JAX )

               @jax.jit
               def eval(**__lowerCamelCase ):
                      return model(**__lowerCamelCase )

               eval(**__lowerCamelCase ).block_until_ready()



 def 				UpperCamelCase			(							self ):
        with self.assertRaisesRegex(
            __lowerCamelCase,'''bert-base is not a local folder and is not a valid model identifier''' ):
               A__											=						FlaxAutoModel.from_pretrained('''bert-base''' )



 def 				UpperCamelCase			(							self ):
        with self.assertRaisesRegex(
            __lowerCamelCase,r'''aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)''' ):
               A__											=						FlaxAutoModel.from_pretrained(__lowerCamelCase,revision='''aaaaaa''' )



 def 				UpperCamelCase			(							self ):
        with self.assertRaisesRegex(
            __lowerCamelCase,'''hf-internal-testing/config-no-model does not appear to have a file named flax_model.msgpack''',):
               A__											=						FlaxAutoModel.from_pretrained('''hf-internal-testing/config-no-model''' )




 def 				UpperCamelCase			(							self ):
        with self.assertRaisesRegex(__lowerCamelCase,'''Use `from_pt=True` to load this model''' ):
               A__											=						FlaxAutoModel.from_pretrained('''hf-internal-testing/tiny-bert-pt-only''' )





