
import argparse
from collections import OrderedDict
from pathlib import Path

import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from torchvision.transforms import functional as F

from transformers import DetrImageProcessor, TableTransformerConfig, TableTransformerForObjectDetection
from transformers.utils import logging


logging.set_verbosity_info()
A		: Any								=     logging.get_logger(__name__)

# here we list all keys to be renamed (original name on the left, our name on the right)
A		: str								=     []
for i in range(6):
		# encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
		rename_keys.append(
		    (F'transformer.encoder.layers.{i}.self_attn.out_proj.weight', F'encoder.layers.{i}.self_attn.out_proj.weight')
		)
		rename_keys.append(
		    (F'transformer.encoder.layers.{i}.self_attn.out_proj.bias', F'encoder.layers.{i}.self_attn.out_proj.bias')
		)
		rename_keys.append((F'transformer.encoder.layers.{i}.linear1.weight', F'encoder.layers.{i}.fc1.weight'))
		rename_keys.append((F'transformer.encoder.layers.{i}.linear1.bias', F'encoder.layers.{i}.fc1.bias'))
		rename_keys.append((F'transformer.encoder.layers.{i}.linear2.weight', F'encoder.layers.{i}.fc2.weight'))
		rename_keys.append((F'transformer.encoder.layers.{i}.linear2.bias', F'encoder.layers.{i}.fc2.bias'))
		rename_keys.append(
		    (F'transformer.encoder.layers.{i}.norm1.weight', F'encoder.layers.{i}.self_attn_layer_norm.weight')
		)
		rename_keys.append((F'transformer.encoder.layers.{i}.norm1.bias', F'encoder.layers.{i}.self_attn_layer_norm.bias'))
		rename_keys.append((F'transformer.encoder.layers.{i}.norm2.weight', F'encoder.layers.{i}.final_layer_norm.weight'))
		rename_keys.append((F'transformer.encoder.layers.{i}.norm2.bias', F'encoder.layers.{i}.final_layer_norm.bias'))
		# decoder layers: 2 times output projection, 2 feedforward neural networks and 3 layernorms
		rename_keys.append(
		    (F'transformer.decoder.layers.{i}.self_attn.out_proj.weight', F'decoder.layers.{i}.self_attn.out_proj.weight')
		)
		rename_keys.append(
		    (F'transformer.decoder.layers.{i}.self_attn.out_proj.bias', F'decoder.layers.{i}.self_attn.out_proj.bias')
		)
		rename_keys.append(
		    (
		        F'transformer.decoder.layers.{i}.multihead_attn.out_proj.weight',
		        F'decoder.layers.{i}.encoder_attn.out_proj.weight',
		    )
		)
		rename_keys.append(
		    (
		        F'transformer.decoder.layers.{i}.multihead_attn.out_proj.bias',
		        F'decoder.layers.{i}.encoder_attn.out_proj.bias',
		    )
		)
		rename_keys.append((F'transformer.decoder.layers.{i}.linear1.weight', F'decoder.layers.{i}.fc1.weight'))
		rename_keys.append((F'transformer.decoder.layers.{i}.linear1.bias', F'decoder.layers.{i}.fc1.bias'))
		rename_keys.append((F'transformer.decoder.layers.{i}.linear2.weight', F'decoder.layers.{i}.fc2.weight'))
		rename_keys.append((F'transformer.decoder.layers.{i}.linear2.bias', F'decoder.layers.{i}.fc2.bias'))
		rename_keys.append(
		    (F'transformer.decoder.layers.{i}.norm1.weight', F'decoder.layers.{i}.self_attn_layer_norm.weight')
		)
		rename_keys.append((F'transformer.decoder.layers.{i}.norm1.bias', F'decoder.layers.{i}.self_attn_layer_norm.bias'))
		rename_keys.append(
		    (F'transformer.decoder.layers.{i}.norm2.weight', F'decoder.layers.{i}.encoder_attn_layer_norm.weight')
		)
		rename_keys.append(
		    (F'transformer.decoder.layers.{i}.norm2.bias', F'decoder.layers.{i}.encoder_attn_layer_norm.bias')
		)
		rename_keys.append((F'transformer.decoder.layers.{i}.norm3.weight', F'decoder.layers.{i}.final_layer_norm.weight'))
		rename_keys.append((F'transformer.decoder.layers.{i}.norm3.bias', F'decoder.layers.{i}.final_layer_norm.bias'))

# convolutional projection + query embeddings + layernorm of encoder + layernorm of decoder + class and bounding box heads
rename_keys.extend(
    [
        ('''input_proj.weight''', '''input_projection.weight'''),
        ('''input_proj.bias''', '''input_projection.bias'''),
        ('''query_embed.weight''', '''query_position_embeddings.weight'''),
        ('''transformer.encoder.norm.weight''', '''encoder.layernorm.weight'''),
        ('''transformer.encoder.norm.bias''', '''encoder.layernorm.bias'''),
        ('''transformer.decoder.norm.weight''', '''decoder.layernorm.weight'''),
        ('''transformer.decoder.norm.bias''', '''decoder.layernorm.bias'''),
        ('''class_embed.weight''', '''class_labels_classifier.weight'''),
        ('''class_embed.bias''', '''class_labels_classifier.bias'''),
        ('''bbox_embed.layers.0.weight''', '''bbox_predictor.layers.0.weight'''),
        ('''bbox_embed.layers.0.bias''', '''bbox_predictor.layers.0.bias'''),
        ('''bbox_embed.layers.1.weight''', '''bbox_predictor.layers.1.weight'''),
        ('''bbox_embed.layers.1.bias''', '''bbox_predictor.layers.1.bias'''),
        ('''bbox_embed.layers.2.weight''', '''bbox_predictor.layers.2.weight'''),
        ('''bbox_embed.layers.2.bias''', '''bbox_predictor.layers.2.bias'''),
    ]
)




def 	UpperCamelCase__  (     SCREAMING_SNAKE_CASE_	:	Any		,   SCREAMING_SNAKE_CASE_	:	Union[str, Any]		,   SCREAMING_SNAKE_CASE_	:	Optional[int]       )  ->     Any:
						_lowercase					       =					state_dict.pop(SCREAMING_SNAKE_CASE_       )
						_lowercase					       =					val




def 	UpperCamelCase__  (     SCREAMING_SNAKE_CASE_	:	Optional[int]       )  ->     str:
						_lowercase					       =					OrderedDict()
						for key, value in state_dict.items():
												if "backbone.0.body" in key:
																		_lowercase					       =					key.replace("""backbone.0.body"""		,   """backbone.conv_encoder.model"""       )
																		_lowercase					       =					value
												else:
																		_lowercase					       =					value

						return new_state_dict




def 	UpperCamelCase__  (     SCREAMING_SNAKE_CASE_	:	int       )  ->     Any:
						_lowercase					       =					""""""

						# first: transformer encoder
						for i in range(6       ):
												# read in weights + bias of input projection layer (in PyTorch's MultiHeadAttention, this is a single matrix + bias)
												_lowercase					       =					state_dict.pop(f"""{prefix}transformer.encoder.layers.{i}.self_attn.in_proj_weight"""       )
												_lowercase					       =					state_dict.pop(f"""{prefix}transformer.encoder.layers.{i}.self_attn.in_proj_bias"""       )
												# next, add query, keys and values (in that order) to the state dict
												_lowercase					       =					in_proj_weight[:2_56, :]
												_lowercase					       =					in_proj_bias[:2_56]
												_lowercase					       =					in_proj_weight[2_56:5_12, :]
												_lowercase					       =					in_proj_bias[2_56:5_12]
												_lowercase					       =					in_proj_weight[-2_56:, :]
												_lowercase					       =					in_proj_bias[-2_56:]
						# next: transformer decoder (which is a bit more complex because it also includes cross-attention)
						for i in range(6       ):
												# read in weights + bias of input projection layer of self-attention
												_lowercase					       =					state_dict.pop(f"""{prefix}transformer.decoder.layers.{i}.self_attn.in_proj_weight"""       )
												_lowercase					       =					state_dict.pop(f"""{prefix}transformer.decoder.layers.{i}.self_attn.in_proj_bias"""       )
												# next, add query, keys and values (in that order) to the state dict
												_lowercase					       =					in_proj_weight[:2_56, :]
												_lowercase					       =					in_proj_bias[:2_56]
												_lowercase					       =					in_proj_weight[2_56:5_12, :]
												_lowercase					       =					in_proj_bias[2_56:5_12]
												_lowercase					       =					in_proj_weight[-2_56:, :]
												_lowercase					       =					in_proj_bias[-2_56:]
												# read in weights + bias of input projection layer of cross-attention
												_lowercase					       =					state_dict.pop(
												    f"""{prefix}transformer.decoder.layers.{i}.multihead_attn.in_proj_weight"""       )
												_lowercase					       =					state_dict.pop(f"""{prefix}transformer.decoder.layers.{i}.multihead_attn.in_proj_bias"""       )
												# next, add query, keys and values (in that order) of cross-attention to the state dict
												_lowercase					       =					in_proj_weight_cross_attn[:2_56, :]
												_lowercase					       =					in_proj_bias_cross_attn[:2_56]
												_lowercase					       =					in_proj_weight_cross_attn[2_56:5_12, :]
												_lowercase					       =					in_proj_bias_cross_attn[2_56:5_12]
												_lowercase					       =					in_proj_weight_cross_attn[-2_56:, :]
												_lowercase					       =					in_proj_bias_cross_attn[-2_56:]




def 	UpperCamelCase__  (     SCREAMING_SNAKE_CASE_	:	List[Any]		,   SCREAMING_SNAKE_CASE_	:	str       )  ->     Union[str, Any]:
						_lowercase       ,						_lowercase					       =					image.size
						_lowercase					       =					max(SCREAMING_SNAKE_CASE_		,   SCREAMING_SNAKE_CASE_       )
						_lowercase					       =					8_00 if """detection""" in checkpoint_url else 10_00
						_lowercase					       =					target_max_size / current_max_size
						_lowercase					       =					image.resize((int(round(scale * width       )       ), int(round(scale * height       )       ))       )

						return resized_image




def 	UpperCamelCase__  (     SCREAMING_SNAKE_CASE_	:	str       )  ->     str:
						_lowercase					       =					F.to_tensor(SCREAMING_SNAKE_CASE_       )
						_lowercase					       =					F.normalize(SCREAMING_SNAKE_CASE_		,   mean=[0.4_85, 0.4_56, 0.4_06]		,   std=[0.2_29, 0.2_24, 0.2_25]       )
						return image




@torch.no_grad()
def 	UpperCamelCase__  (     SCREAMING_SNAKE_CASE_	:	Tuple		,   SCREAMING_SNAKE_CASE_	:	int		,   SCREAMING_SNAKE_CASE_	:	List[str]       )  ->     Union[str, Any]:
						logger.info("""Converting model..."""       )

						# load original state dict
						_lowercase					       =					torch.hub.load_state_dict_from_url(SCREAMING_SNAKE_CASE_		,   map_location="""cpu"""       )
						# rename keys
						for src, dest in rename_keys:
												rename_key(SCREAMING_SNAKE_CASE_		,   SCREAMING_SNAKE_CASE_		,   SCREAMING_SNAKE_CASE_       )
						_lowercase					       =					rename_backbone_keys(SCREAMING_SNAKE_CASE_       )
						# query, key and value matrices need special treatment
						read_in_q_k_v(SCREAMING_SNAKE_CASE_       )
						# important: we need to prepend a prefix to each of the base model keys as the head models use different attributes for them
						_lowercase					       =					"""model."""
						for key in state_dict.copy().keys():
												if not key.startswith("""class_labels_classifier"""       ) and not key.startswith("""bbox_predictor"""       ):
																		_lowercase					       =					state_dict.pop(SCREAMING_SNAKE_CASE_       )
																		_lowercase					       =					val
    # create HuggingFace model and load state dict
						_lowercase					       =					TableTransformerConfig(
						    backbone="""resnet18"""		,   mask_loss_coefficient=1		,   dice_loss_coefficient=1		,   ce_loss_coefficient=1		,   bbox_loss_coefficient=5		,   giou_loss_coefficient=2		,   eos_coefficient=0.4		,   class_cost=1		,   bbox_cost=5		,   giou_cost=2		,   )

						if "detection" in checkpoint_url:
												_lowercase					       =					15
												_lowercase					       =					2
												_lowercase					       =					{0: """table""", 1: """table rotated"""}
												_lowercase					       =					idalabel
												_lowercase					       =					{v: k for k, v in idalabel.items()}
						else:
												_lowercase					       =					1_25
												_lowercase					       =					6
												_lowercase					       =					{
												    0: """table""",
												    1: """table column""",
												    2: """table row""",
												    3: """table column header""",
												    4: """table projected row header""",
												    5: """table spanning cell""",
												}
												_lowercase					       =					idalabel
												_lowercase					       =					{v: k for k, v in idalabel.items()}

						_lowercase					       =					DetrImageProcessor(
						    format="""coco_detection"""		,   max_size=8_00 if """detection""" in checkpoint_url else 10_00       )
						_lowercase					       =					TableTransformerForObjectDetection(SCREAMING_SNAKE_CASE_       )
						model.load_state_dict(SCREAMING_SNAKE_CASE_       )
						model.eval()

						# verify our conversion
						_lowercase					       =					"""example_pdf.png""" if """detection""" in checkpoint_url else """example_table.png"""
						_lowercase					       =					hf_hub_download(repo_id="""nielsr/example-pdf"""		,   repo_type="""dataset"""		,   filename=SCREAMING_SNAKE_CASE_       )
						_lowercase					       =					Image.open(SCREAMING_SNAKE_CASE_       ).convert("""RGB"""       )
						_lowercase					       =					normalize(resize(SCREAMING_SNAKE_CASE_		,   SCREAMING_SNAKE_CASE_       )       ).unsqueeze(0       )

						_lowercase					       =					model(SCREAMING_SNAKE_CASE_       )

						if "detection" in checkpoint_url:
												_lowercase					       =					(1, 15, 3)
												_lowercase					       =					torch.tensor(
												    [[-6.78_97, -16.99_85, 6.79_37], [-8.01_86, -22.21_92, 6.96_77], [-7.31_17, -21.07_08, 7.40_55]]       )
												_lowercase					       =					torch.tensor([[0.48_67, 0.17_67, 0.67_32], [0.67_18, 0.44_79, 0.38_30], [0.47_16, 0.17_60, 0.63_64]]       )

						else:
												_lowercase					       =					(1, 1_25, 7)
												_lowercase					       =					torch.tensor(
												    [[-18.14_30, -8.32_14, 4.82_74], [-18.46_85, -7.13_61, -4.26_67], [-26.36_93, -9.34_29, -4.99_62]]       )
												_lowercase					       =					torch.tensor([[0.49_83, 0.55_95, 0.94_40], [0.49_16, 0.63_15, 0.59_54], [0.61_08, 0.86_37, 0.11_35]]       )

						assert outputs.logits.shape == expected_shape
						assert torch.allclose(outputs.logits[0, :3, :3]		,   SCREAMING_SNAKE_CASE_		,   atol=1E-4       )
						assert torch.allclose(outputs.pred_boxes[0, :3, :3]		,   SCREAMING_SNAKE_CASE_		,   atol=1E-4       )
						print("""Looks ok!"""       )

						if pytorch_dump_folder_path is not None:
												# Save model and image processor
												logger.info(f"""Saving PyTorch model and image processor to {pytorch_dump_folder_path}..."""       )
												Path(SCREAMING_SNAKE_CASE_       ).mkdir(exist_ok=SCREAMING_SNAKE_CASE_       )
												model.save_pretrained(SCREAMING_SNAKE_CASE_       )
												image_processor.save_pretrained(SCREAMING_SNAKE_CASE_       )

						if push_to_hub:
												# Push model to HF hub
												logger.info("""Pushing model to the hub..."""       )
												_lowercase					       =					(
												    """microsoft/table-transformer-detection"""
												    if """detection""" in checkpoint_url
												    else """microsoft/table-transformer-structure-recognition"""
												)
												model.push_to_hub(SCREAMING_SNAKE_CASE_       )
												image_processor.push_to_hub(SCREAMING_SNAKE_CASE_       )


if __name__ == "__main__":
		A		: Tuple								=     argparse.ArgumentParser()

		parser.add_argument(
		    '''--checkpoint_url''',
		    default='''https://pubtables1m.blob.core.windows.net/model/pubtables1m_detection_detr_r18.pth''',
		    type=str,
		    choices=[
		        '''https://pubtables1m.blob.core.windows.net/model/pubtables1m_detection_detr_r18.pth''',
		        '''https://pubtables1m.blob.core.windows.net/model/pubtables1m_structure_detr_r18.pth''',
		    ],
		    help='''URL of the Table Transformer checkpoint you\'d like to convert.''',
		)
		parser.add_argument(
		    '''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the folder to output PyTorch model.'''
		)
		parser.add_argument(
		    '''--push_to_hub''', action='''store_true''', help='''Whether or not to push the converted model to the 🤗 hub.'''
		)
		A		: List[str]								=     parser.parse_args()
		convert_table_transformer_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.push_to_hub)

import inspect
import unittest

from huggingface_hub import hf_hub_download

from transformers import ConvNextConfig, UperNetConfig
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
		import torch

		from transformers import UperNetForSemanticSegmentation
		from transformers.models.upernet.modeling_upernet import UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
		from PIL import Image

		from transformers import AutoImageProcessor



class    a_ :






			def __init__(      self				,			__UpperCamelCase				,			__UpperCamelCase=13				,			__UpperCamelCase=32				,			__UpperCamelCase=3				,			__UpperCamelCase=4				,			__UpperCamelCase=[10, 20, 30, 40]				,			__UpperCamelCase=[2, 2, 3, 2]				,			__UpperCamelCase=True				,			__UpperCamelCase=True				,			__UpperCamelCase=37				,			__UpperCamelCase="gelu"				,			__UpperCamelCase=10				,			__UpperCamelCase=0.02				,			__UpperCamelCase=["stage2", "stage3", "stage4"]				,			__UpperCamelCase=3				,			__UpperCamelCase=None				,			):
									_lowercase					       =					parent
									_lowercase					       =					batch_size
									_lowercase					       =					image_size
									_lowercase					       =					num_channels
									_lowercase					       =					num_stages
									_lowercase					       =					hidden_sizes
									_lowercase					       =					depths
									_lowercase					       =					is_training
									_lowercase					       =					use_labels
									_lowercase					       =					intermediate_size
									_lowercase					       =					hidden_act
									_lowercase					       =					type_sequence_label_size
									_lowercase					       =					initializer_range
									_lowercase					       =					out_features
									_lowercase					       =					num_labels
									_lowercase					       =					scope
									_lowercase					       =					num_stages






			def    UpperCamelCase_				(      self  ):
									_lowercase					       =					floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]  )

									_lowercase					       =					None
									if self.use_labels:
															_lowercase					       =					ids_tensor([self.batch_size]				,			self.type_sequence_label_size  )

									_lowercase					       =					self.get_config()

									return config, pixel_values, labels






			def    UpperCamelCase_				(      self  ):
									return ConvNextConfig(
									    num_channels=self.num_channels				,			num_stages=self.num_stages				,			hidden_sizes=self.hidden_sizes				,			depths=self.depths				,			is_training=self.is_training				,			intermediate_size=self.intermediate_size				,			hidden_act=self.hidden_act				,			out_features=self.out_features				,			)






			def    UpperCamelCase_				(      self  ):
									return UperNetConfig(
									    backbone_config=self.get_backbone_config()				,			hidden_size=512				,			pool_scales=[1, 2, 3, 6]				,			use_auxiliary_head=__UpperCamelCase				,			auxiliary_loss_weight=0.4				,			auxiliary_in_channels=40				,			auxiliary_channels=256				,			auxiliary_num_convs=1				,			auxiliary_concat_input=__UpperCamelCase				,			loss_ignore_index=255				,			num_labels=self.num_labels				,			)






			def    UpperCamelCase_				(      self				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase  ):
									_lowercase					       =					UperNetForSemanticSegmentation(config=__UpperCamelCase  )
									model.to(__UpperCamelCase  )
									model.eval()
									_lowercase					       =					model(__UpperCamelCase  )
									self.parent.assertEqual(
									    result.logits.shape				,			(self.batch_size, self.num_labels, self.image_size, self.image_size)  )






			def    UpperCamelCase_				(      self  ):
									_lowercase					       =					self.prepare_config_and_inputs()
									(
									    (
									    _lowercase
									)       ,						(
									    _lowercase
									)       ,						(
									    _lowercase
									)       ,						
									)					       =					config_and_inputs
									_lowercase					       =					{"""pixel_values""": pixel_values}
									return config, inputs_dict







@require_torch
class    a_ (			_a					,			_a					,			unittest.TestCase    ):
			a	:			Optional[int]					   =							(UperNetForSemanticSegmentation,) if is_torch_available() else ()
			a	:			Optional[Any]					   =							{'''image-segmentation''': UperNetForSemanticSegmentation} if is_torch_available() else {}
			a	:			Any					   =							False
			a	:			Union[str, Any]					   =							False
			a	:			Optional[Any]					   =							False
			a	:			int					   =							False
			a	:			Optional[int]					   =							False
			a	:			List[str]					   =							False






			def    UpperCamelCase_				(      self  ):
									_lowercase					       =					UperNetModelTester(self  )
									_lowercase					       =					ConfigTester(self				,			config_class=__UpperCamelCase				,			has_text_modality=__UpperCamelCase				,			hidden_size=37  )






			def    UpperCamelCase_				(      self  ):
									self.create_and_test_config_common_properties()
									self.config_tester.create_and_test_config_to_json_string()
									self.config_tester.create_and_test_config_to_json_file()
									self.config_tester.create_and_test_config_from_and_save_pretrained()
									self.config_tester.create_and_test_config_with_num_labels()
									self.config_tester.check_config_can_be_init_without_params()
									self.config_tester.check_config_arguments_init()






			def    UpperCamelCase_				(      self  ):
									return






			def    UpperCamelCase_				(      self  ):
									_lowercase       ,						_lowercase					       =					self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes:
															_lowercase					       =					model_class(__UpperCamelCase  )
															_lowercase					       =					inspect.signature(model.forward  )
															# signature.parameters is an OrderedDict => so arg_names order is deterministic
															_lowercase					       =					[*signature.parameters.keys()]

															_lowercase					       =					["""pixel_values"""]
															self.assertListEqual(arg_names[:1]				,			__UpperCamelCase  )






			def    UpperCamelCase_				(      self  ):
									_lowercase					       =					self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_semantic_segmentation(*__UpperCamelCase  )






			@unittest.skip(reason="""UperNet does not use inputs_embeds"""  )
			def    UpperCamelCase_				(      self  ):
									pass






			@unittest.skip(reason="""UperNet does not support input and output embeddings"""  )
			def    UpperCamelCase_				(      self  ):
									pass






			@unittest.skip(reason="""UperNet does not have a base model"""  )
			def    UpperCamelCase_				(      self  ):
									pass






			@unittest.skip(reason="""UperNet does not have a base model"""  )
			def    UpperCamelCase_				(      self  ):
									pass






			@require_torch_multi_gpu
			@unittest.skip(reason="""UperNet has some layers using `add_module` which doesn't work well with `nn.DataParallel`"""  )
			def    UpperCamelCase_				(      self  ):
									pass






			@unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""  )
			def    UpperCamelCase_				(      self  ):
									pass






			def    UpperCamelCase_				(      self  ):
									def check_hidden_states_output(__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase  ):
															_lowercase					       =					model_class(__UpperCamelCase  )
															model.to(__UpperCamelCase  )
															model.eval()

															with torch.no_grad():
																					_lowercase					       =					model(**self._prepare_for_class(__UpperCamelCase				,			__UpperCamelCase  )  )

															_lowercase					       =					outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

															_lowercase					       =					self.model_tester.num_stages
															self.assertEqual(len(__UpperCamelCase  )				,			expected_num_stages + 1  )

															# ConvNext's feature maps are of shape (batch_size, num_channels, height, width)
															self.assertListEqual(
															    list(hidden_states[0].shape[-2:]  )				,			[self.model_tester.image_size // 4, self.model_tester.image_size // 4]				,			)

									_lowercase       ,						_lowercase					       =					self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes:
															_lowercase					       =					True
															check_hidden_states_output(__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase  )

															# check that output_hidden_states also work using config
															del inputs_dict["output_hidden_states"]
															_lowercase					       =					True

															check_hidden_states_output(__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase  )






			def    UpperCamelCase_				(      self  ):
									_lowercase       ,						_lowercase					       =					self.model_tester.prepare_config_and_inputs_for_common()

									_lowercase					       =					_config_zero_init(__UpperCamelCase  )
									_lowercase					       =					_config_zero_init(configs_no_init.backbone_config  )
									for model_class in self.all_model_classes:
															_lowercase					       =					model_class(config=__UpperCamelCase  )
															for name, param in model.named_parameters():
																					if param.requires_grad:
																											self.assertIn(
																											    ((param.data.mean() * 1E9).round() / 1E9).item()				,			[0.0, 1.0]				,			msg=f"""Parameter {name} of model {model_class} seems not properly initialized"""				,			)






			@unittest.skip(reason="""UperNet does not have tied weights"""  )
			def    UpperCamelCase_				(      self  ):
									pass






			@slow
			def    UpperCamelCase_				(      self  ):
									for model_name in UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															_lowercase					       =					UperNetForSemanticSegmentation.from_pretrained(__UpperCamelCase  )
															self.assertIsNotNone(__UpperCamelCase  )







def 	UpperCamelCase__  (     )  ->     str:
						_lowercase					       =					hf_hub_download(
						    repo_id="""hf-internal-testing/fixtures_ade20k"""		,   repo_type="""dataset"""		,   filename="""ADE_val_00000001.jpg"""       )
						_lowercase					       =					Image.open(SCREAMING_SNAKE_CASE_       ).convert("""RGB"""       )
						return image



@require_torch
@require_vision
@slow
class    a_ (			unittest.TestCase    ):






			def    UpperCamelCase_				(      self  ):
									_lowercase					       =					AutoImageProcessor.from_pretrained("""openmmlab/upernet-swin-tiny"""  )
									_lowercase					       =					UperNetForSemanticSegmentation.from_pretrained("""openmmlab/upernet-swin-tiny"""  ).to(__UpperCamelCase  )

									_lowercase					       =					prepare_img()
									_lowercase					       =					processor(images=__UpperCamelCase				,			return_tensors="""pt"""  ).to(__UpperCamelCase  )

									with torch.no_grad():
															_lowercase					       =					model(**__UpperCamelCase  )

									_lowercase					       =					torch.Size((1, model.config.num_labels, 512, 512)  )
									self.assertEqual(outputs.logits.shape				,			__UpperCamelCase  )

									_lowercase					       =					torch.tensor(
									    [[-7.59_58, -7.59_58, -7.43_02], [-7.59_58, -7.59_58, -7.43_02], [-7.47_97, -7.47_97, -7.30_68]]  ).to(__UpperCamelCase  )
									self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3]				,			__UpperCamelCase				,			atol=1E-4  )  )






			def    UpperCamelCase_				(      self  ):
									_lowercase					       =					AutoImageProcessor.from_pretrained("""openmmlab/upernet-convnext-tiny"""  )
									_lowercase					       =					UperNetForSemanticSegmentation.from_pretrained("""openmmlab/upernet-convnext-tiny"""  ).to(__UpperCamelCase  )

									_lowercase					       =					prepare_img()
									_lowercase					       =					processor(images=__UpperCamelCase				,			return_tensors="""pt"""  ).to(__UpperCamelCase  )

									with torch.no_grad():
															_lowercase					       =					model(**__UpperCamelCase  )

									_lowercase					       =					torch.Size((1, model.config.num_labels, 512, 512)  )
									self.assertEqual(outputs.logits.shape				,			__UpperCamelCase  )

									_lowercase					       =					torch.tensor(
									    [[-8.81_10, -8.81_10, -8.65_21], [-8.81_10, -8.81_10, -8.65_21], [-8.77_46, -8.77_46, -8.61_30]]  ).to(__UpperCamelCase  )
									self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3]				,			__UpperCamelCase				,			atol=1E-4  )  )