
"""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    __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.0_2       ,	UpperCAmelCase=1e-5       ,	UpperCAmelCase=True       ,	UpperCAmelCase=None       ,	UpperCAmelCase=False       ,	UpperCAmelCase=2       ,	):
							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(A_      )
							lowerCamelCase_ 		=     embed_dim
							lowerCamelCase_ 		=     (image_size // patch_size) ** 2




				def 					UpperCAmelCase__	(							self      ):
							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 					UpperCAmelCase__	(							self      ):
							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 					UpperCAmelCase__	(							self       ,	UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase      ):
							lowerCamelCase_ 		=     SwinaSRModel(config=A_      )
							model.to(A_      )
							model.eval()
							lowerCamelCase_ 		=     model(A_      )

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




				def 					UpperCAmelCase__	(							self       ,	UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase      ):
							lowerCamelCase_ 		=     SwinaSRForImageSuperResolution(A_      )
							model.to(A_      )
							model.eval()
							lowerCamelCase_ 		=     model(A_      )

							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 					UpperCAmelCase__	(							self      ):
							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 ):
				a__:			Dict									=						(SwinaSRModel, SwinaSRForImageSuperResolution) if is_torch_available() else ()
				a__:			Union[str, Any]									=						{'feature-extraction': SwinaSRModel} if is_torch_available() else {}

				a__:			List[str]									=						False
				a__:			Any									=						False
				a__:			str									=						False
				a__:			Tuple									=						False
				a__:			str									=						False




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_ 		=     SwinaSRModelTester(self      )
							lowerCamelCase_ 		=     ConfigTester(self       ,	config_class=A_       ,	embed_dim=37      )




				def 					UpperCAmelCase__	(							self      ):
							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 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_ 		=     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*A_      )




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_ 		=     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_image_super_resolution(*A_      )




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




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




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




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




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_		,   lowerCamelCase_ 		=     self.model_tester.prepare_config_and_inputs_for_common()

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




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_		,   lowerCamelCase_ 		=     self.model_tester.prepare_config_and_inputs_for_common()

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




				@slow
				def 					UpperCAmelCase__	(							self      ):
							for model_name in SWIN2SR_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										lowerCamelCase_ 		=     SwinaSRModel.from_pretrained(A_      )
										self.assertIsNotNone(A_      )




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_		,   lowerCamelCase_ 		=     self.model_tester.prepare_config_and_inputs_for_common()

							lowerCamelCase_ 		=     _config_zero_init(A_      )
							for model_class in self.all_model_classes:
										lowerCamelCase_ 		=     model_class(config=A_      )
										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 					UpperCAmelCase__	(							self      ):
							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(A_      )
										model.to(A_      )
										model.eval()
										with torch.no_grad():
													lowerCamelCase_ 		=     model(**self._prepare_for_class(A_       ,	A_      )      )
										lowerCamelCase_ 		=     outputs.attentions
										lowerCamelCase_ 		=     len(self.model_tester.depths      )
										self.assertEqual(len(A_      )       ,	A_      )

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

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

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

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

										lowerCamelCase_ 		=     outputs.attentions

										self.assertEqual(len(A_      )       ,	A_      )

										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 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_ 		=     SwinaSRImageProcessor()
							lowerCamelCase_ 		=     SwinaSRForImageSuperResolution.from_pretrained('''caidas/swin2SR-classical-sr-x2-64'''      ).to(A_      )

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

							# forward pass
							with torch.no_grad():
										lowerCamelCase_ 		=     model(**A_      )

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








"""simple docstring"""


import tempfile
import unittest

import torch

from diffusers import UNetaDConditionModel
from diffusers.training_utils import EMAModel
from diffusers.utils.testing_utils import enable_full_determinism, skip_mps, torch_device


enable_full_determinism()




class    __lowerCamelCase     (   unittest.TestCase ):
				a__:			Union[str, Any]									=						'hf-internal-testing/tiny-stable-diffusion-pipe'
				a__:			Tuple									=						1
				a__:			int									=						77
				a__:			List[str]									=						32
				a__:			List[Any]									=						4
				a__:			str									=						64
				a__:			Optional[Any]									=						torch.manual_seed(0 )




				def 					UpperCAmelCase__	(							self       ,	UpperCAmelCase=0.9_9_9_9      ):
							lowerCamelCase_ 		=     UNetaDConditionModel.from_pretrained(self.model_id       ,	subfolder='''unet'''      )
							lowerCamelCase_ 		=     unet.to(UpperCAmelCase      )
							lowerCamelCase_ 		=     EMAModel(unet.parameters()       ,	decay=UpperCAmelCase       ,	model_cls=UpperCAmelCase       ,	model_config=unet.config      )
							return unet, ema_unet




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_ 		=     torch.randn(
							    self.batch_size       ,	self.num_in_channels       ,	self.latent_height       ,	self.latent_width       ,	generator=self.generator      ).to(UpperCAmelCase      )
							lowerCamelCase_ 		=     torch.randint(0       ,	1000       ,	size=(self.batch_size,)       ,	generator=self.generator      ).to(UpperCAmelCase      )
							lowerCamelCase_ 		=     torch.randn(
							    self.batch_size       ,	self.prompt_length       ,	self.text_encoder_hidden_dim       ,	generator=self.generator      ).to(UpperCAmelCase      )
							return noisy_latents, timesteps, encoder_hidden_states




				def 					UpperCAmelCase__	(							self       ,	UpperCAmelCase      ):
							lowerCamelCase_ 		=     {}
							for k, param in unet.state_dict().items():
										lowerCamelCase_ 		=     torch.randn_like(UpperCAmelCase      ) + (param * torch.randn_like(UpperCAmelCase      ))
										updated_state_dict.update({k: updated_param}      )
							unet.load_state_dict(UpperCAmelCase      )
							return unet




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_		,   lowerCamelCase_ 		=     self.get_models()
							# Take the first (hypothetical) EMA step.
							ema_unet.step(unet.parameters()      )
							assert ema_unet.optimization_step == 1

							# Take two more.
							for _ in range(2      ):
										ema_unet.step(unet.parameters()      )
							assert ema_unet.optimization_step == 3




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_		,   lowerCamelCase_ 		=     self.get_models()
							# Since the `unet` is not being updated (i.e., backprop'd)
							# there won't be any difference between the `params` of `unet`
							# and `ema_unet` even if we call `ema_unet.step(unet.parameters())`.
							ema_unet.step(unet.parameters()      )
							lowerCamelCase_ 		=     list(unet.parameters()      )
							for s_param, param in zip(ema_unet.shadow_params       ,	UpperCAmelCase      ):
										assert torch.allclose(UpperCAmelCase       ,	UpperCAmelCase      )

							# The above holds true even if we call `ema.step()` multiple times since
							# `unet` params are still not being updated.
							for _ in range(4      ):
										ema_unet.step(unet.parameters()      )
							for s_param, param in zip(ema_unet.shadow_params       ,	UpperCAmelCase      ):
										assert torch.allclose(UpperCAmelCase       ,	UpperCAmelCase      )




				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_		,   lowerCamelCase_ 		=     self.get_models()
							# Here we simulate the parameter updates for `unet`. Since there might
							# be some parameters which are initialized to zero we take extra care to
							# initialize their values to something non-zero before the multiplication.
							lowerCamelCase_ 		=     self.simulate_backprop(UpperCAmelCase      )

							# Take the EMA step.
							ema_unet.step(unet_pseudo_updated_step_one.parameters()      )

							# Now the EMA'd parameters won't be equal to the original model parameters.
							lowerCamelCase_ 		=     list(unet_pseudo_updated_step_one.parameters()      )
							for s_param, param in zip(ema_unet.shadow_params       ,	UpperCAmelCase      ):
										assert ~torch.allclose(UpperCAmelCase       ,	UpperCAmelCase      )

							# Ensure this is the case when we take multiple EMA steps.
							for _ in range(4      ):
										ema_unet.step(unet.parameters()      )
							for s_param, param in zip(ema_unet.shadow_params       ,	UpperCAmelCase      ):
										assert ~torch.allclose(UpperCAmelCase       ,	UpperCAmelCase      )




				def 					UpperCAmelCase__	(							self      ):
							# If we call EMA step after a backpropagation consecutively for two times,
							# the shadow params from those two steps should be different.
							lowerCamelCase_		,   lowerCamelCase_ 		=     self.get_models()

							# First backprop + EMA
							lowerCamelCase_ 		=     self.simulate_backprop(UpperCAmelCase      )
							ema_unet.step(unet_step_one.parameters()      )
							lowerCamelCase_ 		=     ema_unet.shadow_params

							# Second backprop + EMA
							lowerCamelCase_ 		=     self.simulate_backprop(UpperCAmelCase      )
							ema_unet.step(unet_step_two.parameters()      )
							lowerCamelCase_ 		=     ema_unet.shadow_params

							for step_one, step_two in zip(UpperCAmelCase       ,	UpperCAmelCase      ):
										assert ~torch.allclose(UpperCAmelCase       ,	UpperCAmelCase      )




				def 					UpperCAmelCase__	(							self      ):
							# If there's no decay even if there are backprops, EMA steps
							# won't take any effect i.e., the shadow params would remain the
							# same.
							lowerCamelCase_		,   lowerCamelCase_ 		=     self.get_models(decay=0.0      )
							lowerCamelCase_ 		=     self.simulate_backprop(UpperCAmelCase      )
							ema_unet.step(unet_step_one.parameters()      )
							lowerCamelCase_ 		=     ema_unet.shadow_params

							lowerCamelCase_ 		=     self.simulate_backprop(UpperCAmelCase      )
							ema_unet.step(unet_step_two.parameters()      )
							lowerCamelCase_ 		=     ema_unet.shadow_params

							for step_one, step_two in zip(UpperCAmelCase       ,	UpperCAmelCase      ):
										assert torch.allclose(UpperCAmelCase       ,	UpperCAmelCase      )





				@skip_mps
				def 					UpperCAmelCase__	(							self      ):
							lowerCamelCase_		,   lowerCamelCase_ 		=     self.get_models()
							lowerCamelCase_		,   lowerCamelCase_		,   lowerCamelCase_ 		=     self.get_dummy_inputs()

							with tempfile.TemporaryDirectory() as tmpdir:
										ema_unet.save_pretrained(UpperCAmelCase      )
										lowerCamelCase_ 		=     UNetaDConditionModel.from_pretrained(UpperCAmelCase       ,	model_cls=UpperCAmelCase      )
										lowerCamelCase_ 		=     loaded_unet.to(unet.device      )

							# Since no EMA step has been performed the outputs should match.
							lowerCamelCase_ 		=     unet(UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase      ).sample
							lowerCamelCase_ 		=     loaded_unet(UpperCAmelCase       ,	UpperCAmelCase       ,	UpperCAmelCase      ).sample

							assert torch.allclose(UpperCAmelCase       ,	UpperCAmelCase       ,	atol=1e-4      )



