165 lines
4.9 KiB
Python
165 lines
4.9 KiB
Python
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
|
|
import torch
|
|
import torch.cuda.amp as amp
|
|
from xfuser.core.distributed import (get_sequence_parallel_rank,
|
|
get_sequence_parallel_world_size,
|
|
get_sp_group)
|
|
from xfuser.core.long_ctx_attention import xFuserLongContextAttention
|
|
|
|
from ..modules.model import sinusoidal_embedding_1d
|
|
|
|
|
|
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
|
|
|
|
|
|
def usp_dit_forward(
|
|
self,
|
|
x,
|
|
t,
|
|
context,
|
|
seq_len,
|
|
clip_fea=None,
|
|
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 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)
|
|
|
|
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
|
|
with amp.autocast(dtype=torch.float32):
|
|
e = self.time_embedding(
|
|
sinusoidal_embedding_1d(self.freq_dim, t).float())
|
|
e0 = self.time_projection(e).unflatten(1, (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
|
|
]))
|
|
|
|
if clip_fea is not None:
|
|
context_clip = self.img_emb(clip_fea) # bs x 257 x dim
|
|
context = torch.concat([context_clip, context], dim=1)
|
|
|
|
# arguments
|
|
kwargs = dict(
|
|
e=e0,
|
|
seq_lens=seq_lens,
|
|
grid_sizes=grid_sizes,
|
|
freqs=self.freqs,
|
|
context=context,
|
|
context_lens=context_lens)
|
|
|
|
# Context Parallel
|
|
x = torch.chunk(
|
|
x, get_sequence_parallel_world_size(),
|
|
dim=1)[get_sequence_parallel_rank()]
|
|
# print('x shape:', x.shape)
|
|
|
|
iii = 0
|
|
for block in self.blocks:
|
|
iii+=1
|
|
# print('block :', iii, get_sequence_parallel_rank())
|
|
x = block(x, **kwargs)
|
|
# print('block ok :', iii, get_sequence_parallel_rank())
|
|
|
|
# head
|
|
x = self.head(x, e)
|
|
|
|
# Context Parallel
|
|
x = get_sp_group().all_gather(x, dim=1)
|
|
|
|
# unpatchify
|
|
x = self.unpatchify(x, grid_sizes)
|
|
return [u.float() for u in x]
|
|
|
|
|
|
def usp_attn_forward(self,
|
|
x,
|
|
seq_lens,
|
|
grid_sizes,
|
|
freqs,
|
|
dtype=torch.bfloat16):
|
|
# print('attn forward')
|
|
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):
|
|
# print('com q fn:', x.shape)
|
|
qq = self.q(x)
|
|
# print('qq is ok:', qq.shape)
|
|
q = self.norm_q(qq).view(b, s, n, d)
|
|
# print('q is ok:', q.shape)
|
|
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
|
|
# print('com qkv:', x.shape)
|
|
q, k, v = qkv_fn(x)
|
|
# print('qkv is ok', q.shape, k.shape, v.shape)
|
|
q = rope_apply(q, grid_sizes, freqs)
|
|
k = rope_apply(k, grid_sizes, freqs)
|
|
|
|
# TODO: We should use unpaded q,k,v for attention.
|
|
# k_lens = seq_lens // get_sequence_parallel_world_size()
|
|
# if k_lens is not None:
|
|
# q = torch.cat([u[:l] for u, l in zip(q, k_lens)]).unsqueeze(0)
|
|
# k = torch.cat([u[:l] for u, l in zip(k, k_lens)]).unsqueeze(0)
|
|
# v = torch.cat([u[:l] for u, l in zip(v, k_lens)]).unsqueeze(0)
|
|
|
|
# print('self attn coming')
|
|
x = xFuserLongContextAttention()(
|
|
None,
|
|
query=half(q),
|
|
key=half(k),
|
|
value=half(v),
|
|
window_size=self.window_size)
|
|
# print('self attn is ok')
|
|
|
|
# TODO: padding after attention.
|
|
# x = torch.cat([x, x.new_zeros(b, s - x.size(1), n, d)], dim=1)
|
|
|
|
# output
|
|
x = x.flatten(2)
|
|
x = self.o(x)
|
|
return x
|