| import torch |
| import torch.nn as nn |
|
|
|
|
| class BlockWiseStateSpaceMemory(nn.Module): |
| """ |
| Paper-aligned block-wise recurrent SSM. |
| |
| This module is intentionally separate from VideoSSM hybrid. It performs a |
| recurrent state update along the latent time axis for each spatial token |
| trajectory, and is attached to selected DiT blocks. |
| """ |
|
|
| def __init__(self, dim: int): |
| super().__init__() |
| self.dim = int(dim) |
| self.in_proj = nn.Linear(self.dim, self.dim * 2) |
| self.out_proj = nn.Linear(self.dim, self.dim) |
| self.decay_logit = nn.Parameter(torch.zeros(self.dim)) |
| self.gate = nn.Parameter(torch.zeros(1)) |
|
|
| def forward(self, x: torch.Tensor, f: int, **_kwargs): |
| |
| if x is None or x.ndim != 3: |
| return x |
| b, n, d = x.shape |
| f = int(f or 0) |
| if d != self.dim or f <= 1 or n % f != 0: |
| return x |
|
|
| spatial = n // f |
| x_seq = x.reshape(b, f, spatial, d).permute(0, 2, 1, 3).reshape(b * spatial, f, d) |
| update, update_gate = self.in_proj(x_seq).chunk(2, dim=-1) |
| update = torch.tanh(update) |
| update_gate = torch.sigmoid(update_gate) |
| decay = torch.sigmoid(self.decay_logit).to(dtype=x.dtype, device=x.device).view(1, d) |
|
|
| state = torch.zeros(x_seq.shape[0], d, dtype=x.dtype, device=x.device) |
| outputs = [] |
| for t in range(f): |
| state = decay * state + (1.0 - decay) * update[:, t, :] |
| outputs.append(state * update_gate[:, t, :]) |
| y = torch.stack(outputs, dim=1) |
| y = self.out_proj(y) |
| y = y.reshape(b, spatial, f, d).permute(0, 2, 1, 3).reshape(b, n, d) |
| return x + torch.tanh(self.gate) * y |
|
|
