mtmd: support Qwen3-TTS (note: breaking change to llama-tts binary) (#26254)
* convert text model * main model load ok * convert encoder ok * speaker encoder loading ok * speaker enc graph * adapt vocab for backbone (with some tricks) * add suppress_tokens * poc new mtmd gen api * convert code_predictor to gguf * load gen_code model ok * add clip_encode * wire up * code gen cgraph init version Co-authored-by: Pascal <admin@serveurperso.com> * code2wav convert to gguf * code2wav graph ok * wire up in/out * (wip) subgraph * wire up * wip, correct code2wav * demo (to be removed) * code2wav preserve kv between calls * demo voice clone * llama: add llama_model_get_tok_embd * mtmd_helper_gen_audio API * fix clamp cold prefix Co-authored-by: Pascal <admin@serveurperso.com> * fuse snake op Co-authored-by: Pascal <admin@serveurperso.com> * demo: use proper sampling * update dev docs * polymorphism helper * revamp llama-tts binary * update docs * fix compile * fix lint * nits * add guide + docs * more timings info * clean up code comments * security fixes * update docs * use ggml_build_forward_select, clean up comments * fix ci * use ISO 639-1 language code * rename CODE2WAV --> GEN_WAV, update docs * clean up * clean up tts.cpp * add seq_id * add step_prompt() * mtmd_helper_model_can_chat * clean up comments --------- Co-authored-by: Pascal <admin@serveurperso.com>
This commit is contained in:
Xuan-Son Nguyen
co-authored by
Pascal
Pascal
parent
1c3c9674de
commit
0713275082
@@ -2,6 +2,11 @@
|
||||
|
||||
#include "../clip-graph.h"
|
||||
|
||||
#include <map>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
/*
|
||||
* IMPORTANT: The mtmd module does NOT accept pull requests that are fully or predominantly AI-generated.
|
||||
* We encourage human contributors to ensure the quality and reliability of the codebase.
|
||||
@@ -215,6 +220,111 @@ struct clip_graph_mimo_audio : clip_graph {
|
||||
ggml_cgraph * build() override;
|
||||
};
|
||||
|
||||
struct clip_graph_qwen3tts_spkenc : clip_graph {
|
||||
clip_graph_qwen3tts_spkenc(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {}
|
||||
ggml_cgraph * build() override;
|
||||
|
||||
ggml_tensor * conv1d_same(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int dilation) const;
|
||||
ggml_tensor * res2net(ggml_tensor * x, const clip_layer & layer, int dilation, int scale) const;
|
||||
ggml_tensor * se_block(ggml_tensor * x, const clip_layer & layer) const;
|
||||
ggml_tensor * se_res2net_block(ggml_tensor * x, const clip_layer & layer, int dilation, int scale) const;
|
||||
ggml_tensor * attentive_stats_pool(ggml_tensor * x) const;
|
||||
};
|
||||
|
||||
struct clip_graph_qwen3tts_gen : clip_graph {
|
||||
clip_graph_qwen3tts_gen(clip_ctx * ctx, const clip_image_f32 & img, clip_gen_process_type gen_process, int top_k, float top_p)
|
||||
: clip_graph(ctx, img), gen_process(gen_process), top_k(top_k), top_p(top_p) {}
|
||||
ggml_cgraph * build() override;
|
||||
|
||||
// which sub-graph build() constructs, fixed at graph-build time
|
||||
clip_gen_process_type gen_process;
|
||||
|
||||
// sampling params, fixed at graph-build time (GEN_CODE only)
|
||||
int top_k;
|
||||
float top_p;
|
||||
|
||||
//
|
||||
// code_gen: backbone hidden state + sampled code0 -> 16 RVQ codes
|
||||
// MTP-style code predictor, one token per codebook
|
||||
//
|
||||
struct code_gen : clip_graph {
|
||||
code_gen(const clip_graph & parent, int top_k, float top_p)
|
||||
: clip_graph(parent), top_k(top_k), top_p(top_p) {}
|
||||
ggml_cgraph * build() override { GGML_ABORT("call prefill()/step() instead"); }
|
||||
|
||||
int top_k;
|
||||
float top_p;
|
||||
|
||||
ggml_tensor * cache_set(ggml_tensor * cache, int row_idx, ggml_tensor * value) const;
|
||||
ggml_tensor * do_sampling(ggml_tensor * logits, ggml_tensor * inp_rand) const;
|
||||
|
||||
ggml_tensor * const_i32(ggml_tensor * anchor, float value) const;
|
||||
ggml_tensor * causal_mask_row(int64_t n_kv_pad, int pos) const;
|
||||
ggml_tensor * project_in(ggml_tensor * cur) const;
|
||||
|
||||
ggml_tensor * layer_forward(
|
||||
ggml_tensor * cur,
|
||||
const clip_layer & layer,
|
||||
ggml_tensor * inp_pos,
|
||||
ggml_tensor * kq_mask,
|
||||
ggml_tensor *& k_cache_layer,
|
||||
ggml_tensor *& v_cache_layer,
|
||||
int64_t n_kv_pad,
|
||||
int pos,
|
||||
int il) const;
|
||||
|
||||
void prefill(
|
||||
std::vector<ggml_tensor *> & k_cache,
|
||||
std::vector<ggml_tensor *> & v_cache,
|
||||
ggml_tensor *& out_code_cache,
|
||||
ggml_tensor * h_state,
|
||||
ggml_tensor * code0_embd,
|
||||
ggml_tensor * inp_rand) const;
|
||||
|
||||
ggml_tensor * step(
|
||||
std::vector<ggml_tensor *> & k_cache,
|
||||
std::vector<ggml_tensor *> & v_cache,
|
||||
ggml_tensor * out_code_cache,
|
||||
ggml_tensor * inp_rand,
|
||||
int step_idx) const;
|
||||
};
|
||||
|
||||
//
|
||||
// code2wav: RVQ codes -> raw PCM (quantizer + pre_conv + pre_transformer + upsample + DAC).
|
||||
//
|
||||
struct code2wav : clip_graph {
|
||||
code2wav(const clip_graph & parent) : clip_graph(parent) {}
|
||||
ggml_cgraph * build() override { GGML_ABORT("call decode() instead"); }
|
||||
|
||||
// state_in: previous call's persisted state, by slot name (see list_c2w_state_slots())
|
||||
std::map<std::string, ggml_tensor *> state_in;
|
||||
// state_out: this call's state to persist, added to the graph outputs by build()
|
||||
mutable std::vector<std::pair<std::string, ggml_tensor *>> state_out;
|
||||
|
||||
// stateful conv ops: read/update their state via state_in/state_out[state_name]
|
||||
ggml_tensor * causal_conv1d(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int dilation, const std::string & state_name) const;
|
||||
ggml_tensor * causal_conv1d_dw(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, const std::string & state_name) const;
|
||||
ggml_tensor * causal_conv_transpose1d(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int stride, const std::string & state_name) const;
|
||||
ggml_tensor * snake(ggml_tensor * x, ggml_tensor * alpha, ggml_tensor * beta) const;
|
||||
|
||||
ggml_tensor * quant_decode(ggml_tensor * inp_codes) const;
|
||||
ggml_tensor * tfm_layer_forward(ggml_tensor * cur, const clip_layer & layer, int il) const;
|
||||
ggml_tensor * convnext_block(ggml_tensor * x, const clip_code2wav::upsample_block & blk, const std::string & state_prefix) const;
|
||||
ggml_tensor * dac_res_unit(ggml_tensor * x, const clip_code2wav::dac_res & res, int dilation, const std::string & state_name) const;
|
||||
|
||||
// inp_codes [1, n_codes] I32 -> this frame's audio samples [n_samples] F32, clamped to [-1, 1]
|
||||
ggml_tensor * decode(ggml_tensor * inp_codes) const;
|
||||
};
|
||||
};
|
||||
|
||||
// one persisted state buffer used by code2wav, see qwen3tts-gen.cpp
|
||||
struct c2w_state_slot {
|
||||
std::string name;
|
||||
int64_t ne0;
|
||||
int64_t ne1;
|
||||
};
|
||||
std::vector<c2w_state_slot> list_c2w_state_slots(const clip_hparams & hparams, const clip_model & model);
|
||||
|
||||
struct clip_graph_kimik25 : clip_graph {
|
||||
clip_graph_kimik25(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {}
|
||||
ggml_cgraph * build() override;
|
||||
|
||||
@@ -0,0 +1,766 @@
|
||||
#include "models.h"
|
||||
|
||||
#include <string>
|
||||
|
||||
// on-device sampling: top-k, top-p, then a random draw
|
||||
ggml_tensor * clip_graph_qwen3tts_gen::code_gen::do_sampling(ggml_tensor * logits, ggml_tensor * inp_rand) const {
|
||||
logits = ggml_reshape_1d(ctx0, logits, ggml_nelements(logits));
|
||||
const int64_t n_vocab = logits->ne[0];
|
||||
|
||||
// sort a's rows by idx
|
||||
auto sort_by = [this](ggml_tensor * a, ggml_tensor * idx) {
|
||||
ggml_tensor * a2d = ggml_reshape_2d(ctx0, a, 1, a->ne[0]);
|
||||
return ggml_reshape_1d(ctx0, ggml_get_rows(ctx0, a2d, idx), idx->ne[0]);
|
||||
};
|
||||
|
||||
ggml_tensor * cur = logits;
|
||||
ggml_tensor * candidates = nullptr; // maps row index back to vocab id
|
||||
|
||||
if (top_k > 0 && top_k < n_vocab) {
|
||||
ggml_tensor * idx = ggml_top_k(ctx0, cur, top_k);
|
||||
candidates = idx;
|
||||
cur = sort_by(cur, idx);
|
||||
cb(cur, "sample_top_k_logits", -1);
|
||||
}
|
||||
|
||||
if (top_p < 1.0f) {
|
||||
ggml_tensor * sorted_idx = ggml_argsort(ctx0, cur, GGML_SORT_ORDER_DESC);
|
||||
ggml_tensor * sorted_logits = sort_by(cur, sorted_idx);
|
||||
candidates = candidates ? sort_by(candidates, sorted_idx) : sorted_idx;
|
||||
|
||||
ggml_tensor * probs = ggml_soft_max(ctx0, sorted_logits);
|
||||
ggml_tensor * cdf = ggml_cumsum(ctx0, probs);
|
||||
|
||||
// keep_mask[i] = 1 once cdf[i] crosses top_p
|
||||
ggml_tensor * cdf_scaled = ggml_scale_bias(ctx0, cdf, -1.0f, top_p);
|
||||
ggml_tensor * keep_mask = ggml_step(ctx0, cdf_scaled);
|
||||
ggml_tensor * idxf = ggml_sum(ctx0, keep_mask);
|
||||
idxf = ggml_clamp(ctx0, idxf, 0.0f, (float) keep_mask->ne[0] - 1);
|
||||
ggml_tensor * ones = ggml_scale_bias(ctx0, idxf, 0.0f, 1.0f);
|
||||
|
||||
// top-p must include the crossing element, so force it to 1
|
||||
ggml_tensor * keep_mask_2d = ggml_reshape_2d(ctx0, keep_mask, 1, keep_mask->ne[0]);
|
||||
keep_mask_2d = ggml_set_rows(ctx0, keep_mask_2d, ones, ggml_cast(ctx0, idxf, GGML_TYPE_I32));
|
||||
keep_mask = ggml_reshape_1d(ctx0, keep_mask_2d, keep_mask->ne[0]);
|
||||
|
||||
// log(1) = 0 (keep), log(0) = -inf (drop)
|
||||
ggml_tensor * bias = ggml_log(ctx0, keep_mask);
|
||||
cur = ggml_add(ctx0, sorted_logits, bias);
|
||||
cb(cur, "sample_top_p_logits", -1);
|
||||
}
|
||||
|
||||
// draw one token: find where the cdf crosses inp_rand
|
||||
ggml_tensor * probs = ggml_soft_max(ctx0, cur);
|
||||
ggml_tensor * cumsum = ggml_cumsum(ctx0, probs);
|
||||
|
||||
ggml_tensor * diff = ggml_sub(ctx0, cumsum, inp_rand);
|
||||
ggml_tensor * cross_mask = ggml_step(ctx0, diff);
|
||||
ggml_tensor * idxf = ggml_sum(ctx0, cross_mask);
|
||||
ggml_tensor * idx = ggml_cast(ctx0, ggml_scale_bias(ctx0, idxf, -1.0f, (float) cross_mask->ne[0]), GGML_TYPE_I32);
|
||||
|
||||
if (candidates) {
|
||||
ggml_tensor * cand_2d = ggml_reshape_2d(ctx0, candidates, 1, candidates->ne[0]);
|
||||
idx = ggml_get_rows(ctx0, cand_2d, idx);
|
||||
}
|
||||
cb(idx, "sample_token_id", -1);
|
||||
|
||||
return idx;
|
||||
}
|
||||
|
||||
// returns a new cache with row row_idx set to value
|
||||
ggml_tensor * clip_graph_qwen3tts_gen::code_gen::cache_set(ggml_tensor * cache, int row_idx, ggml_tensor * value) const {
|
||||
const int64_t n_embd = cache->ne[0];
|
||||
const int64_t n_cache = cache->ne[1];
|
||||
GGML_ASSERT(row_idx >= 0 && row_idx < n_cache);
|
||||
|
||||
// append value as the last row, then gather it back into place
|
||||
ggml_tensor * value_2d = ggml_reshape_2d(ctx0, value, n_embd, 1);
|
||||
ggml_tensor * cache_ext = ggml_concat(ctx0, cache, value_2d, 1); // [n_embd, n_cache + 1]
|
||||
|
||||
// gather indices [0..row_idx-1, n_cache, row_idx+1..n_cache-1]
|
||||
// built via concat, since ggml_set_rows needs F32/F16 values, not an I32 index array
|
||||
ggml_tensor * idx = const_i32(cache, (float) n_cache);
|
||||
if (row_idx > 0) {
|
||||
ggml_tensor * prefix = ggml_cast(ctx0, ggml_arange(ctx0, 0.0f, (float) row_idx, 1.0f), GGML_TYPE_I32);
|
||||
idx = ggml_concat(ctx0, prefix, idx, 0);
|
||||
}
|
||||
if (row_idx < n_cache - 1) {
|
||||
ggml_tensor * suffix = ggml_cast(ctx0, ggml_arange(ctx0, (float) (row_idx + 1), (float) n_cache, 1.0f), GGML_TYPE_I32);
|
||||
idx = ggml_concat(ctx0, idx, suffix, 0);
|
||||
}
|
||||
|
||||
ggml_tensor * result = ggml_get_rows(ctx0, cache_ext, idx);
|
||||
cb(result, "cache_set_out", -1);
|
||||
return result;
|
||||
}
|
||||
|
||||
// builds a const i32 with no host upload: view a tensor, zero it via scale, add value, cast to i32
|
||||
ggml_tensor * clip_graph_qwen3tts_gen::code_gen::const_i32(ggml_tensor * anchor, float value) const {
|
||||
ggml_tensor * v = ggml_view_1d(ctx0, anchor, 1, 0);
|
||||
if (v->type != GGML_TYPE_F32) {
|
||||
v = ggml_cast(ctx0, v, GGML_TYPE_F32);
|
||||
}
|
||||
return ggml_cast(ctx0, ggml_scale_bias(ctx0, v, 0.0f, value), GGML_TYPE_I32);
|
||||
}
|
||||
|
||||
// causal keep-mask row for a query at position pos, window size n_kv_pad
|
||||
ggml_tensor * clip_graph_qwen3tts_gen::code_gen::causal_mask_row(int64_t n_kv_pad, int pos) const {
|
||||
ggml_tensor * ones = ggml_fill(ctx0, ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_kv_pad, n_kv_pad), 1.0f);
|
||||
ggml_tensor * keep = ggml_tri(ctx0, ones, GGML_TRI_TYPE_LOWER_DIAG);
|
||||
ggml_tensor * row = ggml_view_1d(ctx0, keep, n_kv_pad, (size_t) pos * keep->nb[1]);
|
||||
ggml_tensor * mask = ggml_log(ctx0, row); // 0 = keep, -inf = masked
|
||||
return ggml_reshape_4d(ctx0, mask, n_kv_pad, 1, 1, 1);
|
||||
}
|
||||
|
||||
// talker hidden size -> predictor hidden size (small_to_mtp_projection)
|
||||
ggml_tensor * clip_graph_qwen3tts_gen::code_gen::project_in(ggml_tensor * cur) const {
|
||||
if (!model.gen_code_proj_in_w) {
|
||||
return cur;
|
||||
}
|
||||
cur = ggml_mul_mat(ctx0, model.gen_code_proj_in_w, cur);
|
||||
if (model.gen_code_proj_in_b) {
|
||||
cur = ggml_add(ctx0, cur, model.gen_code_proj_in_b);
|
||||
}
|
||||
return cur;
|
||||
}
|
||||
|
||||
// one transformer layer at position pos; writes k/v into k_cache_layer/v_cache_layer at row pos
|
||||
ggml_tensor * clip_graph_qwen3tts_gen::code_gen::layer_forward(
|
||||
ggml_tensor * cur,
|
||||
const clip_layer & layer,
|
||||
ggml_tensor * inp_pos,
|
||||
ggml_tensor * kq_mask,
|
||||
ggml_tensor *& k_cache_layer,
|
||||
ggml_tensor *& v_cache_layer,
|
||||
int64_t n_kv_pad,
|
||||
int pos,
|
||||
int il) const {
|
||||
const int n_head = hparams.n_head;
|
||||
const int n_head_kv = hparams.n_head_kv;
|
||||
const int64_t d_head = layer.q_w->ne[1] / n_head; // real head_dim, not n_embd / n_head
|
||||
const float kq_scale = 1.0f / sqrtf((float) d_head);
|
||||
|
||||
ggml_tensor * residual = cur;
|
||||
|
||||
ggml_tensor * h = ggml_rms_norm(ctx0, cur, hparams.eps);
|
||||
h = ggml_mul(ctx0, h, layer.ln_1_w);
|
||||
|
||||
ggml_tensor * q = ggml_mul_mat(ctx0, layer.q_w, h);
|
||||
ggml_tensor * k = ggml_mul_mat(ctx0, layer.k_w, h);
|
||||
ggml_tensor * v = ggml_mul_mat(ctx0, layer.v_w, h);
|
||||
|
||||
q = ggml_reshape_3d(ctx0, q, d_head, n_head, 1);
|
||||
k = ggml_reshape_3d(ctx0, k, d_head, n_head_kv, 1);
|
||||
|
||||
q = ggml_rms_norm(ctx0, q, hparams.eps);
|
||||
q = ggml_mul(ctx0, q, layer.q_norm);
|
||||
k = ggml_rms_norm(ctx0, k, hparams.eps);
|
||||
k = ggml_mul(ctx0, k, layer.k_norm);
|
||||
|
||||
q = ggml_rope_ext(ctx0, q, inp_pos, nullptr, (int) d_head, GGML_ROPE_TYPE_NEOX, 0,
|
||||
hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
|
||||
k = ggml_rope_ext(ctx0, k, inp_pos, nullptr, (int) d_head, GGML_ROPE_TYPE_NEOX, 0,
|
||||
hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
|
||||
|
||||
// write k/v into the cache at row pos, flat layout
|
||||
ggml_tensor * k_flat = ggml_reshape_1d(ctx0, k, d_head * n_head_kv);
|
||||
k_cache_layer = cache_set(k_cache_layer, pos, k_flat);
|
||||
v_cache_layer = cache_set(v_cache_layer, pos, v);
|
||||
|
||||
ggml_tensor * q_cur = ggml_reshape_4d(ctx0, q, d_head, n_head, 1, 1);
|
||||
ggml_tensor * k_cur = ggml_reshape_4d(ctx0, k_cache_layer, d_head, n_head_kv, n_kv_pad, 1);
|
||||
ggml_tensor * v_cur = ggml_reshape_4d(ctx0, v_cache_layer, d_head, n_head_kv, n_kv_pad, 1);
|
||||
|
||||
ggml_tensor * attn_out = build_attn(layer.o_w, layer.o_b, q_cur, k_cur, v_cur, kq_mask, kq_scale, il);
|
||||
|
||||
cur = ggml_add(ctx0, residual, attn_out);
|
||||
|
||||
ggml_tensor * h2 = ggml_rms_norm(ctx0, cur, hparams.eps);
|
||||
h2 = ggml_mul(ctx0, h2, layer.ln_2_w);
|
||||
|
||||
ggml_tensor * gate = ggml_mul_mat(ctx0, layer.ff_gate_w, h2);
|
||||
ggml_tensor * up = ggml_mul_mat(ctx0, layer.ff_up_w, h2);
|
||||
ggml_tensor * gu = ggml_swiglu_split(ctx0, gate, up);
|
||||
ggml_tensor * down = ggml_mul_mat(ctx0, layer.ff_down_w, gu);
|
||||
|
||||
return ggml_add(ctx0, cur, down);
|
||||
}
|
||||
|
||||
// position 0: hidden bridge, seeds the k/v cache, no sampling
|
||||
// position 1: embed(code0), sample with lm_head[0], write out_code_cache[1]
|
||||
void clip_graph_qwen3tts_gen::code_gen::prefill(
|
||||
std::vector<ggml_tensor *> & k_cache,
|
||||
std::vector<ggml_tensor *> & v_cache,
|
||||
ggml_tensor *& out_code_cache,
|
||||
ggml_tensor * h_state,
|
||||
ggml_tensor * code0_embd,
|
||||
ggml_tensor * inp_rand) const {
|
||||
const int64_t n_kv_pad = k_cache[0]->ne[1];
|
||||
|
||||
{
|
||||
ggml_tensor * cur = project_in(h_state);
|
||||
ggml_tensor * kq_mask = causal_mask_row(n_kv_pad, 0);
|
||||
ggml_tensor * inp_pos = const_i32(k_cache[0], 0.0f);
|
||||
for (size_t il = 0; il < model.layers.size(); il++) {
|
||||
cur = layer_forward(cur, model.layers[il], inp_pos, kq_mask, k_cache[il], v_cache[il], n_kv_pad, 0, (int) il);
|
||||
}
|
||||
// position 0's output is unused, it only seeded the cache
|
||||
}
|
||||
|
||||
{
|
||||
ggml_tensor * cur = project_in(code0_embd);
|
||||
ggml_tensor * kq_mask = causal_mask_row(n_kv_pad, 1);
|
||||
ggml_tensor * inp_pos = const_i32(k_cache[0], 1.0f);
|
||||
for (size_t il = 0; il < model.layers.size(); il++) {
|
||||
cur = layer_forward(cur, model.layers[il], inp_pos, kq_mask, k_cache[il], v_cache[il], n_kv_pad, 1, (int) il);
|
||||
}
|
||||
|
||||
cur = ggml_rms_norm(ctx0, cur, hparams.eps);
|
||||
cur = ggml_mul(ctx0, cur, model.gen_code_norm_w);
|
||||
|
||||
ggml_tensor * head_w = model.gen_code_head_w;
|
||||
ggml_tensor * head_g = ggml_view_2d(ctx0, head_w, head_w->ne[0], head_w->ne[1], head_w->nb[1], 0); // lm_head[0]
|
||||
ggml_tensor * logits = ggml_mul_mat(ctx0, head_g, cur);
|
||||
|
||||
ggml_tensor * sampled = do_sampling(logits, inp_rand);
|
||||
out_code_cache = cache_set(out_code_cache, 1, sampled);
|
||||
}
|
||||
}
|
||||
|
||||
// one decode step of code_predictor
|
||||
// at step_idx g:
|
||||
// - read code from out_code_cache[g], then embed it with codebook table g-1
|
||||
// - write new kv at cache row g+1, sample with lm_head[g]
|
||||
// - write result to out_code_cache[g+1]
|
||||
// step_idx must be in [1, n_acoustic - 1]
|
||||
ggml_tensor * clip_graph_qwen3tts_gen::code_gen::step(
|
||||
std::vector<ggml_tensor *> & k_cache,
|
||||
std::vector<ggml_tensor *> & v_cache,
|
||||
ggml_tensor * out_code_cache,
|
||||
ggml_tensor * inp_rand,
|
||||
int step_idx) const {
|
||||
const int64_t n_acoustic = model.gen_code_head_w->ne[2];
|
||||
GGML_ASSERT(step_idx >= 1 && step_idx < n_acoustic);
|
||||
GGML_ASSERT(k_cache.size() == model.layers.size());
|
||||
GGML_ASSERT(v_cache.size() == model.layers.size());
|
||||
|
||||
const int64_t n_kv_pad = k_cache[0]->ne[1];
|
||||
const int pos = step_idx + 1; // new cache row and RoPE position
|
||||
|
||||
// embed the previous code via this step's codebook table (rows are already scalars)
|
||||
ggml_tensor * code_in = ggml_view_1d(ctx0, out_code_cache, 1, (size_t) step_idx * out_code_cache->nb[1]);
|
||||
|
||||
ggml_tensor * embd_w = model.gen_code_embd_w; // [n_embd_talker, vocab, n_acoustic]
|
||||
ggml_tensor * embd_g = ggml_view_2d(ctx0, embd_w, embd_w->ne[0], embd_w->ne[1], embd_w->nb[1],
|
||||
(size_t) (step_idx - 1) * embd_w->nb[2]);
|
||||
ggml_tensor * cur = ggml_get_rows(ctx0, embd_g, code_in);
|
||||
cur = ggml_reshape_1d(ctx0, cur, cur->ne[0]);
|
||||
cb(cur, "step_embd_in", step_idx);
|
||||
|
||||
cur = project_in(cur);
|
||||
cb(cur, "step_proj_in", step_idx);
|
||||
|
||||
ggml_tensor * kq_mask = causal_mask_row(n_kv_pad, pos);
|
||||
ggml_tensor * inp_pos = const_i32(k_cache[0], (float) pos);
|
||||
|
||||
for (size_t il = 0; il < model.layers.size(); il++) {
|
||||
cur = layer_forward(cur, model.layers[il], inp_pos, kq_mask, k_cache[il], v_cache[il], n_kv_pad, pos, (int) il);
|
||||
cb(cur, "step_layer_out", (int) il);
|
||||
}
|
||||
|
||||
// final norm, this step's lm_head, sample, write the result
|
||||
cur = ggml_rms_norm(ctx0, cur, hparams.eps);
|
||||
cur = ggml_mul(ctx0, cur, model.gen_code_norm_w);
|
||||
|
||||
ggml_tensor * head_w = model.gen_code_head_w; // [n_embd_pred, vocab, n_acoustic]
|
||||
ggml_tensor * head_g = ggml_view_2d(ctx0, head_w, head_w->ne[0], head_w->ne[1], head_w->nb[1],
|
||||
(size_t) step_idx * head_w->nb[2]);
|
||||
ggml_tensor * logits = ggml_mul_mat(ctx0, head_g, cur);
|
||||
cb(logits, "step_logits", step_idx);
|
||||
|
||||
ggml_tensor * sampled = do_sampling(logits, inp_rand);
|
||||
cb(sampled, "step_sampled", step_idx);
|
||||
|
||||
return cache_set(out_code_cache, pos, sampled);
|
||||
}
|
||||
|
||||
// causal conv1d, stride 1: prepend persisted left-context instead of zero-padding, then a plain conv
|
||||
// x: [T, IC] (T-first). w: [K, IC, OC]. state_name empty means K == 1 (no left-context). returns [T, OC]
|
||||
ggml_tensor * clip_graph_qwen3tts_gen::code2wav::causal_conv1d(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int dilation, const std::string & state_name) const {
|
||||
const int K = (int) w->ne[0];
|
||||
const int pad = (K - 1) * dilation;
|
||||
|
||||
ggml_tensor * x_full = x;
|
||||
if (pad > 0) {
|
||||
ggml_tensor * left = state_in.at(state_name); // [pad, IC]
|
||||
x_full = ggml_concat(ctx0, left, x, 0);
|
||||
}
|
||||
ggml_tensor * y = ggml_conv_1d(ctx0, w, x_full, 1, 0, dilation); // [T, OC, 1]
|
||||
y = ggml_reshape_2d(ctx0, y, y->ne[0], y->ne[1]);
|
||||
if (b) {
|
||||
y = ggml_add(ctx0, y, ggml_reshape_2d(ctx0, b, 1, b->ne[0]));
|
||||
}
|
||||
if (pad > 0) {
|
||||
ggml_tensor * new_left = ggml_cont(ctx0, ggml_view_2d(ctx0, x_full, pad, x_full->ne[1], x_full->nb[1],
|
||||
(size_t) (x_full->ne[0] - pad) * x_full->nb[0]));
|
||||
state_out.push_back({state_name, new_left});
|
||||
}
|
||||
return y;
|
||||
}
|
||||
|
||||
// causal depthwise conv1d, stride 1, dilation 1, kernel from w's shape.
|
||||
// x: [T, C]. w: [K, 1, C]. returns [T, C]. see causal_conv1d for the state contract.
|
||||
ggml_tensor * clip_graph_qwen3tts_gen::code2wav::causal_conv1d_dw(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, const std::string & state_name) const {
|
||||
const int K = (int) w->ne[0];
|
||||
const int pad = K - 1;
|
||||
|
||||
ggml_tensor * x_full = x;
|
||||
if (pad > 0) {
|
||||
ggml_tensor * left = state_in.at(state_name); // [pad, C]
|
||||
x_full = ggml_concat(ctx0, left, x, 0);
|
||||
}
|
||||
ggml_tensor * y = ggml_conv_1d_dw(ctx0, w, x_full, 1, 0, 1); // [T, C, 1]
|
||||
y = ggml_reshape_2d(ctx0, y, y->ne[0], y->ne[1]);
|
||||
if (b) {
|
||||
y = ggml_add(ctx0, y, ggml_reshape_2d(ctx0, b, 1, b->ne[0]));
|
||||
}
|
||||
if (pad > 0) {
|
||||
ggml_tensor * new_left = ggml_cont(ctx0, ggml_view_2d(ctx0, x_full, pad, x_full->ne[1], x_full->nb[1],
|
||||
(size_t) (x_full->ne[0] - pad) * x_full->nb[0]));
|
||||
state_out.push_back({state_name, new_left});
|
||||
}
|
||||
return y;
|
||||
}
|
||||
|
||||
// causal ConvTranspose1d, the (kernel - stride) overlap tail is kept as state for the next call
|
||||
// x: [T, IC], w: [K, OC, IC]. state_name empty means K == stride (no overlap). returns [T * stride, OC]
|
||||
ggml_tensor * clip_graph_qwen3tts_gen::code2wav::causal_conv_transpose1d(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int stride, const std::string & state_name) const {
|
||||
const int K = (int) w->ne[0];
|
||||
const int OC = (int) w->ne[1];
|
||||
const int trim = K - stride;
|
||||
const int64_t emit_len = x->ne[0] * stride;
|
||||
|
||||
// transposed conv as GEMM + col2im scatter-add, y: [emit_len + trim, OC]
|
||||
ggml_tensor * w2 = ggml_reshape_2d(ctx0, w, (int64_t) K * OC, w->ne[2]);
|
||||
w2 = ggml_cont(ctx0, ggml_transpose(ctx0, w2));
|
||||
ggml_tensor * xt = ggml_cont(ctx0, ggml_transpose(ctx0, x));
|
||||
ggml_tensor * col = ggml_mul_mat(ctx0, w2, xt);
|
||||
ggml_tensor * y = ggml_col2im_1d(ctx0, col, stride, OC, 0);
|
||||
|
||||
ggml_tensor * out = y;
|
||||
if (trim > 0) {
|
||||
ggml_tensor * tail = state_in.at(state_name); // [trim, OC]
|
||||
ggml_tensor * head = ggml_add(ctx0, ggml_view_2d(ctx0, y, trim, y->ne[1], y->nb[1], 0), tail);
|
||||
if (emit_len > trim) {
|
||||
ggml_tensor * middle = ggml_view_2d(ctx0, y, emit_len - trim, y->ne[1], y->nb[1], (size_t) trim * y->nb[0]);
|
||||
out = ggml_concat(ctx0, head, middle, 0);
|
||||
} else {
|
||||
out = head;
|
||||
}
|
||||
ggml_tensor * new_tail = ggml_cont(ctx0, ggml_view_2d(ctx0, y, trim, y->ne[1], y->nb[1], (size_t) emit_len * y->nb[0]));
|
||||
state_out.push_back({state_name, new_tail});
|
||||
}
|
||||
if (b) {
|
||||
out = ggml_add(ctx0, out, ggml_reshape_2d(ctx0, b, 1, b->ne[0]));
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
// SnakeBeta activation: y = x + sin(alpha*x)^2 * inv_beta (alpha/inv_beta folded via exp/reciprocal at conversion time)
|
||||
// x: [T, C]. alpha/beta: [C], broadcasts over T
|
||||
ggml_tensor * clip_graph_qwen3tts_gen::code2wav::snake(ggml_tensor * x, ggml_tensor * alpha, ggml_tensor * beta) const {
|
||||
ggml_tensor * a = ggml_reshape_2d(ctx0, alpha, 1, alpha->ne[0]);
|
||||
ggml_tensor * b = ggml_reshape_2d(ctx0, beta, 1, beta->ne[0]);
|
||||
|
||||
// expand reshapes first so mul/sin/sqr/mul/add lands as consecutive nodes, letting backends fuse them
|
||||
ggml_build_forward_expand(gf, a);
|
||||
ggml_build_forward_expand(gf, b);
|
||||
|
||||
ggml_tensor * s = ggml_sin(ctx0, ggml_mul(ctx0, x, a));
|
||||
s = ggml_sqr(ctx0, s);
|
||||
s = ggml_mul(ctx0, s, b);
|
||||
return ggml_add(ctx0, x, s);
|
||||
}
|
||||
|
||||
// RVQ codebook decode: T frames of 16 codes -> 512-dim hidden (C-first, [512, T])
|
||||
// codebook 0 (semantic) and 1..15 (acoustic) sum within their group, project separately, then add
|
||||
ggml_tensor * clip_graph_qwen3tts_gen::code2wav::quant_decode(ggml_tensor * inp_codes) const {
|
||||
const auto & c2w = model.c2w;
|
||||
const int64_t T = inp_codes->ne[0];
|
||||
|
||||
// ids for codebook group g over all T frames, [T] I32
|
||||
auto group_ids = [&](int g) {
|
||||
return ggml_view_1d(ctx0, inp_codes, T, (size_t) g * inp_codes->nb[1]);
|
||||
};
|
||||
|
||||
ggml_tensor * sem = ggml_get_rows(ctx0, c2w.quant_first_cb_w, group_ids(0)); // [256, T]
|
||||
ggml_tensor * sem_out = ggml_mul_mat(ctx0, c2w.quant_first_out_w, sem); // [512, T]
|
||||
|
||||
ggml_tensor * acc = nullptr;
|
||||
const int64_t n_acoustic = c2w.quant_rest_cb_w->ne[2];
|
||||
for (int g = 1; g <= n_acoustic; g++) {
|
||||
ggml_tensor * cb_g = ggml_view_2d(ctx0, c2w.quant_rest_cb_w, c2w.quant_rest_cb_w->ne[0], c2w.quant_rest_cb_w->ne[1],
|
||||
c2w.quant_rest_cb_w->nb[1], (size_t) (g - 1) * c2w.quant_rest_cb_w->nb[2]);
|
||||
ggml_tensor * embd = ggml_get_rows(ctx0, cb_g, group_ids(g)); // [256, T]
|
||||
acc = acc ? ggml_add(ctx0, acc, embd) : embd;
|
||||
}
|
||||
ggml_tensor * ac_out = ggml_mul_mat(ctx0, c2w.quant_rest_out_w, acc); // [512, T]
|
||||
|
||||
ggml_tensor * hidden = ggml_add(ctx0, sem_out, ac_out);
|
||||
cb(hidden, "wav_quant_hidden", -1);
|
||||
return hidden;
|
||||
}
|
||||
|
||||
// one pre_transformer layer over a batch of N = sliding_window new frames
|
||||
// attention runs over [(W-1)-frame prefix from the last batch] + [N new frames]
|
||||
// RoPE positions come from a persisted counter, so phases line up across batches
|
||||
ggml_tensor * clip_graph_qwen3tts_gen::code2wav::tfm_layer_forward(ggml_tensor * cur, const clip_layer & layer, int il) const {
|
||||
const int n_head = hparams.wav_tfm_n_head;
|
||||
const int n_head_kv = hparams.wav_tfm_n_head_kv;
|
||||
const int64_t d_head = layer.q_w->ne[1] / n_head;
|
||||
const float kq_scale = 1.0f / sqrtf((float) d_head);
|
||||
const int64_t W = hparams.wav_tfm_swa; // == N, frames per batch
|
||||
const int64_t N = cur->ne[1];
|
||||
const int64_t prefix = W - 1;
|
||||
const int64_t total_kv = prefix + N;
|
||||
|
||||
ggml_tensor * residual = cur;
|
||||
ggml_tensor * h = ggml_rms_norm(ctx0, cur, hparams.wav_tfm_eps);
|
||||
h = ggml_mul(ctx0, h, layer.ln_1_w);
|
||||
|
||||
ggml_tensor * q = ggml_mul_mat(ctx0, layer.q_w, h); // [n_head*d_head, N]
|
||||
ggml_tensor * k = ggml_mul_mat(ctx0, layer.k_w, h); // [n_head_kv*d_head, N]
|
||||
ggml_tensor * v = ggml_mul_mat(ctx0, layer.v_w, h); // [n_head_kv*d_head, N]
|
||||
|
||||
q = ggml_reshape_3d(ctx0, q, d_head, n_head, N);
|
||||
k = ggml_reshape_3d(ctx0, k, d_head, n_head_kv, N);
|
||||
|
||||
// real, ever-increasing positions: base (persisted) .. base+N-1
|
||||
ggml_tensor * base = ggml_reshape_1d(ctx0, state_in.at("tfm_pos"), 1);
|
||||
ggml_tensor * offset = ggml_arange(ctx0, 0.0f, (float) N, 1.0f);
|
||||
ggml_tensor * pos = ggml_cast(ctx0, ggml_add(ctx0, offset, base), GGML_TYPE_I32);
|
||||
|
||||
q = ggml_rope_ext(ctx0, q, pos, nullptr, (int) d_head, GGML_ROPE_TYPE_NEOX, 0,
|
||||
hparams.wav_tfm_rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
|
||||
k = ggml_rope_ext(ctx0, k, pos, nullptr, (int) d_head, GGML_ROPE_TYPE_NEOX, 0,
|
||||
hparams.wav_tfm_rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
|
||||
|
||||
// the position counter is the same for all layers, push it once from layer 0
|
||||
if (il == 0) {
|
||||
state_out.push_back({"tfm_pos", ggml_scale_bias(ctx0, state_in.at("tfm_pos"), 1.0f, (float) N)});
|
||||
}
|
||||
|
||||
ggml_tensor * k_new = ggml_reshape_2d(ctx0, k, d_head * n_head_kv, N);
|
||||
ggml_tensor * v_new = ggml_reshape_2d(ctx0, v, d_head * n_head_kv, N);
|
||||
|
||||
ggml_tensor * old_k = state_in.at("tfm_k_" + std::to_string(il)); // [d_head*n_head_kv, W-1]
|
||||
ggml_tensor * old_v = state_in.at("tfm_v_" + std::to_string(il));
|
||||
|
||||
ggml_tensor * k_full = ggml_concat(ctx0, old_k, k_new, 1); // [.., prefix+N]
|
||||
ggml_tensor * v_full = ggml_concat(ctx0, old_v, v_new, 1);
|
||||
|
||||
// next batch's prefix: the last (W-1) frames of this batch
|
||||
state_out.push_back({"tfm_k_" + std::to_string(il),
|
||||
ggml_cont(ctx0, ggml_view_2d(ctx0, k_full, k_full->ne[0], prefix, k_full->nb[1], (size_t) N * k_full->nb[1]))});
|
||||
state_out.push_back({"tfm_v_" + std::to_string(il),
|
||||
ggml_cont(ctx0, ggml_view_2d(ctx0, v_full, v_full->ne[0], prefix, v_full->nb[1], (size_t) N * v_full->nb[1]))});
|
||||
|
||||
// banded causal mask: key j is visible to query i iff 0 <= (prefix+i) - j < W
|
||||
ggml_tensor * pos_k = ggml_reshape_2d(ctx0, ggml_arange(ctx0, 0.0f, (float) total_kv, 1.0f), total_kv, 1);
|
||||
ggml_tensor * pos_q = ggml_reshape_2d(ctx0, ggml_arange(ctx0, (float) prefix, (float) (prefix + N), 1.0f), 1, N);
|
||||
ggml_tensor * pos_q_grid = ggml_repeat_4d(ctx0, pos_q, total_kv, N, 1, 1);
|
||||
ggml_tensor * diff = ggml_sub(ctx0, pos_q_grid, pos_k); // [total_kv, N]
|
||||
|
||||
ggml_tensor * causal_keep = ggml_step(ctx0, ggml_scale_bias(ctx0, diff, 1.0f, 0.5f)); // diff >= 0
|
||||
ggml_tensor * in_window = ggml_step(ctx0, ggml_scale_bias(ctx0, diff, -1.0f, (float) W - 0.5f)); // diff < W
|
||||
ggml_tensor * keep = ggml_mul(ctx0, causal_keep, in_window);
|
||||
|
||||
// on a cold start, key j is real state only when j >= prefix - tfm_pos, mask out the rest
|
||||
ggml_tensor * warm = ggml_step(ctx0, ggml_scale_bias(ctx0, ggml_add(ctx0, pos_k, base),
|
||||
1.0f, 0.5f - (float) prefix)); // j + pos > prefix - 0.5
|
||||
keep = ggml_mul(ctx0, keep, warm);
|
||||
|
||||
ggml_tensor * mask = ggml_reshape_4d(ctx0, ggml_log(ctx0, keep), total_kv, N, 1, 1); // 0 = keep, -inf = masked
|
||||
|
||||
ggml_tensor * q_cur = ggml_reshape_4d(ctx0, q, d_head, n_head, N, 1);
|
||||
ggml_tensor * k_cur = ggml_reshape_4d(ctx0, k_full, d_head, n_head_kv, total_kv, 1);
|
||||
ggml_tensor * v_cur = ggml_reshape_4d(ctx0, v_full, d_head, n_head_kv, total_kv, 1);
|
||||
|
||||
ggml_tensor * attn_out = build_attn(layer.o_w, layer.o_b, q_cur, k_cur, v_cur, mask, kq_scale, il);
|
||||
if (layer.ls_1_w) {
|
||||
attn_out = ggml_mul(ctx0, attn_out, layer.ls_1_w);
|
||||
}
|
||||
cur = ggml_add(ctx0, residual, attn_out);
|
||||
|
||||
ggml_tensor * residual2 = cur;
|
||||
ggml_tensor * h2 = ggml_rms_norm(ctx0, cur, hparams.wav_tfm_eps);
|
||||
h2 = ggml_mul(ctx0, h2, layer.ln_2_w);
|
||||
|
||||
ggml_tensor * gate = ggml_mul_mat(ctx0, layer.ff_gate_w, h2);
|
||||
ggml_tensor * up = ggml_mul_mat(ctx0, layer.ff_up_w, h2);
|
||||
ggml_tensor * gu = ggml_swiglu_split(ctx0, gate, up);
|
||||
ggml_tensor * down = ggml_mul_mat(ctx0, layer.ff_down_w, gu);
|
||||
if (layer.ls_2_w) {
|
||||
down = ggml_mul(ctx0, down, layer.ls_2_w);
|
||||
}
|
||||
return ggml_add(ctx0, residual2, down);
|
||||
}
|
||||
|
||||
// dwconv -> LayerNorm -> pwconv1 -> GELU -> pwconv2 -> layer scale -> residual
|
||||
// x: [T, C] T-first; LayerNorm/pwconv need C on ne0, so this transposes in and back out
|
||||
ggml_tensor * clip_graph_qwen3tts_gen::code2wav::convnext_block(ggml_tensor * x, const clip_code2wav::upsample_block & blk, const std::string & state_prefix) const {
|
||||
ggml_tensor * residual = x;
|
||||
|
||||
ggml_tensor * h = causal_conv1d_dw(x, blk.dwconv_w, blk.dwconv_b, state_prefix + "_dwconv"); // [T, C]
|
||||
ggml_tensor * hc = ggml_cont(ctx0, ggml_transpose(ctx0, h)); // [C, T]
|
||||
|
||||
hc = ggml_norm(ctx0, hc, 1e-6f);
|
||||
hc = ggml_mul(ctx0, hc, blk.norm_w);
|
||||
hc = ggml_add(ctx0, hc, blk.norm_b);
|
||||
|
||||
ggml_tensor * g = ggml_mul_mat(ctx0, blk.pw1_w, hc);
|
||||
g = ggml_add(ctx0, g, blk.pw1_b);
|
||||
g = ggml_gelu(ctx0, g);
|
||||
g = ggml_mul_mat(ctx0, blk.pw2_w, g);
|
||||
g = ggml_add(ctx0, g, blk.pw2_b);
|
||||
g = ggml_mul(ctx0, g, blk.gamma);
|
||||
|
||||
ggml_tensor * g_t = ggml_cont(ctx0, ggml_transpose(ctx0, g)); // back to [T, C]
|
||||
return ggml_add(ctx0, residual, g_t);
|
||||
}
|
||||
|
||||
// SnakeBeta -> dilated causal conv (k=7) -> SnakeBeta -> pointwise causal conv (k=1) -> residual.
|
||||
// x: [T, C]. returns [T, C].
|
||||
ggml_tensor * clip_graph_qwen3tts_gen::code2wav::dac_res_unit(ggml_tensor * x, const clip_code2wav::dac_res & res, int dilation, const std::string & state_name) const {
|
||||
ggml_tensor * residual = x;
|
||||
ggml_tensor * h = snake(x, res.act1_alpha, res.act1_beta);
|
||||
h = causal_conv1d(h, res.conv1_w, res.conv1_b, dilation, state_name);
|
||||
h = snake(h, res.act2_alpha, res.act2_beta);
|
||||
h = causal_conv1d(h, res.conv2_w, res.conv2_b, 1, ""); // k=1, no left-context needed
|
||||
return ggml_add(ctx0, residual, h);
|
||||
}
|
||||
|
||||
// RVQ codes -> raw PCM for a batch of N = sliding_window frames
|
||||
ggml_tensor * clip_graph_qwen3tts_gen::code2wav::decode(ggml_tensor * inp_codes) const {
|
||||
const auto & c2w = model.c2w;
|
||||
|
||||
// 1. quantizer decode: N frames of 16 codes -> [512, N] (C-first)
|
||||
ggml_tensor * hidden = quant_decode(inp_codes);
|
||||
|
||||
// 2. pre_conv: [512, N] -> T-first [N, 512] -> causal conv k=3 -> [N, 1024]
|
||||
ggml_tensor * x = ggml_cont(ctx0, ggml_transpose(ctx0, hidden)); // [N, 512]
|
||||
x = causal_conv1d(x, c2w.pre_conv_w, c2w.pre_conv_b, 1, "pre_conv"); // [N, 1024]
|
||||
cb(x, "wav_pre_conv_out", -1);
|
||||
|
||||
// 3. pre_transformer: back to C-first [1024, N], project down, run the layers, project back up
|
||||
ggml_tensor * cur = ggml_cont(ctx0, ggml_transpose(ctx0, x)); // [1024, N]
|
||||
cur = ggml_mul_mat(ctx0, c2w.tfm_in_proj_w, cur);
|
||||
cur = ggml_add(ctx0, cur, c2w.tfm_in_proj_b); // [512 (tfm hidden), N]
|
||||
|
||||
for (int il = 0; il < hparams.wav_tfm_n_layer; il++) {
|
||||
cur = tfm_layer_forward(cur, c2w.tfm_layers[il], il);
|
||||
}
|
||||
|
||||
cur = ggml_rms_norm(ctx0, cur, hparams.wav_tfm_eps);
|
||||
cur = ggml_mul(ctx0, cur, c2w.tfm_output_norm_w);
|
||||
cur = ggml_mul_mat(ctx0, c2w.tfm_out_proj_w, cur);
|
||||
cur = ggml_add(ctx0, cur, c2w.tfm_out_proj_b); // [1024, N]
|
||||
cb(cur, "wav_tfm_out", -1);
|
||||
|
||||
// 4. upsample: 2x (causal ConvTranspose1d, stride 2 + ConvNeXt block), back to T-first
|
||||
// kernel == stride here, so there is no overlap tail to persist
|
||||
x = ggml_cont(ctx0, ggml_transpose(ctx0, cur)); // [N, 1024]
|
||||
for (size_t il = 0; il < c2w.upsample.size(); il++) {
|
||||
const auto & up = c2w.upsample[il];
|
||||
x = causal_conv_transpose1d(x, up.conv_w, up.conv_b, 2, "");
|
||||
x = convnext_block(x, up, "up" + std::to_string(il));
|
||||
cb(x, "wav_upsample_out", (int) il);
|
||||
}
|
||||
|
||||
// 5. DAC decoder: conv_pre -> n blocks (SnakeBeta -> ConvTranspose1d -> 3 res units) -> conv_post
|
||||
static constexpr int DAC_DILATIONS[3] = { 1, 3, 9 };
|
||||
|
||||
x = causal_conv1d(x, c2w.dac_entry_w, c2w.dac_entry_b, 1, "dac_entry");
|
||||
cb(x, "wav_dac_entry_out", -1);
|
||||
|
||||
for (size_t il = 0; il < c2w.dac.size(); il++) {
|
||||
const auto & blk = c2w.dac[il];
|
||||
const int stride = (int) (blk.conv_w->ne[0] / 2); // kernel == 2*stride for all 4 blocks
|
||||
const std::string blk_name = "dac" + std::to_string(il);
|
||||
x = snake(x, blk.snake_alpha, blk.snake_beta);
|
||||
x = causal_conv_transpose1d(x, blk.conv_w, blk.conv_b, stride, blk_name + "_tail");
|
||||
for (size_t ir = 0; ir < blk.res.size(); ir++) {
|
||||
x = dac_res_unit(x, blk.res[ir], DAC_DILATIONS[ir], blk_name + "_res" + std::to_string(ir));
|
||||
}
|
||||
cb(x, "wav_dac_block_out", (int) il);
|
||||
}
|
||||
|
||||
x = snake(x, c2w.dac_post_snake_alpha, c2w.dac_post_snake_beta);
|
||||
x = causal_conv1d(x, c2w.dac_post_conv_w, c2w.dac_post_conv_b, 1, "dac_post_conv"); // [n_samples, 1]
|
||||
|
||||
x = ggml_clamp(ctx0, x, -1.0f, 1.0f);
|
||||
x = ggml_reshape_1d(ctx0, x, x->ne[0]);
|
||||
cb(x, "wav_audio_out", -1);
|
||||
return x;
|
||||
}
|
||||
|
||||
// code2wav's persisted state buffers: RoPE position counter, K/V per pre_transformer layer,
|
||||
// left-context/tail per stateful conv. shape lookup only, no graph needed
|
||||
std::vector<c2w_state_slot> list_c2w_state_slots(const clip_hparams & hparams, const clip_model & model) {
|
||||
const auto & c2w = model.c2w;
|
||||
std::vector<c2w_state_slot> slots;
|
||||
|
||||
slots.push_back({"tfm_pos", 1, 1});
|
||||
|
||||
// prefix is (W-1) frames, the batch itself gives the other N=W frames (see tfm_layer_forward)
|
||||
const int64_t d_head = c2w.tfm_layers[0].q_w->ne[1] / hparams.wav_tfm_n_head;
|
||||
const int64_t kv_ch = d_head * hparams.wav_tfm_n_head_kv;
|
||||
const int64_t prefix = hparams.wav_tfm_swa - 1;
|
||||
for (int il = 0; il < hparams.wav_tfm_n_layer; il++) {
|
||||
slots.push_back({"tfm_k_" + std::to_string(il), kv_ch, prefix});
|
||||
slots.push_back({"tfm_v_" + std::to_string(il), kv_ch, prefix});
|
||||
}
|
||||
|
||||
slots.push_back({"pre_conv", c2w.pre_conv_w->ne[0] - 1, c2w.pre_conv_w->ne[1]});
|
||||
|
||||
for (size_t il = 0; il < c2w.upsample.size(); il++) {
|
||||
const auto & up = c2w.upsample[il];
|
||||
slots.push_back({"up" + std::to_string(il) + "_dwconv", up.dwconv_w->ne[0] - 1, up.dwconv_w->ne[2]});
|
||||
}
|
||||
|
||||
slots.push_back({"dac_entry", c2w.dac_entry_w->ne[0] - 1, c2w.dac_entry_w->ne[1]});
|
||||
|
||||
static constexpr int DAC_DILATIONS[3] = { 1, 3, 9 };
|
||||
for (size_t il = 0; il < c2w.dac.size(); il++) {
|
||||
const auto & blk = c2w.dac[il];
|
||||
const int64_t stride = blk.conv_w->ne[0] / 2; // kernel == 2*stride for all 4 blocks
|
||||
const std::string blk_name = "dac" + std::to_string(il);
|
||||
slots.push_back({blk_name + "_tail", stride, blk.conv_w->ne[1]});
|
||||
for (size_t ir = 0; ir < blk.res.size(); ir++) {
|
||||
const auto & res = blk.res[ir];
|
||||
slots.push_back({blk_name + "_res" + std::to_string(ir),
|
||||
(res.conv1_w->ne[0] - 1) * DAC_DILATIONS[ir], res.conv1_w->ne[1]});
|
||||
}
|
||||
}
|
||||
|
||||
slots.push_back({"dac_post_conv", c2w.dac_post_conv_w->ne[0] - 1, c2w.dac_post_conv_w->ne[1]});
|
||||
|
||||
return slots;
|
||||
}
|
||||
|
||||
// both sub-graphs are always built, so the topology stays constant
|
||||
// ggml_build_forward_select() then picks the one that actually runs
|
||||
ggml_cgraph * clip_graph_qwen3tts_gen::build() {
|
||||
GGML_ASSERT(n_batch == 1); // this module only ever processes one frame at a time
|
||||
|
||||
int idx;
|
||||
switch (gen_process) {
|
||||
case CLIP_GEN_PROCESS_GEN_CODE: idx = 0; break;
|
||||
case CLIP_GEN_PROCESS_GEN_WAV: idx = 1; break;
|
||||
default: GGML_ABORT("unknown gen_process");
|
||||
}
|
||||
|
||||
// ---- CLIP_GEN_PROCESS_GEN_CODE: backbone hidden state -> 16 RVQ codes + next-step embd ----
|
||||
// not build_inp_raw(), a GEN_WAV call's `img` has no hidden-state data
|
||||
ggml_tensor * h_state = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, n_mmproj_embd);
|
||||
ggml_set_name(h_state, "inp_raw"); // must keep this exact name, clip_encode() sets it by name
|
||||
ggml_set_input(h_state);
|
||||
|
||||
ggml_tensor * code0 = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, 1);
|
||||
ggml_set_name(code0, "inp_code0");
|
||||
ggml_set_input(code0);
|
||||
|
||||
ggml_tensor * code0_embd = ggml_get_rows(ctx0, model.gen_code_out_embd_w, code0);
|
||||
code0_embd = ggml_reshape_1d(ctx0, code0_embd, code0_embd->ne[0]);
|
||||
cb(code0_embd, "code0_embd", -1);
|
||||
|
||||
const int64_t n_acoustic = model.gen_code_head_w->ne[2]; // 15
|
||||
const int n_codes = (int) n_acoustic + 1; // 16
|
||||
const int64_t n_kv_pad = n_codes;
|
||||
const int n_layer = (int) model.layers.size();
|
||||
const int n_head = hparams.n_head;
|
||||
const int n_head_kv = hparams.n_head_kv;
|
||||
const int64_t d_head = model.layers[0].q_w->ne[1] / n_head;
|
||||
|
||||
// zero-filled per layer k/v caches, so masked-out rows can't hold garbage
|
||||
std::vector<ggml_tensor *> k_cache(n_layer), v_cache(n_layer);
|
||||
for (int il = 0; il < n_layer; il++) {
|
||||
k_cache[il] = ggml_fill(ctx0, ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, d_head * n_head_kv, n_kv_pad), 0.0f);
|
||||
v_cache[il] = ggml_fill(ctx0, ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, d_head * n_head_kv, n_kv_pad), 0.0f);
|
||||
}
|
||||
|
||||
code_gen cg(*this, top_k, top_p);
|
||||
|
||||
ggml_tensor * out_code_cache = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, 1, n_codes);
|
||||
out_code_cache = cg.cache_set(out_code_cache, 0, code0);
|
||||
|
||||
ggml_tensor * inp_rand0 = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, 1);
|
||||
ggml_set_name(inp_rand0, "inp_rand_0");
|
||||
ggml_set_input(inp_rand0);
|
||||
|
||||
cg.prefill(k_cache, v_cache, out_code_cache, h_state, code0_embd, inp_rand0);
|
||||
|
||||
for (int g = 1; g < n_acoustic; g++) {
|
||||
ggml_tensor * inp_rand = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, 1);
|
||||
ggml_set_name(inp_rand, ("inp_rand_" + std::to_string(g)).c_str());
|
||||
ggml_set_input(inp_rand);
|
||||
out_code_cache = cg.step(k_cache, v_cache, out_code_cache, inp_rand, g);
|
||||
}
|
||||
|
||||
// output 1: this frame's 16 sampled codes, for the caller's code2wav window
|
||||
ggml_tensor * out_codes = ggml_cont(ctx0, out_code_cache);
|
||||
ggml_set_name(out_codes, "out_codes");
|
||||
ggml_set_output(out_codes);
|
||||
|
||||
// output 2: sum of all 16 codebook embeddings, fed back to the talker for the next frame
|
||||
ggml_tensor * out_embd = code0_embd;
|
||||
for (int g = 1; g <= n_acoustic; g++) {
|
||||
ggml_tensor * code_g = ggml_view_1d(ctx0, out_code_cache, 1, (size_t) g * out_code_cache->nb[1]);
|
||||
|
||||
ggml_tensor * embd_g = ggml_view_2d(ctx0, model.gen_code_embd_w, model.gen_code_embd_w->ne[0], model.gen_code_embd_w->ne[1],
|
||||
model.gen_code_embd_w->nb[1], (size_t) (g - 1) * model.gen_code_embd_w->nb[2]);
|
||||
ggml_tensor * e = ggml_get_rows(ctx0, embd_g, code_g);
|
||||
e = ggml_reshape_1d(ctx0, e, e->ne[0]);
|
||||
|
||||
out_embd = ggml_add(ctx0, out_embd, e);
|
||||
}
|
||||
out_embd = ggml_reshape_2d(ctx0, out_embd, out_embd->ne[0], 1);
|
||||
cb(out_embd, "gen_audio_out", -1);
|
||||
|
||||
// ---- CLIP_GEN_PROCESS_GEN_WAV: 16 RVQ codes -> raw PCM ----
|
||||
const int n_frames = hparams.wav_tfm_swa; // frames per batch, == the attention window
|
||||
|
||||
ggml_tensor * inp_codes = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_frames, n_codes);
|
||||
ggml_set_name(inp_codes, "inp_codes");
|
||||
ggml_set_input(inp_codes);
|
||||
|
||||
code2wav c2w(*this);
|
||||
for (const auto & slot : list_c2w_state_slots(hparams, model)) {
|
||||
ggml_tensor * t = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, slot.ne0, slot.ne1);
|
||||
ggml_set_name(t, ("state_in_" + slot.name).c_str());
|
||||
ggml_set_input(t);
|
||||
c2w.state_in[slot.name] = t;
|
||||
}
|
||||
|
||||
ggml_tensor * out_audio = c2w.decode(inp_codes);
|
||||
ggml_set_name(out_audio, "out_audio");
|
||||
ggml_set_output(out_audio);
|
||||
|
||||
for (auto & slot : c2w.state_out) {
|
||||
ggml_set_name(slot.second, ("state_out_" + slot.first).c_str());
|
||||
ggml_set_output(slot.second);
|
||||
}
|
||||
|
||||
// out_embd goes last, clip_encode() reads it back via ggml_graph_node(gf, -1)
|
||||
ggml_tensor * outs[2];
|
||||
outs[0] = out_codes; outs[1] = out_audio;
|
||||
ggml_build_forward_select(gf, outs, 2, idx);
|
||||
for (auto & slot : c2w.state_out) {
|
||||
outs[0] = out_codes; outs[1] = slot.second;
|
||||
ggml_build_forward_select(gf, outs, 2, idx);
|
||||
}
|
||||
outs[0] = out_embd; outs[1] = out_audio;
|
||||
ggml_build_forward_select(gf, outs, 2, idx);
|
||||
|
||||
return gf;
|
||||
}
|
||||
@@ -0,0 +1,197 @@
|
||||
#include "models.h"
|
||||
|
||||
static constexpr int SPK_RES2NET_SCALE = 8; // enc_res2net_scale
|
||||
static constexpr int SPK_DILATIONS[3] = { 2, 3, 4 }; // enc_dilations[1..3]
|
||||
|
||||
// conv1d, kernel K, padding "same" (reflect), dilation d
|
||||
// x: [C, T] (ne[0]=C, ne[1]=T) -> [out_c, T]
|
||||
ggml_tensor * clip_graph_qwen3tts_spkenc::conv1d_same(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int dilation) const {
|
||||
const int K = (int) w->ne[0];
|
||||
const int IC = (int) w->ne[1];
|
||||
const int OC = (int) w->ne[2];
|
||||
const int pad = ((K - 1) * dilation) / 2;
|
||||
|
||||
// ggml_pad_reflect_1d pads ne[0], so bring T onto ne[0] first, same layout as im2col wants
|
||||
ggml_tensor * x_t = ggml_cont(ctx0, ggml_transpose(ctx0, x)); // [T, IC]
|
||||
if (pad > 0) {
|
||||
x_t = ggml_pad_reflect_1d(ctx0, x_t, pad, pad); // [T + 2*pad, IC]
|
||||
}
|
||||
ggml_tensor * x4d = ggml_reshape_4d(ctx0, x_t, x_t->ne[0], IC, 1, 1);
|
||||
|
||||
// dummy F32 kernel, im2col only reads its shape, so a quantized w does not assert
|
||||
ggml_tensor * dummy = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, K, IC, 1, 1);
|
||||
|
||||
ggml_tensor * col = ggml_im2col(ctx0, dummy, x4d, 1, 1, 0, 0, dilation, 1, false, GGML_TYPE_F32);
|
||||
const int64_t T_out = col->ne[1];
|
||||
col = ggml_reshape_2d(ctx0, col, (int64_t) K * IC, T_out);
|
||||
|
||||
ggml_tensor * w2d = ggml_reshape_2d(ctx0, w, (int64_t) K * IC, OC);
|
||||
ggml_tensor * y = ggml_mul_mat(ctx0, w2d, col); // [OC, T_out]
|
||||
ggml_mul_mat_set_prec(y, GGML_PREC_F32);
|
||||
|
||||
ggml_tensor * b2d = ggml_reshape_2d(ctx0, b, OC, 1);
|
||||
y = ggml_add(ctx0, y, b2d);
|
||||
return y;
|
||||
}
|
||||
|
||||
// Res2Net: split channel axis into `scale` chunks, chain dilated conv1d branches
|
||||
// x: [C, T] -> [C, T]
|
||||
ggml_tensor * clip_graph_qwen3tts_spkenc::res2net(ggml_tensor * x, const clip_layer & layer, int dilation, int scale) const {
|
||||
const int64_t C = x->ne[0];
|
||||
const int64_t T = x->ne[1];
|
||||
const int64_t Cs = C / scale;
|
||||
|
||||
std::vector<ggml_tensor *> outs;
|
||||
outs.reserve(scale);
|
||||
|
||||
auto chunk = [&](int i) -> ggml_tensor * {
|
||||
return ggml_view_2d(ctx0, x, Cs, T, x->nb[1], (size_t) i * Cs * x->nb[0]);
|
||||
};
|
||||
|
||||
ggml_tensor * prev = nullptr;
|
||||
for (int i = 0; i < scale; i++) {
|
||||
ggml_tensor * c = ggml_cont(ctx0, chunk(i));
|
||||
if (i == 0) {
|
||||
outs.push_back(c);
|
||||
continue;
|
||||
}
|
||||
ggml_tensor * inp = (i >= 2) ? ggml_add(ctx0, c, prev) : c;
|
||||
ggml_tensor * y = conv1d_same(inp, layer.res2_conv_w[i - 1], layer.res2_conv_b[i - 1], dilation);
|
||||
y = ggml_relu(ctx0, y);
|
||||
outs.push_back(y);
|
||||
prev = y;
|
||||
}
|
||||
|
||||
ggml_tensor * acc = outs[0];
|
||||
for (int i = 1; i < scale; i++) {
|
||||
acc = ggml_concat(ctx0, acc, outs[i], 0);
|
||||
}
|
||||
return acc;
|
||||
}
|
||||
|
||||
// squeeze-and-excitation gate. x: [C, T] -> [C, T]
|
||||
ggml_tensor * clip_graph_qwen3tts_spkenc::se_block(ggml_tensor * x, const clip_layer & layer) const {
|
||||
// temporal mean, keepdim: transpose so T is on ne[0], reduce, transpose back
|
||||
ggml_tensor * x_t = ggml_cont(ctx0, ggml_transpose(ctx0, x)); // [T, C]
|
||||
ggml_tensor * mean = ggml_mean(ctx0, x_t); // [1, C]
|
||||
mean = ggml_cont(ctx0, ggml_transpose(ctx0, mean)); // [C, 1]
|
||||
|
||||
ggml_tensor * h = conv1d_same(mean, layer.se_conv1_w, layer.se_conv1_b, 1);
|
||||
h = ggml_relu(ctx0, h);
|
||||
h = conv1d_same(h, layer.se_conv2_w, layer.se_conv2_b, 1);
|
||||
h = ggml_sigmoid(ctx0, h); // [C, 1]
|
||||
|
||||
return ggml_mul(ctx0, x, h); // broadcast gate over T
|
||||
}
|
||||
|
||||
// tdnn1 -> res2net -> tdnn2 -> se, plus residual. x: [C, T] -> [C, T]
|
||||
ggml_tensor * clip_graph_qwen3tts_spkenc::se_res2net_block(ggml_tensor * x, const clip_layer & layer, int dilation, int scale) const {
|
||||
ggml_tensor * residual = x;
|
||||
ggml_tensor * h = conv1d_same(x, layer.conv_pw1_w, layer.conv_pw1_b, 1); // tdnn1
|
||||
h = ggml_relu(ctx0, h);
|
||||
h = res2net(h, layer, dilation, scale);
|
||||
h = conv1d_same(h, layer.conv_pw2_w, layer.conv_pw2_b, 1); // tdnn2
|
||||
h = ggml_relu(ctx0, h);
|
||||
h = se_block(h, layer);
|
||||
return ggml_add(ctx0, h, residual);
|
||||
}
|
||||
|
||||
// attentive statistics pooling. x: [C, T] -> [2*C, 1]
|
||||
ggml_tensor * clip_graph_qwen3tts_spkenc::attentive_stats_pool(ggml_tensor * x) const {
|
||||
const int64_t T = x->ne[1];
|
||||
|
||||
// mean over T: [C, 1]
|
||||
ggml_tensor * x_t = ggml_cont(ctx0, ggml_transpose(ctx0, x));
|
||||
ggml_tensor * mean = ggml_mean(ctx0, x_t);
|
||||
mean = ggml_cont(ctx0, ggml_transpose(ctx0, mean));
|
||||
|
||||
// std over T: sqrt(clamp(mean((x - mean)^2), eps))
|
||||
ggml_tensor * mean_rep = ggml_repeat(ctx0, mean, x);
|
||||
ggml_tensor * centered = ggml_sub(ctx0, x, mean_rep);
|
||||
ggml_tensor * var_t = ggml_cont(ctx0, ggml_transpose(ctx0, ggml_sqr(ctx0, centered)));
|
||||
ggml_tensor * var = ggml_mean(ctx0, var_t);
|
||||
var = ggml_cont(ctx0, ggml_transpose(ctx0, var));
|
||||
var = ggml_scale_bias(ctx0, var, 1.0f, 1e-12f);
|
||||
ggml_tensor * std = ggml_sqrt(ctx0, var);
|
||||
|
||||
// attention input: cat([x, mean, std]) along channel axis -> [3C, T]
|
||||
ggml_tensor * std_rep = ggml_repeat(ctx0, std, x);
|
||||
ggml_tensor * cat = ggml_concat(ctx0, x, mean_rep, 0);
|
||||
cat = ggml_concat(ctx0, cat, std_rep, 0);
|
||||
|
||||
// attention TDNN (3C -> attn_c) + ReLU, tanh, then 1x1 conv (attn_c -> C)
|
||||
ggml_tensor * a = conv1d_same(cat, model.spk_asp_tdnn_w, model.spk_asp_tdnn_b, 1);
|
||||
a = ggml_relu(ctx0, a);
|
||||
a = ggml_tanh(ctx0, a);
|
||||
a = conv1d_same(a, model.spk_asp_attn_w, model.spk_asp_attn_b, 1);
|
||||
|
||||
// softmax over T
|
||||
ggml_tensor * a_t = ggml_cont(ctx0, ggml_transpose(ctx0, a)); // [T, C]
|
||||
ggml_tensor * w_t = ggml_soft_max(ctx0, a_t);
|
||||
ggml_tensor * w = ggml_cont(ctx0, ggml_transpose(ctx0, w_t)); // [C, T]
|
||||
|
||||
// weighted mean: sum(w * x) over T, multiply by T to undo ggml_mean's 1/T scaling
|
||||
ggml_tensor * wx = ggml_mul(ctx0, w, x);
|
||||
ggml_tensor * wx_t = ggml_cont(ctx0, ggml_transpose(ctx0, wx));
|
||||
ggml_tensor * w_mean = ggml_mean(ctx0, wx_t);
|
||||
w_mean = ggml_scale(ctx0, w_mean, (float) T);
|
||||
w_mean = ggml_cont(ctx0, ggml_transpose(ctx0, w_mean)); // [C, 1]
|
||||
|
||||
// weighted std: sum(w * (x - w_mean)^2) over T
|
||||
ggml_tensor * w_mean_rep = ggml_repeat(ctx0, w_mean, x);
|
||||
ggml_tensor * dev = ggml_sub(ctx0, x, w_mean_rep);
|
||||
ggml_tensor * w_var_in = ggml_mul(ctx0, w, ggml_sqr(ctx0, dev));
|
||||
ggml_tensor * w_var_t = ggml_cont(ctx0, ggml_transpose(ctx0, w_var_in));
|
||||
ggml_tensor * w_var = ggml_mean(ctx0, w_var_t);
|
||||
w_var = ggml_scale(ctx0, w_var, (float) T);
|
||||
w_var = ggml_cont(ctx0, ggml_transpose(ctx0, w_var));
|
||||
w_var = ggml_scale_bias(ctx0, w_var, 1.0f, 1e-12f);
|
||||
ggml_tensor * w_std = ggml_sqrt(ctx0, w_var);
|
||||
|
||||
return ggml_concat(ctx0, w_mean, w_std, 0); // [2C, 1]
|
||||
}
|
||||
|
||||
ggml_cgraph * clip_graph_qwen3tts_spkenc::build() {
|
||||
// inp_raw: [T, n_mel, 1, 1], from mtmd_audio_preprocessor_qwen3tts_spk
|
||||
ggml_tensor * inp = build_inp_raw(1);
|
||||
inp = ggml_reshape_2d(ctx0, inp, inp->ne[0], inp->ne[1]);
|
||||
|
||||
// this file's convention is [C, T]; the preprocessor delivers [T, C]
|
||||
ggml_tensor * mel = ggml_cont(ctx0, ggml_transpose(ctx0, inp)); // [n_mel, T]
|
||||
cb(mel, "mel", -1);
|
||||
|
||||
// frontend conv0 TDNN k=5, dilation=1: 128 -> 512
|
||||
ggml_tensor * cur = conv1d_same(mel, model.conv1d_1_w, model.conv1d_1_b, 1);
|
||||
cur = ggml_relu(ctx0, cur);
|
||||
cb(cur, "frontend", -1);
|
||||
|
||||
// 3 SE-Res2Net blocks at dilations 2, 3, 4
|
||||
GGML_ASSERT((int) model.layers.size() == 3);
|
||||
std::vector<ggml_tensor *> blk_out(3);
|
||||
for (int il = 0; il < 3; il++) {
|
||||
cur = se_res2net_block(cur, model.layers[il], SPK_DILATIONS[il], SPK_RES2NET_SCALE);
|
||||
blk_out[il] = cur;
|
||||
cb(cur, "block_out", il);
|
||||
}
|
||||
|
||||
// multi-layer feature aggregation: cat blk[0..2] then TDNN k=1 + ReLU
|
||||
ggml_tensor * cat = ggml_concat(ctx0, blk_out[0], blk_out[1], 0);
|
||||
cat = ggml_concat(ctx0, cat, blk_out[2], 0); // [1536, T]
|
||||
ggml_tensor * mfa = conv1d_same(cat, model.conv_out_w, model.conv_out_b, 1);
|
||||
mfa = ggml_relu(ctx0, mfa);
|
||||
cb(mfa, "mfa", -1);
|
||||
|
||||
// attentive statistics pooling: [1536, T] -> [3072, 1]
|
||||
ggml_tensor * stats = attentive_stats_pool(mfa);
|
||||
cb(stats, "asp", -1);
|
||||
|
||||
// final FC k=1: [3072, 1] -> [enc_dim, 1]
|
||||
ggml_tensor * emb = conv1d_same(stats, model.mm_fc_w, model.mm_fc_b, 1);
|
||||
|
||||
emb = ggml_reshape_1d(ctx0, emb, emb->ne[0]);
|
||||
emb = ggml_cont(ctx0, emb);
|
||||
cb(emb, "spk_embedding", -1);
|
||||
|
||||
ggml_build_forward_expand(gf, emb);
|
||||
return gf;
|
||||
}
|
||||
Reference in New Issue
Block a user