diff --git a/CHANGELOG.md b/CHANGELOG.md index c3a77ef..1465b79 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,8 @@ ## Unreleased +- WebGPU text-to-speech applies the GGUF's suppress tokens in the talker sampler, as upstream does. + Qwen3-TTS no longer fails with `code out of range` when a codec control token is sampled. - `createCompletion` can explicitly count prompt tokens in penalties with a grammar using `penaltyPromptTokensWithGrammar`; a matching loaded-core capability is required, and existing grammar defaults stay unchanged. - Large wasm32 streamed model files use bounded size hints to reserve filesystem capacity, avoiding geometric file-growth allocation failures at 1 GiB without treating hidden compression metadata as a decoded-size contract. diff --git a/src/llama_webgpu_tts.cpp b/src/llama_webgpu_tts.cpp index 32815c6..2ea71cd 100644 --- a/src/llama_webgpu_tts.cpp +++ b/src/llama_webgpu_tts.cpp @@ -109,12 +109,32 @@ uint32_t capabilities_for( : 0u); } -llama_sampler * init_sampler(const llama_webgpu_tts_request & request) { +llama_sampler * init_sampler( + const llama_vocab * vocab, + const llama_webgpu_tts_request & request) { llama_sampler * sampler = llama_sampler_chain_init(llama_sampler_chain_default_params()); if (sampler == nullptr) { return nullptr; } + // As in upstream common_sampler, suppressed tokens (for Qwen3-TTS the codec's + // control tokens, which code2wav rejects as out-of-range codes) get a -inf bias first. + int32_t n_suppress = 0; + const llama_token * suppress = + llama_vocab_get_suppress_tokens(vocab, &n_suppress); + if (suppress != nullptr && n_suppress > 0) { + std::vector bias; + bias.reserve(static_cast(n_suppress)); + for (int32_t i = 0; i < n_suppress; ++i) { + bias.push_back({suppress[i], -INFINITY}); + } + llama_sampler_chain_add( + sampler, + llama_sampler_init_logit_bias( + llama_vocab_n_tokens(vocab), + static_cast(bias.size()), + bias.data())); + } llama_sampler_chain_add(sampler, llama_sampler_init_top_k(request.top_k)); llama_sampler_chain_add(sampler, llama_sampler_init_top_p(request.top_p, 1)); llama_sampler_chain_add(sampler, llama_sampler_init_min_p(request.min_p, 1)); @@ -307,7 +327,8 @@ llama_webgpu_tts_status llama_webgpu_tts_start( tts->speaker = wrapper.bitmap; } - tts->sampler = init_sampler(request); + tts->sampler = init_sampler( + llama_model_get_vocab(llama_get_model(tts->llama)), request); if (tts->sampler == nullptr) { return fail( tts, diff --git a/tests/bridge/tts_sampler_suppress_tokens_test.mjs b/tests/bridge/tts_sampler_suppress_tokens_test.mjs new file mode 100644 index 0000000..f88cdc6 --- /dev/null +++ b/tests/bridge/tts_sampler_suppress_tokens_test.mjs @@ -0,0 +1,36 @@ +// The TTS talker sampler must bias the GGUF's suppress tokens to -inf before any +// other stage, as upstream common_sampler does. Qwen3-TTS lists its codec control +// tokens there, and code2wav rejects a sampled control token as out of range. +import assert from 'node:assert/strict'; +import test from 'node:test'; + +import { readRepoText } from './bridge_js_source.mjs'; + +const TTS = readRepoText('src/llama_webgpu_tts.cpp'); + +function samplerBody() { + const start = TTS.indexOf('llama_sampler * init_sampler('); + assert.notEqual(start, -1, 'init_sampler not found'); + const end = TTS.indexOf('\nllama_webgpu_tts_status ', start); + assert.ok(end > start, 'init_sampler end not found'); + return TTS.slice(start, end); +} + +test('TTS sampler biases suppress tokens to -inf before the first sampling stage', () => { + const body = samplerBody(); + const suppress = body.indexOf('llama_vocab_get_suppress_tokens('); + const bias = body.indexOf('llama_sampler_init_logit_bias('); + const firstStage = body.indexOf('llama_sampler_init_top_k('); + assert.ok(suppress !== -1, 'suppress tokens are not read from the vocab'); + assert.ok(bias > suppress, 'logit bias is not built from the suppress tokens'); + assert.ok(firstStage > bias, 'logit bias must be added before top_k'); + assert.match(body, /-INFINITY/); + assert.match(body, /llama_vocab_n_tokens\(vocab\)/); +}); + +test('TTS sampler is built from the loaded model vocab', () => { + assert.match( + TTS, + /init_sampler\(\s*llama_model_get_vocab\(llama_get_model\(tts->llama\)\),\s*request\)/, + ); +});