




'''simple docstring'''




import argparse

import requests
import torch
from PIL import Image

from transformers import (
    CLIPSegConfig,
    CLIPSegForImageSegmentation,
    CLIPSegProcessor,
    CLIPSegTextConfig,
    CLIPSegVisionConfig,
    CLIPTokenizer,
    ViTImageProcessor,
)




def 						snake_case__    ( lowerCamelCase__				:   List[str]			)   ->   Dict:
					A_ :     Optional[int] =							CLIPSegTextConfig()
					A_ :     Union[str, Any] =							CLIPSegVisionConfig(patch_size=1_6			)

					A_ :     List[str] =							True if '''refined''' in model_name else False
					A_ :     Any =							1_6 if '''rd16''' in model_name else 6_4

					A_ :     Any =							CLIPSegConfig.from_text_vision_configs(
					    lowerCamelCase__			,		lowerCamelCase__			,		use_complex_transposed_convolution=lowerCamelCase__			,		reduce_dim=lowerCamelCase__			,		)
					return config




def 						snake_case__    ( lowerCamelCase__				:   Tuple			)   ->   Union[str, Any]:
					# update prefixes
					if "clip_model" in name:
										A_ :     Any =							name.replace('''clip_model'''			,		'''clip'''			)
					if "transformer" in name:
										if "visual" in name:
															A_ :     str =							name.replace('''visual.transformer'''			,		'''vision_model'''			)
										else:
															A_ :     List[Any] =							name.replace('''transformer'''			,		'''text_model'''			)
					if "resblocks" in name:
										A_ :     Union[str, Any] =							name.replace('''resblocks'''			,		'''encoder.layers'''			)
					if "ln_1" in name:
										A_ :     str =							name.replace('''ln_1'''			,		'''layer_norm1'''			)
					if "ln_2" in name:
										A_ :     Union[str, Any] =							name.replace('''ln_2'''			,		'''layer_norm2'''			)
					if "c_fc" in name:
										A_ :     int =							name.replace('''c_fc'''			,		'''fc1'''			)
					if "c_proj" in name:
										A_ :     Optional[int] =							name.replace('''c_proj'''			,		'''fc2'''			)
					if "attn" in name and "self" not in name:
										A_ :     List[Any] =							name.replace('''attn'''			,		'''self_attn'''			)
					# text encoder
					if "token_embedding" in name:
										A_ :     Any =							name.replace('''token_embedding'''			,		'''text_model.embeddings.token_embedding'''			)
					if "positional_embedding" in name and "visual" not in name:
										A_ :     Any =							name.replace('''positional_embedding'''			,		'''text_model.embeddings.position_embedding.weight'''			)
					if "ln_final" in name:
										A_ :     Optional[int] =							name.replace('''ln_final'''			,		'''text_model.final_layer_norm'''			)
					# vision encoder
					if "visual.class_embedding" in name:
										A_ :     Dict =							name.replace('''visual.class_embedding'''			,		'''vision_model.embeddings.class_embedding'''			)
					if "visual.conv1" in name:
										A_ :     Any =							name.replace('''visual.conv1'''			,		'''vision_model.embeddings.patch_embedding'''			)
					if "visual.positional_embedding" in name:
										A_ :     int =							name.replace('''visual.positional_embedding'''			,		'''vision_model.embeddings.position_embedding.weight'''			)
					if "visual.ln_pre" in name:
										A_ :     Optional[Any] =							name.replace('''visual.ln_pre'''			,		'''vision_model.pre_layrnorm'''			)
					if "visual.ln_post" in name:
										A_ :     int =							name.replace('''visual.ln_post'''			,		'''vision_model.post_layernorm'''			)
					# projection layers
					if "visual.proj" in name:
										A_ :     Any =							name.replace('''visual.proj'''			,		'''visual_projection.weight'''			)
					if "text_projection" in name:
										A_ :     List[Any] =							name.replace('''text_projection'''			,		'''text_projection.weight'''			)
					# decoder
					if "trans_conv" in name:
										A_ :     Optional[int] =							name.replace('''trans_conv'''			,		'''transposed_convolution'''			)
					if "film_mul" in name or "film_add" in name or "reduce" in name or "transposed_convolution" in name:
										A_ :     str =							'''decoder.''' + name
					if "blocks" in name:
										A_ :     List[str] =							name.replace('''blocks'''			,		'''decoder.layers'''			)
					if "linear1" in name:
										A_ :     int =							name.replace('''linear1'''			,		'''mlp.fc1'''			)
					if "linear2" in name:
										A_ :     Optional[int] =							name.replace('''linear2'''			,		'''mlp.fc2'''			)
					if "norm1" in name and "layer_" not in name:
										A_ :     Optional[Any] =							name.replace('''norm1'''			,		'''layer_norm1'''			)
					if "norm2" in name and "layer_" not in name:
										A_ :     Optional[Any] =							name.replace('''norm2'''			,		'''layer_norm2'''			)

					return name




def 						snake_case__    ( lowerCamelCase__				:   Optional[int]			,		lowerCamelCase__				:   int			)   ->   Dict:
					for key in orig_state_dict.copy().keys():
										A_ :     Optional[Any] =							orig_state_dict.pop(lowerCamelCase__			)

										if key.startswith('''clip_model'''			) and "attn.in_proj" in key:
															A_ :     str =							key.split('''.'''			)
															if "visual" in key:
																				A_ :     str =							int(key_split[4]			)
																				A_ :     Dict =							config.vision_config.hidden_size
																				A_ :     Optional[Any] =							'''vision_model'''
															else:
																				A_ :     List[Any] =							int(key_split[3]			)
																				A_ :     Optional[Any] =							config.text_config.hidden_size
																				A_ :     Any =							'''text_model'''

															if "weight" in key:
																				A_ :     Dict =							val[:dim, :]
																				A_ :     Any =							val[
																				    dim : dim * 2, :
																				]
																				A_ :     str =							val[-dim:, :]
															else:
																				A_ :     Any =							val[:dim]
																				A_ :     str =							val[dim : dim * 2]
																				A_ :     Optional[Any] =							val[-dim:]
										elif "self_attn" in key and "out_proj" not in key:
															A_ :     Optional[int] =							key.split('''.'''			)
															A_ :     Optional[int] =							int(key_split[1]			)
															A_ :     Optional[Any] =							config.reduce_dim
															if "weight" in key:
																				A_ :     str =							val[:dim, :]
																				A_ :     Any =							val[dim : dim * 2, :]
																				A_ :     Optional[int] =							val[-dim:, :]
															else:
																				A_ :     List[Any] =							val[:dim]
																				A_ :     Tuple =							val[dim : dim * 2]
																				A_ :     List[str] =							val[-dim:]
										else:
															A_ :     int =							rename_key(lowerCamelCase__			)
															if "visual_projection" in new_name or "text_projection" in new_name:
																				A_ :     List[str] =							val.T
															A_ :     List[str] =							val

					return orig_state_dict




def 						snake_case__    ( )   ->   Dict:
					A_ :     Dict =							'''http://images.cocodataset.org/val2017/000000039769.jpg'''
					A_ :     List[Any] =							Image.open(requests.get(lowerCamelCase__			,		stream=lowerCamelCase__			).raw			)
					return image




def 						snake_case__    ( lowerCamelCase__				:   Any			,		lowerCamelCase__				:   Tuple			,		lowerCamelCase__				:   str			,		lowerCamelCase__				:   int			)   ->   Optional[Any]:
					A_ :     Any =							get_clipseg_config(lowerCamelCase__			)
					A_ :     int =							CLIPSegForImageSegmentation(lowerCamelCase__			)
					model.eval()

					A_ :     Optional[int] =							torch.load(lowerCamelCase__			,		map_location='''cpu'''			)

					# remove some keys
					for key in state_dict.copy().keys():
										if key.startswith('''model'''			):
															state_dict.pop(lowerCamelCase__			,		lowerCamelCase__			)

    # rename some keys
					A_ :     Any =							convert_state_dict(lowerCamelCase__			,		lowerCamelCase__			)
					A_ :     Tuple =							model.load_state_dict(lowerCamelCase__			,		strict=lowerCamelCase__			)

					if missing_keys != ["clip.text_model.embeddings.position_ids", "clip.vision_model.embeddings.position_ids"]:
										raise ValueError('''Missing keys that are not expected: {}'''.format(lowerCamelCase__			)			)
					if unexpected_keys != ["decoder.reduce.weight", "decoder.reduce.bias"]:
										raise ValueError(f'Unexpected keys: {unexpected_keys}'			)

					A_ :     Optional[Any] =							ViTImageProcessor(size=3_5_2			)
					A_ :     str =							CLIPTokenizer.from_pretrained('''openai/clip-vit-base-patch32'''			)
					A_ :     Tuple =							CLIPSegProcessor(image_processor=lowerCamelCase__			,		tokenizer=lowerCamelCase__			)

					A_ :     Optional[int] =							prepare_img()
					A_ :     Tuple =							['''a glass''', '''something to fill''', '''wood''', '''a jar''']

					A_ :     str =							processor(text=lowerCamelCase__			,		images=[image] * len(lowerCamelCase__			)			,		padding='''max_length'''			,		return_tensors='''pt'''			)

					with torch.no_grad():
										A_ :     str =							model(**lowerCamelCase__			)

					# verify values
					A_ :     str =							torch.tensor([0.1110, -0.1882, 0.1645]			)
					A_ :     List[Any] =							torch.tensor([0.2692, -0.7197, -0.1328]			)
					if model_name == "clipseg-rd64-refined":
										A_ :     List[Any] =							torch.tensor(
										    [[-10.0407, -9.9431, -10.2646], [-9.9751, -9.7064, -9.9586], [-9.6891, -9.5645, -9.9618]]			)
					elif model_name == "clipseg-rd64":
										A_ :     Dict =							torch.tensor(
										    [[-7.2877, -7.2711, -7.2463], [-7.2652, -7.2780, -7.2520], [-7.2239, -7.2204, -7.2001]]			)
					elif model_name == "clipseg-rd16":
										A_ :     Any =							torch.tensor(
										    [[-6.3955, -6.4055, -6.4151], [-6.3911, -6.4033, -6.4100], [-6.3474, -6.3702, -6.3762]]			)
					else:
										raise ValueError(f'Model name {model_name} not supported.'			)

					assert torch.allclose(outputs.logits[0, :3, :3]			,		lowerCamelCase__			,		atol=1e-3			)
					assert torch.allclose(outputs.conditional_embeddings[0, :3]			,		lowerCamelCase__			,		atol=1e-3			)
					assert torch.allclose(outputs.pooled_output[0, :3]			,		lowerCamelCase__			,		atol=1e-3			)
					print('''Looks ok!'''			)

					if pytorch_dump_folder_path is not None:
										print(f'Saving model and processor to {pytorch_dump_folder_path}'			)
										model.save_pretrained(lowerCamelCase__			)
										processor.save_pretrained(lowerCamelCase__			)

					if push_to_hub:
										print(f'Pushing model and processor for {model_name} to the hub'			)
										model.push_to_hub(f'CIDAS/{model_name}'			)
										processor.push_to_hub(f'CIDAS/{model_name}'			)


if __name__ == "__main__":
							snake_case__          =						argparse.ArgumentParser()
							# Required parameters
							parser.add_argument(
							    """--model_name""",
							    default="""clipseg-rd64""",
							    type=str,
							    choices=["""clipseg-rd16""", """clipseg-rd64""", """clipseg-rd64-refined"""],
							    help=(
							        """Name of the model. Supported models are: clipseg-rd64, clipseg-rd16 and clipseg-rd64-refined (rd meaning"""
							        """ reduce dimension)"""
							    ),
							)
							parser.add_argument(
							    """--checkpoint_path""",
							    default="""/Users/nielsrogge/Documents/CLIPSeg/clip_plus_rd64-uni.pth""",
							    type=str,
							    help=(
							        """Path to the original checkpoint. Note that the script assumes that the checkpoint includes both CLIP and"""
							        """ the decoder weights."""
							    ),
							)
							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__          =						parser.parse_args()
							convert_clipseg_checkpoint(args.model_name, args.checkpoint_path, 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     UpperCamelCase_     :







						"""simple docstring"""


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





											"""simple docstring"""







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







						def       _a    (     self      :       Dict    ):





											"""simple docstring"""







											A_ :     Dict =							floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]    )

											A_ :     int =							None
											if self.use_labels:
																A_ :     Union[str, Any] =							ids_tensor([self.batch_size]							, self.num_labels    )

											A_ :     int =							self.get_config()

											return config, pixel_values, labels







						def       _a    (     self      :       List[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       _a    (     self      :       int							, _lowerCamelCase      :       str							, _lowerCamelCase      :       Dict							, _lowerCamelCase      :       Dict    ):





											"""simple docstring"""







											A_ :     str =							NatModel(config=_lowerCamelCase    )
											model.to(_lowerCamelCase    )
											model.eval()
											A_ :     Union[str, Any] =							model(_lowerCamelCase    )

											A_ :     List[str] =							(config.image_size // config.patch_size) // (2 ** (len(config.depths    ) - 1))
											A_ :     Optional[int] =							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       _a    (     self      :       str							, _lowerCamelCase      :       Tuple							, _lowerCamelCase      :       int							, _lowerCamelCase      :       Union[str, Any]    ):





											"""simple docstring"""







											A_ :     int =							NatForImageClassification(_lowerCamelCase    )
											model.to(_lowerCamelCase    )
											model.eval()
											A_ :     str =							model(_lowerCamelCase							, labels=_lowerCamelCase    )
											self.parent.assertEqual(result.logits.shape							, (self.batch_size, self.num_labels)    )

											# test greyscale images
											A_ :     str =							1
											A_ :     Dict =							NatForImageClassification(_lowerCamelCase    )
											model.to(_lowerCamelCase    )
											model.eval()

											A_ :     str =							floats_tensor([self.batch_size, 1, self.image_size, self.image_size]    )
											A_ :     Tuple =							model(_lowerCamelCase    )
											self.parent.assertEqual(result.logits.shape							, (self.batch_size, self.num_labels)    )







						def       _a    (     self      :       List[Any]							, _lowerCamelCase      :       List[str]							, _lowerCamelCase      :       Any							, _lowerCamelCase      :       Tuple    ):





											"""simple docstring"""







											A_ :     Union[str, Any] =							NatBackbone(config=_lowerCamelCase    )
											model.to(_lowerCamelCase    )
											model.eval()
											A_ :     Optional[int] =							model(_lowerCamelCase    )

											# 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
											A_ :     Tuple =							None
											A_ :     List[Any] =							NatBackbone(config=_lowerCamelCase    )
											model.to(_lowerCamelCase    )
											model.eval()
											A_ :     Optional[int] =							model(_lowerCamelCase    )

											# 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       _a    (     self      :       List[Any]    ):





											"""simple docstring"""







											A_ :     Optional[int] =							self.prepare_config_and_inputs()
											A_ ,A_ ,A_ :     Any =							config_and_inputs
											A_ :     Union[str, Any] =							{'''pixel_values''': pixel_values}
											return config, inputs_dict






@require_natten
@require_torch
class     UpperCamelCase_     (a__,      a__,      unittest.TestCase     ):







						"""simple docstring"""


						_lowerCAmelCase    						=					(
						    (
						        NatModel,
						        NatForImageClassification,
						        NatBackbone,
						    )
						    if is_torch_available()
						    else ()
						)
						_lowerCAmelCase    						=					(
						    {'feature-extraction': NatModel, 'image-classification': NatForImageClassification}
						    if is_torch_available()
						    else {}
						)
						_lowerCAmelCase    						=					False

						_lowerCAmelCase    						=					False
						_lowerCAmelCase    						=					False
						_lowerCAmelCase    						=					False
						_lowerCAmelCase    						=					False







						def       _a    (     self      :       Any    ):





											"""simple docstring"""







											A_ :     Tuple =							NatModelTester(self    )
											A_ :     Any =							ConfigTester(self							, config_class=_lowerCamelCase							, embed_dim=37    )







						def       _a    (     self      :       Union[str, 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       _a    (     self      :       str    ):





											"""simple docstring"""







											return







						def       _a    (     self      :       Union[str, Any]    ):





											"""simple docstring"""







											A_ :     Union[str, Any] =							self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*_lowerCamelCase    )







						def       _a    (     self      :       Optional[int]    ):





											"""simple docstring"""







											A_ :     Dict =							self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_image_classification(*_lowerCamelCase    )







						def       _a    (     self      :       str    ):





											"""simple docstring"""







											A_ :     str =							self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_backbone(*_lowerCamelCase    )







						@unittest.skip(reason='''Nat does not use inputs_embeds'''    )
						def       _a    (     self      :       str    ):





											"""simple docstring"""







											pass







						@unittest.skip(reason='''Nat does not use feedforward chunking'''    )
						def       _a    (     self      :       Optional[Any]    ):





											"""simple docstring"""







											pass







						def       _a    (     self      :       Dict    ):





											"""simple docstring"""







											A_ ,A_ :     List[Any] =							self.model_tester.prepare_config_and_inputs_for_common()

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







						def       _a    (     self      :       Union[str, Any]    ):





											"""simple docstring"""







											A_ ,A_ :     Tuple =							self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																A_ :     List[Any] =							model_class(_lowerCamelCase    )
																A_ :     int =							inspect.signature(model.forward    )
																# signature.parameters is an OrderedDict => so arg_names order is deterministic
																A_ :     Dict =							[*signature.parameters.keys()]

																A_ :     List[Any] =							['''pixel_values''']
																self.assertListEqual(arg_names[:1]							, _lowerCamelCase    )







						def       _a    (     self      :       Optional[int]    ):





											"""simple docstring"""







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







						def       _a    (     self      :       List[Any]							, _lowerCamelCase      :       Optional[int]							, _lowerCamelCase      :       Union[str, Any]							, _lowerCamelCase      :       Tuple							, _lowerCamelCase      :       List[str]    ):





											"""simple docstring"""







											A_ :     Tuple =							model_class(_lowerCamelCase    )
											model.to(_lowerCamelCase    )
											model.eval()

											with torch.no_grad():
																A_ :     List[str] =							model(**self._prepare_for_class(_lowerCamelCase							, _lowerCamelCase    )    )

											A_ :     List[Any] =							outputs.hidden_states

											A_ :     List[str] =							getattr(
											    self.model_tester							, '''expected_num_hidden_layers'''							, len(self.model_tester.depths    ) + 1    )
											self.assertEqual(len(_lowerCamelCase    )							, _lowerCamelCase    )

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

											A_ :     Optional[Any] =							image_size[0] // patch_size[0]
											A_ :     Optional[int] =							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":
																A_ :     Tuple =							outputs.reshaped_hidden_states
																self.assertEqual(len(_lowerCamelCase    )							, _lowerCamelCase    )

																A_ ,A_ ,A_ ,A_ :     Union[str, Any] =							reshaped_hidden_states[0].shape
																A_ :     Optional[Any] =							(
																    reshaped_hidden_states[0].view(_lowerCamelCase							, _lowerCamelCase							, _lowerCamelCase							, _lowerCamelCase    ).permute(0							, 2							, 3							, 1    )
																)
																self.assertListEqual(
																    list(reshaped_hidden_states.shape[-3:]    )							, [height, width, self.model_tester.embed_dim]							, )







						def       _a    (     self      :       str    ):





											"""simple docstring"""







											A_ ,A_ :     List[str] =							self.model_tester.prepare_config_and_inputs_for_common()

											A_ :     List[str] =							(
											    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:
																A_ :     Tuple =							True
																self.check_hidden_states_output(_lowerCamelCase							, _lowerCamelCase							, _lowerCamelCase							, _lowerCamelCase    )

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

																self.check_hidden_states_output(_lowerCamelCase							, _lowerCamelCase							, _lowerCamelCase							, _lowerCamelCase    )







						@slow
						def       _a    (     self      :       Tuple    ):





											"""simple docstring"""







											for model_name in NAT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																A_ :     int =							NatModel.from_pretrained(_lowerCamelCase    )
																self.assertIsNotNone(_lowerCamelCase    )







						def       _a    (     self      :       List[str]    ):





											"""simple docstring"""







											A_ ,A_ :     List[str] =							self.model_tester.prepare_config_and_inputs_for_common()

											A_ :     Tuple =							_config_zero_init(_lowerCamelCase    )
											for model_class in self.all_model_classes:
																A_ :     Union[str, Any] =							model_class(config=_lowerCamelCase    )
																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     UpperCamelCase_     (unittest.TestCase     ):







						"""simple docstring"""


						@cached_property
						def       _a    (     self      :       List[str]    ):





											"""simple docstring"""







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







						@slow
						def       _a    (     self      :       Optional[Any]    ):





											"""simple docstring"""







											A_ :     Optional[Any] =							NatForImageClassification.from_pretrained('''shi-labs/nat-mini-in1k-224'''    ).to(_lowerCamelCase    )
											A_ :     List[Any] =							self.default_image_processor

											A_ :     Dict =							Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''    )
											A_ :     str =							image_processor(images=_lowerCamelCase							, return_tensors='''pt'''    ).to(_lowerCamelCase    )

											# forward pass
											with torch.no_grad():
																A_ :     Union[str, Any] =							model(**_lowerCamelCase    )

											# verify the logits
											A_ :     int =							torch.Size((1, 1000)    )
											self.assertEqual(outputs.logits.shape							, _lowerCamelCase    )
											A_ :     int =							torch.tensor([0.38_05, -0.86_76, -0.39_12]    ).to(_lowerCamelCase    )
											self.assertTrue(torch.allclose(outputs.logits[0, :3]							, _lowerCamelCase							, atol=1E-4    )    )






@require_torch
@require_natten
class     UpperCamelCase_     (unittest.TestCase,      a__     ):







						"""simple docstring"""


						_lowerCAmelCase    						=					(NatBackbone,) if is_torch_available() else ()
						_lowerCAmelCase    						=					NatConfig







						def       _a    (     self      :       Optional[Any]    ):





											"""simple docstring"""







											A_ :     Union[str, Any] =							NatModelTester(self    )



