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

							"""simple docstring"""






							def __init__(	self  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_=13  ,						SCREAMING_SNAKE_CASE_=32  ,						SCREAMING_SNAKE_CASE_=1  ,						SCREAMING_SNAKE_CASE_=3  ,						SCREAMING_SNAKE_CASE_=16  ,						SCREAMING_SNAKE_CASE_=[1, 2, 1]  ,						SCREAMING_SNAKE_CASE_=[2, 2, 4]  ,						SCREAMING_SNAKE_CASE_=2  ,						SCREAMING_SNAKE_CASE_=2.0  ,						SCREAMING_SNAKE_CASE_=True  ,						SCREAMING_SNAKE_CASE_=0.0  ,						SCREAMING_SNAKE_CASE_=0.0  ,						SCREAMING_SNAKE_CASE_=0.1  ,						SCREAMING_SNAKE_CASE_="gelu"  ,						SCREAMING_SNAKE_CASE_=False  ,						SCREAMING_SNAKE_CASE_=True  ,						SCREAMING_SNAKE_CASE_=0.0_2  ,						SCREAMING_SNAKE_CASE_=1E-5  ,						SCREAMING_SNAKE_CASE_=True  ,						SCREAMING_SNAKE_CASE_=None  ,						SCREAMING_SNAKE_CASE_=False  ,						SCREAMING_SNAKE_CASE_=2  ,						)-> List[str]:





													'''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(SCREAMING_SNAKE_CASE_      )
													__UpperCamelCase  	=      embed_dim
													__UpperCamelCase  	=      (image_size // patch_size) ** 2



							def      A__		(	self      )-> Any:





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





													'''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      A__		(	self  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_      )-> Optional[int]:





													'''simple docstring'''






													__UpperCamelCase  	=      SwinaSRModel(config=SCREAMING_SNAKE_CASE_      )
													model.to(SCREAMING_SNAKE_CASE_      )
													model.eval()
													__UpperCamelCase  	=      model(SCREAMING_SNAKE_CASE_      )

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



							def      A__		(	self  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_      )-> Dict:





													'''simple docstring'''






													__UpperCamelCase  	=      SwinaSRForImageSuperResolution(SCREAMING_SNAKE_CASE_      )
													model.to(SCREAMING_SNAKE_CASE_      )
													model.eval()
													__UpperCamelCase  	=      model(SCREAMING_SNAKE_CASE_      )

													__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      A__		(	self      )-> Optional[Any]:





													'''simple docstring'''






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






@require_torch
class 							SCREAMING_SNAKE_CASE__     (				SCREAMING_SNAKE_CASE_ ,		SCREAMING_SNAKE_CASE_ ,		unittest.TestCase     ):

							"""simple docstring"""






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

							_snake_case  		=	False
							_snake_case  		=	False
							_snake_case  		=	False
							_snake_case  		=	False
							_snake_case  		=	False



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





													'''simple docstring'''






													__UpperCamelCase  	=      SwinaSRModelTester(self      )
													__UpperCamelCase  	=      ConfigTester(self  ,						config_class=SCREAMING_SNAKE_CASE_  ,						embed_dim=37      )



							def      A__		(	self      )-> List[Any]:





													'''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      A__		(	self      )-> Optional[int]:





													'''simple docstring'''






													__UpperCamelCase  	=      self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_      )



							def      A__		(	self      )-> Optional[Any]:





													'''simple docstring'''






													__UpperCamelCase  	=      self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_image_super_resolution(*SCREAMING_SNAKE_CASE_      )



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





													'''simple docstring'''






													pass



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





													'''simple docstring'''






													pass



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





													'''simple docstring'''






													pass



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





													'''simple docstring'''






													pass



							def      A__		(	self      )-> Dict:





													'''simple docstring'''






													__UpperCamelCase	,				__UpperCamelCase  	=      self.model_tester.prepare_config_and_inputs_for_common()

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



							def      A__		(	self      )-> str:





													'''simple docstring'''






													__UpperCamelCase	,				__UpperCamelCase  	=      self.model_tester.prepare_config_and_inputs_for_common()

													for model_class in self.all_model_classes:
																			__UpperCamelCase  	=      model_class(SCREAMING_SNAKE_CASE_      )
																			__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]  ,						SCREAMING_SNAKE_CASE_      )



							@slow
							def      A__		(	self      )-> Dict:





													'''simple docstring'''






													for model_name in SWIN2SR_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																			__UpperCamelCase  	=      SwinaSRModel.from_pretrained(SCREAMING_SNAKE_CASE_      )
																			self.assertIsNotNone(SCREAMING_SNAKE_CASE_      )



							def      A__		(	self      )-> List[Any]:





													'''simple docstring'''






													__UpperCamelCase	,				__UpperCamelCase  	=      self.model_tester.prepare_config_and_inputs_for_common()

													__UpperCamelCase  	=      _config_zero_init(SCREAMING_SNAKE_CASE_      )
													for model_class in self.all_model_classes:
																			__UpperCamelCase  	=      model_class(config=SCREAMING_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      A__		(	self      )-> str:





													'''simple docstring'''






													__UpperCamelCase	,				__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(SCREAMING_SNAKE_CASE_      )
																			model.to(SCREAMING_SNAKE_CASE_      )
																			model.eval()
																			with torch.no_grad():
																									__UpperCamelCase  	=      model(**self._prepare_for_class(SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_      )      )
																			__UpperCamelCase  	=      outputs.attentions
																			__UpperCamelCase  	=      len(self.model_tester.depths      )
																			self.assertEqual(len(SCREAMING_SNAKE_CASE_      )  ,						SCREAMING_SNAKE_CASE_      )

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

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

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

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

																			__UpperCamelCase  	=      outputs.attentions

																			self.assertEqual(len(SCREAMING_SNAKE_CASE_      )  ,						SCREAMING_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 							SCREAMING_SNAKE_CASE__     (				unittest.TestCase     ):

							"""simple docstring"""






							def      A__		(	self      )-> int:





													'''simple docstring'''






													__UpperCamelCase  	=      SwinaSRImageProcessor()
													__UpperCamelCase  	=      SwinaSRForImageSuperResolution.from_pretrained('''caidas/swin2SR-classical-sr-x2-64'''      ).to(SCREAMING_SNAKE_CASE_      )

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

													# forward pass
													with torch.no_grad():
																			__UpperCamelCase  	=      model(**SCREAMING_SNAKE_CASE_      )

													# verify the logits
													__UpperCamelCase  	=      torch.Size([1, 3, 976, 1296]      )
													self.assertEqual(outputs.reconstruction.shape  ,						SCREAMING_SNAKE_CASE_      )
													__UpperCamelCase  	=      torch.tensor(
													    [[0.5_4_5_8, 0.5_5_4_6, 0.5_6_3_8], [0.5_5_2_6, 0.5_5_6_5, 0.5_6_5_1], [0.5_3_9_6, 0.5_4_2_6, 0.5_6_2_1]]      ).to(SCREAMING_SNAKE_CASE_      )
													self.assertTrue(torch.allclose(outputs.reconstruction[0, 0, :3, :3]  ,						SCREAMING_SNAKE_CASE_  ,						atol=1E-4      )      )






"""simple docstring"""

import random
import unittest

import torch

from diffusers import IFSuperResolutionPipeline
from diffusers.utils import floats_tensor
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import skip_mps, torch_device

from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin
from . import IFPipelineTesterMixin

@skip_mps
class __lowerCAmelCase						( __SCREAMING_SNAKE_CASE     ,    __SCREAMING_SNAKE_CASE     ,    unittest.TestCase     ):
     lowercase		  =		IFSuperResolutionPipeline
     lowercase		  =		TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"width", "height"}
     lowercase		  =		TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
     lowercase		  =		PipelineTesterMixin.required_optional_params - {"latents"}





     def 		UpperCAmelCase   (      self ):

      '''simple docstring'''




      return self._get_superresolution_dummy_components()





     def 		UpperCAmelCase   (      self      ,	__UpperCAmelCase      ,	__UpperCAmelCase=0 ):

      '''simple docstring'''




      if str(__UpperCAmelCase ).startswith('mps' ):
       __UpperCamelCase							    =			torch.manual_seed(__UpperCAmelCase )
      else:
       __UpperCamelCase							    =			torch.Generator(device=__UpperCAmelCase ).manual_seed(__UpperCAmelCase )

      __UpperCamelCase							    =			floats_tensor((1, 3, 32, 32)      ,	rng=random.Random(__UpperCAmelCase ) ).to(__UpperCAmelCase )

      __UpperCamelCase							    =			{
          'prompt': 'A painting of a squirrel eating a burger',
          'image': image,
          'generator': generator,
          'num_inference_steps': 2,
          'output_type': 'numpy',
      }

      return inputs





     @unittest.skipIf(
         torch_device != 'cuda' or not is_xformers_available()      ,	reason='XFormers attention is only available with CUDA and `xformers` installed'      ,	)
     def 		UpperCAmelCase   (      self ):

      '''simple docstring'''




      self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3 )





     def 		UpperCAmelCase   (      self ):

      '''simple docstring'''




      self._test_save_load_optional_components()





     @unittest.skipIf(torch_device != 'cuda'      ,	reason='float16 requires CUDA' )
     def 		UpperCAmelCase   (      self ):

      '''simple docstring'''




      super().test_save_load_floataa(expected_max_diff=1E-1 )





     def 		UpperCAmelCase   (      self ):

      '''simple docstring'''




      self._test_attention_slicing_forward_pass(expected_max_diff=1E-2 )





     def 		UpperCAmelCase   (      self ):

      '''simple docstring'''




      self._test_save_load_local()







     def 		UpperCAmelCase   (      self ):

      '''simple docstring'''




      self._test_inference_batch_single_identical(
          expected_max_diff=1E-2      ,	)



