(self, x, feat_cache=None, feat_idx=[0])
| 204 | if in_dim != out_dim else nn.Identity() |
| 205 | |
| 206 | def forward(self, x, feat_cache=None, feat_idx=[0]): |
| 207 | h = self.shortcut(x) |
| 208 | for layer in self.residual: |
| 209 | if isinstance(layer, CausalConv3d) and feat_cache is not None: |
| 210 | idx = feat_idx[0] |
| 211 | cache_x = x[:, :, -CACHE_T:, :, :].clone() |
| 212 | if cache_x.shape[2] < 2 and feat_cache[idx] is not None: |
| 213 | # cache last frame of last two chunk |
| 214 | cache_x = torch.cat([ |
| 215 | feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( |
| 216 | cache_x.device), cache_x |
| 217 | ], |
| 218 | dim=2) |
| 219 | x = layer(x, feat_cache[idx]) |
| 220 | feat_cache[idx] = cache_x |
| 221 | feat_idx[0] += 1 |
| 222 | else: |
| 223 | x = layer(x) |
| 224 | return x + h |
| 225 | |
| 226 | |
| 227 | class AttentionBlock(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected