File size: 1,448 Bytes
00c7b31 | 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 | import torch
import torch.nn as nn
class HybridStateSpaceMemory(nn.Module):
"""
Lightweight legacy VideoSSM hybrid block:
depthwise temporal conv over per-spatial token trajectories.
"""
def __init__(self, dim: int, kernel_size: int = 3, expand: int = 2):
super().__init__()
self.dim = int(dim)
self.kernel_size = int(kernel_size)
hidden = int(dim) * max(int(expand), 1)
pad = self.kernel_size - 1 # causal-like left padding
self.in_proj = nn.Linear(self.dim, hidden)
self.dw = nn.Conv1d(hidden, hidden, kernel_size=self.kernel_size, groups=hidden, padding=pad)
self.out_proj = nn.Linear(hidden, self.dim)
self.gate = nn.Parameter(torch.zeros(1))
def forward(self, x: torch.Tensor, f: int, h: int, w: int, **_kwargs):
# x: (B, F*H*W, D)
if x is None or x.ndim != 3:
return x
b, n, d = x.shape
f = int(f)
hw = int(h) * int(w)
if d != self.dim or f <= 1 or n != f * hw:
return x
x4 = x.reshape(b, f, hw, d).permute(0, 2, 1, 3).reshape(b * hw, f, d) # (B*HW, F, D)
y = self.in_proj(x4)
y = y.transpose(1, 2) # (B*HW, hidden, F)
y = self.dw(y)[..., :f] # causal crop
y = y.transpose(1, 2)
y = self.out_proj(y)
y = y.reshape(b, hw, f, d).permute(0, 2, 1, 3).reshape(b, n, d)
return x + torch.tanh(self.gate) * y
|