Args: x: shape (b, c, t) cache: shape (b, c_in+c_out, t=2)
(self, x: torch.Tensor, cache: torch.Tensor = None)
| 148 | return x + h |
| 149 | |
| 150 | def forward_chunk(self, x: torch.Tensor, cache: torch.Tensor = None): |
| 151 | """ |
| 152 | Args: |
| 153 | x: shape (b, c, t) |
| 154 | cache: shape (b, c_in+c_out, t=2) |
| 155 | """ |
| 156 | cache1, cache2 = ( |
| 157 | (None, None) |
| 158 | if cache is None |
| 159 | else cache.split((self.in_channels, self.out_channels), dim=1) |
| 160 | ) |
| 161 | h = x |
| 162 | # block1 |
| 163 | h = self.block1[:4](h) |
| 164 | h, new_cache1 = self.block1[4].forward_chunk(h, cache1) |
| 165 | # block2 |
| 166 | h = self.block2[:5](h) |
| 167 | h, new_cache2 = self.block2[5].forward_chunk(h, cache2) |
| 168 | if self.in_channels != self.out_channels: |
| 169 | x = self.nin_shortcut(x) |
| 170 | new_cache = torch.cat([new_cache1, new_cache2], dim=1) |
| 171 | return x + h, new_cache |
| 172 | |
| 173 | |
| 174 | # Nonstreaming Vocos backbone based on Transformer layers |
nothing calls this directly
no test coverage detected