

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

  '''simple docstring'''





  def __init__(     self   , a_   , a_=1_3   , a_=3_2   , a_=1   , a_=3   , a_=1_6   , a_=[1, 2, 1]   , a_=[2, 2, 4]   , a_=2   , a_=2.0   , a_=True   , a_=0.0   , a_=0.0   , a_=0.1   , a_="gelu"   , a_=False   , a_=True   , a_=0.02   , a_=1E-5   , a_=True   , a_=None   , a_=False   , a_=2   , )							->				Union[str, Any]:


    """simple docstring"""
    UpperCAmelCase													=	parent
    UpperCAmelCase													=	batch_size
    UpperCAmelCase													=	image_size
    UpperCAmelCase													=	patch_size
    UpperCAmelCase													=	num_channels
    UpperCAmelCase													=	embed_dim
    UpperCAmelCase													=	depths
    UpperCAmelCase													=	num_heads
    UpperCAmelCase													=	window_size
    UpperCAmelCase													=	mlp_ratio
    UpperCAmelCase													=	qkv_bias
    UpperCAmelCase													=	hidden_dropout_prob
    UpperCAmelCase													=	attention_probs_dropout_prob
    UpperCAmelCase													=	drop_path_rate
    UpperCAmelCase													=	hidden_act
    UpperCAmelCase													=	use_absolute_embeddings
    UpperCAmelCase													=	patch_norm
    UpperCAmelCase													=	layer_norm_eps
    UpperCAmelCase													=	initializer_range
    UpperCAmelCase													=	is_training
    UpperCAmelCase													=	scope
    UpperCAmelCase													=	use_labels
    UpperCAmelCase													=	upscale

    # here we set some attributes to make tests pass
    UpperCAmelCase													=	len(lowerCamelCase__     )
    UpperCAmelCase													=	embed_dim
    UpperCAmelCase													=	(image_size // patch_size) ** 2
  def 						snake_case_			(     self     )							->				Optional[int]:


    """simple docstring"""
    UpperCAmelCase													=	floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]     )

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

    UpperCAmelCase													=	self.get_config()

    return config, pixel_values, labels
  def 						snake_case_			(     self     )							->				List[Any]:


    """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 						snake_case_			(     self   , a_   , a_   , a_     )							->				List[str]:


    """simple docstring"""
    UpperCAmelCase													=	SwinaSRModel(config=lowerCamelCase__     )
    model.to(lowerCamelCase__     )
    model.eval()
    UpperCAmelCase													=	model(lowerCamelCase__     )

    self.parent.assertEqual(
        result.last_hidden_state.shape   , (self.batch_size, self.embed_dim, self.image_size, self.image_size)     )
  def 						snake_case_			(     self   , a_   , a_   , a_     )							->				Optional[int]:


    """simple docstring"""
    UpperCAmelCase													=	SwinaSRForImageSuperResolution(lowerCamelCase__     )
    model.to(lowerCamelCase__     )
    model.eval()
    UpperCAmelCase													=	model(lowerCamelCase__     )

    UpperCAmelCase													=	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     )							->				List[Any]:


    """simple docstring"""
    UpperCAmelCase													=	self.prepare_config_and_inputs()
    UpperCAmelCase													=	config_and_inputs
    UpperCAmelCase													=	{"pixel_values": pixel_values}
    return config, inputs_dict




@require_torch
class        lowercase_	(		UpperCAmelCase__ ,       UpperCAmelCase__ ,       unittest.TestCase ):

  '''simple docstring'''





  __lowerCAmelCase					: int      =   (SwinaSRModel, SwinaSRForImageSuperResolution) if is_torch_available() else ()
  __lowerCAmelCase					: List[str]      =   {'feature-extraction': SwinaSRModel} if is_torch_available() else {}

  __lowerCAmelCase					: Tuple      =   False
  __lowerCAmelCase					: Any      =   False
  __lowerCAmelCase					: List[str]      =   False
  __lowerCAmelCase					: Tuple      =   False
  __lowerCAmelCase					: Any      =   False
  def 						snake_case_			(     self     )							->				int:


    """simple docstring"""
    UpperCAmelCase													=	SwinaSRModelTester(self     )
    UpperCAmelCase													=	ConfigTester(self   , config_class=lowerCamelCase__   , embed_dim=3_7     )
  def 						snake_case_			(     self     )							->				List[str]:


    """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 						snake_case_			(     self     )							->				str:


    """simple docstring"""
    UpperCAmelCase													=	self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_model(*lowerCamelCase__     )
  def 						snake_case_			(     self     )							->				Union[str, Any]:


    """simple docstring"""
    UpperCAmelCase													=	self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_image_super_resolution(*lowerCamelCase__     )
  @unittest.skip(reason='Got `CUDA error: misaligned address` with PyTorch 2.0.0.'     )
  def 						snake_case_			(     self     )							->				Dict:


    """simple docstring"""
    pass
  @unittest.skip(reason='Swin2SR does not use inputs_embeds'     )
  def 						snake_case_			(     self     )							->				Dict:


    """simple docstring"""
    pass
  @unittest.skip(reason='Swin2SR does not support training yet'     )
  def 						snake_case_			(     self     )							->				List[Any]:


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


    """simple docstring"""
    pass
  def 						snake_case_			(     self     )							->				int:


    """simple docstring"""
    UpperCAmelCase													=	self.model_tester.prepare_config_and_inputs_for_common()

    for model_class in self.all_model_classes:
      UpperCAmelCase													=	model_class(lowerCamelCase__     )
      self.assertIsInstance(model.get_input_embeddings()   , (nn.Module)     )
      UpperCAmelCase													=	model.get_output_embeddings()
      self.assertTrue(x is None or isinstance(lowerCamelCase__   , nn.Linear     )     )
  def 						snake_case_			(     self     )							->				List[Any]:


    """simple docstring"""
    UpperCAmelCase													=	self.model_tester.prepare_config_and_inputs_for_common()

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

      UpperCAmelCase													=	["pixel_values"]
      self.assertListEqual(arg_names[:1]   , lowerCamelCase__     )
  @slow
  def 						snake_case_			(     self     )							->				str:


    """simple docstring"""
    for model_name in SWIN2SR_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
      UpperCAmelCase													=	SwinaSRModel.from_pretrained(lowerCamelCase__     )
      self.assertIsNotNone(lowerCamelCase__     )
  def 						snake_case_			(     self     )							->				Optional[Any]:


    """simple docstring"""
    UpperCAmelCase													=	self.model_tester.prepare_config_and_inputs_for_common()

    UpperCAmelCase													=	_config_zero_init(lowerCamelCase__     )
    for model_class in self.all_model_classes:
      UpperCAmelCase													=	model_class(config=lowerCamelCase__     )
      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:


    """simple docstring"""
    UpperCAmelCase													=	self.model_tester.prepare_config_and_inputs_for_common()
    UpperCAmelCase													=	True

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

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

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

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

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

      UpperCAmelCase													=	outputs.attentions

      self.assertEqual(len(lowerCamelCase__     )   , lowerCamelCase__     )

      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        lowercase_	(		unittest.TestCase ):

  '''simple docstring'''





  def 						snake_case_			(     self     )							->				Tuple:


    """simple docstring"""
    UpperCAmelCase													=	SwinaSRImageProcessor()
    UpperCAmelCase													=	SwinaSRForImageSuperResolution.from_pretrained('caidas/swin2SR-classical-sr-x2-64'     ).to(lowerCamelCase__     )

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

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

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





'''simple docstring'''




from typing import List, Optional, Tuple, Union

from . import logging
from .import_utils import is_torch_available, is_torch_version


if is_torch_available():
  import torch

_a				: List[Any] 							=						logging.get_logger(__name__)  # pylint: disable=invalid-name

try:
  from torch._dynamo import allow_in_graph as maybe_allow_in_graph
except (ImportError, ModuleNotFoundError):


  def       lowerCamelCase__		(						cls      :   Dict						):
    return cls


def       lowerCamelCase__		(						SCREAMING_SNAKE_CASE      :   Union[Tuple, List]       ,   SCREAMING_SNAKE_CASE      :   Optional[Union[List["torch.Generator"], "torch.Generator"]] = None       ,   SCREAMING_SNAKE_CASE      :   Optional["torch.device"] = None       ,   SCREAMING_SNAKE_CASE      :   Optional["torch.dtype"] = None       ,   SCREAMING_SNAKE_CASE      :   Optional["torch.layout"] = None       ,   ):
  UpperCAmelCase													=	device
  UpperCAmelCase													=	shape[0]

  UpperCAmelCase													=	layout or torch.strided
  UpperCAmelCase													=	device or torch.device('cpu'						)

  if generator is not None:
    UpperCAmelCase													=	generator.device.type if not isinstance(SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE						) else generator[0].device.type
    if gen_device_type != device.type and gen_device_type == "cpu":
      UpperCAmelCase													=	'cpu'
      if device != "mps":
        logger.info(
            f'''The passed generator was created on \'cpu\' even though a tensor on {device} was expected.'''
            f''' Tensors will be created on \'cpu\' and then moved to {device}. Note that one can probably'''
            f''' slighly speed up this function by passing a generator that was created on the {device} device.'''						)
    elif gen_device_type != device.type and gen_device_type == "cuda":
      raise ValueError(f'''Cannot generate a {device} tensor from a generator of type {gen_device_type}.'''						)

  if isinstance(SCREAMING_SNAKE_CASE       ,   SCREAMING_SNAKE_CASE						):
    UpperCAmelCase													=	(1,) + shape[1:]
    UpperCAmelCase													=	[
        torch.randn(SCREAMING_SNAKE_CASE       ,   generator=generator[i]       ,   device=SCREAMING_SNAKE_CASE       ,   dtype=SCREAMING_SNAKE_CASE       ,   layout=SCREAMING_SNAKE_CASE						)
        for i in range(SCREAMING_SNAKE_CASE						)
    ]
    UpperCAmelCase													=	torch.cat(SCREAMING_SNAKE_CASE       ,   dim=0						).to(SCREAMING_SNAKE_CASE						)
  else:
    UpperCAmelCase													=	torch.randn(SCREAMING_SNAKE_CASE       ,   generator=SCREAMING_SNAKE_CASE       ,   device=SCREAMING_SNAKE_CASE       ,   dtype=SCREAMING_SNAKE_CASE       ,   layout=SCREAMING_SNAKE_CASE						).to(SCREAMING_SNAKE_CASE						)

  return latents



def       lowerCamelCase__		(						SCREAMING_SNAKE_CASE      :   Tuple						):
  if is_torch_version('<'       ,   '2.0.0'						) or not hasattr(SCREAMING_SNAKE_CASE       ,   '_dynamo'						):
    return False
  return isinstance(SCREAMING_SNAKE_CASE       ,   torch._dynamo.eval_frame.OptimizedModule						)


