

"""simple docstring"""


from dataclasses import dataclass
from typing import List, Optional, Union

import numpy as np
import PIL

from ...utils import BaseOutput, is_invisible_watermark_available, is_torch_available, is_transformers_available





@dataclass
class 					snake_case			(       __snake_case						):
				SCREAMING_SNAKE_CASE_     :   Union[List[PIL.Image.Image], np.ndarray]


if is_transformers_available() and is_torch_available() and is_invisible_watermark_available():
			from .pipeline_stable_diffusion_xl import StableDiffusionXLPipeline
			from .pipeline_stable_diffusion_xl_imgaimg import StableDiffusionXLImgaImgPipeline





"""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 					snake_case			:

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



						'''simple docstring'''


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

						# here we set some attributes to make tests pass
						__lowerCAmelCase: int					   =    len(UpperCamelCase__)
						__lowerCAmelCase: Tuple					   =    embed_dim
						__lowerCAmelCase: Dict					   =    (image_size // patch_size) ** 2

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



						'''simple docstring'''


						__lowerCAmelCase: Any					   =    floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

						__lowerCAmelCase: Union[str, Any]					   =    None
						if self.use_labels:
								__lowerCAmelCase: Union[str, Any]					   =    ids_tensor([self.batch_size]  ,						self.type_sequence_label_size)

						__lowerCAmelCase: Dict					   =    self.get_config()

						return config, pixel_values, labels

				def 							lowercase_  (			self							:       Any)->     Optional[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 							lowercase_  (			self							:       Any  ,						UpperCamelCase__							:       Any  ,						UpperCamelCase__							:       List[str]  ,						UpperCamelCase__							:       Any)->     Any:



						'''simple docstring'''


						__lowerCAmelCase: Optional[int]					   =    SwinaSRModel(config=UpperCamelCase__)
						model.to(UpperCamelCase__)
						model.eval()
						__lowerCAmelCase: Union[str, Any]					   =    model(UpperCamelCase__)

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

				def 							lowercase_  (			self							:       List[Any]  ,						UpperCamelCase__							:       str  ,						UpperCamelCase__							:       List[Any]  ,						UpperCamelCase__							:       Optional[Any])->     List[Any]:



						'''simple docstring'''


						__lowerCAmelCase: List[Any]					   =    SwinaSRForImageSuperResolution(UpperCamelCase__)
						model.to(UpperCamelCase__)
						model.eval()
						__lowerCAmelCase: int					   =    model(UpperCamelCase__)

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




				def 							lowercase_  (			self							:       str)->     str:



						'''simple docstring'''


						__lowerCAmelCase: str					   =    self.prepare_config_and_inputs()
						__lowerCAmelCase   ,				__lowerCAmelCase   ,				__lowerCAmelCase: int					   =    config_and_inputs
						__lowerCAmelCase: Tuple					   =    {"pixel_values": pixel_values}
						return config, inputs_dict





@require_torch
class 					snake_case			(       __snake_case,							__snake_case,							unittest.TestCase						):
				SCREAMING_SNAKE_CASE_     :   Optional[int]        =						(SwinaSRModel, SwinaSRForImageSuperResolution) if is_torch_available() else ()
				SCREAMING_SNAKE_CASE_     :   List[str]        =						{"""feature-extraction""": SwinaSRModel} if is_torch_available() else {}

				SCREAMING_SNAKE_CASE_     :   int        =						False
				SCREAMING_SNAKE_CASE_     :   str        =						False
				SCREAMING_SNAKE_CASE_     :   str        =						False
				SCREAMING_SNAKE_CASE_     :   Any        =						False
				SCREAMING_SNAKE_CASE_     :   int        =						False

				def 							lowercase_  (			self							:       str)->     str:



						'''simple docstring'''


						__lowerCAmelCase: Optional[int]					   =    SwinaSRModelTester(self)
						__lowerCAmelCase: List[Any]					   =    ConfigTester(self  ,						config_class=UpperCamelCase__  ,						embed_dim=3_7)

				def 							lowercase_  (			self							:       Optional[int])->     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 							lowercase_  (			self							:       str)->     Tuple:



						'''simple docstring'''


						__lowerCAmelCase: Union[str, Any]					   =    self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*UpperCamelCase__)

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



						'''simple docstring'''


						__lowerCAmelCase: Dict					   =    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 							lowercase_  (			self							:       List[Any])->     Union[str, Any]:



						'''simple docstring'''


						pass

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



						'''simple docstring'''


						pass

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



						'''simple docstring'''


						pass

				@unittest.skip(reason="Swin2SR does not support training yet")
				def 							lowercase_  (			self							:       Optional[int])->     str:



						'''simple docstring'''


						pass

				def 							lowercase_  (			self							:       Dict)->     int:



						'''simple docstring'''


						__lowerCAmelCase   ,				__lowerCAmelCase: Any					   =    self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
								__lowerCAmelCase: Optional[Any]					   =    model_class(UpperCamelCase__)
								self.assertIsInstance(model.get_input_embeddings()  ,						(nn.Module))
								__lowerCAmelCase: Tuple					   =    model.get_output_embeddings()
								self.assertTrue(x is None or isinstance(UpperCamelCase__  ,						nn.Linear))

				def 							lowercase_  (			self							:       str)->     List[str]:



						'''simple docstring'''


						__lowerCAmelCase   ,				__lowerCAmelCase: str					   =    self.model_tester.prepare_config_and_inputs_for_common()

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

								__lowerCAmelCase: Optional[Any]					   =    ["pixel_values"]
								self.assertListEqual(arg_names[:1]  ,						UpperCamelCase__)

				@slow
				def 							lowercase_  (			self							:       List[Any])->     Union[str, Any]:



						'''simple docstring'''


						for model_name in SWIN2SR_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
								__lowerCAmelCase: Tuple					   =    SwinaSRModel.from_pretrained(UpperCamelCase__)
								self.assertIsNotNone(UpperCamelCase__)

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



						'''simple docstring'''


						__lowerCAmelCase   ,				__lowerCAmelCase: Union[str, Any]					   =    self.model_tester.prepare_config_and_inputs_for_common()

						__lowerCAmelCase: Optional[int]					   =    _config_zero_init(UpperCamelCase__)
						for model_class in self.all_model_classes:
								__lowerCAmelCase: Union[str, Any]					   =    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 							lowercase_  (			self							:       List[Any])->     List[Any]:



						'''simple docstring'''


						__lowerCAmelCase   ,				__lowerCAmelCase: Dict					   =    self.model_tester.prepare_config_and_inputs_for_common()
						__lowerCAmelCase: Optional[int]					   =    True

						for model_class in self.all_model_classes:
								__lowerCAmelCase: Dict					   =    True
								__lowerCAmelCase: Dict					   =    False
								__lowerCAmelCase: Any					   =    True
								__lowerCAmelCase: Optional[Any]					   =    model_class(UpperCamelCase__)
								model.to(UpperCamelCase__)
								model.eval()
								with torch.no_grad():
										__lowerCAmelCase: Union[str, Any]					   =    model(**self._prepare_for_class(UpperCamelCase__  ,						UpperCamelCase__))
								__lowerCAmelCase: Optional[Any]					   =    outputs.attentions
								__lowerCAmelCase: Any					   =    len(self.model_tester.depths)
								self.assertEqual(len(UpperCamelCase__)  ,						UpperCamelCase__)

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

								# Check attention is always last and order is fine
								__lowerCAmelCase: Dict					   =    True
								__lowerCAmelCase: List[Any]					   =    True
								__lowerCAmelCase: int					   =    model_class(UpperCamelCase__)
								model.to(UpperCamelCase__)
								model.eval()
								with torch.no_grad():
										__lowerCAmelCase: Optional[Any]					   =    model(**self._prepare_for_class(UpperCamelCase__  ,						UpperCamelCase__))

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

								__lowerCAmelCase: List[str]					   =    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 					snake_case			(       unittest.TestCase						):

				def 							lowercase_  (			self							:       Tuple)->     Dict:



						'''simple docstring'''


						__lowerCAmelCase: Dict					   =    SwinaSRImageProcessor()
						__lowerCAmelCase: Tuple					   =    SwinaSRForImageSuperResolution.from_pretrained("caidas/swin2SR-classical-sr-x2-64").to(UpperCamelCase__)

						__lowerCAmelCase: Optional[int]					   =    Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png")
						__lowerCAmelCase: List[str]					   =    processor(images=UpperCamelCase__  ,						return_tensors="pt").to(UpperCamelCase__)

						# forward pass
						with torch.no_grad():
								__lowerCAmelCase: Union[str, Any]					   =    model(**UpperCamelCase__)

						# verify the logits
						__lowerCAmelCase: Optional[Any]					   =    torch.Size([1, 3, 9_7_6, 1_2_9_6])
						self.assertEqual(outputs.reconstruction.shape  ,						UpperCamelCase__)
						__lowerCAmelCase: List[Any]					   =    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))


