
"""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							,  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"""
     UpperCAmelCase_   :   Dict     	=       parent
     UpperCAmelCase_   :   Dict     	=       batch_size
     UpperCAmelCase_   :   Optional[Any]     	=       image_size
     UpperCAmelCase_   :   List[Any]     	=       patch_size
     UpperCAmelCase_   :   List[Any]     	=       num_channels
     UpperCAmelCase_   :   str     	=       embed_dim
     UpperCAmelCase_   :   Any     	=       depths
     UpperCAmelCase_   :   Tuple     	=       num_heads
     UpperCAmelCase_   :   str     	=       window_size
     UpperCAmelCase_   :   Optional[Any]     	=       mlp_ratio
     UpperCAmelCase_   :   List[str]     	=       qkv_bias
     UpperCAmelCase_   :   Optional[int]     	=       hidden_dropout_prob
     UpperCAmelCase_   :   Optional[int]     	=       attention_probs_dropout_prob
     UpperCAmelCase_   :   Optional[Any]     	=       drop_path_rate
     UpperCAmelCase_   :   str     	=       hidden_act
     UpperCAmelCase_   :   int     	=       use_absolute_embeddings
     UpperCAmelCase_   :   int     	=       patch_norm
     UpperCAmelCase_   :   List[str]     	=       layer_norm_eps
     UpperCAmelCase_   :   Optional[Any]     	=       initializer_range
     UpperCAmelCase_   :   int     	=       is_training
     UpperCAmelCase_   :   str     	=       scope
     UpperCAmelCase_   :   Tuple     	=       use_labels
     UpperCAmelCase_   :   Union[str, Any]     	=       upscale

     # here we set some attributes to make tests pass
     UpperCAmelCase_   :   Tuple     	=       len(lowercase_    )
     UpperCAmelCase_   :   Tuple     	=       embed_dim
     UpperCAmelCase_   :   int     	=       (image_size // patch_size) ** 2







  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     UpperCAmelCase_   :   Tuple     	=       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]    )

     UpperCAmelCase_   :   int     	=       None
     if self.use_labels:
        UpperCAmelCase_   :   Union[str, Any]     	=       ids_tensor([self.batch_size]							,  self.type_sequence_label_size    )

     UpperCAmelCase_   :   Optional[int]     	=       self.get_config()

     return config, pixel_values, labels







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

     """simple docstring"""
     UpperCAmelCase_   :   str     	=       SwinaSRModel(config=lowercase_    )
     model.to(lowercase_    )
     model.eval()
     UpperCAmelCase_   :   Optional[Any]     	=       model(lowercase_    )

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







  def 			UpperCamelCase__   (						self							,  lowercase_							,  lowercase_							,  lowercase_    ):

     """simple docstring"""
     UpperCAmelCase_   :   Optional[Any]     	=       SwinaSRForImageSuperResolution(lowercase_    )
     model.to(lowercase_    )
     model.eval()
     UpperCAmelCase_   :   List[str]     	=       model(lowercase_    )

     UpperCAmelCase_   :   List[Any]     	=       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    ):

     """simple docstring"""
     UpperCAmelCase_   :   Any     	=       self.prepare_config_and_inputs()
     UpperCAmelCase_	,  UpperCAmelCase_	,  UpperCAmelCase_   :   Any     	=       config_and_inputs
     UpperCAmelCase_   :   Optional[Any]     	=       {"pixel_values": pixel_values}
     return config, inputs_dict




@require_torch
class A_      (lowercase__							,lowercase__							,unittest.TestCase ):




  '''simple docstring'''
  SCREAMING_SNAKE_CASE__     :		int    	=						(SwinaSRModel, SwinaSRForImageSuperResolution) if is_torch_available() else ()
  SCREAMING_SNAKE_CASE__     :		List[Any]    	=						{"""feature-extraction""": SwinaSRModel} if is_torch_available() else {}

  SCREAMING_SNAKE_CASE__     :		Tuple    	=						False
  SCREAMING_SNAKE_CASE__     :		Dict    	=						False
  SCREAMING_SNAKE_CASE__     :		Optional[int]    	=						False
  SCREAMING_SNAKE_CASE__     :		str    	=						False
  SCREAMING_SNAKE_CASE__     :		List[str]    	=						False







  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     UpperCAmelCase_   :   Optional[int]     	=       SwinaSRModelTester(self    )
     UpperCAmelCase_   :   int     	=       ConfigTester(self							,  config_class=lowercase_							,  embed_dim=37    )







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

     """simple docstring"""
     UpperCAmelCase_   :   Union[str, Any]     	=       self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*lowercase_    )







  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     UpperCAmelCase_   :   Optional[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 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     pass







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

     """simple docstring"""
     pass







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

     """simple docstring"""
     pass







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

     """simple docstring"""
     pass







  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     UpperCAmelCase_	,  UpperCAmelCase_   :   List[str]     	=       self.model_tester.prepare_config_and_inputs_for_common()

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







  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     UpperCAmelCase_	,  UpperCAmelCase_   :   Optional[Any]     	=       self.model_tester.prepare_config_and_inputs_for_common()

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

        UpperCAmelCase_   :   Union[str, Any]     	=       ["pixel_values"]
        self.assertListEqual(arg_names[:1]							,  lowercase_    )







  @slow
  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     for model_name in SWIN2SR_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        UpperCAmelCase_   :   Dict     	=       SwinaSRModel.from_pretrained(lowercase_    )
        self.assertIsNotNone(lowercase_    )







  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     UpperCAmelCase_	,  UpperCAmelCase_   :   Optional[int]     	=       self.model_tester.prepare_config_and_inputs_for_common()

     UpperCAmelCase_   :   Union[str, Any]     	=       _config_zero_init(lowercase_    )
     for model_class in self.all_model_classes:
        UpperCAmelCase_   :   Union[str, Any]     	=       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 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     UpperCAmelCase_	,  UpperCAmelCase_   :   Optional[Any]     	=       self.model_tester.prepare_config_and_inputs_for_common()
     UpperCAmelCase_   :   str     	=       True

     for model_class in self.all_model_classes:
        UpperCAmelCase_   :   Tuple     	=       True
        UpperCAmelCase_   :   Dict     	=       False
        UpperCAmelCase_   :   Any     	=       True
        UpperCAmelCase_   :   List[Any]     	=       model_class(lowercase_    )
        model.to(lowercase_    )
        model.eval()
        with torch.no_grad():
           UpperCAmelCase_   :   Tuple     	=       model(**self._prepare_for_class(lowercase_							,  lowercase_    )    )
        UpperCAmelCase_   :   Union[str, Any]     	=       outputs.attentions
        UpperCAmelCase_   :   Any     	=       len(self.model_tester.depths    )
        self.assertEqual(len(lowercase_    )							,  lowercase_    )

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

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

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

        UpperCAmelCase_   :   Optional[int]     	=       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 A_      (unittest.TestCase ):




  '''simple docstring'''







  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     UpperCAmelCase_   :   Union[str, Any]     	=       SwinaSRImageProcessor()
     UpperCAmelCase_   :   Tuple     	=       SwinaSRForImageSuperResolution.from_pretrained("caidas/swin2SR-classical-sr-x2-64"    ).to(lowercase_    )

     UpperCAmelCase_   :   Union[str, Any]     	=       Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"    )
     UpperCAmelCase_   :   List[Any]     	=       processor(images=lowercase_							,  return_tensors="pt"    ).to(lowercase_    )

     # forward pass
     with torch.no_grad():
        UpperCAmelCase_   :   List[Any]     	=       model(**lowercase_    )

     # verify the logits
     UpperCAmelCase_   :   Optional[Any]     	=       torch.Size([1, 3, 976, 1296]    )
     self.assertEqual(outputs.reconstruction.shape							,  lowercase_    )
     UpperCAmelCase_   :   List[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    )    )



"""simple docstring"""







from dataclasses import dataclass
from typing import Any, Dict, Optional

import torch
import torch.nn.functional as F
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..models.embeddings import ImagePositionalEmbeddings
from ..utils import BaseOutput, deprecate
from .attention import BasicTransformerBlock
from .embeddings import PatchEmbed
from .modeling_utils import ModelMixin



@dataclass
class 			_lowerCamelCase      (__lowerCamelCase  ):
				_snake_case  =	42



class 			_lowerCamelCase      (__lowerCamelCase    ,				__lowerCamelCase  ):







				@register_to_config
				def __init__(			self     :							Dict      ,       lowerCamelCase_     :							int = 1_6      ,       lowerCamelCase_     :							int = 8_8      ,       lowerCamelCase_     :							Optional[int] = None      ,       lowerCamelCase_     :							Optional[int] = None      ,       lowerCamelCase_     :							int = 1      ,       lowerCamelCase_     :							float = 0.0      ,       lowerCamelCase_     :							int = 3_2      ,       lowerCamelCase_     :							Optional[int] = None      ,       lowerCamelCase_     :							bool = False      ,       lowerCamelCase_     :							Optional[int] = None      ,       lowerCamelCase_     :							Optional[int] = None      ,       lowerCamelCase_     :							Optional[int] = None      ,       lowerCamelCase_     :							str = "geglu"      ,       lowerCamelCase_     :							Optional[int] = None      ,       lowerCamelCase_     :							bool = False      ,       lowerCamelCase_     :							bool = False      ,       lowerCamelCase_     :							bool = False      ,       lowerCamelCase_     :							str = "layer_norm"      ,       lowerCamelCase_     :							bool = True      ,       ):

					"""simple docstring"""






					super().__init__()
					_lowercase			:     Tuple      					=				use_linear_projection
					_lowercase			:     List[Any]      					=				num_attention_heads
					_lowercase			:     Optional[int]      					=				attention_head_dim
					_lowercase			:     Union[str, Any]      					=				num_attention_heads * attention_head_dim

					# 1. Transformer2DModel can process both standard continuous images of shape `(batch_size, num_channels, width, height)` as well as quantized image embeddings of shape `(batch_size, num_image_vectors)`
					# Define whether input is continuous or discrete depending on configuration
					_lowercase			:     Dict      					=				(in_channels is not None) and (patch_size is None)
					_lowercase			:     Any      					=				num_vector_embeds is not None
					_lowercase			:     List[Any]      					=				in_channels is not None and patch_size is not None

					if norm_type == "layer_norm" and num_embeds_ada_norm is not None:
						_lowercase			:     Tuple      					=				(
						    F'''The configuration file of this model: {self.__class__} is outdated. `norm_type` is either not set or'''
						    ' incorrectly set to `\'layer_norm\'`.Make sure to set `norm_type` to `\'ada_norm\'` in the config.'
						    ' Please make sure to update the config accordingly as leaving `norm_type` might led to incorrect'
						    ' results in future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it'
						    ' would be very nice if you could open a Pull request for the `transformer/config.json` file'
						)
						deprecate('norm_type!=num_embeds_ada_norm'      ,       '1.0.0'      ,       lowerCamelCase_      ,       standard_warn=lowerCamelCase_    )
						_lowercase			:     Dict      					=				'ada_norm'

					if self.is_input_continuous and self.is_input_vectorized:
						raise ValueError(
						    F'''Cannot define both `in_channels`: {in_channels} and `num_vector_embeds`: {num_vector_embeds}. Make'''
						    ' sure that either `in_channels` or `num_vector_embeds` is None.'    )
					elif self.is_input_vectorized and self.is_input_patches:
						raise ValueError(
						    F'''Cannot define both `num_vector_embeds`: {num_vector_embeds} and `patch_size`: {patch_size}. Make'''
						    ' sure that either `num_vector_embeds` or `num_patches` is None.'    )
					elif not self.is_input_continuous and not self.is_input_vectorized and not self.is_input_patches:
						raise ValueError(
						    F'''Has to define `in_channels`: {in_channels}, `num_vector_embeds`: {num_vector_embeds}, or patch_size:'''
						    F''' {patch_size}. Make sure that `in_channels`, `num_vector_embeds` or `num_patches` is not None.'''    )

					# 2. Define input layers
					if self.is_input_continuous:
						_lowercase			:     int      					=				in_channels

						_lowercase			:     List[Any]      					=				torch.nn.GroupNorm(num_groups=lowerCamelCase_      ,       num_channels=lowerCamelCase_      ,       eps=1E-6      ,       affine=lowerCamelCase_    )
						if use_linear_projection:
							_lowercase			:     Any      					=				nn.Linear(lowerCamelCase_      ,       lowerCamelCase_    )
						else:
							_lowercase			:     List[str]      					=				nn.Convad(lowerCamelCase_      ,       lowerCamelCase_      ,       kernel_size=1      ,       stride=1      ,       padding=0    )
					elif self.is_input_vectorized:
						assert sample_size is not None, "Transformer2DModel over discrete input must provide sample_size"
						assert num_vector_embeds is not None, "Transformer2DModel over discrete input must provide num_embed"

						_lowercase			:     int      					=				sample_size
						_lowercase			:     str      					=				sample_size
						_lowercase			:     int      					=				num_vector_embeds
						_lowercase			:     List[str]      					=				self.height * self.width

						_lowercase			:     Optional[int]      					=				ImagePositionalEmbeddings(
						    num_embed=lowerCamelCase_      ,       embed_dim=lowerCamelCase_      ,       height=self.height      ,       width=self.width    )
					elif self.is_input_patches:
						assert sample_size is not None, "Transformer2DModel over patched input must provide sample_size"

						_lowercase			:     Optional[Any]      					=				sample_size
						_lowercase			:     Optional[int]      					=				sample_size

						_lowercase			:     str      					=				patch_size
						_lowercase			:     Union[str, Any]      					=				PatchEmbed(
						    height=lowerCamelCase_      ,       width=lowerCamelCase_      ,       patch_size=lowerCamelCase_      ,       in_channels=lowerCamelCase_      ,       embed_dim=lowerCamelCase_      ,       )

					# 3. Define transformers blocks
					_lowercase			:     Any      					=				nn.ModuleList(
					    [
					        BasicTransformerBlock(
					            lowerCamelCase_      ,       lowerCamelCase_      ,       lowerCamelCase_      ,       dropout=lowerCamelCase_      ,       cross_attention_dim=lowerCamelCase_      ,       activation_fn=lowerCamelCase_      ,       num_embeds_ada_norm=lowerCamelCase_      ,       attention_bias=lowerCamelCase_      ,       only_cross_attention=lowerCamelCase_      ,       upcast_attention=lowerCamelCase_      ,       norm_type=lowerCamelCase_      ,       norm_elementwise_affine=lowerCamelCase_      ,       )
					        for d in range(lowerCamelCase_    )
					    ]    )

					# 4. Define output layers
					_lowercase			:     Dict      					=				in_channels if out_channels is None else out_channels
					if self.is_input_continuous:
						# TODO: should use out_channels for continuous projections
						if use_linear_projection:
							_lowercase			:     str      					=				nn.Linear(lowerCamelCase_      ,       lowerCamelCase_    )
						else:
							_lowercase			:     List[Any]      					=				nn.Convad(lowerCamelCase_      ,       lowerCamelCase_      ,       kernel_size=1      ,       stride=1      ,       padding=0    )
					elif self.is_input_vectorized:
						_lowercase			:     Dict      					=				nn.LayerNorm(lowerCamelCase_    )
						_lowercase			:     Union[str, Any]      					=				nn.Linear(lowerCamelCase_      ,       self.num_vector_embeds - 1    )
					elif self.is_input_patches:
						_lowercase			:     Dict      					=				nn.LayerNorm(lowerCamelCase_      ,       elementwise_affine=lowerCamelCase_      ,       eps=1E-6    )
						_lowercase			:     Union[str, Any]      					=				nn.Linear(lowerCamelCase_      ,       2 * inner_dim    )
						_lowercase			:     List[Any]      					=				nn.Linear(lowerCamelCase_      ,       patch_size * patch_size * self.out_channels    )







				def 				__UpperCAmelCase       (			self     :							Optional[Any]      ,       lowerCamelCase_     :							torch.Tensor      ,       lowerCamelCase_     :							Optional[torch.Tensor] = None      ,       lowerCamelCase_     :							Optional[torch.LongTensor] = None      ,       lowerCamelCase_     :							Optional[torch.LongTensor] = None      ,       lowerCamelCase_     :							Dict[str, Any] = None      ,       lowerCamelCase_     :							Optional[torch.Tensor] = None      ,       lowerCamelCase_     :							Optional[torch.Tensor] = None      ,       lowerCamelCase_     :							bool = True      ,       ):

					"""simple docstring"""






					if attention_mask is not None and attention_mask.ndim == 2:
						# assume that mask is expressed as:
						#   (1 = keep,      0 = discard)
						# convert mask into a bias that can be added to attention scores:
						#       (keep = +0,     discard = -10000.0)
						_lowercase			:     Optional[Any]      					=				(1 - attention_mask.to(hidden_states.dtype    )) * -1_0000.0
						_lowercase			:     Optional[Any]      					=				attention_mask.unsqueeze(1    )

					# convert encoder_attention_mask to a bias the same way we do for attention_mask
					if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2:
						_lowercase			:     List[Any]      					=				(1 - encoder_attention_mask.to(hidden_states.dtype    )) * -1_0000.0
						_lowercase			:     Any      					=				encoder_attention_mask.unsqueeze(1    )

					# 1. Input
					if self.is_input_continuous:
						_lowercase					,   _lowercase					,   _lowercase					,   _lowercase			:     List[Any]      					=				hidden_states.shape
						_lowercase			:     int      					=				hidden_states

						_lowercase			:     Any      					=				self.norm(lowerCamelCase_    )
						if not self.use_linear_projection:
							_lowercase			:     Any      					=				self.proj_in(lowerCamelCase_    )
							_lowercase			:     Any      					=				hidden_states.shape[1]
							_lowercase			:     List[Any]      					=				hidden_states.permute(0      ,       2      ,       3      ,       1    ).reshape(lowerCamelCase_      ,       height * width      ,       lowerCamelCase_    )
						else:
							_lowercase			:     Any      					=				hidden_states.shape[1]
							_lowercase			:     int      					=				hidden_states.permute(0      ,       2      ,       3      ,       1    ).reshape(lowerCamelCase_      ,       height * width      ,       lowerCamelCase_    )
							_lowercase			:     Any      					=				self.proj_in(lowerCamelCase_    )
					elif self.is_input_vectorized:
						_lowercase			:     Optional[Any]      					=				self.latent_image_embedding(lowerCamelCase_    )
					elif self.is_input_patches:
						_lowercase			:     Union[str, Any]      					=				self.pos_embed(lowerCamelCase_    )

					# 2. Blocks
					for block in self.transformer_blocks:
						_lowercase			:     Tuple      					=				block(
						    lowerCamelCase_      ,       attention_mask=lowerCamelCase_      ,       encoder_hidden_states=lowerCamelCase_      ,       encoder_attention_mask=lowerCamelCase_      ,       timestep=lowerCamelCase_      ,       cross_attention_kwargs=lowerCamelCase_      ,       class_labels=lowerCamelCase_      ,       )

					# 3. Output
					if self.is_input_continuous:
						if not self.use_linear_projection:
							_lowercase			:     Union[str, Any]      					=				hidden_states.reshape(lowerCamelCase_      ,       lowerCamelCase_      ,       lowerCamelCase_      ,       lowerCamelCase_    ).permute(0      ,       3      ,       1      ,       2    ).contiguous()
							_lowercase			:     Tuple      					=				self.proj_out(lowerCamelCase_    )
						else:
							_lowercase			:     int      					=				self.proj_out(lowerCamelCase_    )
							_lowercase			:     List[Any]      					=				hidden_states.reshape(lowerCamelCase_      ,       lowerCamelCase_      ,       lowerCamelCase_      ,       lowerCamelCase_    ).permute(0      ,       3      ,       1      ,       2    ).contiguous()

						_lowercase			:     List[str]      					=				hidden_states + residual
					elif self.is_input_vectorized:
						_lowercase			:     Dict      					=				self.norm_out(lowerCamelCase_    )
						_lowercase			:     Dict      					=				self.out(lowerCamelCase_    )
						# (batch, self.num_vector_embeds - 1, self.num_latent_pixels)
						_lowercase			:     Dict      					=				logits.permute(0      ,       2      ,       1    )

						# log(p(x_0))
						_lowercase			:     int      					=				F.log_softmax(logits.double()      ,       dim=1    ).float()
					elif self.is_input_patches:
						# TODO: cleanup!
						_lowercase			:     Tuple      					=				self.transformer_blocks[0].norma.emb(
						    lowerCamelCase_      ,       lowerCamelCase_      ,       hidden_dtype=hidden_states.dtype    )
						_lowercase					,   _lowercase			:     List[Any]      					=				self.proj_out_a(F.silu(lowerCamelCase_    )    ).chunk(2      ,       dim=1    )
						_lowercase			:     Optional[int]      					=				self.norm_out(lowerCamelCase_    ) * (1 + scale[:, None]) + shift[:, None]
						_lowercase			:     Optional[Any]      					=				self.proj_out_a(lowerCamelCase_    )

						# unpatchify
						_lowercase			:     Union[str, Any]      					=				int(hidden_states.shape[1] ** 0.5    )
						_lowercase			:     int      					=				hidden_states.reshape(
						    shape=(-1, height, width, self.patch_size, self.patch_size, self.out_channels)    )
						_lowercase			:     str      					=				torch.einsum('nhwpqc->nchpwq'      ,       lowerCamelCase_    )
						_lowercase			:     List[str]      					=				hidden_states.reshape(
						    shape=(-1, self.out_channels, height * self.patch_size, width * self.patch_size)    )

					if not return_dict:
						return (output,)

					return TransformeraDModelOutput(sample=lowerCamelCase_    )



