



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 	lowerCamelCase__ :


						def __init__( self						:  int      ,       _a						:  Any      ,       _a						:  str=1_3      ,       _a						:  List[Any]=6_4      ,       _a						:  str=4      ,       _a						:  Dict=3      ,       _a						:  str=1_6      ,       _a						:  List[str]=[1, 2, 1]      ,       _a						:  Any=[2, 4, 8]      ,       _a						:  Dict=3      ,       _a						:  Union[str, Any]=2.0      ,       _a						:  Any=True      ,       _a						:  Optional[int]=0.0      ,       _a						:  Dict=0.0      ,       _a						:  Dict=0.1      ,       _a						:  Optional[Any]="gelu"      ,       _a						:  List[Any]=True      ,       _a						:  List[Any]=0.0_2      ,       _a						:  Dict=1e-5      ,       _a						:  Tuple=True      ,       _a						:  List[Any]=None      ,       _a						:  Optional[Any]=True      ,       _a						:  int=1_0      ,       _a						:  Tuple=["stage1", "stage2"]      ,       _a						:  int=[1, 2]      ,       ):
													a__:       Any											=parent
													a__:       Any											=batch_size
													a__:       str											=image_size
													a__:       List[Any]											=patch_size
													a__:       Optional[Any]											=num_channels
													a__:       str											=embed_dim
													a__:       Optional[Any]											=depths
													a__:       Optional[Any]											=num_heads
													a__:       str											=kernel_size
													a__:       int											=mlp_ratio
													a__:       Any											=qkv_bias
													a__:       List[str]											=hidden_dropout_prob
													a__:       Optional[Any]											=attention_probs_dropout_prob
													a__:       Optional[Any]											=drop_path_rate
													a__:       int											=hidden_act
													a__:       List[str]											=patch_norm
													a__:       Optional[int]											=layer_norm_eps
													a__:       int											=initializer_range
													a__:       Dict											=is_training
													a__:       Tuple											=scope
													a__:       Union[str, Any]											=use_labels
													a__:       int											=num_labels
													a__:       Optional[Any]											=out_features
													a__:       List[Any]											=out_indices


						def   _lowerCamelCase			( self						:  Any			):
													a__:       Any											=floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]			)

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

													a__:       List[Any]											=self.get_config()

													return config, pixel_values, labels


						def   _lowerCamelCase			( self						:  Tuple			):
													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   _lowerCamelCase			( self						:  Dict      ,       _a						:  Union[str, Any]      ,       _a						:  List[Any]      ,       _a						:  List[Any]			):
													a__:       str											=NatModel(config=_a			)
													model.to(_a			)
													model.eval()
													a__:       str											=model(_a			)

													a__:       str											=(config.image_size // config.patch_size) // (2 ** (len(config.depths			) - 1))
													a__:       Any											=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   _lowerCamelCase			( self						:  Dict      ,       _a						:  Any      ,       _a						:  Union[str, Any]      ,       _a						:  Union[str, Any]			):
													a__:       List[Any]											=NatForImageClassification(_a			)
													model.to(_a			)
													model.eval()
													a__:       Optional[int]											=model(_a      ,       labels=_a			)
													self.parent.assertEqual(result.logits.shape      ,       (self.batch_size, self.num_labels)			)

													# test greyscale images
													a__:       Union[str, Any]											=1
													a__:       List[str]											=NatForImageClassification(_a			)
													model.to(_a			)
													model.eval()

													a__:       int											=floats_tensor([self.batch_size, 1, self.image_size, self.image_size]			)
													a__:       str											=model(_a			)
													self.parent.assertEqual(result.logits.shape      ,       (self.batch_size, self.num_labels)			)


						def   _lowerCamelCase			( self						:  Tuple      ,       _a						:  Optional[Any]      ,       _a						:  int      ,       _a						:  List[Any]			):
													a__:       List[Any]											=NatBackbone(config=_a			)
													model.to(_a			)
													model.eval()
													a__:       int											=model(_a			)

													# 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], 1_6, 1_6]			)

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

													# verify backbone works with out_features=None
													a__:       Any											=None
													a__:       List[str]											=NatBackbone(config=_a			)
													model.to(_a			)
													model.eval()
													a__:       Tuple											=model(_a			)

													# 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   _lowerCamelCase			( self						:  List[Any]			):
													a__:       Optional[Any]											=self.prepare_config_and_inputs()
													a__:       int											=config_and_inputs
													a__:       List[Any]											={"pixel_values": pixel_values}
													return config, inputs_dict


@require_natten
@require_torch
class 	lowerCamelCase__ (		_a				,					_a				,					unittest.TestCase   ):
						_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   _lowerCamelCase			( self						:  Optional[int]			):
													a__:       List[Any]											=NatModelTester(self			)
													a__:       Dict											=ConfigTester(self      ,       config_class=_a      ,       embed_dim=3_7			)


						def   _lowerCamelCase			( self						:  List[str]			):
													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   _lowerCamelCase			( self						:  List[str]			):
													return


						def   _lowerCamelCase			( self						:  Union[str, Any]			):
													a__:       Tuple											=self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_model(*_a			)


						def   _lowerCamelCase			( self						:  Tuple			):
													a__:       Tuple											=self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_image_classification(*_a			)


						def   _lowerCamelCase			( self						:  int			):
													a__:       int											=self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_backbone(*_a			)


						@unittest.skip(reason="Nat does not use inputs_embeds"			)
						def   _lowerCamelCase			( self						:  Tuple			):
													pass


						@unittest.skip(reason="Nat does not use feedforward chunking"			)
						def   _lowerCamelCase			( self						:  Optional[int]			):
													pass


						def   _lowerCamelCase			( self						:  Union[str, Any]			):
													a__:       str											=self.model_tester.prepare_config_and_inputs_for_common()

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


						def   _lowerCamelCase			( self						:  Optional[Any]			):
													a__:       str											=self.model_tester.prepare_config_and_inputs_for_common()

													for model_class in self.all_model_classes:
																				a__:       Union[str, Any]											=model_class(_a			)
																				a__:       Union[str, Any]											=inspect.signature(model.forward			)
																				# signature.parameters is an OrderedDict => so arg_names order is deterministic
																				a__:       List[Any]											=[*signature.parameters.keys()]

																				a__:       Tuple											=["pixel_values"]
																				self.assertListEqual(arg_names[:1]      ,       _a			)


						def   _lowerCamelCase			( self						:  str			):
													self.skipTest("Nat's attention operation is handled entirely by NATTEN."			)


						def   _lowerCamelCase			( self						:  Any      ,       _a						:  Optional[Any]      ,       _a						:  Any      ,       _a						:  Dict      ,       _a						:  str			):
													a__:       int											=model_class(_a			)
													model.to(_a			)
													model.eval()

													with torch.no_grad():
																				a__:       Tuple											=model(**self._prepare_for_class(_a      ,       _a			)			)

													a__:       str											=outputs.hidden_states

													a__:       List[str]											=getattr(
													    self.model_tester      ,       "expected_num_hidden_layers"      ,       len(self.model_tester.depths			) + 1			)
													self.assertEqual(len(_a			)      ,       _a			)

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

													a__:       Any											=image_size[0] // patch_size[0]
													a__:       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__:       str											=outputs.reshaped_hidden_states
																				self.assertEqual(len(_a			)      ,       _a			)

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


						def   _lowerCamelCase			( self						:  List[Any]			):
													a__:       Tuple											=self.model_tester.prepare_config_and_inputs_for_common()

													a__:       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:
																				a__:       Union[str, Any]											=True
																				self.check_hidden_states_output(_a      ,       _a      ,       _a      ,       _a			)

																				# check that output_hidden_states also work using config
																				del inputs_dict["output_hidden_states"]
																				a__:       Optional[int]											=True

																				self.check_hidden_states_output(_a      ,       _a      ,       _a      ,       _a			)


						@slow
						def   _lowerCamelCase			( self						:  List[str]			):
													for model_name in NAT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																				a__:       int											=NatModel.from_pretrained(_a			)
																				self.assertIsNotNone(_a			)







						def   _lowerCamelCase			( self						:  Optional[int]			):
													a__:       Optional[Any]											=self.model_tester.prepare_config_and_inputs_for_common()

													a__:       List[str]											=_config_zero_init(_a			)
													for model_class in self.all_model_classes:
																				a__:       str											=model_class(config=_a			)
																				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 	lowerCamelCase__ (		unittest.TestCase   ):


						@cached_property
						def   _lowerCamelCase			( self						:  int			):
													return AutoImageProcessor.from_pretrained("shi-labs/nat-mini-in1k-224"			) if is_vision_available() else None







						@slow
						def   _lowerCamelCase			( self						:  int			):
													a__:       List[Any]											=NatForImageClassification.from_pretrained("shi-labs/nat-mini-in1k-224"			).to(_a			)
													a__:       Any											=self.default_image_processor

													a__:       List[str]											=Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"			)
													a__:       Union[str, Any]											=image_processor(images=_a      ,       return_tensors="pt"			).to(_a			)

													# forward pass
													with torch.no_grad():
																				a__:       Dict											=model(**_a			)

													# verify the logits
													a__:       Dict											=torch.Size((1, 1_0_0_0)			)
													self.assertEqual(outputs.logits.shape      ,       _a			)
													a__:       int											=torch.tensor([0.3_8_0_5, -0.8_6_7_6, -0.3_9_1_2]			).to(_a			)
													self.assertTrue(torch.allclose(outputs.logits[0, :3]      ,       _a      ,       atol=1e-4			)			)


@require_torch
@require_natten
class 	lowerCamelCase__ (		unittest.TestCase				,					_a   ):
						_lowerCAmelCase								= (NatBackbone,) if is_torch_available() else ()
						_lowerCAmelCase								= NatConfig


						def   _lowerCamelCase			( self						:  Tuple			):
													a__:       Any											=NatModelTester(self			)











from ...utils import is_transformers_available
from .pipeline_latent_diffusion_superresolution import LDMSuperResolutionPipeline


if is_transformers_available():
		from .pipeline_latent_diffusion import LDMBertModel, LDMTextToImagePipeline





