




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  :

      """simple docstring"""

      def __init__(       self   :	Any	,    snake_case   :	Union[str, Any]	,    snake_case   :	Union[str, Any]=13	,    snake_case   :	str=32	,    snake_case   :	List[Any]=1	,    snake_case   :	Optional[int]=3	,    snake_case   :	Any=16	,    snake_case   :	List[str]=[1, 2, 1]	,    snake_case   :	List[str]=[2, 2, 4]	,    snake_case   :	Any=2	,    snake_case   :	Tuple=2.0	,    snake_case   :	Union[str, Any]=True	,    snake_case   :	int=0.0	,    snake_case   :	Optional[int]=0.0	,    snake_case   :	Any=0.1	,    snake_case   :	List[Any]="gelu"	,    snake_case   :	Union[str, Any]=False	,    snake_case   :	Tuple=True	,    snake_case   :	Dict=0.02	,    snake_case   :	Tuple=1E-5	,    snake_case   :	List[str]=True	,    snake_case   :	str=None	,    snake_case   :	Union[str, Any]=False	,    snake_case   :	Tuple=2	,    ):
          __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(snake_case			)
          __UpperCamelCase											=  embed_dim
          __UpperCamelCase											=  (image_size // patch_size) ** 2







      def  snake_case		(       self   :	int			):
          __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  snake_case		(       self   :	Tuple			):
          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  snake_case		(       self   :	List[str]	,    snake_case   :	Optional[int]	,    snake_case   :	List[Any]	,    snake_case   :	Any			):
          __UpperCamelCase											=  SwinaSRModel(config=snake_case			)
          model.to(snake_case			)
          model.eval()
          __UpperCamelCase											=  model(snake_case			)

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







      def  snake_case		(       self   :	List[Any]	,    snake_case   :	Any	,    snake_case   :	Optional[int]	,    snake_case   :	Optional[int]			):
          __UpperCamelCase											=  SwinaSRForImageSuperResolution(snake_case			)
          model.to(snake_case			)
          model.eval()
          __UpperCamelCase											=  model(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  snake_case		(       self   :	Union[str, Any]			):
          __UpperCamelCase											=  self.prepare_config_and_inputs()
          __UpperCamelCase						, __UpperCamelCase						, __UpperCamelCase											=  config_and_inputs
          __UpperCamelCase											=  {'''pixel_values''': pixel_values}
          return config, inputs_dict







@require_torch
class 	_lowerCamelCase  (     _SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE							,     unittest.TestCase				):

      """simple docstring"""

      lowerCAmelCase__     :	int      =   (SwinaSRModel, SwinaSRForImageSuperResolution) if is_torch_available() else ()
      lowerCAmelCase__     :	Optional[Any]      =   {"feature-extraction": SwinaSRModel} if is_torch_available() else {}

      lowerCAmelCase__     :	Optional[int]      =   False
      lowerCAmelCase__     :	Tuple      =   False
      lowerCAmelCase__     :	List[str]      =   False
      lowerCAmelCase__     :	Optional[int]      =   False
      lowerCAmelCase__     :	Any      =   False







      def  snake_case		(       self   :	int			):
          __UpperCamelCase											=  SwinaSRModelTester(self			)
          __UpperCamelCase											=  ConfigTester(self	,    config_class=snake_case	,    embed_dim=37			)







      def  snake_case		(       self   :	Optional[Any]			):
          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  snake_case		(       self   :	Tuple			):
          __UpperCamelCase											=  self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*snake_case			)







      def  snake_case		(       self   :	int			):
          __UpperCamelCase											=  self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_image_super_resolution(*snake_case			)







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







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







      @unittest.skip(reason='''Swin2SR does not support training yet'''			)
      def  snake_case		(       self   :	Optional[Any]			):
          pass







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







      def  snake_case		(       self   :	int			):
          __UpperCamelCase						, __UpperCamelCase											=  self.model_tester.prepare_config_and_inputs_for_common()

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







      def  snake_case		(       self   :	Optional[int]			):
          __UpperCamelCase						, __UpperCamelCase											=  self.model_tester.prepare_config_and_inputs_for_common()

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







      @slow
      def  snake_case		(       self   :	List[str]			):
          for model_name in SWIN2SR_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
              __UpperCamelCase											=  SwinaSRModel.from_pretrained(snake_case			)
              self.assertIsNotNone(snake_case			)







      def  snake_case		(       self   :	List[Any]			):
          __UpperCamelCase						, __UpperCamelCase											=  self.model_tester.prepare_config_and_inputs_for_common()

          __UpperCamelCase											=  _config_zero_init(snake_case			)
          for model_class in self.all_model_classes:
              __UpperCamelCase											=  model_class(config=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  snake_case		(       self   :	List[Any]			):
          __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(snake_case			)
              model.to(snake_case			)
              model.eval()
              with torch.no_grad():
                  __UpperCamelCase											=  model(**self._prepare_for_class(snake_case	,    snake_case			)			)
              __UpperCamelCase											=  outputs.attentions
              __UpperCamelCase											=  len(self.model_tester.depths			)
              self.assertEqual(len(snake_case			)	,    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(snake_case			)
              model.to(snake_case			)
              model.eval()
              with torch.no_grad():
                  __UpperCamelCase											=  model(**self._prepare_for_class(snake_case	,    snake_case			)			)
              __UpperCamelCase											=  outputs.attentions
              self.assertEqual(len(snake_case			)	,    snake_case			)

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

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

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

              __UpperCamelCase											=  outputs.attentions

              self.assertEqual(len(snake_case			)	,    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 	_lowerCamelCase  (     unittest.TestCase				):

      """simple docstring"""

      def  snake_case		(       self   :	str			):
          __UpperCamelCase											=  SwinaSRImageProcessor()
          __UpperCamelCase											=  SwinaSRForImageSuperResolution.from_pretrained('''caidas/swin2SR-classical-sr-x2-64'''			).to(snake_case			)

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

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

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












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 	_lowerCamelCase  (     _SCREAMING_SNAKE_CASE							,     unittest.TestCase				):

      """simple docstring"""

      lowerCAmelCase__     :	Union[str, Any]      =   RePaintPipeline
      lowerCAmelCase__     :	Optional[int]      =   IMAGE_INPAINTING_PARAMS - {"width", "height", "guidance_scale"}
      lowerCAmelCase__     :	Tuple      =   PipelineTesterMixin.required_optional_params - {
          "latents",
          "num_images_per_prompt",
          "callback",
          "callback_steps",
      }
      lowerCAmelCase__     :	List[str]      =   IMAGE_INPAINTING_BATCH_PARAMS
      lowerCAmelCase__     :	Dict      =   False







      def  snake_case		(       self   :	Optional[Any]			):
          torch.manual_seed(0			)
          torch.manual_seed(0			)
          __UpperCamelCase											=  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''')	,    )
          __UpperCamelCase											=  RePaintScheduler()
          __UpperCamelCase											=  {'''unet''': unet, '''scheduler''': scheduler}
          return components







      def  snake_case		(       self   :	int	,    snake_case   :	str	,    snake_case   :	List[Any]=0			):
          if str(snake_case			).startswith('''mps'''			):
              __UpperCamelCase											=  torch.manual_seed(snake_case			)
          else:
              __UpperCamelCase											=  torch.Generator(device=snake_case			).manual_seed(snake_case			)
          __UpperCamelCase											=  np.random.RandomState(snake_case			).standard_normal((1, 3, 32, 32)			)
          __UpperCamelCase											=  torch.from_numpy(snake_case			).to(device=snake_case	,    dtype=torch.floataa			)
          __UpperCamelCase											=  (image > 0).to(device=snake_case	,    dtype=torch.floataa			)
          __UpperCamelCase											=  {
              '''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  snake_case		(       self   :	str			):
          __UpperCamelCase											=  '''cpu'''  # ensure determinism for the device-dependent torch.Generator
          __UpperCamelCase											=  self.get_dummy_components()
          __UpperCamelCase											=  RePaintPipeline(**snake_case			)
          __UpperCamelCase											=  sd_pipe.to(snake_case			)
          sd_pipe.set_progress_bar_config(disable=snake_case			)

          __UpperCamelCase											=  self.get_dummy_inputs(snake_case			)
          __UpperCamelCase											=  sd_pipe(**snake_case			).images
          __UpperCamelCase											=  image[0, -3:, -3:, -1]

          assert image.shape == (1, 32, 32, 3)
          __UpperCamelCase											=  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  snake_case		(       self   :	str			):
          return super().test_save_load_local()







      @unittest.skip('''non-deterministic pipeline'''			)
      def  snake_case		(       self   :	Any			):
          return super().test_inference_batch_single_identical()







      @skip_mps
      def  snake_case		(       self   :	List[Any]			):
          return super().test_dict_tuple_outputs_equivalent()







      @skip_mps
      def  snake_case		(       self   :	List[str]			):
          return super().test_save_load_optional_components()







      @skip_mps
      def  snake_case		(       self   :	List[str]			):
          return super().test_attention_slicing_forward_pass()







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

      """simple docstring"""

      def  snake_case		(       self   :	Dict			):
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()







      def  snake_case		(       self   :	Any			):
          __UpperCamelCase											=  load_image(
              '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/'''
              '''repaint/celeba_hq_256.png'''			)
          __UpperCamelCase											=  load_image(
              '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/repaint/mask_256.png'''			)
          __UpperCamelCase											=  load_numpy(
              '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/'''
              '''repaint/celeba_hq_256_result.npy'''			)

          __UpperCamelCase											=  '''google/ddpm-ema-celebahq-256'''
          __UpperCamelCase											=  UNetaDModel.from_pretrained(snake_case			)
          __UpperCamelCase											=  RePaintScheduler.from_pretrained(snake_case			)

          __UpperCamelCase											=  RePaintPipeline(unet=snake_case	,    scheduler=snake_case			).to(snake_case			)
          repaint.set_progress_bar_config(disable=snake_case			)
          repaint.enable_attention_slicing()

          __UpperCamelCase											=  torch.manual_seed(0			)
          __UpperCamelCase											=  repaint(
              snake_case	,    snake_case	,    num_inference_steps=250	,    eta=0.0	,    jump_length=10	,    jump_n_sample=10	,    generator=snake_case	,    output_type='''np'''	,    )
          __UpperCamelCase											=  output.images[0]

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






