| 109 | |
| 110 | |
| 111 | class AudioFrameDecoder: |
| 112 | def __init__( |
| 113 | self, |
| 114 | decoder: AudioStreamDecoder, |
| 115 | codebook_size: int, |
| 116 | audio_eos_token: int, |
| 117 | callbacks: Optional[StreamingCallbacks] = None, |
| 118 | ): |
| 119 | self.decoder = decoder |
| 120 | self.codebook_size = codebook_size |
| 121 | self.audio_eos_token = audio_eos_token |
| 122 | self.callbacks = callbacks or StreamingCallbacks() |
| 123 | self._started = False |
| 124 | self._finished = False |
| 125 | |
| 126 | def _mark_started(self) -> None: |
| 127 | if self._started: |
| 128 | return |
| 129 | self._started = True |
| 130 | if self.callbacks.on_audio_stream_start: |
| 131 | self.callbacks.on_audio_stream_start() |
| 132 | |
| 133 | def finish(self) -> None: |
| 134 | if self._finished: |
| 135 | return |
| 136 | self._finished = True |
| 137 | if self._started and self.callbacks.on_audio_stream_stop: |
| 138 | self.callbacks.on_audio_stream_stop() |
| 139 | |
| 140 | def decode_frames(self, audio_frames: list[torch.Tensor]) -> Iterator[np.ndarray]: |
| 141 | for frame in audio_frames: |
| 142 | tokens = frame |
| 143 | if tokens.dim() == 3: |
| 144 | tokens = tokens[0] |
| 145 | if tokens.dim() != 2: |
| 146 | raise ValueError(f"Expected [T, C] audio tokens, got {tuple(tokens.shape)}") |
| 147 | tokens, _ = _sanitize_tokens(tokens, self.codebook_size, self.audio_eos_token) |
| 148 | if tokens.numel() == 0: |
| 149 | continue |
| 150 | self.decoder.push_tokens(tokens.detach()) |
| 151 | for wav in self.decoder.audio_chunks(): |
| 152 | if wav.numel() == 0: |
| 153 | continue |
| 154 | self._mark_started() |
| 155 | yield wav.detach().cpu().numpy().reshape(-1) |
| 156 | |
| 157 | def flush(self) -> Iterator[np.ndarray]: |
| 158 | final_chunk = self.decoder.flush() |
| 159 | if final_chunk is not None and final_chunk.numel() > 0: |
| 160 | self._mark_started() |
| 161 | yield final_chunk.detach().cpu().numpy().reshape(-1) |
| 162 | self.finish() |
| 163 | |
| 164 | |
| 165 | def _maybe_wait_for_buffer(buffer_tracker: BufferedAudioTracker, threshold_seconds: float) -> None: |