mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-08-27 15:41:19 +02:00
Xuan-Son Nguyen
6e62ba5384
* adapt the api * text model ok * working impl, need verify and clean up * mtmd: build the pocket-tts transposed convolutions as GEMM + col2im ggml_conv_transpose_1d has no grouped mode, so the depthwise upsample was built as one convolution and one concat per channel, which floods the graph with small nodes and makes kernel launches dominate the decoder. Fold both cases into the column form the seanet decoder already needs: the general case reshapes the kernel to [IC, K * OC] and matmuls it with the input, the depthwise case batches a matmul over the channels so a step scales its own kernel. A single col2im_1d then scatter-adds the columns back to the signal, with the same shape as before, so the overlap-add tail, the streaming state and the bias are untouched. Generation time per frame drops by 80% on CUDA and by 50% on CPU. The output matches the previous implementation sample for sample, with a correlation of 0.999994 and identical frame counts. * flow_temp + frames_after_eos * chunking * mtmd: carry the remaining pocket-tts per-pack settings The language packs also tune the end-of-speech padding and the padding of short prompts, next to the temperature already carried in the mmproj: french_24l asks for 8 tail frames instead of the guessed 3, english_2026-01 asks for short prompts to be padded with spaces. Write both in the mmproj as clip.gen.audio.frames_after_eos and clip.gen.audio.pad_short_text, keyed on the pack in the conversion script like the temperature. The loader keeps them optional, so a mmproj without them behaves as before. Map semicolons to commas for every pack instead, the reference only asks for it on three of them and it costs nothing elsewhere. Existing mmproj files must be converted again to carry the two keys. On a long french text the port now lands within 2% of the reference: 22.96s against 23.44s, with the same peak level and the same amount of silence. * clip.gen.audio.model_variant * clean up code comments * nit: drop the dead flow_temp hparam, the pack table holds the default * update docs * address security problems * less invasive base.py * lint * add mtmd_gen_inp_default * add docs * rm gen_flow_temp --------- Co-authored-by: Pascal <admin@serveurperso.com>
771 lines
36 KiB
C++
771 lines
36 KiB
C++
#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;
|
|
|
|
if (c2w.pre_conv_w == nullptr) {
|
|
return slots; // not a code2wav model, it keeps no state between calls
|
|
}
|
|
|
|
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;
|
|
}
|