575 lines
22 KiB
Python
575 lines
22 KiB
Python
# -*- coding: utf-8 -*-
|
|
|
|
# Copyright 2024 The CogVideoX team, Tsinghua University & ZhipuAI and The HuggingFace Team.
|
|
# All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
from typing import Optional, Tuple
|
|
|
|
import torch
|
|
from torch import nn
|
|
import torch.nn.functional as F
|
|
|
|
from .utils import get_activation, get_timestep_embedding, get_3d_sincos_pos_embed, apply_rotary_emb
|
|
from .utils import GELU, GEGLU, ApproximateGELU, SwiGLU
|
|
|
|
|
|
class TimestepEmbedding(nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
time_embed_dim: int,
|
|
act_fn: str = "silu",
|
|
out_dim: int = None,
|
|
post_act_fn: Optional[str] = None,
|
|
cond_proj_dim=None,
|
|
sample_proj_bias=True,
|
|
):
|
|
super().__init__()
|
|
|
|
self.linear_1 = nn.Linear(in_channels, time_embed_dim, sample_proj_bias)
|
|
|
|
if cond_proj_dim is not None:
|
|
self.cond_proj = nn.Linear(cond_proj_dim, in_channels, bias=False)
|
|
else:
|
|
self.cond_proj = None
|
|
|
|
self.act = get_activation(act_fn)
|
|
|
|
if out_dim is not None:
|
|
time_embed_dim_out = out_dim
|
|
else:
|
|
time_embed_dim_out = time_embed_dim
|
|
self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim_out, sample_proj_bias)
|
|
|
|
if post_act_fn is None:
|
|
self.post_act = None
|
|
else:
|
|
self.post_act = get_activation(post_act_fn)
|
|
|
|
def forward(self, sample, condition=None):
|
|
if condition is not None:
|
|
sample = sample + self.cond_proj(condition)
|
|
sample = self.linear_1(sample)
|
|
|
|
if self.act is not None:
|
|
sample = self.act(sample)
|
|
|
|
sample = self.linear_2(sample)
|
|
|
|
if self.post_act is not None:
|
|
sample = self.post_act(sample)
|
|
return sample
|
|
|
|
|
|
|
|
class Timesteps(nn.Module):
|
|
def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float, scale: int = 1):
|
|
super().__init__()
|
|
self.num_channels = num_channels
|
|
self.flip_sin_to_cos = flip_sin_to_cos
|
|
self.downscale_freq_shift = downscale_freq_shift
|
|
self.scale = scale
|
|
|
|
def forward(self, timesteps):
|
|
t_emb = get_timestep_embedding(
|
|
timesteps,
|
|
self.num_channels,
|
|
flip_sin_to_cos=self.flip_sin_to_cos,
|
|
downscale_freq_shift=self.downscale_freq_shift,
|
|
scale=self.scale,
|
|
)
|
|
return t_emb
|
|
|
|
|
|
class CogVideoXLayerNormZero(nn.Module):
|
|
def __init__(
|
|
self,
|
|
conditioning_dim: int,
|
|
embedding_dim: int,
|
|
elementwise_affine: bool = True,
|
|
eps: float = 1e-5,
|
|
bias: bool = True,
|
|
) -> None:
|
|
super().__init__()
|
|
|
|
self.silu = nn.SiLU()
|
|
self.linear = nn.Linear(conditioning_dim, 6 * embedding_dim, bias=bias)
|
|
self.norm = nn.LayerNorm(embedding_dim, eps=eps, elementwise_affine=elementwise_affine)
|
|
|
|
def forward(
|
|
self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor
|
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
shift, scale, gate, enc_shift, enc_scale, enc_gate = self.linear(self.silu(temb)).chunk(6, dim=1)
|
|
hidden_states = self.norm(hidden_states) * (1 + scale)[:, None, :] + shift[:, None, :]
|
|
encoder_hidden_states = self.norm(encoder_hidden_states) * (1 + enc_scale)[:, None, :] + enc_shift[:, None, :]
|
|
return hidden_states, encoder_hidden_states, gate[:, None, :], enc_gate[:, None, :]
|
|
|
|
|
|
class AdaLayerNorm(nn.Module):
|
|
r"""
|
|
Norm layer modified to incorporate timestep embeddings.
|
|
|
|
Parameters:
|
|
embedding_dim (`int`): The size of each embedding vector.
|
|
num_embeddings (`int`, *optional*): The size of the embeddings dictionary.
|
|
output_dim (`int`, *optional*):
|
|
norm_elementwise_affine (`bool`, defaults to `False):
|
|
norm_eps (`bool`, defaults to `False`):
|
|
chunk_dim (`int`, defaults to `0`):
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
embedding_dim: int,
|
|
num_embeddings: Optional[int] = None,
|
|
output_dim: Optional[int] = None,
|
|
norm_elementwise_affine: bool = False,
|
|
norm_eps: float = 1e-5,
|
|
chunk_dim: int = 0,
|
|
):
|
|
super().__init__()
|
|
|
|
self.chunk_dim = chunk_dim
|
|
output_dim = output_dim or embedding_dim * 2
|
|
|
|
if num_embeddings is not None:
|
|
self.emb = nn.Embedding(num_embeddings, embedding_dim)
|
|
else:
|
|
self.emb = None
|
|
|
|
self.silu = nn.SiLU()
|
|
self.linear = nn.Linear(embedding_dim, output_dim)
|
|
self.norm = nn.LayerNorm(output_dim // 2, norm_eps, norm_elementwise_affine)
|
|
|
|
def forward(
|
|
self, x: torch.Tensor, timestep: Optional[torch.Tensor] = None, temb: Optional[torch.Tensor] = None
|
|
) -> torch.Tensor:
|
|
if self.emb is not None:
|
|
temb = self.emb(timestep)
|
|
|
|
temb = self.linear(self.silu(temb))
|
|
|
|
if self.chunk_dim == 1:
|
|
# This is a bit weird why we have the order of "shift, scale" here and "scale, shift" in the
|
|
# other if-branch. This branch is specific to CogVideoX for now.
|
|
shift, scale = temb.chunk(2, dim=1)
|
|
shift = shift[:, None, :]
|
|
scale = scale[:, None, :]
|
|
else:
|
|
scale, shift = temb.chunk(2, dim=0)
|
|
|
|
x = self.norm(x) * (1 + scale) + shift
|
|
return x
|
|
|
|
|
|
class CogVideoXPatchEmbed(nn.Module):
|
|
def __init__(
|
|
self,
|
|
patch_size: int = 2,
|
|
patch_size_t: Optional[int] = None,
|
|
in_channels: int = 16,
|
|
embed_dim: int = 1920,
|
|
text_embed_dim: int = 4096,
|
|
bias: bool = True,
|
|
sample_width: int = 90,
|
|
sample_height: int = 60,
|
|
sample_frames: int = 49,
|
|
temporal_compression_ratio: int = 4,
|
|
max_text_seq_length: int = 226,
|
|
spatial_interpolation_scale: float = 1.875,
|
|
temporal_interpolation_scale: float = 1.0,
|
|
use_positional_embeddings: bool = True,
|
|
use_learned_positional_embeddings: bool = True,
|
|
) -> None:
|
|
super().__init__()
|
|
|
|
self.patch_size = patch_size
|
|
self.patch_size_t = patch_size_t
|
|
self.embed_dim = embed_dim
|
|
self.sample_height = sample_height
|
|
self.sample_width = sample_width
|
|
self.sample_frames = sample_frames
|
|
self.temporal_compression_ratio = temporal_compression_ratio
|
|
self.max_text_seq_length = max_text_seq_length
|
|
self.spatial_interpolation_scale = spatial_interpolation_scale
|
|
self.temporal_interpolation_scale = temporal_interpolation_scale
|
|
self.use_positional_embeddings = use_positional_embeddings
|
|
self.use_learned_positional_embeddings = use_learned_positional_embeddings
|
|
|
|
if patch_size_t is None:
|
|
# CogVideoX 1.0 checkpoints
|
|
self.proj = nn.Conv2d(
|
|
in_channels, embed_dim, kernel_size=(patch_size, patch_size), stride=patch_size, bias=bias
|
|
)
|
|
else:
|
|
# CogVideoX 1.5 checkpoints
|
|
self.proj = nn.Linear(in_channels * patch_size * patch_size * patch_size_t, embed_dim)
|
|
|
|
self.text_proj = nn.Linear(text_embed_dim, embed_dim)
|
|
|
|
if use_positional_embeddings or use_learned_positional_embeddings:
|
|
persistent = use_learned_positional_embeddings
|
|
pos_embedding = self._get_positional_embeddings(sample_height, sample_width, sample_frames)
|
|
self.register_buffer("pos_embedding", pos_embedding, persistent=persistent)
|
|
|
|
def _get_positional_embeddings(self, sample_height: int, sample_width: int, sample_frames: int) -> torch.Tensor:
|
|
post_patch_height = sample_height // self.patch_size
|
|
post_patch_width = sample_width // self.patch_size
|
|
post_time_compression_frames = (sample_frames - 1) // self.temporal_compression_ratio + 1
|
|
num_patches = post_patch_height * post_patch_width * post_time_compression_frames
|
|
|
|
pos_embedding = get_3d_sincos_pos_embed(
|
|
self.embed_dim,
|
|
(post_patch_width, post_patch_height),
|
|
post_time_compression_frames,
|
|
self.spatial_interpolation_scale,
|
|
self.temporal_interpolation_scale,
|
|
)
|
|
pos_embedding = torch.from_numpy(pos_embedding).flatten(0, 1)
|
|
joint_pos_embedding = torch.zeros(
|
|
1, self.max_text_seq_length + num_patches, self.embed_dim, requires_grad=False
|
|
)
|
|
joint_pos_embedding.data[:, self.max_text_seq_length :].copy_(pos_embedding)
|
|
|
|
return joint_pos_embedding
|
|
|
|
def forward(self, text_embeds: torch.Tensor, image_embeds: torch.Tensor):
|
|
r"""
|
|
Args:
|
|
text_embeds (`torch.Tensor`):
|
|
Input text embeddings. Expected shape: (batch_size, seq_length, embedding_dim).
|
|
image_embeds (`torch.Tensor`):
|
|
Input image embeddings. Expected shape: (batch_size, num_frames, channels, height, width).
|
|
"""
|
|
text_embeds = self.text_proj(text_embeds)
|
|
|
|
batch_size, num_frames, channels, height, width = image_embeds.shape
|
|
|
|
if self.patch_size_t is None:
|
|
image_embeds = image_embeds.reshape(-1, channels, height, width)
|
|
image_embeds = self.proj(image_embeds)
|
|
image_embeds = image_embeds.view(batch_size, num_frames, *image_embeds.shape[1:])
|
|
image_embeds = image_embeds.flatten(3).transpose(2, 3) # [batch, num_frames, height x width, channels]
|
|
image_embeds = image_embeds.flatten(1, 2) # [batch, num_frames x height x width, channels]
|
|
else:
|
|
p = self.patch_size
|
|
p_t = self.patch_size_t
|
|
|
|
image_embeds = image_embeds.permute(0, 1, 3, 4, 2)
|
|
image_embeds = image_embeds.reshape(
|
|
batch_size, num_frames // p_t, p_t, height // p, p, width // p, p, channels
|
|
)
|
|
image_embeds = image_embeds.permute(0, 1, 3, 5, 7, 2, 4, 6).flatten(4, 7).flatten(1, 3)
|
|
image_embeds = self.proj(image_embeds)
|
|
|
|
embeds = torch.cat(
|
|
[text_embeds, image_embeds], dim=1
|
|
).contiguous() # [batch, seq_length + num_frames x height x width, channels]
|
|
|
|
if self.use_positional_embeddings or self.use_learned_positional_embeddings:
|
|
if self.use_learned_positional_embeddings and (self.sample_width != width or self.sample_height != height):
|
|
raise ValueError(
|
|
"It is currently not possible to generate videos at a different resolution that the defaults. This should only be the case with 'THUDM/CogVideoX-5b-I2V'."
|
|
"If you think this is incorrect, please open an issue at https://github.com/huggingface/diffusers/issues."
|
|
)
|
|
|
|
pre_time_compression_frames = (num_frames - 1) * self.temporal_compression_ratio + 1
|
|
|
|
if (
|
|
self.sample_height != height
|
|
or self.sample_width != width
|
|
or self.sample_frames != pre_time_compression_frames
|
|
):
|
|
pos_embedding = self._get_positional_embeddings(height, width, pre_time_compression_frames)
|
|
pos_embedding = pos_embedding.to(embeds.device, dtype=embeds.dtype)
|
|
else:
|
|
pos_embedding = self.pos_embedding
|
|
|
|
embeds = embeds + pos_embedding
|
|
|
|
return embeds
|
|
|
|
|
|
class FeedForward(nn.Module):
|
|
r"""
|
|
A feed-forward layer.
|
|
|
|
Parameters:
|
|
dim (`int`): The number of channels in the input.
|
|
dim_out (`int`, *optional*): The number of channels in the output. If not given, defaults to `dim`.
|
|
mult (`int`, *optional*, defaults to 4): The multiplier to use for the hidden dimension.
|
|
dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
|
|
activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.
|
|
final_dropout (`bool` *optional*, defaults to False): Apply a final dropout.
|
|
bias (`bool`, defaults to True): Whether to use a bias in the linear layer.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
dim: int,
|
|
dim_out: Optional[int] = None,
|
|
mult: int = 4,
|
|
dropout: float = 0.0,
|
|
activation_fn: str = "geglu",
|
|
final_dropout: bool = False,
|
|
inner_dim=None,
|
|
bias: bool = True,
|
|
):
|
|
super().__init__()
|
|
if inner_dim is None:
|
|
inner_dim = int(dim * mult)
|
|
dim_out = dim_out if dim_out is not None else dim
|
|
|
|
if activation_fn == "gelu":
|
|
act_fn = GELU(dim, inner_dim, bias=bias)
|
|
if activation_fn == "gelu-approximate":
|
|
act_fn = GELU(dim, inner_dim, approximate="tanh", bias=bias)
|
|
elif activation_fn == "geglu":
|
|
act_fn = GEGLU(dim, inner_dim, bias=bias)
|
|
elif activation_fn == "geglu-approximate":
|
|
act_fn = ApproximateGELU(dim, inner_dim, bias=bias)
|
|
elif activation_fn == "swiglu":
|
|
act_fn = SwiGLU(dim, inner_dim, bias=bias)
|
|
|
|
self.net = nn.ModuleList([])
|
|
# project in
|
|
self.net.append(act_fn)
|
|
# project dropout
|
|
self.net.append(nn.Dropout(dropout))
|
|
# project out
|
|
self.net.append(nn.Linear(inner_dim, dim_out, bias=bias))
|
|
# FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout
|
|
if final_dropout:
|
|
self.net.append(nn.Dropout(dropout))
|
|
|
|
def forward(self, hidden_states: torch.Tensor, *args, **kwargs) -> torch.Tensor:
|
|
if len(args) > 0 or kwargs.get("scale", None) is not None:
|
|
deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
|
|
print(deprecation_message)
|
|
for module in self.net:
|
|
hidden_states = module(hidden_states)
|
|
return hidden_states
|
|
|
|
|
|
class Attention(nn.Module):
|
|
def __init__(
|
|
self,
|
|
query_dim: int,
|
|
dim_head: int = 64,
|
|
heads: int = 8,
|
|
kv_heads: Optional[int] = None,
|
|
qk_norm: Optional[str] = None,
|
|
eps: float = 1e-5,
|
|
bias: bool = False,
|
|
out_bias: bool = True,
|
|
dropout: float = 0.0,
|
|
out_dim: int = None,
|
|
cross_attention_dim: Optional[int] = None,
|
|
):
|
|
super().__init__()
|
|
self.inner_dim = out_dim if out_dim is not None else dim_head * heads
|
|
self.inner_kv_dim = self.inner_dim if kv_heads is None else dim_head * kv_heads
|
|
self.query_dim = query_dim
|
|
self.cross_attention_dim = cross_attention_dim if cross_attention_dim is not None else query_dim
|
|
self.is_cross_attention = cross_attention_dim is not None
|
|
self.out_dim = out_dim if out_dim is not None else query_dim
|
|
self.heads = out_dim // dim_head if out_dim is not None else heads
|
|
|
|
if qk_norm is None:
|
|
self.norm_q = None
|
|
self.norm_k = None
|
|
elif qk_norm == "layer_norm":
|
|
self.norm_q = nn.LayerNorm(dim_head, eps=eps)
|
|
self.norm_k = nn.LayerNorm(dim_head, eps=eps)
|
|
|
|
self.to_q = nn.Linear(query_dim, self.inner_dim, bias=bias)
|
|
self.to_k = nn.Linear(self.cross_attention_dim, self.inner_kv_dim, bias=bias)
|
|
self.to_v = nn.Linear(self.cross_attention_dim, self.inner_kv_dim, bias=bias)
|
|
self.to_out = nn.ModuleList([])
|
|
self.to_out.append(nn.Linear(self.inner_dim, self.out_dim, bias=out_bias))
|
|
self.to_out.append(nn.Dropout(dropout))
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
encoder_hidden_states: torch.Tensor,
|
|
attention_mask: Optional[torch.Tensor] = None,
|
|
image_rotary_emb: Optional[torch.Tensor] = None,
|
|
) -> torch.Tensor:
|
|
|
|
text_seq_length = encoder_hidden_states.size(1)
|
|
|
|
hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
|
|
|
|
batch_size, sequence_length, _ = (
|
|
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
|
|
)
|
|
|
|
query = self.to_q(hidden_states)
|
|
key = self.to_k(hidden_states)
|
|
value = self.to_v(hidden_states)
|
|
|
|
inner_dim = key.shape[-1]
|
|
head_dim = inner_dim // self.heads
|
|
|
|
query = query.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
|
|
key = key.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
|
|
value = value.view(batch_size, -1, self.heads, head_dim).transpose(1, 2)
|
|
|
|
if self.norm_q is not None:
|
|
query = self.norm_q(query)
|
|
if self.norm_k is not None:
|
|
key = self.norm_k(key)
|
|
|
|
# Apply RoPE if needed
|
|
if image_rotary_emb is not None:
|
|
query[:, :, text_seq_length:] = apply_rotary_emb(query[:, :, text_seq_length:], image_rotary_emb)
|
|
if not self.is_cross_attention:
|
|
key[:, :, text_seq_length:] = apply_rotary_emb(key[:, :, text_seq_length:], image_rotary_emb)
|
|
|
|
hidden_states = F.scaled_dot_product_attention(
|
|
query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
|
|
)
|
|
|
|
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, self.heads * head_dim)
|
|
|
|
# linear proj
|
|
hidden_states = self.to_out[0](hidden_states)
|
|
# dropout
|
|
hidden_states = self.to_out[1](hidden_states)
|
|
|
|
encoder_hidden_states, hidden_states = hidden_states.split(
|
|
[text_seq_length, hidden_states.size(1) - text_seq_length], dim=1
|
|
)
|
|
return hidden_states, encoder_hidden_states
|
|
|
|
|
|
class CogVideoXBlock(nn.Module):
|
|
r"""
|
|
Transformer block used in [CogVideoX](https://github.com/THUDM/CogVideo) model.
|
|
|
|
Parameters:
|
|
dim (`int`):
|
|
The number of channels in the input and output.
|
|
num_attention_heads (`int`):
|
|
The number of heads to use for multi-head attention.
|
|
attention_head_dim (`int`):
|
|
The number of channels in each head.
|
|
time_embed_dim (`int`):
|
|
The number of channels in timestep embedding.
|
|
dropout (`float`, defaults to `0.0`):
|
|
The dropout probability to use.
|
|
activation_fn (`str`, defaults to `"gelu-approximate"`):
|
|
Activation function to be used in feed-forward.
|
|
attention_bias (`bool`, defaults to `False`):
|
|
Whether or not to use bias in attention projection layers.
|
|
qk_norm (`bool`, defaults to `True`):
|
|
Whether or not to use normalization after query and key projections in Attention.
|
|
norm_elementwise_affine (`bool`, defaults to `True`):
|
|
Whether to use learnable elementwise affine parameters for normalization.
|
|
norm_eps (`float`, defaults to `1e-5`):
|
|
Epsilon value for normalization layers.
|
|
final_dropout (`bool` defaults to `False`):
|
|
Whether to apply a final dropout after the last feed-forward layer.
|
|
ff_inner_dim (`int`, *optional*, defaults to `None`):
|
|
Custom hidden dimension of Feed-forward layer. If not provided, `4 * dim` is used.
|
|
ff_bias (`bool`, defaults to `True`):
|
|
Whether or not to use bias in Feed-forward layer.
|
|
attention_out_bias (`bool`, defaults to `True`):
|
|
Whether or not to use bias in Attention output projection layer.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
dim: int,
|
|
num_attention_heads: int,
|
|
attention_head_dim: int,
|
|
time_embed_dim: int,
|
|
dropout: float = 0.0,
|
|
activation_fn: str = "gelu-approximate",
|
|
attention_bias: bool = False,
|
|
qk_norm: bool = True,
|
|
norm_elementwise_affine: bool = True,
|
|
norm_eps: float = 1e-5,
|
|
final_dropout: bool = True,
|
|
ff_inner_dim: Optional[int] = None,
|
|
ff_bias: bool = True,
|
|
attention_out_bias: bool = True,
|
|
):
|
|
super().__init__()
|
|
|
|
# 1. Self Attention
|
|
self.norm1 = CogVideoXLayerNormZero(time_embed_dim, dim, norm_elementwise_affine, norm_eps, bias=True)
|
|
|
|
self.attn1 = Attention(
|
|
query_dim=dim,
|
|
dim_head=attention_head_dim,
|
|
heads=num_attention_heads,
|
|
qk_norm="layer_norm" if qk_norm else None,
|
|
eps=1e-6,
|
|
bias=attention_bias,
|
|
out_bias=attention_out_bias
|
|
)
|
|
|
|
# 2. Feed Forward
|
|
self.norm2 = CogVideoXLayerNormZero(time_embed_dim, dim, norm_elementwise_affine, norm_eps, bias=True)
|
|
|
|
self.ff = FeedForward(
|
|
dim,
|
|
dropout=dropout,
|
|
activation_fn=activation_fn,
|
|
final_dropout=final_dropout,
|
|
inner_dim=ff_inner_dim,
|
|
bias=ff_bias,
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
encoder_hidden_states: torch.Tensor,
|
|
temb: torch.Tensor,
|
|
image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
|
|
) -> torch.Tensor:
|
|
text_seq_length = encoder_hidden_states.size(1)
|
|
|
|
# norm & modulate
|
|
norm_hidden_states, norm_encoder_hidden_states, gate_msa, enc_gate_msa = self.norm1(
|
|
hidden_states, encoder_hidden_states, temb
|
|
)
|
|
|
|
# attention
|
|
attn_hidden_states, attn_encoder_hidden_states = self.attn1(
|
|
hidden_states=norm_hidden_states,
|
|
encoder_hidden_states=norm_encoder_hidden_states,
|
|
image_rotary_emb=image_rotary_emb,
|
|
)
|
|
|
|
hidden_states = hidden_states + gate_msa * attn_hidden_states
|
|
encoder_hidden_states = encoder_hidden_states + enc_gate_msa * attn_encoder_hidden_states
|
|
|
|
# norm & modulate
|
|
norm_hidden_states, norm_encoder_hidden_states, gate_ff, enc_gate_ff = self.norm2(
|
|
hidden_states, encoder_hidden_states, temb
|
|
)
|
|
|
|
# feed-forward
|
|
norm_hidden_states = torch.cat([norm_encoder_hidden_states, norm_hidden_states], dim=1)
|
|
ff_output = self.ff(norm_hidden_states)
|
|
|
|
hidden_states = hidden_states + gate_ff * ff_output[:, text_seq_length:]
|
|
encoder_hidden_states = encoder_hidden_states + enc_gate_ff * ff_output[:, :text_seq_length]
|
|
|
|
return hidden_states, encoder_hidden_states
|