



'''simple docstring'''



import inspect
import unittest

from transformers import BitConfig
from transformers.testing_utils import 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, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
	import torch
	from torch import nn

	from transformers import BitBackbone, BitForImageClassification, BitImageProcessor, BitModel
	from transformers.models.bit.modeling_bit import BIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
	from PIL import Image






class 					UpperCAmelCase__	:



					def __init__(    self			,				UpperCamelCase			,				UpperCamelCase=3			,				UpperCamelCase=32			,				UpperCamelCase=3			,				UpperCamelCase=10			,				UpperCamelCase=[8, 16, 32, 64]			,				UpperCamelCase=[1, 1, 2, 1]			,				UpperCamelCase=True			,				UpperCamelCase=True			,				UpperCamelCase="relu"			,				UpperCamelCase=3			,				UpperCamelCase=None			,				UpperCamelCase=["stage2", "stage3", "stage4"]			,				UpperCamelCase=[2, 3, 4]			,				UpperCamelCase=1			,				)     ->			List[Any]:
						__lowerCAmelCase												=  parent
						__lowerCAmelCase												=  batch_size
						__lowerCAmelCase												=  image_size
						__lowerCAmelCase												=  num_channels
						__lowerCAmelCase												=  embeddings_size
						__lowerCAmelCase												=  hidden_sizes
						__lowerCAmelCase												=  depths
						__lowerCAmelCase												=  is_training
						__lowerCAmelCase												=  use_labels
						__lowerCAmelCase												=  hidden_act
						__lowerCAmelCase												=  num_labels
						__lowerCAmelCase												=  scope
						__lowerCAmelCase												=  len(UpperCamelCase     )
						__lowerCAmelCase												=  out_features
						__lowerCAmelCase												=  out_indices
						__lowerCAmelCase												=  num_groups



					def        UpperCAmelCase_ (    self     )     ->			List[str]:
						__lowerCAmelCase												=  floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]     )

						__lowerCAmelCase												=  None
						if self.use_labels:
							__lowerCAmelCase												=  ids_tensor([self.batch_size]			,				self.num_labels     )

						__lowerCAmelCase												=  self.get_config()

						return config, pixel_values, labels



					def        UpperCAmelCase_ (    self     )     ->			Dict:
						return BitConfig(
						    num_channels=self.num_channels			,				embeddings_size=self.embeddings_size			,				hidden_sizes=self.hidden_sizes			,				depths=self.depths			,				hidden_act=self.hidden_act			,				num_labels=self.num_labels			,				out_features=self.out_features			,				out_indices=self.out_indices			,				num_groups=self.num_groups			,				)



					def        UpperCAmelCase_ (    self			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase     )     ->			int:
						__lowerCAmelCase												=  BitModel(config=UpperCamelCase     )
						model.to(UpperCamelCase     )
						model.eval()
						__lowerCAmelCase												=  model(UpperCamelCase     )
						self.parent.assertEqual(
						    result.last_hidden_state.shape			,				(self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)			,				)



					def        UpperCAmelCase_ (    self			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase     )     ->			List[str]:
						__lowerCAmelCase												=  self.num_labels
						__lowerCAmelCase												=  BitForImageClassification(UpperCamelCase     )
						model.to(UpperCamelCase     )
						model.eval()
						__lowerCAmelCase												=  model(UpperCamelCase			,				labels=UpperCamelCase     )
						self.parent.assertEqual(result.logits.shape			,				(self.batch_size, self.num_labels)     )



					def        UpperCAmelCase_ (    self			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase     )     ->			Dict:
						__lowerCAmelCase												=  BitBackbone(config=UpperCamelCase     )
						model.to(UpperCamelCase     )
						model.eval()
						__lowerCAmelCase												=  model(UpperCamelCase     )

						# verify feature maps
						self.parent.assertEqual(len(result.feature_maps     )			,				len(config.out_features     )     )
						self.parent.assertListEqual(list(result.feature_maps[0].shape     )			,				[self.batch_size, self.hidden_sizes[1], 4, 4]     )

						# verify channels
						self.parent.assertEqual(len(model.channels     )			,				len(config.out_features     )     )
						self.parent.assertListEqual(model.channels			,				config.hidden_sizes[1:]     )

						# verify backbone works with out_features=None
						__lowerCAmelCase												=  None
						__lowerCAmelCase												=  BitBackbone(config=UpperCamelCase     )
						model.to(UpperCamelCase     )
						model.eval()
						__lowerCAmelCase												=  model(UpperCamelCase     )

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

						# verify channels
						self.parent.assertEqual(len(model.channels     )			,				1     )
						self.parent.assertListEqual(model.channels			,				[config.hidden_sizes[-1]]     )



					def        UpperCAmelCase_ (    self     )     ->			Optional[Any]:
						__lowerCAmelCase												=  self.prepare_config_and_inputs()
						__lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase												=  config_and_inputs
						__lowerCAmelCase												=  {"pixel_values": pixel_values}
						return config, inputs_dict






@require_torch
class 					UpperCAmelCase__	(		UpperCamelCase__							,       UpperCamelCase__							,       unittest.TestCase ):
					a      :     Any        =							(BitModel, BitForImageClassification, BitBackbone) if is_torch_available() else ()
					a      :     int        =							(
					    {"""feature-extraction""": BitModel, """image-classification""": BitForImageClassification}
					    if is_torch_available()
					    else {}
					)

					a      :     Dict        =							False
					a      :     Optional[Any]        =							False
					a      :     str        =							False
					a      :     List[Any]        =							False
					a      :     str        =							False



					def        UpperCAmelCase_ (    self     )     ->			Optional[Any]:
						__lowerCAmelCase												=  BitModelTester(self     )
						__lowerCAmelCase												=  ConfigTester(self			,				config_class=UpperCamelCase			,				has_text_modality=UpperCamelCase     )



					def        UpperCAmelCase_ (    self     )     ->			Any:
						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     )     ->			Dict:
						return



					@unittest.skip(reason="Bit does not output attentions"     )
					def        UpperCAmelCase_ (    self     )     ->			int:
						pass



					@unittest.skip(reason="Bit does not use inputs_embeds"     )
					def        UpperCAmelCase_ (    self     )     ->			int:
						pass



					@unittest.skip(reason="Bit does not support input and output embeddings"     )
					def        UpperCAmelCase_ (    self     )     ->			List[Any]:
						pass



					def        UpperCAmelCase_ (    self     )     ->			Optional[Any]:
						__lowerCAmelCase							,       __lowerCAmelCase												=  self.model_tester.prepare_config_and_inputs_for_common()

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

							__lowerCAmelCase												=  ["pixel_values"]
							self.assertListEqual(arg_names[:1]			,				UpperCamelCase     )



					def        UpperCAmelCase_ (    self     )     ->			Any:
						__lowerCAmelCase												=  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*UpperCamelCase     )



					def        UpperCAmelCase_ (    self     )     ->			Union[str, Any]:
						__lowerCAmelCase												=  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_backbone(*UpperCamelCase     )



					def        UpperCAmelCase_ (    self     )     ->			List[Any]:
						__lowerCAmelCase							,       __lowerCAmelCase												=  self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
							__lowerCAmelCase												=  model_class(config=UpperCamelCase     )
							for name, module in model.named_modules():
								if isinstance(UpperCamelCase			,				(nn.BatchNormad, nn.GroupNorm)     ):
									self.assertTrue(
									    torch.all(module.weight == 1     )			,				msg=F'''Parameter {name} of model {model_class} seems not properly initialized'''			,				)
									self.assertTrue(
									    torch.all(module.bias == 0     )			,				msg=F'''Parameter {name} of model {model_class} seems not properly initialized'''			,				)



					def        UpperCAmelCase_ (    self     )     ->			int:
						def check_hidden_states_output(UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase     ):
							__lowerCAmelCase												=  model_class(UpperCamelCase     )
							model.to(UpperCamelCase     )
							model.eval()

							with torch.no_grad():
								__lowerCAmelCase												=  model(**self._prepare_for_class(UpperCamelCase			,				UpperCamelCase     )     )

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

							__lowerCAmelCase												=  self.model_tester.num_stages
							self.assertEqual(len(UpperCamelCase     )			,				expected_num_stages + 1     )

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

						__lowerCAmelCase							,       __lowerCAmelCase												=  self.model_tester.prepare_config_and_inputs_for_common()
						__lowerCAmelCase												=  ["preactivation", "bottleneck"]
						for model_class in self.all_model_classes:
							for layer_type in layers_type:
								__lowerCAmelCase												=  layer_type
								__lowerCAmelCase												=  True
								check_hidden_states_output(UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase     )

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

								check_hidden_states_output(UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase     )



					@unittest.skip(reason="Bit does not use feedforward chunking"     )
					def        UpperCAmelCase_ (    self     )     ->			List[Any]:
						pass



					def        UpperCAmelCase_ (    self     )     ->			List[str]:
						__lowerCAmelCase												=  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_image_classification(*UpperCamelCase     )



					@slow
					def        UpperCAmelCase_ (    self     )     ->			Tuple:
						for model_name in BIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							__lowerCAmelCase												=  BitModel.from_pretrained(UpperCamelCase     )
							self.assertIsNotNone(UpperCamelCase     )



def   __lowerCAmelCase					(							):






	'''simple docstring'''




	__lowerCAmelCase												=  Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"					)
	return image






@require_torch
@require_vision
class 					UpperCAmelCase__	(		unittest.TestCase ):



					@cached_property
					def        UpperCAmelCase_ (    self     )     ->			str:
						return (
						    BitImageProcessor.from_pretrained(BIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]     ) if is_vision_available() else None
						)



					@slow
					def        UpperCAmelCase_ (    self     )     ->			Union[str, Any]:
						__lowerCAmelCase												=  BitForImageClassification.from_pretrained(BIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]     ).to(UpperCamelCase     )

						__lowerCAmelCase												=  self.default_image_processor
						__lowerCAmelCase												=  prepare_img()
						__lowerCAmelCase												=  image_processor(images=UpperCamelCase			,				return_tensors="pt"     ).to(UpperCamelCase     )

						# forward pass
						with torch.no_grad():
							__lowerCAmelCase												=  model(**UpperCamelCase     )

						# verify the logits
						__lowerCAmelCase												=  torch.Size((1, 1000)     )
						self.assertEqual(outputs.logits.shape			,				UpperCamelCase     )

						__lowerCAmelCase												=  torch.tensor([[-0.65_26, -0.52_63, -1.43_98]]     ).to(UpperCamelCase     )

						self.assertTrue(torch.allclose(outputs.logits[0, :3]			,				UpperCamelCase			,				atol=1E-4     )     )






@require_torch
class 					UpperCAmelCase__	(		UpperCamelCase__							,       unittest.TestCase ):
					a      :     Any        =							(BitBackbone,) if is_torch_available() else ()
					a      :     str        =							BitConfig

					a      :     Optional[Any]        =							False



					def        UpperCAmelCase_ (    self     )     ->			str:
						__lowerCAmelCase												=  BitModelTester(self     )




'''simple docstring'''



from dataclasses import dataclass
from typing import Optional

import torch
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .attention import BasicTransformerBlock
from .modeling_utils import ModelMixin






@dataclass
class 					UpperCAmelCase__	(		UpperCamelCase__ ):
					a      :     torch.FloatTensor






class 					UpperCAmelCase__	(		UpperCamelCase__							,       UpperCamelCase__ ):



					@register_to_config
					def __init__(    self			,				UpperCamelCase = 16			,				UpperCamelCase = 88			,				UpperCamelCase = None			,				UpperCamelCase = None			,				UpperCamelCase = 1			,				UpperCamelCase = 0.0			,				UpperCamelCase = 32			,				UpperCamelCase = None			,				UpperCamelCase = False			,				UpperCamelCase = None			,				UpperCamelCase = "geglu"			,				UpperCamelCase = True			,				UpperCamelCase = True			,				)     ->			List[str]:
						super().__init__()
						__lowerCAmelCase												=  num_attention_heads
						__lowerCAmelCase												=  attention_head_dim
						__lowerCAmelCase												=  num_attention_heads * attention_head_dim

						__lowerCAmelCase												=  in_channels

						__lowerCAmelCase												=  torch.nn.GroupNorm(num_groups=UpperCamelCase			,				num_channels=UpperCamelCase			,				eps=1E-6			,				affine=UpperCamelCase     )
						__lowerCAmelCase												=  nn.Linear(UpperCamelCase			,				UpperCamelCase     )

						# 3. Define transformers blocks
						__lowerCAmelCase												=  nn.ModuleList(
						    [
						        BasicTransformerBlock(
						            UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				dropout=UpperCamelCase			,				cross_attention_dim=UpperCamelCase			,				activation_fn=UpperCamelCase			,				attention_bias=UpperCamelCase			,				double_self_attention=UpperCamelCase			,				norm_elementwise_affine=UpperCamelCase			,				)
						        for d in range(UpperCamelCase     )
						    ]     )

						__lowerCAmelCase												=  nn.Linear(UpperCamelCase			,				UpperCamelCase     )



					def        UpperCAmelCase_ (    self			,				UpperCamelCase			,				UpperCamelCase=None			,				UpperCamelCase=None			,				UpperCamelCase=None			,				UpperCamelCase=1			,				UpperCamelCase=None			,				UpperCamelCase = True			,				)     ->			List[str]:
						__lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase												=  hidden_states.shape
						__lowerCAmelCase												=  batch_frames // num_frames

						__lowerCAmelCase												=  hidden_states

						__lowerCAmelCase												=  hidden_states[None, :].reshape(UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase     )
						__lowerCAmelCase												=  hidden_states.permute(0			,				2			,				1			,				3			,				4     )

						__lowerCAmelCase												=  self.norm(UpperCamelCase     )
						__lowerCAmelCase												=  hidden_states.permute(0			,				3			,				4			,				2			,				1     ).reshape(batch_size * height * width			,				UpperCamelCase			,				UpperCamelCase     )

						__lowerCAmelCase												=  self.proj_in(UpperCamelCase     )

						# 2. Blocks
						for block in self.transformer_blocks:
							__lowerCAmelCase												=  block(
							    UpperCamelCase			,				encoder_hidden_states=UpperCamelCase			,				timestep=UpperCamelCase			,				cross_attention_kwargs=UpperCamelCase			,				class_labels=UpperCamelCase			,				)

						# 3. Output
						__lowerCAmelCase												=  self.proj_out(UpperCamelCase     )
						__lowerCAmelCase												=  (
						    hidden_states[None, None, :]
						    .reshape(UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase     )
						    .permute(0			,				3			,				4			,				1			,				2     )
						    .contiguous()
						)
						__lowerCAmelCase												=  hidden_states.reshape(UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase			,				UpperCamelCase     )

						__lowerCAmelCase												=  hidden_states + residual

						if not return_dict:
							return (output,)

						return TransformerTemporalModelOutput(sample=UpperCamelCase     )