Decode SNAC in single pass to fix audio artifacts
SNAC's convolutional decoder has a receptive field spanning multiple frames. Batched decode (28 tokens at a time) created boundary discontinuities that produced muddled audio. Since pw-play waits for the full WAV anyway, single-pass decode costs only ~2s extra and produces clean audio.
This commit is contained in:
parent
538b8a513e
commit
bf0dfa7a5e
@ -24,8 +24,6 @@ SAMPLE_RATE = 24000
|
|||||||
SNAC_CODEBOOK_SIZE = 4096
|
SNAC_CODEBOOK_SIZE = 4096
|
||||||
TOKEN_PATTERN = re.compile(r"<custom_token_(\d+)>")
|
TOKEN_PATTERN = re.compile(r"<custom_token_(\d+)>")
|
||||||
TOKENS_PER_FRAME = 7
|
TOKENS_PER_FRAME = 7
|
||||||
FRAMES_PER_BATCH = 4
|
|
||||||
BATCH_SIZE = TOKENS_PER_FRAME * FRAMES_PER_BATCH # 28
|
|
||||||
|
|
||||||
ALL_VOICES = ["tara", "leah", "jess", "leo", "dan", "mia", "zac", "zoe"]
|
ALL_VOICES = ["tara", "leah", "jess", "leo", "dan", "mia", "zac", "zoe"]
|
||||||
|
|
||||||
@ -149,13 +147,15 @@ class OrpheusEngine(TTSEngine):
|
|||||||
loop = asyncio.get_running_loop()
|
loop = asyncio.get_running_loop()
|
||||||
snac = await self._get_snac()
|
snac = await self._get_snac()
|
||||||
|
|
||||||
# Stream tokens from llama-server via SSE
|
# Stream tokens from llama-server via SSE, then decode in one pass.
|
||||||
|
# SNAC's convolutional decoder has a receptive field spanning multiple
|
||||||
|
# frames, so batched decode creates boundary artifacts. Since pw-play
|
||||||
|
# waits for the full WAV anyway, single-pass decode costs only ~2s
|
||||||
|
# extra on a 50s generation — and produces clean audio.
|
||||||
t0 = time.time()
|
t0 = time.time()
|
||||||
token_strings: list[str] = []
|
token_strings: list[str] = []
|
||||||
audio_chunks: list[np.ndarray] = []
|
|
||||||
total_tokens = 0
|
total_tokens = 0
|
||||||
dropped_lines = 0
|
dropped_lines = 0
|
||||||
stream_interrupted = False
|
|
||||||
|
|
||||||
try:
|
try:
|
||||||
async with self._client.stream(
|
async with self._client.stream(
|
||||||
@ -202,27 +202,16 @@ class OrpheusEngine(TTSEngine):
|
|||||||
token_strings.append(f"<custom_token_{m}>")
|
token_strings.append(f"<custom_token_{m}>")
|
||||||
total_tokens += 1
|
total_tokens += 1
|
||||||
|
|
||||||
# Batch decode every 28 tokens (4 SNAC frames)
|
|
||||||
while len(token_strings) >= BATCH_SIZE:
|
|
||||||
batch = token_strings[:BATCH_SIZE]
|
|
||||||
token_strings = token_strings[BATCH_SIZE:]
|
|
||||||
chunk_audio = await loop.run_in_executor(
|
|
||||||
None, _tokens_to_audio, batch, snac
|
|
||||||
)
|
|
||||||
if chunk_audio is not None:
|
|
||||||
audio_chunks.append(chunk_audio)
|
|
||||||
|
|
||||||
except httpx.ConnectError as e:
|
except httpx.ConnectError as e:
|
||||||
raise RuntimeError(f"Cannot reach llama-server at {self._url}: {e}") from e
|
raise RuntimeError(f"Cannot reach llama-server at {self._url}: {e}") from e
|
||||||
except (httpx.ReadError, httpx.RemoteProtocolError) as e:
|
except (httpx.ReadError, httpx.RemoteProtocolError) as e:
|
||||||
if not audio_chunks:
|
if total_tokens < TOKENS_PER_FRAME:
|
||||||
raise RuntimeError(
|
raise RuntimeError(
|
||||||
f"llama-server connection lost with no audio decoded: {e}"
|
f"llama-server connection lost with no usable tokens: {e}"
|
||||||
) from e
|
) from e
|
||||||
stream_interrupted = True
|
|
||||||
print(
|
print(
|
||||||
f" Warning: llama-server connection lost after {total_tokens} tokens. "
|
f" Warning: llama-server connection lost after {total_tokens} tokens. "
|
||||||
f"Using {len(audio_chunks)} partial chunks.",
|
"Decoding what we have.",
|
||||||
file=sys.stderr,
|
file=sys.stderr,
|
||||||
)
|
)
|
||||||
|
|
||||||
@ -234,32 +223,29 @@ class OrpheusEngine(TTSEngine):
|
|||||||
file=sys.stderr,
|
file=sys.stderr,
|
||||||
)
|
)
|
||||||
|
|
||||||
# Flush remaining tokens (drop partial frame — at most 0.29ms lost)
|
# Drop partial frame at end (at most 6 tokens / 0.29ms lost)
|
||||||
if token_strings:
|
|
||||||
usable = len(token_strings) - (len(token_strings) % TOKENS_PER_FRAME)
|
usable = len(token_strings) - (len(token_strings) % TOKENS_PER_FRAME)
|
||||||
if usable > 0:
|
|
||||||
chunk_audio = await loop.run_in_executor(
|
|
||||||
None, _tokens_to_audio, token_strings[:usable], snac
|
|
||||||
)
|
|
||||||
if chunk_audio is not None:
|
|
||||||
audio_chunks.append(chunk_audio)
|
|
||||||
|
|
||||||
num_frames = total_tokens // TOKENS_PER_FRAME
|
num_frames = usable // TOKENS_PER_FRAME
|
||||||
tok_per_sec = total_tokens / gen_time if gen_time > 0 else 0
|
tok_per_sec = total_tokens / gen_time if gen_time > 0 else 0
|
||||||
status = " (TRUNCATED)" if stream_interrupted else ""
|
|
||||||
print(
|
print(
|
||||||
f" Orpheus: {total_tokens} tokens ({num_frames} frames) "
|
f" Orpheus: {total_tokens} tokens ({num_frames} frames) "
|
||||||
f"in {gen_time:.1f}s ({tok_per_sec:.1f} tok/s){status}",
|
f"in {gen_time:.1f}s ({tok_per_sec:.1f} tok/s)",
|
||||||
file=sys.stderr,
|
file=sys.stderr,
|
||||||
)
|
)
|
||||||
|
|
||||||
if not audio_chunks:
|
if usable < TOKENS_PER_FRAME:
|
||||||
raise RuntimeError(
|
raise RuntimeError(
|
||||||
f"Orpheus returned insufficient tokens ({total_tokens}). "
|
f"Orpheus returned insufficient tokens ({total_tokens}). "
|
||||||
"Check llama-server logs."
|
"Check llama-server logs."
|
||||||
)
|
)
|
||||||
|
|
||||||
audio = np.concatenate(audio_chunks)
|
# Single-pass SNAC decode — full context across all frames
|
||||||
|
audio = await loop.run_in_executor(
|
||||||
|
None, _tokens_to_audio, token_strings[:usable], snac
|
||||||
|
)
|
||||||
|
if audio is None:
|
||||||
|
raise RuntimeError("SNAC decoding produced no audio")
|
||||||
path = write_wav(audio, SAMPLE_RATE, prefix="orpheus-")
|
path = write_wav(audio, SAMPLE_RATE, prefix="orpheus-")
|
||||||
|
|
||||||
return TTSResult(
|
return TTSResult(
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user