684 lines
26 KiB
Python
684 lines
26 KiB
Python
# Copyright 2025 StepFun Inc. All Rights Reserved.
|
|
#
|
|
# Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
# of this software and associated documentation files (the "Software"), to deal
|
|
# in the Software without restriction, including without limitation the rights
|
|
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
# copies of the Software, and to permit persons to whom the Software is
|
|
# furnished to do so, subject to the following conditions:
|
|
#
|
|
# The above copyright notice and this permission notice shall be included in all
|
|
# copies or substantial portions of the Software.
|
|
# ==============================================================================
|
|
|
|
import torch
|
|
from einops import rearrange, repeat
|
|
from torch import nn
|
|
|
|
from fastvideo.v1.attention import DistributedAttention, LocalAttention
|
|
from fastvideo.v1.configs.models.dits import StepVideoConfig
|
|
from fastvideo.v1.distributed.parallel_state import (
|
|
get_sequence_model_parallel_world_size)
|
|
from fastvideo.v1.layers.layernorm import LayerNormScaleShift
|
|
from fastvideo.v1.layers.linear import ReplicatedLinear
|
|
from fastvideo.v1.layers.mlp import MLP
|
|
from fastvideo.v1.layers.rotary_embedding import (_apply_rotary_emb,
|
|
get_rotary_pos_embed)
|
|
from fastvideo.v1.layers.visual_embedding import TimestepEmbedder
|
|
from fastvideo.v1.models.dits.base import BaseDiT
|
|
from fastvideo.v1.platforms import _Backend
|
|
|
|
|
|
class PatchEmbed2D(nn.Module):
|
|
"""2D Image to Patch Embedding
|
|
|
|
Image to Patch Embedding using Conv2d
|
|
|
|
A convolution based approach to patchifying a 2D image w/ embedding projection.
|
|
|
|
Based on the impl in https://github.com/google-research/vision_transformer
|
|
|
|
Hacked together by / Copyright 2020 Ross Wightman
|
|
|
|
Remove the _assert function in forward function to be compatible with multi-resolution images.
|
|
"""
|
|
|
|
def __init__(self,
|
|
patch_size=16,
|
|
in_chans=3,
|
|
embed_dim=768,
|
|
norm_layer=None,
|
|
flatten=True,
|
|
bias=True,
|
|
dtype=None,
|
|
prefix: str = ""):
|
|
super().__init__()
|
|
# Convert patch_size to 2-tuple
|
|
if isinstance(patch_size, list | tuple):
|
|
if len(patch_size) == 1:
|
|
patch_size = (patch_size[0], patch_size[0])
|
|
else:
|
|
patch_size = (patch_size, patch_size)
|
|
|
|
self.patch_size = patch_size
|
|
self.flatten = flatten
|
|
|
|
self.proj = nn.Conv2d(in_chans,
|
|
embed_dim,
|
|
kernel_size=patch_size,
|
|
stride=patch_size,
|
|
bias=bias,
|
|
dtype=dtype)
|
|
self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
|
|
|
|
def forward(self, x):
|
|
x = self.proj(x)
|
|
if self.flatten:
|
|
x = x.flatten(2).transpose(1, 2) # BCHW -> BNC
|
|
x = self.norm(x)
|
|
return x
|
|
|
|
|
|
class StepVideoRMSNorm(nn.Module):
|
|
|
|
def __init__(
|
|
self,
|
|
dim: int,
|
|
elementwise_affine=True,
|
|
eps: float = 1e-6,
|
|
device=None,
|
|
dtype=None,
|
|
):
|
|
"""
|
|
Initialize the RMSNorm normalization layer.
|
|
|
|
Args:
|
|
dim (int): The dimension of the input tensor.
|
|
eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6.
|
|
|
|
Attributes:
|
|
eps (float): A small value added to the denominator for numerical stability.
|
|
weight (nn.Parameter): Learnable scaling parameter.
|
|
|
|
"""
|
|
factory_kwargs = {"device": device, "dtype": dtype}
|
|
super().__init__()
|
|
self.eps = eps
|
|
if elementwise_affine:
|
|
self.weight = nn.Parameter(torch.ones(dim, **factory_kwargs))
|
|
|
|
def _norm(self, x) -> torch.Tensor:
|
|
"""
|
|
Apply the RMSNorm normalization to the input tensor.
|
|
|
|
Args:
|
|
x (torch.Tensor): The input tensor.
|
|
|
|
Returns:
|
|
torch.Tensor: The normalized tensor.
|
|
|
|
"""
|
|
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
|
|
|
|
def forward(self, x):
|
|
"""
|
|
Forward pass through the RMSNorm layer.
|
|
|
|
Args:
|
|
x (torch.Tensor): The input tensor.
|
|
|
|
Returns:
|
|
torch.Tensor: The output tensor after applying RMSNorm.
|
|
|
|
"""
|
|
output = self._norm(x.float()).type_as(x)
|
|
if hasattr(self, "weight"):
|
|
output = output * self.weight
|
|
return output
|
|
|
|
|
|
class SelfAttention(nn.Module):
|
|
|
|
def __init__(self,
|
|
hidden_dim,
|
|
head_dim,
|
|
rope_split: tuple[int, int, int] = (64, 32, 32),
|
|
bias: bool = False,
|
|
with_rope: bool = True,
|
|
with_qk_norm: bool = True,
|
|
attn_type: str = "torch",
|
|
supported_attention_backends=(_Backend.FLASH_ATTN,
|
|
_Backend.TORCH_SDPA)):
|
|
super().__init__()
|
|
self.head_dim = head_dim
|
|
self.hidden_dim = hidden_dim
|
|
self.rope_split = list(rope_split)
|
|
self.n_heads = hidden_dim // head_dim
|
|
|
|
self.wqkv = ReplicatedLinear(hidden_dim, hidden_dim * 3, bias=bias)
|
|
self.wo = ReplicatedLinear(hidden_dim, hidden_dim, bias=bias)
|
|
|
|
self.with_rope = with_rope
|
|
self.with_qk_norm = with_qk_norm
|
|
if self.with_qk_norm:
|
|
self.q_norm = StepVideoRMSNorm(head_dim, elementwise_affine=True)
|
|
self.k_norm = StepVideoRMSNorm(head_dim, elementwise_affine=True)
|
|
|
|
# self.core_attention = self.attn_processor(attn_type=attn_type)
|
|
self.parallel = attn_type == 'parallel'
|
|
self.attn = DistributedAttention(
|
|
num_heads=self.n_heads,
|
|
head_size=head_dim,
|
|
causal=False,
|
|
supported_attention_backends=supported_attention_backends)
|
|
|
|
def _apply_rope(self, x: torch.Tensor, cos: torch.Tensor,
|
|
sin: torch.Tensor):
|
|
"""
|
|
x: [B, S, H, D]
|
|
cos: [S, D/2] where D = head_dim = sum(self.rope_split)
|
|
sin: [S, D/2]
|
|
returns x with rotary applied exactly as v0 did
|
|
"""
|
|
B, S, H, D = x.shape
|
|
# 1) split cos/sin per chunk
|
|
half_splits = [c // 2
|
|
for c in self.rope_split] # [32,16,16] for [64,32,32]
|
|
cos_splits = cos.split(half_splits, dim=1)
|
|
sin_splits = sin.split(half_splits, dim=1)
|
|
|
|
outs = []
|
|
idx = 0
|
|
for (chunk_size, cos_i, sin_i) in zip(self.rope_split,
|
|
cos_splits,
|
|
sin_splits,
|
|
strict=False):
|
|
# slice the corresponding channels
|
|
x_chunk = x[..., idx:idx + chunk_size] # [B,S,H,chunk_size]
|
|
idx += chunk_size
|
|
|
|
# flatten to [S, B*H, chunk_size]
|
|
x_flat = rearrange(x_chunk, 'b s h d -> s (b h) d')
|
|
|
|
# apply rotary on *that* chunk
|
|
out_flat = _apply_rotary_emb(x_flat,
|
|
cos_i,
|
|
sin_i,
|
|
is_neox_style=True)
|
|
|
|
# restore [B,S,H,chunk_size]
|
|
out = rearrange(out_flat, 's (b h) d -> b s h d', b=B, h=H)
|
|
outs.append(out)
|
|
|
|
# concatenate back to [B,S,H,D]
|
|
return torch.cat(outs, dim=-1)
|
|
|
|
def forward(self,
|
|
x,
|
|
cu_seqlens=None,
|
|
max_seqlen=None,
|
|
rope_positions=None,
|
|
cos_sin=None,
|
|
attn_mask=None,
|
|
mask_strategy=None):
|
|
|
|
B, S, _ = x.shape
|
|
xqkv, _ = self.wqkv(x)
|
|
xqkv = xqkv.view(*x.shape[:-1], self.n_heads, 3 * self.head_dim)
|
|
q, k, v = torch.split(xqkv, [self.head_dim] * 3, dim=-1) # [B,S,H,D]
|
|
|
|
if self.with_qk_norm:
|
|
q = self.q_norm(q)
|
|
k = self.k_norm(k)
|
|
|
|
if self.with_rope:
|
|
if rope_positions is not None:
|
|
F, Ht, W = rope_positions
|
|
assert F * Ht * W == S, "rope_positions mismatches sequence length"
|
|
|
|
cos, sin = cos_sin
|
|
cos = cos.to(x.device, dtype=x.dtype)
|
|
sin = sin.to(x.device, dtype=x.dtype)
|
|
|
|
q = self._apply_rope(q, cos, sin)
|
|
k = self._apply_rope(k, cos, sin)
|
|
|
|
output, _ = self.attn(q, k, v) # [B,heads,S,D]
|
|
|
|
output = rearrange(output, 'b s h d -> b s (h d)')
|
|
output, _ = self.wo(output)
|
|
|
|
return output
|
|
|
|
|
|
class CrossAttention(nn.Module):
|
|
|
|
def __init__(
|
|
self,
|
|
hidden_dim,
|
|
head_dim,
|
|
bias=False,
|
|
with_qk_norm=True,
|
|
supported_attention_backends=(_Backend.FLASH_ATTN, _Backend.TORCH_SDPA)
|
|
) -> None:
|
|
super().__init__()
|
|
self.head_dim = head_dim
|
|
self.n_heads = hidden_dim // head_dim
|
|
|
|
self.wq = ReplicatedLinear(hidden_dim, hidden_dim, bias=bias)
|
|
self.wkv = ReplicatedLinear(hidden_dim, hidden_dim * 2, bias=bias)
|
|
self.wo = ReplicatedLinear(hidden_dim, hidden_dim, bias=bias)
|
|
|
|
self.with_qk_norm = with_qk_norm
|
|
if self.with_qk_norm:
|
|
self.q_norm = StepVideoRMSNorm(head_dim, elementwise_affine=True)
|
|
self.k_norm = StepVideoRMSNorm(head_dim, elementwise_affine=True)
|
|
|
|
self.attn = LocalAttention(
|
|
num_heads=self.n_heads,
|
|
head_size=head_dim,
|
|
causal=False,
|
|
supported_attention_backends=supported_attention_backends)
|
|
|
|
def forward(self,
|
|
x: torch.Tensor,
|
|
encoder_hidden_states: torch.Tensor,
|
|
attn_mask=None) -> torch.Tensor:
|
|
|
|
xq, _ = self.wq(x)
|
|
xq = xq.view(*xq.shape[:-1], self.n_heads, self.head_dim)
|
|
|
|
xkv, _ = self.wkv(encoder_hidden_states)
|
|
xkv = xkv.view(*xkv.shape[:-1], self.n_heads, 2 * self.head_dim)
|
|
|
|
xk, xv = torch.split(xkv, [self.head_dim] * 2,
|
|
dim=-1) ## seq_len, n, dim
|
|
|
|
if self.with_qk_norm:
|
|
xq = self.q_norm(xq)
|
|
xk = self.k_norm(xk)
|
|
|
|
output = self.attn(xq, xk, xv)
|
|
|
|
output = rearrange(output, 'b s h d -> b s (h d)')
|
|
output, _ = self.wo(output)
|
|
|
|
return output
|
|
|
|
|
|
class AdaLayerNormSingle(nn.Module):
|
|
r"""
|
|
Norm layer adaptive layer norm single (adaLN-single).
|
|
|
|
As proposed in PixArt-Alpha (see: https://arxiv.org/abs/2310.00426; Section 2.3).
|
|
|
|
Parameters:
|
|
embedding_dim (`int`): The size of each embedding vector.
|
|
use_additional_conditions (`bool`): To use additional conditions for normalization or not.
|
|
"""
|
|
|
|
def __init__(self, embedding_dim: int, time_step_rescale=1000):
|
|
super().__init__()
|
|
|
|
self.emb = TimestepEmbedder(embedding_dim)
|
|
|
|
self.silu = nn.SiLU()
|
|
self.linear = ReplicatedLinear(embedding_dim,
|
|
6 * embedding_dim,
|
|
bias=True)
|
|
|
|
self.time_step_rescale = time_step_rescale ## timestep usually in [0, 1], we rescale it to [0,1000] for stability
|
|
|
|
def forward(
|
|
self,
|
|
timestep: torch.Tensor,
|
|
added_cond_kwargs: dict[str, torch.Tensor] | None = None,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
embedded_timestep = self.emb(timestep * self.time_step_rescale)
|
|
|
|
out, _ = self.linear(self.silu(embedded_timestep))
|
|
|
|
return out, embedded_timestep
|
|
|
|
|
|
class StepVideoTransformerBlock(nn.Module):
|
|
r"""
|
|
A basic Transformer block.
|
|
|
|
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.
|
|
dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
|
|
cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention.
|
|
activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.
|
|
num_embeds_ada_norm (:
|
|
obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`.
|
|
attention_bias (:
|
|
obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter.
|
|
only_cross_attention (`bool`, *optional*):
|
|
Whether to use only cross-attention layers. In this case two cross attention layers are used.
|
|
double_self_attention (`bool`, *optional*):
|
|
Whether to use two self-attention layers. In this case no cross attention layers are used.
|
|
upcast_attention (`bool`, *optional*):
|
|
Whether to upcast the attention computation to float32. This is useful for mixed precision training.
|
|
norm_elementwise_affine (`bool`, *optional*, defaults to `True`):
|
|
Whether to use learnable elementwise affine parameters for normalization.
|
|
norm_type (`str`, *optional*, defaults to `"layer_norm"`):
|
|
The normalization layer to use. Can be `"layer_norm"`, `"ada_norm"` or `"ada_norm_zero"`.
|
|
final_dropout (`bool` *optional*, defaults to False):
|
|
Whether to apply a final dropout after the last feed-forward layer.
|
|
positional_embeddings (`str`, *optional*, defaults to `None`):
|
|
The type of positional embeddings to apply to.
|
|
num_positional_embeddings (`int`, *optional*, defaults to `None`):
|
|
The maximum number of positional embeddings to apply.
|
|
"""
|
|
|
|
def __init__(self,
|
|
dim: int,
|
|
attention_head_dim: int,
|
|
norm_eps: float = 1e-5,
|
|
ff_inner_dim: int | None = None,
|
|
ff_bias: bool = False,
|
|
attention_type: str = 'torch'):
|
|
super().__init__()
|
|
self.dim = dim
|
|
self.norm1 = LayerNormScaleShift(dim,
|
|
norm_type="layer",
|
|
elementwise_affine=True,
|
|
eps=norm_eps)
|
|
self.attn1 = SelfAttention(
|
|
dim,
|
|
attention_head_dim,
|
|
bias=False,
|
|
with_rope=True,
|
|
with_qk_norm=True,
|
|
)
|
|
|
|
self.norm2 = LayerNormScaleShift(dim,
|
|
norm_type="layer",
|
|
elementwise_affine=True,
|
|
eps=norm_eps)
|
|
self.attn2 = CrossAttention(dim,
|
|
attention_head_dim,
|
|
bias=False,
|
|
with_qk_norm=True)
|
|
|
|
self.ff = MLP(input_dim=dim,
|
|
mlp_hidden_dim=dim *
|
|
4 if ff_inner_dim is None else ff_inner_dim,
|
|
act_type="gelu_pytorch_tanh",
|
|
bias=ff_bias)
|
|
|
|
self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5)
|
|
|
|
@torch.no_grad()
|
|
def forward(self,
|
|
q: torch.Tensor,
|
|
kv: torch.Tensor,
|
|
t_expand: torch.LongTensor,
|
|
attn_mask=None,
|
|
rope_positions: list | None = None,
|
|
cos_sin=None,
|
|
mask_strategy=None) -> torch.Tensor:
|
|
|
|
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
|
|
torch.clone(chunk)
|
|
for chunk in (self.scale_shift_table[None] +
|
|
t_expand.reshape(-1, 6, self.dim)).chunk(6, dim=1))
|
|
|
|
scale_shift_q = self.norm1(q,
|
|
scale=scale_msa.squeeze(1),
|
|
shift=shift_msa.squeeze(1))
|
|
|
|
attn_q = self.attn1(scale_shift_q,
|
|
rope_positions=rope_positions,
|
|
cos_sin=cos_sin,
|
|
mask_strategy=mask_strategy)
|
|
|
|
q = attn_q * gate_msa + q
|
|
|
|
attn_q = self.attn2(q, kv, attn_mask)
|
|
|
|
q = attn_q + q
|
|
|
|
scale_shift_q = self.norm2(q,
|
|
scale=scale_mlp.squeeze(1),
|
|
shift=shift_mlp.squeeze(1))
|
|
|
|
ff_output = self.ff(scale_shift_q)
|
|
|
|
q = ff_output * gate_mlp + q
|
|
|
|
return q
|
|
|
|
|
|
class StepVideoModel(BaseDiT):
|
|
# (Optional) Keep the same attribute for compatibility with splitting, etc.
|
|
_fsdp_shard_conditions = [
|
|
lambda n, m: "transformer_blocks" in n and n.split(".")[-1].isdigit(),
|
|
# lambda n, m: "pos_embed" in n # If needed for the patch embedding.
|
|
]
|
|
_param_names_mapping = StepVideoConfig()._param_names_mapping
|
|
_supported_attention_backends = StepVideoConfig(
|
|
)._supported_attention_backends
|
|
|
|
def __init__(self, config: StepVideoConfig) -> None:
|
|
super().__init__(config=config)
|
|
self.num_attention_heads = config.num_attention_heads
|
|
self.attention_head_dim = config.attention_head_dim
|
|
self.in_channels = config.in_channels
|
|
self.out_channels = config.out_channels
|
|
self.num_layers = config.num_layers
|
|
self.dropout = config.dropout
|
|
self.patch_size = config.patch_size
|
|
self.norm_type = config.norm_type
|
|
self.norm_elementwise_affine = config.norm_elementwise_affine
|
|
self.norm_eps = config.norm_eps
|
|
self.use_additional_conditions = config.use_additional_conditions
|
|
self.caption_channels = config.caption_channels
|
|
self.attention_type = config.attention_type
|
|
self.num_channels_latents = config.num_channels_latents
|
|
# Compute inner dimension.
|
|
self.hidden_size = config.hidden_size
|
|
|
|
# Image/video patch embedding.
|
|
self.pos_embed = PatchEmbed2D(
|
|
patch_size=self.patch_size,
|
|
in_chans=self.in_channels,
|
|
embed_dim=self.hidden_size,
|
|
)
|
|
|
|
self._rope_cache: dict[tuple, tuple[torch.Tensor, torch.Tensor]] = {}
|
|
# Transformer blocks.
|
|
self.transformer_blocks = nn.ModuleList([
|
|
StepVideoTransformerBlock(
|
|
dim=self.hidden_size,
|
|
attention_head_dim=self.attention_head_dim,
|
|
attention_type=self.attention_type)
|
|
for _ in range(self.num_layers)
|
|
])
|
|
|
|
# Output blocks.
|
|
self.norm_out = LayerNormScaleShift(
|
|
self.hidden_size,
|
|
norm_type="layer",
|
|
eps=self.norm_eps,
|
|
elementwise_affine=self.norm_elementwise_affine)
|
|
self.scale_shift_table = nn.Parameter(
|
|
torch.randn(2, self.hidden_size) / (self.hidden_size**0.5))
|
|
self.proj_out = ReplicatedLinear(
|
|
self.hidden_size,
|
|
self.patch_size * self.patch_size * self.out_channels)
|
|
# Time modulation via adaptive layer norm.
|
|
self.adaln_single = AdaLayerNormSingle(self.hidden_size)
|
|
|
|
# Set up caption conditioning.
|
|
if isinstance(self.caption_channels, int):
|
|
caption_channel = self.caption_channels
|
|
else:
|
|
caption_channel, clip_channel = self.caption_channels
|
|
self.clip_projection = ReplicatedLinear(clip_channel,
|
|
self.hidden_size)
|
|
self.caption_norm = nn.LayerNorm(
|
|
caption_channel,
|
|
eps=self.norm_eps,
|
|
elementwise_affine=self.norm_elementwise_affine)
|
|
self.caption_projection = MLP(input_dim=caption_channel,
|
|
mlp_hidden_dim=self.hidden_size,
|
|
act_type="gelu_pytorch_tanh")
|
|
|
|
# Flag to indicate if using parallel attention.
|
|
self.parallel = (self.attention_type == "parallel")
|
|
|
|
self.__post_init__()
|
|
|
|
def patchfy(self, hidden_states) -> torch.Tensor:
|
|
hidden_states = rearrange(hidden_states, 'b f c h w -> (b f) c h w')
|
|
hidden_states = self.pos_embed(hidden_states)
|
|
return hidden_states
|
|
|
|
def prepare_attn_mask(self, encoder_attention_mask, encoder_hidden_states,
|
|
q_seqlen) -> tuple[torch.Tensor, torch.Tensor]:
|
|
kv_seqlens = encoder_attention_mask.sum(dim=1).int()
|
|
mask = torch.zeros([len(kv_seqlens), q_seqlen,
|
|
max(kv_seqlens)],
|
|
dtype=torch.bool,
|
|
device=encoder_attention_mask.device)
|
|
encoder_hidden_states = encoder_hidden_states[:, :max(kv_seqlens)]
|
|
for i, kv_len in enumerate(kv_seqlens):
|
|
mask[i, :, :kv_len] = 1
|
|
return encoder_hidden_states, mask
|
|
|
|
def block_forward(self,
|
|
hidden_states,
|
|
encoder_hidden_states=None,
|
|
t_expand=None,
|
|
rope_positions=None,
|
|
cos_sin=None,
|
|
attn_mask=None,
|
|
parallel=True,
|
|
mask_strategy=None) -> torch.Tensor:
|
|
|
|
for i, block in enumerate(self.transformer_blocks):
|
|
hidden_states = block(hidden_states,
|
|
encoder_hidden_states,
|
|
t_expand=t_expand,
|
|
attn_mask=attn_mask,
|
|
rope_positions=rope_positions,
|
|
cos_sin=cos_sin,
|
|
mask_strategy=mask_strategy[i])
|
|
|
|
return hidden_states
|
|
|
|
def _get_rope(self, rope_positions: tuple[int, int, int],
|
|
dtype: torch.dtype, device: torch.device):
|
|
F, Ht, W = rope_positions
|
|
key = (F, Ht, W, dtype)
|
|
if key not in self._rope_cache:
|
|
cos, sin = get_rotary_pos_embed(
|
|
rope_sizes=(F * get_sequence_model_parallel_world_size(), Ht,
|
|
W),
|
|
hidden_size=self.hidden_size,
|
|
heads_num=self.hidden_size // self.attention_head_dim,
|
|
rope_dim_list=(64, 32, 32), # same split you used
|
|
rope_theta=1.0e4,
|
|
dtype=torch.float32 # build once in fp32
|
|
)
|
|
# move & cast once
|
|
self._rope_cache[key] = (cos.to(device, dtype=dtype),
|
|
sin.to(device, dtype=dtype))
|
|
return self._rope_cache[key]
|
|
|
|
@torch.inference_mode()
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
encoder_hidden_states: torch.Tensor | None = None,
|
|
t_expand: torch.LongTensor | None = None,
|
|
encoder_hidden_states_2: torch.Tensor | None = None,
|
|
added_cond_kwargs: dict[str, torch.Tensor] | None = None,
|
|
encoder_attention_mask: torch.Tensor | None = None,
|
|
fps: torch.Tensor | None = None,
|
|
return_dict: bool = True,
|
|
mask_strategy=None,
|
|
guidance=None,
|
|
):
|
|
assert hidden_states.ndim == 5
|
|
"hidden_states's shape should be (bsz, f, ch, h ,w)"
|
|
frame = hidden_states.shape[2]
|
|
hidden_states = rearrange(hidden_states,
|
|
'b c f h w -> b f c h w',
|
|
f=frame)
|
|
if mask_strategy is None:
|
|
mask_strategy = [None, None]
|
|
bsz, frame, _, height, width = hidden_states.shape
|
|
height, width = height // self.patch_size, width // self.patch_size
|
|
|
|
hidden_states = self.patchfy(hidden_states)
|
|
len_frame = hidden_states.shape[1]
|
|
|
|
t_expand, embedded_timestep = self.adaln_single(t_expand)
|
|
encoder_hidden_states = self.caption_projection(
|
|
self.caption_norm(encoder_hidden_states))
|
|
|
|
if encoder_hidden_states_2 is not None and hasattr(
|
|
self, 'clip_projection'):
|
|
clip_embedding, _ = self.clip_projection(encoder_hidden_states_2)
|
|
encoder_hidden_states = torch.cat(
|
|
[clip_embedding, encoder_hidden_states], dim=1)
|
|
|
|
hidden_states = rearrange(hidden_states,
|
|
'(b f) l d-> b (f l) d',
|
|
b=bsz,
|
|
f=frame,
|
|
l=len_frame).contiguous()
|
|
encoder_hidden_states, attn_mask = self.prepare_attn_mask(
|
|
encoder_attention_mask,
|
|
encoder_hidden_states,
|
|
q_seqlen=frame * len_frame)
|
|
|
|
cos_sin = self._get_rope((frame, height, width), hidden_states.dtype,
|
|
hidden_states.device)
|
|
|
|
hidden_states = self.block_forward(
|
|
hidden_states,
|
|
encoder_hidden_states,
|
|
t_expand=t_expand,
|
|
rope_positions=[frame, height, width],
|
|
cos_sin=cos_sin,
|
|
attn_mask=attn_mask,
|
|
parallel=self.parallel,
|
|
mask_strategy=mask_strategy)
|
|
|
|
hidden_states = rearrange(hidden_states,
|
|
'b (f l) d -> (b f) l d',
|
|
b=bsz,
|
|
f=frame,
|
|
l=len_frame)
|
|
|
|
embedded_timestep = repeat(embedded_timestep, 'b d -> (b f) d',
|
|
f=frame).contiguous()
|
|
|
|
shift, scale = (self.scale_shift_table[None] +
|
|
embedded_timestep[:, None]).chunk(2, dim=1)
|
|
hidden_states = self.norm_out(hidden_states,
|
|
shift=shift.squeeze(1),
|
|
scale=scale.squeeze(1))
|
|
# Modulation
|
|
hidden_states, _ = self.proj_out(hidden_states)
|
|
|
|
# unpatchify
|
|
hidden_states = hidden_states.reshape(shape=(-1, height, width,
|
|
self.patch_size,
|
|
self.patch_size,
|
|
self.out_channels))
|
|
|
|
hidden_states = rearrange(hidden_states, 'n h w p q c -> n c h p w q')
|
|
output = hidden_states.reshape(shape=(-1, self.out_channels,
|
|
height * self.patch_size,
|
|
width * self.patch_size))
|
|
|
|
output = rearrange(output, '(b f) c h w -> b c f h w', f=frame)
|
|
return output
|