| import logging |
| import math |
|
|
| import einops |
| import torch |
| from einops import rearrange |
| from timm.models.layers.drop import DropPath |
| from torch import nn |
| from torch.nn import LayerNorm, Linear, MultiheadAttention |
|
|
| logger = logging.getLogger(__name__) |
|
|
|
|
| class STAdapter(nn.Module): |
| """ST Adapter""" |
|
|
| def __init__( |
| self, |
| kernel_size=(3, 3, 3), |
| input_dim=768, |
| hidden_dim=384, |
| img_size=224, |
| patch_size=16, |
| drop_prob=0.1, |
| ): |
| super(STAdapter, self).__init__() |
| self.kernel_size = kernel_size |
| self.input_dim = input_dim |
| self.hidden_dim = hidden_dim |
|
|
| self.h = self.w = img_size // patch_size |
|
|
| self.linear1 = nn.Linear(input_dim, hidden_dim) |
| self.linear2 = nn.Linear(hidden_dim, input_dim) |
| self.act = nn.ReLU() |
| self.conv = nn.Conv3d( |
| hidden_dim, hidden_dim, kernel_size=kernel_size, padding="same", groups=hidden_dim |
| ) |
| self.droppath = DropPath(drop_prob=drop_prob) |
|
|
| self.scale = nn.parameter.Parameter(torch.zeros([])) |
|
|
| def forward(self, x: torch.Tensor): |
| """forward |
| |
| Args: |
| x (torch.Tensor): input features. Shape: [bs, nframes, l, c]. l = 1 + h*w |
| |
| Returns: features after adapter. The same shape as input. |
| |
| """ |
| if x.shape[1] == 1: |
| return x |
|
|
| shortcut = x |
| x = self.linear1(x) |
| cls = x[:, :, :1, :] |
| tokens = x[:, :, 1:, :] |
| tokens = einops.rearrange(tokens, "b t (h w) c -> b c t h w", h=self.h).contiguous() |
| tokens = self.conv(tokens) |
| tokens = einops.rearrange(tokens, "b c t h w -> b t (h w) c") |
| x = torch.cat([cls, tokens], dim=2) |
| x = self.act(x) |
| x = self.linear2(x) |
|
|
| return shortcut + self.scale * self.droppath(x) |
|
|
|
|
| class TemporalAttention(nn.Module): |
|
|
| """perform temporal self-attention""" |
|
|
| def __init__(self, input_dim=768, droppath_rate=0.1): |
| """ |
| |
| Kwargs: |
| input_dim (int): The input feature dimension. |
| |
| |
| """ |
| super().__init__() |
|
|
| self._input_dim = input_dim |
| self.temporal_attn = MultiheadAttention(input_dim, num_heads=input_dim // 64) |
| self.norm = LayerNorm(input_dim, eps=1e-12) |
| self.linear = Linear(input_dim, input_dim) |
| self.droppath = DropPath(droppath_rate) |
| self.scale = nn.parameter.Parameter(torch.zeros([])) |
|
|
| def forward(self, x: torch.Tensor): |
| """forward |
| |
| Args: |
| x (torch.Tensor): input features. Shape: [bs, nframes, l, c]. l = 1 + h*w |
| |
| Returns: features after adapter. The same shape as input. |
| |
| """ |
| if x.shape[1] == 1: |
| return x |
|
|
| shortcut = x |
| x = einops.rearrange(x, "b t l c -> t (b l) c") |
| x = self.norm(x) |
| x = self.temporal_attn(x, x, x)[0] |
| x = einops.rearrange(x, "t (b l) c -> b t l c", b=shortcut.shape[0]) |
| return shortcut + self.scale * self.droppath(x) |
|
|
|
|
| class WindowTemporalAttention(nn.Module): |
|
|
| """perform windowed temporal self-attention""" |
|
|
| def __init__(self, input_dim=768, droppath_rate=0.1, window_size=(2, 2)): |
| """ |
| |
| Kwargs: |
| input_dim (int): The input feature dimension. |
| |
| |
| """ |
| super().__init__() |
|
|
| self._input_dim = input_dim |
| self.temporal_attn = MultiheadAttention(input_dim, num_heads=input_dim // 64) |
| self.norm = LayerNorm(input_dim, eps=1e-12) |
| self.droppath = DropPath(droppath_rate) |
| self.scale = nn.parameter.Parameter(torch.zeros([])) |
| self.wh, self.ww = window_size |
| |
|
|
| def forward(self, x: torch.Tensor): |
| """forward |
| |
| Args: |
| x (torch.Tensor): input features. Shape: [bs, nframes, l, c]. l = 1 + h*w |
| |
| Returns: features after adapter. The same shape as input. |
| |
| """ |
| if x.shape[1] == 1: |
| return x |
| shortcut = x |
|
|
| h = w = int(math.sqrt(x.shape[2] - 1)) |
| cls_token = x[:, :, :1, :] |
| x = einops.rearrange( |
| x[:, :, 1:, :], |
| "b t (nh wh nw ww) c -> (t wh ww) (b nh nw) c", |
| nh=h // self.wh, |
| wh=self.wh, |
| nw=w // self.ww, |
| ww=self.ww, |
| ) |
| x = self.norm(x) |
| x = self.temporal_attn(x, x, x)[0] |
| x = einops.rearrange( |
| x, |
| "(t wh ww) (b nh nw) c -> b t (nh wh nw ww) c", |
| wh=self.wh, |
| ww=self.ww, |
| nh=h // self.wh, |
| nw=w // self.ww, |
| ) |
| |
| x = torch.concat([cls_token, x], dim=2) |
| return shortcut + self.scale * self.droppath(x) |
|
|
|
|
| class X_CLIP(nn.Module): |
|
|
| """perform windowed temporal self-attention""" |
|
|
| def __init__(self, input_dim=768, droppath_rate=0.1, num_prompts=1): |
| """ |
| |
| Kwargs: |
| input_dim (int): The input feature dimension. |
| |
| |
| """ |
| super().__init__() |
|
|
| d_model = input_dim |
|
|
| self.message_fc = nn.Linear(d_model, d_model) |
| self.message_ln = LayerNorm(d_model, eps=1e-12) |
| self.message_attn = nn.MultiheadAttention(d_model, d_model // 64) |
| self.num_prompts = num_prompts |
|
|
| self.droppath = DropPath(droppath_rate) |
|
|
| def forward(self, x: torch.Tensor): |
| """forward |
| |
| Args: |
| x (torch.Tensor): input features. Shape: [bs, nframes, l, c]. l = 1 + h*w |
| |
| Returns: features after adapter. The same shape as input. |
| |
| """ |
| if x.shape[1] == 1: |
| return x |
| msg_token = self.message_ln(self.message_fc(x[:, :, 0, :])) |
| msg_token = rearrange(msg_token, "b t c -> t b c") |
| msg_token = msg_token + self.droppath( |
| self.message_attn(msg_token, msg_token, msg_token)[0] |
| ) |
| msg_token = rearrange(msg_token, "t b c -> b t c") |
| |
| x = torch.cat([x[:, :, :-1, :], msg_token.unsqueeze(2)], dim=2) |
| return x |
|
|
