

import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionModelEditingPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()



@skip_mps
class 					__snake_case     (			lowerCAmelCase__   ,       lowerCAmelCase__   ,       lowerCAmelCase__   ,       unittest.TestCase				):
  __lowerCAmelCase   :		Tuple         =     StableDiffusionModelEditingPipeline
  __lowerCAmelCase   :		Tuple         =     TEXT_TO_IMAGE_PARAMS
  __lowerCAmelCase   :		List[str]         =     TEXT_TO_IMAGE_BATCH_PARAMS
  __lowerCAmelCase   :		Optional[int]         =     TEXT_TO_IMAGE_IMAGE_PARAMS
  __lowerCAmelCase   :		Dict         =     TEXT_TO_IMAGE_IMAGE_PARAMS

  def   lowerCAmelCase__   (    self):
      torch.manual_seed(0)
      SCREAMING_SNAKE_CASE_													=  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  ,						)
      SCREAMING_SNAKE_CASE_													=  DDIMScheduler()
      torch.manual_seed(0)
      SCREAMING_SNAKE_CASE_													=  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)
      SCREAMING_SNAKE_CASE_													=  CLIPTextConfig(
          bos_token_id=0  ,						eos_token_id=2  ,						hidden_size=32  ,						intermediate_size=37  ,						layer_norm_eps=1E-05  ,						num_attention_heads=4  ,						num_hidden_layers=5  ,						pad_token_id=1  ,						vocab_size=1000  ,						)
      SCREAMING_SNAKE_CASE_													=  CLIPTextModel(_A)
      SCREAMING_SNAKE_CASE_													=  CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip')

      SCREAMING_SNAKE_CASE_													=  {
          'unet': unet,
          'scheduler': scheduler,
          'vae': vae,
          'text_encoder': text_encoder,
          'tokenizer': tokenizer,
          'safety_checker': None,
          'feature_extractor': None,
      }
      return components

  def   lowerCAmelCase__   (    self  ,						_A  ,						_A=0):
      SCREAMING_SNAKE_CASE_													=  torch.manual_seed(_A)
      SCREAMING_SNAKE_CASE_													=  {
          'prompt': 'A field of roses',
          'generator': generator,
          # Setting height and width to None to prevent OOMs on CPU.
          'height': None,
          'width': None,
          'num_inference_steps': 2,
          'guidance_scale': 6.0,
          'output_type': 'numpy',
      }
      return inputs

  def   lowerCAmelCase__   (    self):
      SCREAMING_SNAKE_CASE_													=  'cpu'  # ensure determinism for the device-dependent torch.Generator
      SCREAMING_SNAKE_CASE_													=  self.get_dummy_components()
      SCREAMING_SNAKE_CASE_													=  StableDiffusionModelEditingPipeline(**_A)
      SCREAMING_SNAKE_CASE_													=  sd_pipe.to(_A)
      sd_pipe.set_progress_bar_config(disable=_A)

      SCREAMING_SNAKE_CASE_													=  self.get_dummy_inputs(_A)
      SCREAMING_SNAKE_CASE_													=  sd_pipe(**_A).images
      SCREAMING_SNAKE_CASE_													=  image[0, -3:, -3:, -1]
      assert image.shape == (1, 64, 64, 3)

      SCREAMING_SNAKE_CASE_													=  np.array([0.4_7_5_5, 0.5_1_3_2, 0.4_9_7_6, 0.3_9_0_4, 0.3_5_5_4, 0.4_7_6_5, 0.5_1_3_9, 0.5_1_5_8, 0.4_8_8_9])

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

  def   lowerCAmelCase__   (    self):
      SCREAMING_SNAKE_CASE_													=  'cpu'  # ensure determinism for the device-dependent torch.Generator
      SCREAMING_SNAKE_CASE_													=  self.get_dummy_components()
      SCREAMING_SNAKE_CASE_													=  StableDiffusionModelEditingPipeline(**_A)
      SCREAMING_SNAKE_CASE_													=  sd_pipe.to(_A)
      sd_pipe.set_progress_bar_config(disable=_A)

      SCREAMING_SNAKE_CASE_													=  self.get_dummy_inputs(_A)
      SCREAMING_SNAKE_CASE_													=  'french fries'
      SCREAMING_SNAKE_CASE_													=  sd_pipe(**_A  ,						negative_prompt=_A)
      SCREAMING_SNAKE_CASE_													=  output.images
      SCREAMING_SNAKE_CASE_													=  image[0, -3:, -3:, -1]

      assert image.shape == (1, 64, 64, 3)

      SCREAMING_SNAKE_CASE_													=  np.array([0.4_9_9_2, 0.5_1_0_1, 0.5_0_0_4, 0.3_9_4_9, 0.3_6_0_4, 0.4_7_3_5, 0.5_2_1_6, 0.5_2_0_4, 0.4_9_1_3])

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

  def   lowerCAmelCase__   (    self):
      SCREAMING_SNAKE_CASE_													=  'cpu'  # ensure determinism for the device-dependent torch.Generator
      SCREAMING_SNAKE_CASE_													=  self.get_dummy_components()
      SCREAMING_SNAKE_CASE_													=  EulerAncestralDiscreteScheduler(
          beta_start=0.0_0_0_8_5  ,						beta_end=0.0_1_2  ,						beta_schedule='scaled_linear')
      SCREAMING_SNAKE_CASE_													=  StableDiffusionModelEditingPipeline(**_A)
      SCREAMING_SNAKE_CASE_													=  sd_pipe.to(_A)
      sd_pipe.set_progress_bar_config(disable=_A)

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

      assert image.shape == (1, 64, 64, 3)

      SCREAMING_SNAKE_CASE_													=  np.array([0.4_7_4_7, 0.5_3_7_2, 0.4_7_7_9, 0.4_9_8_2, 0.5_5_4_3, 0.4_8_1_6, 0.5_2_3_8, 0.4_9_0_4, 0.5_0_2_7])

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

  def   lowerCAmelCase__   (    self):
      SCREAMING_SNAKE_CASE_													=  'cpu'  # ensure determinism for the device-dependent torch.Generator
      SCREAMING_SNAKE_CASE_													=  self.get_dummy_components()
      SCREAMING_SNAKE_CASE_													=  PNDMScheduler()
      SCREAMING_SNAKE_CASE_													=  StableDiffusionModelEditingPipeline(**_A)
      SCREAMING_SNAKE_CASE_													=  sd_pipe.to(_A)
      sd_pipe.set_progress_bar_config(disable=_A)

      SCREAMING_SNAKE_CASE_													=  self.get_dummy_inputs(_A)
      # the pipeline does not expect pndm so test if it raises error.
      with self.assertRaises(_A):
          SCREAMING_SNAKE_CASE_													=  sd_pipe(**_A).images

  def   lowerCAmelCase__   (    self):
      super().test_inference_batch_single_identical(expected_max_diff=5E-3)


  def   lowerCAmelCase__   (    self):
      super().test_attention_slicing_forward_pass(expected_max_diff=5E-3)






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

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

  def   lowerCAmelCase__   (    self  ,						_A=0):
      SCREAMING_SNAKE_CASE_													=  torch.manual_seed(_A)
      SCREAMING_SNAKE_CASE_													=  {
          'prompt': 'A field of roses',
          'generator': generator,
          'num_inference_steps': 3,
          'guidance_scale': 7.5,
          'output_type': 'numpy',
      }
      return inputs

  def   lowerCAmelCase__   (    self):
      SCREAMING_SNAKE_CASE_													=  'CompVis/stable-diffusion-v1-4'
      SCREAMING_SNAKE_CASE_													=  StableDiffusionModelEditingPipeline.from_pretrained(_A  ,						safety_checker=_A)
      pipe.to(_A)
      pipe.set_progress_bar_config(disable=_A)
      pipe.enable_attention_slicing()

      SCREAMING_SNAKE_CASE_													=  self.get_inputs()
      SCREAMING_SNAKE_CASE_													=  pipe(**_A).images
      SCREAMING_SNAKE_CASE_													=  image[0, -3:, -3:, -1].flatten()

      assert image.shape == (1, 512, 512, 3)

      SCREAMING_SNAKE_CASE_													=  np.array(
          [0.6_7_4_9_4_9_6, 0.6_3_8_6_4_5_3, 0.5_1_4_4_3_2_6_7, 0.6_6_0_9_4_9_0_5, 0.6_1_9_2_1_2_1_5, 0.5_4_9_1_3_3_2, 0.5_7_4_4_4_1_7, 0.5_8_0_7_5_1_0_6, 0.5_1_7_4_6_5_8])

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

      # make sure image changes after editing
      pipe.edit_model('A pack of roses'  ,						'A pack of blue roses')

      SCREAMING_SNAKE_CASE_													=  pipe(**_A).images
      SCREAMING_SNAKE_CASE_													=  image[0, -3:, -3:, -1].flatten()

      assert image.shape == (1, 512, 512, 3)

      assert np.abs(expected_slice - image_slice).max() > 1E-1


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

      SCREAMING_SNAKE_CASE_													=  'CompVis/stable-diffusion-v1-4'
      SCREAMING_SNAKE_CASE_													=  DDIMScheduler.from_pretrained(_A  ,						subfolder='scheduler')
      SCREAMING_SNAKE_CASE_													=  StableDiffusionModelEditingPipeline.from_pretrained(
          _A  ,						scheduler=_A  ,						safety_checker=_A)
      SCREAMING_SNAKE_CASE_													=  pipe.to(_A)
      pipe.set_progress_bar_config(disable=_A)
      pipe.enable_attention_slicing(1)
      pipe.enable_sequential_cpu_offload()

      SCREAMING_SNAKE_CASE_													=  self.get_inputs()
      SCREAMING_SNAKE_CASE_													=  pipe(**_A)

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





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 					__snake_case     (			lowerCAmelCase__				):
  __lowerCAmelCase   :		torch.FloatTensor






class 					__snake_case     (			lowerCAmelCase__   ,       lowerCAmelCase__				):

  @register_to_config
  def __init__(    self  ,						_A = 16  ,						_A = 88  ,						_A = None  ,						_A = None  ,						_A = 1  ,						_A = 0.0  ,						_A = 32  ,						_A = None  ,						_A = False  ,						_A = None  ,						_A = None  ,						_A = None  ,						_A = "geglu"  ,						_A = None  ,						_A = False  ,						_A = False  ,						_A = False  ,						_A = "layer_norm"  ,						_A = True  ,						):
      super().__init__()
      SCREAMING_SNAKE_CASE_													=  use_linear_projection
      SCREAMING_SNAKE_CASE_													=  num_attention_heads
      SCREAMING_SNAKE_CASE_													=  attention_head_dim
      SCREAMING_SNAKE_CASE_													=  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
      SCREAMING_SNAKE_CASE_													=  (in_channels is not None) and (patch_size is None)
      SCREAMING_SNAKE_CASE_													=  num_vector_embeds is not None
      SCREAMING_SNAKE_CASE_													=  in_channels is not None and patch_size is not None

      if norm_type == "layer_norm" and num_embeds_ada_norm is not None:
          SCREAMING_SNAKE_CASE_													=  (
              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'  ,						_A  ,						standard_warn=_A)
          SCREAMING_SNAKE_CASE_													=  '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:
          SCREAMING_SNAKE_CASE_													=  in_channels

          SCREAMING_SNAKE_CASE_													=  torch.nn.GroupNorm(num_groups=_A  ,						num_channels=_A  ,						eps=1E-6  ,						affine=_A)
          if use_linear_projection:
              SCREAMING_SNAKE_CASE_													=  nn.Linear(_A  ,						_A)
          else:
              SCREAMING_SNAKE_CASE_													=  nn.Convad(_A  ,						_A  ,						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"

          SCREAMING_SNAKE_CASE_													=  sample_size
          SCREAMING_SNAKE_CASE_													=  sample_size
          SCREAMING_SNAKE_CASE_													=  num_vector_embeds
          SCREAMING_SNAKE_CASE_													=  self.height * self.width

          SCREAMING_SNAKE_CASE_													=  ImagePositionalEmbeddings(
              num_embed=_A  ,						embed_dim=_A  ,						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"

          SCREAMING_SNAKE_CASE_													=  sample_size
          SCREAMING_SNAKE_CASE_													=  sample_size

          SCREAMING_SNAKE_CASE_													=  patch_size
          SCREAMING_SNAKE_CASE_													=  PatchEmbed(
              height=_A  ,						width=_A  ,						patch_size=_A  ,						in_channels=_A  ,						embed_dim=_A  ,						)

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

      # 4. Define output layers
      SCREAMING_SNAKE_CASE_													=  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:
              SCREAMING_SNAKE_CASE_													=  nn.Linear(_A  ,						_A)
          else:
              SCREAMING_SNAKE_CASE_													=  nn.Convad(_A  ,						_A  ,						kernel_size=1  ,						stride=1  ,						padding=0)
      elif self.is_input_vectorized:
          SCREAMING_SNAKE_CASE_													=  nn.LayerNorm(_A)
          SCREAMING_SNAKE_CASE_													=  nn.Linear(_A  ,						self.num_vector_embeds - 1)
      elif self.is_input_patches:
          SCREAMING_SNAKE_CASE_													=  nn.LayerNorm(_A  ,						elementwise_affine=_A  ,						eps=1E-6)
          SCREAMING_SNAKE_CASE_													=  nn.Linear(_A  ,						2 * inner_dim)
          SCREAMING_SNAKE_CASE_													=  nn.Linear(_A  ,						patch_size * patch_size * self.out_channels)


  def   lowerCAmelCase__   (    self  ,						_A  ,						_A = None  ,						_A = None  ,						_A = None  ,						_A = None  ,						_A = None  ,						_A = None  ,						_A = True  ,						):
      # ensure attention_mask is a bias, and give it a singleton query_tokens dimension.
      #   we may have done this conversion already, e.g. if we came here via UNet2DConditionModel#forward.
      #   we can tell by counting dims; if ndim == 2: it's a mask rather than a bias.
      # expects mask of shape:
      #   [batch, key_tokens]
      # adds singleton query_tokens dimension:
      #   [batch,                    1, key_tokens]
      # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:
      #   [batch,  heads, query_tokens, key_tokens] (e.g. torch sdp attn)
      #   [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)
      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)
          SCREAMING_SNAKE_CASE_													=  (1 - attention_mask.to(hidden_states.dtype)) * -1_0_0_0_0.0
          SCREAMING_SNAKE_CASE_													=  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:
          SCREAMING_SNAKE_CASE_													=  (1 - encoder_attention_mask.to(hidden_states.dtype)) * -1_0_0_0_0.0
          SCREAMING_SNAKE_CASE_													=  encoder_attention_mask.unsqueeze(1)

      # 1. Input
      if self.is_input_continuous:
          SCREAMING_SNAKE_CASE_							,	SCREAMING_SNAKE_CASE_							,	SCREAMING_SNAKE_CASE_							,	SCREAMING_SNAKE_CASE_													=  hidden_states.shape
          SCREAMING_SNAKE_CASE_													=  hidden_states

          SCREAMING_SNAKE_CASE_													=  self.norm(_A)
          if not self.use_linear_projection:
              SCREAMING_SNAKE_CASE_													=  self.proj_in(_A)
              SCREAMING_SNAKE_CASE_													=  hidden_states.shape[1]
              SCREAMING_SNAKE_CASE_													=  hidden_states.permute(0  ,						2  ,						3  ,						1).reshape(_A  ,						height * width  ,						_A)
          else:
              SCREAMING_SNAKE_CASE_													=  hidden_states.shape[1]
              SCREAMING_SNAKE_CASE_													=  hidden_states.permute(0  ,						2  ,						3  ,						1).reshape(_A  ,						height * width  ,						_A)
              SCREAMING_SNAKE_CASE_													=  self.proj_in(_A)
      elif self.is_input_vectorized:
          SCREAMING_SNAKE_CASE_													=  self.latent_image_embedding(_A)
      elif self.is_input_patches:
          SCREAMING_SNAKE_CASE_													=  self.pos_embed(_A)

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

      # 3. Output
      if self.is_input_continuous:
          if not self.use_linear_projection:
              SCREAMING_SNAKE_CASE_													=  hidden_states.reshape(_A  ,						_A  ,						_A  ,						_A).permute(0  ,						3  ,						1  ,						2).contiguous()
              SCREAMING_SNAKE_CASE_													=  self.proj_out(_A)
          else:
              SCREAMING_SNAKE_CASE_													=  self.proj_out(_A)
              SCREAMING_SNAKE_CASE_													=  hidden_states.reshape(_A  ,						_A  ,						_A  ,						_A).permute(0  ,						3  ,						1  ,						2).contiguous()

          SCREAMING_SNAKE_CASE_													=  hidden_states + residual
      elif self.is_input_vectorized:
          SCREAMING_SNAKE_CASE_													=  self.norm_out(_A)
          SCREAMING_SNAKE_CASE_													=  self.out(_A)
          # (batch, self.num_vector_embeds - 1, self.num_latent_pixels)
          SCREAMING_SNAKE_CASE_													=  logits.permute(0  ,						2  ,						1)

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

          # unpatchify
          SCREAMING_SNAKE_CASE_													=  SCREAMING_SNAKE_CASE_													=  int(hidden_states.shape[1] ** 0.5)
          SCREAMING_SNAKE_CASE_													=  hidden_states.reshape(
              shape=(-1, height, width, self.patch_size, self.patch_size, self.out_channels))
          SCREAMING_SNAKE_CASE_													=  torch.einsum('nhwpqc->nchpwq'  ,						_A)
          SCREAMING_SNAKE_CASE_													=  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=_A)


