




"""simple docstring"""

import inspect
import unittest

from transformers import SwinaSRConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import is_torch_available, is_vision_available

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 SwinaSRForImageSuperResolution, SwinaSRModel
						from transformers.models.swinasr.modeling_swinasr import SWIN2SR_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
						from PIL import Image

						from transformers import SwinaSRImageProcessor



class  _snake_case				:

				'''simple docstring'''




				def __init__(							self      :  Optional[Any]					,							snake_case      :  List[Any]					,							snake_case      :  Optional[Any]=13					,							snake_case      :  Tuple=32					,							snake_case      :  int=1					,							snake_case      :  Optional[int]=3					,							snake_case      :  Dict=16					,							snake_case      :  Dict=[1, 2, 1]					,							snake_case      :  Union[str, Any]=[2, 2, 4]					,							snake_case      :  str=2					,							snake_case      :  Optional[Any]=2.0					,							snake_case      :  int=True					,							snake_case      :  List[Any]=0.0					,							snake_case      :  str=0.0					,							snake_case      :  Union[str, Any]=0.1					,							snake_case      :  Any="gelu"					,							snake_case      :  List[str]=False					,							snake_case      :  List[str]=True					,							snake_case      :  Any=0.02					,							snake_case      :  List[Any]=1e-5					,							snake_case      :  List[str]=True					,							snake_case      :  List[str]=None					,							snake_case      :  Union[str, Any]=False					,							snake_case      :  int=2					,							):
								UpperCAmelCase_					:Union[str, Any]      			=      parent
								UpperCAmelCase_					:Union[str, Any]      			=      batch_size
								UpperCAmelCase_					:int      			=      image_size
								UpperCAmelCase_					:int      			=      patch_size
								UpperCAmelCase_					:Any      			=      num_channels
								UpperCAmelCase_					:str      			=      embed_dim
								UpperCAmelCase_					:int      			=      depths
								UpperCAmelCase_					:Tuple      			=      num_heads
								UpperCAmelCase_					:Tuple      			=      window_size
								UpperCAmelCase_					:Optional[Any]      			=      mlp_ratio
								UpperCAmelCase_					:Dict      			=      qkv_bias
								UpperCAmelCase_					:Any      			=      hidden_dropout_prob
								UpperCAmelCase_					:Union[str, Any]      			=      attention_probs_dropout_prob
								UpperCAmelCase_					:Union[str, Any]      			=      drop_path_rate
								UpperCAmelCase_					:str      			=      hidden_act
								UpperCAmelCase_					:int      			=      use_absolute_embeddings
								UpperCAmelCase_					:Any      			=      patch_norm
								UpperCAmelCase_					:Dict      			=      layer_norm_eps
								UpperCAmelCase_					:Optional[Any]      			=      initializer_range
								UpperCAmelCase_					:int      			=      is_training
								UpperCAmelCase_					:Tuple      			=      scope
								UpperCAmelCase_					:List[Any]      			=      use_labels
								UpperCAmelCase_					:Union[str, Any]      			=      upscale

								# here we set some attributes to make tests pass
								UpperCAmelCase_					:Any      			=      len(snake_case     )
								UpperCAmelCase_					:Optional[Any]      			=      embed_dim
								UpperCAmelCase_					:List[Any]      			=      (image_size // patch_size) ** 2





				def 				snake_case_ (							self      :  Dict     ):
								UpperCAmelCase_					:Tuple      			=      floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]     )

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

								UpperCAmelCase_					:Tuple      			=      self.get_config()

								return config, pixel_values, labels





				def 				snake_case_ (							self      :  Dict     ):
								return SwinaSRConfig(
								    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					,							window_size=self.window_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					,							use_absolute_embeddings=self.use_absolute_embeddings					,							path_norm=self.patch_norm					,							layer_norm_eps=self.layer_norm_eps					,							initializer_range=self.initializer_range					,							upscale=self.upscale					,							)





				def 				snake_case_ (							self      :  Any					,							snake_case      :  Union[str, Any]					,							snake_case      :  Dict					,							snake_case      :  Any     ):
								UpperCAmelCase_					:Dict      			=      SwinaSRModel(config=snake_case     )
								model.to(snake_case     )
								model.eval()
								UpperCAmelCase_					:str      			=      model(snake_case     )

								self.parent.assertEqual(
								    result.last_hidden_state.shape					,							(self.batch_size, self.embed_dim, self.image_size, self.image_size)     )





				def 				snake_case_ (							self      :  str					,							snake_case      :  List[str]					,							snake_case      :  Union[str, Any]					,							snake_case      :  List[str]     ):
								UpperCAmelCase_					:Any      			=      SwinaSRForImageSuperResolution(snake_case     )
								model.to(snake_case     )
								model.eval()
								UpperCAmelCase_					:Tuple      			=      model(snake_case     )

								UpperCAmelCase_					:Tuple      			=      self.image_size * self.upscale
								self.parent.assertEqual(
								    result.reconstruction.shape					,							(self.batch_size, self.num_channels, expected_image_size, expected_image_size)     )





				def 				snake_case_ (							self      :  Union[str, Any]     ):
								UpperCAmelCase_					:Any      			=      self.prepare_config_and_inputs()
								UpperCAmelCase_			,UpperCAmelCase_			,UpperCAmelCase_					:List[Any]      			=      config_and_inputs
								UpperCAmelCase_					:str      			=      {'''pixel_values''': pixel_values}
								return config, inputs_dict



@require_torch
class  _snake_case				(							A__							,						A__							,						unittest.TestCase       ):

				'''simple docstring'''




				UpperCamelCase__		     =(SwinaSRModel, SwinaSRForImageSuperResolution) if is_torch_available() else ()
				UpperCamelCase__		     ={"""feature-extraction""": SwinaSRModel} if is_torch_available() else {}

				UpperCamelCase__		     =False
				UpperCamelCase__		     =False
				UpperCamelCase__		     =False
				UpperCamelCase__		     =False
				UpperCamelCase__		     =False





				def 				snake_case_ (							self      :  Any     ):
								UpperCAmelCase_					:Dict      			=      SwinaSRModelTester(self     )
								UpperCAmelCase_					:Dict      			=      ConfigTester(self					,							config_class=snake_case					,							embed_dim=37     )





				def 				snake_case_ (							self      :  List[Any]     ):
								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 				snake_case_ (							self      :  Optional[Any]     ):
								UpperCAmelCase_					:Union[str, Any]      			=      self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*snake_case     )





				def 				snake_case_ (							self      :  Optional[int]     ):
								UpperCAmelCase_					:List[str]      			=      self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_image_super_resolution(*snake_case     )





				@unittest.skip(reason='''Got `CUDA error: misaligned address` with PyTorch 2.0.0.'''     )
				def 				snake_case_ (							self      :  Optional[Any]     ):
								pass





				@unittest.skip(reason='''Swin2SR does not use inputs_embeds'''     )
				def 				snake_case_ (							self      :  List[Any]     ):
								pass





				@unittest.skip(reason='''Swin2SR does not support training yet'''     )
				def 				snake_case_ (							self      :  Any     ):
								pass





				@unittest.skip(reason='''Swin2SR does not support training yet'''     )
				def 				snake_case_ (							self      :  Any     ):
								pass





				def 				snake_case_ (							self      :  int     ):
								UpperCAmelCase_			,UpperCAmelCase_					:Union[str, Any]      			=      self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
												UpperCAmelCase_					:Tuple      			=      model_class(snake_case     )
												self.assertIsInstance(model.get_input_embeddings()					,							(nn.Module)     )
												UpperCAmelCase_					:str      			=      model.get_output_embeddings()
												self.assertTrue(x is None or isinstance(snake_case					,							nn.Linear     )     )





				def 				snake_case_ (							self      :  Optional[Any]     ):
								UpperCAmelCase_			,UpperCAmelCase_					:Union[str, Any]      			=      self.model_tester.prepare_config_and_inputs_for_common()

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

												UpperCAmelCase_					:Optional[int]      			=      ['''pixel_values''']
												self.assertListEqual(arg_names[:1]					,							snake_case     )





				@slow
				def 				snake_case_ (							self      :  Dict     ):
								for model_name in SWIN2SR_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
												UpperCAmelCase_					:int      			=      SwinaSRModel.from_pretrained(snake_case     )
												self.assertIsNotNone(snake_case     )





				def 				snake_case_ (							self      :  Union[str, Any]     ):
								UpperCAmelCase_			,UpperCAmelCase_					:Optional[Any]      			=      self.model_tester.prepare_config_and_inputs_for_common()

								UpperCAmelCase_					:int      			=      _config_zero_init(snake_case     )
								for model_class in self.all_model_classes:
												UpperCAmelCase_					:Dict      			=      model_class(config=snake_case     )
												for name, param in model.named_parameters():
																if "logit_scale" in name:
																				continue
																if 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'					,							)





				def 				snake_case_ (							self      :  Any     ):
								UpperCAmelCase_			,UpperCAmelCase_					:List[Any]      			=      self.model_tester.prepare_config_and_inputs_for_common()
								UpperCAmelCase_					:Optional[int]      			=      True

								for model_class in self.all_model_classes:
												UpperCAmelCase_					:List[str]      			=      True
												UpperCAmelCase_					:str      			=      False
												UpperCAmelCase_					:Optional[int]      			=      True
												UpperCAmelCase_					:Dict      			=      model_class(snake_case     )
												model.to(snake_case     )
												model.eval()
												with torch.no_grad():
																UpperCAmelCase_					:int      			=      model(**self._prepare_for_class(snake_case					,							snake_case     )     )
												UpperCAmelCase_					:int      			=      outputs.attentions
												UpperCAmelCase_					:Union[str, Any]      			=      len(self.model_tester.depths     )
												self.assertEqual(len(snake_case     )					,							snake_case     )

												# check that output_attentions also work using config
												del inputs_dict["output_attentions"]
												UpperCAmelCase_					:List[str]      			=      True
												UpperCAmelCase_					:int      			=      config.window_size**2
												UpperCAmelCase_					:Union[str, Any]      			=      model_class(snake_case     )
												model.to(snake_case     )
												model.eval()
												with torch.no_grad():
																UpperCAmelCase_					:str      			=      model(**self._prepare_for_class(snake_case					,							snake_case     )     )
												UpperCAmelCase_					:Union[str, Any]      			=      outputs.attentions
												self.assertEqual(len(snake_case     )					,							snake_case     )

												self.assertListEqual(
												    list(attentions[0].shape[-3:]     )					,							[self.model_tester.num_heads[0], window_size_squared, window_size_squared]					,							)
												UpperCAmelCase_					:List[Any]      			=      len(snake_case     )

												# Check attention is always last and order is fine
												UpperCAmelCase_					:Dict      			=      True
												UpperCAmelCase_					:Any      			=      True
												UpperCAmelCase_					:Optional[Any]      			=      model_class(snake_case     )
												model.to(snake_case     )
												model.eval()
												with torch.no_grad():
																UpperCAmelCase_					:List[str]      			=      model(**self._prepare_for_class(snake_case					,							snake_case     )     )

												self.assertEqual(out_len + 1					,							len(snake_case     )     )

												UpperCAmelCase_					:str      			=      outputs.attentions

												self.assertEqual(len(snake_case     )					,							snake_case     )

												self.assertListEqual(
												    list(self_attentions[0].shape[-3:]     )					,							[self.model_tester.num_heads[0], window_size_squared, window_size_squared]					,							)



@require_vision
@require_torch
@slow
class  _snake_case				(							unittest.TestCase       ):

				'''simple docstring'''




				def 				snake_case_ (							self      :  str     ):
								UpperCAmelCase_					:str      			=      SwinaSRImageProcessor()
								UpperCAmelCase_					:Any      			=      SwinaSRForImageSuperResolution.from_pretrained('''caidas/swin2SR-classical-sr-x2-64'''     ).to(snake_case     )

								UpperCAmelCase_					:Any      			=      Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''     )
								UpperCAmelCase_					:int      			=      processor(images=snake_case					,							return_tensors='''pt'''     ).to(snake_case     )

								# forward pass
								with torch.no_grad():
												UpperCAmelCase_					:Dict      			=      model(**snake_case     )

								# verify the logits
								UpperCAmelCase_					:Optional[int]      			=      torch.Size([1, 3, 976, 1_296]     )
								self.assertEqual(outputs.reconstruction.shape					,							snake_case     )
								UpperCAmelCase_					:List[Any]      			=      torch.tensor(
								    [[0.5_458, 0.5_546, 0.5_638], [0.5_526, 0.5_565, 0.5_651], [0.5_396, 0.5_426, 0.5_621]]     ).to(snake_case     )
								self.assertTrue(torch.allclose(outputs.reconstruction[0, 0, :3, :3]					,							snake_case					,							atol=1e-4     )     )








"""simple docstring"""

from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


__lowerCamelCase								=						logging.get_logger(__name__)

__lowerCamelCase								=						{
    "mnaylor/mega-base-wikitext": "https://huggingface.co/mnaylor/mega-base-wikitext/resolve/main/config.json",
}



class  _snake_case				(							A__       ):

				'''simple docstring'''




				UpperCamelCase__		     ="""mega"""





				def __init__(							self      :  Optional[int]					,							snake_case      :  Union[str, Any]=30_522					,							snake_case      :  Optional[Any]=128					,							snake_case      :  str=4					,							snake_case      :  Dict=256					,							snake_case      :  Optional[Any]=16					,							snake_case      :  Any=True					,							snake_case      :  str=64					,							snake_case      :  Optional[Any]=False					,							snake_case      :  Union[str, Any]=-1					,							snake_case      :  List[str]=None					,							snake_case      :  Any=True					,							snake_case      :  Optional[int]="scalenorm"					,							snake_case      :  Tuple=True					,							snake_case      :  str="silu"					,							snake_case      :  Dict="softmax"					,							snake_case      :  Optional[int]=0.1					,							snake_case      :  Tuple=0.1					,							snake_case      :  Dict=0.1					,							snake_case      :  List[str]=False					,							snake_case      :  Any=True					,							snake_case      :  List[Any]=256					,							snake_case      :  Optional[Any]=True					,							snake_case      :  Optional[int]=0.1					,							snake_case      :  str=2_048					,							snake_case      :  Union[str, Any]=False					,							snake_case      :  List[Any]=2					,							snake_case      :  str=0.02					,							snake_case      :  str=0.2					,							snake_case      :  str=0.02					,							snake_case      :  Union[str, Any]=1.0					,							snake_case      :  Optional[Any]=1					,							snake_case      :  List[Any]=0					,							snake_case      :  List[Any]=2					,							snake_case      :  List[Any]="rotary"					,							snake_case      :  Tuple=None					,							snake_case      :  int=True					,							snake_case      :  Dict=True					,							**snake_case      :  List[str]					,							):
								super().__init__(pad_token_id=snake_case					,							bos_token_id=snake_case					,							eos_token_id=snake_case					,							**snake_case     )

								UpperCAmelCase_					:List[Any]      			=      vocab_size
								UpperCAmelCase_					:Union[str, Any]      			=      hidden_size
								UpperCAmelCase_					:str      			=      num_hidden_layers
								UpperCAmelCase_					:List[Any]      			=      activation
								UpperCAmelCase_					:int      			=      attention_activation
								UpperCAmelCase_					:Tuple      			=      intermediate_size
								UpperCAmelCase_					:Optional[Any]      			=      ema_projection_size
								UpperCAmelCase_					:Optional[int]      			=      bidirectional
								UpperCAmelCase_					:Optional[int]      			=      shared_representation_size
								UpperCAmelCase_					:Optional[Any]      			=      use_chunking
								UpperCAmelCase_					:List[str]      			=      chunk_size
								UpperCAmelCase_					:Optional[Any]      			=      truncation
								UpperCAmelCase_					:List[str]      			=      normalize_before_mega
								UpperCAmelCase_					:List[Any]      			=      normalization_type
								UpperCAmelCase_					:Tuple      			=      norm_affine
								UpperCAmelCase_					:Union[str, Any]      			=      dropout_prob
								UpperCAmelCase_					:Tuple      			=      hidden_dropout_prob
								UpperCAmelCase_					:List[Any]      			=      attention_probs_dropout_prob
								UpperCAmelCase_					:str      			=      use_feature_dropout
								UpperCAmelCase_					:List[Any]      			=      use_normalized_ffn
								UpperCAmelCase_					:Any      			=      nffn_hidden_size
								UpperCAmelCase_					:List[Any]      			=      normalize_before_ffn
								UpperCAmelCase_					:Optional[int]      			=      nffn_activation_dropout_prob
								UpperCAmelCase_					:str      			=      max_positions
								UpperCAmelCase_					:int      			=      add_token_type_embeddings
								UpperCAmelCase_					:List[Any]      			=      type_vocab_size
								UpperCAmelCase_					:List[str]      			=      initializer_range
								UpperCAmelCase_					:int      			=      ema_delta_alpha_range
								UpperCAmelCase_					:Optional[int]      			=      ema_beta_range
								UpperCAmelCase_					:int      			=      ema_gamma_omega_range
								UpperCAmelCase_					:int      			=      relative_positional_bias
								UpperCAmelCase_					:List[Any]      			=      use_cache
								UpperCAmelCase_					:Optional[Any]      			=      classifier_dropout
								UpperCAmelCase_					:Optional[Any]      			=      add_lm_hidden_dense_layer
								UpperCAmelCase_					:str      			=      1  # not used but required by Hugging Face



class  _snake_case				(							A__       ):

				'''simple docstring'''




				@property
				def 				snake_case_ (							self      :  int     ):
								if self.task == "multiple-choice":
												UpperCAmelCase_					:List[str]      			=      {0: '''batch''', 1: '''choice''', 2: '''sequence'''}
								else:
												UpperCAmelCase_					:Tuple      			=      {0: '''batch''', 1: '''sequence'''}
								return OrderedDict(
								    [
								        ('''input_ids''', dynamic_axis),
								        ('''attention_mask''', dynamic_axis),
								    ]     )


