chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,477 @@
|
||||
# Adopted from https://github.com/Wan-Video/Wan2.2
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
import torch
|
||||
import torch.cuda.amp as amp
|
||||
|
||||
from ..modules.model import sinusoidal_embedding_1d
|
||||
from .ulysses import distributed_attention, distributed_flex_attention
|
||||
from .util import gather_forward, get_rank, get_world_size
|
||||
import math
|
||||
|
||||
|
||||
def pad_freqs(original_tensor, target_len):
|
||||
seq_len, s1, s2 = original_tensor.shape
|
||||
pad_size = target_len - seq_len
|
||||
padding_tensor = torch.ones(
|
||||
pad_size,
|
||||
s1,
|
||||
s2,
|
||||
dtype=original_tensor.dtype,
|
||||
device=original_tensor.device)
|
||||
padded_tensor = torch.cat([original_tensor, padding_tensor], dim=0)
|
||||
return padded_tensor
|
||||
|
||||
|
||||
from utils.position_embedding_utils import (
|
||||
compute_temporal_freqs as _compute_temporal_freqs,
|
||||
select_temporal_offset_for_sample,
|
||||
)
|
||||
|
||||
|
||||
@torch.amp.autocast('cuda', enabled=False)
|
||||
def sp_rope_apply(
|
||||
x,
|
||||
grid_sizes,
|
||||
freqs,
|
||||
t_scale=1.0,
|
||||
method="linear",
|
||||
original_seq_len=None,
|
||||
temporal_offset=0.0,
|
||||
):
|
||||
"""
|
||||
x: [B, L, N, C].
|
||||
grid_sizes: [B, 3].
|
||||
freqs: [M, C // 2].
|
||||
"""
|
||||
n, c = x.size(2), x.size(3) // 2
|
||||
# split freqs
|
||||
freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
|
||||
|
||||
# loop over samples
|
||||
output = []
|
||||
for i, (f, h, w) in enumerate(grid_sizes.tolist()):
|
||||
local_f = f
|
||||
sp_rank = get_rank()
|
||||
start_frame = sp_rank * local_f
|
||||
seq_len = local_f * h * w
|
||||
|
||||
# precompute multipliers
|
||||
x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float64).reshape(
|
||||
seq_len, n, -1, 2))
|
||||
temporal_offset_i = select_temporal_offset_for_sample(
|
||||
temporal_offset, i, local_f, start_frame=start_frame)
|
||||
temporal_freqs = _compute_temporal_freqs(
|
||||
freqs[0], local_f, start_frame, t_scale, x.device,
|
||||
method=method, original_seq_len=original_seq_len,
|
||||
temporal_offset=temporal_offset_i)
|
||||
freqs_i = torch.cat([
|
||||
temporal_freqs.view(local_f, 1, 1, -1).expand(local_f, h, w, -1),
|
||||
freqs[1][:h].view(1, h, 1, -1).expand(local_f, h, w, -1),
|
||||
freqs[2][:w].view(1, 1, w, -1).expand(local_f, h, w, -1)
|
||||
],
|
||||
dim=-1).reshape(seq_len, 1, -1)
|
||||
|
||||
# apply rotary embedding
|
||||
x_i = torch.view_as_real(x_i * freqs_i).flatten(2)
|
||||
x_i = torch.cat([x_i, x[i, seq_len:]])
|
||||
|
||||
# append to collection
|
||||
output.append(x_i)
|
||||
return torch.stack(output).float()
|
||||
|
||||
|
||||
def sp_dit_forward(
|
||||
self,
|
||||
x,
|
||||
t,
|
||||
context,
|
||||
seq_len,
|
||||
y=None,
|
||||
):
|
||||
"""
|
||||
x: A list of videos each with shape [C, T, H, W].
|
||||
t: [B].
|
||||
context: A list of text embeddings each with shape [L, C].
|
||||
"""
|
||||
if self.model_type == 'i2v':
|
||||
assert y is not None
|
||||
# params
|
||||
device = self.patch_embedding.weight.device
|
||||
if self.freqs.device != device:
|
||||
self.freqs = self.freqs.to(device)
|
||||
|
||||
if y is not None:
|
||||
x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]
|
||||
|
||||
# embeddings
|
||||
x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
|
||||
grid_sizes = torch.stack(
|
||||
[torch.tensor(u.shape[2:], dtype=torch.long) for u in x])
|
||||
x = [u.flatten(2).transpose(1, 2) for u in x]
|
||||
seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)
|
||||
assert seq_lens.max() <= seq_len
|
||||
x = torch.cat([
|
||||
torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], dim=1)
|
||||
for u in x
|
||||
])
|
||||
|
||||
# time embeddings
|
||||
if t.dim() == 1:
|
||||
t = t.expand(t.size(0), seq_len)
|
||||
with torch.amp.autocast('cuda', dtype=torch.float32):
|
||||
bt = t.size(0)
|
||||
t = t.flatten()
|
||||
e = self.time_embedding(
|
||||
sinusoidal_embedding_1d(self.freq_dim,
|
||||
t).unflatten(0, (bt, seq_len)).float())
|
||||
e0 = self.time_projection(e).unflatten(2, (6, self.dim))
|
||||
assert e.dtype == torch.float32 and e0.dtype == torch.float32
|
||||
|
||||
# context
|
||||
context_lens = None
|
||||
context = self.text_embedding(
|
||||
torch.stack([
|
||||
torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))])
|
||||
for u in context
|
||||
]))
|
||||
|
||||
# Context Parallel
|
||||
x = torch.chunk(x, get_world_size(), dim=1)[get_rank()]
|
||||
e = torch.chunk(e, get_world_size(), dim=1)[get_rank()]
|
||||
e0 = torch.chunk(e0, get_world_size(), dim=1)[get_rank()]
|
||||
|
||||
# arguments
|
||||
kwargs = dict(
|
||||
e=e0,
|
||||
seq_lens=seq_lens,
|
||||
grid_sizes=grid_sizes,
|
||||
freqs=self.freqs,
|
||||
context=context,
|
||||
context_lens=context_lens)
|
||||
|
||||
for block in self.blocks:
|
||||
x = block(x, **kwargs)
|
||||
|
||||
# head
|
||||
x = self.head(x, e)
|
||||
|
||||
# Context Parallel
|
||||
x = gather_forward(x, dim=1)
|
||||
|
||||
# unpatchify
|
||||
x = self.unpatchify(x, grid_sizes)
|
||||
return [u.float() for u in x]
|
||||
|
||||
|
||||
def sp_attn_forward(self, x, seq_lens, grid_sizes, freqs, dtype=torch.bfloat16,
|
||||
t_scale=1.0, method="linear", original_seq_len=None,
|
||||
temporal_offset=0.0):
|
||||
b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
|
||||
half_dtypes = (torch.float16, torch.bfloat16)
|
||||
|
||||
def half(x):
|
||||
return x if x.dtype in half_dtypes else x.to(dtype)
|
||||
|
||||
# query, key, value function
|
||||
def qkv_fn(x):
|
||||
q = self.norm_q(self.q(x)).view(b, s, n, d)
|
||||
k = self.norm_k(self.k(x)).view(b, s, n, d)
|
||||
v = self.v(x).view(b, s, n, d)
|
||||
return q, k, v
|
||||
|
||||
q, k, v = qkv_fn(x)
|
||||
q = sp_rope_apply(q, grid_sizes, freqs, t_scale=t_scale,
|
||||
method=method, original_seq_len=original_seq_len,
|
||||
temporal_offset=temporal_offset)
|
||||
k = sp_rope_apply(k, grid_sizes, freqs, t_scale=t_scale,
|
||||
method=method, original_seq_len=original_seq_len,
|
||||
temporal_offset=temporal_offset)
|
||||
|
||||
x = distributed_attention(
|
||||
half(q),
|
||||
half(k),
|
||||
half(v),
|
||||
seq_lens,
|
||||
window_size=self.window_size,
|
||||
)
|
||||
|
||||
# output
|
||||
x = x.flatten(2)
|
||||
x = self.o(x)
|
||||
return x
|
||||
|
||||
|
||||
|
||||
|
||||
def sp_dit_causal_forward_train(
|
||||
self,
|
||||
x,
|
||||
t,
|
||||
context,
|
||||
seq_len,
|
||||
clean_x=None,
|
||||
aug_t=None,
|
||||
clip_fea=None,
|
||||
y=None,
|
||||
):
|
||||
r"""
|
||||
Forward pass through the diffusion model
|
||||
|
||||
Args:
|
||||
x (List[Tensor]):
|
||||
List of input video tensors, each with shape [C_in, F, H, W]
|
||||
t (Tensor):
|
||||
Diffusion timesteps tensor of shape [B]
|
||||
context (List[Tensor]):
|
||||
List of text embeddings each with shape [L, C]
|
||||
seq_len (`int`):
|
||||
Maximum sequence length for positional encoding
|
||||
clip_fea (Tensor, *optional*):
|
||||
CLIP image features for image-to-video mode
|
||||
y (List[Tensor], *optional*):
|
||||
Conditional video inputs for image-to-video mode, same shape as x
|
||||
|
||||
Returns:
|
||||
List[Tensor]:
|
||||
List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8]
|
||||
"""
|
||||
if self.model_type == 'i2v':
|
||||
assert clip_fea is not None and y is not None
|
||||
# params
|
||||
device = self.patch_embedding.weight.device
|
||||
if self.freqs.device != device:
|
||||
self.freqs = self.freqs.to(device)
|
||||
|
||||
# Construct the blockwise causal attention mask. Frames are sharded across
|
||||
# SP ranks, so total frames = local frames per rank * sp_size.
|
||||
sp_size = get_world_size()
|
||||
# Recreate mask when batch size changes to avoid Triton broadcasting bug
|
||||
current_batch_size = x.shape[0]
|
||||
if self.block_mask is None or self._block_mask_batch_size != current_batch_size:
|
||||
self._block_mask_batch_size = current_batch_size
|
||||
if clean_x is not None:
|
||||
if self.independent_first_frame:
|
||||
raise NotImplementedError()
|
||||
else:
|
||||
self.block_mask = self._prepare_teacher_forcing_mask_natural(
|
||||
device, num_frames=x.shape[2] * sp_size,
|
||||
frame_seqlen=x.shape[-2] * x.shape[-1] // (self.patch_size[1] * self.patch_size[2]),
|
||||
num_frame_per_block=self.num_frame_per_block,
|
||||
sp_size=sp_size,
|
||||
batch_size=current_batch_size,
|
||||
)
|
||||
else:
|
||||
if self.independent_first_frame:
|
||||
self.block_mask = self._prepare_blockwise_causal_attn_mask_i2v(
|
||||
device, num_frames=x.shape[2] * sp_size,
|
||||
frame_seqlen=x.shape[-2] * x.shape[-1] // (self.patch_size[1] * self.patch_size[2]),
|
||||
num_frame_per_block=self.num_frame_per_block,
|
||||
batch_size=current_batch_size,
|
||||
)
|
||||
else:
|
||||
self.block_mask = self._prepare_blockwise_causal_attn_mask(
|
||||
device, num_frames=x.shape[2] * sp_size,
|
||||
frame_seqlen=x.shape[-2] * x.shape[-1] // (self.patch_size[1] * self.patch_size[2]),
|
||||
num_frame_per_block=self.num_frame_per_block,
|
||||
batch_size=current_batch_size,
|
||||
)
|
||||
|
||||
if y is not None:
|
||||
x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]
|
||||
|
||||
# embeddings
|
||||
x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
|
||||
|
||||
grid_sizes = torch.stack(
|
||||
[torch.tensor(u.shape[2:], dtype=torch.long) for u in x])
|
||||
x = [u.flatten(2).transpose(1, 2) for u in x]
|
||||
|
||||
seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)
|
||||
max_len = int(seq_lens.max().item())
|
||||
assert max_len > 0, "Token sequence length is zero after patch embedding"
|
||||
# Pad all samples to the batch max length instead of the first sample length
|
||||
x = torch.cat([
|
||||
torch.cat([u, u.new_zeros(1, max_len - u.size(1), u.size(2))], dim=1)
|
||||
for u in x
|
||||
])
|
||||
|
||||
# time embeddings
|
||||
if t.dim() == 1:
|
||||
raise NotImplementedError(f"t.shape should be [B, F], but got {t.shape}")
|
||||
bt = t.size(0)
|
||||
t_len = t.size(1)
|
||||
t_ori_shape = t.shape
|
||||
t = t.flatten()
|
||||
e = self.time_embedding(
|
||||
sinusoidal_embedding_1d(self.freq_dim, t).unflatten(0, (bt, t_len)).type_as(x))
|
||||
e0 = self.time_projection(e).unflatten(2, (6, self.dim)) # B, F, 6, C
|
||||
|
||||
# context
|
||||
context_lens = None
|
||||
context = self.text_embedding(
|
||||
torch.stack([
|
||||
torch.cat(
|
||||
[u, u.new_zeros(self.text_len - u.size(0), u.size(1))])
|
||||
for u in context
|
||||
]))
|
||||
|
||||
if clean_x is not None:
|
||||
clean_x = [self.patch_embedding(u.unsqueeze(0)) for u in clean_x]
|
||||
clean_x = [u.flatten(2).transpose(1, 2) for u in clean_x]
|
||||
|
||||
seq_lens_clean = torch.tensor([u.size(1) for u in clean_x], dtype=torch.long)
|
||||
clean_x = torch.cat([
|
||||
torch.cat([u, u.new_zeros(1, seq_lens_clean[0] - u.size(1), u.size(2))], dim=1) for u in clean_x
|
||||
])
|
||||
|
||||
x = torch.cat([clean_x, x], dim=1)
|
||||
if aug_t is None:
|
||||
aug_t = torch.zeros(t_ori_shape, device=t.device, dtype=t.dtype)
|
||||
bt_clean = aug_t.size(0)
|
||||
t_clean_len = aug_t.size(1)
|
||||
aug_t = aug_t.flatten()
|
||||
e_clean = self.time_embedding(
|
||||
sinusoidal_embedding_1d(self.freq_dim, aug_t).unflatten(0, (bt_clean, t_clean_len)).type_as(x))
|
||||
e0_clean = self.time_projection(e_clean).unflatten(2, (6, self.dim))
|
||||
e0 = torch.cat([e0_clean, e0], dim=1)
|
||||
|
||||
|
||||
# arguments
|
||||
kwargs = dict(
|
||||
e=e0,
|
||||
seq_lens=seq_lens,
|
||||
grid_sizes=grid_sizes,
|
||||
freqs=self.freqs,
|
||||
context=context,
|
||||
context_lens=context_lens,
|
||||
block_mask=self.block_mask,
|
||||
t_scale=self.t_scale,
|
||||
use_relative_rope=getattr(self, "use_relative_rope", False),
|
||||
method=getattr(self, "rope_method", "linear"),
|
||||
original_seq_len=getattr(self, "original_seq_len", None),
|
||||
temporal_offset=getattr(self, "rope_temporal_offset", 0.0),
|
||||
)
|
||||
|
||||
def create_custom_forward(module):
|
||||
def custom_forward(*inputs, **kwargs):
|
||||
return module(*inputs, **kwargs)
|
||||
return custom_forward
|
||||
|
||||
for block in self.blocks:
|
||||
if torch.is_grad_enabled() and self.gradient_checkpointing:
|
||||
x = torch.utils.checkpoint.checkpoint(
|
||||
create_custom_forward(block),
|
||||
x, **kwargs,
|
||||
use_reentrant=False,
|
||||
)
|
||||
else:
|
||||
x = block(x, **kwargs)
|
||||
|
||||
if clean_x is not None:
|
||||
x = x[:, x.shape[1] // 2:]
|
||||
|
||||
# head
|
||||
x = self.head(x, e.unsqueeze(2))
|
||||
|
||||
x = self.unpatchify(x, grid_sizes)
|
||||
return torch.stack(x)
|
||||
|
||||
|
||||
def sp_causal_attn_forward(
|
||||
self,
|
||||
x,
|
||||
seq_lens,
|
||||
grid_sizes,
|
||||
freqs,
|
||||
block_mask,
|
||||
kv_cache=None,
|
||||
current_start=0,
|
||||
cache_start=None,
|
||||
t_scale=1.0,
|
||||
use_relative_rope=False,
|
||||
method="linear",
|
||||
original_seq_len=None,
|
||||
temporal_offset=0.0,
|
||||
**kwargs,
|
||||
):
|
||||
r"""
|
||||
Args:
|
||||
x(Tensor): Shape [B, L, num_heads, C / num_heads]
|
||||
seq_lens(Tensor): Shape [B]
|
||||
grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
|
||||
freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
|
||||
block_mask (BlockMask)
|
||||
t_scale (float): Temporal RoPE interpolation scale. <1.0 compresses positions.
|
||||
method (str): RoPE method. This release supports "linear".
|
||||
original_seq_len (int): Unused by the release linear RoPE path.
|
||||
"""
|
||||
|
||||
b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
|
||||
if cache_start is None:
|
||||
cache_start = current_start
|
||||
|
||||
# query, key, value function
|
||||
def qkv_fn(x):
|
||||
q = self.norm_q(self.q(x)).view(b, s, n, d)
|
||||
k = self.norm_k(self.k(x)).view(b, s, n, d)
|
||||
v = self.v(x).view(b, s, n, d)
|
||||
return q, k, v
|
||||
|
||||
q, k, v = qkv_fn(x)
|
||||
|
||||
if kv_cache is None:
|
||||
# Teacher-forcing training doubles sequence length with clean/noisy halves.
|
||||
is_tf = (s == seq_lens[0].item() * 2)
|
||||
if is_tf:
|
||||
q_chunk = torch.chunk(q, 2, dim=1)
|
||||
k_chunk = torch.chunk(k, 2, dim=1)
|
||||
roped_query = []
|
||||
roped_key = []
|
||||
# rope should be same for clean and noisy parts
|
||||
for ii in range(2):
|
||||
rq = sp_rope_apply(q_chunk[ii], grid_sizes, freqs, t_scale=t_scale,
|
||||
method=method, original_seq_len=original_seq_len,
|
||||
temporal_offset=temporal_offset).type_as(v)
|
||||
rk = sp_rope_apply(k_chunk[ii], grid_sizes, freqs, t_scale=t_scale,
|
||||
method=method, original_seq_len=original_seq_len,
|
||||
temporal_offset=temporal_offset).type_as(v)
|
||||
roped_query.append(rq)
|
||||
roped_key.append(rk)
|
||||
|
||||
roped_query = torch.cat(roped_query, dim=1)
|
||||
roped_key = torch.cat(roped_key, dim=1)
|
||||
|
||||
x = distributed_flex_attention(
|
||||
roped_query,
|
||||
roped_key,
|
||||
v,
|
||||
block_mask,
|
||||
)
|
||||
|
||||
else:
|
||||
roped_query = sp_rope_apply(q, grid_sizes, freqs, t_scale=t_scale,
|
||||
method=method, original_seq_len=original_seq_len,
|
||||
temporal_offset=temporal_offset).type_as(v)
|
||||
roped_key = sp_rope_apply(k, grid_sizes, freqs, t_scale=t_scale,
|
||||
method=method, original_seq_len=original_seq_len,
|
||||
temporal_offset=temporal_offset).type_as(v)
|
||||
|
||||
x = distributed_flex_attention(
|
||||
roped_query,
|
||||
roped_key,
|
||||
v,
|
||||
block_mask,
|
||||
)
|
||||
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
|
||||
# output
|
||||
x = x.flatten(2)
|
||||
x = self.o(x)
|
||||
|
||||
# Return both output and cache update info
|
||||
if kv_cache is not None:
|
||||
raise NotImplementedError()
|
||||
else:
|
||||
return x
|
||||
Reference in New Issue
Block a user