Class AudioFrameDecoder

moss_tts_realtime/app.py:200–234

class AudioFrameDecoder:
    def __init__(
        self,
        decoder: AudioStreamDecoder,
        codebook_size: int,
        audio_eos_token: int,
    ):
        self.decoder = decoder
        self.codebook_size = codebook_size
        self.audio_eos_token = audio_eos_token

    def decode_frames(self, audio_frames: list[torch.Tensor]) -> Iterator[np.ndarray]:
        for frame in audio_frames:
            tokens = frame
            if tokens.dim() == 3:
                tokens = tokens[0]
            if tokens.dim() != 2:
                raise ValueError(f"Expected [T, C] audio tokens, got {tuple(tokens.shape)}")
            tokens, stop = _sanitize_tokens(tokens, self.codebook_size, self.audio_eos_token)
            if tokens.numel() == 0:
                if stop:
                    break
                continue
            self.decoder.push_tokens(tokens.detach())
            for wav in self.decoder.audio_chunks():
                if wav.numel() == 0:
                    continue
                yield wav.detach().cpu().numpy().reshape(-1)
            if stop:
                break

    def flush(self) -> Iterator[np.ndarray]:
        final_chunk = self.decoder.flush()
        if final_chunk is not None and final_chunk.numel() > 0:
            yield final_chunk.detach().cpu().numpy().reshape(-1)


class StreamAudioEmitter: