(self, x, feat_cache=None, feat_idx=[0])
| 316 | CausalConv3d(out_dim, z_dim, 3, padding=1)) |
| 317 | |
| 318 | def forward(self, x, feat_cache=None, feat_idx=[0]): |
| 319 | if feat_cache is not None: |
| 320 | idx = feat_idx[0] |
| 321 | cache_x = x[:, :, -CACHE_T:, :, :].clone() |
| 322 | if cache_x.shape[2] < 2 and feat_cache[idx] is not None: |
| 323 | # cache last frame of last two chunk |
| 324 | cache_x = torch.cat([ |
| 325 | feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( |
| 326 | cache_x.device), cache_x |
| 327 | ], |
| 328 | dim=2) |
| 329 | x = self.conv1(x, feat_cache[idx]) |
| 330 | feat_cache[idx] = cache_x |
| 331 | feat_idx[0] += 1 |
| 332 | else: |
| 333 | x = self.conv1(x) |
| 334 | |
| 335 | ## downsamples |
| 336 | for layer in self.downsamples: |
| 337 | if feat_cache is not None: |
| 338 | x = layer(x, feat_cache, feat_idx) |
| 339 | else: |
| 340 | x = layer(x) |
| 341 | |
| 342 | ## middle |
| 343 | for layer in self.middle: |
| 344 | if isinstance(layer, ResidualBlock) and feat_cache is not None: |
| 345 | x = layer(x, feat_cache, feat_idx) |
| 346 | else: |
| 347 | x = layer(x) |
| 348 | |
| 349 | ## head |
| 350 | for layer in self.head: |
| 351 | if isinstance(layer, CausalConv3d) and feat_cache is not None: |
| 352 | idx = feat_idx[0] |
| 353 | cache_x = x[:, :, -CACHE_T:, :, :].clone() |
| 354 | if cache_x.shape[2] < 2 and feat_cache[idx] is not None: |
| 355 | # cache last frame of last two chunk |
| 356 | cache_x = torch.cat([ |
| 357 | feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( |
| 358 | cache_x.device), cache_x |
| 359 | ], |
| 360 | dim=2) |
| 361 | x = layer(x, feat_cache[idx]) |
| 362 | feat_cache[idx] = cache_x |
| 363 | feat_idx[0] += 1 |
| 364 | else: |
| 365 | x = layer(x) |
| 366 | return x |
| 367 | |
| 368 | |
| 369 | class Decoder3d(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected