v5 release with triple architecture support and prompt enhancer
This commit is contained in:
DeepBeepMeep
26
hyvideo/modules/__init__.py
Normal file
26
hyvideo/modules/__init__.py
Normal file
@@ -0,0 +1,26 @@
|
||||
from .models import HYVideoDiffusionTransformer, HUNYUAN_VIDEO_CONFIG
|
||||
|
||||
|
||||
def load_model(model, i2v_condition_type, in_channels, out_channels, factor_kwargs):
|
||||
"""load hunyuan video model
|
||||
|
||||
Args:
|
||||
args (dict): model args
|
||||
in_channels (int): input channels number
|
||||
out_channels (int): output channels number
|
||||
factor_kwargs (dict): factor kwargs
|
||||
|
||||
Returns:
|
||||
model (nn.Module): The hunyuan video model
|
||||
"""
|
||||
if model in HUNYUAN_VIDEO_CONFIG.keys():
|
||||
model = HYVideoDiffusionTransformer(
|
||||
i2v_condition_type = i2v_condition_type,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
**HUNYUAN_VIDEO_CONFIG[model],
|
||||
**factor_kwargs,
|
||||
)
|
||||
return model
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
23
hyvideo/modules/activation_layers.py
Normal file
23
hyvideo/modules/activation_layers.py
Normal file
@@ -0,0 +1,23 @@
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
def get_activation_layer(act_type):
|
||||
"""get activation layer
|
||||
|
||||
Args:
|
||||
act_type (str): the activation type
|
||||
|
||||
Returns:
|
||||
torch.nn.functional: the activation layer
|
||||
"""
|
||||
if act_type == "gelu":
|
||||
return lambda: nn.GELU()
|
||||
elif act_type == "gelu_tanh":
|
||||
# Approximate `tanh` requires torch >= 1.13
|
||||
return lambda: nn.GELU(approximate="tanh")
|
||||
elif act_type == "relu":
|
||||
return nn.ReLU
|
||||
elif act_type == "silu":
|
||||
return nn.SiLU
|
||||
else:
|
||||
raise ValueError(f"Unknown activation type: {act_type}")
|
||||
362
hyvideo/modules/attenion.py
Normal file
362
hyvideo/modules/attenion.py
Normal file
@@ -0,0 +1,362 @@
|
||||
import importlib.metadata
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from importlib.metadata import version
|
||||
|
||||
def clear_list(l):
|
||||
for i in range(len(l)):
|
||||
l[i] = None
|
||||
|
||||
try:
|
||||
import flash_attn
|
||||
from flash_attn.flash_attn_interface import _flash_attn_forward
|
||||
from flash_attn.flash_attn_interface import flash_attn_varlen_func
|
||||
except ImportError:
|
||||
flash_attn = None
|
||||
flash_attn_varlen_func = None
|
||||
_flash_attn_forward = None
|
||||
|
||||
try:
|
||||
from xformers.ops import memory_efficient_attention
|
||||
except ImportError:
|
||||
memory_efficient_attention = None
|
||||
|
||||
try:
|
||||
from sageattention import sageattn_varlen
|
||||
def sageattn_varlen_wrapper(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
cu_seqlens_q,
|
||||
cu_seqlens_kv,
|
||||
max_seqlen_q,
|
||||
max_seqlen_kv,
|
||||
):
|
||||
return sageattn_varlen(q, k, v, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv)
|
||||
except ImportError:
|
||||
sageattn_varlen_wrapper = None
|
||||
|
||||
try:
|
||||
from sageattention import sageattn
|
||||
@torch.compiler.disable()
|
||||
def sageattn_wrapper(
|
||||
qkv_list,
|
||||
attention_length
|
||||
):
|
||||
q,k, v = qkv_list
|
||||
padding_length = q.shape[1] -attention_length
|
||||
q = q[:, :attention_length, :, : ]
|
||||
k = k[:, :attention_length, :, : ]
|
||||
v = v[:, :attention_length, :, : ]
|
||||
|
||||
o = sageattn(q, k, v, tensor_layout="NHD")
|
||||
del q, k ,v
|
||||
clear_list(qkv_list)
|
||||
|
||||
if padding_length > 0:
|
||||
o = torch.cat([o, torch.empty( (o.shape[0], padding_length, *o.shape[-2:]), dtype= o.dtype, device=o.device ) ], 1)
|
||||
|
||||
return o
|
||||
|
||||
except ImportError:
|
||||
sageattn = None
|
||||
|
||||
|
||||
def get_attention_modes():
|
||||
ret = ["sdpa", "auto"]
|
||||
if flash_attn != None:
|
||||
ret.append("flash")
|
||||
if memory_efficient_attention != None:
|
||||
ret.append("xformers")
|
||||
if sageattn_varlen_wrapper != None:
|
||||
ret.append("sage")
|
||||
if sageattn != None and version("sageattention").startswith("2") :
|
||||
ret.append("sage2")
|
||||
|
||||
return ret
|
||||
|
||||
|
||||
|
||||
MEMORY_LAYOUT = {
|
||||
"sdpa": (
|
||||
lambda x: x.transpose(1, 2),
|
||||
lambda x: x.transpose(1, 2),
|
||||
),
|
||||
"xformers": (
|
||||
lambda x: x,
|
||||
lambda x: x,
|
||||
),
|
||||
"sage2": (
|
||||
lambda x: x,
|
||||
lambda x: x,
|
||||
),
|
||||
"sage": (
|
||||
lambda x: x.view(x.shape[0] * x.shape[1], *x.shape[2:]),
|
||||
lambda x: x,
|
||||
),
|
||||
"flash": (
|
||||
lambda x: x.view(x.shape[0] * x.shape[1], *x.shape[2:]),
|
||||
lambda x: x,
|
||||
),
|
||||
"torch": (
|
||||
lambda x: x.transpose(1, 2),
|
||||
lambda x: x.transpose(1, 2),
|
||||
),
|
||||
"vanilla": (
|
||||
lambda x: x.transpose(1, 2),
|
||||
lambda x: x.transpose(1, 2),
|
||||
),
|
||||
}
|
||||
|
||||
@torch.compiler.disable()
|
||||
def sdpa_wrapper(
|
||||
qkv_list,
|
||||
attention_length
|
||||
):
|
||||
q,k, v = qkv_list
|
||||
padding_length = q.shape[2] -attention_length
|
||||
q = q[:, :, :attention_length, :]
|
||||
k = k[:, :, :attention_length, :]
|
||||
v = v[:, :, :attention_length, :]
|
||||
|
||||
o = F.scaled_dot_product_attention(
|
||||
q, k, v, attn_mask=None, is_causal=False
|
||||
)
|
||||
del q, k ,v
|
||||
clear_list(qkv_list)
|
||||
|
||||
if padding_length > 0:
|
||||
o = torch.cat([o, torch.empty( (*o.shape[:2], padding_length, o.shape[-1]), dtype= o.dtype, device=o.device ) ], 2)
|
||||
|
||||
return o
|
||||
|
||||
def get_cu_seqlens(text_mask, img_len):
|
||||
"""Calculate cu_seqlens_q, cu_seqlens_kv using text_mask and img_len
|
||||
|
||||
Args:
|
||||
text_mask (torch.Tensor): the mask of text
|
||||
img_len (int): the length of image
|
||||
|
||||
Returns:
|
||||
torch.Tensor: the calculated cu_seqlens for flash attention
|
||||
"""
|
||||
batch_size = text_mask.shape[0]
|
||||
text_len = text_mask.sum(dim=1)
|
||||
max_len = text_mask.shape[1] + img_len
|
||||
|
||||
cu_seqlens = torch.zeros([2 * batch_size + 1], dtype=torch.int32, device="cuda")
|
||||
|
||||
for i in range(batch_size):
|
||||
s = text_len[i] + img_len
|
||||
s1 = i * max_len + s
|
||||
s2 = (i + 1) * max_len
|
||||
cu_seqlens[2 * i + 1] = s1
|
||||
cu_seqlens[2 * i + 2] = s2
|
||||
|
||||
return cu_seqlens
|
||||
|
||||
|
||||
def attention(
|
||||
qkv_list,
|
||||
mode="flash",
|
||||
drop_rate=0,
|
||||
attn_mask=None,
|
||||
causal=False,
|
||||
cu_seqlens_q=None,
|
||||
cu_seqlens_kv=None,
|
||||
max_seqlen_q=None,
|
||||
max_seqlen_kv=None,
|
||||
batch_size=1,
|
||||
):
|
||||
"""
|
||||
Perform QKV self attention.
|
||||
|
||||
Args:
|
||||
q (torch.Tensor): Query tensor with shape [b, s, a, d], where a is the number of heads.
|
||||
k (torch.Tensor): Key tensor with shape [b, s1, a, d]
|
||||
v (torch.Tensor): Value tensor with shape [b, s1, a, d]
|
||||
mode (str): Attention mode. Choose from 'self_flash', 'cross_flash', 'torch', and 'vanilla'.
|
||||
drop_rate (float): Dropout rate in attention map. (default: 0)
|
||||
attn_mask (torch.Tensor): Attention mask with shape [b, s1] (cross_attn), or [b, a, s, s1] (torch or vanilla).
|
||||
(default: None)
|
||||
causal (bool): Whether to use causal attention. (default: False)
|
||||
cu_seqlens_q (torch.Tensor): dtype torch.int32. The cumulative sequence lengths of the sequences in the batch,
|
||||
used to index into q.
|
||||
cu_seqlens_kv (torch.Tensor): dtype torch.int32. The cumulative sequence lengths of the sequences in the batch,
|
||||
used to index into kv.
|
||||
max_seqlen_q (int): The maximum sequence length in the batch of q.
|
||||
max_seqlen_kv (int): The maximum sequence length in the batch of k and v.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Output tensor after self attention with shape [b, s, ad]
|
||||
"""
|
||||
pre_attn_layout, post_attn_layout = MEMORY_LAYOUT[mode]
|
||||
q , k , v = qkv_list
|
||||
clear_list(qkv_list)
|
||||
del qkv_list
|
||||
padding_length = 0
|
||||
# if attn_mask == None and mode == "sdpa":
|
||||
# padding_length = q.shape[1] - cu_seqlens_q
|
||||
# q = q[:, :cu_seqlens_q, ... ]
|
||||
# k = k[:, :cu_seqlens_kv, ... ]
|
||||
# v = v[:, :cu_seqlens_kv, ... ]
|
||||
|
||||
q = pre_attn_layout(q)
|
||||
k = pre_attn_layout(k)
|
||||
v = pre_attn_layout(v)
|
||||
|
||||
if mode == "torch":
|
||||
if attn_mask is not None and attn_mask.dtype != torch.bool:
|
||||
attn_mask = attn_mask.to(q.dtype)
|
||||
x = F.scaled_dot_product_attention(
|
||||
q, k, v, attn_mask=attn_mask, dropout_p=drop_rate, is_causal=causal
|
||||
)
|
||||
|
||||
elif mode == "sdpa":
|
||||
# if attn_mask is not None and attn_mask.dtype != torch.bool:
|
||||
# attn_mask = attn_mask.to(q.dtype)
|
||||
# x = F.scaled_dot_product_attention(
|
||||
# q, k, v, attn_mask=attn_mask, dropout_p=drop_rate, is_causal=causal
|
||||
# )
|
||||
assert attn_mask==None
|
||||
qkv_list = [q, k, v]
|
||||
del q, k , v
|
||||
x = sdpa_wrapper( qkv_list, cu_seqlens_q )
|
||||
|
||||
elif mode == "xformers":
|
||||
x = memory_efficient_attention(
|
||||
q, k, v , attn_bias= attn_mask
|
||||
)
|
||||
|
||||
elif mode == "sage2":
|
||||
qkv_list = [q, k, v]
|
||||
del q, k , v
|
||||
x = sageattn_wrapper(qkv_list, cu_seqlens_q)
|
||||
|
||||
elif mode == "sage":
|
||||
x = sageattn_varlen_wrapper(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
cu_seqlens_q,
|
||||
cu_seqlens_kv,
|
||||
max_seqlen_q,
|
||||
max_seqlen_kv,
|
||||
)
|
||||
# x with shape [(bxs), a, d]
|
||||
x = x.view(
|
||||
batch_size, max_seqlen_q, x.shape[-2], x.shape[-1]
|
||||
) # reshape x to [b, s, a, d]
|
||||
|
||||
elif mode == "flash":
|
||||
x = flash_attn_varlen_func(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
cu_seqlens_q,
|
||||
cu_seqlens_kv,
|
||||
max_seqlen_q,
|
||||
max_seqlen_kv,
|
||||
)
|
||||
# x with shape [(bxs), a, d]
|
||||
x = x.view(
|
||||
batch_size, max_seqlen_q, x.shape[-2], x.shape[-1]
|
||||
) # reshape x to [b, s, a, d]
|
||||
elif mode == "vanilla":
|
||||
scale_factor = 1 / math.sqrt(q.size(-1))
|
||||
|
||||
b, a, s, _ = q.shape
|
||||
s1 = k.size(2)
|
||||
attn_bias = torch.zeros(b, a, s, s1, dtype=q.dtype, device=q.device)
|
||||
if causal:
|
||||
# Only applied to self attention
|
||||
assert (
|
||||
attn_mask is None
|
||||
), "Causal mask and attn_mask cannot be used together"
|
||||
temp_mask = torch.ones(b, a, s, s, dtype=torch.bool, device=q.device).tril(
|
||||
diagonal=0
|
||||
)
|
||||
attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))
|
||||
attn_bias.to(q.dtype)
|
||||
|
||||
if attn_mask is not None:
|
||||
if attn_mask.dtype == torch.bool:
|
||||
attn_bias.masked_fill_(attn_mask.logical_not(), float("-inf"))
|
||||
else:
|
||||
attn_bias += attn_mask
|
||||
|
||||
# TODO: Maybe force q and k to be float32 to avoid numerical overflow
|
||||
attn = (q @ k.transpose(-2, -1)) * scale_factor
|
||||
attn += attn_bias
|
||||
attn = attn.softmax(dim=-1)
|
||||
attn = torch.dropout(attn, p=drop_rate, train=True)
|
||||
x = attn @ v
|
||||
else:
|
||||
raise NotImplementedError(f"Unsupported attention mode: {mode}")
|
||||
|
||||
x = post_attn_layout(x)
|
||||
b, s, a, d = x.shape
|
||||
out = x.reshape(b, s, -1)
|
||||
if padding_length > 0 :
|
||||
out = torch.cat([out, torch.empty( (out.shape[0], padding_length, out.shape[2]), dtype= out.dtype, device=out.device ) ], 1)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def parallel_attention(
|
||||
hybrid_seq_parallel_attn,
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
img_q_len,
|
||||
img_kv_len,
|
||||
cu_seqlens_q,
|
||||
cu_seqlens_kv
|
||||
):
|
||||
attn1 = hybrid_seq_parallel_attn(
|
||||
None,
|
||||
q[:, :img_q_len, :, :],
|
||||
k[:, :img_kv_len, :, :],
|
||||
v[:, :img_kv_len, :, :],
|
||||
dropout_p=0.0,
|
||||
causal=False,
|
||||
joint_tensor_query=q[:,img_q_len:cu_seqlens_q[1]],
|
||||
joint_tensor_key=k[:,img_kv_len:cu_seqlens_kv[1]],
|
||||
joint_tensor_value=v[:,img_kv_len:cu_seqlens_kv[1]],
|
||||
joint_strategy="rear",
|
||||
)
|
||||
if flash_attn.__version__ >= '2.7.0':
|
||||
attn2, *_ = _flash_attn_forward(
|
||||
q[:,cu_seqlens_q[1]:],
|
||||
k[:,cu_seqlens_kv[1]:],
|
||||
v[:,cu_seqlens_kv[1]:],
|
||||
dropout_p=0.0,
|
||||
softmax_scale=q.shape[-1] ** (-0.5),
|
||||
causal=False,
|
||||
window_size_left=-1,
|
||||
window_size_right=-1,
|
||||
softcap=0.0,
|
||||
alibi_slopes=None,
|
||||
return_softmax=False,
|
||||
)
|
||||
else:
|
||||
attn2, *_ = _flash_attn_forward(
|
||||
q[:,cu_seqlens_q[1]:],
|
||||
k[:,cu_seqlens_kv[1]:],
|
||||
v[:,cu_seqlens_kv[1]:],
|
||||
dropout_p=0.0,
|
||||
softmax_scale=q.shape[-1] ** (-0.5),
|
||||
causal=False,
|
||||
window_size=(-1, -1),
|
||||
softcap=0.0,
|
||||
alibi_slopes=None,
|
||||
return_softmax=False,
|
||||
)
|
||||
attn = torch.cat([attn1, attn2], dim=1)
|
||||
b, s, a, d = attn.shape
|
||||
attn = attn.reshape(b, s, -1)
|
||||
|
||||
return attn
|
||||
157
hyvideo/modules/embed_layers.py
Normal file
157
hyvideo/modules/embed_layers.py
Normal file
@@ -0,0 +1,157 @@
|
||||
import math
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from einops import rearrange, repeat
|
||||
|
||||
from ..utils.helpers import to_2tuple
|
||||
|
||||
|
||||
class PatchEmbed(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,
|
||||
device=None,
|
||||
):
|
||||
factory_kwargs = {"dtype": dtype, "device": device}
|
||||
super().__init__()
|
||||
patch_size = to_2tuple(patch_size)
|
||||
self.patch_size = patch_size
|
||||
self.flatten = flatten
|
||||
|
||||
self.proj = nn.Conv3d(
|
||||
in_chans,
|
||||
embed_dim,
|
||||
kernel_size=patch_size,
|
||||
stride=patch_size,
|
||||
bias=bias,
|
||||
**factory_kwargs
|
||||
)
|
||||
nn.init.xavier_uniform_(self.proj.weight.view(self.proj.weight.size(0), -1))
|
||||
if bias:
|
||||
nn.init.zeros_(self.proj.bias)
|
||||
|
||||
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 TextProjection(nn.Module):
|
||||
"""
|
||||
Projects text embeddings. Also handles dropout for classifier-free guidance.
|
||||
|
||||
Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/nets/PixArt_blocks.py
|
||||
"""
|
||||
|
||||
def __init__(self, in_channels, hidden_size, act_layer, dtype=None, device=None):
|
||||
factory_kwargs = {"dtype": dtype, "device": device}
|
||||
super().__init__()
|
||||
self.linear_1 = nn.Linear(
|
||||
in_features=in_channels,
|
||||
out_features=hidden_size,
|
||||
bias=True,
|
||||
**factory_kwargs
|
||||
)
|
||||
self.act_1 = act_layer()
|
||||
self.linear_2 = nn.Linear(
|
||||
in_features=hidden_size,
|
||||
out_features=hidden_size,
|
||||
bias=True,
|
||||
**factory_kwargs
|
||||
)
|
||||
|
||||
def forward(self, caption):
|
||||
hidden_states = self.linear_1(caption)
|
||||
hidden_states = self.act_1(hidden_states)
|
||||
hidden_states = self.linear_2(hidden_states)
|
||||
return hidden_states
|
||||
|
||||
|
||||
def timestep_embedding(t, dim, max_period=10000):
|
||||
"""
|
||||
Create sinusoidal timestep embeddings.
|
||||
|
||||
Args:
|
||||
t (torch.Tensor): a 1-D Tensor of N indices, one per batch element. These may be fractional.
|
||||
dim (int): the dimension of the output.
|
||||
max_period (int): controls the minimum frequency of the embeddings.
|
||||
|
||||
Returns:
|
||||
embedding (torch.Tensor): An (N, D) Tensor of positional embeddings.
|
||||
|
||||
.. ref_link: https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
|
||||
"""
|
||||
half = dim // 2
|
||||
freqs = torch.exp(
|
||||
-math.log(max_period)
|
||||
* torch.arange(start=0, end=half, dtype=torch.float32)
|
||||
/ half
|
||||
).to(device=t.device)
|
||||
args = t[:, None].float() * freqs[None]
|
||||
embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
|
||||
if dim % 2:
|
||||
embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
|
||||
return embedding
|
||||
|
||||
|
||||
class TimestepEmbedder(nn.Module):
|
||||
"""
|
||||
Embeds scalar timesteps into vector representations.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size,
|
||||
act_layer,
|
||||
frequency_embedding_size=256,
|
||||
max_period=10000,
|
||||
out_size=None,
|
||||
dtype=None,
|
||||
device=None,
|
||||
):
|
||||
factory_kwargs = {"dtype": dtype, "device": device}
|
||||
super().__init__()
|
||||
self.frequency_embedding_size = frequency_embedding_size
|
||||
self.max_period = max_period
|
||||
if out_size is None:
|
||||
out_size = hidden_size
|
||||
|
||||
self.mlp = nn.Sequential(
|
||||
nn.Linear(
|
||||
frequency_embedding_size, hidden_size, bias=True, **factory_kwargs
|
||||
),
|
||||
act_layer(),
|
||||
nn.Linear(hidden_size, out_size, bias=True, **factory_kwargs),
|
||||
)
|
||||
nn.init.normal_(self.mlp[0].weight, std=0.02)
|
||||
nn.init.normal_(self.mlp[2].weight, std=0.02)
|
||||
|
||||
def forward(self, t):
|
||||
t_freq = timestep_embedding(
|
||||
t, self.frequency_embedding_size, self.max_period
|
||||
).type(self.mlp[0].weight.dtype)
|
||||
t_emb = self.mlp(t_freq)
|
||||
return t_emb
|
||||
131
hyvideo/modules/mlp_layers.py
Normal file
131
hyvideo/modules/mlp_layers.py
Normal file
@@ -0,0 +1,131 @@
|
||||
# Modified from timm library:
|
||||
# https://github.com/huggingface/pytorch-image-models/blob/648aaa41233ba83eb38faf5ba9d415d574823241/timm/layers/mlp.py#L13
|
||||
|
||||
from functools import partial
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from .modulate_layers import modulate_
|
||||
from ..utils.helpers import to_2tuple
|
||||
|
||||
|
||||
class MLP(nn.Module):
|
||||
"""MLP as used in Vision Transformer, MLP-Mixer and related networks"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
hidden_channels=None,
|
||||
out_features=None,
|
||||
act_layer=nn.GELU,
|
||||
norm_layer=None,
|
||||
bias=True,
|
||||
drop=0.0,
|
||||
use_conv=False,
|
||||
device=None,
|
||||
dtype=None,
|
||||
):
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__()
|
||||
out_features = out_features or in_channels
|
||||
hidden_channels = hidden_channels or in_channels
|
||||
bias = to_2tuple(bias)
|
||||
drop_probs = to_2tuple(drop)
|
||||
linear_layer = partial(nn.Conv2d, kernel_size=1) if use_conv else nn.Linear
|
||||
|
||||
self.fc1 = linear_layer(
|
||||
in_channels, hidden_channels, bias=bias[0], **factory_kwargs
|
||||
)
|
||||
self.act = act_layer()
|
||||
self.drop1 = nn.Dropout(drop_probs[0])
|
||||
self.norm = (
|
||||
norm_layer(hidden_channels, **factory_kwargs)
|
||||
if norm_layer is not None
|
||||
else nn.Identity()
|
||||
)
|
||||
self.fc2 = linear_layer(
|
||||
hidden_channels, out_features, bias=bias[1], **factory_kwargs
|
||||
)
|
||||
self.drop2 = nn.Dropout(drop_probs[1])
|
||||
|
||||
def forward(self, x):
|
||||
x = self.fc1(x)
|
||||
x = self.act(x)
|
||||
x = self.drop1(x)
|
||||
x = self.norm(x)
|
||||
x = self.fc2(x)
|
||||
x = self.drop2(x)
|
||||
return x
|
||||
|
||||
def apply_(self, x, divide = 4):
|
||||
x_shape = x.shape
|
||||
x = x.view(-1, x.shape[-1])
|
||||
chunk_size = int(x_shape[1]/divide)
|
||||
x_chunks = torch.split(x, chunk_size)
|
||||
for i, x_chunk in enumerate(x_chunks):
|
||||
mlp_chunk = self.fc1(x_chunk)
|
||||
mlp_chunk = self.act(mlp_chunk)
|
||||
mlp_chunk = self.drop1(mlp_chunk)
|
||||
mlp_chunk = self.norm(mlp_chunk)
|
||||
mlp_chunk = self.fc2(mlp_chunk)
|
||||
x_chunk[...] = self.drop2(mlp_chunk)
|
||||
return x
|
||||
|
||||
#
|
||||
class MLPEmbedder(nn.Module):
|
||||
"""copied from https://github.com/black-forest-labs/flux/blob/main/src/flux/modules/layers.py"""
|
||||
def __init__(self, in_dim: int, hidden_dim: int, device=None, dtype=None):
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__()
|
||||
self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True, **factory_kwargs)
|
||||
self.silu = nn.SiLU()
|
||||
self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True, **factory_kwargs)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return self.out_layer(self.silu(self.in_layer(x)))
|
||||
|
||||
|
||||
class FinalLayer(nn.Module):
|
||||
"""The final layer of DiT."""
|
||||
|
||||
def __init__(
|
||||
self, hidden_size, patch_size, out_channels, act_layer, device=None, dtype=None
|
||||
):
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__()
|
||||
|
||||
# Just use LayerNorm for the final layer
|
||||
self.norm_final = nn.LayerNorm(
|
||||
hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs
|
||||
)
|
||||
if isinstance(patch_size, int):
|
||||
self.linear = nn.Linear(
|
||||
hidden_size,
|
||||
patch_size * patch_size * out_channels,
|
||||
bias=True,
|
||||
**factory_kwargs
|
||||
)
|
||||
else:
|
||||
self.linear = nn.Linear(
|
||||
hidden_size,
|
||||
patch_size[0] * patch_size[1] * patch_size[2] * out_channels,
|
||||
bias=True,
|
||||
)
|
||||
nn.init.zeros_(self.linear.weight)
|
||||
nn.init.zeros_(self.linear.bias)
|
||||
|
||||
# Here we don't distinguish between the modulate types. Just use the simple one.
|
||||
self.adaLN_modulation = nn.Sequential(
|
||||
act_layer(),
|
||||
nn.Linear(hidden_size, 2 * hidden_size, bias=True, **factory_kwargs),
|
||||
)
|
||||
# Zero-initialize the modulation
|
||||
nn.init.zeros_(self.adaLN_modulation[1].weight)
|
||||
nn.init.zeros_(self.adaLN_modulation[1].bias)
|
||||
|
||||
def forward(self, x, c):
|
||||
shift, scale = self.adaLN_modulation(c).chunk(2, dim=1)
|
||||
x = modulate_(self.norm_final(x), shift=shift, scale=scale)
|
||||
x = self.linear(x)
|
||||
return x
|
||||
1020
hyvideo/modules/models.py
Normal file
1020
hyvideo/modules/models.py
Normal file
File diff suppressed because it is too large
Load Diff
136
hyvideo/modules/modulate_layers.py
Normal file
136
hyvideo/modules/modulate_layers.py
Normal file
@@ -0,0 +1,136 @@
|
||||
from typing import Callable
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import math
|
||||
|
||||
class ModulateDiT(nn.Module):
|
||||
"""Modulation layer for DiT."""
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size: int,
|
||||
factor: int,
|
||||
act_layer: Callable,
|
||||
dtype=None,
|
||||
device=None,
|
||||
):
|
||||
factory_kwargs = {"dtype": dtype, "device": device}
|
||||
super().__init__()
|
||||
self.act = act_layer()
|
||||
self.linear = nn.Linear(
|
||||
hidden_size, factor * hidden_size, bias=True, **factory_kwargs
|
||||
)
|
||||
# Zero-initialize the modulation
|
||||
nn.init.zeros_(self.linear.weight)
|
||||
nn.init.zeros_(self.linear.bias)
|
||||
|
||||
def forward(self, x: torch.Tensor, condition_type=None, token_replace_vec=None) -> torch.Tensor:
|
||||
x_out = self.linear(self.act(x))
|
||||
|
||||
if condition_type == "token_replace":
|
||||
x_token_replace_out = self.linear(self.act(token_replace_vec))
|
||||
return x_out, x_token_replace_out
|
||||
else:
|
||||
return x_out
|
||||
|
||||
def modulate(x, shift=None, scale=None):
|
||||
"""modulate by shift and scale
|
||||
|
||||
Args:
|
||||
x (torch.Tensor): input tensor.
|
||||
shift (torch.Tensor, optional): shift tensor. Defaults to None.
|
||||
scale (torch.Tensor, optional): scale tensor. Defaults to None.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: the output tensor after modulate.
|
||||
"""
|
||||
if scale is None and shift is None:
|
||||
return x
|
||||
elif shift is None:
|
||||
return x * (1 + scale.unsqueeze(1))
|
||||
elif scale is None:
|
||||
return x + shift.unsqueeze(1)
|
||||
else:
|
||||
return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
|
||||
|
||||
def modulate_(x, shift=None, scale=None):
|
||||
|
||||
if scale is None and shift is None:
|
||||
return x
|
||||
elif shift is None:
|
||||
scale = scale + 1
|
||||
scale = scale.unsqueeze(1)
|
||||
return x.mul_(scale)
|
||||
elif scale is None:
|
||||
return x + shift.unsqueeze(1)
|
||||
else:
|
||||
scale = scale + 1
|
||||
scale = scale.unsqueeze(1)
|
||||
# return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
|
||||
torch.addcmul(shift.unsqueeze(1), x, scale, out =x )
|
||||
return x
|
||||
|
||||
def modulate(x, shift=None, scale=None, condition_type=None,
|
||||
tr_shift=None, tr_scale=None,
|
||||
frist_frame_token_num=None):
|
||||
if condition_type == "token_replace":
|
||||
x_zero = x[:, :frist_frame_token_num] * (1 + tr_scale.unsqueeze(1)) + tr_shift.unsqueeze(1)
|
||||
x_orig = x[:, frist_frame_token_num:] * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
|
||||
x = torch.concat((x_zero, x_orig), dim=1)
|
||||
return x
|
||||
else:
|
||||
if scale is None and shift is None:
|
||||
return x
|
||||
elif shift is None:
|
||||
return x * (1 + scale.unsqueeze(1))
|
||||
elif scale is None:
|
||||
return x + shift.unsqueeze(1)
|
||||
else:
|
||||
return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
|
||||
|
||||
def apply_gate(x, gate=None, tanh=False, condition_type=None, tr_gate=None, frist_frame_token_num=None):
|
||||
"""AI is creating summary for apply_gate
|
||||
|
||||
Args:
|
||||
x (torch.Tensor): input tensor.
|
||||
gate (torch.Tensor, optional): gate tensor. Defaults to None.
|
||||
tanh (bool, optional): whether to use tanh function. Defaults to False.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: the output tensor after apply gate.
|
||||
"""
|
||||
if condition_type == "token_replace":
|
||||
if gate is None:
|
||||
return x
|
||||
if tanh:
|
||||
x_zero = x[:, :frist_frame_token_num] * tr_gate.unsqueeze(1).tanh()
|
||||
x_orig = x[:, frist_frame_token_num:] * gate.unsqueeze(1).tanh()
|
||||
x = torch.concat((x_zero, x_orig), dim=1)
|
||||
return x
|
||||
else:
|
||||
x_zero = x[:, :frist_frame_token_num] * tr_gate.unsqueeze(1)
|
||||
x_orig = x[:, frist_frame_token_num:] * gate.unsqueeze(1)
|
||||
x = torch.concat((x_zero, x_orig), dim=1)
|
||||
return x
|
||||
else:
|
||||
if gate is None:
|
||||
return x
|
||||
if tanh:
|
||||
return x * gate.unsqueeze(1).tanh()
|
||||
else:
|
||||
return x * gate.unsqueeze(1)
|
||||
|
||||
def apply_gate_and_accumulate_(accumulator, x, gate=None, tanh=False):
|
||||
if gate is None:
|
||||
return accumulator
|
||||
if tanh:
|
||||
return accumulator.addcmul_(x, gate.unsqueeze(1).tanh())
|
||||
else:
|
||||
return accumulator.addcmul_(x, gate.unsqueeze(1))
|
||||
|
||||
def ckpt_wrapper(module):
|
||||
def ckpt_forward(*inputs):
|
||||
outputs = module(*inputs)
|
||||
return outputs
|
||||
|
||||
return ckpt_forward
|
||||
88
hyvideo/modules/norm_layers.py
Normal file
88
hyvideo/modules/norm_layers.py
Normal file
@@ -0,0 +1,88 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class RMSNorm(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):
|
||||
"""
|
||||
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
|
||||
|
||||
def apply_(self, x):
|
||||
y = x.pow(2).mean(-1, keepdim=True)
|
||||
y.add_(self.eps)
|
||||
y.rsqrt_()
|
||||
x.mul_(y)
|
||||
del y
|
||||
if hasattr(self, "weight"):
|
||||
x.mul_(self.weight)
|
||||
return x
|
||||
|
||||
|
||||
def get_norm_layer(norm_layer):
|
||||
"""
|
||||
Get the normalization layer.
|
||||
|
||||
Args:
|
||||
norm_layer (str): The type of normalization layer.
|
||||
|
||||
Returns:
|
||||
norm_layer (nn.Module): The normalization layer.
|
||||
"""
|
||||
if norm_layer == "layer":
|
||||
return nn.LayerNorm
|
||||
elif norm_layer == "rms":
|
||||
return RMSNorm
|
||||
else:
|
||||
raise NotImplementedError(f"Norm layer {norm_layer} is not implemented")
|
||||
760
hyvideo/modules/original models.py
Normal file
760
hyvideo/modules/original models.py
Normal file
@@ -0,0 +1,760 @@
|
||||
from typing import Any, List, Tuple, Optional, Union, Dict
|
||||
from einops import rearrange
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from diffusers.models import ModelMixin
|
||||
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
||||
|
||||
from .activation_layers import get_activation_layer
|
||||
from .norm_layers import get_norm_layer
|
||||
from .embed_layers import TimestepEmbedder, PatchEmbed, TextProjection
|
||||
from .attenion import attention, parallel_attention, get_cu_seqlens
|
||||
from .posemb_layers import apply_rotary_emb
|
||||
from .mlp_layers import MLP, MLPEmbedder, FinalLayer
|
||||
from .modulate_layers import ModulateDiT, modulate, apply_gate
|
||||
from .token_refiner import SingleTokenRefiner
|
||||
|
||||
|
||||
class MMDoubleStreamBlock(nn.Module):
|
||||
"""
|
||||
A multimodal dit block with seperate modulation for
|
||||
text and image/video, see more details (SD3): https://arxiv.org/abs/2403.03206
|
||||
(Flux.1): https://github.com/black-forest-labs/flux
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size: int,
|
||||
heads_num: int,
|
||||
mlp_width_ratio: float,
|
||||
mlp_act_type: str = "gelu_tanh",
|
||||
qk_norm: bool = True,
|
||||
qk_norm_type: str = "rms",
|
||||
qkv_bias: bool = False,
|
||||
dtype: Optional[torch.dtype] = None,
|
||||
device: Optional[torch.device] = None,
|
||||
):
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__()
|
||||
|
||||
self.deterministic = False
|
||||
self.heads_num = heads_num
|
||||
head_dim = hidden_size // heads_num
|
||||
mlp_hidden_dim = int(hidden_size * mlp_width_ratio)
|
||||
|
||||
self.img_mod = ModulateDiT(
|
||||
hidden_size,
|
||||
factor=6,
|
||||
act_layer=get_activation_layer("silu"),
|
||||
**factory_kwargs,
|
||||
)
|
||||
self.img_norm1 = nn.LayerNorm(
|
||||
hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs
|
||||
)
|
||||
|
||||
self.img_attn_qkv = nn.Linear(
|
||||
hidden_size, hidden_size * 3, bias=qkv_bias, **factory_kwargs
|
||||
)
|
||||
qk_norm_layer = get_norm_layer(qk_norm_type)
|
||||
self.img_attn_q_norm = (
|
||||
qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
|
||||
if qk_norm
|
||||
else nn.Identity()
|
||||
)
|
||||
self.img_attn_k_norm = (
|
||||
qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
|
||||
if qk_norm
|
||||
else nn.Identity()
|
||||
)
|
||||
self.img_attn_proj = nn.Linear(
|
||||
hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs
|
||||
)
|
||||
|
||||
self.img_norm2 = nn.LayerNorm(
|
||||
hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs
|
||||
)
|
||||
self.img_mlp = MLP(
|
||||
hidden_size,
|
||||
mlp_hidden_dim,
|
||||
act_layer=get_activation_layer(mlp_act_type),
|
||||
bias=True,
|
||||
**factory_kwargs,
|
||||
)
|
||||
|
||||
self.txt_mod = ModulateDiT(
|
||||
hidden_size,
|
||||
factor=6,
|
||||
act_layer=get_activation_layer("silu"),
|
||||
**factory_kwargs,
|
||||
)
|
||||
self.txt_norm1 = nn.LayerNorm(
|
||||
hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs
|
||||
)
|
||||
|
||||
self.txt_attn_qkv = nn.Linear(
|
||||
hidden_size, hidden_size * 3, bias=qkv_bias, **factory_kwargs
|
||||
)
|
||||
self.txt_attn_q_norm = (
|
||||
qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
|
||||
if qk_norm
|
||||
else nn.Identity()
|
||||
)
|
||||
self.txt_attn_k_norm = (
|
||||
qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
|
||||
if qk_norm
|
||||
else nn.Identity()
|
||||
)
|
||||
self.txt_attn_proj = nn.Linear(
|
||||
hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs
|
||||
)
|
||||
|
||||
self.txt_norm2 = nn.LayerNorm(
|
||||
hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs
|
||||
)
|
||||
self.txt_mlp = MLP(
|
||||
hidden_size,
|
||||
mlp_hidden_dim,
|
||||
act_layer=get_activation_layer(mlp_act_type),
|
||||
bias=True,
|
||||
**factory_kwargs,
|
||||
)
|
||||
self.hybrid_seq_parallel_attn = None
|
||||
|
||||
def enable_deterministic(self):
|
||||
self.deterministic = True
|
||||
|
||||
def disable_deterministic(self):
|
||||
self.deterministic = False
|
||||
|
||||
def forward(
|
||||
self,
|
||||
img: torch.Tensor,
|
||||
txt: torch.Tensor,
|
||||
vec: torch.Tensor,
|
||||
cu_seqlens_q: Optional[torch.Tensor] = None,
|
||||
cu_seqlens_kv: Optional[torch.Tensor] = None,
|
||||
max_seqlen_q: Optional[int] = None,
|
||||
max_seqlen_kv: Optional[int] = None,
|
||||
freqs_cis: tuple = None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
(
|
||||
img_mod1_shift,
|
||||
img_mod1_scale,
|
||||
img_mod1_gate,
|
||||
img_mod2_shift,
|
||||
img_mod2_scale,
|
||||
img_mod2_gate,
|
||||
) = self.img_mod(vec).chunk(6, dim=-1)
|
||||
(
|
||||
txt_mod1_shift,
|
||||
txt_mod1_scale,
|
||||
txt_mod1_gate,
|
||||
txt_mod2_shift,
|
||||
txt_mod2_scale,
|
||||
txt_mod2_gate,
|
||||
) = self.txt_mod(vec).chunk(6, dim=-1)
|
||||
|
||||
# Prepare image for attention.
|
||||
img_modulated = self.img_norm1(img)
|
||||
img_modulated = modulate(
|
||||
img_modulated, shift=img_mod1_shift, scale=img_mod1_scale
|
||||
)
|
||||
img_qkv = self.img_attn_qkv(img_modulated)
|
||||
img_q, img_k, img_v = rearrange(
|
||||
img_qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num
|
||||
)
|
||||
# Apply QK-Norm if needed
|
||||
img_q = self.img_attn_q_norm(img_q).to(img_v)
|
||||
img_k = self.img_attn_k_norm(img_k).to(img_v)
|
||||
|
||||
# Apply RoPE if needed.
|
||||
if freqs_cis is not None:
|
||||
img_qq, img_kk = apply_rotary_emb(img_q, img_k, freqs_cis, head_first=False)
|
||||
assert (
|
||||
img_qq.shape == img_q.shape and img_kk.shape == img_k.shape
|
||||
), f"img_kk: {img_qq.shape}, img_q: {img_q.shape}, img_kk: {img_kk.shape}, img_k: {img_k.shape}"
|
||||
img_q, img_k = img_qq, img_kk
|
||||
|
||||
# Prepare txt for attention.
|
||||
txt_modulated = self.txt_norm1(txt)
|
||||
txt_modulated = modulate(
|
||||
txt_modulated, shift=txt_mod1_shift, scale=txt_mod1_scale
|
||||
)
|
||||
txt_qkv = self.txt_attn_qkv(txt_modulated)
|
||||
txt_q, txt_k, txt_v = rearrange(
|
||||
txt_qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num
|
||||
)
|
||||
# Apply QK-Norm if needed.
|
||||
txt_q = self.txt_attn_q_norm(txt_q).to(txt_v)
|
||||
txt_k = self.txt_attn_k_norm(txt_k).to(txt_v)
|
||||
|
||||
# Run actual attention.
|
||||
q = torch.cat((img_q, txt_q), dim=1)
|
||||
k = torch.cat((img_k, txt_k), dim=1)
|
||||
v = torch.cat((img_v, txt_v), dim=1)
|
||||
assert (
|
||||
cu_seqlens_q.shape[0] == 2 * img.shape[0] + 1
|
||||
), f"cu_seqlens_q.shape:{cu_seqlens_q.shape}, img.shape[0]:{img.shape[0]}"
|
||||
|
||||
# attention computation start
|
||||
if not self.hybrid_seq_parallel_attn:
|
||||
attn = attention(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
cu_seqlens_q=cu_seqlens_q,
|
||||
cu_seqlens_kv=cu_seqlens_kv,
|
||||
max_seqlen_q=max_seqlen_q,
|
||||
max_seqlen_kv=max_seqlen_kv,
|
||||
batch_size=img_k.shape[0],
|
||||
)
|
||||
else:
|
||||
attn = parallel_attention(
|
||||
self.hybrid_seq_parallel_attn,
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
img_q_len=img_q.shape[1],
|
||||
img_kv_len=img_k.shape[1],
|
||||
cu_seqlens_q=cu_seqlens_q,
|
||||
cu_seqlens_kv=cu_seqlens_kv
|
||||
)
|
||||
|
||||
# attention computation end
|
||||
|
||||
img_attn, txt_attn = attn[:, : img.shape[1]], attn[:, img.shape[1] :]
|
||||
|
||||
# Calculate the img bloks.
|
||||
img = img + apply_gate(self.img_attn_proj(img_attn), gate=img_mod1_gate)
|
||||
img = img + apply_gate(
|
||||
self.img_mlp(
|
||||
modulate(
|
||||
self.img_norm2(img), shift=img_mod2_shift, scale=img_mod2_scale
|
||||
)
|
||||
),
|
||||
gate=img_mod2_gate,
|
||||
)
|
||||
|
||||
# Calculate the txt bloks.
|
||||
txt = txt + apply_gate(self.txt_attn_proj(txt_attn), gate=txt_mod1_gate)
|
||||
txt = txt + apply_gate(
|
||||
self.txt_mlp(
|
||||
modulate(
|
||||
self.txt_norm2(txt), shift=txt_mod2_shift, scale=txt_mod2_scale
|
||||
)
|
||||
),
|
||||
gate=txt_mod2_gate,
|
||||
)
|
||||
|
||||
return img, txt
|
||||
|
||||
|
||||
class MMSingleStreamBlock(nn.Module):
|
||||
"""
|
||||
A DiT block with parallel linear layers as described in
|
||||
https://arxiv.org/abs/2302.05442 and adapted modulation interface.
|
||||
Also refer to (SD3): https://arxiv.org/abs/2403.03206
|
||||
(Flux.1): https://github.com/black-forest-labs/flux
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size: int,
|
||||
heads_num: int,
|
||||
mlp_width_ratio: float = 4.0,
|
||||
mlp_act_type: str = "gelu_tanh",
|
||||
qk_norm: bool = True,
|
||||
qk_norm_type: str = "rms",
|
||||
qk_scale: float = None,
|
||||
dtype: Optional[torch.dtype] = None,
|
||||
device: Optional[torch.device] = None,
|
||||
):
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__()
|
||||
|
||||
self.deterministic = False
|
||||
self.hidden_size = hidden_size
|
||||
self.heads_num = heads_num
|
||||
head_dim = hidden_size // heads_num
|
||||
mlp_hidden_dim = int(hidden_size * mlp_width_ratio)
|
||||
self.mlp_hidden_dim = mlp_hidden_dim
|
||||
self.scale = qk_scale or head_dim ** -0.5
|
||||
|
||||
# qkv and mlp_in
|
||||
self.linear1 = nn.Linear(
|
||||
hidden_size, hidden_size * 3 + mlp_hidden_dim, **factory_kwargs
|
||||
)
|
||||
# proj and mlp_out
|
||||
self.linear2 = nn.Linear(
|
||||
hidden_size + mlp_hidden_dim, hidden_size, **factory_kwargs
|
||||
)
|
||||
|
||||
qk_norm_layer = get_norm_layer(qk_norm_type)
|
||||
self.q_norm = (
|
||||
qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
|
||||
if qk_norm
|
||||
else nn.Identity()
|
||||
)
|
||||
self.k_norm = (
|
||||
qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
|
||||
if qk_norm
|
||||
else nn.Identity()
|
||||
)
|
||||
|
||||
self.pre_norm = nn.LayerNorm(
|
||||
hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs
|
||||
)
|
||||
|
||||
self.mlp_act = get_activation_layer(mlp_act_type)()
|
||||
self.modulation = ModulateDiT(
|
||||
hidden_size,
|
||||
factor=3,
|
||||
act_layer=get_activation_layer("silu"),
|
||||
**factory_kwargs,
|
||||
)
|
||||
self.hybrid_seq_parallel_attn = None
|
||||
|
||||
def enable_deterministic(self):
|
||||
self.deterministic = True
|
||||
|
||||
def disable_deterministic(self):
|
||||
self.deterministic = False
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
vec: torch.Tensor,
|
||||
txt_len: int,
|
||||
cu_seqlens_q: Optional[torch.Tensor] = None,
|
||||
cu_seqlens_kv: Optional[torch.Tensor] = None,
|
||||
max_seqlen_q: Optional[int] = None,
|
||||
max_seqlen_kv: Optional[int] = None,
|
||||
freqs_cis: Tuple[torch.Tensor, torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
mod_shift, mod_scale, mod_gate = self.modulation(vec).chunk(3, dim=-1)
|
||||
x_mod = modulate(self.pre_norm(x), shift=mod_shift, scale=mod_scale)
|
||||
qkv, mlp = torch.split(
|
||||
self.linear1(x_mod), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1
|
||||
)
|
||||
|
||||
q, k, v = rearrange(qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num)
|
||||
|
||||
# Apply QK-Norm if needed.
|
||||
q = self.q_norm(q).to(v)
|
||||
k = self.k_norm(k).to(v)
|
||||
|
||||
# Apply RoPE if needed.
|
||||
if freqs_cis is not None:
|
||||
img_q, txt_q = q[:, :-txt_len, :, :], q[:, -txt_len:, :, :]
|
||||
img_k, txt_k = k[:, :-txt_len, :, :], k[:, -txt_len:, :, :]
|
||||
img_qq, img_kk = apply_rotary_emb(img_q, img_k, freqs_cis, head_first=False)
|
||||
assert (
|
||||
img_qq.shape == img_q.shape and img_kk.shape == img_k.shape
|
||||
), f"img_kk: {img_qq.shape}, img_q: {img_q.shape}, img_kk: {img_kk.shape}, img_k: {img_k.shape}"
|
||||
img_q, img_k = img_qq, img_kk
|
||||
q = torch.cat((img_q, txt_q), dim=1)
|
||||
k = torch.cat((img_k, txt_k), dim=1)
|
||||
|
||||
# Compute attention.
|
||||
assert (
|
||||
cu_seqlens_q.shape[0] == 2 * x.shape[0] + 1
|
||||
), f"cu_seqlens_q.shape:{cu_seqlens_q.shape}, x.shape[0]:{x.shape[0]}"
|
||||
|
||||
# attention computation start
|
||||
if not self.hybrid_seq_parallel_attn:
|
||||
attn = attention(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
cu_seqlens_q=cu_seqlens_q,
|
||||
cu_seqlens_kv=cu_seqlens_kv,
|
||||
max_seqlen_q=max_seqlen_q,
|
||||
max_seqlen_kv=max_seqlen_kv,
|
||||
batch_size=x.shape[0],
|
||||
)
|
||||
else:
|
||||
attn = parallel_attention(
|
||||
self.hybrid_seq_parallel_attn,
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
img_q_len=img_q.shape[1],
|
||||
img_kv_len=img_k.shape[1],
|
||||
cu_seqlens_q=cu_seqlens_q,
|
||||
cu_seqlens_kv=cu_seqlens_kv
|
||||
)
|
||||
# attention computation end
|
||||
|
||||
# Compute activation in mlp stream, cat again and run second linear layer.
|
||||
output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2))
|
||||
return x + apply_gate(output, gate=mod_gate)
|
||||
|
||||
|
||||
class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin):
|
||||
"""
|
||||
HunyuanVideo Transformer backbone
|
||||
|
||||
Inherited from ModelMixin and ConfigMixin for compatibility with diffusers' sampler StableDiffusionPipeline.
|
||||
|
||||
Reference:
|
||||
[1] Flux.1: https://github.com/black-forest-labs/flux
|
||||
[2] MMDiT: http://arxiv.org/abs/2403.03206
|
||||
|
||||
Parameters
|
||||
----------
|
||||
args: argparse.Namespace
|
||||
The arguments parsed by argparse.
|
||||
patch_size: list
|
||||
The size of the patch.
|
||||
in_channels: int
|
||||
The number of input channels.
|
||||
out_channels: int
|
||||
The number of output channels.
|
||||
hidden_size: int
|
||||
The hidden size of the transformer backbone.
|
||||
heads_num: int
|
||||
The number of attention heads.
|
||||
mlp_width_ratio: float
|
||||
The ratio of the hidden size of the MLP in the transformer block.
|
||||
mlp_act_type: str
|
||||
The activation function of the MLP in the transformer block.
|
||||
depth_double_blocks: int
|
||||
The number of transformer blocks in the double blocks.
|
||||
depth_single_blocks: int
|
||||
The number of transformer blocks in the single blocks.
|
||||
rope_dim_list: list
|
||||
The dimension of the rotary embedding for t, h, w.
|
||||
qkv_bias: bool
|
||||
Whether to use bias in the qkv linear layer.
|
||||
qk_norm: bool
|
||||
Whether to use qk norm.
|
||||
qk_norm_type: str
|
||||
The type of qk norm.
|
||||
guidance_embed: bool
|
||||
Whether to use guidance embedding for distillation.
|
||||
text_projection: str
|
||||
The type of the text projection, default is single_refiner.
|
||||
use_attention_mask: bool
|
||||
Whether to use attention mask for text encoder.
|
||||
dtype: torch.dtype
|
||||
The dtype of the model.
|
||||
device: torch.device
|
||||
The device of the model.
|
||||
"""
|
||||
|
||||
@register_to_config
|
||||
def __init__(
|
||||
self,
|
||||
args: Any,
|
||||
patch_size: list = [1, 2, 2],
|
||||
in_channels: int = 4, # Should be VAE.config.latent_channels.
|
||||
out_channels: int = None,
|
||||
hidden_size: int = 3072,
|
||||
heads_num: int = 24,
|
||||
mlp_width_ratio: float = 4.0,
|
||||
mlp_act_type: str = "gelu_tanh",
|
||||
mm_double_blocks_depth: int = 20,
|
||||
mm_single_blocks_depth: int = 40,
|
||||
rope_dim_list: List[int] = [16, 56, 56],
|
||||
qkv_bias: bool = True,
|
||||
qk_norm: bool = True,
|
||||
qk_norm_type: str = "rms",
|
||||
guidance_embed: bool = False, # For modulation.
|
||||
text_projection: str = "single_refiner",
|
||||
use_attention_mask: bool = True,
|
||||
dtype: Optional[torch.dtype] = None,
|
||||
device: Optional[torch.device] = None,
|
||||
):
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__()
|
||||
|
||||
self.patch_size = patch_size
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = in_channels if out_channels is None else out_channels
|
||||
self.unpatchify_channels = self.out_channels
|
||||
self.guidance_embed = guidance_embed
|
||||
self.rope_dim_list = rope_dim_list
|
||||
|
||||
# Text projection. Default to linear projection.
|
||||
# Alternative: TokenRefiner. See more details (LI-DiT): http://arxiv.org/abs/2406.11831
|
||||
self.use_attention_mask = use_attention_mask
|
||||
self.text_projection = text_projection
|
||||
|
||||
self.text_states_dim = args.text_states_dim
|
||||
self.text_states_dim_2 = args.text_states_dim_2
|
||||
|
||||
if hidden_size % heads_num != 0:
|
||||
raise ValueError(
|
||||
f"Hidden size {hidden_size} must be divisible by heads_num {heads_num}"
|
||||
)
|
||||
pe_dim = hidden_size // heads_num
|
||||
if sum(rope_dim_list) != pe_dim:
|
||||
raise ValueError(
|
||||
f"Got {rope_dim_list} but expected positional dim {pe_dim}"
|
||||
)
|
||||
self.hidden_size = hidden_size
|
||||
self.heads_num = heads_num
|
||||
|
||||
# image projection
|
||||
self.img_in = PatchEmbed(
|
||||
self.patch_size, self.in_channels, self.hidden_size, **factory_kwargs
|
||||
)
|
||||
|
||||
# text projection
|
||||
if self.text_projection == "linear":
|
||||
self.txt_in = TextProjection(
|
||||
self.text_states_dim,
|
||||
self.hidden_size,
|
||||
get_activation_layer("silu"),
|
||||
**factory_kwargs,
|
||||
)
|
||||
elif self.text_projection == "single_refiner":
|
||||
self.txt_in = SingleTokenRefiner(
|
||||
self.text_states_dim, hidden_size, heads_num, depth=2, **factory_kwargs
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f"Unsupported text_projection: {self.text_projection}"
|
||||
)
|
||||
|
||||
# time modulation
|
||||
self.time_in = TimestepEmbedder(
|
||||
self.hidden_size, get_activation_layer("silu"), **factory_kwargs
|
||||
)
|
||||
|
||||
# text modulation
|
||||
self.vector_in = MLPEmbedder(
|
||||
self.text_states_dim_2, self.hidden_size, **factory_kwargs
|
||||
)
|
||||
|
||||
# guidance modulation
|
||||
self.guidance_in = (
|
||||
TimestepEmbedder(
|
||||
self.hidden_size, get_activation_layer("silu"), **factory_kwargs
|
||||
)
|
||||
if guidance_embed
|
||||
else None
|
||||
)
|
||||
|
||||
# double blocks
|
||||
self.double_blocks = nn.ModuleList(
|
||||
[
|
||||
MMDoubleStreamBlock(
|
||||
self.hidden_size,
|
||||
self.heads_num,
|
||||
mlp_width_ratio=mlp_width_ratio,
|
||||
mlp_act_type=mlp_act_type,
|
||||
qk_norm=qk_norm,
|
||||
qk_norm_type=qk_norm_type,
|
||||
qkv_bias=qkv_bias,
|
||||
**factory_kwargs,
|
||||
)
|
||||
for _ in range(mm_double_blocks_depth)
|
||||
]
|
||||
)
|
||||
|
||||
# single blocks
|
||||
self.single_blocks = nn.ModuleList(
|
||||
[
|
||||
MMSingleStreamBlock(
|
||||
self.hidden_size,
|
||||
self.heads_num,
|
||||
mlp_width_ratio=mlp_width_ratio,
|
||||
mlp_act_type=mlp_act_type,
|
||||
qk_norm=qk_norm,
|
||||
qk_norm_type=qk_norm_type,
|
||||
**factory_kwargs,
|
||||
)
|
||||
for _ in range(mm_single_blocks_depth)
|
||||
]
|
||||
)
|
||||
|
||||
self.final_layer = FinalLayer(
|
||||
self.hidden_size,
|
||||
self.patch_size,
|
||||
self.out_channels,
|
||||
get_activation_layer("silu"),
|
||||
**factory_kwargs,
|
||||
)
|
||||
|
||||
def enable_deterministic(self):
|
||||
for block in self.double_blocks:
|
||||
block.enable_deterministic()
|
||||
for block in self.single_blocks:
|
||||
block.enable_deterministic()
|
||||
|
||||
def disable_deterministic(self):
|
||||
for block in self.double_blocks:
|
||||
block.disable_deterministic()
|
||||
for block in self.single_blocks:
|
||||
block.disable_deterministic()
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
t: torch.Tensor, # Should be in range(0, 1000).
|
||||
text_states: torch.Tensor = None,
|
||||
text_mask: torch.Tensor = None, # Now we don't use it.
|
||||
text_states_2: Optional[torch.Tensor] = None, # Text embedding for modulation.
|
||||
freqs_cos: Optional[torch.Tensor] = None,
|
||||
freqs_sin: Optional[torch.Tensor] = None,
|
||||
guidance: torch.Tensor = None, # Guidance for modulation, should be cfg_scale x 1000.
|
||||
return_dict: bool = True,
|
||||
) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
|
||||
out = {}
|
||||
img = x
|
||||
txt = text_states
|
||||
_, _, ot, oh, ow = x.shape
|
||||
tt, th, tw = (
|
||||
ot // self.patch_size[0],
|
||||
oh // self.patch_size[1],
|
||||
ow // self.patch_size[2],
|
||||
)
|
||||
|
||||
# Prepare modulation vectors.
|
||||
vec = self.time_in(t)
|
||||
|
||||
# text modulation
|
||||
vec = vec + self.vector_in(text_states_2)
|
||||
|
||||
# guidance modulation
|
||||
if self.guidance_embed:
|
||||
if guidance is None:
|
||||
raise ValueError(
|
||||
"Didn't get guidance strength for guidance distilled model."
|
||||
)
|
||||
|
||||
# our timestep_embedding is merged into guidance_in(TimestepEmbedder)
|
||||
vec = vec + self.guidance_in(guidance)
|
||||
|
||||
# Embed image and text.
|
||||
img = self.img_in(img)
|
||||
if self.text_projection == "linear":
|
||||
txt = self.txt_in(txt)
|
||||
elif self.text_projection == "single_refiner":
|
||||
txt = self.txt_in(txt, t, text_mask if self.use_attention_mask else None)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f"Unsupported text_projection: {self.text_projection}"
|
||||
)
|
||||
|
||||
txt_seq_len = txt.shape[1]
|
||||
img_seq_len = img.shape[1]
|
||||
|
||||
# Compute cu_squlens and max_seqlen for flash attention
|
||||
cu_seqlens_q = get_cu_seqlens(text_mask, img_seq_len)
|
||||
cu_seqlens_kv = cu_seqlens_q
|
||||
max_seqlen_q = img_seq_len + txt_seq_len
|
||||
max_seqlen_kv = max_seqlen_q
|
||||
|
||||
freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None
|
||||
# --------------------- Pass through DiT blocks ------------------------
|
||||
for _, block in enumerate(self.double_blocks):
|
||||
double_block_args = [
|
||||
img,
|
||||
txt,
|
||||
vec,
|
||||
cu_seqlens_q,
|
||||
cu_seqlens_kv,
|
||||
max_seqlen_q,
|
||||
max_seqlen_kv,
|
||||
freqs_cis,
|
||||
]
|
||||
|
||||
img, txt = block(*double_block_args)
|
||||
|
||||
# Merge txt and img to pass through single stream blocks.
|
||||
x = torch.cat((img, txt), 1)
|
||||
if len(self.single_blocks) > 0:
|
||||
for _, block in enumerate(self.single_blocks):
|
||||
single_block_args = [
|
||||
x,
|
||||
vec,
|
||||
txt_seq_len,
|
||||
cu_seqlens_q,
|
||||
cu_seqlens_kv,
|
||||
max_seqlen_q,
|
||||
max_seqlen_kv,
|
||||
(freqs_cos, freqs_sin),
|
||||
]
|
||||
|
||||
x = block(*single_block_args)
|
||||
|
||||
img = x[:, :img_seq_len, ...]
|
||||
|
||||
# ---------------------------- Final layer ------------------------------
|
||||
img = self.final_layer(img, vec) # (N, T, patch_size ** 2 * out_channels)
|
||||
|
||||
img = self.unpatchify(img, tt, th, tw)
|
||||
if return_dict:
|
||||
out["x"] = img
|
||||
return out
|
||||
return img
|
||||
|
||||
def unpatchify(self, x, t, h, w):
|
||||
"""
|
||||
x: (N, T, patch_size**2 * C)
|
||||
imgs: (N, H, W, C)
|
||||
"""
|
||||
c = self.unpatchify_channels
|
||||
pt, ph, pw = self.patch_size
|
||||
assert t * h * w == x.shape[1]
|
||||
|
||||
x = x.reshape(shape=(x.shape[0], t, h, w, c, pt, ph, pw))
|
||||
x = torch.einsum("nthwcopq->nctohpwq", x)
|
||||
imgs = x.reshape(shape=(x.shape[0], c, t * pt, h * ph, w * pw))
|
||||
|
||||
return imgs
|
||||
|
||||
def params_count(self):
|
||||
counts = {
|
||||
"double": sum(
|
||||
[
|
||||
sum(p.numel() for p in block.img_attn_qkv.parameters())
|
||||
+ sum(p.numel() for p in block.img_attn_proj.parameters())
|
||||
+ sum(p.numel() for p in block.img_mlp.parameters())
|
||||
+ sum(p.numel() for p in block.txt_attn_qkv.parameters())
|
||||
+ sum(p.numel() for p in block.txt_attn_proj.parameters())
|
||||
+ sum(p.numel() for p in block.txt_mlp.parameters())
|
||||
for block in self.double_blocks
|
||||
]
|
||||
),
|
||||
"single": sum(
|
||||
[
|
||||
sum(p.numel() for p in block.linear1.parameters())
|
||||
+ sum(p.numel() for p in block.linear2.parameters())
|
||||
for block in self.single_blocks
|
||||
]
|
||||
),
|
||||
"total": sum(p.numel() for p in self.parameters()),
|
||||
}
|
||||
counts["attn+mlp"] = counts["double"] + counts["single"]
|
||||
return counts
|
||||
|
||||
|
||||
#################################################################################
|
||||
# HunyuanVideo Configs #
|
||||
#################################################################################
|
||||
|
||||
HUNYUAN_VIDEO_CONFIG = {
|
||||
"HYVideo-T/2": {
|
||||
"mm_double_blocks_depth": 20,
|
||||
"mm_single_blocks_depth": 40,
|
||||
"rope_dim_list": [16, 56, 56],
|
||||
"hidden_size": 3072,
|
||||
"heads_num": 24,
|
||||
"mlp_width_ratio": 4,
|
||||
},
|
||||
"HYVideo-T/2-cfgdistill": {
|
||||
"mm_double_blocks_depth": 20,
|
||||
"mm_single_blocks_depth": 40,
|
||||
"rope_dim_list": [16, 56, 56],
|
||||
"hidden_size": 3072,
|
||||
"heads_num": 24,
|
||||
"mlp_width_ratio": 4,
|
||||
"guidance_embed": True,
|
||||
},
|
||||
}
|
||||
389
hyvideo/modules/placement.py
Normal file
389
hyvideo/modules/placement.py
Normal file
@@ -0,0 +1,389 @@
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
def hunyuan_token_reorder_to_token_major(tensor, fix_len, reorder_len, reorder_num_frame, frame_size):
|
||||
"""Reorder it from frame major to token major!"""
|
||||
assert reorder_len == reorder_num_frame * frame_size
|
||||
assert tensor.shape[2] == fix_len + reorder_len
|
||||
|
||||
tensor[:, :, :-fix_len, :] = tensor[:, :, :-fix_len:, :].reshape(tensor.shape[0], tensor.shape[1], reorder_num_frame, frame_size, tensor.shape[3]) \
|
||||
.transpose(2, 3).reshape(tensor.shape[0], tensor.shape[1], reorder_len, tensor.shape[3])
|
||||
return tensor
|
||||
|
||||
def hunyuan_token_reorder_to_frame_major(tensor, fix_len, reorder_len, reorder_num_frame, frame_size):
|
||||
"""Reorder it from token major to frame major!"""
|
||||
assert reorder_len == reorder_num_frame * frame_size
|
||||
assert tensor.shape[2] == fix_len + reorder_len
|
||||
|
||||
tensor[:, :, :-fix_len:, :] = tensor[:, :, :-fix_len:, :].reshape(tensor.shape[0], tensor.shape[1], frame_size, reorder_num_frame, tensor.shape[3]) \
|
||||
.transpose(2, 3).reshape(tensor.shape[0], tensor.shape[1], reorder_len, tensor.shape[3])
|
||||
return tensor
|
||||
|
||||
|
||||
@triton.jit
|
||||
def hunyuan_sparse_head_placement_kernel(
|
||||
query_ptr, key_ptr, value_ptr, # [cfg, num_heads, seq_len, head_dim] seq_len = context_length + num_frame * frame_size
|
||||
query_out_ptr, key_out_ptr, value_out_ptr, # [cfg, num_heads, seq_len, head_dim]
|
||||
best_mask_idx_ptr, # [cfg, num_heads]
|
||||
query_stride_b, query_stride_h, query_stride_s, query_stride_d,
|
||||
mask_idx_stride_b, mask_idx_stride_h,
|
||||
seq_len: tl.constexpr,
|
||||
head_dim: tl.constexpr,
|
||||
context_length: tl.constexpr,
|
||||
num_frame: tl.constexpr,
|
||||
frame_size: tl.constexpr,
|
||||
BLOCK_SIZE: tl.constexpr
|
||||
):
|
||||
# Copy query, key, value to output
|
||||
# range: [b, h, block_id * block_size: block_id * block_size + block_size, :]
|
||||
cfg = tl.program_id(0)
|
||||
head = tl.program_id(1)
|
||||
block_id = tl.program_id(2)
|
||||
|
||||
start_id = block_id * BLOCK_SIZE
|
||||
end_id = start_id + BLOCK_SIZE
|
||||
end_id = tl.where(end_id > seq_len, seq_len, end_id)
|
||||
|
||||
# Load best mask idx (0 is spatial, 1 is temporal)
|
||||
is_temporal = tl.load(best_mask_idx_ptr + cfg * mask_idx_stride_b + head * mask_idx_stride_h)
|
||||
|
||||
offset_token = tl.arange(0, BLOCK_SIZE) + start_id
|
||||
offset_mask = offset_token < seq_len
|
||||
offset_d = tl.arange(0, head_dim)
|
||||
|
||||
if is_temporal:
|
||||
frame_id = offset_token // frame_size
|
||||
patch_id = offset_token - frame_id * frame_size
|
||||
offset_store_token = tl.where(offset_token >= seq_len - context_length, offset_token, patch_id * num_frame + frame_id)
|
||||
|
||||
offset_load = (cfg * query_stride_b + head * query_stride_h + offset_token[:,None] * query_stride_s) + offset_d[None,:] * query_stride_d
|
||||
offset_query = query_ptr + offset_load
|
||||
offset_key = key_ptr + offset_load
|
||||
offset_value = value_ptr + offset_load
|
||||
|
||||
offset_store = (cfg * query_stride_b + head * query_stride_h + offset_store_token[:,None] * query_stride_s) + offset_d[None,:] * query_stride_d
|
||||
offset_query_out = query_out_ptr + offset_store
|
||||
offset_key_out = key_out_ptr + offset_store
|
||||
offset_value_out = value_out_ptr + offset_store
|
||||
|
||||
# Maybe tune the pipeline here
|
||||
query = tl.load(offset_query, mask=offset_mask[:,None])
|
||||
tl.store(offset_query_out, query, mask=offset_mask[:,None])
|
||||
key = tl.load(offset_key, mask=offset_mask[:,None])
|
||||
tl.store(offset_key_out, key, mask=offset_mask[:,None])
|
||||
value = tl.load(offset_value, mask=offset_mask[:,None])
|
||||
tl.store(offset_value_out, value, mask=offset_mask[:,None])
|
||||
|
||||
|
||||
else:
|
||||
offset_load = (cfg * query_stride_b + head * query_stride_h + offset_token[:,None] * query_stride_s) + offset_d[None,:] * query_stride_d
|
||||
offset_query = query_ptr + offset_load
|
||||
offset_key = key_ptr + offset_load
|
||||
offset_value = value_ptr + offset_load
|
||||
|
||||
offset_store = offset_load
|
||||
offset_query_out = query_out_ptr + offset_store
|
||||
offset_key_out = key_out_ptr + offset_store
|
||||
offset_value_out = value_out_ptr + offset_store
|
||||
|
||||
# Maybe tune the pipeline here
|
||||
query = tl.load(offset_query, mask=offset_mask[:,None])
|
||||
tl.store(offset_query_out, query, mask=offset_mask[:,None])
|
||||
key = tl.load(offset_key, mask=offset_mask[:,None])
|
||||
tl.store(offset_key_out, key, mask=offset_mask[:,None])
|
||||
value = tl.load(offset_value, mask=offset_mask[:,None])
|
||||
tl.store(offset_value_out, value, mask=offset_mask[:,None])
|
||||
|
||||
|
||||
def hunyuan_sparse_head_placement(query, key, value, query_out, key_out, value_out, best_mask_idx, context_length, num_frame, frame_size):
|
||||
cfg, num_heads, seq_len, head_dim = query.shape
|
||||
BLOCK_SIZE = 128
|
||||
assert seq_len == context_length + num_frame * frame_size
|
||||
|
||||
grid = (cfg, num_heads, (seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE)
|
||||
|
||||
hunyuan_sparse_head_placement_kernel[grid](
|
||||
query, key, value,
|
||||
query_out, key_out, value_out,
|
||||
best_mask_idx,
|
||||
query.stride(0), query.stride(1), query.stride(2), query.stride(3),
|
||||
best_mask_idx.stride(0), best_mask_idx.stride(1),
|
||||
seq_len, head_dim, context_length, num_frame, frame_size,
|
||||
BLOCK_SIZE
|
||||
)
|
||||
|
||||
|
||||
def ref_hunyuan_sparse_head_placement(query, key, value, best_mask_idx, context_length, num_frame, frame_size):
|
||||
cfg, num_heads, seq_len, head_dim = query.shape
|
||||
assert seq_len == context_length + num_frame * frame_size
|
||||
|
||||
query_out = query.clone()
|
||||
key_out = key.clone()
|
||||
value_out = value.clone()
|
||||
|
||||
# Spatial
|
||||
query_out[best_mask_idx == 0], key_out[best_mask_idx == 0], value_out[best_mask_idx == 0] = \
|
||||
query[best_mask_idx == 0], key[best_mask_idx == 0], value[best_mask_idx == 0]
|
||||
|
||||
# Temporal
|
||||
query_out[best_mask_idx == 1], key_out[best_mask_idx == 1], value_out[best_mask_idx == 1] = \
|
||||
hunyuan_token_reorder_to_token_major(query[best_mask_idx == 1].unsqueeze(0), context_length, num_frame * frame_size, num_frame, frame_size).squeeze(0), \
|
||||
hunyuan_token_reorder_to_token_major(key[best_mask_idx == 1].unsqueeze(0), context_length, num_frame * frame_size, num_frame, frame_size).squeeze(0), \
|
||||
hunyuan_token_reorder_to_token_major(value[best_mask_idx == 1].unsqueeze(0), context_length, num_frame * frame_size, num_frame, frame_size).squeeze(0)
|
||||
|
||||
return query_out, key_out, value_out
|
||||
|
||||
|
||||
def test_hunyuan_sparse_head_placement():
|
||||
|
||||
context_length = 226
|
||||
num_frame = 11
|
||||
frame_size = 4080
|
||||
|
||||
cfg = 2
|
||||
num_heads = 48
|
||||
|
||||
seq_len = context_length + num_frame * frame_size
|
||||
head_dim = 64
|
||||
|
||||
dtype = torch.bfloat16
|
||||
device = torch.device("cuda")
|
||||
|
||||
query = torch.randn(cfg, num_heads, seq_len, head_dim, dtype=dtype, device=device)
|
||||
key = torch.randn(cfg, num_heads, seq_len, head_dim, dtype=dtype, device=device)
|
||||
value = torch.randn(cfg, num_heads, seq_len, head_dim, dtype=dtype, device=device)
|
||||
|
||||
best_mask_idx = torch.randint(0, 2, (cfg, num_heads), device=device)
|
||||
|
||||
query_out = torch.empty_like(query)
|
||||
key_out = torch.empty_like(key)
|
||||
value_out = torch.empty_like(value)
|
||||
|
||||
hunyuan_sparse_head_placement(query, key, value, query_out, key_out, value_out, best_mask_idx, context_length, num_frame, frame_size)
|
||||
ref_query_out, ref_key_out, ref_value_out = ref_hunyuan_sparse_head_placement(query, key, value, best_mask_idx, context_length, num_frame, frame_size)
|
||||
|
||||
torch.testing.assert_close(query_out, ref_query_out)
|
||||
torch.testing.assert_close(key_out, ref_key_out)
|
||||
torch.testing.assert_close(value_out, ref_value_out)
|
||||
|
||||
|
||||
def benchmark_hunyuan_sparse_head_placement():
|
||||
import time
|
||||
|
||||
context_length = 226
|
||||
num_frame = 11
|
||||
frame_size = 4080
|
||||
|
||||
cfg = 2
|
||||
num_heads = 48
|
||||
|
||||
seq_len = context_length + num_frame * frame_size
|
||||
head_dim = 64
|
||||
|
||||
dtype = torch.bfloat16
|
||||
device = torch.device("cuda")
|
||||
|
||||
query = torch.randn(cfg, num_heads, seq_len, head_dim, dtype=dtype, device=device)
|
||||
key = torch.randn(cfg, num_heads, seq_len, head_dim, dtype=dtype, device=device)
|
||||
value = torch.randn(cfg, num_heads, seq_len, head_dim, dtype=dtype, device=device)
|
||||
best_mask_idx = torch.randint(0, 2, (cfg, num_heads), device=device)
|
||||
|
||||
query_out = torch.empty_like(query)
|
||||
key_out = torch.empty_like(key)
|
||||
value_out = torch.empty_like(value)
|
||||
|
||||
warmup = 10
|
||||
all_iter = 1000
|
||||
|
||||
# warmup
|
||||
for _ in range(warmup):
|
||||
hunyuan_sparse_head_placement(query, key, value, query_out, key_out, value_out, best_mask_idx, context_length, num_frame, frame_size)
|
||||
|
||||
torch.cuda.synchronize()
|
||||
start = time.time()
|
||||
for _ in range(all_iter):
|
||||
hunyuan_sparse_head_placement(query, key, value, query_out, key_out, value_out, best_mask_idx, context_length, num_frame, frame_size)
|
||||
torch.cuda.synchronize()
|
||||
end = time.time()
|
||||
|
||||
print(f"Triton Elapsed Time: {(end - start) / all_iter * 1e3:.2f} ms")
|
||||
print(f"Triton Total Bandwidth: {query.nelement() * query.element_size() * 3 * 2 * all_iter / (end - start) / 1e9:.2f} GB/s")
|
||||
|
||||
torch.cuda.synchronize()
|
||||
start = time.time()
|
||||
for _ in range(all_iter):
|
||||
ref_hunyuan_sparse_head_placement(query, key, value, best_mask_idx, context_length, num_frame, frame_size)
|
||||
torch.cuda.synchronize()
|
||||
end = time.time()
|
||||
|
||||
print(f"Reference Elapsed Time: {(end - start) / all_iter * 1e3:.2f} ms")
|
||||
print(f"Reference Total Bandwidth: {query.nelement() * query.element_size() * 3 * 2 * all_iter / (end - start) / 1e9:.2f} GB/s")
|
||||
|
||||
|
||||
@triton.jit
|
||||
def hunyuan_hidden_states_placement_kernel(
|
||||
hidden_states_ptr, # [cfg, num_heads, seq_len, head_dim] seq_len = context_length + num_frame * frame_size
|
||||
hidden_states_out_ptr, # [cfg, num_heads, seq_len, head_dim]
|
||||
best_mask_idx_ptr, # [cfg, num_heads]
|
||||
hidden_states_stride_b, hidden_states_stride_h, hidden_states_stride_s, hidden_states_stride_d,
|
||||
mask_idx_stride_b, mask_idx_stride_h,
|
||||
seq_len: tl.constexpr,
|
||||
head_dim: tl.constexpr,
|
||||
context_length: tl.constexpr,
|
||||
num_frame: tl.constexpr,
|
||||
frame_size: tl.constexpr,
|
||||
BLOCK_SIZE: tl.constexpr
|
||||
):
|
||||
# Copy hidden_states to output
|
||||
# range: [b, h, block_id * block_size: block_id * block_size + block_size, :]
|
||||
cfg = tl.program_id(0)
|
||||
head = tl.program_id(1)
|
||||
block_id = tl.program_id(2)
|
||||
|
||||
start_id = block_id * BLOCK_SIZE
|
||||
end_id = start_id + BLOCK_SIZE
|
||||
end_id = tl.where(end_id > seq_len, seq_len, end_id)
|
||||
|
||||
# Load best mask idx (0 is spatial, 1 is temporal)
|
||||
is_temporal = tl.load(best_mask_idx_ptr + cfg * mask_idx_stride_b + head * mask_idx_stride_h)
|
||||
|
||||
offset_token = tl.arange(0, BLOCK_SIZE) + start_id
|
||||
offset_mask = offset_token < seq_len
|
||||
offset_d = tl.arange(0, head_dim)
|
||||
|
||||
if is_temporal:
|
||||
patch_id = offset_token // num_frame
|
||||
frame_id = offset_token - patch_id * num_frame
|
||||
offset_store_token = tl.where(offset_token >= seq_len - context_length, offset_token, frame_id * frame_size + patch_id)
|
||||
|
||||
offset_load = (cfg * hidden_states_stride_b + head * hidden_states_stride_h + offset_token[:,None] * hidden_states_stride_s) + offset_d[None,:] * hidden_states_stride_d
|
||||
offset_hidden_states = hidden_states_ptr + offset_load
|
||||
|
||||
offset_store = (cfg * hidden_states_stride_b + head * hidden_states_stride_h + offset_store_token[:,None] * hidden_states_stride_s) + offset_d[None,:] * hidden_states_stride_d
|
||||
offset_hidden_states_out = hidden_states_out_ptr + offset_store
|
||||
|
||||
# Maybe tune the pipeline here
|
||||
hidden_states = tl.load(offset_hidden_states, mask=offset_mask[:,None])
|
||||
tl.store(offset_hidden_states_out, hidden_states, mask=offset_mask[:,None])
|
||||
else:
|
||||
offset_load = (cfg * hidden_states_stride_b + head * hidden_states_stride_h + offset_token[:,None] * hidden_states_stride_s) + offset_d[None,:] * hidden_states_stride_d
|
||||
offset_hidden_states = hidden_states_ptr + offset_load
|
||||
|
||||
offset_store = offset_load
|
||||
offset_hidden_states_out = hidden_states_out_ptr + offset_store
|
||||
|
||||
# Maybe tune the pipeline here
|
||||
hidden_states = tl.load(offset_hidden_states, mask=offset_mask[:,None])
|
||||
tl.store(offset_hidden_states_out, hidden_states, mask=offset_mask[:,None])
|
||||
|
||||
|
||||
def hunyuan_hidden_states_placement(hidden_states, hidden_states_out, best_mask_idx, context_length, num_frame, frame_size):
|
||||
cfg, num_heads, seq_len, head_dim = hidden_states.shape
|
||||
BLOCK_SIZE = 128
|
||||
assert seq_len == context_length + num_frame * frame_size
|
||||
|
||||
grid = (cfg, num_heads, (seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE)
|
||||
|
||||
|
||||
hunyuan_hidden_states_placement_kernel[grid](
|
||||
hidden_states,
|
||||
hidden_states_out,
|
||||
best_mask_idx,
|
||||
hidden_states.stride(0), hidden_states.stride(1), hidden_states.stride(2), hidden_states.stride(3),
|
||||
best_mask_idx.stride(0), best_mask_idx.stride(1),
|
||||
seq_len, head_dim, context_length, num_frame, frame_size,
|
||||
BLOCK_SIZE
|
||||
)
|
||||
|
||||
return hidden_states_out
|
||||
|
||||
def ref_hunyuan_hidden_states_placement(hidden_states, output_hidden_states, best_mask_idx, context_length, num_frame, frame_size):
|
||||
cfg, num_heads, seq_len, head_dim = hidden_states.shape
|
||||
assert seq_len == context_length + num_frame * frame_size
|
||||
|
||||
# Spatial
|
||||
output_hidden_states[best_mask_idx == 0] = hidden_states[best_mask_idx == 0]
|
||||
# Temporal
|
||||
output_hidden_states[best_mask_idx == 1] = hunyuan_token_reorder_to_frame_major(hidden_states[best_mask_idx == 1].unsqueeze(0), context_length, num_frame * frame_size, num_frame, frame_size).squeeze(0)
|
||||
|
||||
def test_hunyuan_hidden_states_placement():
|
||||
|
||||
context_length = 226
|
||||
num_frame = 11
|
||||
frame_size = 4080
|
||||
|
||||
cfg = 2
|
||||
num_heads = 48
|
||||
|
||||
seq_len = context_length + num_frame * frame_size
|
||||
head_dim = 64
|
||||
|
||||
dtype = torch.bfloat16
|
||||
device = torch.device("cuda")
|
||||
|
||||
hidden_states = torch.randn(cfg, num_heads, seq_len, head_dim, dtype=dtype, device=device)
|
||||
best_mask_idx = torch.randint(0, 2, (cfg, num_heads), device=device)
|
||||
|
||||
hidden_states_out1 = torch.empty_like(hidden_states)
|
||||
hidden_states_out2 = torch.empty_like(hidden_states)
|
||||
|
||||
hunyuan_hidden_states_placement(hidden_states, hidden_states_out1, best_mask_idx, context_length, num_frame, frame_size)
|
||||
ref_hunyuan_hidden_states_placement(hidden_states, hidden_states_out2, best_mask_idx, context_length, num_frame, frame_size)
|
||||
|
||||
torch.testing.assert_close(hidden_states_out1, hidden_states_out2)
|
||||
|
||||
def benchmark_hunyuan_hidden_states_placement():
|
||||
import time
|
||||
|
||||
context_length = 226
|
||||
num_frame = 11
|
||||
frame_size = 4080
|
||||
|
||||
cfg = 2
|
||||
num_heads = 48
|
||||
|
||||
seq_len = context_length + num_frame * frame_size
|
||||
head_dim = 64
|
||||
|
||||
dtype = torch.bfloat16
|
||||
device = torch.device("cuda")
|
||||
|
||||
hidden_states = torch.randn(cfg, num_heads, seq_len, head_dim, dtype=dtype, device=device)
|
||||
best_mask_idx = torch.randint(0, 2, (cfg, num_heads), device=device)
|
||||
|
||||
hidden_states_out = torch.empty_like(hidden_states)
|
||||
|
||||
warmup = 10
|
||||
all_iter = 1000
|
||||
|
||||
# warmup
|
||||
for _ in range(warmup):
|
||||
hunyuan_hidden_states_placement(hidden_states, hidden_states_out, best_mask_idx, context_length, num_frame, frame_size)
|
||||
|
||||
torch.cuda.synchronize()
|
||||
start = time.time()
|
||||
for _ in range(all_iter):
|
||||
hunyuan_hidden_states_placement(hidden_states, hidden_states_out, best_mask_idx, context_length, num_frame, frame_size)
|
||||
torch.cuda.synchronize()
|
||||
end = time.time()
|
||||
|
||||
print(f"Triton Elapsed Time: {(end - start) / all_iter * 1e3:.2f} ms")
|
||||
print(f"Triton Total Bandwidth: {hidden_states.nelement() * hidden_states.element_size() * 2 * all_iter / (end - start) / 1e9:.2f} GB/s")
|
||||
|
||||
torch.cuda.synchronize()
|
||||
start = time.time()
|
||||
for _ in range(all_iter):
|
||||
ref_hunyuan_hidden_states_placement(hidden_states, hidden_states.clone(), best_mask_idx, context_length, num_frame, frame_size)
|
||||
torch.cuda.synchronize()
|
||||
end = time.time()
|
||||
|
||||
print(f"Reference Elapsed Time: {(end - start) / all_iter * 1e3:.2f} ms")
|
||||
print(f"Reference Total Bandwidth: {hidden_states.nelement() * hidden_states.element_size() * 2 * all_iter / (end - start) / 1e9:.2f} GB/s")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_hunyuan_sparse_head_placement()
|
||||
benchmark_hunyuan_sparse_head_placement()
|
||||
test_hunyuan_hidden_states_placement()
|
||||
benchmark_hunyuan_hidden_states_placement()
|
||||
475
hyvideo/modules/posemb_layers.py
Normal file
475
hyvideo/modules/posemb_layers.py
Normal file
@@ -0,0 +1,475 @@
|
||||
import torch
|
||||
from typing import Union, Tuple, List, Optional
|
||||
import numpy as np
|
||||
|
||||
|
||||
###### Thanks to the RifleX project (https://github.com/thu-ml/RIFLEx/) for this alternative pos embed for long videos
|
||||
#
|
||||
def get_1d_rotary_pos_embed_riflex(
|
||||
dim: int,
|
||||
pos: Union[np.ndarray, int],
|
||||
theta: float = 10000.0,
|
||||
use_real=False,
|
||||
k: Optional[int] = None,
|
||||
L_test: Optional[int] = None,
|
||||
):
|
||||
"""
|
||||
RIFLEx: Precompute the frequency tensor for complex exponentials (cis) with given dimensions.
|
||||
|
||||
This function calculates a frequency tensor with complex exponentials using the given dimension 'dim' and the end
|
||||
index 'end'. The 'theta' parameter scales the frequencies. The returned tensor contains complex values in complex64
|
||||
data type.
|
||||
|
||||
Args:
|
||||
dim (`int`): Dimension of the frequency tensor.
|
||||
pos (`np.ndarray` or `int`): Position indices for the frequency tensor. [S] or scalar
|
||||
theta (`float`, *optional*, defaults to 10000.0):
|
||||
Scaling factor for frequency computation. Defaults to 10000.0.
|
||||
use_real (`bool`, *optional*):
|
||||
If True, return real part and imaginary part separately. Otherwise, return complex numbers.
|
||||
k (`int`, *optional*, defaults to None): the index for the intrinsic frequency in RoPE
|
||||
L_test (`int`, *optional*, defaults to None): the number of frames for inference
|
||||
Returns:
|
||||
`torch.Tensor`: Precomputed frequency tensor with complex exponentials. [S, D/2]
|
||||
"""
|
||||
assert dim % 2 == 0
|
||||
|
||||
if isinstance(pos, int):
|
||||
pos = torch.arange(pos)
|
||||
if isinstance(pos, np.ndarray):
|
||||
pos = torch.from_numpy(pos) # type: ignore # [S]
|
||||
|
||||
freqs = 1.0 / (
|
||||
theta ** (torch.arange(0, dim, 2, device=pos.device)[: (dim // 2)].float() / dim)
|
||||
) # [D/2]
|
||||
|
||||
# === Riflex modification start ===
|
||||
# Reduce the intrinsic frequency to stay within a single period after extrapolation (see Eq. (8)).
|
||||
# Empirical observations show that a few videos may exhibit repetition in the tail frames.
|
||||
# To be conservative, we multiply by 0.9 to keep the extrapolated length below 90% of a single period.
|
||||
if k is not None:
|
||||
freqs[k-1] = 0.9 * 2 * torch.pi / L_test
|
||||
# === Riflex modification end ===
|
||||
|
||||
freqs = torch.outer(pos, freqs) # type: ignore # [S, D/2]
|
||||
if use_real:
|
||||
freqs_cos = freqs.cos().repeat_interleave(2, dim=1).float() # [S, D]
|
||||
freqs_sin = freqs.sin().repeat_interleave(2, dim=1).float() # [S, D]
|
||||
return freqs_cos, freqs_sin
|
||||
else:
|
||||
# lumina
|
||||
freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64 # [S, D/2]
|
||||
return freqs_cis
|
||||
|
||||
def identify_k( b: float, d: int, N: int):
|
||||
"""
|
||||
This function identifies the index of the intrinsic frequency component in a RoPE-based pre-trained diffusion transformer.
|
||||
|
||||
Args:
|
||||
b (`float`): The base frequency for RoPE.
|
||||
d (`int`): Dimension of the frequency tensor
|
||||
N (`int`): the first observed repetition frame in latent space
|
||||
Returns:
|
||||
k (`int`): the index of intrinsic frequency component
|
||||
N_k (`int`): the period of intrinsic frequency component in latent space
|
||||
Example:
|
||||
In HunyuanVideo, b=256 and d=16, the repetition occurs approximately 8s (N=48 in latent space).
|
||||
k, N_k = identify_k(b=256, d=16, N=48)
|
||||
In this case, the intrinsic frequency index k is 4, and the period N_k is 50.
|
||||
"""
|
||||
|
||||
# Compute the period of each frequency in RoPE according to Eq.(4)
|
||||
periods = []
|
||||
for j in range(1, d // 2 + 1):
|
||||
theta_j = 1.0 / (b ** (2 * (j - 1) / d))
|
||||
N_j = round(2 * torch.pi / theta_j)
|
||||
periods.append(N_j)
|
||||
|
||||
# Identify the intrinsic frequency whose period is closed to N（see Eq.(7)）
|
||||
diffs = [abs(N_j - N) for N_j in periods]
|
||||
k = diffs.index(min(diffs)) + 1
|
||||
N_k = periods[k-1]
|
||||
return k, N_k
|
||||
|
||||
def _to_tuple(x, dim=2):
|
||||
if isinstance(x, int):
|
||||
return (x,) * dim
|
||||
elif len(x) == dim:
|
||||
return x
|
||||
else:
|
||||
raise ValueError(f"Expected length {dim} or int, but got {x}")
|
||||
|
||||
|
||||
def get_meshgrid_nd(start, *args, dim=2):
|
||||
"""
|
||||
Get n-D meshgrid with start, stop and num.
|
||||
|
||||
Args:
|
||||
start (int or tuple): If len(args) == 0, start is num; If len(args) == 1, start is start, args[0] is stop,
|
||||
step is 1; If len(args) == 2, start is start, args[0] is stop, args[1] is num. For n-dim, start/stop/num
|
||||
should be int or n-tuple. If n-tuple is provided, the meshgrid will be stacked following the dim order in
|
||||
n-tuples.
|
||||
*args: See above.
|
||||
dim (int): Dimension of the meshgrid. Defaults to 2.
|
||||
|
||||
Returns:
|
||||
grid (np.ndarray): [dim, ...]
|
||||
"""
|
||||
if len(args) == 0:
|
||||
# start is grid_size
|
||||
num = _to_tuple(start, dim=dim)
|
||||
start = (0,) * dim
|
||||
stop = num
|
||||
elif len(args) == 1:
|
||||
# start is start, args[0] is stop, step is 1
|
||||
start = _to_tuple(start, dim=dim)
|
||||
stop = _to_tuple(args[0], dim=dim)
|
||||
num = [stop[i] - start[i] for i in range(dim)]
|
||||
elif len(args) == 2:
|
||||
# start is start, args[0] is stop, args[1] is num
|
||||
start = _to_tuple(start, dim=dim) # Left-Top eg: 12,0
|
||||
stop = _to_tuple(args[0], dim=dim) # Right-Bottom eg: 20,32
|
||||
num = _to_tuple(args[1], dim=dim) # Target Size eg: 32,124
|
||||
else:
|
||||
raise ValueError(f"len(args) should be 0, 1 or 2, but got {len(args)}")
|
||||
|
||||
# PyTorch implement of np.linspace(start[i], stop[i], num[i], endpoint=False)
|
||||
axis_grid = []
|
||||
for i in range(dim):
|
||||
a, b, n = start[i], stop[i], num[i]
|
||||
g = torch.linspace(a, b, n + 1, dtype=torch.float32)[:n]
|
||||
axis_grid.append(g)
|
||||
grid = torch.meshgrid(*axis_grid, indexing="ij") # dim x [W, H, D]
|
||||
grid = torch.stack(grid, dim=0) # [dim, W, H, D]
|
||||
|
||||
return grid
|
||||
|
||||
|
||||
#################################################################################
|
||||
# Rotary Positional Embedding Functions #
|
||||
#################################################################################
|
||||
# https://github.com/meta-llama/llama/blob/be327c427cc5e89cc1d3ab3d3fec4484df771245/llama/model.py#L80
|
||||
|
||||
|
||||
def reshape_for_broadcast(
|
||||
freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]],
|
||||
x: torch.Tensor,
|
||||
head_first=False,
|
||||
):
|
||||
"""
|
||||
Reshape frequency tensor for broadcasting it with another tensor.
|
||||
|
||||
This function reshapes the frequency tensor to have the same shape as the target tensor 'x'
|
||||
for the purpose of broadcasting the frequency tensor during element-wise operations.
|
||||
|
||||
Notes:
|
||||
When using FlashMHAModified, head_first should be False.
|
||||
When using Attention, head_first should be True.
|
||||
|
||||
Args:
|
||||
freqs_cis (Union[torch.Tensor, Tuple[torch.Tensor]]): Frequency tensor to be reshaped.
|
||||
x (torch.Tensor): Target tensor for broadcasting compatibility.
|
||||
head_first (bool): head dimension first (except batch dim) or not.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Reshaped frequency tensor.
|
||||
|
||||
Raises:
|
||||
AssertionError: If the frequency tensor doesn't match the expected shape.
|
||||
AssertionError: If the target tensor 'x' doesn't have the expected number of dimensions.
|
||||
"""
|
||||
ndim = x.ndim
|
||||
assert 0 <= 1 < ndim
|
||||
|
||||
if isinstance(freqs_cis, tuple):
|
||||
# freqs_cis: (cos, sin) in real space
|
||||
if head_first:
|
||||
assert freqs_cis[0].shape == (
|
||||
x.shape[-2],
|
||||
x.shape[-1],
|
||||
), f"freqs_cis shape {freqs_cis[0].shape} does not match x shape {x.shape}"
|
||||
shape = [
|
||||
d if i == ndim - 2 or i == ndim - 1 else 1
|
||||
for i, d in enumerate(x.shape)
|
||||
]
|
||||
else:
|
||||
assert freqs_cis[0].shape == (
|
||||
x.shape[1],
|
||||
x.shape[-1],
|
||||
), f"freqs_cis shape {freqs_cis[0].shape} does not match x shape {x.shape}"
|
||||
shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
|
||||
return freqs_cis[0].view(*shape), freqs_cis[1].view(*shape)
|
||||
else:
|
||||
# freqs_cis: values in complex space
|
||||
if head_first:
|
||||
assert freqs_cis.shape == (
|
||||
x.shape[-2],
|
||||
x.shape[-1],
|
||||
), f"freqs_cis shape {freqs_cis.shape} does not match x shape {x.shape}"
|
||||
shape = [
|
||||
d if i == ndim - 2 or i == ndim - 1 else 1
|
||||
for i, d in enumerate(x.shape)
|
||||
]
|
||||
else:
|
||||
assert freqs_cis.shape == (
|
||||
x.shape[1],
|
||||
x.shape[-1],
|
||||
), f"freqs_cis shape {freqs_cis.shape} does not match x shape {x.shape}"
|
||||
shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
|
||||
return freqs_cis.view(*shape)
|
||||
|
||||
|
||||
def rotate_half(x):
|
||||
x_real, x_imag = (
|
||||
x.float().reshape(*x.shape[:-1], -1, 2).unbind(-1)
|
||||
) # [B, S, H, D//2]
|
||||
return torch.stack([-x_imag, x_real], dim=-1).flatten(3)
|
||||
|
||||
|
||||
def apply_rotary_emb( qklist,
|
||||
freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
|
||||
head_first: bool = False,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Apply rotary embeddings to input tensors using the given frequency tensor.
|
||||
|
||||
This function applies rotary embeddings to the given query 'xq' and key 'xk' tensors using the provided
|
||||
frequency tensor 'freqs_cis'. The input tensors are reshaped as complex numbers, and the frequency tensor
|
||||
is reshaped for broadcasting compatibility. The resulting tensors contain rotary embeddings and are
|
||||
returned as real tensors.
|
||||
|
||||
Args:
|
||||
xq (torch.Tensor): Query tensor to apply rotary embeddings. [B, S, H, D]
|
||||
xk (torch.Tensor): Key tensor to apply rotary embeddings. [B, S, H, D]
|
||||
freqs_cis (torch.Tensor or tuple): Precomputed frequency tensor for complex exponential.
|
||||
head_first (bool): head dimension first (except batch dim) or not.
|
||||
|
||||
Returns:
|
||||
Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.
|
||||
|
||||
"""
|
||||
xq, xk = qklist
|
||||
qklist.clear()
|
||||
xk_out = None
|
||||
if isinstance(freqs_cis, tuple):
|
||||
cos, sin = reshape_for_broadcast(freqs_cis, xq, head_first) # [S, D]
|
||||
cos, sin = cos.to(xq.device), sin.to(xq.device)
|
||||
# real * cos - imag * sin
|
||||
# imag * cos + real * sin
|
||||
xq_dtype = xq.dtype
|
||||
xq_out = xq.to(torch.float)
|
||||
xq = None
|
||||
xq_rot = rotate_half(xq_out)
|
||||
xq_out *= cos
|
||||
xq_rot *= sin
|
||||
xq_out += xq_rot
|
||||
del xq_rot
|
||||
xq_out = xq_out.to(xq_dtype)
|
||||
|
||||
xk_out = xk.to(torch.float)
|
||||
xk = None
|
||||
xk_rot = rotate_half(xk_out)
|
||||
xk_out *= cos
|
||||
xk_rot *= sin
|
||||
xk_out += xk_rot
|
||||
del xk_rot
|
||||
xk_out = xk_out.to(xq_dtype)
|
||||
else:
|
||||
# view_as_complex will pack [..., D/2, 2](real) to [..., D/2](complex)
|
||||
xq_ = torch.view_as_complex(
|
||||
xq.float().reshape(*xq.shape[:-1], -1, 2)
|
||||
) # [B, S, H, D//2]
|
||||
freqs_cis = reshape_for_broadcast(freqs_cis, xq_, head_first).to(
|
||||
xq.device
|
||||
) # [S, D//2] --> [1, S, 1, D//2]
|
||||
# (real, imag) * (cos, sin) = (real * cos - imag * sin, imag * cos + real * sin)
|
||||
# view_as_real will expand [..., D/2](complex) to [..., D/2, 2](real)
|
||||
xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3).type_as(xq)
|
||||
xk_ = torch.view_as_complex(
|
||||
xk.float().reshape(*xk.shape[:-1], -1, 2)
|
||||
) # [B, S, H, D//2]
|
||||
xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3).type_as(xk)
|
||||
|
||||
return xq_out, xk_out
|
||||
|
||||
def get_nd_rotary_pos_embed_new(rope_dim_list, start, *args, theta=10000., use_real=False,
|
||||
theta_rescale_factor: Union[float, List[float]]=1.0,
|
||||
interpolation_factor: Union[float, List[float]]=1.0,
|
||||
concat_dict={}
|
||||
):
|
||||
|
||||
grid = get_meshgrid_nd(start, *args, dim=len(rope_dim_list)) # [3, W, H, D] / [2, W, H]
|
||||
if len(concat_dict)<1:
|
||||
pass
|
||||
else:
|
||||
if concat_dict['mode']=='timecat':
|
||||
bias = grid[:,:1].clone()
|
||||
bias[0] = concat_dict['bias']*torch.ones_like(bias[0])
|
||||
grid = torch.cat([bias, grid], dim=1)
|
||||
|
||||
elif concat_dict['mode']=='timecat-w':
|
||||
bias = grid[:,:1].clone()
|
||||
bias[0] = concat_dict['bias']*torch.ones_like(bias[0])
|
||||
bias[2] += start[-1] ## ref https://github.com/Yuanshi9815/OminiControl/blob/main/src/generate.py#L178
|
||||
grid = torch.cat([bias, grid], dim=1)
|
||||
if isinstance(theta_rescale_factor, int) or isinstance(theta_rescale_factor, float):
|
||||
theta_rescale_factor = [theta_rescale_factor] * len(rope_dim_list)
|
||||
elif isinstance(theta_rescale_factor, list) and len(theta_rescale_factor) == 1:
|
||||
theta_rescale_factor = [theta_rescale_factor[0]] * len(rope_dim_list)
|
||||
assert len(theta_rescale_factor) == len(rope_dim_list), "len(theta_rescale_factor) should equal to len(rope_dim_list)"
|
||||
|
||||
if isinstance(interpolation_factor, int) or isinstance(interpolation_factor, float):
|
||||
interpolation_factor = [interpolation_factor] * len(rope_dim_list)
|
||||
elif isinstance(interpolation_factor, list) and len(interpolation_factor) == 1:
|
||||
interpolation_factor = [interpolation_factor[0]] * len(rope_dim_list)
|
||||
assert len(interpolation_factor) == len(rope_dim_list), "len(interpolation_factor) should equal to len(rope_dim_list)"
|
||||
|
||||
# use 1/ndim of dimensions to encode grid_axis
|
||||
embs = []
|
||||
for i in range(len(rope_dim_list)):
|
||||
emb = get_1d_rotary_pos_embed(rope_dim_list[i], grid[i].reshape(-1), theta, use_real=use_real,
|
||||
theta_rescale_factor=theta_rescale_factor[i],
|
||||
interpolation_factor=interpolation_factor[i]) # 2 x [WHD, rope_dim_list[i]]
|
||||
|
||||
embs.append(emb)
|
||||
|
||||
if use_real:
|
||||
cos = torch.cat([emb[0] for emb in embs], dim=1) # (WHD, D/2)
|
||||
sin = torch.cat([emb[1] for emb in embs], dim=1) # (WHD, D/2)
|
||||
return cos, sin
|
||||
else:
|
||||
emb = torch.cat(embs, dim=1) # (WHD, D/2)
|
||||
return emb
|
||||
|
||||
def get_nd_rotary_pos_embed(
|
||||
rope_dim_list,
|
||||
start,
|
||||
*args,
|
||||
theta=10000.0,
|
||||
use_real=False,
|
||||
theta_rescale_factor: Union[float, List[float]] = 1.0,
|
||||
interpolation_factor: Union[float, List[float]] = 1.0,
|
||||
k = 4,
|
||||
L_test = 66,
|
||||
enable_riflex = True
|
||||
):
|
||||
"""
|
||||
This is a n-d version of precompute_freqs_cis, which is a RoPE for tokens with n-d structure.
|
||||
|
||||
Args:
|
||||
rope_dim_list (list of int): Dimension of each rope. len(rope_dim_list) should equal to n.
|
||||
sum(rope_dim_list) should equal to head_dim of attention layer.
|
||||
start (int | tuple of int | list of int): If len(args) == 0, start is num; If len(args) == 1, start is start,
|
||||
args[0] is stop, step is 1; If len(args) == 2, start is start, args[0] is stop, args[1] is num.
|
||||
*args: See above.
|
||||
theta (float): Scaling factor for frequency computation. Defaults to 10000.0.
|
||||
use_real (bool): If True, return real part and imaginary part separately. Otherwise, return complex numbers.
|
||||
Some libraries such as TensorRT does not support complex64 data type. So it is useful to provide a real
|
||||
part and an imaginary part separately.
|
||||
theta_rescale_factor (float): Rescale factor for theta. Defaults to 1.0.
|
||||
|
||||
Returns:
|
||||
pos_embed (torch.Tensor): [HW, D/2]
|
||||
"""
|
||||
|
||||
grid = get_meshgrid_nd(
|
||||
start, *args, dim=len(rope_dim_list)
|
||||
) # [3, W, H, D] / [2, W, H]
|
||||
|
||||
if isinstance(theta_rescale_factor, int) or isinstance(theta_rescale_factor, float):
|
||||
theta_rescale_factor = [theta_rescale_factor] * len(rope_dim_list)
|
||||
elif isinstance(theta_rescale_factor, list) and len(theta_rescale_factor) == 1:
|
||||
theta_rescale_factor = [theta_rescale_factor[0]] * len(rope_dim_list)
|
||||
assert len(theta_rescale_factor) == len(
|
||||
rope_dim_list
|
||||
), "len(theta_rescale_factor) should equal to len(rope_dim_list)"
|
||||
|
||||
if isinstance(interpolation_factor, int) or isinstance(interpolation_factor, float):
|
||||
interpolation_factor = [interpolation_factor] * len(rope_dim_list)
|
||||
elif isinstance(interpolation_factor, list) and len(interpolation_factor) == 1:
|
||||
interpolation_factor = [interpolation_factor[0]] * len(rope_dim_list)
|
||||
assert len(interpolation_factor) == len(
|
||||
rope_dim_list
|
||||
), "len(interpolation_factor) should equal to len(rope_dim_list)"
|
||||
|
||||
# use 1/ndim of dimensions to encode grid_axis
|
||||
embs = []
|
||||
for i in range(len(rope_dim_list)):
|
||||
# emb = get_1d_rotary_pos_embed(
|
||||
# rope_dim_list[i],
|
||||
# grid[i].reshape(-1),
|
||||
# theta,
|
||||
# use_real=use_real,
|
||||
# theta_rescale_factor=theta_rescale_factor[i],
|
||||
# interpolation_factor=interpolation_factor[i],
|
||||
# ) # 2 x [WHD, rope_dim_list[i]]
|
||||
|
||||
|
||||
# === RIFLEx modification start ===
|
||||
# apply RIFLEx for time dimension
|
||||
if i == 0 and enable_riflex:
|
||||
emb = get_1d_rotary_pos_embed_riflex(rope_dim_list[i], grid[i].reshape(-1), theta, use_real=True, k=k, L_test=L_test)
|
||||
# === RIFLEx modification end ===
|
||||
else:
|
||||
emb = get_1d_rotary_pos_embed(rope_dim_list[i], grid[i].reshape(-1), theta, use_real=True, theta_rescale_factor=theta_rescale_factor[i],interpolation_factor=interpolation_factor[i],)
|
||||
embs.append(emb)
|
||||
|
||||
if use_real:
|
||||
cos = torch.cat([emb[0] for emb in embs], dim=1) # (WHD, D/2)
|
||||
sin = torch.cat([emb[1] for emb in embs], dim=1) # (WHD, D/2)
|
||||
return cos, sin
|
||||
else:
|
||||
emb = torch.cat(embs, dim=1) # (WHD, D/2)
|
||||
return emb
|
||||
|
||||
|
||||
def get_1d_rotary_pos_embed(
|
||||
dim: int,
|
||||
pos: Union[torch.FloatTensor, int],
|
||||
theta: float = 10000.0,
|
||||
use_real: bool = False,
|
||||
theta_rescale_factor: float = 1.0,
|
||||
interpolation_factor: float = 1.0,
|
||||
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
|
||||
"""
|
||||
Precompute the frequency tensor for complex exponential (cis) with given dimensions.
|
||||
(Note: `cis` means `cos + i * sin`, where i is the imaginary unit.)
|
||||
|
||||
This function calculates a frequency tensor with complex exponential using the given dimension 'dim'
|
||||
and the end index 'end'. The 'theta' parameter scales the frequencies.
|
||||
The returned tensor contains complex values in complex64 data type.
|
||||
|
||||
Args:
|
||||
dim (int): Dimension of the frequency tensor.
|
||||
pos (int or torch.FloatTensor): Position indices for the frequency tensor. [S] or scalar
|
||||
theta (float, optional): Scaling factor for frequency computation. Defaults to 10000.0.
|
||||
use_real (bool, optional): If True, return real part and imaginary part separately.
|
||||
Otherwise, return complex numbers.
|
||||
theta_rescale_factor (float, optional): Rescale factor for theta. Defaults to 1.0.
|
||||
|
||||
Returns:
|
||||
freqs_cis: Precomputed frequency tensor with complex exponential. [S, D/2]
|
||||
freqs_cos, freqs_sin: Precomputed frequency tensor with real and imaginary parts separately. [S, D]
|
||||
"""
|
||||
if isinstance(pos, int):
|
||||
pos = torch.arange(pos).float()
|
||||
|
||||
# proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning
|
||||
# has some connection to NTK literature
|
||||
if theta_rescale_factor != 1.0:
|
||||
theta *= theta_rescale_factor ** (dim / (dim - 2))
|
||||
|
||||
freqs = 1.0 / (
|
||||
theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)
|
||||
) # [D/2]
|
||||
# assert interpolation_factor == 1.0, f"interpolation_factor: {interpolation_factor}"
|
||||
freqs = torch.outer(pos * interpolation_factor, freqs) # [S, D/2]
|
||||
if use_real:
|
||||
freqs_cos = freqs.cos().repeat_interleave(2, dim=1) # [S, D]
|
||||
freqs_sin = freqs.sin().repeat_interleave(2, dim=1) # [S, D]
|
||||
return freqs_cos, freqs_sin
|
||||
else:
|
||||
freqs_cis = torch.polar(
|
||||
torch.ones_like(freqs), freqs
|
||||
) # complex64 # [S, D/2]
|
||||
return freqs_cis
|
||||
237
hyvideo/modules/token_refiner.py
Normal file
237
hyvideo/modules/token_refiner.py
Normal file
@@ -0,0 +1,237 @@
|
||||
from typing import Optional
|
||||
|
||||
from einops import rearrange
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from .activation_layers import get_activation_layer
|
||||
from .attenion import attention
|
||||
from .norm_layers import get_norm_layer
|
||||
from .embed_layers import TimestepEmbedder, TextProjection
|
||||
from .attenion import attention
|
||||
from .mlp_layers import MLP
|
||||
from .modulate_layers import modulate, apply_gate
|
||||
|
||||
|
||||
class IndividualTokenRefinerBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size,
|
||||
heads_num,
|
||||
mlp_width_ratio: str = 4.0,
|
||||
mlp_drop_rate: float = 0.0,
|
||||
act_type: str = "silu",
|
||||
qk_norm: bool = False,
|
||||
qk_norm_type: str = "layer",
|
||||
qkv_bias: bool = True,
|
||||
dtype: Optional[torch.dtype] = None,
|
||||
device: Optional[torch.device] = None,
|
||||
):
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__()
|
||||
self.heads_num = heads_num
|
||||
head_dim = hidden_size // heads_num
|
||||
mlp_hidden_dim = int(hidden_size * mlp_width_ratio)
|
||||
|
||||
self.norm1 = nn.LayerNorm(
|
||||
hidden_size, elementwise_affine=True, eps=1e-6, **factory_kwargs
|
||||
)
|
||||
self.self_attn_qkv = nn.Linear(
|
||||
hidden_size, hidden_size * 3, bias=qkv_bias, **factory_kwargs
|
||||
)
|
||||
qk_norm_layer = get_norm_layer(qk_norm_type)
|
||||
self.self_attn_q_norm = (
|
||||
qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
|
||||
if qk_norm
|
||||
else nn.Identity()
|
||||
)
|
||||
self.self_attn_k_norm = (
|
||||
qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
|
||||
if qk_norm
|
||||
else nn.Identity()
|
||||
)
|
||||
self.self_attn_proj = nn.Linear(
|
||||
hidden_size, hidden_size, bias=qkv_bias, **factory_kwargs
|
||||
)
|
||||
|
||||
self.norm2 = nn.LayerNorm(
|
||||
hidden_size, elementwise_affine=True, eps=1e-6, **factory_kwargs
|
||||
)
|
||||
act_layer = get_activation_layer(act_type)
|
||||
self.mlp = MLP(
|
||||
in_channels=hidden_size,
|
||||
hidden_channels=mlp_hidden_dim,
|
||||
act_layer=act_layer,
|
||||
drop=mlp_drop_rate,
|
||||
**factory_kwargs,
|
||||
)
|
||||
|
||||
self.adaLN_modulation = nn.Sequential(
|
||||
act_layer(),
|
||||
nn.Linear(hidden_size, 2 * hidden_size, bias=True, **factory_kwargs),
|
||||
)
|
||||
# Zero-initialize the modulation
|
||||
nn.init.zeros_(self.adaLN_modulation[1].weight)
|
||||
nn.init.zeros_(self.adaLN_modulation[1].bias)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
c: torch.Tensor, # timestep_aware_representations + context_aware_representations
|
||||
attn_mask: torch.Tensor = None,
|
||||
):
|
||||
gate_msa, gate_mlp = self.adaLN_modulation(c).chunk(2, dim=1)
|
||||
|
||||
norm_x = self.norm1(x)
|
||||
qkv = self.self_attn_qkv(norm_x)
|
||||
q, k, v = rearrange(qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num)
|
||||
# Apply QK-Norm if needed
|
||||
q = self.self_attn_q_norm(q).to(v)
|
||||
k = self.self_attn_k_norm(k).to(v)
|
||||
qkv_list = [q, k, v]
|
||||
del q,k
|
||||
# Self-Attention
|
||||
attn = attention( qkv_list, mode="torch", attn_mask=attn_mask)
|
||||
|
||||
x = x + apply_gate(self.self_attn_proj(attn), gate_msa)
|
||||
|
||||
# FFN Layer
|
||||
x = x + apply_gate(self.mlp(self.norm2(x)), gate_mlp)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class IndividualTokenRefiner(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size,
|
||||
heads_num,
|
||||
depth,
|
||||
mlp_width_ratio: float = 4.0,
|
||||
mlp_drop_rate: float = 0.0,
|
||||
act_type: str = "silu",
|
||||
qk_norm: bool = False,
|
||||
qk_norm_type: str = "layer",
|
||||
qkv_bias: bool = True,
|
||||
dtype: Optional[torch.dtype] = None,
|
||||
device: Optional[torch.device] = None,
|
||||
):
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__()
|
||||
self.blocks = nn.ModuleList(
|
||||
[
|
||||
IndividualTokenRefinerBlock(
|
||||
hidden_size=hidden_size,
|
||||
heads_num=heads_num,
|
||||
mlp_width_ratio=mlp_width_ratio,
|
||||
mlp_drop_rate=mlp_drop_rate,
|
||||
act_type=act_type,
|
||||
qk_norm=qk_norm,
|
||||
qk_norm_type=qk_norm_type,
|
||||
qkv_bias=qkv_bias,
|
||||
**factory_kwargs,
|
||||
)
|
||||
for _ in range(depth)
|
||||
]
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
c: torch.LongTensor,
|
||||
mask: Optional[torch.Tensor] = None,
|
||||
):
|
||||
self_attn_mask = None
|
||||
if mask is not None:
|
||||
batch_size = mask.shape[0]
|
||||
seq_len = mask.shape[1]
|
||||
mask = mask.to(x.device)
|
||||
# batch_size x 1 x seq_len x seq_len
|
||||
self_attn_mask_1 = mask.view(batch_size, 1, 1, seq_len).repeat(
|
||||
1, 1, seq_len, 1
|
||||
)
|
||||
# batch_size x 1 x seq_len x seq_len
|
||||
self_attn_mask_2 = self_attn_mask_1.transpose(2, 3)
|
||||
# batch_size x 1 x seq_len x seq_len, 1 for broadcasting of heads_num
|
||||
self_attn_mask = (self_attn_mask_1 & self_attn_mask_2).bool()
|
||||
# avoids self-attention weight being NaN for padding tokens
|
||||
self_attn_mask[:, :, :, 0] = True
|
||||
|
||||
for block in self.blocks:
|
||||
x = block(x, c, self_attn_mask)
|
||||
return x
|
||||
|
||||
|
||||
class SingleTokenRefiner(nn.Module):
|
||||
"""
|
||||
A single token refiner block for llm text embedding refine.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
hidden_size,
|
||||
heads_num,
|
||||
depth,
|
||||
mlp_width_ratio: float = 4.0,
|
||||
mlp_drop_rate: float = 0.0,
|
||||
act_type: str = "silu",
|
||||
qk_norm: bool = False,
|
||||
qk_norm_type: str = "layer",
|
||||
qkv_bias: bool = True,
|
||||
attn_mode: str = "torch",
|
||||
dtype: Optional[torch.dtype] = None,
|
||||
device: Optional[torch.device] = None,
|
||||
):
|
||||
factory_kwargs = {"device": device, "dtype": dtype}
|
||||
super().__init__()
|
||||
self.attn_mode = attn_mode
|
||||
assert self.attn_mode == "torch", "Only support 'torch' mode for token refiner."
|
||||
|
||||
self.input_embedder = nn.Linear(
|
||||
in_channels, hidden_size, bias=True, **factory_kwargs
|
||||
)
|
||||
|
||||
act_layer = get_activation_layer(act_type)
|
||||
# Build timestep embedding layer
|
||||
self.t_embedder = TimestepEmbedder(hidden_size, act_layer, **factory_kwargs)
|
||||
# Build context embedding layer
|
||||
self.c_embedder = TextProjection(
|
||||
in_channels, hidden_size, act_layer, **factory_kwargs
|
||||
)
|
||||
|
||||
self.individual_token_refiner = IndividualTokenRefiner(
|
||||
hidden_size=hidden_size,
|
||||
heads_num=heads_num,
|
||||
depth=depth,
|
||||
mlp_width_ratio=mlp_width_ratio,
|
||||
mlp_drop_rate=mlp_drop_rate,
|
||||
act_type=act_type,
|
||||
qk_norm=qk_norm,
|
||||
qk_norm_type=qk_norm_type,
|
||||
qkv_bias=qkv_bias,
|
||||
**factory_kwargs,
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
t: torch.LongTensor,
|
||||
mask: Optional[torch.LongTensor] = None,
|
||||
):
|
||||
timestep_aware_representations = self.t_embedder(t)
|
||||
|
||||
if mask is None:
|
||||
context_aware_representations = x.mean(dim=1)
|
||||
else:
|
||||
mask_float = mask.float().unsqueeze(-1) # [b, s1, 1]
|
||||
context_aware_representations = (x * mask_float).sum(
|
||||
dim=1
|
||||
) / mask_float.sum(dim=1)
|
||||
context_aware_representations = self.c_embedder(context_aware_representations.to(x.dtype))
|
||||
c = timestep_aware_representations + context_aware_representations
|
||||
|
||||
x = self.input_embedder(x)
|
||||
|
||||
x = self.individual_token_refiner(x, c, mask)
|
||||
|
||||
return x
|
||||
43
hyvideo/modules/utils.py
Normal file
43
hyvideo/modules/utils.py
Normal file
@@ -0,0 +1,43 @@
|
||||
"""Mask Mod for Image2Video"""
|
||||
|
||||
from math import floor
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
|
||||
from functools import lru_cache
|
||||
from typing import Optional, List
|
||||
|
||||
import torch
|
||||
from torch.nn.attention.flex_attention import (
|
||||
create_block_mask,
|
||||
)
|
||||
|
||||
|
||||
@lru_cache
|
||||
def create_block_mask_cached(score_mod, B, H, M, N, device="cuda", _compile=False):
|
||||
block_mask = create_block_mask(score_mod, B, H, M, N, device=device, _compile=_compile)
|
||||
return block_mask
|
||||
|
||||
def generate_temporal_head_mask_mod(context_length: int = 226, prompt_length: int = 226, num_frames: int = 13, token_per_frame: int = 1350, mul: int = 2):
|
||||
|
||||
def round_to_multiple(idx):
|
||||
return floor(idx / 128) * 128
|
||||
|
||||
real_length = num_frames * token_per_frame + prompt_length
|
||||
def temporal_mask_mod(b, h, q_idx, kv_idx):
|
||||
real_mask = (kv_idx < real_length) & (q_idx < real_length)
|
||||
fake_mask = (kv_idx >= real_length) & (q_idx >= real_length)
|
||||
|
||||
two_frame = round_to_multiple(mul * token_per_frame)
|
||||
temporal_head_mask = (torch.abs(q_idx - kv_idx) < two_frame)
|
||||
|
||||
text_column_mask = (num_frames * token_per_frame <= kv_idx) & (kv_idx < real_length)
|
||||
text_row_mask = (num_frames * token_per_frame <= q_idx) & (q_idx < real_length)
|
||||
|
||||
video_mask = temporal_head_mask | text_column_mask | text_row_mask
|
||||
real_mask = real_mask & video_mask
|
||||
|
||||
return real_mask | fake_mask
|
||||
|
||||
return temporal_mask_mod
|
||||
Reference in New Issue
Block a user