





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:

  """simple docstring"""







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

     # here we set some attributes to make tests pass
     __A				 =       len(UpperCAmelCase     )
     __A				 =       embed_dim
     __A				 =       (image_size // patch_size) ** 2

  def 		SCREAMING_SNAKE_CASE__	( self     )-> Dict:
     __A				 =       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]     )

     __A				 =       None
     if self.use_labels:
        __A				 =       ids_tensor([self.batch_size]  ,    self.type_sequence_label_size     )

     __A				 =       self.get_config()

     return config, pixel_values, labels

  def 		SCREAMING_SNAKE_CASE__	( self     )-> Union[str, Any]:
     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 		SCREAMING_SNAKE_CASE__	( self  ,    UpperCAmelCase  ,    UpperCAmelCase  ,    UpperCAmelCase     )-> Dict:
     __A				 =       SwinaSRModel(config=UpperCAmelCase     )
     model.to(UpperCAmelCase     )
     model.eval()
     __A				 =       model(UpperCAmelCase     )

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

  def 		SCREAMING_SNAKE_CASE__	( self  ,    UpperCAmelCase  ,    UpperCAmelCase  ,    UpperCAmelCase     )-> Optional[int]:
     __A				 =       SwinaSRForImageSuperResolution(UpperCAmelCase     )
     model.to(UpperCAmelCase     )
     model.eval()
     __A				 =       model(UpperCAmelCase     )

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





  def 		SCREAMING_SNAKE_CASE__	( self     )-> Optional[int]:
     __A				 =       self.prepare_config_and_inputs()
     __A ,						__A ,						__A				 =       config_and_inputs
     __A				 =       {'''pixel_values''': pixel_values}
     return config, inputs_dict







@require_torch
class 						_lowerCAmelCase(						_a		,   _a		,   unittest.TestCase):

  """simple docstring"""







  lowerCamelCase__      					=   (SwinaSRModel, SwinaSRForImageSuperResolution) if is_torch_available() else ()
  lowerCamelCase__      					=   {'''feature-extraction''': SwinaSRModel} if is_torch_available() else {}

  lowerCamelCase__      					=   False
  lowerCamelCase__      					=   False
  lowerCamelCase__      					=   False
  lowerCamelCase__      					=   False
  lowerCamelCase__      					=   False

  def 		SCREAMING_SNAKE_CASE__	( self     )-> Union[str, Any]:
     __A				 =       SwinaSRModelTester(self     )
     __A				 =       ConfigTester(self  ,    config_class=UpperCAmelCase  ,    embed_dim=37     )

  def 		SCREAMING_SNAKE_CASE__	( self     )-> Optional[int]:
     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     )-> int:
     __A				 =       self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*UpperCAmelCase     )

  def 		SCREAMING_SNAKE_CASE__	( self     )-> Optional[Any]:
     __A				 =       self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_image_super_resolution(*UpperCAmelCase     )

  @unittest.skip(reason='''Got `CUDA error: misaligned address` with PyTorch 2.0.0.'''     )
  def 		SCREAMING_SNAKE_CASE__	( self     )-> Dict:
     pass

  @unittest.skip(reason='''Swin2SR does not use inputs_embeds'''     )
  def 		SCREAMING_SNAKE_CASE__	( self     )-> Optional[int]:
     pass

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

  @unittest.skip(reason='''Swin2SR does not support training yet'''     )
  def 		SCREAMING_SNAKE_CASE__	( self     )-> Dict:
     pass

  def 		SCREAMING_SNAKE_CASE__	( self     )-> Dict:
     __A ,						__A				 =       self.model_tester.prepare_config_and_inputs_for_common()

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

  def 		SCREAMING_SNAKE_CASE__	( self     )-> Dict:
     __A ,						__A				 =       self.model_tester.prepare_config_and_inputs_for_common()

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

        __A				 =       ['''pixel_values''']
        self.assertListEqual(arg_names[:1]  ,    UpperCAmelCase     )

  @slow
  def 		SCREAMING_SNAKE_CASE__	( self     )-> List[str]:
     for model_name in SWIN2SR_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        __A				 =       SwinaSRModel.from_pretrained(UpperCAmelCase     )
        self.assertIsNotNone(UpperCAmelCase     )

  def 		SCREAMING_SNAKE_CASE__	( self     )-> Tuple:
     __A ,						__A				 =       self.model_tester.prepare_config_and_inputs_for_common()

     __A				 =       _config_zero_init(UpperCAmelCase     )
     for model_class in self.all_model_classes:
        __A				 =       model_class(config=UpperCAmelCase     )
        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 		SCREAMING_SNAKE_CASE__	( self     )-> List[Any]:
     __A ,						__A				 =       self.model_tester.prepare_config_and_inputs_for_common()
     __A				 =       True

     for model_class in self.all_model_classes:
        __A				 =       True
        __A				 =       False
        __A				 =       True
        __A				 =       model_class(UpperCAmelCase     )
        model.to(UpperCAmelCase     )
        model.eval()
        with torch.no_grad():
           __A				 =       model(**self._prepare_for_class(UpperCAmelCase  ,    UpperCAmelCase     )     )
        __A				 =       outputs.attentions
        __A				 =       len(self.model_tester.depths     )
        self.assertEqual(len(UpperCAmelCase     )  ,    UpperCAmelCase     )

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

        self.assertListEqual(
            list(attentions[0].shape[-3:]     )  ,    [self.model_tester.num_heads[0], window_size_squared, window_size_squared]  ,    )
        __A				 =       len(UpperCAmelCase     )

        # Check attention is always last and order is fine
        __A				 =       True
        __A				 =       True
        __A				 =       model_class(UpperCAmelCase     )
        model.to(UpperCAmelCase     )
        model.eval()
        with torch.no_grad():
           __A				 =       model(**self._prepare_for_class(UpperCAmelCase  ,    UpperCAmelCase     )     )

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

        __A				 =       outputs.attentions

        self.assertEqual(len(UpperCAmelCase     )  ,    UpperCAmelCase     )

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

  """simple docstring"""







  def 		SCREAMING_SNAKE_CASE__	( self     )-> Dict:
     __A				 =       SwinaSRImageProcessor()
     __A				 =       SwinaSRForImageSuperResolution.from_pretrained('''caidas/swin2SR-classical-sr-x2-64'''     ).to(UpperCAmelCase     )

     __A				 =       Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''     )
     __A				 =       processor(images=UpperCAmelCase  ,    return_tensors='''pt'''     ).to(UpperCAmelCase     )

     # forward pass
     with torch.no_grad():
        __A				 =       model(**UpperCAmelCase     )

     # verify the logits
     __A				 =       torch.Size([1, 3, 9_76, 12_96]     )
     self.assertEqual(outputs.reconstruction.shape  ,    UpperCAmelCase     )
     __A				 =       torch.tensor(
         [[0.5458, 0.5546, 0.5638], [0.5526, 0.5565, 0.5651], [0.5396, 0.5426, 0.5621]]     ).to(UpperCAmelCase     )
     self.assertTrue(torch.allclose(outputs.reconstruction[0, 0, :3, :3]  ,    UpperCAmelCase  ,    atol=1e-4     )     )








from ...configuration_utils import PretrainedConfig
from ...utils import logging


_UpperCamelCase   :							Dict     =	logging.get_logger(__name__)

_UpperCamelCase   :							Optional[int]     =	{
    """MCG-NJU/videomae-base""": """https://huggingface.co/MCG-NJU/videomae-base/resolve/main/config.json""",
}







class 						_lowerCAmelCase(						_a):

  """simple docstring"""







  lowerCamelCase__      					=   '''videomae'''

  def __init__( self  ,    UpperCAmelCase=2_24  ,    UpperCAmelCase=16  ,    UpperCAmelCase=3  ,    UpperCAmelCase=16  ,    UpperCAmelCase=2  ,    UpperCAmelCase=7_68  ,    UpperCAmelCase=12  ,    UpperCAmelCase=12  ,    UpperCAmelCase=30_72  ,    UpperCAmelCase="gelu"  ,    UpperCAmelCase=0.0  ,    UpperCAmelCase=0.0  ,    UpperCAmelCase=0.02  ,    UpperCAmelCase=1e-12  ,    UpperCAmelCase=True  ,    UpperCAmelCase=True  ,    UpperCAmelCase=6  ,    UpperCAmelCase=3_84  ,    UpperCAmelCase=4  ,    UpperCAmelCase=15_36  ,    UpperCAmelCase=True  ,    **UpperCAmelCase  ,    )-> Any:
     super().__init__(**UpperCAmelCase     )

     __A				 =       image_size
     __A				 =       patch_size
     __A				 =       num_channels
     __A				 =       num_frames
     __A				 =       tubelet_size

     __A				 =       hidden_size
     __A				 =       num_hidden_layers
     __A				 =       num_attention_heads
     __A				 =       intermediate_size
     __A				 =       hidden_act
     __A				 =       hidden_dropout_prob
     __A				 =       attention_probs_dropout_prob
     __A				 =       initializer_range
     __A				 =       layer_norm_eps
     __A				 =       qkv_bias
     __A				 =       use_mean_pooling

     __A				 =       decoder_num_attention_heads
     __A				 =       decoder_hidden_size
     __A				 =       decoder_num_hidden_layers
     __A				 =       decoder_intermediate_size
     __A				 =       norm_pix_loss

