(self, x, feat_cache=None, feat_idx=[0])
| 787 | CausalConv3d(out_dim, 3, 3, padding=1)) |
| 788 | |
| 789 | def forward(self, x, feat_cache=None, feat_idx=[0]): |
| 790 | ## conv1 |
| 791 | if feat_cache is not None: |
| 792 | idx = feat_idx[0] |
| 793 | cache_x = x[:, :, -CACHE_T:, :, :].clone() |
| 794 | if cache_x.shape[2] < 2 and feat_cache[idx] is not None: |
| 795 | # cache last frame of last two chunk |
| 796 | cache_x = torch.cat([ |
| 797 | feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( |
| 798 | cache_x.device), cache_x |
| 799 | ], |
| 800 | dim=2) |
| 801 | x = self.conv1(x, feat_cache[idx]) |
| 802 | feat_cache[idx] = cache_x |
| 803 | feat_idx[0] += 1 |
| 804 | else: |
| 805 | x = self.conv1(x) |
| 806 | |
| 807 | ## middle |
| 808 | for layer in self.middle: |
| 809 | if check_is_instance(layer, ResidualBlock) and feat_cache is not None: |
| 810 | x = layer(x, feat_cache, feat_idx) |
| 811 | else: |
| 812 | x = layer(x) |
| 813 | |
| 814 | ## upsamples |
| 815 | for layer in self.upsamples: |
| 816 | if feat_cache is not None: |
| 817 | x = layer(x, feat_cache, feat_idx) |
| 818 | else: |
| 819 | x = layer(x) |
| 820 | |
| 821 | ## head |
| 822 | for layer in self.head: |
| 823 | if check_is_instance(layer, CausalConv3d) and feat_cache is not None: |
| 824 | idx = feat_idx[0] |
| 825 | cache_x = x[:, :, -CACHE_T:, :, :].clone() |
| 826 | if cache_x.shape[2] < 2 and feat_cache[idx] is not None: |
| 827 | # cache last frame of last two chunk |
| 828 | cache_x = torch.cat([ |
| 829 | feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( |
| 830 | cache_x.device), cache_x |
| 831 | ], |
| 832 | dim=2) |
| 833 | x = layer(x, feat_cache[idx]) |
| 834 | feat_cache[idx] = cache_x |
| 835 | feat_idx[0] += 1 |
| 836 | else: |
| 837 | x = layer(x) |
| 838 | return x |
| 839 | |
| 840 | |
| 841 |
nothing calls this directly
no test coverage detected