Spaces:
Running on Zero
Running on Zero
File size: 5,727 Bytes
6e4b62e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 | import torch
import torch.nn as nn
import torch.nn.functional as F
from diffsynth.models.memory.block_wise_ssm import BlockWiseStateSpaceMemory
from diffsynth.models.memory.videossm_hybrid import HybridStateSpaceMemory
from diffsynth.models.wan_video_dit import SelfAttention, CrossAttention, GateModule, modulate
class MLP_Action(nn.Module):
def __init__(self, out_dim, sliding_window_size=3, r=4):
super().__init__()
self.proj_action = nn.Linear(r * sliding_window_size * 10, out_dim)
nn.init.zeros_(self.proj_action.weight)
nn.init.zeros_(self.proj_action.bias)
self.sliding_window_size = sliding_window_size
self.r = r
def forward(self, x):
bs, nr, act_dim = x.shape
r = self.r
n = nr // r
actions = x.reshape(bs, n, r, act_dim)
actions = F.pad(actions, (0, 0, 0, 0, self.sliding_window_size - 1, 1), mode="replicate")
action_windows = []
for i in range(self.sliding_window_size):
action_windows.append(actions[:, i:i + n + 1])
actions = torch.cat(action_windows, dim=2)
actions = actions.reshape(bs, n + 1, -1)
actions = self.proj_action(actions)
return actions
class MLP_CamPose(nn.Module):
def __init__(self, out_dim, pose_dim=12):
super().__init__()
self.proj = nn.Linear(pose_dim, out_dim)
nn.init.zeros_(self.proj.weight)
nn.init.zeros_(self.proj.bias)
def forward(self, x):
return self.proj(x)
class DiTBlock_w_Action(nn.Module):
def __init__(self, has_image_input: bool, dim: int, num_heads: int, ffn_dim: int,
eps: float = 1e-6, add_action_attn=False,
action_use_temporal_attention: bool = True, use_cam_pose: bool = False,
use_block_wise_ssm: bool = False, use_videossm_hybrid: bool = False,
videossm_kernel_size: int = 3, videossm_expand: int = 2):
super().__init__()
self.dim = dim
self.num_heads = num_heads
self.ffn_dim = ffn_dim
if add_action_attn:
self.self_attn_with_action = SelfAttention(dim, num_heads, eps)
nn.init.zeros_(self.self_attn_with_action.o.weight)
nn.init.zeros_(self.self_attn_with_action.o.bias)
if use_cam_pose:
self.action_mlp = MLP_CamPose(dim)
else:
self.action_mlp = MLP_Action(dim)
self.self_attn = SelfAttention(dim, num_heads, eps)
self.cross_attn = CrossAttention(dim, num_heads, eps, has_image_input=has_image_input)
self.norm1 = nn.LayerNorm(dim, eps=eps, elementwise_affine=False)
self.norm2 = nn.LayerNorm(dim, eps=eps, elementwise_affine=False)
self.norm3 = nn.LayerNorm(dim, eps=eps)
self.ffn = nn.Sequential(nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'), nn.Linear(ffn_dim, dim))
self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
self.gate = GateModule()
self.action_use_temporal_attention = action_use_temporal_attention
self.use_block_wise_ssm = bool(use_block_wise_ssm)
self.use_videossm_hybrid = bool(use_videossm_hybrid)
if use_block_wise_ssm:
self.block_wise_ssm = BlockWiseStateSpaceMemory(dim)
if use_videossm_hybrid:
self.videossm_hybrid = HybridStateSpaceMemory(
dim, kernel_size=videossm_kernel_size, expand=videossm_expand
)
def forward(self, x, context, t_mod, freqs, actions=None):
has_seq = len(t_mod.shape) == 4
chunk_dim = 2 if has_seq else 1
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod).chunk(6, dim=chunk_dim)
if has_seq:
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
shift_msa.squeeze(2), scale_msa.squeeze(2), gate_msa.squeeze(2),
shift_mlp.squeeze(2), scale_mlp.squeeze(2), gate_mlp.squeeze(2),
)
num_frames = None
if actions is not None:
original_x = x
actions = self.action_mlp(actions.to(x.dtype)).to(x.dtype)
bs, num_frames, dim = actions.shape
actions = actions.reshape(bs, num_frames, 1, dim)
x = x.reshape(bs, num_frames, -1, dim)
x = x + actions
if hasattr(self, "self_attn_with_action"):
if not self.action_use_temporal_attention:
x = x.reshape(bs, -1, dim)
x = original_x + self.self_attn_with_action(x, freqs)
else:
from einops import rearrange
x = rearrange(x, "b f p d -> (b p) f d")
attn_out = self.self_attn_with_action(x)
attn_out = rearrange(attn_out, "(b p) f d -> b f p d", b=bs)
x = original_x + attn_out.reshape(bs, -1, dim)
else:
x = x.reshape(bs, -1, dim)
input_x = modulate(self.norm1(x), shift_msa, scale_msa)
x = self.gate(x, gate_msa, self.self_attn(input_x, freqs))
if num_frames is not None:
if hasattr(self, "block_wise_ssm"):
x = self.block_wise_ssm(x, f=num_frames)
if hasattr(self, "videossm_hybrid"):
spatial = x.shape[1] // int(num_frames) if int(num_frames) > 0 else 0
x = self.videossm_hybrid(x, f=num_frames, h=1, w=spatial)
x = x + self.cross_attn(self.norm3(x), context)
input_x = modulate(self.norm2(x), shift_mlp, scale_mlp)
x = self.gate(x, gate_mlp, self.ffn(input_x))
return x
|