mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-09-20 01:31:31 +02:00
wip, correct code2wav
This commit is contained in:
@@ -145,6 +145,7 @@ struct clip_hparams {
|
||||
int32_t wav_upsample_n_block = 0;
|
||||
int32_t wav_dac_n_block = 0;
|
||||
int32_t wav_dac_n_res = 0;
|
||||
int32_t wav_c2w_window_frames = 0; // number of frames code2wav decodes per call
|
||||
|
||||
// mimo-v2.5: LLM-side connector (input_local_transformer)
|
||||
int32_t audio_local_n_layer = 0;
|
||||
|
||||
+19
-5
@@ -1723,6 +1723,8 @@ struct clip_model_loader {
|
||||
hparams.wav_upsample_n_block = 2;
|
||||
hparams.wav_dac_n_block = 4;
|
||||
hparams.wav_dac_n_res = 3;
|
||||
// code2wav keeps no state between calls, so it decodes this many frames each time to get enough left context
|
||||
hparams.wav_c2w_window_frames = 24;
|
||||
} break;
|
||||
case PROJECTOR_TYPE_PADDLEOCR:
|
||||
{
|
||||
@@ -4165,8 +4167,7 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
|
||||
set_input_f32("inp_raw", inp_raw);
|
||||
|
||||
} else if (!(ctx->proj_type() == PROJECTOR_TYPE_QWEN3TTS_GEN && params->gen_process == CLIP_GEN_PROCESS_CODE2WAV)) {
|
||||
// audio input (code2wav has no hidden-state/raw input at all, its
|
||||
// only input is the "inp_codes" tensor handled in the switch below)
|
||||
// audio input (code2wav has no hidden-state/raw input at all, its only input is the "inp_codes" tensor handled in the switch below)
|
||||
GGML_ASSERT(imgs.entries.size() == 1);
|
||||
|
||||
const auto & mel_inp = imgs.entries[0];
|
||||
@@ -4726,7 +4727,21 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
|
||||
{
|
||||
if (params->gen_process == CLIP_GEN_PROCESS_CODE2WAV) {
|
||||
GGML_ASSERT(params->codes != nullptr);
|
||||
std::vector<int32_t> codes = *params->codes;
|
||||
|
||||
// the caller sends codes frame-major (frame 0's codes, then frame 1's, ...). the graph wants them group-major (all frames of codebook 0, then all frames of codebook 1, ...). pad the front with code 0 if fewer frames than the window.
|
||||
const int64_t n_codes = model.gen_code_head_w->ne[2] + 1;
|
||||
const int64_t n_frames_w = hparams.wav_c2w_window_frames;
|
||||
const int64_t n_frames = (int64_t) params->codes->size() / n_codes;
|
||||
const int64_t n_use = std::min(n_frames, n_frames_w);
|
||||
const int64_t dst0 = n_frames_w - n_use; // front padding
|
||||
const int64_t src0 = n_frames - n_use; // newest frames
|
||||
|
||||
std::vector<int32_t> codes(n_frames_w * n_codes, 0);
|
||||
for (int64_t f = 0; f < n_use; f++) {
|
||||
for (int64_t g = 0; g < n_codes; g++) {
|
||||
codes[g * n_frames_w + dst0 + f] = (*params->codes)[(src0 + f) * n_codes + g];
|
||||
}
|
||||
}
|
||||
set_input_i32("inp_codes", codes);
|
||||
} else {
|
||||
std::vector<int32_t> code0 = { params->code0 };
|
||||
@@ -5158,8 +5173,7 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// the last node is the embedding tensor (not produced by the code2wav
|
||||
// sub-graph, which has no out_embd at all)
|
||||
// the last node is the embedding tensor (not produced by the code2wav sub-graph, which has no out_embd at all)
|
||||
ggml_tensor * embeddings = params->out_embd ? ggml_graph_node(gf, -1) : nullptr;
|
||||
|
||||
if (embeddings != nullptr) {
|
||||
|
||||
@@ -286,7 +286,7 @@ struct clip_graph_qwen3tts_gen : clip_graph {
|
||||
|
||||
//
|
||||
// code2wav: RVQ codes -> raw PCM (quantizer + pre_conv + pre_transformer + upsample + DAC).
|
||||
// Single-frame only for now: no cross-call state, RoPE position is always 0.
|
||||
// Each call decodes a window of T frames from scratch, RoPE positions 0..T-1. No state is kept between calls. A window of about 24 frames gives enough left context without a persistent KV cache.
|
||||
//
|
||||
struct code2wav : clip_graph {
|
||||
code2wav(const clip_graph & parent) : clip_graph(parent) {}
|
||||
@@ -298,12 +298,12 @@ struct clip_graph_qwen3tts_gen : clip_graph {
|
||||
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, ggml_tensor * pos0, ggml_tensor * mask) const;
|
||||
ggml_tensor * tfm_layer_forward(ggml_tensor * cur, const clip_layer & layer, ggml_tensor * pos, ggml_tensor * mask) const;
|
||||
ggml_tensor * convnext_block(ggml_tensor * x, const clip_code2wav::upsample_block & blk) const;
|
||||
ggml_tensor * dac_res_unit(ggml_tensor * x, const clip_code2wav::dac_res & res, int dilation) const;
|
||||
|
||||
// inp_codes: [1, n_codes] I32, one frame's RVQ codes.
|
||||
// returns audio samples, [n_samples] F32, clamped to [-1, 1].
|
||||
// inp_codes: [T, n_codes] I32, T frames of RVQ codes (group-major: all T frames of codebook 0, then all T frames of codebook 1, etc.).
|
||||
// returns audio samples for all T frames, [n_samples] F32, clamped to [-1, 1].
|
||||
ggml_tensor * decode(ggml_tensor * inp_codes) const;
|
||||
};
|
||||
};
|
||||
|
||||
@@ -354,41 +354,41 @@ ggml_tensor * clip_graph_qwen3tts_gen::code2wav::snake(ggml_tensor * x, ggml_ten
|
||||
return ggml_add(ctx0, x, s);
|
||||
}
|
||||
|
||||
// RVQ codebook decode: 16 codes -> 512-dim hidden (C-first, [512, 1]).
|
||||
// codebook 0 (semantic) and 1..15 (acoustic) are summed within their own
|
||||
// group, projected out_proj'd separately, then the two projections added.
|
||||
// RVQ codebook decode: T frames of 16 codes -> 512-dim hidden (C-first, [512, T]). codebook 0 (semantic) and 1..15 (acoustic) are summed within their own group, projected out_proj'd separately, then the two projections added.
|
||||
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];
|
||||
|
||||
ggml_tensor * code0 = ggml_view_1d(ctx0, inp_codes, 1, 0);
|
||||
ggml_tensor * sem = ggml_get_rows(ctx0, c2w.quant_first_cb_w, code0); // [256, 1]
|
||||
ggml_tensor * sem_out = ggml_mul_mat(ctx0, c2w.quant_first_out_w, sem); // [512, 1]
|
||||
// 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 * codeg = ggml_view_1d(ctx0, inp_codes, 1, (size_t) g * inp_codes->nb[1]);
|
||||
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, codeg); // [256, 1]
|
||||
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, 1]
|
||||
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. Single position only: pos0/mask are shared
|
||||
// constants across all layers/calls (self-attention on one token always
|
||||
// has softmax weight 1, but the ops are still built out in full so this
|
||||
// slots into a real KV cache later without restructuring).
|
||||
ggml_tensor * clip_graph_qwen3tts_gen::code2wav::tfm_layer_forward(ggml_tensor * cur, const clip_layer & layer, ggml_tensor * pos0, ggml_tensor * mask) const {
|
||||
// one pre_transformer layer, causal over the T-frame window (RoPE positions 0..T-1)
|
||||
ggml_tensor * clip_graph_qwen3tts_gen::code2wav::tfm_layer_forward(ggml_tensor * cur, const clip_layer & layer, ggml_tensor * pos, ggml_tensor * mask) 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 T = cur->ne[1];
|
||||
|
||||
ggml_tensor * residual = cur;
|
||||
ggml_tensor * h = ggml_rms_norm(ctx0, cur, hparams.wav_tfm_eps);
|
||||
@@ -398,17 +398,17 @@ ggml_tensor * clip_graph_qwen3tts_gen::code2wav::tfm_layer_forward(ggml_tensor *
|
||||
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_reshape_3d(ctx0, q, d_head, n_head, T);
|
||||
k = ggml_reshape_3d(ctx0, k, d_head, n_head_kv, T);
|
||||
|
||||
q = ggml_rope_ext(ctx0, q, pos0, nullptr, (int) d_head, GGML_ROPE_TYPE_NEOX, 0,
|
||||
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, pos0, nullptr, (int) d_head, GGML_ROPE_TYPE_NEOX, 0,
|
||||
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);
|
||||
|
||||
ggml_tensor * q_cur = ggml_reshape_4d(ctx0, q, d_head, n_head, 1, 1);
|
||||
ggml_tensor * k_cur = ggml_reshape_4d(ctx0, k, d_head, n_head_kv, 1, 1);
|
||||
ggml_tensor * v_cur = ggml_reshape_4d(ctx0, v, d_head, n_head_kv, 1, 1);
|
||||
ggml_tensor * q_cur = ggml_reshape_4d(ctx0, q, d_head, n_head, T, 1);
|
||||
ggml_tensor * k_cur = ggml_reshape_4d(ctx0, k, d_head, n_head_kv, T, 1);
|
||||
ggml_tensor * v_cur = ggml_reshape_4d(ctx0, v, d_head, n_head_kv, T, 1);
|
||||
|
||||
ggml_tensor * attn_out = build_attn(layer.o_w, layer.o_b, q_cur, k_cur, v_cur, mask, kq_scale, 0);
|
||||
if (layer.ls_1_w) {
|
||||
@@ -465,42 +465,40 @@ ggml_tensor * clip_graph_qwen3tts_gen::code2wav::dac_res_unit(ggml_tensor * x, c
|
||||
return ggml_add(ctx0, residual, h);
|
||||
}
|
||||
|
||||
// RVQ codes -> raw PCM. Single frame only: no cross-call state.
|
||||
// RVQ codes -> raw PCM for the whole T-frame window
|
||||
ggml_tensor * clip_graph_qwen3tts_gen::code2wav::decode(ggml_tensor * inp_codes) const {
|
||||
const auto & c2w = model.c2w;
|
||||
const int64_t T = inp_codes->ne[0];
|
||||
|
||||
// 1. quantizer decode: 16 codes -> [512, 1] (C-first)
|
||||
// 1. quantizer decode: T frames of 16 codes -> [512, T] (C-first)
|
||||
ggml_tensor * hidden = quant_decode(inp_codes);
|
||||
|
||||
// 2. pre_conv: [512, 1] -> T-first [1, 512] -> causal conv k=3 -> [1, 1024]
|
||||
ggml_tensor * x = ggml_cont(ctx0, ggml_transpose(ctx0, hidden)); // [1, 512]
|
||||
x = causal_conv1d(x, c2w.pre_conv_w, c2w.pre_conv_b, 1); // [1, 1024]
|
||||
// 2. pre_conv: [512, T] -> T-first [T, 512] -> causal conv k=3 -> [T, 1024]
|
||||
ggml_tensor * x = ggml_cont(ctx0, ggml_transpose(ctx0, hidden)); // [T, 512]
|
||||
x = causal_conv1d(x, c2w.pre_conv_w, c2w.pre_conv_b, 1); // [T, 1024]
|
||||
cb(x, "wav_pre_conv_out", -1);
|
||||
|
||||
// 3. pre_transformer: back to C-first [1024, 1], project down to hidden_size,
|
||||
// run 8 layers (single position, pos=0, self-attention only), project back up
|
||||
ggml_tensor * cur = ggml_cont(ctx0, ggml_transpose(ctx0, x)); // [1024, 1]
|
||||
// 3. pre_transformer: back to C-first [1024, T], project down to hidden_size, run 8 layers causal over the T-frame window, project back up
|
||||
ggml_tensor * cur = ggml_cont(ctx0, ggml_transpose(ctx0, x)); // [1024, T]
|
||||
cur = ggml_mul_mat(ctx0, c2w.tfm_in_proj_w, cur);
|
||||
cur = ggml_add(ctx0, cur, c2w.tfm_in_proj_b); // [512 (tfm hidden), 1]
|
||||
cur = ggml_add(ctx0, cur, c2w.tfm_in_proj_b); // [512 (tfm hidden), T]
|
||||
|
||||
// graph-constant 0 (position / mask), built without a host upload -- same
|
||||
// view+scale_bias+cast trick used elsewhere, anchored on a guaranteed-F32 tensor
|
||||
ggml_tensor * zero_anchor = ggml_view_1d(ctx0, c2w.tfm_layers[0].ln_1_w, 1, 0);
|
||||
zero_anchor = ggml_scale_bias(ctx0, zero_anchor, 0.0f, 0.0f);
|
||||
ggml_tensor * pos0 = ggml_cast(ctx0, zero_anchor, GGML_TYPE_I32);
|
||||
ggml_tensor * mask = ggml_reshape_4d(ctx0, zero_anchor, 1, 1, 1, 1);
|
||||
ggml_tensor * pos = ggml_cast(ctx0, ggml_arange(ctx0, 0.0f, (float) T, 1.0f), GGML_TYPE_I32);
|
||||
ggml_tensor * tri = ggml_tri(ctx0, ggml_fill(ctx0, ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, T, T), 1.0f),
|
||||
GGML_TRI_TYPE_LOWER_DIAG);
|
||||
ggml_tensor * mask = ggml_reshape_4d(ctx0, ggml_log(ctx0, tri), T, T, 1, 1); // 0 = keep, -inf = masked
|
||||
|
||||
for (int il = 0; il < hparams.wav_tfm_n_layer; il++) {
|
||||
cur = tfm_layer_forward(cur, c2w.tfm_layers[il], pos0, mask);
|
||||
cur = tfm_layer_forward(cur, c2w.tfm_layers[il], pos, mask);
|
||||
}
|
||||
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, 1]
|
||||
cur = ggml_add(ctx0, cur, c2w.tfm_out_proj_b); // [1024, T]
|
||||
cb(cur, "wav_tfm_out", -1);
|
||||
|
||||
// 4. upsample: 2x (causal ConvTranspose1d, stride 2 + ConvNeXt block), back to T-first
|
||||
x = ggml_cont(ctx0, ggml_transpose(ctx0, cur)); // [1, 1024]
|
||||
x = ggml_cont(ctx0, ggml_transpose(ctx0, cur)); // [T, 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);
|
||||
@@ -534,17 +532,16 @@ ggml_tensor * clip_graph_qwen3tts_gen::code2wav::decode(ggml_tensor * inp_codes)
|
||||
return x;
|
||||
}
|
||||
|
||||
// master build(): switches on gen_process to construct either the code_gen
|
||||
// sub-graph (backbone hidden state -> 16 RVQ codes + next-step embd) or the
|
||||
// code2wav sub-graph (16 RVQ codes -> raw PCM), both hosted in this one clip_ctx.
|
||||
// master build(): switches on gen_process to construct either the code_gen sub-graph (backbone hidden state -> 16 RVQ codes + next-step embd) or the code2wav sub-graph (16 RVQ codes -> raw PCM), both hosted in this one clip_ctx.
|
||||
ggml_cgraph * clip_graph_qwen3tts_gen::build() {
|
||||
GGML_ASSERT(n_batch == 1); // this module only ever processes one frame at a time
|
||||
|
||||
if (gen_process == CLIP_GEN_PROCESS_CODE2WAV) {
|
||||
const int64_t n_acoustic = model.gen_code_head_w->ne[2]; // 15
|
||||
const int n_codes = (int) n_acoustic + 1; // 16
|
||||
const int n_frames = hparams.wav_c2w_window_frames;
|
||||
|
||||
ggml_tensor * inp_codes = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, 1, n_codes);
|
||||
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);
|
||||
|
||||
|
||||
Reference in New Issue
Block a user