(self, x, feat_cache=None, feat_idx=[0])
| 421 | CausalConv3d(out_dim, 3, 3, padding=1)) |
| 422 | |
| 423 | def forward(self, x, feat_cache=None, feat_idx=[0]): |
| 424 | ## conv1 |
| 425 | if feat_cache is not None: |
| 426 | idx = feat_idx[0] |
| 427 | cache_x = x[:, :, -CACHE_T:, :, :].clone() |
| 428 | if cache_x.shape[2] < 2 and feat_cache[idx] is not None: |
| 429 | # cache last frame of last two chunk |
| 430 | cache_x = torch.cat([ |
| 431 | feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( |
| 432 | cache_x.device), cache_x |
| 433 | ], |
| 434 | dim=2) |
| 435 | x = self.conv1(x, feat_cache[idx]) |
| 436 | feat_cache[idx] = cache_x |
| 437 | feat_idx[0] += 1 |
| 438 | else: |
| 439 | x = self.conv1(x) |
| 440 | |
| 441 | ## middle |
| 442 | for layer in self.middle: |
| 443 | if isinstance(layer, ResidualBlock) and feat_cache is not None: |
| 444 | x = layer(x, feat_cache, feat_idx) |
| 445 | else: |
| 446 | x = layer(x) |
| 447 | |
| 448 | ## upsamples |
| 449 | for layer in self.upsamples: |
| 450 | if feat_cache is not None: |
| 451 | x = layer(x, feat_cache, feat_idx) |
| 452 | else: |
| 453 | x = layer(x) |
| 454 | |
| 455 | ## head |
| 456 | for layer in self.head: |
| 457 | if isinstance(layer, CausalConv3d) and feat_cache is not None: |
| 458 | idx = feat_idx[0] |
| 459 | cache_x = x[:, :, -CACHE_T:, :, :].clone() |
| 460 | if cache_x.shape[2] < 2 and feat_cache[idx] is not None: |
| 461 | # cache last frame of last two chunk |
| 462 | cache_x = torch.cat([ |
| 463 | feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( |
| 464 | cache_x.device), cache_x |
| 465 | ], |
| 466 | dim=2) |
| 467 | x = layer(x, feat_cache[idx]) |
| 468 | feat_cache[idx] = cache_x |
| 469 | feat_idx[0] += 1 |
| 470 | else: |
| 471 | x = layer(x) |
| 472 | return x |
| 473 | |
| 474 | |
| 475 | def count_conv3d(model): |
nothing calls this directly
no outgoing calls
no test coverage detected