



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



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

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




 def     lowerCamelCase (  self			):



    '''simple docstring'''
    __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.type_sequence_label_size			)

    __lowerCamelCase  							=      self.get_config()

    return config, pixel_values, labels




 def     lowerCamelCase (  self			):



    '''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     lowerCamelCase (  self		,			__UpperCAmelCase		,			__UpperCAmelCase		,			__UpperCAmelCase			):



    '''simple docstring'''
    __lowerCamelCase  							=      SwinaSRModel(config=__UpperCAmelCase			)
    model.to(__UpperCAmelCase			)
    model.eval()
    __lowerCamelCase  							=      model(__UpperCAmelCase			)

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




 def     lowerCamelCase (  self		,			__UpperCAmelCase		,			__UpperCAmelCase		,			__UpperCAmelCase			):



    '''simple docstring'''
    __lowerCamelCase  							=      SwinaSRForImageSuperResolution(__UpperCAmelCase			)
    model.to(__UpperCAmelCase			)
    model.eval()
    __lowerCamelCase  							=      model(__UpperCAmelCase			)

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




 def     lowerCamelCase (  self			):



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



@require_torch
class 	__lowerCAmelCase       (		lowerCAmelCase__					,		lowerCAmelCase__					,		unittest.TestCase				):
 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     lowerCamelCase (  self			):



    '''simple docstring'''
    __lowerCamelCase  							=      SwinaSRModelTester(self			)
    __lowerCamelCase  							=      ConfigTester(self		,			config_class=__UpperCAmelCase		,			embed_dim=37			)




 def     lowerCamelCase (  self			):



    '''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     lowerCamelCase (  self			):



    '''simple docstring'''
    __lowerCamelCase  							=      self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_model(*__UpperCAmelCase			)




 def     lowerCamelCase (  self			):



    '''simple docstring'''
    __lowerCamelCase  							=      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     lowerCamelCase (  self			):



    '''simple docstring'''
    pass




 @unittest.skip(reason='''Swin2SR does not use inputs_embeds'''			)
 def     lowerCamelCase (  self			):



    '''simple docstring'''
    pass




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



    '''simple docstring'''
    pass




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



    '''simple docstring'''
    pass




 def     lowerCamelCase (  self			):



    '''simple docstring'''
    __lowerCamelCase							,__lowerCamelCase  							=      self.model_tester.prepare_config_and_inputs_for_common()

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




 def     lowerCamelCase (  self			):



    '''simple docstring'''
    __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			)




 @slow
 def     lowerCamelCase (  self			):



    '''simple docstring'''
    for model_name in SWIN2SR_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       __lowerCamelCase  							=      SwinaSRModel.from_pretrained(__UpperCAmelCase			)
       self.assertIsNotNone(__UpperCAmelCase			)




 def     lowerCamelCase (  self			):



    '''simple docstring'''
    __lowerCamelCase							,__lowerCamelCase  							=      self.model_tester.prepare_config_and_inputs_for_common()

    __lowerCamelCase  							=      _config_zero_init(__UpperCAmelCase			)
    for model_class in self.all_model_classes:
       __lowerCamelCase  							=      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     lowerCamelCase (  self			):



    '''simple docstring'''
    __lowerCamelCase							,__lowerCamelCase  							=      self.model_tester.prepare_config_and_inputs_for_common()
    __lowerCamelCase  							=      True

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

       # check that output_attentions also work using config
       del inputs_dict["output_attentions"]
       __lowerCamelCase  							=      True
       __lowerCamelCase  							=      config.window_size**2
       __lowerCamelCase  							=      model_class(__UpperCAmelCase			)
       model.to(__UpperCAmelCase			)
       model.eval()
       with torch.no_grad():
          __lowerCamelCase  							=      model(**self._prepare_for_class(__UpperCAmelCase		,			__UpperCAmelCase			)			)
       __lowerCamelCase  							=      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]		,			)
       __lowerCamelCase  							=      len(__UpperCAmelCase			)

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

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

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




 def     lowerCamelCase (  self			):



    '''simple docstring'''
    __lowerCamelCase  							=      SwinaSRImageProcessor()
    __lowerCamelCase  							=      SwinaSRForImageSuperResolution.from_pretrained('''caidas/swin2SR-classical-sr-x2-64'''			).to(__UpperCAmelCase			)

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

    # forward pass
    with torch.no_grad():
       __lowerCamelCase  							=      model(**__UpperCAmelCase			)

    # verify the logits
    __lowerCamelCase  							=      torch.Size([1, 3, 976, 1296]			)
    self.assertEqual(outputs.reconstruction.shape		,			__UpperCAmelCase			)
    __lowerCamelCase  							=      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(__UpperCAmelCase			)
    self.assertTrue(torch.allclose(outputs.reconstruction[0, 0, :3, :3]		,			__UpperCAmelCase		,			atol=1E-4			)			)











import math



class 	__lowerCAmelCase       :




 def __init__(  self		,			__UpperCAmelCase		,			__UpperCAmelCase		,			__UpperCAmelCase		,			__UpperCAmelCase			):



    '''simple docstring'''
    __lowerCamelCase  							=      x
    __lowerCamelCase  							=      y
    __lowerCamelCase  							=      step_size
    __lowerCamelCase  							=      function_to_optimize




 def     lowerCamelCase (  self			):



    '''simple docstring'''
    return self.function(self.x		,			self.y			)




 def     lowerCamelCase (  self			):



    '''simple docstring'''
    __lowerCamelCase  							=      self.step_size
    return [
        SearchProblem(__UpperCAmelCase		,			__UpperCAmelCase		,			__UpperCAmelCase		,			self.function			)
        for x, y in (
            (self.x - step_size, self.y - step_size),
            (self.x - step_size, self.y),
            (self.x - step_size, self.y + step_size),
            (self.x, self.y - step_size),
            (self.x, self.y + step_size),
            (self.x + step_size, self.y - step_size),
            (self.x + step_size, self.y),
            (self.x + step_size, self.y + step_size),
        )
    ]




 def __hash__(  self			):



    '''simple docstring'''
    return hash(str(self			)			)




 def __eq__(  self		,			__UpperCAmelCase			):



    '''simple docstring'''
    if isinstance(__UpperCAmelCase		,			__UpperCAmelCase			):
       return hash(str(self			)			) == hash(str(__UpperCAmelCase			)			)
    return False




 def __str__(  self			):



    '''simple docstring'''
    return F"""x: {self.x} y: {self.y}"""





def 						a__     (						_UpperCamelCase     :	Tuple	,_UpperCamelCase     :	bool = True	,_UpperCamelCase     :	float = math.inf	,_UpperCamelCase     :	float = -math.inf	,_UpperCamelCase     :	float = math.inf	,_UpperCamelCase     :	float = -math.inf	,_UpperCamelCase     :	bool = False	,_UpperCamelCase     :	int = 1_00_00	,):
   __lowerCamelCase  							=      search_prob
   __lowerCamelCase  							=      []  # list to store the current score at each iteration
   __lowerCamelCase  							=      0
   __lowerCamelCase  							=      False
   __lowerCamelCase  							=      set()
   while not solution_found and iterations < max_iter:
      visited.add(_UpperCamelCase							)
      iterations += 1
      __lowerCamelCase  							=      current_state.score()
      scores.append(_UpperCamelCase							)
      __lowerCamelCase  							=      current_state.get_neighbors()
      __lowerCamelCase  							=      -math.inf
      __lowerCamelCase  							=      math.inf
      __lowerCamelCase  							=      None  # to hold the next best neighbor
      for neighbor in neighbors:
         if neighbor in visited:
            continue  # do not want to visit the same state again
         if (
             neighbor.x > max_x
             or neighbor.x < min_x
             or neighbor.y > max_y
             or neighbor.y < min_y
         ):
            continue  # neighbor outside our bounds
         __lowerCamelCase  							=      neighbor.score() - current_score
         if find_max:  # finding max
            # going to direction with greatest ascent
            if change > max_change and change > 0:
               __lowerCamelCase  							=      change
               __lowerCamelCase  							=      neighbor
         else:  # finding min
            # to direction with greatest descent
            if change < min_change and change < 0:
               __lowerCamelCase  							=      change
               __lowerCamelCase  							=      neighbor
      if next_state is not None:
         # we found at least one neighbor which improved the current state
         __lowerCamelCase  							=      next_state
      else:
         # since we have no neighbor that improves the solution we stop the search
         __lowerCamelCase  							=      True

   if visualization:
      from matplotlib import pyplot as plt

      plt.plot(range(_UpperCamelCase							)	,_UpperCamelCase							)
      plt.xlabel('''Iterations'''							)
      plt.ylabel('''Function values'''							)
      plt.show()

   return current_state


if __name__ == "__main__":
  import doctest

  doctest.testmod()





  def 						a__     (						_UpperCamelCase     :	Union[str, Any]	,_UpperCamelCase     :	Union[str, Any]							):
     return (x**2) + (y**2)

  # starting the problem with initial coordinates (3, 4)
  a_								=   SearchProblem(x=3, y=4, step_size=1, function_to_optimize=test_fa)
  a_								=   hill_climbing(prob, find_max=False)
  print(
      """The minimum score for f(x, y) = x^2 + y^2 found via hill climbing: """
      f"{local_min.score()}"
  )

  # starting the problem with initial coordinates (12, 47)
  a_								=   SearchProblem(x=12, y=47, step_size=1, function_to_optimize=test_fa)
  a_								=   hill_climbing(
      prob, find_max=False, max_x=100, min_x=5, max_y=50, min_y=-5, visualization=True
  )
  print(
      """The minimum score for f(x, y) = x^2 + y^2 with the domain 100 > x > 5 """
      f"and 50 > y > - 5 found via hill climbing: {local_min.score()}"
  )





  def 						a__     (						_UpperCamelCase     :	int	,_UpperCamelCase     :	Dict							):
     return (3 * x**2) - (6 * y)

  a_								=   SearchProblem(x=3, y=4, step_size=1, function_to_optimize=test_fa)
  a_								=   hill_climbing(prob, find_max=True)





  print(
      """The maximum score for f(x, y) = x^2 + y^2 found via hill climbing: """
      f"{local_min.score()}"
  )






