codec: add chunked decode with rolling left context

This commit is contained in:
Pascal
2026-05-14 22:53:28 +02:00
parent d99eececc8
commit dda50c2225
7 changed files with 473 additions and 137 deletions
+124 -23
View File
@@ -11,6 +11,7 @@
#include <cstdlib>
#include <cstring>
#include <string>
#include <vector>
#if defined(_WIN32)
# include <fcntl.h>
@@ -332,44 +333,144 @@ static std::string audio_encode_wav(const float * audio, int T_audio, int sr, Wa
return {};
}
// Write mono float audio to WAV file in the requested format. path "-"
// streams the encoded WAV to stdout (pipe friendly). S16/S24 hard clip
// to [-1, +1], F32 preserves the full range.
// Write mono float audio to WAV file in the requested format.
// S16/S24 hard clip to [-1, +1], F32 preserves the full range.
static bool audio_write_wav(const char * path, const float * audio, int T_audio, int sr, WavFormat fmt = WAV_S16) {
std::string wav = audio_encode_wav(audio, T_audio, sr, fmt);
if (wav.empty()) {
return false;
}
const bool to_stdout = (path[0] == '-' && path[1] == '\0');
FILE * fp = to_stdout ? stdout : utf8_fopen(path, "wb");
FILE * fp = utf8_fopen(path, "wb");
if (!fp) {
fprintf(stderr, "[WAV] Cannot open %s for writing\n", path);
return false;
}
#if defined(_WIN32)
// stdout defaults to text mode on Windows; binary mode is mandatory
// for WAV bytes to survive without CRLF translation. The mode is set
// once per process and is harmless on the second call.
if (to_stdout) {
_setmode(_fileno(stdout), _O_BINARY);
}
#endif
if (fwrite(wav.data(), 1, wav.size(), fp) != wav.size()) {
fprintf(stderr, "[WAV] Failed to write %s\n", path);
if (!to_stdout) {
fclose(fp);
}
fclose(fp);
return false;
}
if (to_stdout) {
fflush(fp);
} else {
fclose(fp);
}
fclose(fp);
const char * fmt_name = (fmt == WAV_S16) ? "S16" : (fmt == WAV_S24) ? "S24" : "F32";
fprintf(stderr, "[WAV] Wrote %s: %d samples, %d Hz, mono %s\n", to_stdout ? "<stdout>" : path, T_audio, sr,
fmt_name);
fprintf(stderr, "[WAV] Wrote %s: %d samples, %d Hz, mono %s\n", path, T_audio, sr, fmt_name);
return true;
}
// Minimal streaming WAV sink. Writes a wide RIFF / data size at open and
// never updates them: the stream is one shot, non seekable, suitable for
// stdout pipes where the player reads until EOF. Use audio_write_wav for
// seekable file output (the file there has accurate sizes in headers).
struct wav_stream {
FILE * fp;
WavFormat fmt;
int sr;
};
// Open a streaming WAV sink on stdout. Switches stdout to binary mode on
// Windows. Header advertises 0x7FFFFFFF for both RIFF chunk size and data
// chunk size, the conventional "unknown / live" marker that aplay, ffmpeg
// and most players accept by reading until EOF.
static bool wav_stream_open_stdout(wav_stream * ws, int sr, WavFormat fmt) {
ws->fp = stdout;
ws->fmt = fmt;
ws->sr = sr;
#if defined(_WIN32)
_setmode(_fileno(stdout), _O_BINARY);
#endif
int bits = (fmt == WAV_S16) ? 16 : (fmt == WAV_S24) ? 24 : 32;
uint16_t fmt_tag = (fmt == WAV_F32) ? 3 : 1;
int n_channels = 1;
uint32_t bytes_per_sample = (uint32_t) bits / 8;
uint32_t byte_rate = (uint32_t) sr * (uint32_t) n_channels * bytes_per_sample;
uint16_t block_align = (uint16_t) (n_channels * (int) bytes_per_sample);
uint32_t data_size = 0x7FFFFFFFu;
uint32_t file_size = 0x7FFFFFFFu;
char header[44];
char * p = header;
memcpy(p, "RIFF", 4);
p += 4;
wav_write_u32le(p, file_size);
memcpy(p, "WAVE", 4);
p += 4;
memcpy(p, "fmt ", 4);
p += 4;
wav_write_u32le(p, 16);
wav_write_u16le(p, fmt_tag);
wav_write_u16le(p, (uint16_t) n_channels);
wav_write_u32le(p, (uint32_t) sr);
wav_write_u32le(p, byte_rate);
wav_write_u16le(p, block_align);
wav_write_u16le(p, (uint16_t) bits);
memcpy(p, "data", 4);
p += 4;
wav_write_u32le(p, data_size);
if (fwrite(header, 1, 44, ws->fp) != 44) {
fprintf(stderr, "[WAV-Stream] header write failed\n");
return false;
}
fflush(ws->fp);
const char * fmt_name = (fmt == WAV_S16) ? "S16" : (fmt == WAV_S24) ? "S24" : "F32";
fprintf(stderr, "[WAV-Stream] stdout: %d Hz, mono %s\n", sr, fmt_name);
return true;
}
// Encode and write n mono samples to the streaming sink. NaN and Inf coerce
// to zero. S16 / S24 clamp to [-1, +1] before quantisation. Flushes after
// every write so a downstream pipe sees the bytes immediately.
static bool wav_stream_write(wav_stream * ws, const float * audio, int n) {
if (n <= 0) {
return true;
}
if (ws->fmt == WAV_S16) {
std::vector<uint8_t> out((size_t) n * 2);
char * p = (char *) out.data();
for (int t = 0; t < n; t++) {
int16_t s = (int16_t) (wav_clamp1(wav_sanitize(audio[t])) * 32767.0f);
wav_write_u16le(p, (uint16_t) s);
}
if (fwrite(out.data(), 1, out.size(), ws->fp) != out.size()) {
return false;
}
} else if (ws->fmt == WAV_S24) {
std::vector<uint8_t> out((size_t) n * 3);
char * p = (char *) out.data();
for (int t = 0; t < n; t++) {
int32_t s = (int32_t) (wav_clamp1(wav_sanitize(audio[t])) * 8388607.0f);
wav_write_u24le(p, (uint32_t) s);
}
if (fwrite(out.data(), 1, out.size(), ws->fp) != out.size()) {
return false;
}
} else {
std::vector<uint8_t> out((size_t) n * 4);
char * p = (char *) out.data();
for (int t = 0; t < n; t++) {
float f = wav_sanitize(audio[t]);
uint32_t u;
memcpy(&u, &f, 4);
wav_write_u32le(p, u);
}
if (fwrite(out.data(), 1, out.size(), ws->fp) != out.size()) {
return false;
}
}
fflush(ws->fp);
return true;
}
// Final flush. The sink does not own the stdout FILE so no fclose is issued.
static void wav_stream_close(wav_stream * ws) {
fflush(ws->fp);
}
+169
View File
@@ -0,0 +1,169 @@
#pragma once
// codec-chunked-decode.h: bounded VRAM codec decode with rolling left
// context. Strict equivalent of the upstream Qwen3-TTS 12 Hz tokenizer
// chunked_decode entry
// (qwen_tts/core/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py
// line 886).
//
// The codec decoder runs a causal Conv1d, a sliding window causal
// transformer, an upsample stage and a DAC decoder. Decoding a chunk
// of frames in isolation introduces edge artefacts at the chunk
// boundary because the causal conv kernel and the transformer attention
// have no left context to draw from. Prepending left_ctx_frames
// previously decoded frames and stripping the resulting samples after
// the decode restores continuity.
//
// Two entry points:
//
// codec_chunked_decode : one shot decode of a full codes buffer.
// Bit perfect equivalent of pipeline_codec_decode when the audio
// fits in a single chunk_frames sized window. Bounds VRAM beyond
// that, mirrors the upstream chunked_decode loop frame for frame.
//
// codec_chunked_decoder_stream : rolling state for AR streaming.
// The pipeline pushes one frame at a time as the talker produces
// them ; push_frame decodes and emits a fresh chunk_frames sized
// audio block through the on_chunk callback as soon as enough new
// frames have accumulated. flush drains the tail at EOS.
#include "pipeline-codec.h"
#include "qwen.h"
#include <cstdint>
#include <cstring>
#include <vector>
// One shot chunked decode. codes is K major [K, T] row major (T fastest).
// Returns audio of length T * TOKENIZER_HOP_LENGTH on success, empty on
// failure. chunk_frames clamps to 1, left_ctx_frames clamps to 0.
static inline std::vector<float> codec_chunked_decode(PipelineCodec * pc,
const int32_t * codes,
int K,
int T,
int chunk_frames,
int left_ctx_frames) {
std::vector<float> out;
if (T <= 0) {
return out;
}
if (chunk_frames < 1) {
chunk_frames = 1;
}
if (left_ctx_frames < 0) {
left_ctx_frames = 0;
}
out.reserve((size_t) T * (size_t) TOKENIZER_HOP_LENGTH);
int start = 0;
while (start < T) {
int end = start + chunk_frames;
if (end > T) {
end = T;
}
// Upstream rule : context_size collapses to start when
// left_ctx_frames would underflow before frame 0.
int ctx = (start - left_ctx_frames > 0) ? left_ctx_frames : start;
int slice_start = start - ctx;
int slice_T = end - slice_start;
std::vector<int32_t> slice((size_t) K * (size_t) slice_T);
for (int k = 0; k < K; k++) {
std::memcpy(slice.data() + (size_t) k * (size_t) slice_T,
codes + (size_t) k * (size_t) T + (size_t) slice_start, (size_t) slice_T * sizeof(int32_t));
}
std::vector<float> wav = pipeline_codec_decode(pc, slice.data(), K, slice_T);
if (wav.empty()) {
return std::vector<float>();
}
const size_t drop = (size_t) ctx * (size_t) TOKENIZER_HOP_LENGTH;
if (wav.size() > drop) {
out.insert(out.end(), wav.begin() + drop, wav.end());
}
start = end;
}
return out;
}
// Rolling streaming decoder. Stores codes K major as K parallel vectors
// (by_k[k][t]) so emit_one can memcpy a contiguous K major slice into
// pipeline_codec_decode without a transpose. push_frame triggers as
// many emits as possible after appending one frame ; flush emits the
// tail at EOS.
struct codec_chunked_decoder_stream {
std::vector<std::vector<int32_t>> by_k;
int K;
int T_so_far;
int chunk_frames;
int left_ctx_frames;
int emit_start_frame;
// Set true when an emit returned false because the on_chunk callback
// requested a cancel. Stays false on decode failures so the caller
// can route to QT_STATUS_CANCELLED vs QT_STATUS_GENERATE_FAILED on
// a push_frame / flush negative return.
bool cancelled;
void init(int K_, int chunk_frames_, int left_ctx_frames_) {
K = K_;
T_so_far = 0;
chunk_frames = chunk_frames_ < 1 ? 1 : chunk_frames_;
left_ctx_frames = left_ctx_frames_ < 0 ? 0 : left_ctx_frames_;
emit_start_frame = 0;
cancelled = false;
by_k.assign((size_t) K, {});
}
// Append one frame (K int32 codes, one per codebook). Drain any
// chunks that became emittable. Returns false on decode failure or
// when cb returns false (cancellation).
bool push_frame(PipelineCodec * pc, const int32_t * frame_codes, qt_audio_chunk_cb cb, void * cb_ud) {
for (int k = 0; k < K; k++) {
by_k[(size_t) k].push_back(frame_codes[k]);
}
T_so_far++;
while (T_so_far - emit_start_frame >= chunk_frames) {
if (!emit_one(pc, emit_start_frame + chunk_frames, cb, cb_ud)) {
return false;
}
}
return true;
}
// Drain the tail. If frames remain past emit_start_frame, decode
// them with left context and emit one final short chunk. Idempotent
// on empty tail.
bool flush(PipelineCodec * pc, qt_audio_chunk_cb cb, void * cb_ud) {
if (T_so_far > emit_start_frame) {
return emit_one(pc, T_so_far, cb, cb_ud);
}
return true;
}
private:
// Decode [emit_start_frame - ctx .. end_frame] with left context
// stripped from the emitted samples, then advance emit_start_frame.
bool emit_one(PipelineCodec * pc, int end_frame, qt_audio_chunk_cb cb, void * cb_ud) {
int ctx = (emit_start_frame - left_ctx_frames > 0) ? left_ctx_frames : emit_start_frame;
int slice_start = emit_start_frame - ctx;
int slice_T = end_frame - slice_start;
std::vector<int32_t> slice((size_t) K * (size_t) slice_T);
for (int k = 0; k < K; k++) {
std::memcpy(slice.data() + (size_t) k * (size_t) slice_T, by_k[(size_t) k].data() + (size_t) slice_start,
(size_t) slice_T * sizeof(int32_t));
}
std::vector<float> wav = pipeline_codec_decode(pc, slice.data(), K, slice_T);
if (wav.empty()) {
return false;
}
const size_t drop = (size_t) ctx * (size_t) TOKENIZER_HOP_LENGTH;
const float * emit_first = wav.data() + drop;
int emit_n = (int) (wav.size() - drop);
if (emit_n > 0 && !cb(emit_first, emit_n, cb_ud)) {
cancelled = true;
return false;
}
emit_start_frame = end_frame;
return true;
}
};
+43 -46
View File
@@ -7,6 +7,7 @@
#include "audio-io.h"
#include "bpe.h"
#include "code-predictor-forward.h"
#include "codec-chunked-decode.h"
#include "debug.h"
#include "ggml.h"
#include "pipeline-codec.h"
@@ -429,13 +430,18 @@ qt_status pipeline_tts_synthesize(PipelineTTS * pt,
const float subtk_T = params->subtalker_do_sample ? params->subtalker_temperature : 0.0f;
const float talker_rp = params->repetition_penalty;
// Streaming config: when on_chunk is set, decode the codec every
// chunk_frames AR frames and emit through the callback. The buffered
// path keeps every code, decodes once at the very end and copies the
// resulting audio into out. chunk_frames <= 0 falls back to 1 second.
const bool streaming = (params->on_chunk != NULL);
const float chunk_sec = params->chunk_duration_sec > 0.0f ? params->chunk_duration_sec : 1.0f;
const int chunk_frames = pipeline_tts_duration_sec_to_tokens(pt, chunk_sec);
// Codec decode framing. Both the streaming path and the buffered
// path route through codec_chunked_decode, with a rolling left
// context window that mirrors the upstream Qwen3-TTS 12 Hz tokenizer
// chunked_decode rule : every chunk re uses up to left_ctx_frames
// previously decoded frames as left context, then the matching audio
// samples are stripped from the head of the decoded chunk. The first
// chunk has its left context collapsed to whatever is available.
const bool streaming = (params->on_chunk != NULL);
const float chunk_sec = params->codec_chunk_sec > 0.0f ? params->codec_chunk_sec : 24.0f;
const float left_ctx_sec = params->codec_left_context_sec >= 0.0f ? params->codec_left_context_sec : 2.0f;
const int chunk_frames = pipeline_tts_duration_sec_to_tokens(pt, chunk_sec);
const int left_ctx_frames = pipeline_tts_duration_sec_to_tokens(pt, left_ctx_sec);
std::vector<std::vector<int32_t>> all_codes;
all_codes.reserve((size_t) params->max_new_tokens);
@@ -450,36 +456,14 @@ qt_status pipeline_tts_synthesize(PipelineTTS * pt,
std::vector<float> next_emb((size_t) hidden, 0.0f);
// Streaming bookkeeping: pending_codes accumulates frames since the
// last emit; emit_pending decodes them through the codec and feeds
// on_chunk. Returns false to abort with QT_STATUS_CANCELLED.
std::vector<std::vector<int32_t>> pending_codes;
pending_codes.reserve((size_t) chunk_frames);
auto emit_pending = [&]() -> bool {
if (pending_codes.empty()) {
return true;
}
const int T_frames = (int) pending_codes.size();
std::vector<int32_t> codes_kt((size_t) num_codebooks * (size_t) T_frames);
for (int t = 0; t < T_frames; t++) {
for (int k = 0; k < num_codebooks; k++) {
codes_kt[(size_t) k * (size_t) T_frames + (size_t) t] = pending_codes[(size_t) t][(size_t) k];
}
}
std::vector<float> chunk_audio = pipeline_codec_decode(&pt->codec, codes_kt.data(), num_codebooks, T_frames);
if (chunk_audio.empty()) {
qt_set_error("pipeline_tts_synthesize: streaming codec decode returned no audio");
qt_log(QT_LOG_ERROR, "[Pipeline] streaming codec decode returned no audio");
return false;
}
if (!params->on_chunk(chunk_audio.data(), (int) chunk_audio.size(), params->on_chunk_user_data)) {
qt_log(QT_LOG_INFO, "[Pipeline] on_chunk callback aborted the synthesis");
return false;
}
pending_codes.clear();
return true;
};
// Streaming rolling decoder. Holds the K major codes buffer, the
// emit cursor and the left context window. push_frame triggers an
// emit as soon as chunk_frames new frames have accumulated since
// the previous emit boundary ; flush drains the tail at EOS.
codec_chunked_decoder_stream stream;
if (streaming) {
stream.init(num_codebooks, chunk_frames, left_ctx_frames);
}
for (int step = 0; step < params->max_new_tokens; step++) {
// Cooperative cancellation, polled at every step. Granularity is
@@ -555,11 +539,14 @@ qt_status pipeline_tts_synthesize(PipelineTTS * pt,
all_codes.push_back(cp.codes);
talker_history.push_back(c0);
if (streaming) {
pending_codes.push_back(cp.codes);
if ((int) pending_codes.size() >= chunk_frames) {
if (!emit_pending()) {
if (!stream.push_frame(&pt->codec, cp.codes.data(), params->on_chunk, params->on_chunk_user_data)) {
if (stream.cancelled) {
qt_log(QT_LOG_INFO, "[Pipeline] on_chunk callback aborted the synthesis");
return QT_STATUS_CANCELLED;
}
qt_set_error("pipeline_tts_synthesize: streaming codec decode failed at frame %d", step);
qt_log(QT_LOG_ERROR, "[Pipeline] streaming codec decode failed at frame %d", step);
return QT_STATUS_GENERATE_FAILED;
}
}
@@ -628,8 +615,14 @@ qt_status pipeline_tts_synthesize(PipelineTTS * pt,
// buffered output stays empty in this branch; the caller already
// received every sample through on_chunk.
if (streaming) {
if (!emit_pending()) {
return QT_STATUS_CANCELLED;
if (!stream.flush(&pt->codec, params->on_chunk, params->on_chunk_user_data)) {
if (stream.cancelled) {
qt_log(QT_LOG_INFO, "[Pipeline] on_chunk callback aborted the synthesis on tail flush");
return QT_STATUS_CANCELLED;
}
qt_set_error("pipeline_tts_synthesize: streaming codec decode failed on tail flush");
qt_log(QT_LOG_ERROR, "[Pipeline] streaming codec decode failed on tail flush");
return QT_STATUS_GENERATE_FAILED;
}
out->samples = NULL;
out->n_samples = 0;
@@ -649,9 +642,12 @@ qt_status pipeline_tts_synthesize(PipelineTTS * pt,
return QT_STATUS_OK;
}
// Codec decode: transpose codes from [T_frames, K] to [K, T_frames]
// because pipeline_codec_decode expects K-major layout (codebooks
// first, frames second), then materialise the 24 kHz mono audio.
// Buffered codec decode through the chunked path : same framing as
// the streaming branch (chunk_frames + left_ctx_frames), bit perfect
// equivalent to a single pipeline_codec_decode call when T_frames
// fits in one chunk, bounded VRAM beyond that. Transpose codes from
// [T_frames, K] to [K, T_frames] because codec_chunked_decode
// expects K major layout.
const int T_frames = (int) all_codes.size();
std::vector<int32_t> codes_kt((size_t) num_codebooks * (size_t) T_frames);
for (int t = 0; t < T_frames; t++) {
@@ -659,7 +655,8 @@ qt_status pipeline_tts_synthesize(PipelineTTS * pt,
codes_kt[(size_t) k * (size_t) T_frames + (size_t) t] = all_codes[(size_t) t][(size_t) k];
}
}
std::vector<float> audio = pipeline_codec_decode(&pt->codec, codes_kt.data(), num_codebooks, T_frames);
std::vector<float> audio =
codec_chunked_decode(&pt->codec, codes_kt.data(), num_codebooks, T_frames, chunk_frames, left_ctx_frames);
if (audio.empty()) {
qt_set_error("pipeline_tts_synthesize: codec decode returned no audio");
qt_log(QT_LOG_ERROR, "[Pipeline] codec decode returned no audio");
+26 -25
View File
@@ -192,31 +192,32 @@ void qt_init_default_params(struct qt_init_params * p) {
}
void qt_tts_default_params(struct qt_tts_params * p) {
p->abi_version = QT_ABI_VERSION;
p->text = nullptr;
p->lang = nullptr;
p->instruct = nullptr;
p->speaker = nullptr;
p->ref_audio_24k = nullptr;
p->ref_n_samples = 0;
p->ref_text = nullptr;
p->seed = -1;
p->max_new_tokens = 2048;
p->do_sample = true;
p->temperature = 0.9f;
p->top_k = 50;
p->top_p = 1.0f;
p->repetition_penalty = 1.05f;
p->subtalker_do_sample = true;
p->subtalker_temperature = 0.9f;
p->subtalker_top_k = 50;
p->subtalker_top_p = 1.0f;
p->dump_dir = nullptr;
p->cancel = nullptr;
p->cancel_user_data = nullptr;
p->on_chunk = nullptr;
p->on_chunk_user_data = nullptr;
p->chunk_duration_sec = 1.0f;
p->abi_version = QT_ABI_VERSION;
p->text = nullptr;
p->lang = nullptr;
p->instruct = nullptr;
p->speaker = nullptr;
p->ref_audio_24k = nullptr;
p->ref_n_samples = 0;
p->ref_text = nullptr;
p->seed = -1;
p->max_new_tokens = 2048;
p->do_sample = true;
p->temperature = 0.9f;
p->top_k = 50;
p->top_p = 1.0f;
p->repetition_penalty = 1.05f;
p->subtalker_do_sample = true;
p->subtalker_temperature = 0.9f;
p->subtalker_top_k = 50;
p->subtalker_top_p = 1.0f;
p->dump_dir = nullptr;
p->cancel = nullptr;
p->cancel_user_data = nullptr;
p->on_chunk = nullptr;
p->on_chunk_user_data = nullptr;
p->codec_chunk_sec = 24.0f;
p->codec_left_context_sec = 2.0f;
}
struct qt_context * qt_init(const struct qt_init_params * params) {
+19 -5
View File
@@ -248,18 +248,32 @@ struct qt_tts_params {
// Streaming output. When on_chunk is non NULL, qt_synthesize runs
// the streaming pipeline: audio chunks emit through on_chunk and
// `out` stays empty on success. on_chunk NULL keeps the buffered
// path. chunk_duration_sec drives the chunk size at codec sample
// rate; values <= 0 fall back to 1.0 second. The last chunk on
// EOS or max_new flushes whatever frames remain.
// path. The last chunk on EOS or max_new flushes whatever frames
// remain.
qt_audio_chunk_cb on_chunk;
void * on_chunk_user_data;
float chunk_duration_sec;
// Codec decode framing. Applied to both the streaming path (chunk
// by chunk emission) and the buffered path (one shot decode at the
// end) : the chunked decode rolls a left context window across the
// codec frames to avoid edge artefacts at chunk boundaries. The
// first chunk has its left context collapsed to whatever is
// available, matching the upstream Qwen3-TTS 12 Hz tokenizer
// chunked_decode rule. Defaults match the upstream reference :
// codec_chunk_sec 24.0 (300 frames at 12.5 Hz) and
// codec_left_context_sec 2.0 (25 frames at 12.5 Hz). Values are
// converted internally to integer frame counts via the codec frame
// rate ; codec_chunk_sec clamps to >= 1 frame, codec_left_context_sec
// clamps to >= 0 frames.
float codec_chunk_sec;
float codec_left_context_sec;
};
// Initialise to the standard defaults. Strings NULL, seed -1,
// max_new_tokens 2048, do_sample true, temperature 0.9, top_k 50,
// top_p 1.0, repetition_penalty 1.05, subtalker mirrors talker,
// dump_dir NULL, cancel NULL, on_chunk NULL, chunk_duration_sec 1.0.
// dump_dir NULL, cancel NULL, on_chunk NULL, codec_chunk_sec 24.0,
// codec_left_context_sec 2.0.
QT_API void qt_tts_default_params(struct qt_tts_params * p);
// Run the full TTS synthesis. Validates the params against the loaded