| 670 | ) |
| 671 | |
| 672 | def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False): |
| 673 | if feat_cache is not None: |
| 674 | idx = feat_idx[0] |
| 675 | cache_x = x[:, :, -CACHE_T:, :, :].clone() |
| 676 | if cache_x.shape[2] < 2 and feat_cache[idx] is not None: |
| 677 | cache_x = torch.cat( |
| 678 | [ |
| 679 | feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( |
| 680 | cache_x.device), |
| 681 | cache_x, |
| 682 | ], |
| 683 | dim=2, |
| 684 | ) |
| 685 | x = self.conv1(x, feat_cache[idx]) |
| 686 | feat_cache[idx] = cache_x |
| 687 | feat_idx[0] += 1 |
| 688 | else: |
| 689 | x = self.conv1(x) |
| 690 | |
| 691 | for layer in self.middle: |
| 692 | if isinstance(layer, ResidualBlock) and feat_cache is not None: |
| 693 | x = layer(x, feat_cache, feat_idx) |
| 694 | else: |
| 695 | x = layer(x) |
| 696 | |
| 697 | ## upsamples |
| 698 | for layer in self.upsamples: |
| 699 | if feat_cache is not None: |
| 700 | x = layer(x, feat_cache, feat_idx, first_chunk) |
| 701 | else: |
| 702 | x = layer(x) |
| 703 | |
| 704 | ## head |
| 705 | for layer in self.head: |
| 706 | if isinstance(layer, CausalConv3d) and feat_cache is not None: |
| 707 | idx = feat_idx[0] |
| 708 | cache_x = x[:, :, -CACHE_T:, :, :].clone() |
| 709 | if cache_x.shape[2] < 2 and feat_cache[idx] is not None: |
| 710 | cache_x = torch.cat( |
| 711 | [ |
| 712 | feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( |
| 713 | cache_x.device), |
| 714 | cache_x, |
| 715 | ], |
| 716 | dim=2, |
| 717 | ) |
| 718 | x = layer(x, feat_cache[idx]) |
| 719 | feat_cache[idx] = cache_x |
| 720 | feat_idx[0] += 1 |
| 721 | else: |
| 722 | x = layer(x) |
| 723 | return x |
| 724 | |
| 725 | |
| 726 | def count_conv3d(model): |