Prior Transformer
Last updated
Last updated
A Transformer model for video-like data.
( num_attention_heads: int = 16attention_head_dim: int = 88in_channels: typing.Optional[int] = Noneout_channels: typing.Optional[int] = Nonenum_layers: int = 1dropout: float = 0.0norm_num_groups: int = 32cross_attention_dim: typing.Optional[int] = Noneattention_bias: bool = Falsesample_size: typing.Optional[int] = Noneactivation_fn: str = 'geglu'norm_elementwise_affine: bool = Truedouble_self_attention: bool = True )
Parameters
num_attention_heads (int, optional, defaults to 16) — The number of heads to use for multi-head attention.
attention_head_dim (int, optional, defaults to 88) — The number of channels in each head.
in_channels (int, optional) — The number of channels in the input and output (specify if the input is continuous).
num_layers (int, optional, defaults to 1) — The number of layers of Transformer blocks to use.
dropout (float, optional, defaults to 0.0) — The dropout probability to use.
cross_attention_dim (int, optional) — The number of encoder_hidden_states dimensions to use.
sample_size (int, optional) — The width of the latent images (specify if the input is discrete). This is fixed during training since it is used to learn a number of position embeddings.
activation_fn (str, optional, defaults to "geglu") — Activation function to use in feed-forward.
attention_bias (bool, optional) — Configure if the TransformerBlock attention should contain a bias parameter.
double_self_attention (bool, optional) — Configure if each TransformerBlock should contain two self-attention layers.
A Transformer model for video-like data.
forward
Parameters
hidden_states (torch.LongTensor of shape (batch size, num latent pixels) if discrete, torch.FloatTensor of shape (batch size, channel, height, width) if continuous) — Input hidden_states.
encoder_hidden_states ( torch.LongTensor of shape (batch size, encoder_hidden_states dim), optional) — Conditional embeddings for cross attention layer. If not given, cross-attention defaults to self-attention.
timestep ( torch.long, optional) — Used to indicate denoising step. Optional timestep to be applied as an embedding in AdaLayerNorm.
class_labels ( torch.LongTensor of shape (batch size, num classes), optional) — Used to indicate class labels conditioning. Optional class labels to be applied as an embedding in AdaLayerZeroNorm.
Returns
The TransformerTemporal forward method.
( sample: FloatTensor )
Parameters
sample (torch.FloatTensor of shape (batch_size x num_frames, num_channels, height, width)) — The hidden states output conditioned on encoder_hidden_states input.
The output of TransformerTemporalModel.
( hidden_statesencoder_hidden_states = Nonetimestep = Noneclass_labels = Nonenum_frames = 1cross_attention_kwargs = Nonereturn_dict: bool = True ) → or tuple
return_dict (bool, optional, defaults to True) — Whether or not to return a instead of a plain tuple.
or tuple
If return_dict is True, an is returned, otherwise a tuple where the first element is the sample tensor.