| 570 | |
| 571 | |
| 572 | class StatefulCodecDecoder: |
| 573 | def __init__(self, audio_tokenizer: Any, *, n_vq: int) -> None: |
| 574 | self.audio_tokenizer = audio_tokenizer |
| 575 | self.n_vq = int(n_vq) |
| 576 | self._ctx = None |
| 577 | |
| 578 | @property |
| 579 | def device(self) -> torch.device: |
| 580 | try: |
| 581 | return next(self.audio_tokenizer.parameters()).device |
| 582 | except StopIteration: |
| 583 | return torch.device("cpu") |
| 584 | |
| 585 | def __enter__(self) -> "StatefulCodecDecoder": |
| 586 | self._ctx = self.audio_tokenizer.streaming(batch_size=1) |
| 587 | self._ctx.__enter__() |
| 588 | return self |
| 589 | |
| 590 | def __exit__(self, exc_type, exc, tb) -> None: |
| 591 | if self._ctx is not None: |
| 592 | self._ctx.__exit__(exc_type, exc, tb) |
| 593 | self._ctx = None |
| 594 | |
| 595 | @torch.inference_mode() |
| 596 | def decode_codes(self, codes: torch.LongTensor) -> torch.Tensor: |
| 597 | if codes.numel() == 0: |
| 598 | return torch.empty((2, 0), dtype=torch.float32, device=self.device) |
| 599 | if codes.ndim != 2 or int(codes.shape[1]) != self.n_vq: |
| 600 | raise ValueError(f"Expected codes shape [T, {self.n_vq}], got {tuple(codes.shape)}.") |
| 601 | codes_qbt = codes.transpose(0, 1).contiguous().unsqueeze(1).to(device=self.device, dtype=torch.long) |
| 602 | codes_lengths = torch.tensor([codes_qbt.shape[-1]], device=self.device, dtype=torch.long) |
| 603 | active_mask = torch.tensor([codes_qbt.shape[-1] > 0], device=self.device, dtype=torch.bool) |
| 604 | self.audio_tokenizer._set_streaming_exec_mask(active_mask) |
| 605 | decoded = self.audio_tokenizer._decode_frame(codes_qbt, codes_lengths) |
| 606 | if decoded.audio is None or decoded.audio_lengths is None: |
| 607 | raise RuntimeError("audio tokenizer did not return audio/audio_lengths.") |
| 608 | audio_length = int(decoded.audio_lengths[0].item()) |
| 609 | if audio_length <= 0: |
| 610 | return torch.empty( |
| 611 | (int(getattr(self.audio_tokenizer, "number_channels", 2)), 0), |
| 612 | dtype=torch.float32, |
| 613 | device=self.device, |
| 614 | ) |
| 615 | return decoded.audio[0, :, :audio_length].detach().to(torch.float32) |
| 616 | |
| 617 | |
| 618 | @torch.inference_mode() |
no outgoing calls
no test coverage detected