






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






      def __init__(self					:    Dict		,		_snake_case					:    Tuple		,		_snake_case					:    Tuple=13		,		_snake_case					:    Any=32		,		_snake_case					:    List[str]=1		,		_snake_case					:    List[str]=3		,		_snake_case					:    List[Any]=16		,		_snake_case					:    List[str]=[1, 2, 1]		,		_snake_case					:    Optional[Any]=[2, 2, 4]		,		_snake_case					:    List[str]=2		,		_snake_case					:    Optional[int]=2.0		,		_snake_case					:    Optional[int]=True		,		_snake_case					:    Any=0.0		,		_snake_case					:    Optional[int]=0.0		,		_snake_case					:    Tuple=0.1		,		_snake_case					:    Dict="gelu"		,		_snake_case					:    int=False		,		_snake_case					:    List[Any]=True		,		_snake_case					:    str=0.02		,		_snake_case					:    Tuple=1e-5		,		_snake_case					:    Optional[int]=True		,		_snake_case					:    str=None		,		_snake_case					:    Union[str, Any]=False		,		_snake_case					:    List[Any]=2		,		)       ->   int:


       """simple docstring"""






       lowerCamelCase_							:    int							=     parent
       lowerCamelCase_							:    Union[str, Any]							=     batch_size
       lowerCamelCase_							:    Optional[int]							=     image_size
       lowerCamelCase_							:    Dict							=     patch_size
       lowerCamelCase_							:    Optional[int]							=     num_channels
       lowerCamelCase_							:    Optional[Any]							=     embed_dim
       lowerCamelCase_							:    Optional[Any]							=     depths
       lowerCamelCase_							:    List[str]							=     num_heads
       lowerCamelCase_							:    Dict							=     window_size
       lowerCamelCase_							:    List[str]							=     mlp_ratio
       lowerCamelCase_							:    List[Any]							=     qkv_bias
       lowerCamelCase_							:    int							=     hidden_dropout_prob
       lowerCamelCase_							:    Union[str, Any]							=     attention_probs_dropout_prob
       lowerCamelCase_							:    Tuple							=     drop_path_rate
       lowerCamelCase_							:    int							=     hidden_act
       lowerCamelCase_							:    List[str]							=     use_absolute_embeddings
       lowerCamelCase_							:    Tuple							=     patch_norm
       lowerCamelCase_							:    Any							=     layer_norm_eps
       lowerCamelCase_							:    str							=     initializer_range
       lowerCamelCase_							:    Tuple							=     is_training
       lowerCamelCase_							:    Dict							=     scope
       lowerCamelCase_							:    Tuple							=     use_labels
       lowerCamelCase_							:    Optional[Any]							=     upscale

       # here we set some attributes to make tests pass
       lowerCamelCase_							:    Optional[Any]							=     len(_snake_case		)
       lowerCamelCase_							:    Dict							=     embed_dim
       lowerCamelCase_							:    Union[str, Any]							=     (image_size // patch_size) ** 2






      def      UpperCAmelCase_     (self					:    str		)       ->   List[Any]:


       """simple docstring"""






       lowerCamelCase_							:    str							=     floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]		)

       lowerCamelCase_							:    str							=     None
       if self.use_labels:
        lowerCamelCase_							:    List[str]							=     ids_tensor([self.batch_size]		,		self.type_sequence_label_size		)

       lowerCamelCase_							:    Any							=     self.get_config()

       return config, pixel_values, labels






      def      UpperCAmelCase_     (self					:    Union[str, Any]		)       ->   List[str]:


       """simple docstring"""






       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      UpperCAmelCase_     (self					:    Any		,		_snake_case					:    Optional[int]		,		_snake_case					:    str		,		_snake_case					:    Optional[Any]		)       ->   str:


       """simple docstring"""






       lowerCamelCase_							:    Optional[int]							=     SwinaSRModel(config=_snake_case		)
       model.to(_snake_case		)
       model.eval()
       lowerCamelCase_							:    int							=     model(_snake_case		)

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






      def      UpperCAmelCase_     (self					:    Dict		,		_snake_case					:    int		,		_snake_case					:    Optional[Any]		,		_snake_case					:    Optional[Any]		)       ->   Any:


       """simple docstring"""






       lowerCamelCase_							:    Optional[Any]							=     SwinaSRForImageSuperResolution(_snake_case		)
       model.to(_snake_case		)
       model.eval()
       lowerCamelCase_							:    str							=     model(_snake_case		)

       lowerCamelCase_							:    List[str]							=     self.image_size * self.upscale
       self.parent.assertEqual(
           result.reconstruction.shape		,		(self.batch_size, self.num_channels, expected_image_size, expected_image_size)		)






      def      UpperCAmelCase_     (self					:    Optional[Any]		)       ->   Dict:


       """simple docstring"""






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




@require_torch
class 	lowerCamelCase__				(							UpperCAmelCase,   UpperCAmelCase,   unittest.TestCase ):
      lowerCamelCase_							:   str											=	(SwinaSRModel, SwinaSRForImageSuperResolution) if is_torch_available() else ()
      lowerCamelCase_							:   int											=	{'feature-extraction': SwinaSRModel} if is_torch_available() else {}

      lowerCamelCase_							:   int											=	False
      lowerCamelCase_							:   List[Any]											=	False
      lowerCamelCase_							:   Optional[Any]											=	False
      lowerCamelCase_							:   Dict											=	False
      lowerCamelCase_							:   Union[str, Any]											=	False






      def      UpperCAmelCase_     (self					:    int		)       ->   int:


       """simple docstring"""






       lowerCamelCase_							:    Optional[int]							=     SwinaSRModelTester(self		)
       lowerCamelCase_							:    List[Any]							=     ConfigTester(self		,		config_class=_snake_case		,		embed_dim=37		)






      def      UpperCAmelCase_     (self					:    Any		)       ->   Dict:


       """simple docstring"""






       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					:    str		)       ->   Dict:


       """simple docstring"""






       lowerCamelCase_							:    int							=     self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*_snake_case		)






      def      UpperCAmelCase_     (self					:    int		)       ->   Dict:


       """simple docstring"""






       lowerCamelCase_							:    Optional[Any]							=     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      UpperCAmelCase_     (self					:    List[Any]		)       ->   Any:


       """simple docstring"""






       pass






      @unittest.skip(reason='Swin2SR does not use inputs_embeds'		)
      def      UpperCAmelCase_     (self					:    Tuple		)       ->   Tuple:


       """simple docstring"""






       pass






      @unittest.skip(reason='Swin2SR does not support training yet'		)
      def      UpperCAmelCase_     (self					:    Optional[int]		)       ->   List[Any]:


       """simple docstring"""






       pass






      @unittest.skip(reason='Swin2SR does not support training yet'		)
      def      UpperCAmelCase_     (self					:    Optional[Any]		)       ->   Union[str, Any]:


       """simple docstring"""






       pass






      def      UpperCAmelCase_     (self					:    Union[str, Any]		)       ->   Dict:


       """simple docstring"""






       lowerCamelCase_		,      lowerCamelCase_							:    Optional[int]							=     self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
        lowerCamelCase_							:    Any							=     model_class(_snake_case		)
        self.assertIsInstance(model.get_input_embeddings()		,		(nn.Module)		)
        lowerCamelCase_							:    Union[str, Any]							=     model.get_output_embeddings()
        self.assertTrue(x is None or isinstance(_snake_case		,		nn.Linear		)		)






      def      UpperCAmelCase_     (self					:    Union[str, Any]		)       ->   Dict:


       """simple docstring"""






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

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

        lowerCamelCase_							:    str							=     ['pixel_values']
        self.assertListEqual(arg_names[:1]		,		_snake_case		)






      @slow
      def      UpperCAmelCase_     (self					:    str		)       ->   str:


       """simple docstring"""






       for model_name in SWIN2SR_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        lowerCamelCase_							:    List[str]							=     SwinaSRModel.from_pretrained(_snake_case		)
        self.assertIsNotNone(_snake_case		)






      def      UpperCAmelCase_     (self					:    Optional[int]		)       ->   int:


       """simple docstring"""






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

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


       """simple docstring"""






       lowerCamelCase_		,      lowerCamelCase_							:    int							=     self.model_tester.prepare_config_and_inputs_for_common()
       lowerCamelCase_							:    int							=     True

       for model_class in self.all_model_classes:
        lowerCamelCase_							:    Dict							=     True
        lowerCamelCase_							:    Dict							=     False
        lowerCamelCase_							:    List[str]							=     True
        lowerCamelCase_							:    Dict							=     model_class(_snake_case		)
        model.to(_snake_case		)
        model.eval()
        with torch.no_grad():
         lowerCamelCase_							:    Union[str, Any]							=     model(**self._prepare_for_class(_snake_case		,		_snake_case		)		)
        lowerCamelCase_							:    Optional[int]							=     outputs.attentions
        lowerCamelCase_							:    int							=     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"]
        lowerCamelCase_							:    List[Any]							=     True
        lowerCamelCase_							:    int							=     config.window_size**2
        lowerCamelCase_							:    Union[str, Any]							=     model_class(_snake_case		)
        model.to(_snake_case		)
        model.eval()
        with torch.no_grad():
         lowerCamelCase_							:    Optional[int]							=     model(**self._prepare_for_class(_snake_case		,		_snake_case		)		)
        lowerCamelCase_							:    str							=     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]		,		)
        lowerCamelCase_							:    Dict							=     len(_snake_case		)

        # Check attention is always last and order is fine
        lowerCamelCase_							:    Dict							=     True
        lowerCamelCase_							:    Dict							=     True
        lowerCamelCase_							:    Union[str, Any]							=     model_class(_snake_case		)
        model.to(_snake_case		)
        model.eval()
        with torch.no_grad():
         lowerCamelCase_							:    Tuple							=     model(**self._prepare_for_class(_snake_case		,		_snake_case		)		)

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

        lowerCamelCase_							:    Union[str, Any]							=     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 	lowerCamelCase__				(							unittest.TestCase ):






      def      UpperCAmelCase_     (self					:    List[Any]		)       ->   int:


       """simple docstring"""






       lowerCamelCase_							:    Dict							=     SwinaSRImageProcessor()
       lowerCamelCase_							:    Optional[int]							=     SwinaSRForImageSuperResolution.from_pretrained('caidas/swin2SR-classical-sr-x2-64'		).to(_snake_case		)

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

       # forward pass
       with torch.no_grad():
        lowerCamelCase_							:    List[str]							=     model(**_snake_case		)

       # verify the logits
       lowerCamelCase_							:    Tuple							=     torch.Size([1, 3, 976, 1296]		)
       self.assertEqual(outputs.reconstruction.shape		,		_snake_case		)
       lowerCamelCase_							:    Optional[int]							=     torch.tensor(
           [[0.5458, 0.5546, 0.5638], [0.5526, 0.5565, 0.5651], [0.5396, 0.5426, 0.5621]]		).to(_snake_case		)
       self.assertTrue(torch.allclose(outputs.reconstruction[0, 0, :3, :3]		,		_snake_case		,		atol=1e-4		)		)








from functools import partial
from typing import Any, Callable, Dict, List, Type, TypeVar, Union, overload

import torch
import torch.nn as nn
import torch.types




def       _a   (							lowerCamelCase__ ,			lowerCamelCase__ ,			lowerCamelCase__				)  ->      torch.Tensor:
 # The first operation in a checkpoint can't be in-place, but it's
 # nice to have in-place addition during inference. Thus...
 if not inplace:
  lowerCamelCase_							:    Tuple							=     ma + ma
 else:
  ma += ma

 return ma




def       _a   (							lowerCamelCase__ ,			lowerCamelCase__				)  ->      torch.Tensor:
 lowerCamelCase_							:    Union[str, Any]							=     -1 * len(lowerCamelCase__				)
 lowerCamelCase_							:    Optional[Any]							=     list(range(len(tensor.shape[:zero_index]				)				)				)
 return tensor.permute(first_inds + [zero_index + i for i in inds]				)




def       _a   (							lowerCamelCase__ ,			lowerCamelCase__				)  ->      torch.Tensor:
 return t.reshape(t.shape[:-no_dims] + (-1,)				)




def       _a   (							lowerCamelCase__ ,			lowerCamelCase__ ,			lowerCamelCase__ ,			lowerCamelCase__ = 1e-4				)  ->      torch.Tensor:
 lowerCamelCase_							:    Optional[int]							=     mask.expand(*value.shape				)
 return torch.sum(mask * value ,			dim=lowerCamelCase__				) / (eps + torch.sum(lowerCamelCase__ ,			dim=lowerCamelCase__				))




def       _a   (							lowerCamelCase__ ,			lowerCamelCase__ = 2.3125 ,			lowerCamelCase__ = 21.6875 ,			lowerCamelCase__ = 64				)  ->      torch.Tensor:
 lowerCamelCase_							:    int							=     torch.linspace(lowerCamelCase__ ,			lowerCamelCase__ ,			no_bins - 1 ,			device=pts.device				)
 lowerCamelCase_							:    Tuple							=     torch.sqrt(torch.sum((pts.unsqueeze(-2				) - pts.unsqueeze(-3				)) ** 2 ,			dim=-1				)				)
 return torch.bucketize(lowerCamelCase__ ,			lowerCamelCase__				)




def       _a   (							lowerCamelCase__ ,			lowerCamelCase__				)  ->      dict:
 lowerCamelCase_							:    Optional[Any]							=     dicts[0]
 lowerCamelCase_							:    Optional[Any]							=     {}
 for k, v in first.items():
  lowerCamelCase_							:    Tuple							=     [d[k] for d in dicts]
  if isinstance(lowerCamelCase__ ,			lowerCamelCase__				):
   lowerCamelCase_							:    Optional[int]							=     dict_multimap(lowerCamelCase__ ,			lowerCamelCase__				)
  else:
   lowerCamelCase_							:    List[Any]							=     fn(lowerCamelCase__				)

 return new_dict




def       _a   (							lowerCamelCase__ ,			lowerCamelCase__				)  ->      torch.Tensor:
 lowerCamelCase_							:    str							=     v_bins.view(((1,) * len(x.shape				)) + (len(lowerCamelCase__				),)				)
 lowerCamelCase_							:    List[Any]							=     x[..., None] - reshaped_bins
 lowerCamelCase_							:    List[Any]							=     torch.argmin(torch.abs(lowerCamelCase__				) ,			dim=-1				)
 return nn.functional.one_hot(lowerCamelCase__ ,			num_classes=len(lowerCamelCase__				)				).float()




def       _a   (							lowerCamelCase__ ,			lowerCamelCase__ ,			lowerCamelCase__ = 0 ,			lowerCamelCase__ = 0				)  ->      torch.Tensor:
 lowerCamelCase_							:    List[Union[slice, torch.Tensor]]							=     []
 for i, s in enumerate(data.shape[:no_batch_dims]				):
  lowerCamelCase_							:    int							=     torch.arange(lowerCamelCase__				)
  lowerCamelCase_							:    Any							=     r.view(*(*((1,) * i), -1, *((1,) * (len(inds.shape				) - i - 1)))				)
  ranges.append(lowerCamelCase__				)

 lowerCamelCase_							:    List[Union[slice, torch.Tensor]]							=     [slice(lowerCamelCase__				) for _ in range(len(data.shape				) - no_batch_dims				)]
 lowerCamelCase_							:    Optional[int]							=     inds
 ranges.extend(lowerCamelCase__				)
 # Matt note: Editing this to get around the behaviour of using a list as an array index changing
 # in recent Numpy versions
 return data[tuple(lowerCamelCase__				)]


UpperCamelCase											=      TypeVar('''T''')




def       _a   (							lowerCamelCase__ ,			lowerCamelCase__ ,			lowerCamelCase__				)  ->      Dict[Any, Union[dict, list, tuple, Any]]:
 lowerCamelCase_							:    Dict[Any, Union[dict, list, tuple, Any]]							=     {}
 for k, v in dic.items():
  if isinstance(lowerCamelCase__ ,			lowerCamelCase__				):
   lowerCamelCase_							:    Optional[Any]							=     dict_map(lowerCamelCase__ ,			lowerCamelCase__ ,			lowerCamelCase__				)
  else:
   lowerCamelCase_							:    Dict							=     tree_map(lowerCamelCase__ ,			lowerCamelCase__ ,			lowerCamelCase__				)

 return new_dict




@overload
def       _a   (							lowerCamelCase__ ,			lowerCamelCase__ ,			lowerCamelCase__				)  ->      Any:
 ...




@overload
def       _a   (							lowerCamelCase__ ,			lowerCamelCase__ ,			lowerCamelCase__				)  ->      dict:
 ...




@overload
def       _a   (							lowerCamelCase__ ,			lowerCamelCase__ ,			lowerCamelCase__				)  ->      list:
 ...




@overload
def       _a   (							lowerCamelCase__ ,			lowerCamelCase__ ,			lowerCamelCase__				)  ->      tuple:
 ...




def       _a   (							lowerCamelCase__ ,			lowerCamelCase__ ,			lowerCamelCase__				)  ->      Any:
 if isinstance(lowerCamelCase__ ,			lowerCamelCase__				):
  return dict_map(lowerCamelCase__ ,			lowerCamelCase__ ,			lowerCamelCase__				)
 elif isinstance(lowerCamelCase__ ,			lowerCamelCase__				):
  return [tree_map(lowerCamelCase__ ,			lowerCamelCase__ ,			lowerCamelCase__				) for x in tree]
 elif isinstance(lowerCamelCase__ ,			lowerCamelCase__				):
  return tuple(tree_map(lowerCamelCase__ ,			lowerCamelCase__ ,			lowerCamelCase__				) for x in tree				)
 elif isinstance(lowerCamelCase__ ,			lowerCamelCase__				):
  return fn(lowerCamelCase__				)
 else:
  print(type(lowerCamelCase__				)				)
  raise ValueError('Not supported'				)





UpperCamelCase											=      partial(tree_map, leaf_type=torch.Tensor)
