Brunobkr/llama.cpp_AlgMor24_github

1#include "models.h"2 3static constexpr int SPK_RES2NET_SCALE = 8; // enc_res2net_scale4static constexpr int SPK_DILATIONS[3]  = { 2, 3, 4 }; // enc_dilations[1..3]5 6// conv1d, kernel K, padding "same" (reflect), dilation d7// x: [C, T] (ne[0]=C, ne[1]=T) -> [out_c, T]8ggml_tensor * clip_graph_qwen3tts_spkenc::conv1d_same(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int dilation) const {9    const int K   = (int) w->ne[0];10    const int IC  = (int) w->ne[1];11    const int OC  = (int) w->ne[2];12    const int pad = ((K - 1) * dilation) / 2;13 14    // ggml_pad_reflect_1d pads ne[0], so bring T onto ne[0] first, same layout as im2col wants15    ggml_tensor * x_t = ggml_cont(ctx0, ggml_transpose(ctx0, x)); // [T, IC]16    if (pad > 0) {17        x_t = ggml_pad_reflect_1d(ctx0, x_t, pad, pad); // [T + 2*pad, IC]18    }19    ggml_tensor * x4d = ggml_reshape_4d(ctx0, x_t, x_t->ne[0], IC, 1, 1);20 21    // dummy F32 kernel, im2col only reads its shape, so a quantized w does not assert22    ggml_tensor * dummy = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, K, IC, 1, 1);23 24    ggml_tensor * col = ggml_im2col(ctx0, dummy, x4d, 1, 1, 0, 0, dilation, 1, false, GGML_TYPE_F32);25    const int64_t T_out = col->ne[1];26    col = ggml_reshape_2d(ctx0, col, (int64_t) K * IC, T_out);27 28    ggml_tensor * w2d = ggml_reshape_2d(ctx0, w, (int64_t) K * IC, OC);29    ggml_tensor * y   = ggml_mul_mat(ctx0, w2d, col); // [OC, T_out]30    ggml_mul_mat_set_prec(y, GGML_PREC_F32);31 32    ggml_tensor * b2d = ggml_reshape_2d(ctx0, b, OC, 1);33    y = ggml_add(ctx0, y, b2d);34    return y;35}36 37// Res2Net: split channel axis into `scale` chunks, chain dilated conv1d branches38// x: [C, T] -> [C, T]39ggml_tensor * clip_graph_qwen3tts_spkenc::res2net(ggml_tensor * x, const clip_layer & layer, int dilation, int scale) const {40    const int64_t C  = x->ne[0];41    const int64_t T  = x->ne[1];42    const int64_t Cs = C / scale;43 44    std::vector<ggml_tensor *> outs;45    outs.reserve(scale);46 47    auto chunk = [&](int i) -> ggml_tensor * {48        return ggml_view_2d(ctx0, x, Cs, T, x->nb[1], (size_t) i * Cs * x->nb[0]);49    };50 51    ggml_tensor * prev = nullptr;52    for (int i = 0; i < scale; i++) {53        ggml_tensor * c = ggml_cont(ctx0, chunk(i));54        if (i == 0) {55            outs.push_back(c);56            continue;57        }58        ggml_tensor * inp = (i >= 2) ? ggml_add(ctx0, c, prev) : c;59        ggml_tensor * y   = conv1d_same(inp, layer.res2_conv_w[i - 1], layer.res2_conv_b[i - 1], dilation);60        y                 = ggml_relu(ctx0, y);61        outs.push_back(y);62        prev = y;63    }64 65    ggml_tensor * acc = outs[0];66    for (int i = 1; i < scale; i++) {67        acc = ggml_concat(ctx0, acc, outs[i], 0);68    }69    return acc;70}71 72// squeeze-and-excitation gate. x: [C, T] -> [C, T]73ggml_tensor * clip_graph_qwen3tts_spkenc::se_block(ggml_tensor * x, const clip_layer & layer) const {74    // temporal mean, keepdim: transpose so T is on ne[0], reduce, transpose back75    ggml_tensor * x_t  = ggml_cont(ctx0, ggml_transpose(ctx0, x));    // [T, C]76    ggml_tensor * mean = ggml_mean(ctx0, x_t);                        // [1, C]77    mean               = ggml_cont(ctx0, ggml_transpose(ctx0, mean)); // [C, 1]78 79    ggml_tensor * h = conv1d_same(mean, layer.se_conv1_w, layer.se_conv1_b, 1);80    h = ggml_relu(ctx0, h);81    h = conv1d_same(h, layer.se_conv2_w, layer.se_conv2_b, 1);82    h = ggml_sigmoid(ctx0, h); // [C, 1]83 84    return ggml_mul(ctx0, x, h); // broadcast gate over T85}86 87// tdnn1 -> res2net -> tdnn2 -> se, plus residual. x: [C, T] -> [C, T]88ggml_tensor * clip_graph_qwen3tts_spkenc::se_res2net_block(ggml_tensor * x, const clip_layer & layer, int dilation, int scale) const {89    ggml_tensor * residual = x;90    ggml_tensor * h  = conv1d_same(x, layer.conv_pw1_w, layer.conv_pw1_b, 1); // tdnn191    h = ggml_relu(ctx0, h);92    h = res2net(h, layer, dilation, scale);93    h = conv1d_same(h, layer.conv_pw2_w, layer.conv_pw2_b, 1); // tdnn294    h = ggml_relu(ctx0, h);95    h = se_block(h, layer);96    return ggml_add(ctx0, h, residual);97}98 99// attentive statistics pooling. x: [C, T] -> [2*C, 1]100ggml_tensor * clip_graph_qwen3tts_spkenc::attentive_stats_pool(ggml_tensor * x) const {101    const int64_t T = x->ne[1];102 103    // mean over T: [C, 1]104    ggml_tensor * x_t  = ggml_cont(ctx0, ggml_transpose(ctx0, x));105    ggml_tensor * mean = ggml_mean(ctx0, x_t);106    mean               = ggml_cont(ctx0, ggml_transpose(ctx0, mean));107 108    // std over T: sqrt(clamp(mean((x - mean)^2), eps))109    ggml_tensor * mean_rep = ggml_repeat(ctx0, mean, x);110    ggml_tensor * centered = ggml_sub(ctx0, x, mean_rep);111    ggml_tensor * var_t    = ggml_cont(ctx0, ggml_transpose(ctx0, ggml_sqr(ctx0, centered)));112    ggml_tensor * var      = ggml_mean(ctx0, var_t);113    var                    = ggml_cont(ctx0, ggml_transpose(ctx0, var));114    var                    = ggml_scale_bias(ctx0, var, 1.0f, 1e-12f);115    ggml_tensor * std      = ggml_sqrt(ctx0, var);116 117    // attention input: cat([x, mean, std]) along channel axis -> [3C, T]118    ggml_tensor * std_rep = ggml_repeat(ctx0, std, x);119    ggml_tensor * cat     = ggml_concat(ctx0, x, mean_rep, 0);120    cat                   = ggml_concat(ctx0, cat, std_rep, 0);121 122    // attention TDNN (3C -> attn_c) + ReLU, tanh, then 1x1 conv (attn_c -> C)123    ggml_tensor * a = conv1d_same(cat, model.spk_asp_tdnn_w, model.spk_asp_tdnn_b, 1);124    a = ggml_relu(ctx0, a);125    a = ggml_tanh(ctx0, a);126    a = conv1d_same(a, model.spk_asp_attn_w, model.spk_asp_attn_b, 1);127 128    // softmax over T129    ggml_tensor * a_t = ggml_cont(ctx0, ggml_transpose(ctx0, a));   // [T, C]130    ggml_tensor * w_t = ggml_soft_max(ctx0, a_t);131    ggml_tensor * w   = ggml_cont(ctx0, ggml_transpose(ctx0, w_t)); // [C, T]132 133    // weighted mean: sum(w * x) over T, multiply by T to undo ggml_mean's 1/T scaling134    ggml_tensor * wx     = ggml_mul(ctx0, w, x);135    ggml_tensor * wx_t   = ggml_cont(ctx0, ggml_transpose(ctx0, wx));136    ggml_tensor * w_mean = ggml_mean(ctx0, wx_t);137    w_mean = ggml_scale(ctx0, w_mean, (float) T);138    w_mean = ggml_cont(ctx0, ggml_transpose(ctx0, w_mean)); // [C, 1]139 140    // weighted std: sum(w * (x - w_mean)^2) over T141    ggml_tensor * w_mean_rep = ggml_repeat(ctx0, w_mean, x);142    ggml_tensor * dev        = ggml_sub(ctx0, x, w_mean_rep);143    ggml_tensor * w_var_in   = ggml_mul(ctx0, w, ggml_sqr(ctx0, dev));144    ggml_tensor * w_var_t    = ggml_cont(ctx0, ggml_transpose(ctx0, w_var_in));145    ggml_tensor * w_var      = ggml_mean(ctx0, w_var_t);146    w_var                    = ggml_scale(ctx0, w_var, (float) T);147    w_var                    = ggml_cont(ctx0, ggml_transpose(ctx0, w_var));148    w_var                    = ggml_scale_bias(ctx0, w_var, 1.0f, 1e-12f);149    ggml_tensor * w_std      = ggml_sqrt(ctx0, w_var);150 151    return ggml_concat(ctx0, w_mean, w_std, 0); // [2C, 1]152}153 154ggml_cgraph * clip_graph_qwen3tts_spkenc::build() {155    // inp_raw: [T, n_mel, 1, 1], from mtmd_audio_preprocessor_qwen3tts_spk156    ggml_tensor * inp = build_inp_raw(1);157    inp = ggml_reshape_2d(ctx0, inp, inp->ne[0], inp->ne[1]);158 159    // this file's convention is [C, T]; the preprocessor delivers [T, C]160    ggml_tensor * mel = ggml_cont(ctx0, ggml_transpose(ctx0, inp)); // [n_mel, T]161    cb(mel, "mel", -1);162 163    // frontend conv0 TDNN k=5, dilation=1: 128 -> 512164    ggml_tensor * cur = conv1d_same(mel, model.conv1d_1_w, model.conv1d_1_b, 1);165    cur = ggml_relu(ctx0, cur);166    cb(cur, "frontend", -1);167 168    // 3 SE-Res2Net blocks at dilations 2, 3, 4169    GGML_ASSERT((int) model.layers.size() == 3);170    std::vector<ggml_tensor *> blk_out(3);171    for (int il = 0; il < 3; il++) {172        cur = se_res2net_block(cur, model.layers[il], SPK_DILATIONS[il], SPK_RES2NET_SCALE);173        blk_out[il] = cur;174        cb(cur, "block_out", il);175    }176 177    // multi-layer feature aggregation: cat blk[0..2] then TDNN k=1 + ReLU178    ggml_tensor * cat = ggml_concat(ctx0, blk_out[0], blk_out[1], 0);179    cat = ggml_concat(ctx0, cat, blk_out[2], 0); // [1536, T]180    ggml_tensor * mfa = conv1d_same(cat, model.conv_out_w, model.conv_out_b, 1);181    mfa = ggml_relu(ctx0, mfa);182    cb(mfa, "mfa", -1);183 184    // attentive statistics pooling: [1536, T] -> [3072, 1]185    ggml_tensor * stats = attentive_stats_pool(mfa);186    cb(stats, "asp", -1);187 188    // final FC k=1: [3072, 1] -> [enc_dim, 1]189    ggml_tensor * emb = conv1d_same(stats, model.mm_fc_w, model.mm_fc_b, 1);190 191    emb = ggml_reshape_1d(ctx0, emb, emb->ne[0]);192    emb = ggml_cont(ctx0, emb);193    cb(emb, "spk_embedding", -1);194 195    ggml_build_forward_expand(gf, emb);196    return gf;197}198