Args: x: input sequence at 50Hz, length should be multiples of 4
(
self,
x: torch.Tensor,
# Upsample conv cache
up_conv_cache: torch.Tensor = None,
# Backbone conv cache
bb_conv_cache1: torch.Tensor = None,
bb_conv_cache2: torch.Tensor = None,
# Backbone attention cache
bb_kv_cache: torch.Tensor = None,
# iSTFT cache
is_cache: torch.Tensor = None,
last_chunk: bool = False,
)
| 655 | return x, new_cache |
| 656 | |
| 657 | def forward_chunk( |
| 658 | self, |
| 659 | x: torch.Tensor, |
| 660 | # Upsample conv cache |
| 661 | up_conv_cache: torch.Tensor = None, |
| 662 | # Backbone conv cache |
| 663 | bb_conv_cache1: torch.Tensor = None, |
| 664 | bb_conv_cache2: torch.Tensor = None, |
| 665 | # Backbone attention cache |
| 666 | bb_kv_cache: torch.Tensor = None, |
| 667 | # iSTFT cache |
| 668 | is_cache: torch.Tensor = None, |
| 669 | last_chunk: bool = False, |
| 670 | ): |
| 671 | """ |
| 672 | Args: |
| 673 | x: input sequence at 50Hz, length should be multiples of 4 |
| 674 | """ |
| 675 | assert ( |
| 676 | self.causal |
| 677 | ), "Only AcousticDecoder with causal=True supports forward_chunk method." |
| 678 | |
| 679 | x = x.transpose(1, 2) |
| 680 | x, new_up_conv_cache = self.forward_upsample_conv_chunk(x, up_conv_cache) |
| 681 | x = x.transpose(1, 2) |
| 682 | # Backbone |
| 683 | x, new_bb_conv_cache1, new_bb_conv_cache2, new_bb_kv_cache = ( |
| 684 | self.backbone.forward_chunk( |
| 685 | x, |
| 686 | bb_conv_cache1, |
| 687 | bb_conv_cache2, |
| 688 | bb_kv_cache, |
| 689 | ) |
| 690 | ) |
| 691 | # iSTFT |
| 692 | y, new_is_cache = self.isift.forward_chunk(x, is_cache, last_chunk) |
| 693 | return ( |
| 694 | y, |
| 695 | new_up_conv_cache, |
| 696 | new_bb_conv_cache1, |
| 697 | new_bb_conv_cache2, |
| 698 | new_bb_kv_cache, |
| 699 | new_is_cache, |
| 700 | ) |
nothing calls this directly
no test coverage detected