Spaces:
Sleeping
Sleeping
| 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 | |