



from __future__ import annotations

import os
import tempfile
import unittest

from transformers import ConvBertConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
				import tensorflow as tf

				from transformers import (
				    TFConvBertForMaskedLM,
				    TFConvBertForMultipleChoice,
				    TFConvBertForQuestionAnswering,
				    TFConvBertForSequenceClassification,
				    TFConvBertForTokenClassification,
				    TFConvBertModel,
				)



class        lowerCAmelCase__					:


				def __init__(							self      :  Any						, __UpperCamelCase      :  int						, __UpperCamelCase      :  Optional[Any]=13						, __UpperCamelCase      :  Optional[int]=7						, __UpperCamelCase      :  str=True						, __UpperCamelCase      :  Tuple=True						, __UpperCamelCase      :  Optional[int]=True						, __UpperCamelCase      :  Union[str, Any]=True						, __UpperCamelCase      :  Optional[Any]=99						, __UpperCamelCase      :  Any=32						, __UpperCamelCase      :  List[str]=2						, __UpperCamelCase      :  int=4						, __UpperCamelCase      :  Dict=37						, __UpperCamelCase      :  int="gelu"						, __UpperCamelCase      :  Optional[Any]=0.1						, __UpperCamelCase      :  Union[str, Any]=0.1						, __UpperCamelCase      :  List[Any]=512						, __UpperCamelCase      :  Union[str, Any]=16						, __UpperCamelCase      :  List[str]=2						, __UpperCamelCase      :  Union[str, Any]=0.0_2						, __UpperCamelCase      :  Tuple=3						, __UpperCamelCase      :  Any=4						, __UpperCamelCase      :  Dict=None						, )		->     str:
											A      						=    parent
											A      						=    13
											A      						=    7
											A      						=    True
											A      						=    True
											A      						=    True
											A      						=    True
											A      						=    99
											A      						=    384
											A      						=    2
											A      						=    4
											A      						=    37
											A      						=    'gelu'
											A      						=    0.1
											A      						=    0.1
											A      						=    512
											A      						=    16
											A      						=    2
											A      						=    0.0_2
											A      						=    3
											A      						=    4
											A      						=    128
											A      						=    2
											A      						=    9
											A      						=    1
											A      						=    None


				def __UpperCamelCase						(							self      :  str       )		->     List[str]:
											A      						=    ids_tensor([self.batch_size, self.seq_length]						, self.vocab_size       )

											A      						=    None
											if self.use_input_mask:
																		A      						=    random_attention_mask([self.batch_size, self.seq_length]       )

											A      						=    None
											if self.use_token_type_ids:
																		A      						=    ids_tensor([self.batch_size, self.seq_length]						, self.type_vocab_size       )

											A      						=    None
											A      						=    None
											A      						=    None
											if self.use_labels:
																		A      						=    ids_tensor([self.batch_size]						, self.type_sequence_label_size       )
																		A      						=    ids_tensor([self.batch_size, self.seq_length]						, self.num_labels       )
																		A      						=    ids_tensor([self.batch_size]						, self.num_choices       )

											A      						=    ConvBertConfig(
											    vocab_size=self.vocab_size						, hidden_size=self.hidden_size						, num_hidden_layers=self.num_hidden_layers						, num_attention_heads=self.num_attention_heads						, intermediate_size=self.intermediate_size						, hidden_act=self.hidden_act						, hidden_dropout_prob=self.hidden_dropout_prob						, attention_probs_dropout_prob=self.attention_probs_dropout_prob						, max_position_embeddings=self.max_position_embeddings						, type_vocab_size=self.type_vocab_size						, initializer_range=self.initializer_range						, return_dict=__UpperCamelCase						, )

											return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels


				def __UpperCamelCase						(							self      :  Union[str, Any]						, __UpperCamelCase      :  str						, __UpperCamelCase      :  List[str]						, __UpperCamelCase      :  Dict						, __UpperCamelCase      :  str						, __UpperCamelCase      :  Optional[int]						, __UpperCamelCase      :  List[Any]						, __UpperCamelCase      :  List[str]       )		->     Optional[int]:
											A      						=    TFConvBertModel(config=__UpperCamelCase       )
											A      						=    {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}

											A      						=    [input_ids, input_mask]
											A      						=    model(__UpperCamelCase       )

											A      						=    model(__UpperCamelCase       )

											self.parent.assertEqual(result.last_hidden_state.shape						, (self.batch_size, self.seq_length, self.hidden_size)       )


				def __UpperCamelCase						(							self      :  Union[str, Any]						, __UpperCamelCase      :  List[str]						, __UpperCamelCase      :  Union[str, Any]						, __UpperCamelCase      :  Union[str, Any]						, __UpperCamelCase      :  List[str]						, __UpperCamelCase      :  str						, __UpperCamelCase      :  Dict						, __UpperCamelCase      :  Any       )		->     List[str]:
											A      						=    TFConvBertForMaskedLM(config=__UpperCamelCase       )
											A      						=    {
											    'input_ids': input_ids,
											    'attention_mask': input_mask,
											    'token_type_ids': token_type_ids,
											}
											A      						=    model(__UpperCamelCase       )
											self.parent.assertEqual(result.logits.shape						, (self.batch_size, self.seq_length, self.vocab_size)       )


				def __UpperCamelCase						(							self      :  Union[str, Any]						, __UpperCamelCase      :  Dict						, __UpperCamelCase      :  Optional[int]						, __UpperCamelCase      :  List[str]						, __UpperCamelCase      :  List[str]						, __UpperCamelCase      :  List[Any]						, __UpperCamelCase      :  Union[str, Any]						, __UpperCamelCase      :  List[str]       )		->     List[str]:
											A      						=    self.num_labels
											A      						=    TFConvBertForSequenceClassification(config=__UpperCamelCase       )
											A      						=    {
											    'input_ids': input_ids,
											    'attention_mask': input_mask,
											    'token_type_ids': token_type_ids,
											}

											A      						=    model(__UpperCamelCase       )
											self.parent.assertEqual(result.logits.shape						, (self.batch_size, self.num_labels)       )


				def __UpperCamelCase						(							self      :  int						, __UpperCamelCase      :  Union[str, Any]						, __UpperCamelCase      :  Optional[Any]						, __UpperCamelCase      :  Dict						, __UpperCamelCase      :  List[Any]						, __UpperCamelCase      :  Optional[int]						, __UpperCamelCase      :  Optional[int]						, __UpperCamelCase      :  Optional[Any]       )		->     Any:
											A      						=    self.num_choices
											A      						=    TFConvBertForMultipleChoice(config=__UpperCamelCase       )
											A      						=    tf.tile(tf.expand_dims(__UpperCamelCase						, 1       )						, (1, self.num_choices, 1)       )
											A      						=    tf.tile(tf.expand_dims(__UpperCamelCase						, 1       )						, (1, self.num_choices, 1)       )
											A      						=    tf.tile(tf.expand_dims(__UpperCamelCase						, 1       )						, (1, self.num_choices, 1)       )
											A      						=    {
											    'input_ids': multiple_choice_inputs_ids,
											    'attention_mask': multiple_choice_input_mask,
											    'token_type_ids': multiple_choice_token_type_ids,
											}
											A      						=    model(__UpperCamelCase       )
											self.parent.assertEqual(result.logits.shape						, (self.batch_size, self.num_choices)       )


				def __UpperCamelCase						(							self      :  int						, __UpperCamelCase      :  List[str]						, __UpperCamelCase      :  int						, __UpperCamelCase      :  Dict						, __UpperCamelCase      :  Any						, __UpperCamelCase      :  List[Any]						, __UpperCamelCase      :  List[Any]						, __UpperCamelCase      :  Any       )		->     Union[str, Any]:
											A      						=    self.num_labels
											A      						=    TFConvBertForTokenClassification(config=__UpperCamelCase       )
											A      						=    {
											    'input_ids': input_ids,
											    'attention_mask': input_mask,
											    'token_type_ids': token_type_ids,
											}
											A      						=    model(__UpperCamelCase       )
											self.parent.assertEqual(result.logits.shape						, (self.batch_size, self.seq_length, self.num_labels)       )


				def __UpperCamelCase						(							self      :  Optional[int]						, __UpperCamelCase      :  List[Any]						, __UpperCamelCase      :  Optional[int]						, __UpperCamelCase      :  Optional[Any]						, __UpperCamelCase      :  Dict						, __UpperCamelCase      :  Dict						, __UpperCamelCase      :  Any						, __UpperCamelCase      :  Tuple       )		->     Dict:
											A      						=    TFConvBertForQuestionAnswering(config=__UpperCamelCase       )
											A      						=    {
											    'input_ids': input_ids,
											    'attention_mask': input_mask,
											    'token_type_ids': token_type_ids,
											}

											A      						=    model(__UpperCamelCase       )
											self.parent.assertEqual(result.start_logits.shape						, (self.batch_size, self.seq_length)       )
											self.parent.assertEqual(result.end_logits.shape						, (self.batch_size, self.seq_length)       )


				def __UpperCamelCase						(							self      :  Tuple       )		->     int:
											A      						=    self.prepare_config_and_inputs()
											(
											    (
											    A
											)						,   (
											    A
											)						,   (
											    A
											)						,   (
											    A
											)						,   (
											    A
											)						,   (
											    A
											)						,   (
											    A
											)						,   
											)      						=    config_and_inputs
											A      						=    {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
											return config, inputs_dict





@require_tf
class        lowerCAmelCase__					(  _lowerCamelCase ,				_lowerCamelCase ,				unittest.TestCase       ):
				A_   :     int								=     (
				    (
				        TFConvBertModel,
				        TFConvBertForMaskedLM,
				        TFConvBertForQuestionAnswering,
				        TFConvBertForSequenceClassification,
				        TFConvBertForTokenClassification,
				        TFConvBertForMultipleChoice,
				    )
				    if is_tf_available()
				    else ()
				)
				A_   :     Union[str, Any]								=     (
				    {
				        'feature-extraction': TFConvBertModel,
				        'fill-mask': TFConvBertForMaskedLM,
				        'question-answering': TFConvBertForQuestionAnswering,
				        'text-classification': TFConvBertForSequenceClassification,
				        'token-classification': TFConvBertForTokenClassification,
				        'zero-shot': TFConvBertForSequenceClassification,
				    }
				    if is_tf_available()
				    else {}
				)
				A_   :     Tuple								=     False
				A_   :     Optional[int]								=     False
				A_   :     List[str]								=     False


				def __UpperCamelCase						(							self      :  Optional[Any]       )		->     Optional[Any]:
											A      						=    TFConvBertModelTester(self       )
											A      						=    ConfigTester(self						, config_class=__UpperCamelCase						, hidden_size=37       )


				def __UpperCamelCase						(							self      :  int       )		->     List[Any]:
											self.config_tester.run_common_tests()


				def __UpperCamelCase						(							self      :  Any       )		->     Optional[Any]:
											A      						=    self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*__UpperCamelCase       )


				def __UpperCamelCase						(							self      :  int       )		->     int:
											A      						=    self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_masked_lm(*__UpperCamelCase       )


				def __UpperCamelCase						(							self      :  Optional[Any]       )		->     int:
											A      						=    self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_multiple_choice(*__UpperCamelCase       )


				def __UpperCamelCase						(							self      :  Dict       )		->     Dict:
											A      						=    self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_question_answering(*__UpperCamelCase       )


				def __UpperCamelCase						(							self      :  str       )		->     str:
											A      						=    self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_sequence_classification(*__UpperCamelCase       )


				def __UpperCamelCase						(							self      :  Optional[int]       )		->     Optional[int]:
											A      						=    self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_token_classification(*__UpperCamelCase       )


				@slow
				def __UpperCamelCase						(							self      :  int       )		->     Tuple:
											A						,   A      						=    self.model_tester.prepare_config_and_inputs_for_common()
											A      						=    True
											A      						=    True

											if hasattr(__UpperCamelCase						, 'use_cache'       ):
																		A      						=    True

											A      						=    getattr(self.model_tester						, 'encoder_seq_length'						, self.model_tester.seq_length       )
											A      						=    getattr(self.model_tester						, 'key_length'						, __UpperCamelCase       )

											for model_class in self.all_model_classes:
																		A      						=    self._prepare_for_class(__UpperCamelCase						, __UpperCamelCase       )
																		A      						=    model_class(__UpperCamelCase       )
																		A      						=    len(model(__UpperCamelCase       )       )

																		with tempfile.TemporaryDirectory() as tmpdirname:
																									model.save_pretrained(__UpperCamelCase						, saved_model=__UpperCamelCase       )
																									A      						=    os.path.join(__UpperCamelCase						, 'saved_model'						, '1'       )
																									A      						=    tf.keras.models.load_model(__UpperCamelCase       )
																									A      						=    model(__UpperCamelCase       )

																									if self.is_encoder_decoder:
																																A      						=    outputs['encoder_hidden_states']
																																A      						=    outputs['encoder_attentions']
																									else:
																																A      						=    outputs['hidden_states']
																																A      						=    outputs['attentions']

																									self.assertEqual(len(__UpperCamelCase       )						, __UpperCamelCase       )

																									A      						=    getattr(
																									    self.model_tester						, 'expected_num_hidden_layers'						, self.model_tester.num_hidden_layers + 1       )

																									self.assertEqual(len(__UpperCamelCase       )						, __UpperCamelCase       )
																									self.assertListEqual(
																									    list(output_hidden_states[0].shape[-2:]       )						, [self.model_tester.seq_length, self.model_tester.hidden_size]						, )

																									self.assertEqual(len(__UpperCamelCase       )						, self.model_tester.num_hidden_layers       )
																									self.assertListEqual(
																									    list(output_attentions[0].shape[-3:]       )						, [self.model_tester.num_attention_heads / 2, encoder_seq_length, encoder_key_length]						, )


				@slow
				def __UpperCamelCase						(							self      :  List[str]       )		->     List[Any]:
											A      						=    TFConvBertModel.from_pretrained('YituTech/conv-bert-base'       )
											self.assertIsNotNone(__UpperCamelCase       )


				def __UpperCamelCase						(							self      :  List[Any]       )		->     Dict:
											A						,   A      						=    self.model_tester.prepare_config_and_inputs_for_common()
											A      						=    True
											A      						=    getattr(self.model_tester						, 'decoder_seq_length'						, self.model_tester.seq_length       )
											A      						=    getattr(self.model_tester						, 'encoder_seq_length'						, self.model_tester.seq_length       )
											A      						=    getattr(self.model_tester						, 'key_length'						, __UpperCamelCase       )
											A      						=    getattr(self.model_tester						, 'key_length'						, __UpperCamelCase       )

											def check_decoder_attentions_output(__UpperCamelCase      :  Optional[int]       ):
																		A      						=    len(__UpperCamelCase       )
																		self.assertEqual(out_len % 2						, 0       )
																		A      						=    outputs.decoder_attentions
																		self.assertEqual(len(__UpperCamelCase       )						, self.model_tester.num_hidden_layers       )
																		self.assertListEqual(
																		    list(decoder_attentions[0].shape[-3:]       )						, [self.model_tester.num_attention_heads / 2, decoder_seq_length, decoder_key_length]						, )

											def check_encoder_attentions_output(__UpperCamelCase      :  Optional[Any]       ):
																		A      						=    [
																		    t.numpy() for t in (outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions)
																		]
																		self.assertEqual(len(__UpperCamelCase       )						, self.model_tester.num_hidden_layers       )
																		self.assertListEqual(
																		    list(attentions[0].shape[-3:]       )						, [self.model_tester.num_attention_heads / 2, encoder_seq_length, encoder_key_length]						, )

											for model_class in self.all_model_classes:
																		A      						=    True
																		A      						=    False
																		A      						=    model_class(__UpperCamelCase       )
																		A      						=    model(self._prepare_for_class(__UpperCamelCase						, __UpperCamelCase       )       )
																		A      						=    len(__UpperCamelCase       )
																		self.assertEqual(config.output_hidden_states						, __UpperCamelCase       )
																		check_encoder_attentions_output(__UpperCamelCase       )

																		if self.is_encoder_decoder:
																									A      						=    model_class(__UpperCamelCase       )
																									A      						=    model(self._prepare_for_class(__UpperCamelCase						, __UpperCamelCase       )       )
																									self.assertEqual(config.output_hidden_states						, __UpperCamelCase       )
																									check_decoder_attentions_output(__UpperCamelCase       )

																		# Check that output attentions can also be changed via the config
																		del inputs_dict["output_attentions"]
																		A      						=    True
																		A      						=    model_class(__UpperCamelCase       )
																		A      						=    model(self._prepare_for_class(__UpperCamelCase						, __UpperCamelCase       )       )
																		self.assertEqual(config.output_hidden_states						, __UpperCamelCase       )
																		check_encoder_attentions_output(__UpperCamelCase       )

																		# Check attention is always last and order is fine
																		A      						=    True
																		A      						=    True
																		A      						=    model_class(__UpperCamelCase       )
																		A      						=    model(self._prepare_for_class(__UpperCamelCase						, __UpperCamelCase       )       )

																		self.assertEqual(out_len + (2 if self.is_encoder_decoder else 1)						, len(__UpperCamelCase       )       )
																		self.assertEqual(model.config.output_hidden_states						, __UpperCamelCase       )
																		check_encoder_attentions_output(__UpperCamelCase       )





@require_tf
class        lowerCAmelCase__					(  unittest.TestCase       ):


				@slow
				def __UpperCamelCase						(							self      :  Tuple       )		->     Optional[Any]:
											A      						=    TFConvBertModel.from_pretrained('YituTech/conv-bert-base'       )
											A      						=    tf.constant([[0, 1, 2, 3, 4, 5]]       )
											A      						=    model(__UpperCamelCase       )[0]

											A      						=    [1, 6, 768]
											self.assertEqual(output.shape						, __UpperCamelCase       )

											A      						=    tf.constant(
											    [
											        [
											            [-0.0_3_4_7_5_4_9_3, -0.4_6_8_6_0_3_4, -0.3_0_6_3_8_8_3_2],
											            [0.2_2_6_3_7_2_4_8, -0.2_6_9_8_8_6_4_6, -0.7_4_2_3_4_2_4],
											            [0.1_0_3_2_4_8_6_8, -0.4_5_0_1_3_5_0_8, -0.5_8_2_8_0_7_8_4],
											        ]
											    ]       )
											tf.debugging.assert_near(output[:, :3, :3]						, __UpperCamelCase						, atol=1e-4       )



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   (      A__						: int					)     ->    List[str]:
   lowerCamelCase_						:   int		= 384
   if "tiny" in model_name:
      lowerCamelCase_						:   str		= [3, 3, 9, 3]
      lowerCamelCase_						:   List[str]		= [96, 192, 384, 768]
   if "small" in model_name:
      lowerCamelCase_						:   Any		= [3, 3, 27, 3]
      lowerCamelCase_						:   List[Any]		= [96, 192, 384, 768]
   if "base" in model_name:
      lowerCamelCase_						:   Optional[Any]		= [3, 3, 27, 3]
      lowerCamelCase_						:   int		= [128, 256, 512, 1024]
      lowerCamelCase_						:   int		= 512
   if "large" in model_name:
      lowerCamelCase_						:   Union[str, Any]		= [3, 3, 27, 3]
      lowerCamelCase_						:   Dict		= [192, 384, 768, 1536]
      lowerCamelCase_						:   Union[str, Any]		= 768
   if "xlarge" in model_name:
      lowerCamelCase_						:   Union[str, Any]		= [3, 3, 27, 3]
      lowerCamelCase_						:   Optional[int]		= [256, 512, 1024, 2048]
      lowerCamelCase_						:   List[str]		= 1024

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

   lowerCamelCase_						:   Dict		= ConvNextConfig(
       depths=A__						,						hidden_sizes=A__						,						out_features=["""stage1""", """stage2""", """stage3""", """stage4"""]					)
   lowerCamelCase_						:   str		= UperNetConfig(
       backbone_config=A__						,						auxiliary_in_channels=A__						,						num_labels=A__						,						idalabel=A__						,						labelaid=A__						,						)

   return config


def    __lowerCamelCase   (      A__						: Optional[Any]					)     ->    Optional[Any]:
   lowerCamelCase_						:   Tuple		= []

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


def    __lowerCamelCase   (      A__						: Tuple						,						A__						: str						,						A__						: Any					)     ->    Optional[int]:
   lowerCamelCase_						:   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""",
   }
   lowerCamelCase_						:   List[str]		= model_name_to_url[model_name]
   lowerCamelCase_						:   Dict		= torch.hub.load_state_dict_from_url(A__						,						map_location="""cpu"""					)["""state_dict"""]

   lowerCamelCase_						:   List[Any]		= get_upernet_config(A__					)
   lowerCamelCase_						:   Any		= UperNetForSemanticSegmentation(A__					)
   model.eval()

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

   # rename keys
   lowerCamelCase_						:   List[Any]		= create_rename_keys(A__					)
   for src, dest in rename_keys:
      rename_key(A__						,						A__						,						A__					)

   model.load_state_dict(A__					)

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

   lowerCamelCase_						:   List[str]		= SegformerImageProcessor()
   lowerCamelCase_						:   Union[str, Any]		= processor(A__						,						return_tensors="""pt"""					).pixel_values

   with torch.no_grad():
      lowerCamelCase_						:   List[str]		= model(A__					)

   if model_name == "upernet-convnext-tiny":
      lowerCamelCase_						:   Dict		= 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":
      lowerCamelCase_						:   Optional[Any]		= 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":
      lowerCamelCase_						:   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":
      lowerCamelCase_						:   List[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":
      lowerCamelCase_						:   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]						,						A__						,						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(A__					)
      print(f'''Saving processor to {pytorch_dump_folder_path}'''					)
      processor.save_pretrained(A__					)

   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__":
      snake_case__				:	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.'
      )

      snake_case__				:	List[Any]					       =					parser.parse_args()
      convert_upernet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)





