code: 
  
  
  from dataclasses import dataclass
  from typing import List, Optional, Tuple, Union
  
  import numpy as np
  import torch
  
  from ..configuration_utils import ConfigMixin, register_to_config
  from ..utils import BaseOutput, logging, randn_tensor
  from .scheduling_utils import SchedulerMixin
  
  
  SCREAMING_SNAKE_CASE : Dict          =    logging.get_logger(__name__)  # pylint: disable=invalid-name
  
  
  
  
  
  @dataclass
  class 					snake_case__  ( snake_case__       ):
   UpperCAmelCase   :							Tuple         = 42
  
  
  
  
  
  class 					snake_case__  ( snake_case__      ,    snake_case__       ):
   UpperCAmelCase   :							List[str]         = 1
  
  
  
  
  
   @register_to_config
   def __init__(  self				,    UpperCamelCase_ = 40				,    UpperCamelCase_ = 0.002				,    UpperCamelCase_ = 80.0				,    UpperCamelCase_ = 0.5				,    UpperCamelCase_ = 1.0				,    UpperCamelCase_ = 7.0				,    UpperCamelCase_ = True				,    )		->			List[str]:
          """simple docstring"""
  
  
  
  
  
  
          a_       :   List[str]			=	sigma_max
  
          a_       :   Dict			=	np.linspace(0				,    1				,    __snake_case							)
          a_       :   Any			=	self._convert_to_karras(__snake_case							)
          a_       :   List[str]			=	self.sigma_to_t(__snake_case							)
  
          # setable values
          a_       :   str			=	None
          a_       :   List[Any]			=	torch.from_numpy(__snake_case							)
          a_       :   Any			=	torch.from_numpy(__snake_case							)
          a_       :   Optional[int]			=	False
          a_       :   Union[str, Any]			=	False
  
  
  
  
  
   def     A				(  self				,    UpperCamelCase_				,    UpperCamelCase_=None							)		->			Dict:
          """simple docstring"""
  
  
  
  
  
  
          if schedule_timesteps is None:
                 a_       :   Union[str, Any]			=	self.timesteps
  
          a_       :   Dict			=	(schedule_timesteps == timestep).nonzero()
          return indices.item()
  
  
  
  
  
   def     A				(  self				,    UpperCamelCase_				,    UpperCamelCase_							)		->			Optional[int]:
          """simple docstring"""
  
  
  
  
  
  
          if isinstance(__snake_case				,    torch.Tensor							):
                 a_       :   Tuple			=	timestep.to(self.timesteps.device							)
          a_       :   Any			=	self.index_for_timestep(__snake_case							)
          a_       :   Any			=	self.sigmas[step_idx]
  
          a_       :   Dict			=	sample / ((sigma**2 + self.config.sigma_data**2) ** 0.5)
  
          a_       :   Optional[Any]			=	True
          return sample
  
  
  
  
  
   def     A				(  self				,    UpperCamelCase_							)		->			Dict:
          """simple docstring"""
  
  
  
  
  
  
          if not isinstance(__snake_case				,    np.ndarray							):
                 a_       :   List[str]			=	np.array(__snake_case				,    dtype=np.floataa							)
  
          a_       :   Any			=	1000 * 0.25 * np.log(sigmas + 1e-44							)
  
          return timesteps
  
  
  
  
  
   def     A				(  self				,    UpperCamelCase_ = None				,    UpperCamelCase_ = None				,    UpperCamelCase_ = None				,    )		->			Tuple:
          """simple docstring"""
  
  
  
  
  
  
          if num_inference_steps is None and timesteps is None:
                 raise ValueError("""Exactly one of `num_inference_steps` or `timesteps` must be supplied."""							)
  
          if num_inference_steps is not None and timesteps is not None:
                 raise ValueError("""Can only pass one of `num_inference_steps` or `timesteps`."""							)
  
          # Follow DDPMScheduler custom timesteps logic
          if timesteps is not None:
                 for i in range(1				,    len(__snake_case							)							):
                        if timesteps[i] >= timesteps[i - 1]:
                               raise ValueError("""`timesteps` must be in descending order."""							)
  
                 if timesteps[0] >= self.config.num_train_timesteps:
                        raise ValueError(
                            f"""`timesteps` must start before `self.config.train_timesteps`:"""
                            f""" {self.config.num_train_timesteps}."""							)
  
                 a_       :   Optional[Any]			=	np.array(__snake_case				,    dtype=np.intaa							)
                 a_       :   List[Any]			=	True
          else:
                 if num_inference_steps > self.config.num_train_timesteps:
                        raise ValueError(
                            f"""`num_inference_steps`: {num_inference_steps} cannot be larger than `self.config.train_timesteps`:"""
                            f""" {self.config.num_train_timesteps} as the unet model trained with this scheduler can only handle"""
                            f""" maximal {self.config.num_train_timesteps} timesteps."""							)
  
                 a_       :   Dict			=	num_inference_steps
  
                 a_       :   Any			=	self.config.num_train_timesteps // self.num_inference_steps
                 a_       :   Optional[Any]			=	(np.arange(0				,    __snake_case							) * step_ratio).round()[::-1].copy().astype(np.intaa							)
                 a_       :   int			=	False
  
          # Map timesteps to Karras sigmas directly for multistep sampling
          # See https://github.com/openai/consistency_models/blob/main/cm/karras_diffusion.py#L675
          a_       :   List[Any]			=	self.config.num_train_timesteps
          a_       :   Union[str, Any]			=	timesteps[::-1].copy()
          a_       :   List[str]			=	ramp / (num_train_timesteps - 1)
          a_       :   str			=	self._convert_to_karras(__snake_case							)
          a_       :   Any			=	self.sigma_to_t(__snake_case							)
  
          a_       :   Optional[int]			=	np.concatenate([sigmas, [self.sigma_min]]							).astype(np.floataa							)
          a_       :   Tuple			=	torch.from_numpy(__snake_case							).to(device=__snake_case							)
  
          if str(__snake_case							).startswith("""mps"""							):
                 # mps does not support float64
                 a_       :   List[Any]			=	torch.from_numpy(__snake_case							).to(__snake_case				,    dtype=torch.floataa							)
          else:
                 a_       :   List[Any]			=	torch.from_numpy(__snake_case							).to(device=__snake_case							)
  
  
  
  
  
   def     A				(  self				,    UpperCamelCase_							)		->			Optional[Any]:
          """simple docstring"""
  
  
  
  
  
  
          a_       :   str			=	self.config.sigma_min
          a_       :   int			=	self.config.sigma_max
  
          a_       :   Union[str, Any]			=	self.config.rho
          a_       :   Optional[int]			=	sigma_min ** (1 / rho)
          a_       :   Dict			=	sigma_max ** (1 / rho)
          a_       :   Dict			=	(max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
          return sigmas
  
  
  
  
  
   def     A				(  self				,    UpperCamelCase_							)		->			str:
          """simple docstring"""
  
  
  
  
  
  
          a_       :   List[Any]			=	self.config.sigma_data
  
          a_       :   str			=	sigma_data**2 / (sigma**2 + sigma_data**2)
          a_       :   Optional[Any]			=	sigma * sigma_data / (sigma**2 + sigma_data**2) ** 0.5
          return c_skip, c_out
  
  
  
  
  
   def     A				(  self				,    UpperCamelCase_							)		->			List[Any]:
          """simple docstring"""
  
  
  
  
  
  
          a_       :   Optional[int]			=	self.config.sigma_min
          a_       :   List[Any]			=	self.config.sigma_data
  
          a_       :   Optional[int]			=	sigma_data**2 / ((sigma - sigma_min) ** 2 + sigma_data**2)
          a_       :   List[Any]			=	(sigma - sigma_min) * sigma_data / (sigma**2 + sigma_data**2) ** 0.5
          return c_skip, c_out
  
  
  
  
  
   def     A				(  self				,    UpperCamelCase_				,    UpperCamelCase_				,    UpperCamelCase_				,    UpperCamelCase_ = None				,    UpperCamelCase_ = True				,    )		->			Dict:
          """simple docstring"""
  
  
  
  
  
  
          if (
              isinstance(__snake_case				,    __snake_case							)
              or isinstance(__snake_case				,    torch.IntTensor							)
              or isinstance(__snake_case				,    torch.LongTensor							)
          ):
                 raise ValueError(
                     (
                         """Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"""
                         f""" `{self.__class__}.step()` is not supported. Make sure to pass"""
                         """ one of the `scheduler.timesteps` as a timestep."""
                     )				,    )
  
          if not self.is_scale_input_called:
                 logger.warning(
                     """The `scale_model_input` function should be called before `step` to ensure correct denoising. """
                     """See `StableDiffusionPipeline` for a usage example."""							)
  
          if isinstance(__snake_case				,    torch.Tensor							):
                 a_       :   Tuple			=	timestep.to(self.timesteps.device							)
  
          a_       :   List[str]			=	self.config.sigma_min
          a_       :   Union[str, Any]			=	self.config.sigma_max
  
          a_       :   Optional[Any]			=	self.index_for_timestep(__snake_case							)
  
          # sigma_next corresponds to next_t in original implementation
          a_       :   Optional[Any]			=	self.sigmas[step_index]
          if step_index + 1 < self.config.num_train_timesteps:
                 a_       :   int			=	self.sigmas[step_index + 1]
          else:
                 # Set sigma_next to sigma_min
                 a_       :   List[Any]			=	self.sigmas[-1]
  
          # Get scalings for boundary conditions
          a_		,				a_       :   Tuple			=	self.get_scalings_for_boundary_condition(__snake_case							)
  
          # 1. Denoise model output using boundary conditions
          a_       :   Any			=	c_out * model_output + c_skip * sample
          if self.config.clip_denoised:
                 a_       :   Optional[int]			=	denoised.clamp(-1				,    1							)
  
          # 2. Sample z ~ N(0, s_noise^2 * I)
          # Noise is not used for onestep sampling.
          if len(self.timesteps							) > 1:
                 a_       :   Dict			=	randn_tensor(
                     model_output.shape				,    dtype=model_output.dtype				,    device=model_output.device				,    generator=__snake_case							)
          else:
                 a_       :   Tuple			=	torch.zeros_like(__snake_case							)
          a_       :   Optional[int]			=	noise * self.config.s_noise
  
          a_       :   Optional[int]			=	sigma_next.clamp(min=__snake_case				,    max=__snake_case							)
  
          # 3. Return noisy sample
          # tau = sigma_hat, eps = sigma_min
          a_       :   Any			=	denoised + z * (sigma_hat**2 - sigma_min**2) ** 0.5
  
          if not return_dict:
                 return (prev_sample,)
  
          return CMStochasticIterativeSchedulerOutput(prev_sample=__snake_case							)
  
  
  
  
  
   def     A				(  self				,    UpperCamelCase_				,    UpperCamelCase_				,    UpperCamelCase_				,    )		->			List[Any]:
          """simple docstring"""
  
  
  
  
  
  
          a_       :   List[str]			=	self.sigmas.to(device=original_samples.device				,    dtype=original_samples.dtype							)
          if original_samples.device.type == "mps" and torch.is_floating_point(__snake_case							):
                 # mps does not support float64
                 a_       :   List[str]			=	self.timesteps.to(original_samples.device				,    dtype=torch.floataa							)
                 a_       :   List[str]			=	timesteps.to(original_samples.device				,    dtype=torch.floataa							)
          else:
                 a_       :   Optional[int]			=	self.timesteps.to(original_samples.device							)
                 a_       :   Dict			=	timesteps.to(original_samples.device							)
  
          a_       :   Union[str, Any]			=	[(schedule_timesteps == t).nonzero().item() for t in timesteps]
  
          a_       :   Tuple			=	sigmas[step_indices].flatten()
          while len(sigma.shape							) < len(original_samples.shape							):
                 a_       :   Union[str, Any]			=	sigma.unsqueeze(-1							)
  
          a_       :   Optional[int]			=	original_samples + noise * sigma
          return noisy_samples
  
  
  
  
  
   def __len__(  self							)		->			Union[str, Any]:
          """simple docstring"""
  
  
  
  
  
  
          return self.config.num_train_timesteps
  
  
  
style_context: 
  
  
  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__  ( __A       ):
  	UpperCAmelCase   :							torch.FloatTensor
  
  
  
  
  
  class 					snake_case__  ( __A      ,    __A       ):
  
  
  
  
  
  	@register_to_config
  	def __init__(  self				,    UpperCamelCase_ = 16				,    UpperCamelCase_ = 88				,    UpperCamelCase_ = None				,    UpperCamelCase_ = None				,    UpperCamelCase_ = 1				,    UpperCamelCase_ = 0.0				,    UpperCamelCase_ = 32				,    UpperCamelCase_ = None				,    UpperCamelCase_ = False				,    UpperCamelCase_ = None				,    UpperCamelCase_ = None				,    UpperCamelCase_ = None				,    UpperCamelCase_ = "geglu"				,    UpperCamelCase_ = None				,    UpperCamelCase_ = False				,    UpperCamelCase_ = False				,    UpperCamelCase_ = False				,    UpperCamelCase_ = "layer_norm"				,    UpperCamelCase_ = True				,    )		->			Union[str, Any]:
  								"""simple docstring"""
  
  
  
  
  
  
  								super().__init__()
  								a_       :   Dict			=	use_linear_projection
  								a_       :   Tuple			=	num_attention_heads
  								a_       :   List[Any]			=	attention_head_dim
  								a_       :   Dict			=	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
  								a_       :   List[Any]			=	(in_channels is not None) and (patch_size is None)
  								a_       :   Dict			=	num_vector_embeds is not None
  								a_       :   Union[str, 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:
  															a_       :   int			=	(
  															    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"""				,    UpperCamelCase_				,    standard_warn=UpperCamelCase_							)
  															a_       :   List[Any]			=	"""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:
  															a_       :   List[Any]			=	in_channels
  
  															a_       :   List[str]			=	torch.nn.GroupNorm(num_groups=UpperCamelCase_				,    num_channels=UpperCamelCase_				,    eps=1e-6				,    affine=UpperCamelCase_							)
  															if use_linear_projection:
  																						a_       :   int			=	nn.Linear(UpperCamelCase_				,    UpperCamelCase_							)
  															else:
  																						a_       :   Optional[Any]			=	nn.Convad(UpperCamelCase_				,    UpperCamelCase_				,    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"
  
  															a_       :   Optional[Any]			=	sample_size
  															a_       :   str			=	sample_size
  															a_       :   str			=	num_vector_embeds
  															a_       :   str			=	self.height * self.width
  
  															a_       :   int			=	ImagePositionalEmbeddings(
  															    num_embed=UpperCamelCase_				,    embed_dim=UpperCamelCase_				,    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"
  
  															a_       :   Dict			=	sample_size
  															a_       :   List[Any]			=	sample_size
  
  															a_       :   Tuple			=	patch_size
  															a_       :   Optional[Any]			=	PatchEmbed(
  															    height=UpperCamelCase_				,    width=UpperCamelCase_				,    patch_size=UpperCamelCase_				,    in_channels=UpperCamelCase_				,    embed_dim=UpperCamelCase_				,    )
  
  								# 3. Define transformers blocks
  								a_       :   List[Any]			=	nn.ModuleList(
  								    [
  								        BasicTransformerBlock(
  								            UpperCamelCase_				,    UpperCamelCase_				,    UpperCamelCase_				,    dropout=UpperCamelCase_				,    cross_attention_dim=UpperCamelCase_				,    activation_fn=UpperCamelCase_				,    num_embeds_ada_norm=UpperCamelCase_				,    attention_bias=UpperCamelCase_				,    only_cross_attention=UpperCamelCase_				,    upcast_attention=UpperCamelCase_				,    norm_type=UpperCamelCase_				,    norm_elementwise_affine=UpperCamelCase_				,    )
  								        for d in range(UpperCamelCase_							)
  								    ]							)
  
  								# 4. Define output layers
  								a_       :   str			=	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:
  																						a_       :   Optional[int]			=	nn.Linear(UpperCamelCase_				,    UpperCamelCase_							)
  															else:
  																						a_       :   Any			=	nn.Convad(UpperCamelCase_				,    UpperCamelCase_				,    kernel_size=1				,    stride=1				,    padding=0							)
  								elif self.is_input_vectorized:
  															a_       :   Optional[int]			=	nn.LayerNorm(UpperCamelCase_							)
  															a_       :   str			=	nn.Linear(UpperCamelCase_				,    self.num_vector_embeds - 1							)
  								elif self.is_input_patches:
  															a_       :   int			=	nn.LayerNorm(UpperCamelCase_				,    elementwise_affine=UpperCamelCase_				,    eps=1e-6							)
  															a_       :   str			=	nn.Linear(UpperCamelCase_				,    2 * inner_dim							)
  															a_       :   Optional[int]			=	nn.Linear(UpperCamelCase_				,    patch_size * patch_size * self.out_channels							)
  
  
  
  
  
  	def     A				(  self				,    UpperCamelCase_				,    UpperCamelCase_ = None				,    UpperCamelCase_ = None				,    UpperCamelCase_ = None				,    UpperCamelCase_ = None				,    UpperCamelCase_ = None				,    UpperCamelCase_ = None				,    UpperCamelCase_ = True				,    )		->			str:
  								"""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)
  															a_       :   str			=	(1 - attention_mask.to(hidden_states.dtype							)) * -1_0000.0
  															a_       :   Union[str, 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:
  															a_       :   Any			=	(1 - encoder_attention_mask.to(hidden_states.dtype							)) * -1_0000.0
  															a_       :   List[str]			=	encoder_attention_mask.unsqueeze(1							)
  
  								# 1. Input
  								if self.is_input_continuous:
  															a_		,				a_		,				a_		,				a_       :   str			=	hidden_states.shape
  															a_       :   Optional[Any]			=	hidden_states
  
  															a_       :   Optional[int]			=	self.norm(UpperCamelCase_							)
  															if not self.use_linear_projection:
  																						a_       :   int			=	self.proj_in(UpperCamelCase_							)
  																						a_       :   str			=	hidden_states.shape[1]
  																						a_       :   List[str]			=	hidden_states.permute(0				,    2				,    3				,    1							).reshape(UpperCamelCase_				,    height * width				,    UpperCamelCase_							)
  															else:
  																						a_       :   int			=	hidden_states.shape[1]
  																						a_       :   Optional[int]			=	hidden_states.permute(0				,    2				,    3				,    1							).reshape(UpperCamelCase_				,    height * width				,    UpperCamelCase_							)
  																						a_       :   Tuple			=	self.proj_in(UpperCamelCase_							)
  								elif self.is_input_vectorized:
  															a_       :   List[str]			=	self.latent_image_embedding(UpperCamelCase_							)
  								elif self.is_input_patches:
  															a_       :   Optional[int]			=	self.pos_embed(UpperCamelCase_							)
  
  								# 2. Blocks
  								for block in self.transformer_blocks:
  															a_       :   int			=	block(
  															    UpperCamelCase_				,    attention_mask=UpperCamelCase_				,    encoder_hidden_states=UpperCamelCase_				,    encoder_attention_mask=UpperCamelCase_				,    timestep=UpperCamelCase_				,    cross_attention_kwargs=UpperCamelCase_				,    class_labels=UpperCamelCase_				,    )
  
  								# 3. Output
  								if self.is_input_continuous:
  															if not self.use_linear_projection:
  																						a_       :   Dict			=	hidden_states.reshape(UpperCamelCase_				,    UpperCamelCase_				,    UpperCamelCase_				,    UpperCamelCase_							).permute(0				,    3				,    1				,    2							).contiguous()
  																						a_       :   List[Any]			=	self.proj_out(UpperCamelCase_							)
  															else:
  																						a_       :   Optional[int]			=	self.proj_out(UpperCamelCase_							)
  																						a_       :   int			=	hidden_states.reshape(UpperCamelCase_				,    UpperCamelCase_				,    UpperCamelCase_				,    UpperCamelCase_							).permute(0				,    3				,    1				,    2							).contiguous()
  
  															a_       :   int			=	hidden_states + residual
  								elif self.is_input_vectorized:
  															a_       :   str			=	self.norm_out(UpperCamelCase_							)
  															a_       :   str			=	self.out(UpperCamelCase_							)
  															# (batch, self.num_vector_embeds - 1, self.num_latent_pixels)
  															a_       :   int			=	logits.permute(0				,    2				,    1							)
  
  															# log(p(x_0))
  															a_       :   int			=	F.log_softmax(logits.double()				,    dim=1							).float()
  								elif self.is_input_patches:
  															# TODO: cleanup!
  															a_       :   Union[str, Any]			=	self.transformer_blocks[0].norma.emb(
  															    UpperCamelCase_				,    UpperCamelCase_				,    hidden_dtype=hidden_states.dtype							)
  															a_		,				a_       :   int			=	self.proj_out_a(F.silu(UpperCamelCase_							)							).chunk(2				,    dim=1							)
  															a_       :   Union[str, Any]			=	self.norm_out(UpperCamelCase_							) * (1 + scale[:, None]) + shift[:, None]
  															a_       :   Union[str, Any]			=	self.proj_out_a(UpperCamelCase_							)
  
  															# unpatchify
  															a_       :   Any			=	int(hidden_states.shape[1] ** 0.5							)
  															a_       :   str			=	hidden_states.reshape(
  															    shape=(-1, height, width, self.patch_size, self.patch_size, self.out_channels)							)
  															a_       :   Union[str, Any]			=	torch.einsum("""nhwpqc->nchpwq"""				,    UpperCamelCase_							)
  															a_       :   List[Any]			=	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=UpperCamelCase_							)
  
  
  