class_docstr: A Prior Transformer model.
  
  Parameters:
      num_attention_heads (`int`, *optional*, defaults to 32): The number of heads to use for multi-head attention.
      attention_head_dim (`int`, *optional*, defaults to 64): The number of channels in each head.
      num_layers (`int`, *optional*, defaults to 20): The number of layers of Transformer blocks to use.
      embedding_dim (`int`, *optional*, defaults to 768): The dimension of the model input `hidden_states`
      num_embeddings (`int`, *optional*, defaults to 77):
          The number of embeddings of the model input `hidden_states`
      additional_embeddings (`int`, *optional*, defaults to 4): The number of additional tokens appended to the
          projected `hidden_states`. The actual length of the used `hidden_states` is `num_embeddings +
          additional_embeddings`.
      dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
      time_embed_act_fn (`str`, *optional*, defaults to 'silu'):
          The activation function to use to create timestep embeddings.
      norm_in_type (`str`, *optional*, defaults to None): The normalization layer to apply on hidden states before
          passing to Transformer blocks. Set it to `None` if normalization is not needed.
      embedding_proj_norm_type (`str`, *optional*, defaults to None):
          The normalization layer to apply on the input `proj_embedding`. Set it to `None` if normalization is not
          needed.
      encoder_hid_proj_type (`str`, *optional*, defaults to `linear`):
          The projection layer to apply on the input `encoder_hidden_states`. Set it to `None` if
          `encoder_hidden_states` is `None`.
      added_emb_type (`str`, *optional*, defaults to `prd`): Additional embeddings to condition the model.
          Choose from `prd` or `None`. if choose `prd`, it will prepend a token indicating the (quantized) dot
          product between the text embedding and image embedding as proposed in the unclip paper
          https://arxiv.org/abs/2204.06125 If it is `None`, no additional embeddings will be prepended.
      time_embed_dim (`int, *optional*, defaults to None): The dimension of timestep embeddings.
          If None, will be set to `num_attention_heads * attention_head_dim`
      embedding_proj_dim (`int`, *optional*, default to None):
          The dimension of `proj_embedding`. If None, will be set to `embedding_dim`.
      clip_embed_dim (`int`, *optional*, default to None):
          The dimension of the output. If None, will be set to `embedding_dim`.
docstr: The [`PriorTransformer`] forward method.
  
  Args:
      hidden_states (`torch.FloatTensor` of shape `(batch_size, embedding_dim)`):
          The currently predicted image embeddings.
      timestep (`torch.LongTensor`):
          Current denoising step.
      proj_embedding (`torch.FloatTensor` of shape `(batch_size, embedding_dim)`):
          Projected embedding vector the denoising process is conditioned on.
      encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, num_embeddings, embedding_dim)`):
          Hidden states of the text embeddings the denoising process is conditioned on.
      attention_mask (`torch.BoolTensor` of shape `(batch_size, num_embeddings)`):
          Text mask for the text embeddings.
      return_dict (`bool`, *optional*, defaults to `True`):
          Whether or not to return a [`~models.prior_transformer.PriorTransformerOutput`] instead of a plain
          tuple.
  
  Returns:
      [`~models.prior_transformer.PriorTransformerOutput`] or `tuple`:
          If return_dict is True, a [`~models.prior_transformer.PriorTransformerOutput`] is returned, otherwise a
          tuple is returned where the first element is the sample tensor.
docstr_tok: ["The","[","`","PriorTransformer","`","]","forward","method",".","Args",":","hidden_states","(","`","torch.FloatTensor","`","of","shape","`","(","batch_size",",","embedding_dim",")","`",")",":","The","currently","predicted","image","embeddings",".","timestep","(","`","torch.LongTensor","`",")",":","Current","denoising","step",".","proj_embedding","(","`","torch.FloatTensor","`","of","shape","`","(","batch_size",",","embedding_dim",")","`",")",":","Projected","embedding","vector","the","denoising","process","is","conditioned","on",".","encoder_hidden_states","(","`","torch.FloatTensor","`","of","shape","`","(","batch_size",",","num_embeddings",",","embedding_dim",")","`",")",":","Hidden","states","of","the","text","embeddings","the","denoising","process","is","conditioned","on",".","attention_mask","(","`","torch.BoolTensor","`","of","shape","`","(","batch_size",",","num_embeddings",")","`",")",":","Text","mask","for","the","text","embeddings",".","return_dict","(","`","bool","`",",","*","optional","*",",","defaults","to","`","True","`",")",":","Whether","or","not","to","return","a","[","`","~models.prior_transformer.PriorTransformerOutput","`","]","instead","of","a","plain","tuple",".","Returns",":","[","`","~models.prior_transformer.PriorTransformerOutput","`","]","or","`","tuple","`",":","If","return_dict","is","True",",","a","[","`","~models.prior_transformer.PriorTransformerOutput","`","]","is","returned",",","otherwise","a","tuple","is","returned","where","the","first","element","is","the","sample","tensor","."]
code: def forward(
      self,
      hidden_states,
      timestep: Union[torch.Tensor, float, int],
      proj_embedding: torch.FloatTensor,
      encoder_hidden_states: Optional[torch.FloatTensor] = None,
      attention_mask: Optional[torch.BoolTensor] = None,
      return_dict: bool = True,
  ):
      """
      The [`PriorTransformer`] forward method.
  
      Args:
          hidden_states (`torch.FloatTensor` of shape `(batch_size, embedding_dim)`):
              The currently predicted image embeddings.
          timestep (`torch.LongTensor`):
              Current denoising step.
          proj_embedding (`torch.FloatTensor` of shape `(batch_size, embedding_dim)`):
              Projected embedding vector the denoising process is conditioned on.
          encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, num_embeddings, embedding_dim)`):
              Hidden states of the text embeddings the denoising process is conditioned on.
          attention_mask (`torch.BoolTensor` of shape `(batch_size, num_embeddings)`):
              Text mask for the text embeddings.
          return_dict (`bool`, *optional*, defaults to `True`):
              Whether or not to return a [`~models.prior_transformer.PriorTransformerOutput`] instead of a plain
              tuple.
  
      Returns:
          [`~models.prior_transformer.PriorTransformerOutput`] or `tuple`:
              If return_dict is True, a [`~models.prior_transformer.PriorTransformerOutput`] is returned, otherwise a
              tuple is returned where the first element is the sample tensor.
      """
      batch_size = hidden_states.shape[0]
  
      timesteps = timestep
      if not torch.is_tensor(timesteps):
          timesteps = torch.tensor(
              [timesteps], dtype=torch.long, device=hidden_states.device
          )
      elif torch.is_tensor(timesteps) and len(timesteps.shape) == 0:
          timesteps = timesteps[None].to(hidden_states.device)
  
      # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
      timesteps = timesteps * torch.ones(
          batch_size, dtype=timesteps.dtype, device=timesteps.device
      )
  
      timesteps_projected = self.time_proj(timesteps)
  
      # timesteps does not contain any weights and will always return f32 tensors
      # but time_embedding might be fp16, so we need to cast here.
      timesteps_projected = timesteps_projected.to(dtype=self.dtype)
      time_embeddings = self.time_embedding(timesteps_projected)
  
      if self.embedding_proj_norm is not None:
          proj_embedding = self.embedding_proj_norm(proj_embedding)
  
      proj_embeddings = self.embedding_proj(proj_embedding)
      if (
          self.encoder_hidden_states_proj is not None
          and encoder_hidden_states is not None
      ):
          encoder_hidden_states = self.encoder_hidden_states_proj(
              encoder_hidden_states
          )
      elif (
          self.encoder_hidden_states_proj is not None
          and encoder_hidden_states is None
      ):
          raise ValueError(
              "`encoder_hidden_states_proj` requires `encoder_hidden_states` to be set"
          )
  
      hidden_states = self.proj_in(hidden_states)
  
      positional_embeddings = self.positional_embedding.to(
          hidden_states.dtype
      )
  
      additional_embeds = []
      additional_embeddings_len = 0
  
      if encoder_hidden_states is not None:
          additional_embeds.append(encoder_hidden_states)
          additional_embeddings_len += encoder_hidden_states.shape[1]
  
      if len(proj_embeddings.shape) == 2:
          proj_embeddings = proj_embeddings[:, None, :]
  
      if len(hidden_states.shape) == 2:
          hidden_states = hidden_states[:, None, :]
  
      additional_embeds = additional_embeds + [
          proj_embeddings,
          time_embeddings[:, None, :],
          hidden_states,
      ]
  
      if self.prd_embedding is not None:
          prd_embedding = self.prd_embedding.to(
              hidden_states.dtype
          ).expand(batch_size, -1, -1)
          additional_embeds.append(prd_embedding)
  
      hidden_states = torch.cat(
          additional_embeds,
          dim=1,
      )
  
      # Allow positional_embedding to not include the `addtional_embeddings` and instead pad it with zeros for these additional tokens
      additional_embeddings_len = (
          additional_embeddings_len + proj_embeddings.shape[1] + 1
      )
      if positional_embeddings.shape[1] < hidden_states.shape[1]:
          positional_embeddings = F.pad(
              positional_embeddings,
              (
                  0,
                  0,
                  additional_embeddings_len,
                  self.prd_embedding.shape[1]
                  if self.prd_embedding is not None
                  else 0,
              ),
              value=0.0,
          )
  
      hidden_states = hidden_states + positional_embeddings
  
      if attention_mask is not None:
          attention_mask = (
              1 - attention_mask.to(hidden_states.dtype)
          ) * -10000.0
          attention_mask = F.pad(
              attention_mask, (0, self.additional_embeddings), value=0.0
          )
          attention_mask = (
              attention_mask[:, None, :] + self.causal_attention_mask
          ).to(hidden_states.dtype)
          attention_mask = attention_mask.repeat_interleave(
              self.config.num_attention_heads, dim=0
          )
  
      if self.norm_in is not None:
          hidden_states = self.norm_in(hidden_states)
  
      for block in self.transformer_blocks:
          hidden_states = block(
              hidden_states, attention_mask=attention_mask
          )
  
      hidden_states = self.norm_out(hidden_states)
  
      if self.prd_embedding is not None:
          hidden_states = hidden_states[:, -1]
      else:
          hidden_states = hidden_states[:, additional_embeddings_len:]
  
      predicted_image_embedding = self.proj_to_clip_embeddings(
          hidden_states
      )
  
      if not return_dict:
          return (predicted_image_embedding,)
  
      return PriorTransformerOutput(
          predicted_image_embedding=predicted_image_embedding
      )
  
imports: from dataclasses import dataclass
  from typing import Dict, Optional, Union
  import torch
  import torch.nn.functional
  from torch import nn
  from ..configuration_utils import ConfigMixin, register_to_config
  from ..loaders import PeftAdapterMixin, UNet2DConditionLoadersMixin
  from ..utils import BaseOutput
  from .attention import BasicTransformerBlock
  from .attention_processor import ADDED_KV_ATTENTION_PROCESSORS, CROSS_ATTENTION_PROCESSORS, AttentionProcessor, AttnAddedKVProcessor, AttnProcessor
  from .embeddings import TimestepEmbedding, Timesteps
  from .modeling_utils import ModelMixin
label_desc: Use image node_id 5 for calling the PriorTransformer obj's underlying member method code with example usage: obj.forward(hidden_states, timestep, proj_embedding, encoder_hidden_states, attention_mask, return_dict) and returns: PriorTransformerOutput, predicted_image_embedding