enne2/qwentts.cpp
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix speaker encoder ECAPA forward cossim 0.86 -> 0.996

Browse Source 
mel-spk and mel-mag dumps in speaker-encoder-extract.h applied an extra
ggml_transpose plus cont before write. Raw ggml ne=(C, T) already
streams as numpy [T, C], so the transpose was inverting axes vs the
python upstream. Removed it.
MelMag 0.04 -> 0.999, MelSpk 0.92 -> 0.998

spk_conv1d_same passed ggml_im2col a kernel ne=(K, 1, IC, 1) and an
input ne=(T_pad, 1, IC, 1) with IC in ne[2]. But the im2col impl reads
IC = b->ne[1] when is_2D=false, so it saw IC=1, wrote OW*K floats into
a buffer declared for OW*IC*K floats, and mul_mat consumed 99% garbage.
Moved IC into ne[1] for both kernel and input, which makes the impl
read the real IC and writes a buffer coherent with the declared ne. The
permute and the retranspose after pad become unnecessary, dropped both.
SpkFrontend 0.74 -> 0.994, SpeakerEmb 0.86 -> 0.996

Adds ECAPA bisection infrastructure : 4 stage out params in
speaker_encoder_forward (frontend, block3, mfa, asp), codec encoder
intermediate dumps in pipeline-codec.cpp (seanet-out, enc-transformer
out, codec-pre-fsq), matching pytorch hooks in debug-clone-cossim.py.
This commit is contained in:
Pascal
2026-05-10 20:56:14 +02:00
parent acb75fca36
commit 7e89929a70
7 changed files with 431 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/audio-mel.h
+8 -1
View File
@@ -157,6 +157,9 @@ static void audio_mel_compute_constants(const AudioMelConfig & cfg, AudioMelCons
// dft_real [n_fft, n_freq] f32, host constant
// dft_imag [n_fft, n_freq] f32, host constant
// mel_basis [n_freq, n_mels] f32, host constant
// mag_out optional pointer that receives the post-STFT magnitude
// tensor [n_freq, T_frames] for debug bisection. NULL by
// default. Caller marks it as graph output if needed.
//
// Output : [n_mels, T_frames] f32 log mel.
//
@@ -169,7 +172,8 @@ static struct ggml_tensor * audio_mel_build_graph(struct ggml_context * ctx,
struct ggml_tensor * dft_real,
struct ggml_tensor * dft_imag,
struct ggml_tensor * mel_basis,
const AudioMelConfig & cfg) {
const AudioMelConfig & cfg,
struct ggml_tensor ** mag_out = NULL) {
const int n_fft = cfg.n_fft;
const int hop = cfg.hop;
@@ -208,6 +212,9 @@ static struct ggml_tensor * audio_mel_build_graph(struct ggml_context * ctx,
struct ggml_tensor * mag2 = ggml_add(ctx, ggml_sqr(ctx, spec_re), ggml_sqr(ctx, spec_im));
mag2 = ggml_scale_bias(ctx, mag2, 1.0f, 1e-9f);
struct ggml_tensor * mag = ggml_sqrt(ctx, mag2);
if (mag_out) {
*mag_out = mag;
}
// mel_basis [n_freq, n_mels] @ mag [n_freq, T_frames] -> [n_mels, T_frames].
struct ggml_tensor * mel = ggml_mul_mat(ctx, mel_basis, mag);
src/pipeline-codec.cpp
+28 -19
View File
@@ -9,6 +9,7 @@
#include "pipeline-codec.h"
#include "causal-trans-conv.h"
#include "debug.h"
#include "qt-error.h"
#include <cmath>
@@ -219,7 +220,8 @@ std::vector<float> pipeline_codec_decode(PipelineCodec * pc, const int32_t * cod
return audio;
}
std::vector<int32_t> pipeline_codec_encode(PipelineCodec * pc, const float * audio, int n_samples) {
std::vector<int32_t> pipeline_codec_encode(PipelineCodec * pc, const float * audio, int n_samples,
const char * dump_dir) {
if (n_samples <= 0 || (n_samples % QWEN_TOKENIZER_HOP_LENGTH) != 0) {
qt_log(QT_LOG_ERROR, "[Pipeline] n_samples must be a positive multiple of %d (got %d)",
QWEN_TOKENIZER_HOP_LENGTH, n_samples);
@@ -276,10 +278,16 @@ std::vector<int32_t> pipeline_codec_encode(PipelineCodec * pc, const float * aud
// Transpose to get the buffer layout we want once read back to host.
h = ggml_cont(gctx, ggml_transpose(gctx, h)); // ne=(512, T)
const char * dump_dir = getenv("QWENTTS_DEBUG_DUMP");
if (dump_dir) {
ggml_set_output(h_seanet);
ggml_set_name(h_seanet, "seanet_out");
const char * dump = dump_dir;
struct ggml_tensor * h_seanet_dump = NULL;
if (dump) {
// SEANet output naturally lands as channel-first ggml ne=(T, hidden).
// The encoder_transformer and downsample dumps further down are
// T-first numpy [T, hidden], so we transpose the SEANet view to
// match before pinning it as a graph output.
h_seanet_dump = ggml_cont(gctx, ggml_transpose(gctx, h_seanet));
ggml_set_output(h_seanet_dump);
ggml_set_name(h_seanet_dump, "seanet_out_dump");
ggml_set_output(h_et);
ggml_set_name(h_et, "enc_transformer_out");
}
@@ -289,6 +297,9 @@ std::vector<int32_t> pipeline_codec_encode(PipelineCodec * pc, const float * aud
struct ggml_cgraph * graph = ggml_new_graph_custom(gctx, n_max_nodes, false);
ggml_build_forward_expand(graph, h);
if (h_seanet_dump) {
ggml_build_forward_expand(graph, h_seanet_dump);
}
if (!ggml_backend_sched_alloc_graph(pc->sched, graph)) {
qt_log(QT_LOG_ERROR, "[Pipeline] encode sched_alloc_graph failed");
@@ -315,24 +326,22 @@ std::vector<int32_t> pipeline_codec_encode(PipelineCodec * pc, const float * aud
return {};
}
if (dump_dir) {
auto dump = [&](const char * fname, struct ggml_tensor * t) {
if (dump) {
DebugDumper d;
debug_init(&d, dump);
// ggml ne layout matches numpy's last-dim-fastest, so a [d0, d1]
// tensor in ggml dumps as a [d1, d0] numpy array. We emit the
// shape ggml-side (ne[1], ne[0]) so numpy reshapes it correctly
// on read. Values themselves are the same memory order.
auto dump2 = [&](const char * name, struct ggml_tensor * t) {
size_t n = ggml_nelements(t);
std::vector<float> buf(n);
ggml_backend_tensor_get(t, buf.data(), 0, n * sizeof(float));
char path[512];
snprintf(path, sizeof(path), "%s/%s.f32", dump_dir, fname);
FILE * f = fopen(path, "wb");
if (f) {
fwrite(buf.data(), sizeof(float), n, f);
fclose(f);
qt_log(QT_LOG_INFO, "[Pipeline] Dumped %s: %zu floats, ne=(%lld, %lld, %lld, %lld)", path, n,
(long long) t->ne[0], (long long) t->ne[1], (long long) t->ne[2], (long long) t->ne[3]);
}
debug_dump_2d(&d, name, buf.data(), (int) t->ne[1], (int) t->ne[0]);
};
dump("seanet_out", h_seanet);
dump("enc_transformer_out", h_et);
dump("enc_downsample_out", h);
dump2("seanet-out", h_seanet_dump);
dump2("enc-transformer-out", h_et);
dump2("codec-pre-fsq", h);
}
// Read back the post-downsample hidden buffer for CPU-side RVQ encode.
src/pipeline-codec.h
+9 -4
View File
@@ -86,12 +86,17 @@ bool pipeline_codec_load(PipelineCodec * pc, const char * gguf_path, BackendPair
std::vector<float> pipeline_codec_decode(PipelineCodec * pc, const int32_t * codes, int K, int T);
// Encode a 24 kHz mono waveform into RVQ codes.
// audio : [n_samples] f32 mono 24 kHz. Must be a multiple of
// QWEN_TOKENIZER_HOP_LENGTH (1920); the caller is expected
// to pad with zeros if needed.
// audio : [n_samples] f32 mono 24 kHz. Must be a multiple of
// QWEN_TOKENIZER_HOP_LENGTH (1920); the caller is expected
// to pad with zeros if needed.
// dump_dir : optional path. When non NULL, dumps the SEANet, encoder
// transformer and post-downsample (pre-FSQ latents) buffers
// into seanet-out.bin, enc-transformer-out.bin and
// codec-pre-fsq.bin under that directory. Quiet otherwise.
// Returns codes flat as [K, T] row-major, K = QWEN_TOKENIZER_NUM_CODEBOOKS,
// T = n_samples / 1920. Empty on failure.
std::vector<int32_t> pipeline_codec_encode(PipelineCodec * pc, const float * audio, int n_samples);
std::vector<int32_t> pipeline_codec_encode(PipelineCodec * pc, const float * audio, int n_samples,
const char * dump_dir = NULL);
// Free every backend buffer and ggml context. Safe to call on a zeroed struct.
void pipeline_codec_free(PipelineCodec * pc);
src/pipeline-tts.cpp
+2 -2
View File
@@ -256,7 +256,7 @@ bool pipeline_tts_synthesize(PipelineTTS * pt,
"[Pipeline] FATAL: --ref-audio requires a model with a loaded speaker encoder (Base only)\n");
return false;
}
if (!speaker_encoder_extract(&pt->speaker_encoder, pt->sched, params.ref_audio, ref_spk_emb)) {
if (!speaker_encoder_extract(&pt->speaker_encoder, pt->sched, params.ref_audio, ref_spk_emb, params.dump_dir)) {
return false;
}
if ((int) ref_spk_emb.size() != pt->talker.hidden_size) {
@@ -292,7 +292,7 @@ bool pipeline_tts_synthesize(PipelineTTS * pt,
return false;
}
int aligned_T = (T_codec_audio / QWEN_TOKENIZER_HOP_LENGTH) * QWEN_TOKENIZER_HOP_LENGTH;
ref_codes = pipeline_codec_encode(&pt->codec, raw, aligned_T);
ref_codes = pipeline_codec_encode(&pt->codec, raw, aligned_T, params.dump_dir);
std::free(raw);
if (ref_codes.empty()) {
fprintf(stderr, "[Pipeline] FATAL: pipeline_codec_encode returned empty codes\n");
src/speaker-encoder-extract.h
+144 -3
View File
@@ -18,6 +18,7 @@
#include "audio-io.h"
#include "audio-mel.h"
#include "debug.h"
#include "ggml-alloc.h"
#include "ggml-backend.h"
#include "ggml.h"
@@ -30,11 +31,14 @@
#include <vector>
// Public entry point. Returns true on success, fills emb_out with the
// 2048-dim f32 embedding. Returns false on any IO or graph failure.
// 2048-dim f32 embedding. When dump_dir is non NULL, also writes the post
// mel_spectrogram tensor to mel-spk.bin under that directory using the
// debug.h header format. Quiet otherwise.
static bool speaker_encoder_extract(const SpeakerEncoderWeights * sw,
ggml_backend_sched_t sched,
const char * wav_path,
std::vector<float> & emb_out) {
std::vector<float> & emb_out,
const char * dump_dir = NULL) {
if (sw->weight_buf == NULL) {
fprintf(stderr, "[SpkExtract] FATAL: speaker encoder weights not loaded\n");
return false;
@@ -114,12 +118,91 @@ static bool speaker_encoder_extract(const SpeakerEncoderWeights * sw,
ggml_set_input(dft_im_in);
ggml_set_input(mel_b_in);
struct ggml_tensor * mel_t = NULL;
struct ggml_tensor * mel_dump = NULL;
struct ggml_tensor * mag_t = NULL;
struct ggml_tensor * mag_dump = NULL;
struct ggml_tensor * front_t = NULL;
struct ggml_tensor * front_dump = NULL;
struct ggml_tensor * blk3_t = NULL;
struct ggml_tensor * blk3_dump = NULL;
struct ggml_tensor * mfa_t = NULL;
struct ggml_tensor * mfa_dump = NULL;
struct ggml_tensor * asp_t = NULL;
struct ggml_tensor * asp_dump = NULL;
struct ggml_tensor * emb =
speaker_encoder_forward(gctx, sw, audio_in, hann_in, dft_re_in, dft_im_in, mel_b_in, mel_cfg);
speaker_encoder_forward(gctx, sw, audio_in, hann_in, dft_re_in, dft_im_in, mel_b_in, mel_cfg,
&mel_t, &mag_t, &front_t, &blk3_t, &mfa_t, &asp_t);
ggml_set_output(emb);
if (dump_dir && mel_t) {
// mel_t has ggml ne=(n_mels, T_frames), which streams row-major
// as numpy [T_frames, n_mels], the exact layout the upstream
// mel_spectrogram() exposes after its .transpose(1, 2). The
// cont call forces materialization in case the scheduler fuses
// ggml_log with anything downstream.
mel_dump = ggml_cont(gctx, mel_t);
ggml_set_output(mel_dump);
ggml_set_name(mel_dump, "spk.mel_dump");
}
if (dump_dir && mag_t) {
// mag_t has ggml ne=(n_freq, T_frames), which streams row-major
// as numpy [T_frames, n_freq], matching the upstream STFT
// magnitude after its transpose. Cont forces materialization.
mag_dump = ggml_cont(gctx, mag_t);
ggml_set_output(mag_dump);
ggml_set_name(mag_dump, "spk.mag_dump");
}
if (dump_dir && front_t) {
// front_t has ggml ne=(512, T_frames), reads row-major as numpy
// [T_frames, 512]. Same layout the Python forward sees after
// blocks[0] (TimeDelayNetBlock) when transposed (1, 2).
front_dump = ggml_cont(gctx, front_t);
ggml_set_output(front_dump);
ggml_set_name(front_dump, "spk.frontend_dump");
}
if (dump_dir && blk3_t) {
// blk3_t has ggml ne=(512, T_frames) = numpy [T_frames, 512],
// matching Python blocks[3] output transposed (1, 2).
blk3_dump = ggml_cont(gctx, blk3_t);
ggml_set_output(blk3_dump);
ggml_set_name(blk3_dump, "spk.block3_dump");
}
if (dump_dir && mfa_t) {
// mfa_t has ggml ne=(1536, T_frames) = numpy [T_frames, 1536],
// matching Python mfa output transposed (1, 2).
mfa_dump = ggml_cont(gctx, mfa_t);
ggml_set_output(mfa_dump);
ggml_set_name(mfa_dump, "spk.mfa_dump");
}
if (dump_dir && asp_t) {
// asp_t has ggml ne=(3072, 1), reads row-major as numpy [1, 3072].
// Python asp returns [B=1, 3072, 1] which we slice with [0].T to
// obtain [1, 3072] for a direct shape-aligned compare.
asp_dump = ggml_cont(gctx, asp_t);
ggml_set_output(asp_dump);
ggml_set_name(asp_dump, "spk.asp_dump");
}
struct ggml_cgraph * graph = ggml_new_graph_custom(gctx, 2048, false);
ggml_build_forward_expand(graph, emb);
if (mel_dump) {
ggml_build_forward_expand(graph, mel_dump);
}
if (mag_dump) {
ggml_build_forward_expand(graph, mag_dump);
}
if (front_dump) {
ggml_build_forward_expand(graph, front_dump);
}
if (blk3_dump) {
ggml_build_forward_expand(graph, blk3_dump);
}
if (mfa_dump) {
ggml_build_forward_expand(graph, mfa_dump);
}
if (asp_dump) {
ggml_build_forward_expand(graph, asp_dump);
}
// Reset the shared sched before allocating : the talker may have left
// a residual graph state from a previous synthesis call.
@@ -147,6 +230,64 @@ static bool speaker_encoder_extract(const SpeakerEncoderWeights * sw,
emb_out.assign((size_t) sw->enc_dim, 0.0f);
ggml_backend_tensor_get(emb, emb_out.data(), 0, (size_t) sw->enc_dim * sizeof(float));
if (dump_dir && mel_dump) {
size_t n = ggml_nelements(mel_dump);
std::vector<float> buf(n);
ggml_backend_tensor_get(mel_dump, buf.data(), 0, n * sizeof(float));
DebugDumper d;
debug_init(&d, dump_dir);
// mel_dump has ggml ne=(n_mels, T_frames). debug_dump_2d takes
// (rows, cols) in numpy convention so we pass (ne[1]=T_frames,
// ne[0]=n_mels), which writes shape [T_frames, n_mels] over the
// raw ggml memory layout, matching the Python side dump.
debug_dump_2d(&d, "mel-spk", buf.data(), (int) mel_dump->ne[1], (int) mel_dump->ne[0]);
// CPU side mel constants : audit against torch.hann_window and
// librosa.filters.mel produced by the Python upstream. Layouts
// are kept as numpy [n_fft] for hann and [n_mels, n_freq] for
// mel_basis, matching the librosa convention.
debug_dump_1d(&d, "mel-hann", mel_c.hann.data(), mel_cfg.n_fft);
debug_dump_2d(&d, "mel-basis", mel_c.mel_basis.data(), mel_cfg.n_mels, mel_c.n_freq);
if (mag_dump) {
size_t nm = ggml_nelements(mag_dump);
std::vector<float> bm(nm);
ggml_backend_tensor_get(mag_dump, bm.data(), 0, nm * sizeof(float));
// mag_dump has ggml ne=(n_freq, T_frames). Same dumping
// convention as mel-spk : passing (ne[1], ne[0]) writes
// shape [T_frames, n_freq] over the raw memory layout.
debug_dump_2d(&d, "mel-mag", bm.data(), (int) mag_dump->ne[1], (int) mag_dump->ne[0]);
}
// ECAPA forward bisection points. Each tensor has ggml ne=(C, T)
// and dumps as numpy [T, C] using the (ne[1], ne[0]) convention.
// ASP collapses T to 1 so it lands as [1, 3072].
if (front_dump) {
size_t nf = ggml_nelements(front_dump);
std::vector<float> bf(nf);
ggml_backend_tensor_get(front_dump, bf.data(), 0, nf * sizeof(float));
debug_dump_2d(&d, "spk-frontend", bf.data(), (int) front_dump->ne[1], (int) front_dump->ne[0]);
}
if (blk3_dump) {
size_t nb = ggml_nelements(blk3_dump);
std::vector<float> bb(nb);
ggml_backend_tensor_get(blk3_dump, bb.data(), 0, nb * sizeof(float));
debug_dump_2d(&d, "spk-block3", bb.data(), (int) blk3_dump->ne[1], (int) blk3_dump->ne[0]);
}
if (mfa_dump) {
size_t nf = ggml_nelements(mfa_dump);
std::vector<float> bf(nf);
ggml_backend_tensor_get(mfa_dump, bf.data(), 0, nf * sizeof(float));
debug_dump_2d(&d, "spk-mfa", bf.data(), (int) mfa_dump->ne[1], (int) mfa_dump->ne[0]);
}
if (asp_dump) {
size_t na = ggml_nelements(asp_dump);
std::vector<float> ba(na);
ggml_backend_tensor_get(asp_dump, ba.data(), 0, na * sizeof(float));
debug_dump_2d(&d, "spk-asp", ba.data(), (int) asp_dump->ne[1], (int) asp_dump->ne[0]);
}
}
ggml_backend_sched_reset(sched);
ggml_free(gctx);
src/speaker-encoder-forward.h
+54 -19
View File
@@ -49,35 +49,39 @@ static struct ggml_tensor * spk_conv1d_same(struct ggml_context * ctx,
struct ggml_tensor * b,
int dilation) {
const int K = (int) w->ne[0];
const int IC = (int) w->ne[1];
const int OC = (int) w->ne[2];
const int pad = ((K - 1) * dilation) / 2;
// ggml_pad_reflect_1d pads ne[0]. Our temporal axis is ne[1], so
// transpose first, pad, transpose back.
struct ggml_tensor * x_t = ggml_cont(ctx, ggml_transpose(ctx, x)); // [T, in_c]
// ggml_pad_reflect_1d pads the innermost axis ne[0]. Our temporal
// axis is ne[1], so we transpose to bring T to ne[0], pad, and keep
// it that way : the im2col downstream expects ne[0]=T_pad, ne[1]=IC,
// which is exactly the layout we end up with here.
struct ggml_tensor * x_t = ggml_cont(ctx, ggml_transpose(ctx, x)); // ne=(T, IC)
if (pad > 0) {
x_t = ggml_pad_reflect_1d(ctx, x_t, pad, pad); // [T+2*pad, in_c]
x_t = ggml_pad_reflect_1d(ctx, x_t, pad, pad); // ne=(T+2*pad, IC)
}
x_t = ggml_cont(ctx, ggml_transpose(ctx, x_t)); // [in_c, T+2*pad]
// Reshape as [W=T_pad, H=1, IC=in_c, N=1] for ggml_im2col 1D.
struct ggml_tensor * x4d = ggml_reshape_4d(ctx, x_t, x_t->ne[1], 1, x_t->ne[0], 1);
struct ggml_tensor * dummy = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, K, 1, x_t->ne[0], 1);
// Reshape to 4D for ggml_im2col 1D : ne=(T_pad, IC, 1, 1).
struct ggml_tensor * x4d = ggml_reshape_4d(ctx, x_t, x_t->ne[0], IC, 1, 1);
// Dummy F32 kernel with the (K, IC) shape ggml_im2col needs to read
// off both axes. Borrowing only the ne and not the real weight data
// avoids the impl's src0 type assert when w is quantized.
struct ggml_tensor * dummy = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, K, IC, 1, 1);
ggml_set_name(dummy, "spk.im2col_kernel");
// im2col output shape with is_2D=false : ne = (a.ne[1]*a.ne[0], OW, b.ne[2], 1)
// = (K, T_out, in_c, 1) here. To matmul against the [K*in_c, out_c]
// weight we need [K*in_c, T_out] which means permuting (K, IC, T)
// before flattening. ggml_permute(0, 2, 1, 3) swaps axes 1 and 2.
// im2col with is_2D=false : the constructor declares ne[0]=IC*K and
// ne[1]=OW, and the impl writes the buffer in (k inner, ic middle,
// t outer) order, which matches that ne directly. A reshape_2d to
// (IC*K, T_out) reads col_2d[ic*K + k, t], lining up with the
// weight reshape w_2d[ic*K + k, oc] for the mul_mat below.
struct ggml_tensor * col = ggml_im2col(ctx, dummy, x4d, 1, 1, 0, 0, dilation, 1, false, GGML_TYPE_F32);
int T_out = (int) col->ne[1];
int IC = (int) col->ne[2];
col = ggml_cont(ctx, ggml_permute(ctx, col, 0, 2, 1, 3)); // [K, IC, T, 1]
col = ggml_reshape_2d(ctx, col, K * IC, T_out);
// weight [K, in_c, out_c] reshape as [K * in_c, out_c]. mul_mat
// returns [out_c, T_out].
struct ggml_tensor * w2d = ggml_reshape_2d(ctx, w, K * (int) w->ne[1], OC);
// Weight reshape : [K, IC, OC] -> [K*IC, OC]. mul_mat returns [OC, T_out].
struct ggml_tensor * w2d = ggml_reshape_2d(ctx, w, K * IC, OC);
struct ggml_tensor * y = ggml_mul_mat(ctx, w2d, col);
ggml_mul_mat_set_prec(y, GGML_PREC_F32);
@@ -271,6 +275,16 @@ static struct ggml_tensor * spk_asp(struct ggml_context * ctx, const SpkEncASP &
// Inputs :
// audio_padded [T_pad] f32, host or backend tensor
// mel constants hann/dft_real/dft_imag/mel_basis backend tensors
// mel_out optional out param. When non NULL, receives the post
// mel_spectrogram tensor [n_mels, T_frames] so the caller
// can mark it as a graph output and pull its values back.
// mag_out optional out param. When non NULL, receives the post
// STFT magnitude tensor [n_freq, T_frames] for debug
// bisection between the STFT and the mel filtering.
// frontend_out optional. Post conv0 TDNN k=5 + ReLU output [512, T].
// block3_out optional. Post third SE-Res2Net block output [512, T].
// mfa_out optional. Post multi-layer feature aggregation [1536, T].
// asp_out optional. Post attentive statistical pooling [3072, 1].
// Output : [enc_dim] f32, the speaker embedding (typically 2048 dims).
static struct ggml_tensor * speaker_encoder_forward(struct ggml_context * ctx,
const SpeakerEncoderWeights * sw,
@@ -279,25 +293,46 @@ static struct ggml_tensor * speaker_encoder_forward(struct ggml_context *
struct ggml_tensor * dft_real,
struct ggml_tensor * dft_imag,
struct ggml_tensor * mel_basis,
const AudioMelConfig & mel_cfg) {
const AudioMelConfig & mel_cfg,
struct ggml_tensor ** mel_out = NULL,
struct ggml_tensor ** mag_out = NULL,
struct ggml_tensor ** frontend_out = NULL,
struct ggml_tensor ** block3_out = NULL,
struct ggml_tensor ** mfa_out = NULL,
struct ggml_tensor ** asp_out = NULL) {
// Mel : [n_mels=128, T_frames]
struct ggml_tensor * mel = audio_mel_build_graph(ctx, audio_padded, hann, dft_real, dft_imag, mel_basis, mel_cfg);
struct ggml_tensor * mel = audio_mel_build_graph(ctx, audio_padded, hann, dft_real, dft_imag, mel_basis, mel_cfg, mag_out);
if (mel_out) {
*mel_out = mel;
}
// Frontend conv0 TDNN k=5 + ReLU : 128 -> 512, T preserved.
struct ggml_tensor * h = spk_tdnn(ctx, sw->conv0, mel, 1);
if (frontend_out) {
*frontend_out = h;
}
// Three SE-Res2Net blocks at dilations 2, 3, 4.
struct ggml_tensor * b1 = spk_block(ctx, sw->blocks[0], h, sw->res2net_scale);
struct ggml_tensor * b2 = spk_block(ctx, sw->blocks[1], b1, sw->res2net_scale);
struct ggml_tensor * b3 = spk_block(ctx, sw->blocks[2], b2, sw->res2net_scale);
if (block3_out) {
*block3_out = b3;
}
// Multi-layer feature aggregation : cat blk1..3 then 1x1 TDNN + ReLU.
struct ggml_tensor * cat = ggml_concat(ctx, b1, b2, 0);
cat = ggml_concat(ctx, cat, b3, 0); // [1536, T]
struct ggml_tensor * mfa = spk_tdnn(ctx, sw->mfa, cat, 1); // [1536, T]
if (mfa_out) {
*mfa_out = mfa;
}
// Attentive statistical pooling : [1536, T] -> [3072, 1].
struct ggml_tensor * stats = spk_asp(ctx, sw->asp, mfa);
if (asp_out) {
*asp_out = stats;
}
// Final FC k=1 : [3072, 1] -> [enc_dim, 1].
struct ggml_tensor * emb = spk_conv1d_same(ctx, stats, sw->fc_w, sw->fc_b, 1);
tests/debug-clone-cossim.py
+186 -4
View File
@@ -43,11 +43,171 @@ DEFAULT_REF_TEXT = "../examples/freeman.txt"
# Mode B adds two pre-talker stages to the standard list : the speaker
# embedding extracted from the reference audio (ECAPA forward, projected to
# talker hidden), and the reference codec frames at 12.5 Hz.
# talker hidden), and the reference codec frames at 12.5 Hz. Plus three
# bisection stages for the 12Hz codec encoder (SEANet output, encoder
# transformer output, post-downsample = pre-FSQ latents), the mel front end
# (mel-mag and mel-spk), and four ECAPA forward bisection stages (frontend
# conv0 output, third SE-Res2Net block output, MFA output, ASP output).
STAGES_CLONE = cc.STAGES_STANDARD + [
("SpeakerEmb", "speaker-emb.bin"),
("MelHann", "mel-hann.bin"),
("MelBasis", "mel-basis.bin"),
("MelMag", "mel-mag.bin"),
("MelSpk", "mel-spk.bin"),
("SeanetOut", "seanet-out.bin"),
("EncTransformer", "enc-transformer-out.bin"),
("CodecPreFSQ", "codec-pre-fsq.bin"),
("SpkFrontend", "spk-frontend.bin"),
("SpkBlock3", "spk-block3.bin"),
("SpkMFA", "spk-mfa.bin"),
("SpkASP", "spk-asp.bin"),
("SpeakerEmb", "speaker-emb.bin"),
]
def install_clone_hooks(model, dump_dir):
"""Capture the codec encoder bisection points (SEANet, encoder_transformer,
downsample = pre-FSQ latents), the ECAPA mel front end input, and four
ECAPA forward bisection points (frontend conv0 output, third SE-Res2Net
block output, MFA output, ASP output). Mirrors exactly what
pipeline-codec.cpp and speaker-encoder-extract.h dump on the C++ side,
with matching shapes : [T, 512] for the codec stages, [T_frames, 128]
for the speaker mel, [T_frames, 512] for spk-frontend / spk-block3,
[T_frames, 1536] for spk-mfa, and [1, 3072] for spk-asp."""
enc = model.speech_tokenizer.model.encoder
seen_seanet = {"done": False}
def hook_seanet(module, args, output):
if seen_seanet["done"]:
return
out = output[0] if isinstance(output, tuple) else output
# output shape : [B=1, C=512, T_emb] channel-first from MimiEncoder.
cc.save_dump(os.path.join(dump_dir, "seanet-out.bin"), out[0].transpose(0, 1).contiguous())
seen_seanet["done"] = True
enc.encoder.register_forward_hook(hook_seanet)
seen_enct = {"done": False}
def hook_enct(module, args, output):
if seen_enct["done"]:
return
out = output[0] if isinstance(output, tuple) else output
# encoder_transformer is fed [B, T, 512] T-first and returns the
# same shape, so no transpose needed before the [0] slice.
cc.save_dump(os.path.join(dump_dir, "enc-transformer-out.bin"), out[0])
seen_enct["done"] = True
enc.encoder_transformer.register_forward_hook(hook_enct)
seen_down = {"done": False}
def hook_down(module, args, output):
if seen_down["done"]:
return
out = output[0] if isinstance(output, tuple) else output
# downsample output : [B=1, C=512, T] channel-first, transpose to
# [T, 512] to match the C++ post-downsample dump.
cc.save_dump(os.path.join(dump_dir, "codec-pre-fsq.bin"), out[0].transpose(0, 1).contiguous())
seen_down["done"] = True
enc.downsample.register_forward_hook(hook_down)
seen_mel = {"done": False}
def hook_spk_pre(module, args, kwargs):
if seen_mel["done"]:
return
# mels arrives as args[0] with shape [B=1, T_frames, n_mels=128]
# post the .transpose(1, 2) inside extract_speaker_embedding. The
# C++ side now dumps the same T-first layout, so we keep mels[0]
# as is to preserve [T_frames, n_mels].
mels = args[0] if args else kwargs.get("mels", None)
if mels is None or mels.dim() != 3:
return
cc.save_dump(os.path.join(dump_dir, "mel-spk.bin"), mels[0])
seen_mel["done"] = True
model.speaker_encoder.register_forward_pre_hook(hook_spk_pre, with_kwargs=True)
# ECAPA forward bisection. blocks[0] is the frontend TimeDelayNetBlock
# mapped to spk_tdnn(conv0) on the C++ side. blocks[3] is the third
# SE-Res2Net block, mapped to the C++ blocks[2] output. mfa and asp
# speak for themselves. All these modules ingest channel-first
# [B, C, T] tensors so we transpose to [T, C] before save_dump for a
# direct compare against the C++ ne=(C, T) raw memory dumps.
spk = model.speaker_encoder
seen_front = {"done": False}
def hook_frontend(module, args, output):
if seen_front["done"]:
return
out = output[0] if isinstance(output, tuple) else output
# output shape : [B=1, 512, T_frames] channel-first.
cc.save_dump(os.path.join(dump_dir, "spk-frontend.bin"), out[0].transpose(0, 1).contiguous())
seen_front["done"] = True
spk.blocks[0].register_forward_hook(hook_frontend)
seen_blk3 = {"done": False}
def hook_block3(module, args, output):
if seen_blk3["done"]:
return
out = output[0] if isinstance(output, tuple) else output
# output shape : [B=1, 512, T_frames] channel-first.
cc.save_dump(os.path.join(dump_dir, "spk-block3.bin"), out[0].transpose(0, 1).contiguous())
seen_blk3["done"] = True
spk.blocks[3].register_forward_hook(hook_block3)
seen_mfa = {"done": False}
def hook_mfa(module, args, output):
if seen_mfa["done"]:
return
out = output[0] if isinstance(output, tuple) else output
# output shape : [B=1, 1536, T_frames] channel-first.
cc.save_dump(os.path.join(dump_dir, "spk-mfa.bin"), out[0].transpose(0, 1).contiguous())
seen_mfa["done"] = True
spk.mfa.register_forward_hook(hook_mfa)
seen_asp = {"done": False}
def hook_asp(module, args, output):
if seen_asp["done"]:
return
out = output[0] if isinstance(output, tuple) else output
# output shape : [B=1, 3072, 1] from AttentiveStatisticsPooling.
# Transpose to [1, 3072] to match the C++ ne=(3072, 1) raw layout.
cc.save_dump(os.path.join(dump_dir, "spk-asp.bin"), out[0].transpose(0, 1).contiguous())
seen_asp["done"] = True
spk.asp.register_forward_hook(hook_asp)
def dump_mel_constants(dump_dir):
"""Reproduce the speaker encoder mel front end CPU constants the same
way the upstream mel_spectrogram() builds them (torch.hann_window for
the window and librosa.filters.mel for the Slaney filterbank), and
save them under dump_dir/mel-hann.bin and dump_dir/mel-basis.bin so
they can be paired with the C++ side dumps."""
import librosa
n_fft = 1024
n_mels = 128
sr = 24000
fmin = 0.0
fmax = 12000.0
hann = torch.hann_window(n_fft, periodic=True).numpy().astype(np.float32)
cc.save_dump(os.path.join(dump_dir, "mel-hann.bin"), hann)
mel_basis = librosa.filters.mel(sr=sr, n_fft=n_fft, n_mels=n_mels, fmin=fmin, fmax=fmax)
cc.save_dump(os.path.join(dump_dir, "mel-basis.bin"), mel_basis.astype(np.float32))
def dump_mel_mag_python(ref_wav, dump_dir):
"""Reproduce the upstream mel_spectrogram STFT path (same n_fft / hop /
window / pad as modeling_qwen3_tts.mel_spectrogram) and dump the post
magnitude tensor [T_frames, n_freq] for direct pairing with the C++
spk.mag_dump output. This isolates the STFT step from the mel filter."""
n_fft = 1024
hop = 256
win = 1024
padding = (n_fft - hop) // 2
y = torch.from_numpy(ref_wav).unsqueeze(0)
y = torch.nn.functional.pad(y.unsqueeze(1), (padding, padding), mode="reflect").squeeze(1)
spec = torch.stft(
y, n_fft, hop_length=hop, win_length=win,
window=torch.hann_window(win, periodic=True),
center=False, pad_mode="reflect", normalized=False,
onesided=True, return_complex=True,
)
mag = torch.sqrt(torch.view_as_real(spec).pow(2).sum(-1) + 1e-9)
# mag shape : [B=1, n_freq=513, T_frames]. Transpose to [T_frames, n_freq].
cc.save_dump(os.path.join(dump_dir, "mel-mag.bin"), mag[0].transpose(0, 1).contiguous())
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--prompt", default="../examples/prompt.txt")
@@ -70,6 +230,11 @@ def main():
cc.ensure_dir(DUMP_CPP)
os.makedirs(os.path.dirname(args.out_pt) or ".", exist_ok=True)
# Reproduce the upstream mel front end CPU constants (torch.hann_window
# + librosa.filters.mel) and dump them so they pair with the C++ side
# dumps emitted by speaker-encoder-extract.h.
dump_mel_constants(DUMP_PT)
with open(args.prompt, "r", encoding="utf-8") as f:
text = f.read().strip()
with open(args.ref_text_file, "r", encoding="utf-8") as f:
@@ -96,6 +261,11 @@ def main():
).eval()
processor = cc.AutoProcessor.from_pretrained(CKPT, fix_mistral_regex=True)
# Install codec encoder + ECAPA front end hooks before any encode call,
# so the freshly captured intermediates land in DUMP_PT/*.bin alongside
# the talker stages installed further down by cc.install_hooks.
install_clone_hooks(model, DUMP_PT)
# Load reference WAV. Resample to 24 kHz if needed since both the speaker
# encoder and the codec tokenizer expect 24 kHz mono input.
ref_wav, ref_sr = sf.read(args.ref_audio, always_2d=False)
@@ -108,6 +278,12 @@ def main():
ref_sr = target_sr
print(f"[Python] RefWav: {ref_wav.shape[0]} samples {ref_sr} Hz {ref_wav.shape[0]/ref_sr:.2f}s")
# Reproduce the upstream STFT magnitude on the same ref_wav so the
# mel-mag.bin pair scopes whether the divergence sits in the STFT or
# in the mel filter. This runs before the model speaker_encoder hook
# fires so both intermediates land in DUMP_PT before the test compare.
dump_mel_mag_python(ref_wav, DUMP_PT)
# Extract speaker embedding via ECAPA forward, projected to talker hidden.
spk_emb = model.extract_speaker_embedding(audio=ref_wav, sr=ref_sr)
print(f"[Python] SpeakerEmb shape: {tuple(spk_emb.shape)} dtype: {spk_emb.dtype}")
@@ -116,8 +292,14 @@ def main():
# Encode the reference audio to 16 codebook codes at 12.5 Hz. The encode
# call returns shape [T_codec, K=16] after the internal transpose, while
# the C++ side dumps [K=16, T_codec] row major. We transpose here for a
# straight exact match comparison.
enc = model.speech_tokenizer.encode([ref_wav], sr=int(ref_sr))
# straight exact match comparison. The C++ side aligns the number of
# samples to a multiple of the codec hop length (1920) before feeding
# the tokenizer, so we apply the same truncation upstream to keep T_codec
# comparable across the codec encoder bisection stages.
HOP = 1920
aligned_T = (ref_wav.shape[0] // HOP) * HOP
ref_wav_aln = ref_wav[:aligned_T]
enc = model.speech_tokenizer.encode([ref_wav_aln], sr=int(ref_sr))
ref_code_pt = enc.audio_codes[0]
ref_code_kt = ref_code_pt.transpose(0, 1).contiguous()
print(f"[Python] RefCodes shape: {tuple(ref_code_kt.shape)} (K, T_codec)")