diff --git a/src/code-predictor-forward.h b/src/code-predictor-forward.h index 94814d1..fd2e55a 100644 --- a/src/code-predictor-forward.h +++ b/src/code-predictor-forward.h @@ -193,7 +193,9 @@ static bool code_predictor_run(const CodePredictorWeights * cw, KVCache * kv, ggml_backend_sched_t sched, GraphArena * arena, - const float * fresh_input, + struct ggml_tensor * embd_table, + const float * hidden_row, + int32_t code_id, int T, int n_past, int talker_hidden, @@ -213,13 +215,28 @@ static bool code_predictor_run(const CodePredictorWeights * cw, const int max_nodes = code_predictor_graph_max_nodes(n_layers); struct ggml_context * gctx = graph_arena_begin(arena); - // Inputs: fresh embeddings (talker_hidden), positions, attention mask - struct ggml_tensor * x_in = ggml_new_tensor_2d(gctx, GGML_TYPE_F32, talker_hidden, T); + // Inputs: one code id gathered in graph from embd_table, positions, + // attention mask, plus the raw talker hidden row on the prefill + // path (T == 2, hidden_row non NULL) where the sequence is + // [talker_hidden, embed(c0)]. Steps (T == 1) are pure gathers: the + // only per step upload is 4 bytes of code id. + struct ggml_tensor * ids_in = ggml_new_tensor_1d(gctx, GGML_TYPE_I32, 1); struct ggml_tensor * pos_in = ggml_new_tensor_1d(gctx, GGML_TYPE_I32, T); struct ggml_tensor * mask_in = ggml_new_tensor_2d(gctx, GGML_TYPE_F16, n_kv_pad, T); - ggml_set_name(x_in, "sub_input"); + ggml_set_name(ids_in, "sub_code_id"); ggml_set_name(pos_in, "positions"); ggml_set_name(mask_in, "causal_mask"); + ggml_set_input(ids_in); + + struct ggml_tensor * x_in = ggml_get_rows(gctx, embd_table, ids_in); + struct ggml_tensor * hid_in = NULL; + if (T == 2) { + hid_in = ggml_new_tensor_2d(gctx, GGML_TYPE_F32, talker_hidden, 1); + ggml_set_name(hid_in, "talker_hidden_row"); + ggml_set_input(hid_in); + x_in = ggml_concat(gctx, hid_in, x_in, 1); + } + ggml_set_name(x_in, "sub_input"); struct ggml_cgraph * gf = ggml_new_graph_custom(gctx, max_nodes, false); @@ -254,7 +271,10 @@ static bool code_predictor_run(const CodePredictorWeights * cw, return false; } - ggml_backend_tensor_set(x_in, fresh_input, 0, (size_t) T * (size_t) talker_hidden * sizeof(float)); + ggml_backend_tensor_set(ids_in, &code_id, 0, sizeof(int32_t)); + if (hid_in) { + ggml_backend_tensor_set(hid_in, hidden_row, 0, (size_t) talker_hidden * sizeof(float)); + } { std::vector pos((size_t) T); @@ -294,30 +314,6 @@ static bool code_predictor_run(const CodePredictorWeights * cw, return true; } -// Read one row of an embedding table to f32. Reads from the backend -// (the predictor weights live there) via ggml_backend_tensor_get, -// dispatched through ggml_get_type_traits so quants are accepted. -static void embed_row_from_backend(struct ggml_tensor * t, int row_id, int dim, float * dst) { - if (t->ne[0] != dim) { - qt_throw("[CodePredictor] embed dim mismatch %lld vs %d", (long long) t->ne[0], dim); - } - if (row_id < 0 || row_id >= (int) t->ne[1]) { - qt_throw("[CodePredictor] row %d out of range (vocab=%lld)", row_id, (long long) t->ne[1]); - } - const size_t row_bytes = ggml_row_size(t->type, dim); - if (t->type == GGML_TYPE_F32) { - ggml_backend_tensor_get(t, dst, (size_t) row_id * row_bytes, row_bytes); - return; - } - const struct ggml_type_traits * tt = ggml_get_type_traits(t->type); - if (!tt || !tt->to_float) { - qt_throw("[CodePredictor] unsupported embed dtype %d", (int) t->type); - } - std::vector tmp(row_bytes); - ggml_backend_tensor_get(t, tmp.data(), (size_t) row_id * row_bytes, row_bytes); - tt->to_float(tmp.data(), dst, dim); -} - // Run the predictor for one audio frame. Caller passes the talker hidden // state for the current frame and the already-sampled c0. Sampling // parameters control greedy (temperature <= 0) vs stochastic. subseq_base @@ -358,15 +354,14 @@ static bool code_predictor_step(const TalkerWeights * tw, out->codes.assign((size_t) (n_acoustic + 1), 0); out->codes[0] = c0; - // Prefill: two positions, talker_hidden_last and embed_talker(c0). + // Prefill: two positions, [talker_hidden_last, embed_talker(c0)]. + // The hidden row uploads raw, c0 gathers in graph from the talker + // codec embedding table. kv_cache_reset(kv); - std::vector prefill_input((size_t) 2 * (size_t) talker_hidden, 0.0f); - std::memcpy(prefill_input.data(), talker_hidden_last, (size_t) talker_hidden * sizeof(float)); - embed_row_from_backend(tw->codec_embedding, c0, talker_hidden, prefill_input.data() + (size_t) talker_hidden); std::vector logits; - if (!code_predictor_run(cw, kv, sched, arena_prefill, prefill_input.data(), 2, 0, talker_hidden, 0, use_flash_attn, - clamp_fp16, &logits)) { + if (!code_predictor_run(cw, kv, sched, arena_prefill, tw->codec_embedding, talker_hidden_last, c0, 2, 0, + talker_hidden, 0, use_flash_attn, clamp_fp16, &logits)) { return false; } { @@ -385,13 +380,12 @@ static bool code_predictor_step(const TalkerWeights * tw, } // Decode loop: 14 single-token steps. At step g (g=1..14) we feed - // the embedding of the code we just sampled and read lm_head[g]. - std::vector step_input((size_t) talker_hidden); + // the id of the code we just sampled, gathered in graph from the + // group's private embedding table, and read lm_head[g]. for (int g = 1; g < n_acoustic; g++) { - embed_row_from_backend(cw->codec_embedding[(size_t) (g - 1)], out->codes[(size_t) g], talker_hidden, - step_input.data()); - if (!code_predictor_run(cw, kv, sched, arena_step, step_input.data(), 1, kv->cur_len, talker_hidden, g, - use_flash_attn, clamp_fp16, &logits)) { + if (!code_predictor_run(cw, kv, sched, arena_step, cw->codec_embedding[(size_t) (g - 1)], NULL, + out->codes[(size_t) g], 1, kv->cur_len, talker_hidden, g, use_flash_attn, clamp_fp16, + &logits)) { return false; } float u_g = 0.0f; diff --git a/src/pipeline-codec.cpp b/src/pipeline-codec.cpp index 59ff080..9cbcf6d 100644 --- a/src/pipeline-codec.cpp +++ b/src/pipeline-codec.cpp @@ -11,6 +11,7 @@ #include "causal-trans-conv.h" #include "debug.h" #include "qt-error.h" +#include "timer.h" #include #include @@ -73,51 +74,9 @@ bool pipeline_codec_load(PipelineCodec * pc, const char * gguf_path, BackendPair pc->pre_conv_wctx = std::move(wctx); } - if (!seanet_encoder_load(&pc->seanet, pc->gguf, pc->backend)) { - wctx_free(&pc->pre_conv_wctx); - dac_decoder_free(&pc->dac); - upsample_stage_free(&pc->upsample); - tok_trans_free(&pc->transformer); - quant_decoder_free(&pc->qdec); - gf_close(&pc->gguf); - return false; - } - - if (!enc_trans_load(&pc->enc_transformer, pc->gguf, pc->backend)) { - seanet_encoder_free(&pc->seanet); - wctx_free(&pc->pre_conv_wctx); - dac_decoder_free(&pc->dac); - upsample_stage_free(&pc->upsample); - tok_trans_free(&pc->transformer); - quant_decoder_free(&pc->qdec); - gf_close(&pc->gguf); - return false; - } - - if (!enc_down_load(&pc->enc_downsample, pc->gguf, pc->backend)) { - enc_trans_free(&pc->enc_transformer); - seanet_encoder_free(&pc->seanet); - wctx_free(&pc->pre_conv_wctx); - dac_decoder_free(&pc->dac); - upsample_stage_free(&pc->upsample); - tok_trans_free(&pc->transformer); - quant_decoder_free(&pc->qdec); - gf_close(&pc->gguf); - return false; - } - - if (!quant_encode_load(&pc->qenc, pc->gguf, pc->backend)) { - enc_down_free(&pc->enc_downsample); - enc_trans_free(&pc->enc_transformer); - seanet_encoder_free(&pc->seanet); - wctx_free(&pc->pre_conv_wctx); - dac_decoder_free(&pc->dac); - upsample_stage_free(&pc->upsample); - tok_trans_free(&pc->transformer); - quant_decoder_free(&pc->qdec); - gf_close(&pc->gguf); - return false; - } + // Encoder half (seanet, enc_transformer, enc_downsample, qenc) + // stays on disk until the first encode request. + pc->enc_loaded = false; pc->sched = backend_sched_new(bp, 4096); @@ -217,10 +176,44 @@ std::vector pipeline_codec_decode(PipelineCodec * pc, const int32_t * cod return audio; } +bool pipeline_codec_ensure_encoder(PipelineCodec * pc) { + if (pc->enc_loaded) { + return true; + } + + Timer t_load; + if (!seanet_encoder_load(&pc->seanet, pc->gguf, pc->backend)) { + return false; + } + if (!enc_trans_load(&pc->enc_transformer, pc->gguf, pc->backend)) { + seanet_encoder_free(&pc->seanet); + return false; + } + if (!enc_down_load(&pc->enc_downsample, pc->gguf, pc->backend)) { + enc_trans_free(&pc->enc_transformer); + seanet_encoder_free(&pc->seanet); + return false; + } + if (!quant_encode_load(&pc->qenc, pc->gguf, pc->backend)) { + enc_down_free(&pc->enc_downsample); + enc_trans_free(&pc->enc_transformer); + seanet_encoder_free(&pc->seanet); + return false; + } + + pc->enc_loaded = true; + qt_log(QT_LOG_INFO, "[Pipeline] Codec encoder lazy loaded in %.0f ms", t_load.ms()); + return true; +} + std::vector pipeline_codec_encode(PipelineCodec * pc, const float * audio, int n_samples, const char * dump_dir) { + if (!pipeline_codec_ensure_encoder(pc)) { + qt_log(QT_LOG_ERROR, "[Pipeline] codec encoder load failed"); + return {}; + } if (n_samples <= 0 || (n_samples % TOKENIZER_HOP_LENGTH) != 0) { qt_log(QT_LOG_ERROR, "[Pipeline] n_samples must be a positive multiple of %d (got %d)", TOKENIZER_HOP_LENGTH, n_samples); @@ -434,10 +427,13 @@ void pipeline_codec_free(PipelineCodec * pc) { ggml_backend_sched_free(pc->sched); pc->sched = NULL; } - quant_encode_free(&pc->qenc); - enc_down_free(&pc->enc_downsample); - enc_trans_free(&pc->enc_transformer); - seanet_encoder_free(&pc->seanet); + if (pc->enc_loaded) { + quant_encode_free(&pc->qenc); + enc_down_free(&pc->enc_downsample); + enc_trans_free(&pc->enc_transformer); + seanet_encoder_free(&pc->seanet); + pc->enc_loaded = false; + } wctx_free(&pc->pre_conv_wctx); dac_decoder_free(&pc->dac); upsample_stage_free(&pc->upsample); diff --git a/src/pipeline-codec.h b/src/pipeline-codec.h index aaa7388..50e0bdf 100644 --- a/src/pipeline-codec.h +++ b/src/pipeline-codec.h @@ -63,7 +63,12 @@ struct PipelineCodec { QwenSEANetEncoder seanet; QwenEncoderTransformer enc_transformer; QwenEncoderDownsample enc_downsample; - QwenQuantizerEncode qenc; + + // Encoder weights (seanet, enc_transformer, enc_downsample, qenc) + // load lazily on the first pipeline_codec_encode call: synthesis + // from pre encoded reference codes never pays for them. + bool enc_loaded; + QwenQuantizerEncode qenc; // CPU mirror of the RVQ encode side, lazy-loaded on first encode call. QwenQuantizerEncodeHost qenc_sem_host; @@ -79,6 +84,10 @@ struct PipelineCodec { // On failure leaves the struct in a clean state and returns false. bool pipeline_codec_load(PipelineCodec * pc, const char * gguf_path, BackendPair bp); +// Load the encoder weights on demand. Idempotent, called by +// pipeline_codec_encode; harmless to call when already resident. +bool pipeline_codec_ensure_encoder(PipelineCodec * pc); + // Decode RVQ codes into a 24 kHz mono waveform. // codes: flat int32 buffer, [K, T] row-major (T fastest). // Returns audio of length T * TOKENIZER_HOP_LENGTH, empty on failure. diff --git a/src/pipeline-tts.cpp b/src/pipeline-tts.cpp index 20756d1..76f47a0 100644 --- a/src/pipeline-tts.cpp +++ b/src/pipeline-tts.cpp @@ -158,23 +158,14 @@ bool pipeline_tts_load(PipelineTTS * pt, return false; } - // Speaker encoder is only present in Base checkpoints. Treat absence - // as a soft condition: voice clone path stays disabled, base-direct - // synthesis still works. - if (pt->model_type == "base") { - if (!speaker_encoder_weights_load(&pt->speaker_encoder, pt->gguf_talker, pt->backend)) { - code_predictor_weights_free(&pt->code_predictor); - talker_weights_free(&pt->talker); - gf_close(&pt->gguf_talker); - return false; - } - pt->has_speaker_encoder = (pt->speaker_encoder.weight_buf != NULL); - } + // Speaker encoder tensors are only present in Base checkpoints. The + // weights load lazily on the first --ref-wav request: synthesis from + // pre extracted embeddings never pays for them. has_speaker_encoder + // advertises the capability, spk_enc_loaded tracks residency. + pt->has_speaker_encoder = (pt->model_type == "base"); + pt->spk_enc_loaded = false; if (!pipeline_codec_load(&pt->codec, codec_gguf_path, bp)) { - if (pt->has_speaker_encoder) { - speaker_encoder_weights_free(&pt->speaker_encoder); - } code_predictor_weights_free(&pt->code_predictor); talker_weights_free(&pt->talker); gf_close(&pt->gguf_talker); @@ -189,9 +180,6 @@ bool pipeline_tts_load(PipelineTTS * pt, pt->sched = backend_sched_new(bp, 4096); if (!pt->sched) { pipeline_codec_free(&pt->codec); - if (pt->has_speaker_encoder) { - speaker_encoder_weights_free(&pt->speaker_encoder); - } code_predictor_weights_free(&pt->code_predictor); talker_weights_free(&pt->talker); gf_close(&pt->gguf_talker); @@ -204,9 +192,6 @@ bool pipeline_tts_load(PipelineTTS * pt, ggml_backend_sched_free(pt->sched); pt->sched = NULL; pipeline_codec_free(&pt->codec); - if (pt->has_speaker_encoder) { - speaker_encoder_weights_free(&pt->speaker_encoder); - } code_predictor_weights_free(&pt->code_predictor); talker_weights_free(&pt->talker); gf_close(&pt->gguf_talker); @@ -221,9 +206,6 @@ bool pipeline_tts_load(PipelineTTS * pt, ggml_backend_sched_free(pt->sched); pt->sched = NULL; pipeline_codec_free(&pt->codec); - if (pt->has_speaker_encoder) { - speaker_encoder_weights_free(&pt->speaker_encoder); - } code_predictor_weights_free(&pt->code_predictor); talker_weights_free(&pt->talker); gf_close(&pt->gguf_talker); @@ -236,9 +218,6 @@ bool pipeline_tts_load(PipelineTTS * pt, ggml_backend_sched_free(pt->sched); pt->sched = NULL; pipeline_codec_free(&pt->codec); - if (pt->has_speaker_encoder) { - speaker_encoder_weights_free(&pt->speaker_encoder); - } code_predictor_weights_free(&pt->code_predictor); talker_weights_free(&pt->talker); gf_close(&pt->gguf_talker); @@ -259,9 +238,6 @@ bool pipeline_tts_load(PipelineTTS * pt, ggml_backend_sched_free(pt->sched); pt->sched = NULL; pipeline_codec_free(&pt->codec); - if (pt->has_speaker_encoder) { - speaker_encoder_weights_free(&pt->speaker_encoder); - } code_predictor_weights_free(&pt->code_predictor); talker_weights_free(&pt->talker); gf_close(&pt->gguf_talker); @@ -272,7 +248,7 @@ bool pipeline_tts_load(PipelineTTS * pt, "[Pipeline] Loaded: arch=%s variant=%s tokenizer=%s codebooks=%d speaker_encoder=%s speakers=%zu fa=%s " "clamp_fp16=%s", pt->model_size.c_str(), pt->model_type.c_str(), pt->tokenizer_type.c_str(), pt->num_code_groups, - pt->has_speaker_encoder ? "loaded" : "absent", pt->speakers.size(), pt->use_flash_attn ? "on" : "off", + pt->has_speaker_encoder ? "deferred" : "absent", pt->speakers.size(), pt->use_flash_attn ? "on" : "off", pt->clamp_fp16 ? "on" : "off"); return true; } @@ -288,8 +264,9 @@ void pipeline_tts_free(PipelineTTS * pt) { pt->sched = NULL; } pipeline_codec_free(&pt->codec); - if (pt->has_speaker_encoder) { + if (pt->spk_enc_loaded) { speaker_encoder_weights_free(&pt->speaker_encoder); + pt->spk_enc_loaded = false; } code_predictor_weights_free(&pt->code_predictor); talker_weights_free(&pt->talker); @@ -447,11 +424,24 @@ qt_status pipeline_tts_synthesize(PipelineTTS * pt, qt_log(QT_LOG_INFO, "[Pipeline] Latent speaker embedding: %d values", lat_spk_dim); } else if (has_ref_audio) { if (!pt->has_speaker_encoder) { - qt_set_error( - "pipeline_tts_synthesize: --ref-wav requires a model with a loaded speaker encoder (Base only)"); - qt_log(QT_LOG_ERROR, "[Pipeline] --ref-wav requires a model with a loaded speaker encoder (Base only)"); + qt_set_error("pipeline_tts_synthesize: --ref-wav requires a model with a speaker encoder (Base only)"); + qt_log(QT_LOG_ERROR, "[Pipeline] --ref-wav requires a model with a speaker encoder (Base only)"); return QT_STATUS_GENERATE_FAILED; } + // Lazy residency: the first reference audio request pays the + // weight load once, pre extracted paths never do. + if (!pt->spk_enc_loaded) { + Timer t_spk_load; + if (!speaker_encoder_weights_load(&pt->speaker_encoder, pt->gguf_talker, pt->backend) || + pt->speaker_encoder.weight_buf == NULL) { + pt->has_speaker_encoder = false; + qt_set_error("pipeline_tts_synthesize: speaker encoder load failed"); + qt_log(QT_LOG_ERROR, "[Pipeline] speaker encoder load failed"); + return QT_STATUS_GENERATE_FAILED; + } + pt->spk_enc_loaded = true; + qt_log(QT_LOG_INFO, "[Pipeline] Speaker encoder lazy loaded in %.0f ms", t_spk_load.ms()); + } if (!speaker_encoder_extract(&pt->speaker_encoder, pt->sched, params->ref_audio_24k, params->ref_n_samples, ref_spk_emb, params->dump_dir)) { return QT_STATUS_GENERATE_FAILED; @@ -580,7 +570,11 @@ qt_status pipeline_tts_synthesize(PipelineTTS * pt, // sample (one for c0 of each step, then 15 for the predictor codes). int64_t subseq_counter = 0; - std::vector next_emb((size_t) hidden, 0.0f); + // Decode input state: the codes sampled at the previous frame plus + // the trailing text / pad overlay row for that frame. The talker + // decode graph gathers and sums the 16 embeddings on device. + std::vector prev_ids((size_t) num_codebooks, 0); + const float * prev_overlay = NULL; // Streaming rolling decoder. Holds the K major codes buffer, the // emit cursor and the left context window. push_frame triggers an @@ -608,8 +602,10 @@ qt_status pipeline_tts_synthesize(PipelineTTS * pt, ok = talker_forward_prefill(&pt->talker, &pt->talker_kv, pt->sched, &pt->talker_arena, prompt.input_embed.data(), prompt.T_ctx, use_fa, clamp_fp16, step_dump, &fw); } else { - ok = talker_forward_decode(&pt->talker, &pt->talker_kv, pt->sched, &pt->talker_arena, next_emb.data(), - use_fa, clamp_fp16, &fw); + ok = + talker_forward_decode(&pt->talker, &pt->talker_kv, pt->sched, &pt->talker_arena, prev_ids.data(), + pt->code_predictor.codec_embedding.data(), + pt->code_predictor.num_acoustic_codebooks, prev_overlay, use_fa, clamp_fp16, &fw); } if (!ok) { return QT_STATUS_GENERATE_FAILED; @@ -693,43 +689,41 @@ qt_status pipeline_tts_synthesize(PipelineTTS * pt, } } - // Build next-token embedding: sum of 16 codebook embeddings. - // codebook 0 uses talker.codec_embedding, the 15 acoustic - // codebooks use the predictor's private embedding tables. - Timer t_emb; - std::fill(next_emb.begin(), next_emb.end(), 0.0f); - std::vector tmp((size_t) hidden); - - embed_row_from_gguf(pt->gguf_talker, "talker.codec_embd.weight", c0, hidden, tmp.data()); - for (int i = 0; i < hidden; i++) { - next_emb[(size_t) i] += tmp[(size_t) i]; + // Next decode input: the 16 frame codes gather and sum in graph + // (codebook 0 from talker.codec_embedding, the 15 acoustic + // groups from the predictor's private tables). The overlay row + // adds the next utterance text hidden while any remains, the + // tts_pad embedding afterwards. + prev_ids[0] = c0; + for (int g = 1; g < num_codebooks; g++) { + prev_ids[(size_t) g] = cp.codes[(size_t) g]; } - for (int g = 0; g < num_codebooks - 1; g++) { - int cg = cp.codes[(size_t) (g + 1)]; - char name[64]; - snprintf(name, sizeof(name), "code_pred.codec_embd.%d.weight", g); - embed_row_from_gguf(pt->gguf_talker, name, cg, hidden, tmp.data()); + prev_overlay = (step < prompt.T_trailing) ? + prompt.trailing_text_hidden.data() + (size_t) step * (size_t) hidden : + prompt.tts_pad_embed.data(); + + // Bisection dump: reproduce the in graph composition on host so + // the step 0 next embedding stays byte comparable against the + // Python hook (codebook sums plus trailing text overlay). + if (params->dump_dir && step == 0) { + std::vector next_emb((size_t) hidden, 0.0f); + std::vector tmp((size_t) hidden); + embed_row_from_gguf(pt->gguf_talker, "talker.codec_embd.weight", c0, hidden, tmp.data()); for (int i = 0; i < hidden; i++) { next_emb[(size_t) i] += tmp[(size_t) i]; } - } - - // Trailing text overlay: while we still have utterance text - // hiddens to consume, add the next one; otherwise add the - // tts_pad embedding. - const float * overlay = (step < prompt.T_trailing) ? - prompt.trailing_text_hidden.data() + (size_t) step * (size_t) hidden : - prompt.tts_pad_embed.data(); - for (int i = 0; i < hidden; i++) { - next_emb[(size_t) i] += overlay[(size_t) i]; - } - perf.host_ms += t_emb.ms(); - - // Bisection dump: the next-token embedding produced at step 0 - // is the only thing controlling the talker forward at step 1, so - // matching it bit-exact against Python pinpoints any drift in - // the codebook embedding sums or the trailing text overlay. - if (params->dump_dir && step == 0) { + for (int g = 0; g < num_codebooks - 1; g++) { + int cg = cp.codes[(size_t) (g + 1)]; + char name[64]; + snprintf(name, sizeof(name), "code_pred.codec_embd.%d.weight", g); + embed_row_from_gguf(pt->gguf_talker, name, cg, hidden, tmp.data()); + for (int i = 0; i < hidden; i++) { + next_emb[(size_t) i] += tmp[(size_t) i]; + } + } + for (int i = 0; i < hidden; i++) { + next_emb[(size_t) i] += prev_overlay[(size_t) i]; + } DebugDumper d; debug_init(&d, params->dump_dir); debug_dump_1d(&d, "next-emb-step0", next_emb.data(), hidden); diff --git a/src/pipeline-tts.h b/src/pipeline-tts.h index 740f53d..f451516 100644 --- a/src/pipeline-tts.h +++ b/src/pipeline-tts.h @@ -95,6 +95,10 @@ struct PipelineTTS { SpeakerEncoderWeights speaker_encoder; bool has_speaker_encoder; + // Speaker encoder weights residency: loaded lazily on the first + // reference audio request, see pipeline-tts.cpp. + bool spk_enc_loaded; + PipelineCodec codec; std::string tokenizer_type; diff --git a/src/talker-forward.h b/src/talker-forward.h index 115fbae..7ddbd89 100644 --- a/src/talker-forward.h +++ b/src/talker-forward.h @@ -254,21 +254,28 @@ static struct ggml_tensor * talker_layer_forward(struct ggml_context * ctx, // Shared core that builds the graph, allocates, uploads inputs, runs // it and pulls out the last position hidden + logits. T tokens are -// appended to the cache starting at n_past. When n_past == 0 and -// dump_dir is set, the bisect taps fire on the prefill path. use_fa / -// clamp_fp16 are forwarded as is to every layer. The graph metadata -// lives in the caller owned persistent arena. -static bool talker_forward_core(const TalkerWeights * tw, - KVCache * kv, - ggml_backend_sched_t sched, - GraphArena * arena, - const float * input_embed, - int T, - int n_past, - bool use_flash_attn, - bool clamp_fp16, - const char * dump_dir, - TalkerForwardOutput * out) { +// appended to the cache starting at n_past. Input is either a raw +// embedding upload (input_embed, prefill) or the previous frame code +// ids plus overlay row assembled in graph (frame_ids, decode hot +// path). When n_past == 0 and dump_dir is set, the bisect taps fire on +// the prefill path. use_fa / clamp_fp16 are forwarded as is to every +// layer. The graph metadata lives in the caller owned persistent +// arena. +static bool talker_forward_core(const TalkerWeights * tw, + KVCache * kv, + ggml_backend_sched_t sched, + GraphArena * arena, + const float * input_embed, + const int32_t * frame_ids, + struct ggml_tensor * const * acoustic_embd, + int n_acoustic, + const float * overlay, + int T, + int n_past, + bool use_flash_attn, + bool clamp_fp16, + const char * dump_dir, + TalkerForwardOutput * out) { const int hidden = tw->hidden_size; const int n_layers = tw->num_hidden_layers; const int vocab = tw->vocab_size; @@ -285,17 +292,44 @@ static bool talker_forward_core(const TalkerWeights * tw, const int max_nodes = talker_graph_max_nodes(n_layers); struct ggml_context * gctx = graph_arena_begin(arena); - // IO tensors: input embedding, positions, causal mask. The mask - // spans [n_kv_pad, T]: for each fresh query q in [0, T) keys k in + // IO tensors: positions and causal mask. The mask spans + // [n_kv_pad, T]: for each fresh query q in [0, T) keys k in // [0, n_past + q] carry 0 and every other slot carries neg inf, // including the padded tail beyond T_full. - struct ggml_tensor * x_in = ggml_new_tensor_2d(gctx, GGML_TYPE_F32, hidden, T); struct ggml_tensor * pos_in = ggml_new_tensor_1d(gctx, GGML_TYPE_I32, T); struct ggml_tensor * mask_in = ggml_new_tensor_2d(gctx, GGML_TYPE_F16, n_kv_pad, T); - ggml_set_name(x_in, "input_embed"); ggml_set_name(pos_in, "positions"); ggml_set_name(mask_in, "causal_mask"); + // Input: either a raw embedding upload (prefill path) or, on the + // decode hot path, the frame codes of the previous step gathered + // and summed in graph. x = get_rows(codec_embd, ids[0]) plus the 15 + // acoustic group gathers plus the trailing text / pad overlay row. + // The only per step uploads are 16 code ids and one overlay row. + struct ggml_tensor * x_in = NULL; + struct ggml_tensor * ids_in = NULL; + struct ggml_tensor * overlay_in = NULL; + if (frame_ids) { + ids_in = ggml_new_tensor_1d(gctx, GGML_TYPE_I32, 1 + n_acoustic); + ggml_set_name(ids_in, "frame_code_ids"); + ggml_set_input(ids_in); + overlay_in = ggml_new_tensor_2d(gctx, GGML_TYPE_F32, hidden, 1); + ggml_set_name(overlay_in, "overlay_row"); + ggml_set_input(overlay_in); + + struct ggml_tensor * id0 = ggml_view_1d(gctx, ids_in, 1, 0); + x_in = ggml_get_rows(gctx, tw->codec_embedding, id0); + for (int g = 0; g < n_acoustic; g++) { + struct ggml_tensor * idg = ggml_view_1d(gctx, ids_in, 1, (size_t) (g + 1) * sizeof(int32_t)); + x_in = ggml_add(gctx, x_in, ggml_get_rows(gctx, acoustic_embd[g], idg)); + } + x_in = ggml_add(gctx, x_in, overlay_in); + } else { + x_in = ggml_new_tensor_2d(gctx, GGML_TYPE_F32, hidden, T); + ggml_set_input(x_in); + } + ggml_set_name(x_in, "input_embed"); + struct ggml_cgraph * gf = ggml_new_graph_custom(gctx, max_nodes, false); // Build the layer stack. Bisect taps fire on prefill only. @@ -346,7 +380,12 @@ static bool talker_forward_core(const TalkerWeights * tw, } // Upload input embedding (host [T, hidden] -> ggml [hidden, T]). - ggml_backend_tensor_set(x_in, input_embed, 0, (size_t) T * (size_t) hidden * sizeof(float)); + if (frame_ids) { + ggml_backend_tensor_set(ids_in, frame_ids, 0, (size_t) (1 + n_acoustic) * sizeof(int32_t)); + ggml_backend_tensor_set(overlay_in, overlay, 0, (size_t) hidden * sizeof(float)); + } else { + ggml_backend_tensor_set(x_in, input_embed, 0, (size_t) T * (size_t) hidden * sizeof(float)); + } // Positions: n_past .. n_past + T - 1 { @@ -444,25 +483,33 @@ static bool talker_forward_prefill(const TalkerWeights * tw, fprintf(stderr, "[TalkerForward] FATAL: prefill T=%d exceeds cache max_seq_len=%d\n", T, kv->max_seq_len); return false; } - return talker_forward_core(tw, kv, sched, arena, input_embed, T, 0, use_flash_attn, clamp_fp16, dump_dir, out); + return talker_forward_core(tw, kv, sched, arena, input_embed, NULL, NULL, 0, NULL, T, 0, use_flash_attn, clamp_fp16, + dump_dir, out); } // Decode: feed exactly one embedding and append one position to the // cache. Reads positions [0, kv->cur_len + 1). Caller is responsible // for ensuring kv->cur_len + 1 <= kv->max_seq_len. -static bool talker_forward_decode(const TalkerWeights * tw, - KVCache * kv, - ggml_backend_sched_t sched, - GraphArena * arena, - const float * input_embed_1, - bool use_flash_attn, - bool clamp_fp16, - TalkerForwardOutput * out) { +// Append one position from the previous frame's codes. frame_ids holds +// [c0, c1..c15], acoustic_embd the 15 group tables owned by the code +// predictor, overlay the trailing text / pad row summed on top. The +// input embedding assembles entirely in graph. +static bool talker_forward_decode(const TalkerWeights * tw, + KVCache * kv, + ggml_backend_sched_t sched, + GraphArena * arena, + const int32_t * frame_ids, + struct ggml_tensor * const * acoustic_embd, + int n_acoustic, + const float * overlay, + bool use_flash_attn, + bool clamp_fp16, + TalkerForwardOutput * out) { if (kv->cur_len + 1 > kv->max_seq_len) { fprintf(stderr, "[TalkerForward] FATAL: decode would overflow cache (%d + 1 > %d)\n", kv->cur_len, kv->max_seq_len); return false; } - return talker_forward_core(tw, kv, sched, arena, input_embed_1, 1, kv->cur_len, use_flash_attn, clamp_fp16, NULL, - out); + return talker_forward_core(tw, kv, sched, arena, NULL, frame_ids, acoustic_embd, n_acoustic, overlay, 1, + kv->cur_len, use_flash_attn, clamp_fp16, NULL, out); }