mtmd: support pocket-tts (#26871)
* 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>
This commit is contained in:
Xuan-Son Nguyen
co-authored by
Pascal
Pascal
parent
8d274dd7c6
commit
6e62ba5384
@@ -318,6 +318,59 @@ struct clip_graph_qwen3tts_gen : clip_graph {
|
||||
};
|
||||
};
|
||||
|
||||
//
|
||||
// pocket-tts: SEANet convolution stack, shared by the voice encoder and the mimi decoder.
|
||||
// stateless unless state_in is populated: convs then pad instead of carrying left-context.
|
||||
//
|
||||
struct clip_graph_pockettts_seanet : clip_graph {
|
||||
clip_graph_pockettts_seanet(const clip_graph & parent) : clip_graph(parent) {}
|
||||
ggml_cgraph * build() override { GGML_ABORT("call encode()/decode() instead"); }
|
||||
|
||||
// per-call streaming state, keyed by slot name (see list_pockettts_state_slots)
|
||||
std::map<std::string, ggml_tensor *> state_in;
|
||||
mutable std::vector<std::pair<std::string, ggml_tensor *>> state_out;
|
||||
|
||||
ggml_tensor * conv1d(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int stride, int dilation,
|
||||
bool pad_replicate = false, const std::string & state_name = "") const;
|
||||
ggml_tensor * conv_transpose1d(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int stride,
|
||||
const std::string & state_name = "") const;
|
||||
ggml_tensor * res_unit(ggml_tensor * x, const clip_seanet::stage & stage, int dilation,
|
||||
const std::string & state_prefix = "") const;
|
||||
|
||||
// x: [T, C] -> [T / hop, dim]
|
||||
ggml_tensor * encode(ggml_tensor * x) const;
|
||||
// x: [T, dim] -> [T * hop, 1], streams when state_in is populated
|
||||
ggml_tensor * decode(ggml_tensor * x) const;
|
||||
};
|
||||
|
||||
// mimi encoder + speaker_proj: reference waveform -> voice conditioning rows
|
||||
struct clip_graph_pockettts_spkenc : clip_graph {
|
||||
clip_graph_pockettts_spkenc(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {}
|
||||
ggml_cgraph * build() override;
|
||||
|
||||
ggml_tensor * tfm_layer_forward(ggml_tensor * cur, const clip_layer & layer, ggml_tensor * inp_pos, ggml_tensor * kq_mask, int il) const;
|
||||
};
|
||||
|
||||
//
|
||||
// pocket-tts generation:
|
||||
// GEN_CODE = flow-matching decoder + end-of-speech head, one latent per call
|
||||
// GEN_WAV = mimi decoder, a window of latents -> PCM
|
||||
//
|
||||
struct clip_graph_pockettts_gen : clip_graph {
|
||||
clip_graph_pockettts_gen(clip_ctx * ctx, const clip_image_f32 & img, clip_gen_process_type gen_process, int n_step, int n_frames)
|
||||
: clip_graph(ctx, img), gen_process(gen_process), n_step(n_step), n_frames(n_frames) {}
|
||||
ggml_cgraph * build() override;
|
||||
|
||||
clip_gen_process_type gen_process;
|
||||
int n_step; // lsd_decode steps, fixed at graph-build time
|
||||
int n_frames; // GEN_WAV only: number of latents to decode
|
||||
|
||||
// AdaLN modulation: x * (1 + scale) + shift
|
||||
ggml_tensor * modulate(ggml_tensor * x, ggml_tensor * shift, ggml_tensor * scale) const;
|
||||
ggml_tensor * time_embed(const clip_flow_net::time_embd & te, float t) const;
|
||||
ggml_tensor * flow_forward(ggml_tensor * cond, ggml_tensor * x, float s, float t) const;
|
||||
};
|
||||
|
||||
// one persisted state buffer used by code2wav, see qwen3tts-gen.cpp
|
||||
struct c2w_state_slot {
|
||||
std::string name;
|
||||
@@ -326,6 +379,9 @@ struct c2w_state_slot {
|
||||
};
|
||||
std::vector<c2w_state_slot> list_c2w_state_slots(const clip_hparams & hparams, const clip_model & model);
|
||||
|
||||
// same, for the streaming mimi decoder (pocket-tts GEN_WAV)
|
||||
std::vector<c2w_state_slot> list_pockettts_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,291 @@
|
||||
#include "models.h"
|
||||
|
||||
#include <cmath>
|
||||
|
||||
// pocket-tts generation stages
|
||||
//
|
||||
// GEN_CODE: backbone hidden state -> next 32-d latent (flow matching) + end-of-speech score
|
||||
// GEN_WAV : a window of latents -> PCM, through the mimi decoder
|
||||
//
|
||||
// there is no codebook anywhere, "codes" in the mtmd API are continuous features here
|
||||
|
||||
ggml_tensor * clip_graph_pockettts_gen::modulate(ggml_tensor * x, ggml_tensor * shift, ggml_tensor * scale) const {
|
||||
ggml_tensor * cur = ggml_mul(ctx0, x, ggml_scale_bias(ctx0, scale, 1.0f, 1.0f));
|
||||
return ggml_add(ctx0, cur, shift);
|
||||
}
|
||||
|
||||
// see TimestepEmbedder in the reference
|
||||
ggml_tensor * clip_graph_pockettts_gen::time_embed(const clip_flow_net::time_embd & te, float t) const {
|
||||
// t is a graph-build constant, so the cos/sin table can be folded into a scaled copy
|
||||
ggml_tensor * args = ggml_scale(ctx0, te.freqs, t);
|
||||
ggml_tensor * emb = ggml_concat(ctx0, ggml_cos(ctx0, args), ggml_sin(ctx0, args), 0);
|
||||
|
||||
ggml_tensor * cur = build_mm(te.up_w, emb);
|
||||
cur = ggml_add(ctx0, cur, te.up_b);
|
||||
cur = ggml_silu(ctx0, cur);
|
||||
cur = build_mm(te.down_w, cur);
|
||||
cur = ggml_add(ctx0, cur, te.down_b);
|
||||
|
||||
// this "RMSNorm" divides by the unbiased variance, not the mean square
|
||||
// it also rescales the input, not the centered value, see _rms_norm() in mlp.py
|
||||
{
|
||||
const int64_t n = cur->ne[0];
|
||||
ggml_tensor * mean = ggml_mean(ctx0, cur);
|
||||
ggml_tensor * dev = ggml_sub(ctx0, cur, mean);
|
||||
ggml_tensor * var = ggml_mean(ctx0, ggml_sqr(ctx0, dev));
|
||||
var = ggml_scale_bias(ctx0, var, (float) n / (float) (n - 1), 1e-5f);
|
||||
cur = ggml_div(ctx0, cur, ggml_sqrt(ctx0, var));
|
||||
cur = ggml_mul(ctx0, cur, te.norm);
|
||||
}
|
||||
|
||||
return cur;
|
||||
}
|
||||
|
||||
// one velocity evaluation: v(cond, s, t, x)
|
||||
ggml_tensor * clip_graph_pockettts_gen::flow_forward(ggml_tensor * cond, ggml_tensor * x, float s, float t) const {
|
||||
const auto & flow = model.flow;
|
||||
|
||||
ggml_tensor * cur = build_mm(flow.input_proj_w, x);
|
||||
cur = ggml_add(ctx0, cur, flow.input_proj_b);
|
||||
|
||||
// the two time conditions are averaged, then added to the projected backbone state
|
||||
ggml_tensor * ts = ggml_add(ctx0, time_embed(flow.time[0], s), time_embed(flow.time[1], t));
|
||||
ts = ggml_scale(ctx0, ts, 1.0f / (float) flow.time.size());
|
||||
|
||||
ggml_tensor * c = build_mm(flow.cond_embd_w, cond);
|
||||
c = ggml_add(ctx0, c, flow.cond_embd_b);
|
||||
|
||||
ggml_tensor * y = ggml_add(ctx0, ts, c);
|
||||
cb(y, "flow_cond", -1);
|
||||
|
||||
const int64_t n_ch = flow.blocks.empty() ? 0 : flow.blocks[0].norm_w->ne[0];
|
||||
|
||||
for (size_t il = 0; il < flow.blocks.size(); il++) {
|
||||
const auto & blk = flow.blocks[il];
|
||||
|
||||
ggml_tensor * mod = build_mm(blk.ada_w, ggml_silu(ctx0, y));
|
||||
mod = ggml_add(ctx0, mod, blk.ada_b);
|
||||
|
||||
ggml_tensor * shift = ggml_view_1d(ctx0, mod, n_ch, 0);
|
||||
ggml_tensor * scale = ggml_view_1d(ctx0, mod, n_ch, (size_t) n_ch * mod->nb[0]);
|
||||
ggml_tensor * gate = ggml_view_1d(ctx0, mod, n_ch, (size_t) 2 * n_ch * mod->nb[0]);
|
||||
|
||||
ggml_tensor * h = build_norm(cur, blk.norm_w, blk.norm_b, NORM_TYPE_NORMAL, 1e-6f, (int) il);
|
||||
h = modulate(h, shift, scale);
|
||||
h = build_mm(blk.up_w, h);
|
||||
h = ggml_add(ctx0, h, blk.up_b);
|
||||
h = ggml_silu(ctx0, h);
|
||||
h = build_mm(blk.down_w, h);
|
||||
h = ggml_add(ctx0, h, blk.down_b);
|
||||
|
||||
cur = ggml_add(ctx0, cur, ggml_mul(ctx0, gate, h));
|
||||
cb(cur, "flow_blk", (int) il);
|
||||
}
|
||||
|
||||
// final layer: the norm has no weights, only the AdaLN modulation
|
||||
ggml_tensor * mod = build_mm(flow.final_ada_w, ggml_silu(ctx0, y));
|
||||
mod = ggml_add(ctx0, mod, flow.final_ada_b);
|
||||
|
||||
ggml_tensor * shift = ggml_view_1d(ctx0, mod, n_ch, 0);
|
||||
ggml_tensor * scale = ggml_view_1d(ctx0, mod, n_ch, (size_t) n_ch * mod->nb[0]);
|
||||
|
||||
cur = build_norm(cur, nullptr, nullptr, NORM_TYPE_NORMAL, 1e-6f, -1);
|
||||
cur = modulate(cur, shift, scale);
|
||||
cur = build_mm(flow.final_proj_w, cur);
|
||||
cur = ggml_add(ctx0, cur, flow.final_proj_b);
|
||||
|
||||
return cur;
|
||||
}
|
||||
|
||||
// state carried between GEN_WAV calls: rope offset, per-layer KV window, conv left context
|
||||
// and the transposed-conv overlap tails
|
||||
std::vector<c2w_state_slot> list_pockettts_state_slots(const clip_hparams & hparams, const clip_model & model) {
|
||||
std::vector<c2w_state_slot> slots;
|
||||
if (model.gen_upsample_w == nullptr) {
|
||||
return slots; // not a pocket-tts decoder
|
||||
}
|
||||
const auto & seanet = model.seanet;
|
||||
|
||||
// the slots below are sized from these
|
||||
GGML_ASSERT(!model.gen_tfm_layers.empty());
|
||||
GGML_ASSERT((int) seanet.stages.size() >= hparams.seanet_n_stage);
|
||||
GGML_ASSERT((int) hparams.seanet_ratios.size() >= hparams.seanet_n_stage);
|
||||
GGML_ASSERT(hparams.mimi_tfm_context > 1 && hparams.mimi_downsample > 0);
|
||||
|
||||
slots.push_back({"tfm_pos", 1, 1});
|
||||
|
||||
const int64_t n_embd_a = model.gen_tfm_layers[0].q_w->ne[1];
|
||||
const int64_t prefix = hparams.mimi_tfm_context - 1;
|
||||
for (size_t il = 0; il < model.gen_tfm_layers.size(); il++) {
|
||||
slots.push_back({"tfm_k_" + std::to_string(il), n_embd_a, prefix});
|
||||
slots.push_back({"tfm_v_" + std::to_string(il), n_embd_a, prefix});
|
||||
}
|
||||
|
||||
// upsample is depthwise, its output channel count is the input one
|
||||
slots.push_back({"up", model.gen_upsample_w->ne[0] - hparams.mimi_downsample, model.gen_upsample_w->ne[2]});
|
||||
|
||||
slots.push_back({"dec_in", seanet.conv_in_w->ne[0] - 1, seanet.conv_in_w->ne[1]});
|
||||
for (int i = 0; i < hparams.seanet_n_stage; i++) {
|
||||
const auto & stage = seanet.stages[i];
|
||||
const int stride = hparams.seanet_ratios[hparams.seanet_n_stage - 1 - i];
|
||||
slots.push_back({"dec_up_" + std::to_string(i), stage.scale_conv_w->ne[0] - stride, stage.scale_conv_w->ne[1]});
|
||||
slots.push_back({"dec_res_" + std::to_string(i), stage.res_conv1_w->ne[0] - 1, stage.res_conv1_w->ne[1]});
|
||||
}
|
||||
slots.push_back({"dec_out", seanet.conv_out_w->ne[0] - 1, seanet.conv_out_w->ne[1]});
|
||||
|
||||
return slots;
|
||||
}
|
||||
|
||||
ggml_cgraph * clip_graph_pockettts_gen::build() {
|
||||
if (gen_process == CLIP_GEN_PROCESS_GEN_CODE) {
|
||||
// the backbone hidden state arrives as the single batch entry
|
||||
ggml_tensor * h_state = build_inp_raw(1);
|
||||
h_state = ggml_reshape_2d(ctx0, h_state, n_mmproj_embd, 1);
|
||||
|
||||
// end-of-speech probe, thresholded on the host side
|
||||
ggml_tensor * eos = build_mm(model.gen_out_eos_w, h_state);
|
||||
eos = ggml_add(ctx0, eos, model.gen_out_eos_b);
|
||||
ggml_set_name(eos, "out_eos_score");
|
||||
ggml_set_output(eos);
|
||||
ggml_build_forward_expand(gf, eos);
|
||||
|
||||
const int64_t n_latent = model.gen_input_lin_w->ne[0];
|
||||
|
||||
ggml_tensor * noise = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_latent, 1);
|
||||
ggml_set_name(noise, "inp_noise");
|
||||
ggml_set_input(noise);
|
||||
|
||||
// lsd_decode: integrate the velocity field from the noise sample
|
||||
ggml_tensor * cur = noise;
|
||||
for (int i = 0; i < n_step; i++) {
|
||||
const float s = (float) i / (float) n_step;
|
||||
const float t = (float) (i + 1) / (float) n_step;
|
||||
ggml_tensor * v = flow_forward(h_state, cur, s, t);
|
||||
cur = ggml_add(ctx0, cur, ggml_scale(ctx0, v, 1.0f / (float) n_step));
|
||||
}
|
||||
cb(cur, "flow_latent", -1);
|
||||
|
||||
ggml_set_name(cur, "out_feats");
|
||||
ggml_set_output(cur);
|
||||
ggml_build_forward_expand(gf, cur);
|
||||
|
||||
// the same latent, projected into the backbone's input space for the next step
|
||||
ggml_tensor * embd = build_mm(model.gen_input_lin_w, cur);
|
||||
cb(embd, "gen_embd", -1);
|
||||
ggml_build_forward_expand(gf, embd);
|
||||
|
||||
return gf;
|
||||
}
|
||||
|
||||
// GEN_WAV: [32, n_frames] latents -> PCM
|
||||
ggml_tensor * feats = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32,
|
||||
model.gen_input_lin_w->ne[0], n_frames);
|
||||
ggml_set_name(feats, "inp_feats");
|
||||
ggml_set_input(feats);
|
||||
|
||||
// denormalize, then the DummyQuantizer up-projection
|
||||
ggml_tensor * cur = ggml_add(ctx0, ggml_mul(ctx0, feats, model.gen_emb_std), model.gen_emb_mean);
|
||||
cur = build_mm(model.gen_quant_out_w, cur);
|
||||
cb(cur, "quant_out", -1);
|
||||
|
||||
clip_graph_pockettts_seanet seanet(*this);
|
||||
for (const auto & slot : list_pockettts_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);
|
||||
seanet.state_in[slot.name] = t;
|
||||
}
|
||||
|
||||
// model frame rate -> encoder frame rate, depthwise transposed conv
|
||||
cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur));
|
||||
cur = seanet.conv_transpose1d(cur, model.gen_upsample_w, nullptr, hparams.mimi_downsample, "up");
|
||||
cb(cur, "mimi_upsample", -1);
|
||||
|
||||
cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur));
|
||||
|
||||
// positions continue across calls, the counter lives in the state
|
||||
const int64_t n_pos = cur->ne[1];
|
||||
const int64_t prefix = hparams.mimi_tfm_context - 1;
|
||||
const int64_t n_kv = prefix + n_pos;
|
||||
|
||||
ggml_tensor * base = ggml_reshape_1d(ctx0, seanet.state_in.at("tfm_pos"), 1);
|
||||
ggml_tensor * inp_pos = ggml_cast(ctx0, ggml_add(ctx0, ggml_arange(ctx0, 0.0f, (float) n_pos, 1.0f), base),
|
||||
GGML_TYPE_I32);
|
||||
seanet.state_out.push_back({"tfm_pos", ggml_scale_bias(ctx0, seanet.state_in.at("tfm_pos"), 1.0f, (float) n_pos)});
|
||||
|
||||
// banded causal mask over [cached prefix | this chunk]
|
||||
// the last factor masks out cache rows that hold no real frame yet
|
||||
ggml_tensor * pos_k = ggml_reshape_2d(ctx0, ggml_arange(ctx0, 0.0f, (float) n_kv, 1.0f), n_kv, 1);
|
||||
ggml_tensor * pos_q = ggml_reshape_2d(ctx0, ggml_arange(ctx0, (float) prefix, (float) (prefix + n_pos), 1.0f), 1, n_pos);
|
||||
ggml_tensor * diff = ggml_sub(ctx0, ggml_repeat_4d(ctx0, pos_q, n_kv, n_pos, 1, 1), pos_k);
|
||||
|
||||
ggml_tensor * keep = ggml_mul(ctx0,
|
||||
ggml_step(ctx0, ggml_scale_bias(ctx0, diff, 1.0f, 0.5f)), // delta >= 0
|
||||
ggml_step(ctx0, ggml_scale_bias(ctx0, diff, -1.0f, (float) hparams.mimi_tfm_context - 0.5f))); // delta < context
|
||||
keep = ggml_mul(ctx0, keep,
|
||||
ggml_step(ctx0, ggml_scale_bias(ctx0, ggml_add(ctx0, pos_k, base), 1.0f, 0.5f - (float) prefix)));
|
||||
ggml_tensor * kq_mask = ggml_reshape_4d(ctx0, ggml_log(ctx0, keep), n_kv, n_pos, 1, 1);
|
||||
|
||||
for (int il = 0; il < n_layer; il++) {
|
||||
const auto & layer = model.gen_tfm_layers[il];
|
||||
ggml_tensor * inp = cur;
|
||||
|
||||
cur = build_norm(cur, layer.ln_1_w, layer.ln_1_b, NORM_TYPE_NORMAL, eps, il);
|
||||
|
||||
ggml_tensor * Qcur = build_mm(layer.q_w, cur);
|
||||
ggml_tensor * Kcur = build_mm(layer.k_w, cur);
|
||||
ggml_tensor * Vcur = build_mm(layer.v_w, cur);
|
||||
|
||||
Qcur = ggml_reshape_3d(ctx0, Qcur, d_head, n_head, n_pos);
|
||||
Kcur = ggml_reshape_3d(ctx0, Kcur, d_head, n_head, n_pos);
|
||||
|
||||
Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, d_head, GGML_ROPE_TYPE_NORMAL, 0,
|
||||
hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
|
||||
Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, d_head, GGML_ROPE_TYPE_NORMAL, 0,
|
||||
hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
|
||||
|
||||
// prepend the cached window, then keep this chunk's tail for the next call
|
||||
const std::string k_name = "tfm_k_" + std::to_string(il);
|
||||
const std::string v_name = "tfm_v_" + std::to_string(il);
|
||||
ggml_tensor * k_full = ggml_concat(ctx0, seanet.state_in.at(k_name),
|
||||
ggml_reshape_2d(ctx0, Kcur, d_head * n_head, n_pos), 1);
|
||||
ggml_tensor * v_full = ggml_concat(ctx0, seanet.state_in.at(v_name), Vcur, 1);
|
||||
seanet.state_out.push_back({k_name, ggml_cont(ctx0, ggml_view_2d(ctx0, k_full, k_full->ne[0], prefix,
|
||||
k_full->nb[1], (size_t) n_pos * k_full->nb[1]))});
|
||||
seanet.state_out.push_back({v_name, ggml_cont(ctx0, ggml_view_2d(ctx0, v_full, v_full->ne[0], prefix,
|
||||
v_full->nb[1], (size_t) n_pos * v_full->nb[1]))});
|
||||
|
||||
ggml_tensor * q_cur = ggml_reshape_4d(ctx0, Qcur, d_head, n_head, n_pos, 1);
|
||||
ggml_tensor * k_cur = ggml_reshape_4d(ctx0, k_full, d_head, n_head, n_kv, 1);
|
||||
ggml_tensor * v_cur = ggml_reshape_4d(ctx0, v_full, d_head, n_head, n_kv, 1);
|
||||
|
||||
cur = build_attn(layer.o_w, nullptr, q_cur, k_cur, v_cur, kq_mask, kq_scale, il);
|
||||
cur = ggml_mul(ctx0, cur, layer.ls_1_w);
|
||||
cur = ggml_add(ctx0, cur, inp);
|
||||
|
||||
inp = cur;
|
||||
cur = build_norm(cur, layer.ln_2_w, layer.ln_2_b, NORM_TYPE_NORMAL, eps, il);
|
||||
cur = build_ffn(cur, layer.ff_up_w, nullptr, nullptr, nullptr, layer.ff_down_w, nullptr, FFN_GELU, il);
|
||||
cur = ggml_mul(ctx0, cur, layer.ls_2_w);
|
||||
cur = ggml_add(ctx0, cur, inp);
|
||||
}
|
||||
cb(cur, "mimi_dec_tfm", -1);
|
||||
|
||||
cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur));
|
||||
cur = seanet.decode(cur);
|
||||
|
||||
for (const auto & s : seanet.state_out) {
|
||||
ggml_set_name(s.second, ("state_out_" + s.first).c_str());
|
||||
ggml_set_output(s.second);
|
||||
ggml_build_forward_expand(gf, s.second);
|
||||
}
|
||||
|
||||
// [n_samples, 1] -> [n_samples], clamped like the reference output
|
||||
cur = ggml_reshape_1d(ctx0, cur, cur->ne[0]);
|
||||
cur = ggml_clamp(ctx0, cur, -1.0f, 1.0f);
|
||||
ggml_set_name(cur, "out_audio");
|
||||
ggml_set_output(cur);
|
||||
ggml_build_forward_expand(gf, cur);
|
||||
|
||||
return gf;
|
||||
}
|
||||
@@ -0,0 +1,162 @@
|
||||
#include "models.h"
|
||||
|
||||
// SEANet convolution stack of the mimi codec, see pocket_tts/modules/seanet.py
|
||||
//
|
||||
// tensors are T-first here: [T, C]
|
||||
// the convs are causal: left context comes from a state slot, or from padding on a cold start
|
||||
|
||||
static int64_t div_ceil(int64_t a, int64_t b) {
|
||||
return a / b + (a % b ? 1 : 0);
|
||||
}
|
||||
|
||||
// x: [T, IC], w: [K, IC, OC] -> [T / stride, OC]
|
||||
// the convs are causal, so the whole K - stride padding goes on the left
|
||||
ggml_tensor * clip_graph_pockettts_seanet::conv1d(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int stride, int dilation,
|
||||
bool pad_replicate, const std::string & state_name) const {
|
||||
const int64_t k_size = (w->ne[0] - 1) * dilation + 1;
|
||||
const int64_t p_total = k_size - stride;
|
||||
|
||||
// trailing padding so the last frame is not dropped, see pad_for_conv1d() in conv.py
|
||||
const int64_t n_frames = div_ceil(x->ne[0] - k_size + p_total, stride);
|
||||
const int64_t ideal_len = n_frames * stride + k_size - p_total;
|
||||
const int64_t p_extra = ideal_len - x->ne[0];
|
||||
|
||||
if (!state_name.empty() && p_total > 0) {
|
||||
// streaming: the left context is the tail of the previous call
|
||||
ggml_tensor * left = state_in.at(state_name); // [p_total, IC]
|
||||
x = ggml_concat(ctx0, left, x, 0);
|
||||
state_out.push_back({state_name,
|
||||
ggml_cont(ctx0, ggml_view_2d(ctx0, x, p_total, x->ne[1], x->nb[1],
|
||||
(size_t) (x->ne[0] - p_total) * x->nb[0]))});
|
||||
} else if (pad_replicate && p_total > 0) {
|
||||
// the resamplers repeat the first frame instead of zero-padding
|
||||
ggml_tensor * first = ggml_view_2d(ctx0, x, 1, x->ne[1], x->nb[1], 0);
|
||||
ggml_tensor * left = ggml_repeat_4d(ctx0, first, p_total, x->ne[1], 1, 1);
|
||||
x = ggml_concat(ctx0, left, x, 0);
|
||||
x = ggml_pad_ext(ctx0, x, 0, p_extra, 0, 0, 0, 0, 0, 0);
|
||||
} else {
|
||||
x = ggml_pad_ext(ctx0, x, p_total, p_extra, 0, 0, 0, 0, 0, 0);
|
||||
}
|
||||
|
||||
ggml_tensor * y = ggml_conv_1d(ctx0, w, x, stride, 0, dilation);
|
||||
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]));
|
||||
}
|
||||
return y;
|
||||
}
|
||||
|
||||
// x: [T, IC], w: [K, OC/groups, IC] -> [T * stride, OC]
|
||||
// the K - stride overlap tail belongs to the next call: added to its head when streaming, else dropped
|
||||
ggml_tensor * clip_graph_pockettts_seanet::conv_transpose1d(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int stride,
|
||||
const std::string & state_name) const {
|
||||
const int64_t K = w->ne[0];
|
||||
const int64_t T = x->ne[0];
|
||||
const int64_t p_total = K - stride;
|
||||
const bool depthwise = w->ne[1] == 1 && w->ne[2] > 1;
|
||||
const int64_t OC = depthwise ? w->ne[2] : w->ne[1];
|
||||
const int64_t emit_len = T * stride;
|
||||
|
||||
// one column per input step, holding the [K, OC] window that col2im scatter-adds at t * stride
|
||||
ggml_tensor * col;
|
||||
if (depthwise) {
|
||||
// one group per channel: a batched matmul over the channels scales the kernel by each step
|
||||
ggml_tensor * krn = ggml_reshape_3d(ctx0, w, 1, K, OC); // [1, K, OC]
|
||||
ggml_tensor * xs = ggml_reshape_3d(ctx0, x, 1, T, OC); // [1, T, OC]
|
||||
col = ggml_mul_mat(ctx0, krn, xs); // [K, T, OC]
|
||||
col = ggml_cont(ctx0, ggml_permute(ctx0, col, 0, 2, 1, 3)); // [K, OC, T]
|
||||
col = ggml_reshape_2d(ctx0, col, K * OC, T);
|
||||
} else {
|
||||
ggml_tensor * w2 = ggml_reshape_2d(ctx0, w, K * OC, w->ne[2]);
|
||||
w2 = ggml_cont(ctx0, ggml_transpose(ctx0, w2)); // [IC, K * OC]
|
||||
ggml_tensor * xt = ggml_cont(ctx0, ggml_transpose(ctx0, x)); // [IC, T]
|
||||
col = ggml_mul_mat(ctx0, w2, xt);
|
||||
}
|
||||
ggml_tensor * full = ggml_col2im_1d(ctx0, col, stride, OC, 0); // [emit_len + p_total, OC]
|
||||
|
||||
ggml_tensor * out;
|
||||
if (state_name.empty() || p_total == 0) {
|
||||
out = ggml_cont(ctx0, ggml_view_2d(ctx0, full, emit_len, full->ne[1], full->nb[1], 0));
|
||||
} else {
|
||||
// overlap-add the tail the previous call held back
|
||||
ggml_tensor * prev = state_in.at(state_name); // [p_total, OC]
|
||||
ggml_tensor * head = ggml_add(ctx0, ggml_view_2d(ctx0, full, p_total, full->ne[1], full->nb[1], 0), prev);
|
||||
if (emit_len > p_total) {
|
||||
ggml_tensor * rest = ggml_view_2d(ctx0, full, emit_len - p_total, full->ne[1], full->nb[1],
|
||||
(size_t) p_total * full->nb[0]);
|
||||
out = ggml_concat(ctx0, head, rest, 0);
|
||||
} else {
|
||||
out = head;
|
||||
}
|
||||
state_out.push_back({state_name,
|
||||
ggml_cont(ctx0, ggml_view_2d(ctx0, full, p_total, full->ne[1], full->nb[1],
|
||||
(size_t) emit_len * full->nb[0]))});
|
||||
}
|
||||
|
||||
if (b) {
|
||||
out = ggml_add(ctx0, out, ggml_reshape_2d(ctx0, b, 1, b->ne[0]));
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
ggml_tensor * clip_graph_pockettts_seanet::res_unit(ggml_tensor * x, const clip_seanet::stage & stage, int dilation,
|
||||
const std::string & state_prefix) const {
|
||||
ggml_tensor * h = ggml_elu(ctx0, x);
|
||||
h = conv1d(h, stage.res_conv1_w, stage.res_conv1_b, 1, dilation, false, state_prefix);
|
||||
h = ggml_elu(ctx0, h);
|
||||
// the second conv is pointwise, it needs no left context
|
||||
h = conv1d(h, stage.res_conv2_w, stage.res_conv2_b, 1, 1);
|
||||
return ggml_add(ctx0, x, h);
|
||||
}
|
||||
|
||||
ggml_tensor * clip_graph_pockettts_seanet::encode(ggml_tensor * x) const {
|
||||
const auto & seanet = model.seanet;
|
||||
|
||||
ggml_tensor * cur = conv1d(x, seanet.conv_in_w, seanet.conv_in_b, 1, 1);
|
||||
cb(cur, "seanet_enc_in", -1);
|
||||
|
||||
for (int i = 0; i < hparams.seanet_n_stage; i++) {
|
||||
const auto & stage = seanet.stages[i];
|
||||
const int stride = hparams.seanet_ratios[i];
|
||||
|
||||
cur = res_unit(cur, stage, 1);
|
||||
cur = ggml_elu(ctx0, cur);
|
||||
cur = conv1d(cur, stage.scale_conv_w, stage.scale_conv_b, stride, 1);
|
||||
cb(cur, "seanet_enc_stage", i);
|
||||
}
|
||||
|
||||
cur = ggml_elu(ctx0, cur);
|
||||
cur = conv1d(cur, seanet.conv_out_w, seanet.conv_out_b, 1, 1);
|
||||
cb(cur, "seanet_enc_out", -1);
|
||||
|
||||
return cur;
|
||||
}
|
||||
|
||||
ggml_tensor * clip_graph_pockettts_seanet::decode(ggml_tensor * x) const {
|
||||
const auto & seanet = model.seanet;
|
||||
const bool stream = !state_in.empty();
|
||||
|
||||
ggml_tensor * cur = conv1d(x, seanet.conv_in_w, seanet.conv_in_b, 1, 1, false,
|
||||
stream ? "dec_in" : "");
|
||||
cb(cur, "seanet_dec_in", -1);
|
||||
|
||||
for (int i = 0; i < hparams.seanet_n_stage; i++) {
|
||||
const auto & stage = seanet.stages[i];
|
||||
// the decoder mirrors the encoder, so the ratios are walked backwards
|
||||
const int stride = hparams.seanet_ratios[hparams.seanet_n_stage - 1 - i];
|
||||
const std::string id = std::to_string(i);
|
||||
|
||||
cur = ggml_elu(ctx0, cur);
|
||||
cur = conv_transpose1d(cur, stage.scale_conv_w, stage.scale_conv_b, stride,
|
||||
stream ? "dec_up_" + id : "");
|
||||
cur = res_unit(cur, stage, 1, stream ? "dec_res_" + id : "");
|
||||
cb(cur, "seanet_dec_stage", i);
|
||||
}
|
||||
|
||||
cur = ggml_elu(ctx0, cur);
|
||||
cur = conv1d(cur, seanet.conv_out_w, seanet.conv_out_b, 1, 1, false,
|
||||
stream ? "dec_out" : "");
|
||||
cb(cur, "seanet_dec_out", -1);
|
||||
|
||||
return cur;
|
||||
}
|
||||
@@ -0,0 +1,77 @@
|
||||
#include "models.h"
|
||||
|
||||
// voice-prompt encoder: raw 24kHz waveform -> one conditioning row per 12.5Hz frame
|
||||
// mimi encoder (SEANet + transformer + downsample), then flow_lm.speaker_proj_weight
|
||||
|
||||
// pre-norm block with layer scale on both residual paths, see mimi_transformer.py
|
||||
ggml_tensor * clip_graph_pockettts_spkenc::tfm_layer_forward(ggml_tensor * cur, const clip_layer & layer, ggml_tensor * inp_pos, ggml_tensor * kq_mask, int il) const {
|
||||
ggml_tensor * inp = cur;
|
||||
|
||||
cur = build_norm(cur, layer.ln_1_w, layer.ln_1_b, NORM_TYPE_NORMAL, eps, il);
|
||||
|
||||
ggml_tensor * Qcur = build_mm(layer.q_w, cur);
|
||||
ggml_tensor * Kcur = build_mm(layer.k_w, cur);
|
||||
ggml_tensor * Vcur = build_mm(layer.v_w, cur);
|
||||
|
||||
const int64_t n_pos = cur->ne[1];
|
||||
Qcur = ggml_reshape_3d(ctx0, Qcur, d_head, n_head, n_pos);
|
||||
Kcur = ggml_reshape_3d(ctx0, Kcur, d_head, n_head, n_pos);
|
||||
Vcur = ggml_reshape_3d(ctx0, Vcur, d_head, n_head, n_pos);
|
||||
|
||||
Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, d_head, GGML_ROPE_TYPE_NORMAL, 0,
|
||||
hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
|
||||
Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, d_head, GGML_ROPE_TYPE_NORMAL, 0,
|
||||
hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
|
||||
|
||||
cur = build_attn(layer.o_w, nullptr, Qcur, Kcur, Vcur, kq_mask, kq_scale, il);
|
||||
cur = ggml_mul(ctx0, cur, layer.ls_1_w);
|
||||
cur = ggml_add(ctx0, cur, inp);
|
||||
|
||||
inp = cur;
|
||||
cur = build_norm(cur, layer.ln_2_w, layer.ln_2_b, NORM_TYPE_NORMAL, eps, il);
|
||||
cur = build_ffn(cur, layer.ff_up_w, nullptr, nullptr, nullptr, layer.ff_down_w, nullptr, FFN_GELU, il);
|
||||
cur = ggml_mul(ctx0, cur, layer.ls_2_w);
|
||||
cur = ggml_add(ctx0, cur, inp);
|
||||
|
||||
return cur;
|
||||
}
|
||||
|
||||
ggml_cgraph * clip_graph_pockettts_spkenc::build() {
|
||||
// the preprocessor hands over the waveform as a single-row "mel", already [n_samples, 1]
|
||||
ggml_tensor * inp_raw = build_inp_raw(1);
|
||||
ggml_tensor * cur = ggml_reshape_2d(ctx0, inp_raw, inp_raw->ne[0], inp_raw->ne[1]);
|
||||
|
||||
clip_graph_pockettts_seanet seanet(*this);
|
||||
cur = seanet.encode(cur);
|
||||
cb(cur, "mimi_enc", -1);
|
||||
|
||||
// [T, 512] -> transformer works on [512, T]
|
||||
cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur));
|
||||
|
||||
ggml_tensor * inp_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, cur->ne[1]);
|
||||
ggml_set_name(inp_pos, "inp_pos");
|
||||
ggml_set_input(inp_pos);
|
||||
|
||||
// the mimi transformer is causal with a sliding window, see _build_attention_mask()
|
||||
ggml_tensor * kq_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, cur->ne[1], cur->ne[1]);
|
||||
ggml_set_name(kq_mask, "kq_mask");
|
||||
ggml_set_input(kq_mask);
|
||||
|
||||
for (int il = 0; il < n_layer; il++) {
|
||||
cur = tfm_layer_forward(cur, model.layers[il], inp_pos, kq_mask, il);
|
||||
}
|
||||
cb(cur, "mimi_enc_tfm", -1);
|
||||
|
||||
// downsample to the model frame rate, [512, T] -> [T, 512] -> [T / 16, 32]
|
||||
cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur));
|
||||
cur = seanet.conv1d(cur, model.downsample_w, nullptr, hparams.mimi_downsample, 1, true);
|
||||
cb(cur, "mimi_downsample", -1);
|
||||
|
||||
// voice latent -> backbone embd
|
||||
cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur));
|
||||
cur = build_mm(model.spk_proj_w, cur);
|
||||
cb(cur, "spk_proj", -1);
|
||||
|
||||
ggml_build_forward_expand(gf, cur);
|
||||
return gf;
|
||||
}
|
||||
@@ -610,6 +610,10 @@ std::vector<c2w_state_slot> list_c2w_state_slots(const clip_hparams & hparams, c
|
||||
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)
|
||||
|
||||
Reference in New Issue
Block a user