





'''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       lowerCamelCase       :


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





											'''simple docstring'''





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


					def  SCREAMING_SNAKE_CASE_		(       self						:	Tuple							):





											'''simple docstring'''





											_snake_case: List[Any]	       =		floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]							)

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

											_snake_case: Union[str, Any]	       =		self.get_config()

											return config, pixel_values, labels


					def  SCREAMING_SNAKE_CASE_		(       self						:	Union[str, 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  SCREAMING_SNAKE_CASE_		(       self						:	Dict							,  __snake_case						:	Union[str, Any]							,  __snake_case						:	Tuple							,  __snake_case						:	int							):





											'''simple docstring'''





											_snake_case: Optional[Any]	       =		NatModel(config=__snake_case							)
											model.to(__snake_case							)
											model.eval()
											_snake_case: Any	       =		model(__snake_case							)

											_snake_case: List[Any]	       =		(config.image_size // config.patch_size) // (2 ** (len(config.depths							) - 1))
											_snake_case: Dict	       =		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  SCREAMING_SNAKE_CASE_		(       self						:	Any							,  __snake_case						:	int							,  __snake_case						:	Optional[int]							,  __snake_case						:	str							):





											'''simple docstring'''





											_snake_case: str	       =		NatForImageClassification(__snake_case							)
											model.to(__snake_case							)
											model.eval()
											_snake_case: Any	       =		model(__snake_case							,  labels=__snake_case							)
											self.parent.assertEqual(result.logits.shape							,  (self.batch_size, self.num_labels)							)

											# test greyscale images
											_snake_case: Tuple	       =		1
											_snake_case: Union[str, Any]	       =		NatForImageClassification(__snake_case							)
											model.to(__snake_case							)
											model.eval()

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


					def  SCREAMING_SNAKE_CASE_		(       self						:	Optional[Any]							,  __snake_case						:	str							,  __snake_case						:	Optional[int]							,  __snake_case						:	str							):





											'''simple docstring'''





											_snake_case: Any	       =		NatBackbone(config=__snake_case							)
											model.to(__snake_case							)
											model.eval()
											_snake_case: Dict	       =		model(__snake_case							)

											# 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
											_snake_case: Any	       =		None
											_snake_case: str	       =		NatBackbone(config=__snake_case							)
											model.to(__snake_case							)
											model.eval()
											_snake_case: Optional[int]	       =		model(__snake_case							)

											# 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  SCREAMING_SNAKE_CASE_		(       self						:	int							):





											'''simple docstring'''





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







@require_natten
@require_torch
class       lowerCamelCase       (    __UpperCAmelCase  ,				__UpperCAmelCase  ,				unittest.TestCase	):
					_SCREAMING_SNAKE_CASE				 =     (
					    (
					        NatModel,
					        NatForImageClassification,
					        NatBackbone,
					    )
					    if is_torch_available()
					    else ()
					)
					_SCREAMING_SNAKE_CASE				 =     (
					    {"feature-extraction": NatModel, "image-classification": NatForImageClassification}
					    if is_torch_available()
					    else {}
					)
					_SCREAMING_SNAKE_CASE				 =     False

					_SCREAMING_SNAKE_CASE				 =     False
					_SCREAMING_SNAKE_CASE				 =     False
					_SCREAMING_SNAKE_CASE				 =     False
					_SCREAMING_SNAKE_CASE				 =     False


					def  SCREAMING_SNAKE_CASE_		(       self						:	List[str]							):





											'''simple docstring'''





											_snake_case: str	       =		NatModelTester(self							)
											_snake_case: Tuple	       =		ConfigTester(self							,  config_class=__snake_case							,  embed_dim=37							)


					def  SCREAMING_SNAKE_CASE_		(       self						:	Dict							):





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





											'''simple docstring'''





											return


					def  SCREAMING_SNAKE_CASE_		(       self						:	Any							):





											'''simple docstring'''





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


					def  SCREAMING_SNAKE_CASE_		(       self						:	List[str]							):





											'''simple docstring'''





											_snake_case: Optional[Any]	       =		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_image_classification(*__snake_case							)


					def  SCREAMING_SNAKE_CASE_		(       self						:	Dict							):





											'''simple docstring'''





											_snake_case: Tuple	       =		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_backbone(*__snake_case							)


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





											'''simple docstring'''





											pass


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





											'''simple docstring'''





											pass


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





											'''simple docstring'''





											_snake_case       ,    _snake_case: Optional[Any]	       =		self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																	_snake_case: Any	       =		model_class(__snake_case							)
																	self.assertIsInstance(model.get_input_embeddings()							,  (nn.Module)							)
																	_snake_case: int	       =		model.get_output_embeddings()
																	self.assertTrue(x is None or isinstance(__snake_case							,  nn.Linear							)							)


					def  SCREAMING_SNAKE_CASE_		(       self						:	Optional[Any]							):





											'''simple docstring'''





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

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

																	_snake_case: Tuple	       =		['pixel_values']
																	self.assertListEqual(arg_names[:1]							,  __snake_case							)


					def  SCREAMING_SNAKE_CASE_		(       self						:	List[str]							):





											'''simple docstring'''





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


					def  SCREAMING_SNAKE_CASE_		(       self						:	Dict							,  __snake_case						:	Optional[Any]							,  __snake_case						:	List[str]							,  __snake_case						:	Any							,  __snake_case						:	Union[str, Any]							):





											'''simple docstring'''





											_snake_case: Optional[Any]	       =		model_class(__snake_case							)
											model.to(__snake_case							)
											model.eval()

											with torch.no_grad():
																	_snake_case: Optional[int]	       =		model(**self._prepare_for_class(__snake_case							,  __snake_case							)							)

											_snake_case: Any	       =		outputs.hidden_states

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

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

											_snake_case: Optional[Any]	       =		image_size[0] // patch_size[0]
											_snake_case: 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":
																	_snake_case: List[str]	       =		outputs.reshaped_hidden_states
																	self.assertEqual(len(__snake_case							)							,  __snake_case							)

																	_snake_case       ,    _snake_case       ,    _snake_case       ,    _snake_case: Dict	       =		reshaped_hidden_states[0].shape
																	_snake_case: List[Any]	       =		(
																	    reshaped_hidden_states[0].view(__snake_case							,  __snake_case							,  __snake_case							,  __snake_case							).permute(0							,  2							,  3							,  1							)
																	)
																	self.assertListEqual(
																	    list(reshaped_hidden_states.shape[-3:]							)							,  [height, width, self.model_tester.embed_dim]							,  )


					def  SCREAMING_SNAKE_CASE_		(       self						:	Optional[int]							):





											'''simple docstring'''





											_snake_case       ,    _snake_case: Dict	       =		self.model_tester.prepare_config_and_inputs_for_common()

											_snake_case: 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:
																	_snake_case: Optional[Any]	       =		True
																	self.check_hidden_states_output(__snake_case							,  __snake_case							,  __snake_case							,  __snake_case							)

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

																	self.check_hidden_states_output(__snake_case							,  __snake_case							,  __snake_case							,  __snake_case							)


					@slow
					def  SCREAMING_SNAKE_CASE_		(       self						:	List[Any]							):





											'''simple docstring'''





											for model_name in NAT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																	_snake_case: List[Any]	       =		NatModel.from_pretrained(__snake_case							)
																	self.assertIsNotNone(__snake_case							)




					def  SCREAMING_SNAKE_CASE_		(       self						:	Any							):





											'''simple docstring'''





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

											_snake_case: Tuple	       =		_config_zero_init(__snake_case							)
											for model_class in self.all_model_classes:
																	_snake_case: Union[str, Any]	       =		model_class(config=__snake_case							)
																	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  SCREAMING_SNAKE_CASE_		(       self						:	List[str]							):





											'''simple docstring'''





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




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





											'''simple docstring'''





											_snake_case: List[str]	       =		NatForImageClassification.from_pretrained('shi-labs/nat-mini-in1k-224'							).to(__snake_case							)
											_snake_case: Union[str, Any]	       =		self.default_image_processor

											_snake_case: Optional[int]	       =		Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'							)
											_snake_case: Dict	       =		image_processor(images=__snake_case							,  return_tensors='pt'							).to(__snake_case							)

											# forward pass
											with torch.no_grad():
																	_snake_case: Tuple	       =		model(**__snake_case							)

											# verify the logits
											_snake_case: Optional[int]	       =		torch.Size((1, 10_00)							)
											self.assertEqual(outputs.logits.shape							,  __snake_case							)
											_snake_case: Optional[Any]	       =		torch.tensor([0.3_805, -0.8_676, -0.3_912]							).to(__snake_case							)
											self.assertTrue(torch.allclose(outputs.logits[0, :3]							,  __snake_case							,  atol=1e-4							)							)







@require_torch
@require_natten
class       lowerCamelCase       (    unittest.TestCase  ,				__UpperCAmelCase	):
					_SCREAMING_SNAKE_CASE				 =     (NatBackbone,) if is_torch_available() else ()
					_SCREAMING_SNAKE_CASE				 =     NatConfig


					def  SCREAMING_SNAKE_CASE_		(       self						:	List[Any]							):





											'''simple docstring'''





											_snake_case: int	       =		NatModelTester(self							)












'''simple docstring'''






import tempfile
import unittest

import numpy as np
import torch

from diffusers import DiffusionPipeline
from diffusers.models.attention_processor import Attention, AttnAddedKVProcessor







class       lowerCamelCase       (    unittest.TestCase	):


					def  SCREAMING_SNAKE_CASE_		(       self						:	List[Any]							,  __snake_case						:	bool = False							):





											'''simple docstring'''





											_snake_case: Any	       =		10

											if only_cross_attention:
																	_snake_case: Optional[int]	       =		12
											else:
																	# when only cross attention is not set, the cross attention dim must be the same as the query dim
																	_snake_case: Optional[int]	       =		query_dim

											return {
											    "query_dim": query_dim,
											    "cross_attention_dim": cross_attention_dim,
											    "heads": 2,
											    "dim_head": 4,
											    "added_kv_proj_dim": 6,
											    "norm_num_groups": 1,
											    "only_cross_attention": only_cross_attention,
											    "processor": AttnAddedKVProcessor(),
											}


					def  SCREAMING_SNAKE_CASE_		(       self						:	List[Any]							,  __snake_case						:	List[Any]							,  __snake_case						:	Tuple							):





											'''simple docstring'''





											_snake_case: Optional[int]	       =		2

											_snake_case: Union[str, Any]	       =		torch.rand(__snake_case							,  __snake_case							,  3							,  2							)
											_snake_case: Optional[int]	       =		torch.rand(__snake_case							,  4							,  __snake_case							)
											_snake_case: int	       =		None

											return {
											    "hidden_states": hidden_states,
											    "encoder_hidden_states": encoder_hidden_states,
											    "attention_mask": attention_mask,
											}




					def  SCREAMING_SNAKE_CASE_		(       self						:	List[Any]							):





											'''simple docstring'''





											torch.manual_seed(0							)

											_snake_case: int	       =		self.get_constructor_arguments(only_cross_attention=__snake_case							)
											_snake_case: Union[str, Any]	       =		Attention(**__snake_case							)

											self.assertTrue(attn.to_k is not None							)
											self.assertTrue(attn.to_v is not None							)

											_snake_case: Optional[int]	       =		self.get_forward_arguments(
											    query_dim=constructor_args['query_dim']							,  added_kv_proj_dim=constructor_args['added_kv_proj_dim']							)

											_snake_case: str	       =		attn(**__snake_case							)

											# only self attention

											torch.manual_seed(0							)

											_snake_case: str	       =		self.get_constructor_arguments(only_cross_attention=__snake_case							)
											_snake_case: Optional[Any]	       =		Attention(**__snake_case							)

											self.assertTrue(attn.to_k is None							)
											self.assertTrue(attn.to_v is None							)

											_snake_case: Dict	       =		self.get_forward_arguments(
											    query_dim=constructor_args['query_dim']							,  added_kv_proj_dim=constructor_args['added_kv_proj_dim']							)

											_snake_case: str	       =		attn(**__snake_case							)

											self.assertTrue((only_cross_attn_out != self_and_cross_attn_out).all()							)







class       lowerCamelCase       (    unittest.TestCase	):


					def  SCREAMING_SNAKE_CASE_		(       self						:	Tuple							):





											'''simple docstring'''





											_snake_case: str	       =		DiffusionPipeline.from_pretrained('hf-internal-testing/tiny-stable-diffusion-pipe'							,  safety_checker=__snake_case							)

											_snake_case: int	       =		pipe(
											    'foo'							,  num_inference_steps=2							,  generator=torch.Generator('cpu'							).manual_seed(0							)							,  output_type='np'							,  ).images

											# the initial conversion succeeds
											_snake_case: Dict	       =		DiffusionPipeline.from_pretrained(
											    'hf-internal-testing/tiny-stable-diffusion-pipe'							,  device_map='sequential'							,  safety_checker=__snake_case							)

											_snake_case: Any	       =		pipe(
											    'foo'							,  num_inference_steps=2							,  generator=torch.Generator('cpu'							).manual_seed(0							)							,  output_type='np'							,  ).images

											with tempfile.TemporaryDirectory() as tmpdir:
																	# save the converted model
																	pipe.save_pretrained(__snake_case							)

																	# can also load the converted weights
																	_snake_case: Union[str, Any]	       =		DiffusionPipeline.from_pretrained(__snake_case							,  device_map='sequential'							,  safety_checker=__snake_case							)

											_snake_case: Union[str, Any]	       =		pipe(
											    'foo'							,  num_inference_steps=2							,  generator=torch.Generator('cpu'							).manual_seed(0							)							,  output_type='np'							,  ).images

											self.assertTrue(np.allclose(__snake_case							,  __snake_case							,  atol=1e-5							)							)
											self.assertTrue(np.allclose(__snake_case							,  __snake_case							,  atol=1e-5							)							)





