






"""simple docstring"""







import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPImageProcessor, CLIPVisionConfig, CLIPVisionModelWithProjection

from diffusers import (
    AutoencoderKL,
    DPMSolverMultistepScheduler,
    PNDMScheduler,
    StableDiffusionImageVariationPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, nightly, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import IMAGE_VARIATION_BATCH_PARAMS, IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()







class 					lowerCamelCase__					(       _lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase    ,     unittest.TestCase ):
    SCREAMING_SNAKE_CASE		=			StableDiffusionImageVariationPipeline
    SCREAMING_SNAKE_CASE		=			IMAGE_VARIATION_PARAMS
    SCREAMING_SNAKE_CASE		=			IMAGE_VARIATION_BATCH_PARAMS
    SCREAMING_SNAKE_CASE		=			frozenset([] )
    # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
    SCREAMING_SNAKE_CASE		=			frozenset([] )






    def 				_UpperCamelCase							( self					):
          torch.manual_seed(0					)
          UpperCAmelCase							=       UNetaDConditionModel(
              block_out_channels=(32, 64)	,layers_per_block=2	,sample_size=32	,in_channels=4	,out_channels=4	,down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")	,up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")	,cross_attention_dim=32	,)
          UpperCAmelCase							=       PNDMScheduler(skip_prk_steps=A__					)
          torch.manual_seed(0					)
          UpperCAmelCase							=       AutoencoderKL(
              block_out_channels=[32, 64]	,in_channels=3	,out_channels=3	,down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]	,up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]	,latent_channels=4	,)
          torch.manual_seed(0					)
          UpperCAmelCase							=       CLIPVisionConfig(
              hidden_size=32	,projection_dim=32	,intermediate_size=37	,layer_norm_eps=1e-0_5	,num_attention_heads=4	,num_hidden_layers=5	,image_size=32	,patch_size=4	,)
          UpperCAmelCase							=       CLIPVisionModelWithProjection(A__					)
          UpperCAmelCase							=       CLIPImageProcessor(crop_size=32	,size=32					)

          UpperCAmelCase							=       {
              """unet""": unet,
              """scheduler""": scheduler,
              """vae""": vae,
              """image_encoder""": image_encoder,
              """feature_extractor""": feature_extractor,
              """safety_checker""": None,
          }
          return components






    def 				_UpperCamelCase							( self	,A	,A=0					):
          UpperCAmelCase							=       floats_tensor((1, 3, 32, 32)	,rng=random.Random(A__					)					)
          UpperCAmelCase							=       image.cpu().permute(0	,2	,3	,1					)[0]
          UpperCAmelCase							=       Image.fromarray(np.uinta(A__					)					).convert("""RGB"""					).resize((32, 32)					)
          if str(A__					).startswith("""mps"""					):
                UpperCAmelCase							=       torch.manual_seed(A__					)
          else:
                UpperCAmelCase							=       torch.Generator(device=A__					).manual_seed(A__					)
          UpperCAmelCase							=       {
              """image""": image,
              """generator""": generator,
              """num_inference_steps""": 2,
              """guidance_scale""": 6.0,
              """output_type""": """numpy""",
          }
          return inputs






    def 				_UpperCamelCase							( self					):
          UpperCAmelCase							=       """cpu"""  # ensure determinism for the device-dependent torch.Generator
          UpperCAmelCase							=       self.get_dummy_components()
          UpperCAmelCase							=       StableDiffusionImageVariationPipeline(**A__					)
          UpperCAmelCase							=       sd_pipe.to(A__					)
          sd_pipe.set_progress_bar_config(disable=A__					)

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

          assert image.shape == (1, 64, 64, 3)
          UpperCAmelCase							=       np.array([0.5239, 0.5723, 0.4796, 0.5049, 0.5550, 0.4685, 0.5329, 0.4891, 0.4921]					)

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






    def 				_UpperCamelCase							( self					):
          UpperCAmelCase							=       """cpu"""  # ensure determinism for the device-dependent torch.Generator
          UpperCAmelCase							=       self.get_dummy_components()
          UpperCAmelCase							=       StableDiffusionImageVariationPipeline(**A__					)
          UpperCAmelCase							=       sd_pipe.to(A__					)
          sd_pipe.set_progress_bar_config(disable=A__					)

          UpperCAmelCase							=       self.get_dummy_inputs(A__					)
          UpperCAmelCase							=       2 * [inputs["""image"""]]
          UpperCAmelCase							=       sd_pipe(**A__					)

          UpperCAmelCase							=       output.images

          UpperCAmelCase							=       image[-1, -3:, -3:, -1]

          assert image.shape == (2, 64, 64, 3)
          UpperCAmelCase							=       np.array([0.6892, 0.5637, 0.5836, 0.5771, 0.6254, 0.6409, 0.5580, 0.5569, 0.5289]					)

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






    def 				_UpperCamelCase							( self					):
          super().test_inference_batch_single_identical(expected_max_diff=3e-3					)







@slow
@require_torch_gpu
class 					lowerCamelCase__					(       unittest.TestCase ):






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






    def 				_UpperCamelCase							( self	,A	,A="cpu"	,A=torch.floataa	,A=0					):
          UpperCAmelCase							=       torch.Generator(device=A__					).manual_seed(A__					)
          UpperCAmelCase							=       load_image(
              """https://huggingface.co/datasets/diffusers/test-arrays/resolve/main"""
              """/stable_diffusion_imgvar/input_image_vermeer.png"""					)
          UpperCAmelCase							=       np.random.RandomState(A__					).standard_normal((1, 4, 64, 64)					)
          UpperCAmelCase							=       torch.from_numpy(A__					).to(device=A__	,dtype=A__					)
          UpperCAmelCase							=       {
              """image""": init_image,
              """latents""": latents,
              """generator""": generator,
              """num_inference_steps""": 3,
              """guidance_scale""": 7.5,
              """output_type""": """numpy""",
          }
          return inputs






    def 				_UpperCamelCase							( self					):
          UpperCAmelCase							=       StableDiffusionImageVariationPipeline.from_pretrained(
              """lambdalabs/sd-image-variations-diffusers"""	,safety_checker=A__					)
          UpperCAmelCase							=       sd_pipe.to(A__					)
          sd_pipe.set_progress_bar_config(disable=A__					)

          UpperCAmelCase							=       self.get_inputs(A__					)
          UpperCAmelCase							=       sd_pipe(**A__					).images
          UpperCAmelCase							=       image[0, -3:, -3:, -1].flatten()

          assert image.shape == (1, 512, 512, 3)
          UpperCAmelCase							=       np.array([0.84491, 0.90789, 0.75708, 0.78734, 0.83485, 0.70099, 0.66938, 0.68727, 0.61379]					)
          assert np.abs(image_slice - expected_slice					).max() < 6e-3






    def 				_UpperCamelCase							( self					):
          UpperCAmelCase							=       0

          def callback_fn(A	,A	,A					) -> None:
                UpperCAmelCase							=       True
                nonlocal number_of_steps
                number_of_steps += 1
                if step == 1:
                      UpperCAmelCase							=       latents.detach().cpu().numpy()
                      assert latents.shape == (1, 4, 64, 64)
                      UpperCAmelCase							=       latents[0, -3:, -3:, -1]
                      UpperCAmelCase							=       np.array(
                          [-0.1621, 0.2837, -0.7979, -0.1221, -1.3057, 0.7681, -2.1191, 0.0464, 1.6309]					)

                      assert np.abs(latents_slice.flatten() - expected_slice					).max() < 5e-2
                elif step == 2:
                      UpperCAmelCase							=       latents.detach().cpu().numpy()
                      assert latents.shape == (1, 4, 64, 64)
                      UpperCAmelCase							=       latents[0, -3:, -3:, -1]
                      UpperCAmelCase							=       np.array([0.6299, 1.7500, 1.1992, -2.1582, -1.8994, 0.7334, -0.7090, 1.0137, 1.5273]					)

                      assert np.abs(latents_slice.flatten() - expected_slice					).max() < 5e-2

          UpperCAmelCase							=       False

          UpperCAmelCase							=       StableDiffusionImageVariationPipeline.from_pretrained(
              """fusing/sd-image-variations-diffusers"""	,safety_checker=A__	,torch_dtype=torch.floataa	,)
          pipe.to(A__					)
          pipe.set_progress_bar_config(disable=A__					)
          pipe.enable_attention_slicing()

          UpperCAmelCase							=       self.get_inputs(A__	,dtype=torch.floataa					)
          pipe(**A__	,callback=A__	,callback_steps=1					)
          assert callback_fn.has_been_called
          assert number_of_steps == inputs["num_inference_steps"]






    def 				_UpperCamelCase							( self					):
          torch.cuda.empty_cache()
          torch.cuda.reset_max_memory_allocated()
          torch.cuda.reset_peak_memory_stats()

          UpperCAmelCase							=       """fusing/sd-image-variations-diffusers"""
          UpperCAmelCase							=       StableDiffusionImageVariationPipeline.from_pretrained(
              A__	,safety_checker=A__	,torch_dtype=torch.floataa					)
          UpperCAmelCase							=       pipe.to(A__					)
          pipe.set_progress_bar_config(disable=A__					)
          pipe.enable_attention_slicing(1					)
          pipe.enable_sequential_cpu_offload()

          UpperCAmelCase							=       self.get_inputs(A__	,dtype=torch.floataa					)
          UpperCAmelCase							=       pipe(**A__					)

          UpperCAmelCase							=       torch.cuda.max_memory_allocated()
          # make sure that less than 2.6 GB is allocated
          assert mem_bytes < 2.6 * 10**9







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






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






    def 				_UpperCamelCase							( self	,A	,A="cpu"	,A=torch.floataa	,A=0					):
          UpperCAmelCase							=       torch.Generator(device=A__					).manual_seed(A__					)
          UpperCAmelCase							=       load_image(
              """https://huggingface.co/datasets/diffusers/test-arrays/resolve/main"""
              """/stable_diffusion_imgvar/input_image_vermeer.png"""					)
          UpperCAmelCase							=       np.random.RandomState(A__					).standard_normal((1, 4, 64, 64)					)
          UpperCAmelCase							=       torch.from_numpy(A__					).to(device=A__	,dtype=A__					)
          UpperCAmelCase							=       {
              """image""": init_image,
              """latents""": latents,
              """generator""": generator,
              """num_inference_steps""": 50,
              """guidance_scale""": 7.5,
              """output_type""": """numpy""",
          }
          return inputs






    def 				_UpperCamelCase							( self					):
          UpperCAmelCase							=       StableDiffusionImageVariationPipeline.from_pretrained("""fusing/sd-image-variations-diffusers"""					)
          sd_pipe.to(A__					)
          sd_pipe.set_progress_bar_config(disable=A__					)

          UpperCAmelCase							=       self.get_inputs(A__					)
          UpperCAmelCase							=       sd_pipe(**A__					).images[0]

          UpperCAmelCase							=       load_numpy(
              """https://huggingface.co/datasets/diffusers/test-arrays/resolve/main"""
              """/stable_diffusion_imgvar/lambdalabs_variations_pndm.npy"""					)
          UpperCAmelCase							=       np.abs(expected_image - image					).max()
          assert max_diff < 1e-3






    def 				_UpperCamelCase							( self					):
          UpperCAmelCase							=       StableDiffusionImageVariationPipeline.from_pretrained("""fusing/sd-image-variations-diffusers"""					)
          UpperCAmelCase							=       DPMSolverMultistepScheduler.from_config(sd_pipe.scheduler.config					)
          sd_pipe.to(A__					)
          sd_pipe.set_progress_bar_config(disable=A__					)

          UpperCAmelCase							=       self.get_inputs(A__					)
          UpperCAmelCase							=       25
          UpperCAmelCase							=       sd_pipe(**A__					).images[0]

          UpperCAmelCase							=       load_numpy(
              """https://huggingface.co/datasets/diffusers/test-arrays/resolve/main"""
              """/stable_diffusion_imgvar/lambdalabs_variations_dpm_multi.npy"""					)
          UpperCAmelCase							=       np.abs(expected_image - image					).max()
          assert max_diff < 1e-3





'''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       _a :





       """simple docstring"""




       def __init__(							self				,							A__				,							A__=13				,							A__=32				,							A__=1				,							A__=3				,							A__=16				,							A__=[1, 2, 1]				,							A__=[2, 2, 4]				,							A__=2				,							A__=2.0				,							A__=True				,							A__=0.0				,							A__=0.0				,							A__=0.1				,							A__="gelu"				,							A__=False				,							A__=True				,							A__=0.02				,							A__=1E-5				,							A__=True				,							A__=None				,							A__=False				,							A__=2				,							)			->	Optional[Any]:
          _SCREAMING_SNAKE_CASE          = parent
          _SCREAMING_SNAKE_CASE          = batch_size
          _SCREAMING_SNAKE_CASE          = image_size
          _SCREAMING_SNAKE_CASE          = patch_size
          _SCREAMING_SNAKE_CASE          = num_channels
          _SCREAMING_SNAKE_CASE          = embed_dim
          _SCREAMING_SNAKE_CASE          = depths
          _SCREAMING_SNAKE_CASE          = num_heads
          _SCREAMING_SNAKE_CASE          = window_size
          _SCREAMING_SNAKE_CASE          = mlp_ratio
          _SCREAMING_SNAKE_CASE          = qkv_bias
          _SCREAMING_SNAKE_CASE          = hidden_dropout_prob
          _SCREAMING_SNAKE_CASE          = attention_probs_dropout_prob
          _SCREAMING_SNAKE_CASE          = drop_path_rate
          _SCREAMING_SNAKE_CASE          = hidden_act
          _SCREAMING_SNAKE_CASE          = use_absolute_embeddings
          _SCREAMING_SNAKE_CASE          = patch_norm
          _SCREAMING_SNAKE_CASE          = layer_norm_eps
          _SCREAMING_SNAKE_CASE          = initializer_range
          _SCREAMING_SNAKE_CASE          = is_training
          _SCREAMING_SNAKE_CASE          = scope
          _SCREAMING_SNAKE_CASE          = use_labels
          _SCREAMING_SNAKE_CASE          = upscale

          # here we set some attributes to make tests pass
          _SCREAMING_SNAKE_CASE          = len(A__					)
          _SCREAMING_SNAKE_CASE          = embed_dim
          _SCREAMING_SNAKE_CASE          = (image_size // patch_size) ** 2






       def 				UpperCamelCase (							self					)			->	Any:
          _SCREAMING_SNAKE_CASE          = floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]					)

          _SCREAMING_SNAKE_CASE          = None
          if self.use_labels:
             _SCREAMING_SNAKE_CASE          = ids_tensor([self.batch_size]				,							self.type_sequence_label_size					)

          _SCREAMING_SNAKE_CASE          = self.get_config()

          return config, pixel_values, labels






       def 				UpperCamelCase (							self					)			->	Optional[Any]:
          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				,							A__				,							A__				,							A__					)			->	List[str]:
          _SCREAMING_SNAKE_CASE          = SwinaSRModel(config=A__					)
          model.to(A__					)
          model.eval()
          _SCREAMING_SNAKE_CASE          = 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				,							A__				,							A__				,							A__					)			->	Dict:
          _SCREAMING_SNAKE_CASE          = SwinaSRForImageSuperResolution(A__					)
          model.to(A__					)
          model.eval()
          _SCREAMING_SNAKE_CASE          = model(A__					)

          _SCREAMING_SNAKE_CASE          = 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					)			->	Optional[Any]:
          _SCREAMING_SNAKE_CASE          = self.prepare_config_and_inputs()
          _SCREAMING_SNAKE_CASE       ,							_SCREAMING_SNAKE_CASE       ,							_SCREAMING_SNAKE_CASE          = config_and_inputs
          _SCREAMING_SNAKE_CASE          = {"""pixel_values""": pixel_values}
          return config, inputs_dict






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





       """simple docstring"""




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

       SCREAMING_SNAKE_CASE					=     False
       SCREAMING_SNAKE_CASE					=     False
       SCREAMING_SNAKE_CASE					=     False
       SCREAMING_SNAKE_CASE					=     False
       SCREAMING_SNAKE_CASE					=     False






       def 				UpperCamelCase (							self					)			->	List[str]:
          _SCREAMING_SNAKE_CASE          = SwinaSRModelTester(self					)
          _SCREAMING_SNAKE_CASE          = ConfigTester(self				,							config_class=A__				,							embed_dim=37					)






       def 				UpperCamelCase (							self					)			->	Union[str, 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 				UpperCamelCase (							self					)			->	int:
          _SCREAMING_SNAKE_CASE          = self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*A__					)






       def 				UpperCamelCase (							self					)			->	List[Any]:
          _SCREAMING_SNAKE_CASE          = 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					)			->	List[Any]:
          pass






       @unittest.skip(reason="""Swin2SR does not use inputs_embeds"""					)
       def 				UpperCamelCase (							self					)			->	Union[str, Any]:
          pass






       @unittest.skip(reason="""Swin2SR does not support training yet"""					)
       def 				UpperCamelCase (							self					)			->	Any:
          pass






       @unittest.skip(reason="""Swin2SR does not support training yet"""					)
       def 				UpperCamelCase (							self					)			->	Tuple:
          pass






       def 				UpperCamelCase (							self					)			->	Dict:
          _SCREAMING_SNAKE_CASE       ,							_SCREAMING_SNAKE_CASE          = self.model_tester.prepare_config_and_inputs_for_common()

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






       def 				UpperCamelCase (							self					)			->	Union[str, Any]:
          _SCREAMING_SNAKE_CASE       ,							_SCREAMING_SNAKE_CASE          = self.model_tester.prepare_config_and_inputs_for_common()

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

             _SCREAMING_SNAKE_CASE          = ["""pixel_values"""]
             self.assertListEqual(arg_names[:1]				,							A__					)






       @slow
       def 				UpperCamelCase (							self					)			->	Optional[int]:
          for model_name in SWIN2SR_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             _SCREAMING_SNAKE_CASE          = SwinaSRModel.from_pretrained(A__					)
             self.assertIsNotNone(A__					)






       def 				UpperCamelCase (							self					)			->	List[str]:
          _SCREAMING_SNAKE_CASE       ,							_SCREAMING_SNAKE_CASE          = self.model_tester.prepare_config_and_inputs_for_common()

          _SCREAMING_SNAKE_CASE          = _config_zero_init(A__					)
          for model_class in self.all_model_classes:
             _SCREAMING_SNAKE_CASE          = 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					)			->	Any:
          _SCREAMING_SNAKE_CASE       ,							_SCREAMING_SNAKE_CASE          = self.model_tester.prepare_config_and_inputs_for_common()
          _SCREAMING_SNAKE_CASE          = True

          for model_class in self.all_model_classes:
             _SCREAMING_SNAKE_CASE          = True
             _SCREAMING_SNAKE_CASE          = False
             _SCREAMING_SNAKE_CASE          = True
             _SCREAMING_SNAKE_CASE          = model_class(A__					)
             model.to(A__					)
             model.eval()
             with torch.no_grad():
                _SCREAMING_SNAKE_CASE          = model(**self._prepare_for_class(A__				,							A__					)					)
             _SCREAMING_SNAKE_CASE          = outputs.attentions
             _SCREAMING_SNAKE_CASE          = len(self.model_tester.depths					)
             self.assertEqual(len(A__					)				,							A__					)

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

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

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

             _SCREAMING_SNAKE_CASE          = 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       _a (unittest.TestCase):





       """simple docstring"""




       def 				UpperCamelCase (							self					)			->	Dict:
          _SCREAMING_SNAKE_CASE          = SwinaSRImageProcessor()
          _SCREAMING_SNAKE_CASE          = SwinaSRForImageSuperResolution.from_pretrained("""caidas/swin2SR-classical-sr-x2-64"""					).to(A__					)

          _SCREAMING_SNAKE_CASE          = Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""					)
          _SCREAMING_SNAKE_CASE          = processor(images=A__				,							return_tensors="""pt"""					).to(A__					)

          # forward pass
          with torch.no_grad():
             _SCREAMING_SNAKE_CASE          = model(**A__					)

          # verify the logits
          _SCREAMING_SNAKE_CASE          = torch.Size([1, 3, 9_76, 12_96]					)
          self.assertEqual(outputs.reconstruction.shape				,							A__					)
          _SCREAMING_SNAKE_CASE          = torch.tensor(
              [[0.5458, 0.5546, 0.5638], [0.5526, 0.5565, 0.5651], [0.5396, 0.5426, 0.5621]]					).to(A__					)
          self.assertTrue(torch.allclose(outputs.reconstruction[0, 0, :3, :3]				,							A__				,							atol=1E-4					)					)






