

"""simple docstring"""


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

import flax
import jax.numpy as jnp

from ..configuration_utils import ConfigMixin, register_to_config
from .scheduling_utils_flax import (
    CommonSchedulerState,
    FlaxKarrasDiffusionSchedulers,
    FlaxSchedulerMixin,
    FlaxSchedulerOutput,
    add_noise_common,
    get_velocity_common,
)





@flax.struct.dataclass
class 					snake_case			:
				SCREAMING_SNAKE_CASE_     :   int        =						42
				SCREAMING_SNAKE_CASE_     :   Union[str, Any]        =						42

				# setable values
				SCREAMING_SNAKE_CASE_     :   Tuple        =						42
				SCREAMING_SNAKE_CASE_     :   Dict        =						42
				SCREAMING_SNAKE_CASE_     :   List[str]        =						None

				@classmethod
				def 							lowercase_  (			cls							:       int  ,						UpperCamelCase__							:       Optional[Any]  ,						UpperCamelCase__							:       Union[str, Any]  ,						UpperCamelCase__							:       str  ,						UpperCamelCase__							:       str  ,						)->     List[Any]:



						'''simple docstring'''


						return cls(
						    common=UpperCamelCase__  ,						final_alpha_cumprod=UpperCamelCase__  ,						init_noise_sigma=UpperCamelCase__  ,						timesteps=UpperCamelCase__  ,						)





@dataclass
class 					snake_case			(       snake_case_						):
				SCREAMING_SNAKE_CASE_     :   Optional[int]        =						42





class 					snake_case			(       snake_case_,							snake_case_						):
				SCREAMING_SNAKE_CASE_     :   List[str]        =						[e.name for e in FlaxKarrasDiffusionSchedulers]

				SCREAMING_SNAKE_CASE_     :   List[str]        =						42

				@property
				def 							lowercase_  (			self							:       Any)->     List[Any]:



						'''simple docstring'''


						return True

				@register_to_config
				def __init__(			self							:       Dict  ,						UpperCamelCase__							:       List[Any] = 1_0_0_0  ,						UpperCamelCase__							:       Union[str, Any] = 0.0001  ,						UpperCamelCase__							:       List[str] = 0.02  ,						UpperCamelCase__							:       Optional[Any] = "linear"  ,						UpperCamelCase__							:       List[str] = None  ,						UpperCamelCase__							:       Dict = True  ,						UpperCamelCase__							:       int = 0  ,						UpperCamelCase__							:       Optional[int] = "epsilon"  ,						UpperCamelCase__							:       Any = jnp.floataa  ,						)->     int:



						'''simple docstring'''


						__lowerCAmelCase: str					   =    dtype

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



						'''simple docstring'''


						if common is None:
								__lowerCAmelCase: List[str]					   =    CommonSchedulerState.create(self)

						# At every step in ddim, we are looking into the previous alphas_cumprod
						# For the final step, there is no previous alphas_cumprod because we are already at 0
						# `set_alpha_to_one` decides whether we set this parameter simply to one or
						# whether we use the final alpha of the "non-previous" one.
						__lowerCAmelCase: Dict					   =    (
						    jnp.array(1.0  ,						dtype=self.dtype) if self.config.set_alpha_to_one else common.alphas_cumprod[0]
						)

						# standard deviation of the initial noise distribution
						__lowerCAmelCase: Tuple					   =    jnp.array(1.0  ,						dtype=self.dtype)

						__lowerCAmelCase: Union[str, Any]					   =    jnp.arange(0  ,						self.config.num_train_timesteps).round()[::-1]

						return DDIMSchedulerState.create(
						    common=UpperCamelCase__  ,						final_alpha_cumprod=UpperCamelCase__  ,						init_noise_sigma=UpperCamelCase__  ,						timesteps=UpperCamelCase__  ,						)

				def 							lowercase_  (			self							:       Optional[Any]  ,						UpperCamelCase__							:       Optional[Any]  ,						UpperCamelCase__							:       str  ,						UpperCamelCase__							:       List[str] = None)->     jnp.ndarray:



						'''simple docstring'''


						return sample

				def 							lowercase_  (			self							:       int  ,						UpperCamelCase__							:       int  ,						UpperCamelCase__							:       List[Any]  ,						UpperCamelCase__							:       Union[str, Any] = ())->     DDIMSchedulerState:



						'''simple docstring'''


						__lowerCAmelCase: Dict					   =    self.config.num_train_timesteps // num_inference_steps
						# creates integer timesteps by multiplying by ratio
						# rounding to avoid issues when num_inference_step is power of 3
						__lowerCAmelCase: str					   =    (jnp.arange(0  ,						UpperCamelCase__) * step_ratio).round()[::-1] + self.config.steps_offset

						return state.replace(
						    num_inference_steps=UpperCamelCase__  ,						timesteps=UpperCamelCase__  ,						)

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



						'''simple docstring'''


						__lowerCAmelCase: Union[str, Any]					   =    state.common.alphas_cumprod[timestep]
						__lowerCAmelCase: List[Any]					   =    jnp.where(
						    prev_timestep >= 0  ,						state.common.alphas_cumprod[prev_timestep]  ,						state.final_alpha_cumprod)
						__lowerCAmelCase: Union[str, Any]					   =    1 - alpha_prod_t
						__lowerCAmelCase: str					   =    1 - alpha_prod_t_prev

						__lowerCAmelCase: str					   =    (beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev)

						return variance

				def 							lowercase_  (			self							:       Union[str, Any]  ,						UpperCamelCase__							:       Union[str, Any]  ,						UpperCamelCase__							:       Tuple  ,						UpperCamelCase__							:       Union[str, Any]  ,						UpperCamelCase__							:       Optional[Any]  ,						UpperCamelCase__							:       Optional[int] = 0.0  ,						UpperCamelCase__							:       Dict = True  ,						)->     Union[FlaxDDIMSchedulerOutput, Tuple]:



						'''simple docstring'''


						if state.num_inference_steps is None:
								raise ValueError(
								    "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler")

						# See formulas (12) and (16) of DDIM paper https://arxiv.org/pdf/2010.02502.pdf
						# Ideally, read DDIM paper in-detail understanding

						# Notation (<variable name> -> <name in paper>
						# - pred_noise_t -> e_theta(x_t, t)
						# - pred_original_sample -> f_theta(x_t, t) or x_0
						# - std_dev_t -> sigma_t
						# - eta -> η
						# - pred_sample_direction -> "direction pointing to x_t"
						# - pred_prev_sample -> "x_t-1"

						# 1. get previous step value (=t-1)
						__lowerCAmelCase: Tuple					   =    timestep - self.config.num_train_timesteps // state.num_inference_steps

						__lowerCAmelCase: Union[str, Any]					   =    state.common.alphas_cumprod
						__lowerCAmelCase: Tuple					   =    state.final_alpha_cumprod

						# 2. compute alphas, betas
						__lowerCAmelCase: int					   =    alphas_cumprod[timestep]
						__lowerCAmelCase: List[Any]					   =    jnp.where(prev_timestep >= 0  ,						alphas_cumprod[prev_timestep]  ,						UpperCamelCase__)

						__lowerCAmelCase: List[Any]					   =    1 - alpha_prod_t

						# 3. compute predicted original sample from predicted noise also called
						# "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
						if self.config.prediction_type == "epsilon":
								__lowerCAmelCase: int					   =    (sample - beta_prod_t ** 0.5 * model_output) / alpha_prod_t ** 0.5
								__lowerCAmelCase: Any					   =    model_output
						elif self.config.prediction_type == "sample":
								__lowerCAmelCase: List[Any]					   =    model_output
								__lowerCAmelCase: List[Any]					   =    (sample - alpha_prod_t ** 0.5 * pred_original_sample) / beta_prod_t ** 0.5
						elif self.config.prediction_type == "v_prediction":
								__lowerCAmelCase: List[str]					   =    (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output
								__lowerCAmelCase: Optional[Any]					   =    (alpha_prod_t**0.5) * model_output + (beta_prod_t**0.5) * sample
						else:
								raise ValueError(
								    f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or"
								    " `v_prediction`")

						# 4. compute variance: "sigma_t(η)" -> see formula (16)
						# σ_t = sqrt((1 − α_t−1)/(1 − α_t)) * sqrt(1 − α_t/α_t−1)
						__lowerCAmelCase: Dict					   =    self._get_variance(UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__)
						__lowerCAmelCase: Any					   =    eta * variance ** 0.5

						# 5. compute "direction pointing to x_t" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
						__lowerCAmelCase: List[str]					   =    (1 - alpha_prod_t_prev - std_dev_t**2) ** 0.5 * pred_epsilon

						# 6. compute x_t without "random noise" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
						__lowerCAmelCase: List[Any]					   =    alpha_prod_t_prev ** 0.5 * pred_original_sample + pred_sample_direction

						if not return_dict:
								return (prev_sample, state)

						return FlaxDDIMSchedulerOutput(prev_sample=UpperCamelCase__  ,						state=UpperCamelCase__)

				def 							lowercase_  (			self							:       Tuple  ,						UpperCamelCase__							:       Any  ,						UpperCamelCase__							:       Optional[int]  ,						UpperCamelCase__							:       Union[str, Any]  ,						UpperCamelCase__							:       str  ,						)->     jnp.ndarray:



						'''simple docstring'''


						return add_noise_common(state.common  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__)

				def 							lowercase_  (			self							:       Any  ,						UpperCamelCase__							:       Optional[Any]  ,						UpperCamelCase__							:       Union[str, Any]  ,						UpperCamelCase__							:       Union[str, Any]  ,						UpperCamelCase__							:       Dict  ,						)->     jnp.ndarray:



						'''simple docstring'''


						return get_velocity_common(state.common  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__)




				def __len__(			self							:       Optional[Any])->     Any:



						'''simple docstring'''


						return self.config.num_train_timesteps





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


