734 lines
32 KiB
Plaintext
734 lines
32 KiB
Plaintext
// Forward complet du décodeur Qwen3-TTS en pur ggml (CPU baseline).
|
||
// Approche staging strict : chaque sous-module = un graphe ggml indépendant
|
||
// avec son propre ggml_init, son ggml_graph_compute_with_ctx, et le résultat
|
||
// recopié en host array entre étages. C'est verbeux mais reprend exactement
|
||
// le code des tests par stage qui sont validés numériquement individuellement.
|
||
|
||
#include "decoder.h"
|
||
#include <ggml-backend.h>
|
||
#include <ggml-cpu.h>
|
||
#include <cassert>
|
||
#include <chrono>
|
||
#include <cstdio>
|
||
#include <cstdlib>
|
||
#include <cstring>
|
||
#include <cmath>
|
||
#include <vector>
|
||
|
||
namespace kazeia::tts {
|
||
|
||
// ----- Stage 1 : Quantizer (codes → hidden [512, T]) -----
|
||
// codes_flat layout : [n_q=16][T] flat = codes_flat[layer*T + t]
|
||
// hidden output layout (host) : [c][t] flat with stride compatible test_quantizer's ggml output :
|
||
// ggml ne=[512, T] → memory data[c + 512*t] (col-major).
|
||
// On copie memory tel quel dans le host array (480KB pour T=59).
|
||
static std::vector<float> stage_quantizer(Decoder & d, const std::vector<int32_t> & codes_flat, int T) {
|
||
const int n_q = d.cfg.num_quantizers; // 16
|
||
const int dim = d.cfg.codebook_dim / 2; // 256 (rvq dimension)
|
||
const int hidden_dim = d.cfg.codebook_dim; // 512
|
||
|
||
size_t mem_size = 64 * 1024 * 1024;
|
||
struct ggml_init_params p = { mem_size, nullptr, false };
|
||
struct ggml_context * ctx = ggml_init(p);
|
||
|
||
auto build_codes = [&](int layer_idx) {
|
||
ggml_tensor * t = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, T);
|
||
memcpy(t->data, &codes_flat[layer_idx * T], T * sizeof(int32_t));
|
||
return t;
|
||
};
|
||
|
||
// Sum_first : 1 codebook (rvq_first.cb.0)
|
||
ggml_tensor * emb_first = d.tensors["quantizer.rvq_first.cb.0.embedding"];
|
||
ggml_tensor * sum_first = ggml_get_rows(ctx, emb_first, build_codes(0)); // [256, T]
|
||
|
||
// Sum_rest : 15 codebooks (rvq_rest.cb.0..14)
|
||
ggml_tensor * sum_rest = nullptr;
|
||
for (int i = 0; i < 15; i++) {
|
||
char key[64]; snprintf(key, sizeof(key), "quantizer.rvq_rest.cb.%d.embedding", i);
|
||
ggml_tensor * emb = d.tensors[key];
|
||
ggml_tensor * x = ggml_get_rows(ctx, emb, build_codes(1 + i));
|
||
sum_rest = sum_rest ? ggml_add(ctx, sum_rest, x) : x;
|
||
}
|
||
|
||
auto apply_proj = [&](const std::string & key, ggml_tensor * x_KT) -> ggml_tensor * {
|
||
ggml_tensor * w = d.tensors[key];
|
||
ggml_tensor * w2 = ggml_reshape_2d(ctx, w, w->ne[1], w->ne[2]);
|
||
return ggml_mul_mat(ctx, w2, x_KT);
|
||
};
|
||
|
||
ggml_tensor * out_first = apply_proj("quantizer.rvq_first.output_proj.weight", sum_first);
|
||
ggml_tensor * out_rest = apply_proj("quantizer.rvq_rest.output_proj.weight", sum_rest);
|
||
ggml_tensor * out = ggml_add(ctx, out_first, out_rest); // [512, T]
|
||
|
||
struct ggml_cgraph * gf = ggml_new_graph(ctx);
|
||
ggml_build_forward_expand(gf, out);
|
||
ggml_build_forward_expand(gf, out_first);
|
||
ggml_build_forward_expand(gf, out_rest);
|
||
ggml_build_forward_expand(gf, sum_first);
|
||
ggml_build_forward_expand(gf, sum_rest);
|
||
int n_threads = 8;
|
||
if (const char * t_env = std::getenv("GGML_NUM_THREADS")) {
|
||
int n = std::atoi(t_env); if (n >= 1) n_threads = n;
|
||
}
|
||
ggml_graph_compute_with_ctx(ctx, gf, n_threads);
|
||
|
||
// Debug dumps
|
||
if (const char * dir = std::getenv("DECODER_DEBUG_DUMP_DIR")) {
|
||
auto dump_tensor = [&](const std::string & name, ggml_tensor * t) {
|
||
std::string path = std::string(dir) + "/" + name + ".bin";
|
||
FILE * f = fopen(path.c_str(), "wb");
|
||
if (f) {
|
||
size_t n_elem = ggml_nelements(t);
|
||
fwrite(t->data, sizeof(float), n_elem, f);
|
||
fclose(f);
|
||
fprintf(stderr, " dump: %s ne=[%lld,%lld] (%zu floats)\n",
|
||
path.c_str(), (long long)t->ne[0], (long long)t->ne[1], n_elem);
|
||
}
|
||
};
|
||
dump_tensor("q_sum_first", sum_first);
|
||
dump_tensor("q_sum_rest", sum_rest);
|
||
dump_tensor("q_out_first", out_first);
|
||
dump_tensor("q_out_rest", out_rest);
|
||
}
|
||
|
||
std::vector<float> result(hidden_dim * T);
|
||
memcpy(result.data(), out->data, hidden_dim * T * sizeof(float));
|
||
|
||
ggml_free(ctx);
|
||
(void)dim; (void)n_q;
|
||
return result;
|
||
}
|
||
|
||
// ----- Stage 2 : pre_conv (Conv1d k=3 causal pad 2-left + 1-right) -----
|
||
// Reproduit exactement test_pre_conv. Input layout : ggml [Tp=T+3, in=512]
|
||
// avec dst[(t+2) + Tp*c] = src_input[c + 512*t]. Output : [out=1024, T] dans
|
||
// le sens ggml-natif après slice T_out=T+1 → premier T par channel.
|
||
static std::vector<float> stage_pre_conv(Decoder & d, const std::vector<float> & hidden, int T) {
|
||
const int Cin = d.cfg.codebook_dim; // 512
|
||
const int Cout = d.cfg.latent_dim; // 1024
|
||
const int Tp = T + 3;
|
||
|
||
size_t mem_size = 64 * 1024 * 1024;
|
||
struct ggml_init_params p = { mem_size, nullptr, false };
|
||
struct ggml_context * ctx = ggml_init(p);
|
||
|
||
ggml_tensor * x = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, Tp, Cin);
|
||
{
|
||
float * dst = (float*) x->data;
|
||
for (int i = 0; i < Tp * Cin; i++) dst[i] = 0.0f;
|
||
// hidden in [c + 512*t] layout (from quantizer output memcpy).
|
||
// Map to ggml[t, c] = pytorch[c, t-2] → dst[(t+2) + Tp*c] = hidden[c + 512*t]
|
||
for (int c = 0; c < Cin; c++) {
|
||
for (int t = 0; t < T; t++) {
|
||
dst[(t + 2) + Tp * c] = hidden[c + Cin * t];
|
||
}
|
||
}
|
||
}
|
||
|
||
ggml_tensor * w = d.tensors["pre_conv.weight"]; // F16
|
||
ggml_tensor * b = d.tensors["pre_conv.bias"];
|
||
ggml_tensor * out = ggml_conv_1d(ctx, w, x, 1, 0, 1);
|
||
// out ne=[T+1, Cout]
|
||
ggml_tensor * b2 = ggml_reshape_2d(ctx, b, 1, Cout);
|
||
out = ggml_add(ctx, out, b2);
|
||
|
||
struct ggml_cgraph * gf = ggml_new_graph(ctx);
|
||
ggml_build_forward_expand(gf, out);
|
||
int n_threads = 8;
|
||
if (const char * t_env = std::getenv("GGML_NUM_THREADS")) {
|
||
int n = std::atoi(t_env); if (n >= 1) n_threads = n;
|
||
}
|
||
ggml_graph_compute_with_ctx(ctx, gf, n_threads);
|
||
|
||
// out ne=[T+1, Cout] → take first T per channel.
|
||
// Output layout (matches test_pre_conv) : memory [t + T_out*c] but slice to T.
|
||
// Pour la stage suivante (pre_transformer), on veut input dans le format
|
||
// que test_pre_transformer attend : ggml ne=[latent=Cout, T] avec
|
||
// dst[c + latent*t] = src[t + T*c] (= pytorch [B, latent, T]).
|
||
// Or notre out actuel a memory[t + (T+1)*c] = pytorch_pre_conv[c, t].
|
||
// On extrait de manière à produire host[c + Cout*t] = pytorch[c, t].
|
||
const int T_out_y = (int)out->ne[0];
|
||
const float * data = (const float*) out->data;
|
||
std::vector<float> result(Cout * T);
|
||
for (int c = 0; c < Cout; c++) {
|
||
for (int t = 0; t < T; t++) {
|
||
result[c + Cout * t] = data[t + T_out_y * c];
|
||
}
|
||
}
|
||
|
||
ggml_free(ctx);
|
||
return result;
|
||
}
|
||
|
||
// ----- Stage 3 : pre_transformer (input_proj + 8 DiTLayers + norm + output_proj) -----
|
||
// Input layout (host) : pc_out[c + Cout*t] = pytorch_pre_conv[c, t]
|
||
// = ggml ne=[latent, T] data[c + latent*t]
|
||
// Output layout (host) : pt_out[c + latent*t] = pytorch_pre_transformer_out[t, c]
|
||
// (PyTorch last_hidden_state is [B, T, latent], les valeurs sont les mêmes
|
||
// que [B, latent, T] indexées par (c, t) au lieu de (t, c) → mêmes valeurs scalaires).
|
||
static std::vector<float> stage_pre_transformer(Decoder & d, const std::vector<float> & pc_out, int T) {
|
||
const int latent = d.cfg.latent_dim; // 1024
|
||
const int hidden = d.cfg.hidden_size; // 512
|
||
const int head_dim = d.cfg.head_dim; // 64
|
||
const int n_heads = d.cfg.num_attention_heads; // 16
|
||
const int qkv_dim = n_heads * head_dim;
|
||
const int n_layers = d.cfg.num_hidden_layers; // 8
|
||
const float rope_base = d.cfg.rope_theta;
|
||
const float rms_eps = d.cfg.rms_norm_eps;
|
||
|
||
size_t mem_size = 512ULL * 1024 * 1024;
|
||
struct ggml_init_params p = { mem_size, nullptr, false };
|
||
struct ggml_context * ctx = ggml_init(p);
|
||
|
||
// Input ggml ne=[latent, T] avec dst[c + latent*t] = pc_out[c + latent*t]
|
||
ggml_tensor * x = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, latent, T);
|
||
memcpy(x->data, pc_out.data(), latent * T * sizeof(float));
|
||
|
||
auto get = [&](const std::string & key) -> ggml_tensor * { return d.tensors[key]; };
|
||
|
||
auto linear = [&](ggml_tensor * x, const std::string & w_key, const std::string & b_key = "") -> ggml_tensor * {
|
||
ggml_tensor * y = ggml_mul_mat(ctx, get(w_key), x);
|
||
if (!b_key.empty()) {
|
||
ggml_tensor * b = get(b_key);
|
||
ggml_tensor * b2 = ggml_reshape_2d(ctx, b, b->ne[0], 1);
|
||
y = ggml_add(ctx, y, b2);
|
||
}
|
||
return y;
|
||
};
|
||
|
||
x = linear(x, "pre_transformer.input_proj.weight", "pre_transformer.input_proj.bias");
|
||
|
||
ggml_tensor * pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, T);
|
||
{
|
||
int32_t * pp = (int32_t*) pos->data;
|
||
for (int t = 0; t < T; t++) pp[t] = t;
|
||
}
|
||
ggml_tensor * mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, T, T);
|
||
{
|
||
float * mm = (float*) mask->data;
|
||
for (int q = 0; q < T; q++)
|
||
for (int k = 0; k < T; k++)
|
||
mm[k + T*q] = (k > q) ? -INFINITY : 0.0f;
|
||
}
|
||
const float scale = 1.0f / std::sqrt((float)head_dim);
|
||
|
||
for (int L = 0; L < n_layers; L++) {
|
||
char prefix[64]; snprintf(prefix, sizeof(prefix), "pre_transformer.%d", L);
|
||
ggml_tensor * residual = x;
|
||
|
||
ggml_tensor * x_n = ggml_rms_norm(ctx, x, rms_eps);
|
||
x_n = ggml_mul(ctx, x_n, get(std::string(prefix) + ".input_norm.weight"));
|
||
|
||
ggml_tensor * Q = linear(x_n, std::string(prefix) + ".attn.q.weight");
|
||
ggml_tensor * K = linear(x_n, std::string(prefix) + ".attn.k.weight");
|
||
ggml_tensor * V = linear(x_n, std::string(prefix) + ".attn.v.weight");
|
||
|
||
Q = ggml_reshape_3d(ctx, Q, head_dim, n_heads, T);
|
||
K = ggml_reshape_3d(ctx, K, head_dim, n_heads, T);
|
||
V = ggml_reshape_3d(ctx, V, head_dim, n_heads, T);
|
||
Q = ggml_rope_ext(ctx, Q, pos, NULL, head_dim, GGML_ROPE_TYPE_NEOX, 0,
|
||
rope_base, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
|
||
K = ggml_rope_ext(ctx, K, pos, NULL, head_dim, GGML_ROPE_TYPE_NEOX, 0,
|
||
rope_base, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
|
||
|
||
Q = ggml_cont(ctx, ggml_permute(ctx, Q, 0, 2, 1, 3));
|
||
K = ggml_cont(ctx, ggml_permute(ctx, K, 0, 2, 1, 3));
|
||
V = ggml_cont(ctx, ggml_permute(ctx, V, 1, 2, 0, 3));
|
||
|
||
ggml_tensor * KQ = ggml_mul_mat(ctx, K, Q);
|
||
KQ = ggml_soft_max_ext(ctx, KQ, mask, scale, 0.0f);
|
||
ggml_tensor * KQV = ggml_mul_mat(ctx, V, KQ);
|
||
KQV = ggml_cont(ctx, ggml_permute(ctx, KQV, 0, 2, 1, 3));
|
||
KQV = ggml_reshape_2d(ctx, KQV, qkv_dim, T);
|
||
|
||
ggml_tensor * attn_out = linear(KQV, std::string(prefix) + ".attn.o.weight");
|
||
attn_out = ggml_mul(ctx, attn_out, get(std::string(prefix) + ".attn_layerscale"));
|
||
x = ggml_add(ctx, residual, attn_out);
|
||
|
||
residual = x;
|
||
ggml_tensor * x_n2 = ggml_rms_norm(ctx, x, rms_eps);
|
||
x_n2 = ggml_mul(ctx, x_n2, get(std::string(prefix) + ".post_attn_norm.weight"));
|
||
ggml_tensor * gate = linear(x_n2, std::string(prefix) + ".mlp.gate.weight");
|
||
ggml_tensor * up = linear(x_n2, std::string(prefix) + ".mlp.up.weight");
|
||
gate = ggml_silu(ctx, gate);
|
||
ggml_tensor * mid = ggml_mul(ctx, gate, up);
|
||
ggml_tensor * mlp_out = linear(mid, std::string(prefix) + ".mlp.down.weight");
|
||
mlp_out = ggml_mul(ctx, mlp_out, get(std::string(prefix) + ".mlp_layerscale"));
|
||
x = ggml_add(ctx, residual, mlp_out);
|
||
}
|
||
x = ggml_rms_norm(ctx, x, rms_eps);
|
||
x = ggml_mul(ctx, x, get("pre_transformer.norm.weight"));
|
||
x = linear(x, "pre_transformer.output_proj.weight", "pre_transformer.output_proj.bias");
|
||
// x ne=[latent, T]
|
||
(void)hidden;
|
||
|
||
struct ggml_cgraph * gf = ggml_new_graph_custom(ctx, 8192, false);
|
||
ggml_build_forward_expand(gf, x);
|
||
int n_threads = 8;
|
||
if (const char * t_env = std::getenv("GGML_NUM_THREADS")) {
|
||
int n = std::atoi(t_env); if (n >= 1) n_threads = n;
|
||
}
|
||
ggml_graph_compute_with_ctx(ctx, gf, n_threads);
|
||
|
||
std::vector<float> result(latent * T);
|
||
memcpy(result.data(), x->data, latent * T * sizeof(float));
|
||
// result[c + latent*t] = pytorch_pre_transformer_out at (logical c, logical t)
|
||
// = pytorch [B, T, latent] at (t, c) = pytorch [B, latent, T] at (c, t) — valeur scalaire.
|
||
|
||
ggml_free(ctx);
|
||
return result;
|
||
}
|
||
|
||
// ----- Stage 4 : upsample (2 stages : ConvTranspose1d + ConvNeXtBlock) -----
|
||
// Input host : pt_out[c + latent*t] (sortie de pre_transformer en CT layout ggml)
|
||
// Output host : us_out[c + latent*t_out] avec t_out = T*4 (après 2× upsample par 2)
|
||
//
|
||
// Reproduit exactement test_upsample (deux stages cumulés) sauf que l'input
|
||
// vient d'un host array au lieu d'un .npy.
|
||
static std::vector<float> stage_upsample(Decoder & d, const std::vector<float> & pt_out, int T) {
|
||
const int C = d.cfg.latent_dim; // 1024
|
||
const int T_final = T * 4; // après 2 upsampling ratio 2
|
||
|
||
size_t mem_size = 1024ULL * 1024 * 1024;
|
||
struct ggml_init_params p = { mem_size, nullptr, false };
|
||
struct ggml_context * ctx = ggml_init(p);
|
||
|
||
auto get = [&](const std::string & key) { return d.tensors[key]; };
|
||
|
||
// Input ggml ne=[C, T_in] : memcpy direct depuis pt_out (même layout)
|
||
ggml_tensor * x = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, C, T);
|
||
memcpy(x->data, pt_out.data(), C * T * sizeof(float));
|
||
|
||
int T_cur = T;
|
||
for (int s = 0; s < 2; s++) {
|
||
char prefix[64]; snprintf(prefix, sizeof(prefix), "upsample.%d", s);
|
||
|
||
// Stage A : ConvTranspose1d k=2 s=2
|
||
ggml_tensor * x_tch = ggml_cont(ctx, ggml_transpose(ctx, x)); // [T, C]
|
||
ggml_tensor * y = ggml_conv_transpose_1d(ctx, get(std::string(prefix) + ".convtr.weight"),
|
||
x_tch, 2, 0, 1);
|
||
ggml_tensor * b_ct = get(std::string(prefix) + ".convtr.bias");
|
||
y = ggml_add(ctx, y, ggml_reshape_2d(ctx, b_ct, 1, C));
|
||
y = ggml_cont(ctx, ggml_transpose(ctx, y)); // [C, T*2]
|
||
T_cur *= 2;
|
||
|
||
// Stage B : ConvNeXtBlock
|
||
ggml_tensor * residual = y;
|
||
|
||
// dwconv k=7 causal pad gauche=6
|
||
ggml_tensor * y_tch = ggml_cont(ctx, ggml_transpose(ctx, y)); // [T*2, C]
|
||
ggml_tensor * zeros = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 6, C);
|
||
memset(zeros->data, 0, 6 * C * sizeof(float));
|
||
ggml_tensor * y_pad = ggml_concat(ctx, zeros, y_tch, 0); // [pad+T, C]
|
||
ggml_tensor * dw = ggml_conv_1d_dw(ctx, get(std::string(prefix) + ".cnx.dwconv.weight"),
|
||
y_pad, 1, 0, 1);
|
||
dw = ggml_add(ctx, dw, ggml_reshape_2d(ctx, get(std::string(prefix) + ".cnx.dwconv.bias"), 1, C));
|
||
dw = ggml_cont(ctx, ggml_transpose(ctx, dw)); // [C, T]
|
||
|
||
ggml_tensor * h = ggml_norm(ctx, dw, 1e-6f);
|
||
h = ggml_mul(ctx, h, get(std::string(prefix) + ".cnx.norm.weight"));
|
||
h = ggml_add(ctx, h, get(std::string(prefix) + ".cnx.norm.bias"));
|
||
h = ggml_mul_mat(ctx, get(std::string(prefix) + ".cnx.pwconv1.weight"), h);
|
||
ggml_tensor * pw1_b = get(std::string(prefix) + ".cnx.pwconv1.bias");
|
||
h = ggml_add(ctx, h, ggml_reshape_2d(ctx, pw1_b, pw1_b->ne[0], 1));
|
||
h = ggml_gelu_erf(ctx, h);
|
||
h = ggml_mul_mat(ctx, get(std::string(prefix) + ".cnx.pwconv2.weight"), h);
|
||
ggml_tensor * pw2_b = get(std::string(prefix) + ".cnx.pwconv2.bias");
|
||
h = ggml_add(ctx, h, ggml_reshape_2d(ctx, pw2_b, pw2_b->ne[0], 1));
|
||
h = ggml_mul(ctx, h, get(std::string(prefix) + ".cnx.gamma"));
|
||
x = ggml_add(ctx, residual, h); // [C, T*2]
|
||
}
|
||
|
||
struct ggml_cgraph * gf = ggml_new_graph_custom(ctx, 4096, false);
|
||
ggml_build_forward_expand(gf, x);
|
||
int n_threads = 8;
|
||
if (const char * t_env = std::getenv("GGML_NUM_THREADS")) {
|
||
int n = std::atoi(t_env); if (n >= 1) n_threads = n;
|
||
}
|
||
ggml_graph_compute_with_ctx(ctx, gf, n_threads);
|
||
|
||
std::vector<float> result(C * T_final);
|
||
memcpy(result.data(), x->data, C * T_final * sizeof(float));
|
||
|
||
ggml_free(ctx);
|
||
return result;
|
||
}
|
||
|
||
// ----- Stage 5 : BigVGAN (initial Conv k=7 + 4 blocks + final Snake + final Conv k=7) -----
|
||
// Reproduit test_decoder. Input : us_out[c + C*t] (CT layout, T = T_in * 4).
|
||
// Output : wav[s] for s=0..T_in*1920-1.
|
||
static std::vector<float> stage_bigvgan(Decoder & d, const std::vector<float> & us_out, int T_in_us) {
|
||
const int latent = d.cfg.latent_dim; // 1024 (= C of upsample output)
|
||
const int decoder_dim = d.cfg.decoder_dim; // 1536
|
||
const std::vector<int> & up_rates = d.cfg.upsample_rates;
|
||
|
||
size_t mem_size = 12ULL * 1024 * 1024 * 1024;
|
||
struct ggml_init_params p = { mem_size, nullptr, false };
|
||
struct ggml_context * ctx = ggml_init(p);
|
||
|
||
auto get = [&](const std::string & key) { return d.tensors[key]; };
|
||
|
||
const float SNAKE_EPS = 1e-9f;
|
||
auto precompute_snake = [&](const std::string & ka, const std::string & kb,
|
||
ggml_tensor ** oa, ggml_tensor ** ob) -> bool {
|
||
ggml_tensor * a = d.tensors[ka];
|
||
ggml_tensor * b = d.tensors[kb];
|
||
if (!a || !b) return false;
|
||
const int C = (int)a->ne[0];
|
||
ggml_tensor * a_eff = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, C);
|
||
ggml_tensor * inv_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, C);
|
||
const float * pa = (const float*) a->data;
|
||
const float * pb = (const float*) b->data;
|
||
float * da = (float*) a_eff->data;
|
||
float * db = (float*) inv_b->data;
|
||
for (int c = 0; c < C; c++) {
|
||
da[c] = std::exp(pa[c]);
|
||
db[c] = 1.0f / (std::exp(pb[c]) + SNAKE_EPS);
|
||
}
|
||
*oa = a_eff; *ob = inv_b;
|
||
return true;
|
||
};
|
||
|
||
auto apply_snake_TC = [&](ggml_tensor * x_TC, ggml_tensor * a_eff, ggml_tensor * inv_b) {
|
||
const int C = (int)a_eff->ne[0];
|
||
ggml_tensor * a2 = ggml_reshape_2d(ctx, a_eff, 1, C);
|
||
ggml_tensor * b2 = ggml_reshape_2d(ctx, inv_b, 1, C);
|
||
ggml_tensor * xa = ggml_mul(ctx, x_TC, a2);
|
||
ggml_tensor * s = ggml_sin(ctx, xa);
|
||
ggml_tensor * s2 = ggml_mul(ctx, s, s);
|
||
ggml_tensor * mod = ggml_mul(ctx, s2, b2);
|
||
return ggml_add(ctx, x_TC, mod);
|
||
};
|
||
|
||
auto causal_conv1d_TC = [&](ggml_tensor * x_TC, const std::string & w_key, const std::string & b_key,
|
||
int K, int dilation) -> ggml_tensor * {
|
||
ggml_tensor * w = d.tensors[w_key];
|
||
ggml_tensor * b = d.tensors[b_key];
|
||
const int C_out = (int)w->ne[2];
|
||
const int T_cur = (int)x_TC->ne[0];
|
||
const int C_in = (int)x_TC->ne[1];
|
||
ggml_tensor * x_pad;
|
||
if (K == 1) {
|
||
x_pad = x_TC;
|
||
} else {
|
||
const int pad_left = (K - 1) * dilation;
|
||
ggml_tensor * zeros = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, pad_left, C_in);
|
||
memset(zeros->data, 0, pad_left * C_in * sizeof(float));
|
||
x_pad = ggml_concat(ctx, zeros, x_TC, 0);
|
||
}
|
||
ggml_tensor * y = ggml_conv_1d(ctx, w, x_pad, 1, 0, dilation);
|
||
const int T_out_y = (int)y->ne[0];
|
||
if (T_out_y > T_cur) {
|
||
y = ggml_view_2d(ctx, y, T_cur, C_out, y->nb[1], 0);
|
||
}
|
||
ggml_tensor * b2 = ggml_reshape_2d(ctx, b, 1, C_out);
|
||
return ggml_add(ctx, y, b2);
|
||
};
|
||
|
||
auto causal_convtr1d_TC = [&](ggml_tensor * x_TC, const std::string & w_key, const std::string & b_key,
|
||
int s) -> ggml_tensor * {
|
||
ggml_tensor * w = d.tensors[w_key];
|
||
ggml_tensor * b = d.tensors[b_key];
|
||
const int C_out = (int)w->ne[1];
|
||
const int T_cur = (int)x_TC->ne[0];
|
||
ggml_tensor * y = ggml_conv_transpose_1d(ctx, w, x_TC, s, 0, 1);
|
||
const int T_full = (int)y->ne[0];
|
||
const int T_target = T_cur * s;
|
||
if (T_full > T_target) {
|
||
y = ggml_view_2d(ctx, y, T_target, C_out, y->nb[1], 0);
|
||
}
|
||
ggml_tensor * b2 = ggml_reshape_2d(ctx, b, 1, C_out);
|
||
return ggml_add(ctx, y, b2);
|
||
};
|
||
|
||
// Input : us_out [C=latent, T] (CT layout host) → ggml ne=[C, T] memcpy.
|
||
ggml_tensor * x = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, latent, T_in_us);
|
||
memcpy(x->data, us_out.data(), latent * T_in_us * sizeof(float));
|
||
|
||
// Switch to TC layout for whole BigVGAN (Optim #3)
|
||
ggml_tensor * x_TC = ggml_cont(ctx, ggml_transpose(ctx, x)); // [T, latent]
|
||
|
||
// initial Conv k=7
|
||
x_TC = causal_conv1d_TC(x_TC, "decoder.initial.weight", "decoder.initial.bias", 7, 1);
|
||
|
||
// 4 DecoderBlocks
|
||
for (int B = 0; B < 4; B++) {
|
||
char prefix[64]; snprintf(prefix, sizeof(prefix), "decoder.block.%d", B);
|
||
const int up = up_rates[B];
|
||
|
||
ggml_tensor *a0 = nullptr, *b0 = nullptr;
|
||
precompute_snake(std::string(prefix) + ".block.0.alpha",
|
||
std::string(prefix) + ".block.0.beta", &a0, &b0);
|
||
x_TC = apply_snake_TC(x_TC, a0, b0);
|
||
x_TC = causal_convtr1d_TC(x_TC,
|
||
std::string(prefix) + ".block.1.conv.weight",
|
||
std::string(prefix) + ".block.1.conv.bias", up);
|
||
|
||
const int dilations[3] = {1, 3, 9};
|
||
for (int R = 0; R < 3; R++) {
|
||
char rprefix[80]; snprintf(rprefix, sizeof(rprefix), "%s.block.%d", prefix, 2 + R);
|
||
ggml_tensor * res = x_TC;
|
||
ggml_tensor *a1 = nullptr, *b1 = nullptr;
|
||
precompute_snake(std::string(rprefix) + ".act1.alpha",
|
||
std::string(rprefix) + ".act1.beta", &a1, &b1);
|
||
x_TC = apply_snake_TC(x_TC, a1, b1);
|
||
x_TC = causal_conv1d_TC(x_TC, std::string(rprefix) + ".conv1.conv.weight",
|
||
std::string(rprefix) + ".conv1.conv.bias", 7, dilations[R]);
|
||
ggml_tensor *a2 = nullptr, *b2 = nullptr;
|
||
precompute_snake(std::string(rprefix) + ".act2.alpha",
|
||
std::string(rprefix) + ".act2.beta", &a2, &b2);
|
||
x_TC = apply_snake_TC(x_TC, a2, b2);
|
||
x_TC = causal_conv1d_TC(x_TC, std::string(rprefix) + ".conv2.conv.weight",
|
||
std::string(rprefix) + ".conv2.conv.bias", 1, 1);
|
||
x_TC = ggml_add(ctx, res, x_TC);
|
||
}
|
||
}
|
||
|
||
// final SnakeBeta + Conv k=7
|
||
ggml_tensor *af = nullptr, *bf = nullptr;
|
||
precompute_snake("decoder.final_act.alpha", "decoder.final_act.beta", &af, &bf);
|
||
x_TC = apply_snake_TC(x_TC, af, bf);
|
||
x_TC = causal_conv1d_TC(x_TC, "decoder.final_conv.weight", "decoder.final_conv.bias", 7, 1);
|
||
|
||
struct ggml_cgraph * gf = ggml_new_graph_custom(ctx, 16384, false);
|
||
ggml_build_forward_expand(gf, x_TC);
|
||
int n_threads = 8;
|
||
if (const char * t_env = std::getenv("GGML_NUM_THREADS")) {
|
||
int n = std::atoi(t_env); if (n >= 1) n_threads = n;
|
||
}
|
||
ggml_graph_compute_with_ctx(ctx, gf, n_threads);
|
||
|
||
// x_TC ne=[T_final, 1] — single channel audio
|
||
const int T_final = (int)x_TC->ne[0];
|
||
std::vector<float> wav(T_final);
|
||
const float * data = (const float*) x_TC->data;
|
||
for (int t = 0; t < T_final; t++) {
|
||
float v = data[t];
|
||
if (v > 1.0f) v = 1.0f;
|
||
if (v < -1.0f) v = -1.0f;
|
||
wav[t] = v;
|
||
}
|
||
|
||
ggml_free(ctx);
|
||
(void)decoder_dim;
|
||
return wav;
|
||
}
|
||
|
||
// ============== Variante sched (Phase 2b BigVGAN HMX) ==============
|
||
// stage_bigvgan_sched : pattern sched-friendly pour activer le routage MUL_MAT
|
||
// vers HTP HMX. Diffs avec stage_bigvgan :
|
||
// - ggml_init no_alloc=true + mem_size réduit (~4 MB)
|
||
// - ggml_set_input/output sur entrée/sortie
|
||
// - causal_conv1d_TC : ggml_pad_ext au lieu de memset zeros + ggml_concat
|
||
// - precompute_snake : lookup dans d.snake_consts (précalculé au load)
|
||
// - exécution : sched_reset + sched_alloc_graph + tensor_set + sched_compute + tensor_get
|
||
// Le sched split automatiquement : MUL_MAT/ADD/MUL/PAD -> HTP HMX,
|
||
// IM2COL/SIN/CONV_TRANSPOSE_1D -> CPU. Copies CPU<->HTP gérées par sched.
|
||
static std::vector<float> stage_bigvgan_sched(Decoder & d, const std::vector<float> & us_out, int T_in_us) {
|
||
const int latent = d.cfg.latent_dim; // 1024
|
||
const int decoder_dim = d.cfg.decoder_dim; // 1536
|
||
const std::vector<int> & up_rates = d.cfg.upsample_rates;
|
||
|
||
// mem_size : juste métadonnées tensors (no_alloc=true). BigVGAN a ~30 snake + 26 conv1d
|
||
// + 4 convtr + adds → ~1000 nœuds max. ~256 octets/nœud → ~256 KB ; on prend 16 MB safe.
|
||
size_t mem_size = 16ULL * 1024 * 1024;
|
||
struct ggml_init_params p = { mem_size, nullptr, /*no_alloc=*/true };
|
||
struct ggml_context * ctx = ggml_init(p);
|
||
|
||
auto get = [&](const std::string & key) { return d.tensors[key]; };
|
||
|
||
auto apply_snake_TC = [&](ggml_tensor * x_TC, const std::string & prefix) -> ggml_tensor * {
|
||
auto it = d.snake_consts.find(prefix);
|
||
if (it == d.snake_consts.end()) {
|
||
fprintf(stderr, "bigvgan_sched: snake const not found for %s\n", prefix.c_str());
|
||
return x_TC;
|
||
}
|
||
ggml_tensor * a_eff = it->second.first; // [C], précalculé exp(α)
|
||
ggml_tensor * inv_b = it->second.second; // [C], précalculé 1/(exp(β)+eps)
|
||
const int C = (int)a_eff->ne[0];
|
||
ggml_tensor * a2 = ggml_reshape_2d(ctx, a_eff, 1, C);
|
||
ggml_tensor * b2 = ggml_reshape_2d(ctx, inv_b, 1, C);
|
||
ggml_tensor * xa = ggml_mul(ctx, x_TC, a2);
|
||
ggml_tensor * s = ggml_sin(ctx, xa); // SIN -> CPU fallback via sched
|
||
ggml_tensor * s2 = ggml_mul(ctx, s, s);
|
||
ggml_tensor * mod = ggml_mul(ctx, s2, b2);
|
||
return ggml_add(ctx, x_TC, mod);
|
||
};
|
||
|
||
auto causal_conv1d_TC = [&](ggml_tensor * x_TC, const std::string & w_key, const std::string & b_key,
|
||
int K, int dilation) -> ggml_tensor * {
|
||
ggml_tensor * w = get(w_key);
|
||
ggml_tensor * b = get(b_key);
|
||
const int C_out = (int)w->ne[2];
|
||
const int T_cur = (int)x_TC->ne[0];
|
||
ggml_tensor * x_pad;
|
||
if (K == 1) {
|
||
x_pad = x_TC;
|
||
} else {
|
||
const int pad_left = (K - 1) * dilation;
|
||
// Pad gauche sur dim 0 (T) par pad_left zéros — équivalent à concat(zeros, x_TC)
|
||
// mais en pur graphe (compatible no_alloc=true).
|
||
x_pad = ggml_pad_ext(ctx, x_TC, pad_left, 0, 0, 0, 0, 0, 0, 0);
|
||
}
|
||
ggml_tensor * y = ggml_conv_1d(ctx, w, x_pad, 1, 0, dilation);
|
||
const int T_out_y = (int)y->ne[0];
|
||
if (T_out_y > T_cur) {
|
||
y = ggml_view_2d(ctx, y, T_cur, C_out, y->nb[1], 0);
|
||
}
|
||
ggml_tensor * b2 = ggml_reshape_2d(ctx, b, 1, C_out);
|
||
return ggml_add(ctx, y, b2);
|
||
};
|
||
|
||
auto causal_convtr1d_TC = [&](ggml_tensor * x_TC, const std::string & w_key, const std::string & b_key,
|
||
int s) -> ggml_tensor * {
|
||
ggml_tensor * w = get(w_key);
|
||
ggml_tensor * b = get(b_key);
|
||
const int C_out = (int)w->ne[1];
|
||
const int T_cur = (int)x_TC->ne[0];
|
||
ggml_tensor * y = ggml_conv_transpose_1d(ctx, w, x_TC, s, 0, 1);
|
||
const int T_full = (int)y->ne[0];
|
||
const int T_target = T_cur * s;
|
||
if (T_full > T_target) {
|
||
y = ggml_view_2d(ctx, y, T_target, C_out, y->nb[1], 0);
|
||
}
|
||
ggml_tensor * b2 = ggml_reshape_2d(ctx, b, 1, C_out);
|
||
return ggml_add(ctx, y, b2);
|
||
};
|
||
|
||
// --- Input ---
|
||
ggml_tensor * x = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, latent, T_in_us);
|
||
ggml_set_name(x, "bigvgan_input");
|
||
ggml_set_input(x); // CRITIQUE : marque comme input pour sched_alloc_graph
|
||
|
||
// Switch to TC layout
|
||
ggml_tensor * x_TC = ggml_cont(ctx, ggml_transpose(ctx, x));
|
||
|
||
// initial Conv k=7
|
||
x_TC = causal_conv1d_TC(x_TC, "decoder.initial.weight", "decoder.initial.bias", 7, 1);
|
||
|
||
// 4 DecoderBlocks
|
||
for (int B = 0; B < 4; B++) {
|
||
char prefix[64]; snprintf(prefix, sizeof(prefix), "decoder.block.%d", B);
|
||
const int up = up_rates[B];
|
||
|
||
x_TC = apply_snake_TC(x_TC, std::string(prefix) + ".block.0");
|
||
x_TC = causal_convtr1d_TC(x_TC,
|
||
std::string(prefix) + ".block.1.conv.weight",
|
||
std::string(prefix) + ".block.1.conv.bias", up);
|
||
|
||
const int dilations[3] = {1, 3, 9};
|
||
for (int R = 0; R < 3; R++) {
|
||
char rprefix[80]; snprintf(rprefix, sizeof(rprefix), "%s.block.%d", prefix, 2 + R);
|
||
ggml_tensor * res = x_TC;
|
||
x_TC = apply_snake_TC(x_TC, std::string(rprefix) + ".act1");
|
||
x_TC = causal_conv1d_TC(x_TC, std::string(rprefix) + ".conv1.conv.weight",
|
||
std::string(rprefix) + ".conv1.conv.bias", 7, dilations[R]);
|
||
x_TC = apply_snake_TC(x_TC, std::string(rprefix) + ".act2");
|
||
x_TC = causal_conv1d_TC(x_TC, std::string(rprefix) + ".conv2.conv.weight",
|
||
std::string(rprefix) + ".conv2.conv.bias", 1, 1);
|
||
x_TC = ggml_add(ctx, res, x_TC);
|
||
}
|
||
}
|
||
|
||
// final SnakeBeta + Conv k=7
|
||
x_TC = apply_snake_TC(x_TC, "decoder.final_act");
|
||
x_TC = causal_conv1d_TC(x_TC, "decoder.final_conv.weight", "decoder.final_conv.bias", 7, 1);
|
||
|
||
ggml_set_name(x_TC, "bigvgan_output");
|
||
ggml_set_output(x_TC);
|
||
|
||
struct ggml_cgraph * gf = ggml_new_graph_custom(ctx, 16384, false);
|
||
ggml_build_forward_expand(gf, x_TC);
|
||
|
||
// --- Sched execution ---
|
||
ggml_backend_sched_reset(d.sched);
|
||
if (!ggml_backend_sched_alloc_graph(d.sched, gf)) {
|
||
fprintf(stderr, "bigvgan_sched: sched_alloc_graph FAIL\n");
|
||
ggml_free(ctx); return {};
|
||
}
|
||
// Set input data AFTER alloc (le tensor a maintenant un backend buffer)
|
||
ggml_backend_tensor_set(x, us_out.data(), 0, (size_t)latent * T_in_us * sizeof(float));
|
||
if (ggml_backend_sched_graph_compute(d.sched, gf) != GGML_STATUS_SUCCESS) {
|
||
fprintf(stderr, "bigvgan_sched: sched_graph_compute FAIL\n");
|
||
ggml_free(ctx); return {};
|
||
}
|
||
// Get output
|
||
const int T_final = (int)x_TC->ne[0];
|
||
std::vector<float> wav(T_final);
|
||
ggml_backend_tensor_get(x_TC, wav.data(), 0, (size_t)T_final * sizeof(float));
|
||
// Clamp [-1, 1]
|
||
for (int t = 0; t < T_final; t++) {
|
||
if (wav[t] > 1.0f) wav[t] = 1.0f;
|
||
if (wav[t] < -1.0f) wav[t] = -1.0f;
|
||
}
|
||
|
||
ggml_free(ctx);
|
||
(void)decoder_dim;
|
||
return wav;
|
||
}
|
||
|
||
// ============== Forward complet : chaîne staging strict ==============
|
||
|
||
// Dump an array of floats to disk if env var DECODER_DEBUG_DUMP_DIR is set.
|
||
static void debug_dump(const std::string & name, const std::vector<float> & data) {
|
||
const char * dir = std::getenv("DECODER_DEBUG_DUMP_DIR");
|
||
if (!dir) return;
|
||
std::string path = std::string(dir) + "/" + name + ".bin";
|
||
FILE * f = fopen(path.c_str(), "wb");
|
||
if (!f) {
|
||
fprintf(stderr, "debug_dump: cannot open %s\n", path.c_str());
|
||
return;
|
||
}
|
||
fwrite(data.data(), sizeof(float), data.size(), f);
|
||
fclose(f);
|
||
fprintf(stderr, "debug dump: %s (%zu floats)\n", path.c_str(), data.size());
|
||
}
|
||
|
||
std::vector<float> Decoder::forward(const std::vector<int32_t>& codes_flat, int T) {
|
||
assert((int)codes_flat.size() == cfg.num_quantizers * T);
|
||
// KZTTS_DECODER_PROFILE=1 -> imprime le temps pris par chaque stage.
|
||
const bool prof = (std::getenv("KZTTS_DECODER_PROFILE") &&
|
||
std::atoi(std::getenv("KZTTS_DECODER_PROFILE")) != 0);
|
||
auto nows = []() {
|
||
return std::chrono::duration<double>(
|
||
std::chrono::steady_clock::now().time_since_epoch()).count();
|
||
};
|
||
double t0 = nows(), t1, t2, t3, t4, t5;
|
||
|
||
auto hidden = stage_quantizer(*this, codes_flat, T); // [512, T]
|
||
t1 = nows();
|
||
debug_dump("stage1_quantizer", hidden);
|
||
auto pc_out = stage_pre_conv(*this, hidden, T); // [1024, T]
|
||
t2 = nows();
|
||
debug_dump("stage2_pre_conv", pc_out);
|
||
auto pt_out = stage_pre_transformer(*this, pc_out, T); // [1024, T]
|
||
t3 = nows();
|
||
debug_dump("stage3_pre_transformer", pt_out);
|
||
auto us_out = stage_upsample(*this, pt_out, T); // [1024, T*4]
|
||
t4 = nows();
|
||
debug_dump("stage4_upsample", us_out);
|
||
// BigVGAN : si sched actif (KZTTS_DECODER_SCHED=1), variante sched-friendly
|
||
// qui route les MUL_MAT sur HTP HMX. Sinon, path legacy CPU pur.
|
||
std::vector<float> wav = sched
|
||
? stage_bigvgan_sched(*this, us_out, T * 4)
|
||
: stage_bigvgan (*this, us_out, T * 4);
|
||
t5 = nows();
|
||
debug_dump("stage5_wav", wav);
|
||
|
||
if (prof) {
|
||
fprintf(stderr,
|
||
"decoder: T=%d total=%.3fs quant=%.3f preconv=%.3f pretr=%.3f"
|
||
" upsample=%.3f bigvgan=%.3f\n",
|
||
T, t5 - t0, t1 - t0, t2 - t1, t3 - t2, t4 - t3, t5 - t4);
|
||
}
|
||
return wav;
|
||
}
|
||
|
||
// Stub legacy
|
||
std::vector<float> Decoder::test_quantizer(const std::vector<int32_t>& codes_flat, int T) {
|
||
(void)codes_flat; (void)T;
|
||
return {};
|
||
}
|
||
|
||
} // namespace kazeia::tts
|