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