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`.
class_docstr_tok: ["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","`","."]
code: def __init__(
      self,
      num_attention_heads: int = 32,
      attention_head_dim: int = 64,
      num_layers: int = 20,
      embedding_dim: int = 768,
      num_embeddings=77,
      additional_embeddings=4,
      dropout: float = 0.0,
      time_embed_act_fn: str = "silu",
      norm_in_type: Optional[str] = None,  # layer
      embedding_proj_norm_type: Optional[str] = None,  # layer
      encoder_hid_proj_type: Optional[str] = "linear",  # linear
      added_emb_type: Optional[str] = "prd",  # prd
      time_embed_dim: Optional[int] = None,
      embedding_proj_dim: Optional[int] = None,
      clip_embed_dim: Optional[int] = None,
  ):
      super().__init__()
      self.num_attention_heads = num_attention_heads
      self.attention_head_dim = attention_head_dim
      inner_dim = num_attention_heads * attention_head_dim
      self.additional_embeddings = additional_embeddings
  
      time_embed_dim = time_embed_dim or inner_dim
      embedding_proj_dim = embedding_proj_dim or embedding_dim
      clip_embed_dim = clip_embed_dim or embedding_dim
  
      self.time_proj = Timesteps(inner_dim, True, 0)
      self.time_embedding = TimestepEmbedding(
          inner_dim,
          time_embed_dim,
          out_dim=inner_dim,
          act_fn=time_embed_act_fn,
      )
  
      self.proj_in = nn.Linear(embedding_dim, inner_dim)
  
      if embedding_proj_norm_type is None:
          self.embedding_proj_norm = None
      elif embedding_proj_norm_type == "layer":
          self.embedding_proj_norm = nn.LayerNorm(embedding_proj_dim)
      else:
          raise ValueError(
              f"unsupported embedding_proj_norm_type: {embedding_proj_norm_type}"
          )
  
      self.embedding_proj = nn.Linear(embedding_proj_dim, inner_dim)
  
      if encoder_hid_proj_type is None:
          self.encoder_hidden_states_proj = None
      elif encoder_hid_proj_type == "linear":
          self.encoder_hidden_states_proj = nn.Linear(
              embedding_dim, inner_dim
          )
      else:
          raise ValueError(
              f"unsupported encoder_hid_proj_type: {encoder_hid_proj_type}"
          )
  
      self.positional_embedding = nn.Parameter(
          torch.zeros(
              1, num_embeddings + additional_embeddings, inner_dim
          )
      )
  
      if added_emb_type == "prd":
          self.prd_embedding = nn.Parameter(
              torch.zeros(1, 1, inner_dim)
          )
      elif added_emb_type is None:
          self.prd_embedding = None
      else:
          raise ValueError(
              f"`added_emb_type`: {added_emb_type} is not supported. Make sure to choose one of `'prd'` or `None`."
          )
  
      self.transformer_blocks = nn.ModuleList(
          [
              BasicTransformerBlock(
                  inner_dim,
                  num_attention_heads,
                  attention_head_dim,
                  dropout=dropout,
                  activation_fn="gelu",
                  attention_bias=True,
              )
              for d in range(num_layers)
          ]
      )
  
      if norm_in_type == "layer":
          self.norm_in = nn.LayerNorm(inner_dim)
      elif norm_in_type is None:
          self.norm_in = None
      else:
          raise ValueError(f"Unsupported norm_in_type: {norm_in_type}.")
  
      self.norm_out = nn.LayerNorm(inner_dim)
  
      self.proj_to_clip_embeddings = nn.Linear(
          inner_dim, clip_embed_dim
      )
  
      causal_attention_mask = torch.full(
          [
              num_embeddings + additional_embeddings,
              num_embeddings + additional_embeddings,
          ],
          -10000.0,
      )
      causal_attention_mask.triu_(1)
      causal_attention_mask = causal_attention_mask[None, ...]
      self.register_buffer(
          "causal_attention_mask",
          causal_attention_mask,
          persistent=False,
      )
  
      self.clip_mean = nn.Parameter(torch.zeros(1, clip_embed_dim))
      self.clip_std = nn.Parameter(torch.zeros(1, clip_embed_dim))
  
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 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 1 to create a new PriorTransformer object from inherited base classes: ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin with example: obj = PriorTransformer(num_attention_heads, attention_head_dim, num_layers, embedding_dim, num_embeddings, additional_embeddings, dropout, time_embed_act_fn, norm_in_type, embedding_proj_norm_type, encoder_hid_proj_type, added_emb_type, time_embed_dim, embedding_proj_dim, clip_embed_dim)