




import gc
import unittest

import numpy as np
import torch

from diffusers import RePaintPipeline, RePaintScheduler, UNetaDModel
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    load_image,
    load_numpy,
    nightly,
    require_torch_gpu,
    skip_mps,
    torch_device,
)

from ..pipeline_params import IMAGE_INPAINTING_BATCH_PARAMS, IMAGE_INPAINTING_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()


class 				__SCREAMING_SNAKE_CASE							(		lowerCamelCase   ,     unittest.TestCase  ):
					_lowerCamelCase		  = RePaintPipeline
					_lowerCamelCase		  = IMAGE_INPAINTING_PARAMS - {"""width""", """height""", """guidance_scale"""}
					_lowerCamelCase		  = PipelineTesterMixin.required_optional_params - {
					    """latents""",
					    """num_images_per_prompt""",
					    """callback""",
					    """callback_steps""",
					}
					_lowerCamelCase		  = IMAGE_INPAINTING_BATCH_PARAMS
					_lowerCamelCase		  = False

					def 	lowercase_							(				self					):







						'''simple docstring'''




						torch.manual_seed(0					)
						torch.manual_seed(0					)
						lowerCamelCase_							:		str							   =							UNetaDModel(
						    block_out_channels=(32, 64) ,       layers_per_block=2 ,       sample_size=32 ,       in_channels=3 ,       out_channels=3 ,       down_block_types=('DownBlock2D', 'AttnDownBlock2D') ,       up_block_types=('AttnUpBlock2D', 'UpBlock2D') ,       )
						lowerCamelCase_							:		Optional[Any]							   =							RePaintScheduler()
						lowerCamelCase_							:		List[str]							   =							{'unet': unet, 'scheduler': scheduler}
						return components

					def 	lowercase_							(				self ,       __lowercase ,       __lowercase=0					):







						'''simple docstring'''




						if str(__lowercase					).startswith('mps'					):
							lowerCamelCase_							:		str							   =							torch.manual_seed(__lowercase					)
						else:
							lowerCamelCase_							:		Union[str, Any]							   =							torch.Generator(device=__lowercase					).manual_seed(__lowercase					)
						lowerCamelCase_							:		Dict							   =							np.random.RandomState(__lowercase					).standard_normal((1, 3, 32, 32)					)
						lowerCamelCase_							:		Union[str, Any]							   =							torch.from_numpy(__lowercase					).to(device=__lowercase ,       dtype=torch.floataa					)
						lowerCamelCase_							:		int							   =							(image > 0).to(device=__lowercase ,       dtype=torch.floataa					)
						lowerCamelCase_							:		List[str]							   =							{
						    'image': image,
						    'mask_image': mask,
						    'generator': generator,
						    'num_inference_steps': 5,
						    'eta': 0.0,
						    'jump_length': 2,
						    'jump_n_sample': 2,
						    'output_type': 'numpy',
						}
						return inputs

					def 	lowercase_							(				self					):







						'''simple docstring'''




						lowerCamelCase_							:		Optional[Any]							   =							'cpu'  # ensure determinism for the device-dependent torch.Generator
						lowerCamelCase_							:		int							   =							self.get_dummy_components()
						lowerCamelCase_							:		Dict							   =							RePaintPipeline(**__lowercase					)
						lowerCamelCase_							:		Dict							   =							sd_pipe.to(__lowercase					)
						sd_pipe.set_progress_bar_config(disable=__lowercase					)

						lowerCamelCase_							:		Tuple							   =							self.get_dummy_inputs(__lowercase					)
						lowerCamelCase_							:		Dict							   =							sd_pipe(**__lowercase					).images
						lowerCamelCase_							:		List[Any]							   =							image[0, -3:, -3:, -1]

						assert image.shape == (1, 32, 32, 3)
						lowerCamelCase_							:		Any							   =							np.array([1.00_00, 0.54_26, 0.54_97, 0.22_00, 1.00_00, 1.00_00, 0.56_23, 1.00_00, 0.62_74]					)

						assert np.abs(image_slice.flatten() - expected_slice					).max() < 1E-3

					@skip_mps
					def 	lowercase_							(				self					):







						'''simple docstring'''




						return super().test_save_load_local()

					@unittest.skip('non-deterministic pipeline'					)
					def 	lowercase_							(				self					):







						'''simple docstring'''




						return super().test_inference_batch_single_identical()

					@skip_mps
					def 	lowercase_							(				self					):







						'''simple docstring'''




						return super().test_dict_tuple_outputs_equivalent()

					@skip_mps
					def 	lowercase_							(				self					):







						'''simple docstring'''




						return super().test_save_load_optional_components()





					@skip_mps
					def 	lowercase_							(				self					):







						'''simple docstring'''




						return super().test_attention_slicing_forward_pass()


@nightly
@require_torch_gpu
class 				__SCREAMING_SNAKE_CASE							(		unittest.TestCase  ):

					def 	lowercase_							(				self					):







						'''simple docstring'''




						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()





					def 	lowercase_							(				self					):







						'''simple docstring'''




						lowerCamelCase_							:		List[str]							   =							load_image(
						    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/'
						    'repaint/celeba_hq_256.png'					)
						lowerCamelCase_							:		str							   =							load_image(
						    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/repaint/mask_256.png'					)
						lowerCamelCase_							:		Any							   =							load_numpy(
						    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/'
						    'repaint/celeba_hq_256_result.npy'					)

						lowerCamelCase_							:		Tuple							   =							'google/ddpm-ema-celebahq-256'
						lowerCamelCase_							:		Dict							   =							UNetaDModel.from_pretrained(__lowercase					)
						lowerCamelCase_							:		Union[str, Any]							   =							RePaintScheduler.from_pretrained(__lowercase					)

						lowerCamelCase_							:		Tuple							   =							RePaintPipeline(unet=__lowercase ,       scheduler=__lowercase					).to(__lowercase					)
						repaint.set_progress_bar_config(disable=__lowercase					)
						repaint.enable_attention_slicing()

						lowerCamelCase_							:		Optional[int]							   =							torch.manual_seed(0					)
						lowerCamelCase_							:		Optional[int]							   =							repaint(
						    __lowercase ,       __lowercase ,       num_inference_steps=250 ,       eta=0.0 ,       jump_length=10 ,       jump_n_sample=10 ,       generator=__lowercase ,       output_type='np' ,       )
						lowerCamelCase_							:		Dict							   =							output.images[0]

						assert image.shape == (256, 256, 3)
						assert np.abs(expected_image - image					).mean() < 1E-2






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							:

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







						'''simple docstring'''




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

						# here we set some attributes to make tests pass
						lowerCamelCase_							:		int							   =							len(__lowercase					)
						lowerCamelCase_							:		Optional[int]							   =							embed_dim
						lowerCamelCase_							:		Tuple							   =							(image_size // patch_size) ** 2

					def 	lowercase_							(				self					):







						'''simple docstring'''




						lowerCamelCase_							:		int							   =							floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]					)

						lowerCamelCase_							:		List[Any]							   =							None
						if self.use_labels:
							lowerCamelCase_							:		Dict							   =							ids_tensor([self.batch_size] ,       self.type_sequence_label_size					)

						lowerCamelCase_							:		int							   =							self.get_config()

						return config, pixel_values, labels

					def 	lowercase_							(				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 	lowercase_							(				self ,       __lowercase ,       __lowercase ,       __lowercase					):







						'''simple docstring'''




						lowerCamelCase_							:		Any							   =							SwinaSRModel(config=__lowercase					)
						model.to(__lowercase					)
						model.eval()
						lowerCamelCase_							:		int							   =							model(__lowercase					)

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

					def 	lowercase_							(				self ,       __lowercase ,       __lowercase ,       __lowercase					):







						'''simple docstring'''




						lowerCamelCase_							:		Optional[Any]							   =							SwinaSRForImageSuperResolution(__lowercase					)
						model.to(__lowercase					)
						model.eval()
						lowerCamelCase_							:		Any							   =							model(__lowercase					)

						lowerCamelCase_							:		int							   =							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					):







						'''simple docstring'''




						lowerCamelCase_							:		Dict							   =							self.prepare_config_and_inputs()
						lowerCamelCase_   ,			lowerCamelCase_   ,			lowerCamelCase_							:		Union[str, Any]							   =							config_and_inputs
						lowerCamelCase_							:		Any							   =							{'pixel_values': pixel_values}
						return config, inputs_dict


@require_torch
class 				__SCREAMING_SNAKE_CASE							(		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 	lowercase_							(				self					):







						'''simple docstring'''




						lowerCamelCase_							:		Tuple							   =							SwinaSRModelTester(self					)
						lowerCamelCase_							:		Optional[Any]							   =							ConfigTester(self ,       config_class=__lowercase ,       embed_dim=37					)

					def 	lowercase_							(				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 	lowercase_							(				self					):







						'''simple docstring'''




						lowerCamelCase_							:		Optional[Any]							   =							self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*__lowercase					)

					def 	lowercase_							(				self					):







						'''simple docstring'''




						lowerCamelCase_							:		Union[str, Any]							   =							self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_image_super_resolution(*__lowercase					)

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







						'''simple docstring'''




						pass

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







						'''simple docstring'''




						pass

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







						'''simple docstring'''




						pass

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







						'''simple docstring'''




						pass

					def 	lowercase_							(				self					):







						'''simple docstring'''




						lowerCamelCase_   ,			lowerCamelCase_							:		Dict							   =							self.model_tester.prepare_config_and_inputs_for_common()

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

					def 	lowercase_							(				self					):







						'''simple docstring'''




						lowerCamelCase_   ,			lowerCamelCase_							:		Optional[int]							   =							self.model_tester.prepare_config_and_inputs_for_common()

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

							lowerCamelCase_							:		Tuple							   =							['pixel_values']
							self.assertListEqual(arg_names[:1] ,       __lowercase					)

					@slow
					def 	lowercase_							(				self					):







						'''simple docstring'''




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

					def 	lowercase_							(				self					):







						'''simple docstring'''




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

						lowerCamelCase_							:		Optional[Any]							   =							_config_zero_init(__lowercase					)
						for model_class in self.all_model_classes:
							lowerCamelCase_							:		List[str]							   =							model_class(config=__lowercase					)
							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					):







						'''simple docstring'''




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

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

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

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

							# Check attention is always last and order is fine
							lowerCamelCase_							:		Tuple							   =							True
							lowerCamelCase_							:		List[Any]							   =							True
							lowerCamelCase_							:		Tuple							   =							model_class(__lowercase					)
							model.to(__lowercase					)
							model.eval()
							with torch.no_grad():
								lowerCamelCase_							:		List[str]							   =							model(**self._prepare_for_class(__lowercase ,       __lowercase					)					)

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

							lowerCamelCase_							:		str							   =							outputs.attentions

							self.assertEqual(len(__lowercase					) ,       __lowercase					)

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

					def 	lowercase_							(				self					):







						'''simple docstring'''




						lowerCamelCase_							:		Union[str, Any]							   =							SwinaSRImageProcessor()
						lowerCamelCase_							:		Optional[int]							   =							SwinaSRForImageSuperResolution.from_pretrained('caidas/swin2SR-classical-sr-x2-64'					).to(__lowercase					)

						lowerCamelCase_							:		int							   =							Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'					)
						lowerCamelCase_							:		Optional[int]							   =							processor(images=__lowercase ,       return_tensors='pt'					).to(__lowercase					)

						# forward pass
						with torch.no_grad():
							lowerCamelCase_							:		Optional[int]							   =							model(**__lowercase					)

						# verify the logits
						lowerCamelCase_							:		Dict							   =							torch.Size([1, 3, 976, 1_296]					)
						self.assertEqual(outputs.reconstruction.shape ,       __lowercase					)
						lowerCamelCase_							:		Union[str, Any]							   =							torch.tensor(
						    [[0.54_58, 0.55_46, 0.56_38], [0.55_26, 0.55_65, 0.56_51], [0.53_96, 0.54_26, 0.56_21]]					).to(__lowercase					)
						self.assertTrue(torch.allclose(outputs.reconstruction[0, 0, :3, :3] ,       __lowercase ,       atol=1E-4					)					)
