





'''simple docstring'''







import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import cached_download, hf_hub_download, hf_hub_url
from PIL import Image

from transformers import DetaConfig, DetaForObjectDetection, DetaImageProcessor
from transformers.utils import logging


logging.set_verbosity_info()
snake_case_       :    Optional[Any]          =					logging.get_logger(__name__)





def     lowerCamelCase_ (					)	->		Dict:
						UpperCAmelCase_     : Dict         =					DetaConfig(
						    num_queries=900,     encoder_ffn_dim=2048,     decoder_ffn_dim=2048,     num_feature_levels=5,     assign_first_stage=_lowerCamelCase,     with_box_refine=_lowerCamelCase,     two_stage=_lowerCamelCase,     )

						# set labels
						UpperCAmelCase_     : List[str]         =					91
						UpperCAmelCase_     : Optional[int]         =					"""huggingface/label-files"""
						UpperCAmelCase_     : Union[str, Any]         =					"""coco-detection-id2label.json"""
						UpperCAmelCase_     : List[Any]         =					json.load(open(cached_download(hf_hub_url(_lowerCamelCase,     _lowerCamelCase,     repo_type='''dataset'''				)				),     '''r'''				)				)
						UpperCAmelCase_     : Tuple         =					{int(_lowerCamelCase				): v for k, v in idalabel.items()}
						UpperCAmelCase_     : List[Any]         =					idalabel
						UpperCAmelCase_     : Dict         =					{v: k for k, v in idalabel.items()}

						return config





def     lowerCamelCase_ (					SCREAMING_SNAKE_CASE__ :     List[str]				)	->		Tuple:
						UpperCAmelCase_     : Optional[int]         =					[]

						# stem
						# fmt: off
						rename_keys.append(('''backbone.0.body.conv1.weight''', '''model.backbone.model.embedder.embedder.convolution.weight''')				)
						rename_keys.append(('''backbone.0.body.bn1.weight''', '''model.backbone.model.embedder.embedder.normalization.weight''')				)
						rename_keys.append(('''backbone.0.body.bn1.bias''', '''model.backbone.model.embedder.embedder.normalization.bias''')				)
						rename_keys.append(('''backbone.0.body.bn1.running_mean''', '''model.backbone.model.embedder.embedder.normalization.running_mean''')				)
						rename_keys.append(('''backbone.0.body.bn1.running_var''', '''model.backbone.model.embedder.embedder.normalization.running_var''')				)
						# stages
						for stage_idx in range(len(config.backbone_config.depths				)				):
												for layer_idx in range(config.backbone_config.depths[stage_idx]				):
																		# shortcut
																		if layer_idx == 0:
																								rename_keys.append(
																								    (
																								        F"""backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.0.weight""",
																								        F"""model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.convolution.weight""",
																								    )				)
																								rename_keys.append(
																								    (
																								        F"""backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.weight""",
																								        F"""model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.weight""",
																								    )				)
																								rename_keys.append(
																								    (
																								        F"""backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.bias""",
																								        F"""model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.bias""",
																								    )				)
																								rename_keys.append(
																								    (
																								        F"""backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.running_mean""",
																								        F"""model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.running_mean""",
																								    )				)
																								rename_keys.append(
																								    (
																								        F"""backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.running_var""",
																								        F"""model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.running_var""",
																								    )				)
																		# 3 convs
																		for i in range(3				):
																								rename_keys.append(
																								    (
																								        F"""backbone.0.body.layer{stage_idx + 1}.{layer_idx}.conv{i+1}.weight""",
																								        F"""model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.convolution.weight""",
																								    )				)
																								rename_keys.append(
																								    (
																								        F"""backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.weight""",
																								        F"""model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.weight""",
																								    )				)
																								rename_keys.append(
																								    (
																								        F"""backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.bias""",
																								        F"""model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.bias""",
																								    )				)
																								rename_keys.append(
																								    (
																								        F"""backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.running_mean""",
																								        F"""model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.running_mean""",
																								    )				)
																								rename_keys.append(
																								    (
																								        F"""backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.running_var""",
																								        F"""model.backbone.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.running_var""",
																								    )				)
												# transformer encoder
						for i in range(config.encoder_layers				):
												rename_keys.append((F"""transformer.encoder.layers.{i}.self_attn.sampling_offsets.weight""", F"""model.encoder.layers.{i}.self_attn.sampling_offsets.weight""")				)
												rename_keys.append((F"""transformer.encoder.layers.{i}.self_attn.sampling_offsets.bias""", F"""model.encoder.layers.{i}.self_attn.sampling_offsets.bias""")				)
												rename_keys.append((F"""transformer.encoder.layers.{i}.self_attn.attention_weights.weight""", F"""model.encoder.layers.{i}.self_attn.attention_weights.weight""")				)
												rename_keys.append((F"""transformer.encoder.layers.{i}.self_attn.attention_weights.bias""", F"""model.encoder.layers.{i}.self_attn.attention_weights.bias""")				)
												rename_keys.append((F"""transformer.encoder.layers.{i}.self_attn.value_proj.weight""", F"""model.encoder.layers.{i}.self_attn.value_proj.weight""")				)
												rename_keys.append((F"""transformer.encoder.layers.{i}.self_attn.value_proj.bias""", F"""model.encoder.layers.{i}.self_attn.value_proj.bias""")				)
												rename_keys.append((F"""transformer.encoder.layers.{i}.self_attn.output_proj.weight""", F"""model.encoder.layers.{i}.self_attn.output_proj.weight""")				)
												rename_keys.append((F"""transformer.encoder.layers.{i}.self_attn.output_proj.bias""", F"""model.encoder.layers.{i}.self_attn.output_proj.bias""")				)
												rename_keys.append((F"""transformer.encoder.layers.{i}.norm1.weight""", F"""model.encoder.layers.{i}.self_attn_layer_norm.weight""")				)
												rename_keys.append((F"""transformer.encoder.layers.{i}.norm1.bias""", F"""model.encoder.layers.{i}.self_attn_layer_norm.bias""")				)
												rename_keys.append((F"""transformer.encoder.layers.{i}.linear1.weight""", F"""model.encoder.layers.{i}.fc1.weight""")				)
												rename_keys.append((F"""transformer.encoder.layers.{i}.linear1.bias""", F"""model.encoder.layers.{i}.fc1.bias""")				)
												rename_keys.append((F"""transformer.encoder.layers.{i}.linear2.weight""", F"""model.encoder.layers.{i}.fc2.weight""")				)
												rename_keys.append((F"""transformer.encoder.layers.{i}.linear2.bias""", F"""model.encoder.layers.{i}.fc2.bias""")				)
												rename_keys.append((F"""transformer.encoder.layers.{i}.norm2.weight""", F"""model.encoder.layers.{i}.final_layer_norm.weight""")				)
												rename_keys.append((F"""transformer.encoder.layers.{i}.norm2.bias""", F"""model.encoder.layers.{i}.final_layer_norm.bias""")				)

						# transformer decoder
						for i in range(config.decoder_layers				):
												rename_keys.append((F"""transformer.decoder.layers.{i}.cross_attn.sampling_offsets.weight""", F"""model.decoder.layers.{i}.encoder_attn.sampling_offsets.weight""")				)
												rename_keys.append((F"""transformer.decoder.layers.{i}.cross_attn.sampling_offsets.bias""", F"""model.decoder.layers.{i}.encoder_attn.sampling_offsets.bias""")				)
												rename_keys.append((F"""transformer.decoder.layers.{i}.cross_attn.attention_weights.weight""", F"""model.decoder.layers.{i}.encoder_attn.attention_weights.weight""")				)
												rename_keys.append((F"""transformer.decoder.layers.{i}.cross_attn.attention_weights.bias""", F"""model.decoder.layers.{i}.encoder_attn.attention_weights.bias""")				)
												rename_keys.append((F"""transformer.decoder.layers.{i}.cross_attn.value_proj.weight""", F"""model.decoder.layers.{i}.encoder_attn.value_proj.weight""")				)
												rename_keys.append((F"""transformer.decoder.layers.{i}.cross_attn.value_proj.bias""", F"""model.decoder.layers.{i}.encoder_attn.value_proj.bias""")				)
												rename_keys.append((F"""transformer.decoder.layers.{i}.cross_attn.output_proj.weight""", F"""model.decoder.layers.{i}.encoder_attn.output_proj.weight""")				)
												rename_keys.append((F"""transformer.decoder.layers.{i}.cross_attn.output_proj.bias""", F"""model.decoder.layers.{i}.encoder_attn.output_proj.bias""")				)
												rename_keys.append((F"""transformer.decoder.layers.{i}.norm1.weight""", F"""model.decoder.layers.{i}.encoder_attn_layer_norm.weight""")				)
												rename_keys.append((F"""transformer.decoder.layers.{i}.norm1.bias""", F"""model.decoder.layers.{i}.encoder_attn_layer_norm.bias""")				)
												rename_keys.append((F"""transformer.decoder.layers.{i}.self_attn.out_proj.weight""", F"""model.decoder.layers.{i}.self_attn.out_proj.weight""")				)
												rename_keys.append((F"""transformer.decoder.layers.{i}.self_attn.out_proj.bias""", F"""model.decoder.layers.{i}.self_attn.out_proj.bias""")				)
												rename_keys.append((F"""transformer.decoder.layers.{i}.norm2.weight""", F"""model.decoder.layers.{i}.self_attn_layer_norm.weight""")				)
												rename_keys.append((F"""transformer.decoder.layers.{i}.norm2.bias""", F"""model.decoder.layers.{i}.self_attn_layer_norm.bias""")				)
												rename_keys.append((F"""transformer.decoder.layers.{i}.linear1.weight""", F"""model.decoder.layers.{i}.fc1.weight""")				)
												rename_keys.append((F"""transformer.decoder.layers.{i}.linear1.bias""", F"""model.decoder.layers.{i}.fc1.bias""")				)
												rename_keys.append((F"""transformer.decoder.layers.{i}.linear2.weight""", F"""model.decoder.layers.{i}.fc2.weight""")				)
												rename_keys.append((F"""transformer.decoder.layers.{i}.linear2.bias""", F"""model.decoder.layers.{i}.fc2.bias""")				)
												rename_keys.append((F"""transformer.decoder.layers.{i}.norm3.weight""", F"""model.decoder.layers.{i}.final_layer_norm.weight""")				)
												rename_keys.append((F"""transformer.decoder.layers.{i}.norm3.bias""", F"""model.decoder.layers.{i}.final_layer_norm.bias""")				)

						# fmt: on

						return rename_keys





def     lowerCamelCase_ (					SCREAMING_SNAKE_CASE__ :     int,     SCREAMING_SNAKE_CASE__ :     str,     SCREAMING_SNAKE_CASE__ :     Tuple				)	->		Dict:
						UpperCAmelCase_     : Optional[int]         =					dct.pop(_lowerCamelCase				)
						UpperCAmelCase_     : Dict         =					val





def     lowerCamelCase_ (					SCREAMING_SNAKE_CASE__ :     Optional[Any],     SCREAMING_SNAKE_CASE__ :     str				)	->		str:
						# transformer decoder self-attention layers
						UpperCAmelCase_     : int         =					config.d_model
						for i in range(config.decoder_layers				):
												# read in weights + bias of input projection layer of self-attention
												UpperCAmelCase_     : Tuple         =					state_dict.pop(F"""transformer.decoder.layers.{i}.self_attn.in_proj_weight"""				)
												UpperCAmelCase_     : Tuple         =					state_dict.pop(F"""transformer.decoder.layers.{i}.self_attn.in_proj_bias"""				)
												# next, add query, keys and values (in that order) to the state dict
												UpperCAmelCase_     : List[str]         =					in_proj_weight[:hidden_size, :]
												UpperCAmelCase_     : Union[str, Any]         =					in_proj_bias[:hidden_size]
												UpperCAmelCase_     : str         =					in_proj_weight[
												    hidden_size : hidden_size * 2, :
												]
												UpperCAmelCase_     : Dict         =					in_proj_bias[hidden_size : hidden_size * 2]
												UpperCAmelCase_     : Optional[Any]         =					in_proj_weight[-hidden_size:, :]
												UpperCAmelCase_     : Tuple         =					in_proj_bias[-hidden_size:]





def     lowerCamelCase_ (					)	->		List[str]:
						UpperCAmelCase_     : Union[str, Any]         =					"""http://images.cocodataset.org/val2017/000000039769.jpg"""
						UpperCAmelCase_     : Dict         =					Image.open(requests.get(_lowerCamelCase,     stream=_lowerCamelCase				).raw				)

						return im





@torch.no_grad()
def     lowerCamelCase_ (					SCREAMING_SNAKE_CASE__ :     int,     SCREAMING_SNAKE_CASE__ :     Union[str, Any],     SCREAMING_SNAKE_CASE__ :     Tuple				)	->		Optional[int]:
						UpperCAmelCase_     : str         =					get_deta_config()

						# load original state dict
						if model_name == "deta-resnet-50":
												UpperCAmelCase_     : List[str]         =					"""adet_checkpoint0011.pth"""
						elif model_name == "deta-resnet-50-24-epochs":
												UpperCAmelCase_     : List[str]         =					"""adet_2x_checkpoint0023.pth"""
						else:
												raise ValueError(F"""Model name {model_name} not supported"""				)
						UpperCAmelCase_     : List[str]         =					hf_hub_download(repo_id='''nielsr/deta-checkpoints''',     filename=_lowerCamelCase				)
						UpperCAmelCase_     : Union[str, Any]         =					torch.load(_lowerCamelCase,     map_location='''cpu'''				)["""model"""]

						# rename keys
						UpperCAmelCase_     : List[str]         =					create_rename_keys(_lowerCamelCase				)
						for src, dest in rename_keys:
												rename_key(_lowerCamelCase,     _lowerCamelCase,     _lowerCamelCase				)
						read_in_decoder_q_k_v(_lowerCamelCase,     _lowerCamelCase				)

						# fix some prefixes
						for key in state_dict.copy().keys():
												if "transformer.decoder.class_embed" in key or "transformer.decoder.bbox_embed" in key:
																		UpperCAmelCase_     : Optional[Any]         =					state_dict.pop(_lowerCamelCase				)
																		UpperCAmelCase_     : str         =					val
												if "input_proj" in key:
																		UpperCAmelCase_     : List[Any]         =					state_dict.pop(_lowerCamelCase				)
																		UpperCAmelCase_     : Tuple         =					val
												if "level_embed" in key or "pos_trans" in key or "pix_trans" in key or "enc_output" in key:
																		UpperCAmelCase_     : Optional[Any]         =					state_dict.pop(_lowerCamelCase				)
																		UpperCAmelCase_     : Dict         =					val

												# finally, create HuggingFace model and load state dict
						UpperCAmelCase_     : Union[str, Any]         =					DetaForObjectDetection(_lowerCamelCase				)
						model.load_state_dict(_lowerCamelCase				)
						model.eval()

						UpperCAmelCase_     : Dict         =					"""cuda""" if torch.cuda.is_available() else """cpu"""
						model.to(_lowerCamelCase				)

						# load image processor
						UpperCAmelCase_     : Tuple         =					DetaImageProcessor(format='''coco_detection'''				)

						# verify our conversion on image
						UpperCAmelCase_     : List[Any]         =					prepare_img()
						UpperCAmelCase_     : int         =					processor(images=_lowerCamelCase,     return_tensors='''pt'''				)
						UpperCAmelCase_     : Dict         =					encoding["""pixel_values"""]
						UpperCAmelCase_     : Optional[int]         =					model(pixel_values.to(_lowerCamelCase				)				)

						# verify logits
						if model_name == "deta-resnet-50":
												UpperCAmelCase_     : Any         =					torch.tensor(
												    [[-7.39_78, -2.54_06, -4.16_68], [-8.26_84, -3.99_33, -3.80_96], [-7.05_15, -3.79_73, -5.85_16]]				)
												UpperCAmelCase_     : str         =					torch.tensor([[0.50_43, 0.49_73, 0.99_98], [0.25_42, 0.54_89, 0.47_48], [0.54_90, 0.27_65, 0.05_70]]				)
						elif model_name == "deta-resnet-50-24-epochs":
												UpperCAmelCase_     : str         =					torch.tensor(
												    [[-7.16_88, -2.48_57, -4.86_69], [-7.86_30, -3.81_54, -4.26_74], [-7.27_30, -4.18_65, -5.53_23]]				)
												UpperCAmelCase_     : List[Any]         =					torch.tensor([[0.50_21, 0.49_71, 0.99_94], [0.25_46, 0.54_86, 0.47_31], [0.16_86, 0.19_86, 0.21_42]]				)

						assert torch.allclose(outputs.logits[0, :3, :3],     expected_logits.to(_lowerCamelCase				),     atol=1E-4				)
						assert torch.allclose(outputs.pred_boxes[0, :3, :3],     expected_boxes.to(_lowerCamelCase				),     atol=1E-4				)
						print('''Everything ok!'''				)

						if pytorch_dump_folder_path:
												# Save model and processor
												logger.info(F"""Saving PyTorch model and processor to {pytorch_dump_folder_path}..."""				)
												Path(_lowerCamelCase				).mkdir(exist_ok=_lowerCamelCase				)
												model.save_pretrained(_lowerCamelCase				)
												processor.save_pretrained(_lowerCamelCase				)

						# Push to hub
						if push_to_hub:
												print('''Pushing model and processor to hub...'''				)
												model.push_to_hub(F"""jozhang97/{model_name}"""				)
												processor.push_to_hub(F"""jozhang97/{model_name}"""				)


if __name__ == "__main__":
					snake_case_       :    List[Any]          =					argparse.ArgumentParser()

					parser.add_argument(
					    "--model_name",
					    type=str,
					    default="deta-resnet-50",
					    choices=["deta-resnet-50", "deta-resnet-50-24-epochs"],
					    help="Name of the model 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."
					)
					snake_case_       :    Optional[int]          =					parser.parse_args()
					convert_deta_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)






'''simple docstring'''







import collections
import inspect
import unittest

from transformers import NatConfig
from transformers.testing_utils import require_natten, require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_backbone_common import BackboneTesterMixin
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 torch import nn

					from transformers import NatBackbone, NatForImageClassification, NatModel
					from transformers.models.nat.modeling_nat import NAT_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
					from PIL import Image

					from transformers import AutoImageProcessor


class    __a :

							def __init__(							self							:       List[str]    ,       __magic_name__							:       str    ,       __magic_name__							:       str=13    ,       __magic_name__							:       Any=64    ,       __magic_name__							:       str=4    ,       __magic_name__							:       int=3    ,       __magic_name__							:       List[Any]=16    ,       __magic_name__							:       Union[str, Any]=[1, 2, 1]    ,       __magic_name__							:       Optional[Any]=[2, 4, 8]    ,       __magic_name__							:       List[Any]=3    ,       __magic_name__							:       List[str]=2.0    ,       __magic_name__							:       str=True    ,       __magic_name__							:       int=0.0    ,       __magic_name__							:       List[str]=0.0    ,       __magic_name__							:       Optional[int]=0.1    ,       __magic_name__							:       Any="gelu"    ,       __magic_name__							:       List[Any]=True    ,       __magic_name__							:       Optional[int]=0.0_2    ,       __magic_name__							:       str=1E-5    ,       __magic_name__							:       str=True    ,       __magic_name__							:       Any=None    ,       __magic_name__							:       Any=True    ,       __magic_name__							:       str=10    ,       __magic_name__							:       Optional[Any]=["stage1", "stage2"]    ,       __magic_name__							:       List[Any]=[1, 2]    ,       )   ->    Tuple:







													"""simple docstring"""
													UpperCAmelCase_     : Optional[Any]         =					parent
													UpperCAmelCase_     : Optional[int]         =					batch_size
													UpperCAmelCase_     : Dict         =					image_size
													UpperCAmelCase_     : Union[str, Any]         =					patch_size
													UpperCAmelCase_     : Tuple         =					num_channels
													UpperCAmelCase_     : List[str]         =					embed_dim
													UpperCAmelCase_     : str         =					depths
													UpperCAmelCase_     : Optional[Any]         =					num_heads
													UpperCAmelCase_     : Dict         =					kernel_size
													UpperCAmelCase_     : Optional[Any]         =					mlp_ratio
													UpperCAmelCase_     : Dict         =					qkv_bias
													UpperCAmelCase_     : Optional[int]         =					hidden_dropout_prob
													UpperCAmelCase_     : Dict         =					attention_probs_dropout_prob
													UpperCAmelCase_     : Dict         =					drop_path_rate
													UpperCAmelCase_     : List[str]         =					hidden_act
													UpperCAmelCase_     : str         =					patch_norm
													UpperCAmelCase_     : List[str]         =					layer_norm_eps
													UpperCAmelCase_     : Tuple         =					initializer_range
													UpperCAmelCase_     : List[str]         =					is_training
													UpperCAmelCase_     : str         =					scope
													UpperCAmelCase_     : Dict         =					use_labels
													UpperCAmelCase_     : List[str]         =					num_labels
													UpperCAmelCase_     : Union[str, Any]         =					out_features
													UpperCAmelCase_     : Any         =					out_indices

							def       UpperCAmelCase__		(							self							:       List[Any]				)   ->    Optional[Any]:







													"""simple docstring"""
													UpperCAmelCase_     : Any         =					floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]				)

													UpperCAmelCase_     : Dict         =					None
													if self.use_labels:
																			UpperCAmelCase_     : Any         =					ids_tensor([self.batch_size]    ,       self.num_labels				)

													UpperCAmelCase_     : str         =					self.get_config()

													return config, pixel_values, labels

							def       UpperCAmelCase__		(							self							:       Optional[Any]				)   ->    Any:







													"""simple docstring"""
													return NatConfig(
													    num_labels=self.num_labels    ,       image_size=self.image_size    ,       patch_size=self.patch_size    ,       num_channels=self.num_channels    ,       embed_dim=self.embed_dim    ,       depths=self.depths    ,       num_heads=self.num_heads    ,       kernel_size=self.kernel_size    ,       mlp_ratio=self.mlp_ratio    ,       qkv_bias=self.qkv_bias    ,       hidden_dropout_prob=self.hidden_dropout_prob    ,       attention_probs_dropout_prob=self.attention_probs_dropout_prob    ,       drop_path_rate=self.drop_path_rate    ,       hidden_act=self.hidden_act    ,       patch_norm=self.patch_norm    ,       layer_norm_eps=self.layer_norm_eps    ,       initializer_range=self.initializer_range    ,       out_features=self.out_features    ,       out_indices=self.out_indices    ,       )

							def       UpperCAmelCase__		(							self							:       Tuple    ,       __magic_name__							:       int    ,       __magic_name__							:       int    ,       __magic_name__							:       Optional[int]				)   ->    List[str]:







													"""simple docstring"""
													UpperCAmelCase_     : List[Any]         =					NatModel(config=__magic_name__				)
													model.to(__magic_name__				)
													model.eval()
													UpperCAmelCase_     : str         =					model(__magic_name__				)

													UpperCAmelCase_     : Any         =					(config.image_size // config.patch_size) // (2 ** (len(config.depths				) - 1))
													UpperCAmelCase_     : Tuple         =					int(config.embed_dim * 2 ** (len(config.depths				) - 1)				)

													self.parent.assertEqual(
													    result.last_hidden_state.shape    ,       (self.batch_size, expected_height, expected_width, expected_dim)				)

							def       UpperCAmelCase__		(							self							:       Union[str, Any]    ,       __magic_name__							:       Optional[int]    ,       __magic_name__							:       int    ,       __magic_name__							:       Dict				)   ->    Dict:







													"""simple docstring"""
													UpperCAmelCase_     : str         =					NatForImageClassification(__magic_name__				)
													model.to(__magic_name__				)
													model.eval()
													UpperCAmelCase_     : List[Any]         =					model(__magic_name__    ,       labels=__magic_name__				)
													self.parent.assertEqual(result.logits.shape    ,       (self.batch_size, self.num_labels)				)

													# test greyscale images
													UpperCAmelCase_     : Tuple         =					1
													UpperCAmelCase_     : int         =					NatForImageClassification(__magic_name__				)
													model.to(__magic_name__				)
													model.eval()

													UpperCAmelCase_     : List[str]         =					floats_tensor([self.batch_size, 1, self.image_size, self.image_size]				)
													UpperCAmelCase_     : Union[str, Any]         =					model(__magic_name__				)
													self.parent.assertEqual(result.logits.shape    ,       (self.batch_size, self.num_labels)				)

							def       UpperCAmelCase__		(							self							:       Tuple    ,       __magic_name__							:       str    ,       __magic_name__							:       Any    ,       __magic_name__							:       Any				)   ->    List[str]:







													"""simple docstring"""
													UpperCAmelCase_     : Any         =					NatBackbone(config=__magic_name__				)
													model.to(__magic_name__				)
													model.eval()
													UpperCAmelCase_     : int         =					model(__magic_name__				)

													# verify hidden states
													self.parent.assertEqual(len(result.feature_maps				)    ,       len(config.out_features				)				)
													self.parent.assertListEqual(list(result.feature_maps[0].shape				)    ,       [self.batch_size, model.channels[0], 16, 16]				)

													# verify channels
													self.parent.assertEqual(len(model.channels				)    ,       len(config.out_features				)				)

													# verify backbone works with out_features=None
													UpperCAmelCase_     : Optional[int]         =					None
													UpperCAmelCase_     : Any         =					NatBackbone(config=__magic_name__				)
													model.to(__magic_name__				)
													model.eval()
													UpperCAmelCase_     : Optional[int]         =					model(__magic_name__				)

													# verify feature maps
													self.parent.assertEqual(len(result.feature_maps				)    ,       1				)
													self.parent.assertListEqual(list(result.feature_maps[0].shape				)    ,       [self.batch_size, model.channels[-1], 4, 4]				)

													# verify channels
													self.parent.assertEqual(len(model.channels				)    ,       1				)



							def       UpperCAmelCase__		(							self							:       Union[str, Any]				)   ->    int:







													"""simple docstring"""
													UpperCAmelCase_     : Optional[Any]         =					self.prepare_config_and_inputs()
													UpperCAmelCase_  ,    UpperCAmelCase_  ,    UpperCAmelCase_     : Any         =					config_and_inputs
													UpperCAmelCase_     : Tuple         =					{'''pixel_values''': pixel_values}
													return config, inputs_dict





@require_natten
@require_torch
class    __a (lowerCamelCase    ,						lowerCamelCase    ,						unittest.TestCase       ):
							__a   :				List[Any]       =    (
							    (
							        NatModel,
							        NatForImageClassification,
							        NatBackbone,
							    )
							    if is_torch_available()
							    else ()
							)
							__a   :				Dict       =    (
							    {"feature-extraction": NatModel, "image-classification": NatForImageClassification}
							    if is_torch_available()
							    else {}
							)
							__a   :				List[Any]       =    False

							__a   :				Optional[Any]       =    False
							__a   :				int       =    False
							__a   :				List[str]       =    False
							__a   :				Optional[int]       =    False

							def       UpperCAmelCase__		(							self							:       Optional[Any]				)   ->    Tuple:







													"""simple docstring"""
													UpperCAmelCase_     : Any         =					NatModelTester(self				)
													UpperCAmelCase_     : List[str]         =					ConfigTester(self    ,       config_class=__magic_name__    ,       embed_dim=37				)

							def       UpperCAmelCase__		(							self							:       List[str]				)   ->    Optional[Any]:







													"""simple docstring"""
													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							:       str				)   ->    List[str]:







													"""simple docstring"""
													return

							def       UpperCAmelCase__		(							self							:       Any				)   ->    Union[str, Any]:







													"""simple docstring"""
													UpperCAmelCase_     : List[str]         =					self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_model(*__magic_name__				)

							def       UpperCAmelCase__		(							self							:       Optional[int]				)   ->    Any:







													"""simple docstring"""
													UpperCAmelCase_     : str         =					self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_image_classification(*__magic_name__				)

							def       UpperCAmelCase__		(							self							:       Optional[Any]				)   ->    int:







													"""simple docstring"""
													UpperCAmelCase_     : Optional[int]         =					self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_backbone(*__magic_name__				)

							@unittest.skip(reason='''Nat does not use inputs_embeds'''				)
							def       UpperCAmelCase__		(							self							:       int				)   ->    Optional[int]:







													"""simple docstring"""
													pass

							@unittest.skip(reason='''Nat does not use feedforward chunking'''				)
							def       UpperCAmelCase__		(							self							:       List[str]				)   ->    Dict:







													"""simple docstring"""
													pass

							def       UpperCAmelCase__		(							self							:       Any				)   ->    Optional[Any]:







													"""simple docstring"""
													UpperCAmelCase_  ,    UpperCAmelCase_     : Union[str, Any]         =					self.model_tester.prepare_config_and_inputs_for_common()

													for model_class in self.all_model_classes:
																			UpperCAmelCase_     : List[str]         =					model_class(__magic_name__				)
																			self.assertIsInstance(model.get_input_embeddings()    ,       (nn.Module)				)
																			UpperCAmelCase_     : Optional[Any]         =					model.get_output_embeddings()
																			self.assertTrue(x is None or isinstance(__magic_name__    ,       nn.Linear				)				)

							def       UpperCAmelCase__		(							self							:       Optional[int]				)   ->    str:







													"""simple docstring"""
													UpperCAmelCase_  ,    UpperCAmelCase_     : int         =					self.model_tester.prepare_config_and_inputs_for_common()

													for model_class in self.all_model_classes:
																			UpperCAmelCase_     : Optional[Any]         =					model_class(__magic_name__				)
																			UpperCAmelCase_     : List[str]         =					inspect.signature(model.forward				)
																			# signature.parameters is an OrderedDict => so arg_names order is deterministic
																			UpperCAmelCase_     : Tuple         =					[*signature.parameters.keys()]

																			UpperCAmelCase_     : Dict         =					['''pixel_values''']
																			self.assertListEqual(arg_names[:1]    ,       __magic_name__				)

							def       UpperCAmelCase__		(							self							:       List[Any]				)   ->    str:







													"""simple docstring"""
													self.skipTest('''Nat\'s attention operation is handled entirely by NATTEN.'''				)

							def       UpperCAmelCase__		(							self							:       Optional[Any]    ,       __magic_name__							:       Any    ,       __magic_name__							:       List[Any]    ,       __magic_name__							:       List[Any]    ,       __magic_name__							:       Dict				)   ->    Tuple:







													"""simple docstring"""
													UpperCAmelCase_     : Union[str, Any]         =					model_class(__magic_name__				)
													model.to(__magic_name__				)
													model.eval()

													with torch.no_grad():
																			UpperCAmelCase_     : List[Any]         =					model(**self._prepare_for_class(__magic_name__    ,       __magic_name__				)				)

													UpperCAmelCase_     : int         =					outputs.hidden_states

													UpperCAmelCase_     : Tuple         =					getattr(
													    self.model_tester    ,       '''expected_num_hidden_layers'''    ,       len(self.model_tester.depths				) + 1				)
													self.assertEqual(len(__magic_name__				)    ,       __magic_name__				)

													# Nat has a different seq_length
													UpperCAmelCase_     : Dict         =					(
													    config.patch_size
													    if isinstance(config.patch_size    ,       collections.abc.Iterable				)
													    else (config.patch_size, config.patch_size)
													)

													UpperCAmelCase_     : Optional[int]         =					image_size[0] // patch_size[0]
													UpperCAmelCase_     : List[Any]         =					image_size[1] // patch_size[1]

													self.assertListEqual(
													    list(hidden_states[0].shape[-3:]				)    ,       [height, width, self.model_tester.embed_dim]    ,       )

													if model_class.__name__ != "NatBackbone":
																			UpperCAmelCase_     : str         =					outputs.reshaped_hidden_states
																			self.assertEqual(len(__magic_name__				)    ,       __magic_name__				)

																			UpperCAmelCase_  ,    UpperCAmelCase_  ,    UpperCAmelCase_  ,    UpperCAmelCase_     : int         =					reshaped_hidden_states[0].shape
																			UpperCAmelCase_     : str         =					(
																			    reshaped_hidden_states[0].view(__magic_name__    ,       __magic_name__    ,       __magic_name__    ,       __magic_name__				).permute(0    ,       2    ,       3    ,       1				)
																			)
																			self.assertListEqual(
																			    list(reshaped_hidden_states.shape[-3:]				)    ,       [height, width, self.model_tester.embed_dim]    ,       )

							def       UpperCAmelCase__		(							self							:       Optional[Any]				)   ->    List[str]:







													"""simple docstring"""
													UpperCAmelCase_  ,    UpperCAmelCase_     : List[Any]         =					self.model_tester.prepare_config_and_inputs_for_common()

													UpperCAmelCase_     : Dict         =					(
													    self.model_tester.image_size
													    if isinstance(self.model_tester.image_size    ,       collections.abc.Iterable				)
													    else (self.model_tester.image_size, self.model_tester.image_size)
													)

													for model_class in self.all_model_classes:
																			UpperCAmelCase_     : str         =					True
																			self.check_hidden_states_output(__magic_name__    ,       __magic_name__    ,       __magic_name__    ,       __magic_name__				)

																			# check that output_hidden_states also work using config
																			del inputs_dict["output_hidden_states"]
																			UpperCAmelCase_     : Tuple         =					True

																			self.check_hidden_states_output(__magic_name__    ,       __magic_name__    ,       __magic_name__    ,       __magic_name__				)

							@slow
							def       UpperCAmelCase__		(							self							:       List[Any]				)   ->    Optional[Any]:







													"""simple docstring"""
													for model_name in NAT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																			UpperCAmelCase_     : List[Any]         =					NatModel.from_pretrained(__magic_name__				)
																			self.assertIsNotNone(__magic_name__				)



							def       UpperCAmelCase__		(							self							:       Tuple				)   ->    str:







													"""simple docstring"""
													UpperCAmelCase_  ,    UpperCAmelCase_     : str         =					self.model_tester.prepare_config_and_inputs_for_common()

													UpperCAmelCase_     : Optional[Any]         =					_config_zero_init(__magic_name__				)
													for model_class in self.all_model_classes:
																			UpperCAmelCase_     : List[str]         =					model_class(config=__magic_name__				)
																			for name, param in model.named_parameters():
																									if "embeddings" not in name and 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"""    ,       )





@require_natten
@require_vision
@require_torch
class    __a (unittest.TestCase       ):

							@cached_property
							def       UpperCAmelCase__		(							self							:       str				)   ->    Optional[Any]:







													"""simple docstring"""
													return AutoImageProcessor.from_pretrained('''shi-labs/nat-mini-in1k-224'''				) if is_vision_available() else None



							@slow
							def       UpperCAmelCase__		(							self							:       Tuple				)   ->    str:







													"""simple docstring"""
													UpperCAmelCase_     : int         =					NatForImageClassification.from_pretrained('''shi-labs/nat-mini-in1k-224'''				).to(__magic_name__				)
													UpperCAmelCase_     : Union[str, Any]         =					self.default_image_processor

													UpperCAmelCase_     : Dict         =					Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''				)
													UpperCAmelCase_     : List[Any]         =					image_processor(images=__magic_name__    ,       return_tensors='''pt'''				).to(__magic_name__				)

													# forward pass
													with torch.no_grad():
																			UpperCAmelCase_     : List[str]         =					model(**__magic_name__				)

													# verify the logits
													UpperCAmelCase_     : int         =					torch.Size((1, 10_00)				)
													self.assertEqual(outputs.logits.shape    ,       __magic_name__				)
													UpperCAmelCase_     : List[str]         =					torch.tensor([0.3_8_0_5, -0.8_6_7_6, -0.3_9_1_2]				).to(__magic_name__				)
													self.assertTrue(torch.allclose(outputs.logits[0, :3]    ,       __magic_name__    ,       atol=1E-4				)				)





@require_torch
@require_natten
class    __a (unittest.TestCase    ,						lowerCamelCase       ):
							__a   :				List[str]       =    (NatBackbone,) if is_torch_available() else ()
							__a   :				Optional[Any]       =    NatConfig

							def       UpperCAmelCase__		(							self							:       Any				)   ->    Optional[Any]:







													"""simple docstring"""
													UpperCAmelCase_     : str         =					NatModelTester(self				)





