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

clone: mode B fix, librosa to torchaudio resample, plus SEANet bisection tooling

Browse Source
This commit is contained in:
Pascal
2026-05-10 22:11:45 +02:00
parent 186ee91c70
commit 7b9435c886
3 changed files with 176 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/pipeline-codec.cpp
+70 -5
View File
@@ -265,7 +265,15 @@ std::vector<int32_t> pipeline_codec_encode(PipelineCodec * pc, const float * aud
ggml_set_input(mask);
// Forward chain.
struct ggml_tensor * h_seanet = qwen_seanet_encoder_forward(gctx, &pc->seanet, audio_in); // [T_emb, 512]
struct ggml_tensor * sn_init_t = NULL;
struct ggml_tensor * sn_resnet0_t = NULL;
struct ggml_tensor * sn_stage0_t = NULL;
struct ggml_tensor * sn_stage1_t = NULL;
struct ggml_tensor * sn_stage3_t = NULL;
struct ggml_tensor * h_seanet =
qwen_seanet_encoder_forward(gctx, &pc->seanet, audio_in,
&sn_init_t, &sn_resnet0_t, &sn_stage0_t,
&sn_stage1_t, &sn_stage3_t); // [T_emb, 512]
struct ggml_tensor * h = ggml_cont(gctx, ggml_transpose(gctx, h_seanet)); // [512, T_emb]
struct ggml_tensor * h_et =
qwen_encoder_transformer_forward(gctx, &pc->enc_transformer, h, positions, mask); // [512, T_emb]
@@ -279,17 +287,50 @@ std::vector<int32_t> pipeline_codec_encode(PipelineCodec * pc, const float * aud
h = ggml_cont(gctx, ggml_transpose(gctx, h)); // ne=(512, T)
const char * dump = dump_dir;
struct ggml_tensor * h_seanet_dump = NULL;
struct ggml_tensor * h_seanet_dump = NULL;
struct ggml_tensor * sn_init_dump = NULL;
struct ggml_tensor * sn_resnet0_dump = NULL;
struct ggml_tensor * sn_stage0_dump = NULL;
struct ggml_tensor * sn_stage1_dump = NULL;
struct ggml_tensor * sn_stage3_dump = NULL;
if (dump) {
// SEANet output naturally lands as channel-first ggml ne=(T, hidden).
// SEANet output naturally lands as ggml ne=(T, hidden) (T innermost).
// 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.
// T-first numpy [T, hidden], so we transpose+cont to bring hidden
// innermost before pinning as a graph output. The dump_2d then
// emits shape (ne[1], ne[0]) = (T, hidden) on the numpy side.
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");
// SEANet bisection points. Same transpose convention as h_seanet.
if (sn_init_t) {
sn_init_dump = ggml_cont(gctx, ggml_transpose(gctx, sn_init_t));
ggml_set_output(sn_init_dump);
ggml_set_name(sn_init_dump, "seanet_init_dump");
}
if (sn_resnet0_t) {
sn_resnet0_dump = ggml_cont(gctx, ggml_transpose(gctx, sn_resnet0_t));
ggml_set_output(sn_resnet0_dump);
ggml_set_name(sn_resnet0_dump, "seanet_resnet0_dump");
}
if (sn_stage0_t) {
sn_stage0_dump = ggml_cont(gctx, ggml_transpose(gctx, sn_stage0_t));
ggml_set_output(sn_stage0_dump);
ggml_set_name(sn_stage0_dump, "seanet_stage0_dump");
}
if (sn_stage1_t) {
sn_stage1_dump = ggml_cont(gctx, ggml_transpose(gctx, sn_stage1_t));
ggml_set_output(sn_stage1_dump);
ggml_set_name(sn_stage1_dump, "seanet_stage1_dump");
}
if (sn_stage3_t) {
sn_stage3_dump = ggml_cont(gctx, ggml_transpose(gctx, sn_stage3_t));
ggml_set_output(sn_stage3_dump);
ggml_set_name(sn_stage3_dump, "seanet_stage3_dump");
}
}
ggml_set_name(h, "enc_hidden_out");
@@ -300,6 +341,21 @@ std::vector<int32_t> pipeline_codec_encode(PipelineCodec * pc, const float * aud
if (h_seanet_dump) {
ggml_build_forward_expand(graph, h_seanet_dump);
}
if (sn_init_dump) {
ggml_build_forward_expand(graph, sn_init_dump);
}
if (sn_resnet0_dump) {
ggml_build_forward_expand(graph, sn_resnet0_dump);
}
if (sn_stage0_dump) {
ggml_build_forward_expand(graph, sn_stage0_dump);
}
if (sn_stage1_dump) {
ggml_build_forward_expand(graph, sn_stage1_dump);
}
if (sn_stage3_dump) {
ggml_build_forward_expand(graph, sn_stage3_dump);
}
if (!ggml_backend_sched_alloc_graph(pc->sched, graph)) {
qt_log(QT_LOG_ERROR, "[Pipeline] encode sched_alloc_graph failed");
@@ -329,6 +385,10 @@ std::vector<int32_t> pipeline_codec_encode(PipelineCodec * pc, const float * aud
if (dump) {
DebugDumper d;
debug_init(&d, dump);
// Raw audio input dump : the SEANet sees this, and any divergence
// in the resampler (torchaudio reimpl C++ vs librosa Python) shows
// up here as a phase or amplitude drift.
debug_dump_1d(&d, "audio-input", audio, n_samples);
// 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
@@ -342,6 +402,11 @@ std::vector<int32_t> pipeline_codec_encode(PipelineCodec * pc, const float * aud
dump2("seanet-out", h_seanet_dump);
dump2("enc-transformer-out", h_et);
dump2("codec-pre-fsq", h);
if (sn_init_dump) { dump2("seanet-init", sn_init_dump); }
if (sn_resnet0_dump) { dump2("seanet-resnet0", sn_resnet0_dump); }
if (sn_stage0_dump) { dump2("seanet-stage0", sn_stage0_dump); }
if (sn_stage1_dump) { dump2("seanet-stage1", sn_stage1_dump); }
if (sn_stage3_dump) { dump2("seanet-stage3", sn_stage3_dump); }
}
// Read back the post-downsample hidden buffer for CPU-side RVQ encode.
src/seanet-encoder.h
+35 -8
View File
@@ -8,17 +8,16 @@
//
// Structure:
// init : MimiConv1d k=7, 1 -> 64, causal stride=1
// for ratio in [4, 5, 6, 8] (Python loop reversed: ratios 8,6,5,4 in
// downsampling order applied to the audio):
// for ratio in reversed(upsampling_ratios) i.e. iter [4, 5, 6, 8] :
// resnet block: ELU -> Conv1d k=3 d=1 dim/2 -> ELU -> Conv1d k=1 dim
// ELU
// Conv1d k=2*ratio, stride=ratio, channels x2
// last : MimiConv1d k=3, 1024 -> 512, causal stride=1
//
// Apply order on a 24 kHz waveform:
// audio -> init (1->64) -> stage 0 (8x, 64->128) -> stage 1 (6x, 128->256)
// -> stage 2 (5x, 256->512) -> stage 3 (4x, 512->1024) -> last (1024->512)
// Total downsample = 8 * 6 * 5 * 4 = 960. The 12.5 Hz rate is reached after
// Apply order on a 24 kHz waveform (matches Python MimiEncoder forward):
// audio -> init (1->64) -> stage 0 (4x, 64->128) -> stage 1 (5x, 128->256)
// -> stage 2 (6x, 256->512) -> stage 3 (8x, 512->1024) -> last (1024->512)
// Total downsample = 4 * 5 * 6 * 8 = 960. The 12.5 Hz rate is reached after
// the final downsample conv (factor 2 more in encoder-downsample.h).
#include "causal-trans-conv.h"
@@ -187,17 +186,45 @@ static struct ggml_tensor * qwen_seanet_resnet_forward(struct ggml_context *
// Full SEANet forward.
// x: [T_audio, 1] f32 T-first (mono waveform)
// returns [T_audio / 960, 512] f32 T-first.
// Optional out-params capture intermediate stage outputs in T-first ggml
// layout ne=(C, T_out) for debug bisection. Each is NULL by default and
// the caller decides whether to mark them as graph outputs.
// init_out : post init MimiConv1d k=7, [T_audio, 64]
// resnet0_out : post stage 0 resnet block, before ELU+downsample, [T_audio, 64]
// stage0_out : post stage 0 (resnet + ELU + downsample 4x), [T_audio/4, 128]
// stage1_out : post stage 1 (resnet + ELU + downsample 5x), [T_audio/20, 256]
// stage3_out : post stage 3 (resnet + ELU + downsample 8x), [T_audio/960, 1024]
// Returns [T_audio / 960, 512] f32 T-first.
static struct ggml_tensor * qwen_seanet_encoder_forward(struct ggml_context * ctx,
const QwenSEANetEncoder * s,
struct ggml_tensor * x) {
struct ggml_tensor * x,
struct ggml_tensor ** init_out = NULL,
struct ggml_tensor ** resnet0_out = NULL,
struct ggml_tensor ** stage0_out = NULL,
struct ggml_tensor ** stage1_out = NULL,
struct ggml_tensor ** stage3_out = NULL) {
x = qwen_causal_conv1d(ctx, s->init_w, s->init_b, x, s->kernel_size, 1, 1);
if (init_out) {
*init_out = x;
}
for (int i = 0; i < QWEN_SEANET_NUM_STAGES; i++) {
const QwenSEANetStage & stg = s->stages[i];
x = qwen_seanet_resnet_forward(ctx, &stg.resnet, x, s->residual_kernel_size);
if (i == 0 && resnet0_out) {
*resnet0_out = x;
}
x = ggml_elu(ctx, x);
x = qwen_causal_conv1d(ctx, stg.down_w, stg.down_b, x, 2 * stg.ratio, 1, stg.ratio);
if (i == 0 && stage0_out) {
*stage0_out = x;
}
if (i == 1 && stage1_out) {
*stage1_out = x;
}
if (i == 3 && stage3_out) {
*stage3_out = x;
}
}
x = ggml_elu(ctx, x);
tests/debug-clone-cossim.py
+71 -1
View File
@@ -53,6 +53,11 @@ STAGES_CLONE = cc.STAGES_STANDARD + [
("MelBasis", "mel-basis.bin"),
("MelMag", "mel-mag.bin"),
("MelSpk", "mel-spk.bin"),
("SeanetInit", "seanet-init.bin"),
("SeanetResnet0", "seanet-resnet0.bin"),
("SeanetStage0", "seanet-stage0.bin"),
("SeanetStage1", "seanet-stage1.bin"),
("SeanetStage3", "seanet-stage3.bin"),
("SeanetOut", "seanet-out.bin"),
("EncTransformer", "enc-transformer-out.bin"),
("CodecPreFSQ", "codec-pre-fsq.bin"),
@@ -106,6 +111,61 @@ def install_clone_hooks(model, dump_dir):
seen_down["done"] = True
enc.downsample.register_forward_hook(hook_down)
# SEANet bisection. enc.encoder is a MimiEncoder whose .layers ModuleList
# holds, in order : [0] init MimiConv1d, [1] resnet, [2] ELU, [3] down 4x,
# [4] resnet, [5] ELU, [6] down 5x, [7] resnet, [8] ELU, [9] down 6x,
# [10] resnet, [11] ELU, [12] down 8x, [13] ELU, [14] last MimiConv1d.
# We hook the init conv and the three downsample convs the C++ side
# exposes as out-params in qwen_seanet_encoder_forward.
sn_layers = enc.encoder.layers
seen_sn_init = {"done": False}
def hook_sn_init(module, args, output):
if seen_sn_init["done"]:
return
out = output[0] if isinstance(output, tuple) else output
# MimiConv1d output : [B=1, OC, T] channel-first -> [T, OC] T-first.
cc.save_dump(os.path.join(dump_dir, "seanet-init.bin"), out[0].transpose(0, 1).contiguous())
seen_sn_init["done"] = True
sn_layers[0].register_forward_hook(hook_sn_init)
seen_sn_r0 = {"done": False}
def hook_sn_resnet0(module, args, output):
if seen_sn_r0["done"]:
return
out = output[0] if isinstance(output, tuple) else output
# MimiResnetBlock output : [B=1, OC, T] channel-first -> [T, OC] T-first.
cc.save_dump(os.path.join(dump_dir, "seanet-resnet0.bin"), out[0].transpose(0, 1).contiguous())
seen_sn_r0["done"] = True
sn_layers[1].register_forward_hook(hook_sn_resnet0)
seen_sn_s0 = {"done": False}
def hook_sn_stage0(module, args, output):
if seen_sn_s0["done"]:
return
out = output[0] if isinstance(output, tuple) else output
cc.save_dump(os.path.join(dump_dir, "seanet-stage0.bin"), out[0].transpose(0, 1).contiguous())
seen_sn_s0["done"] = True
sn_layers[3].register_forward_hook(hook_sn_stage0)
seen_sn_s1 = {"done": False}
def hook_sn_stage1(module, args, output):
if seen_sn_s1["done"]:
return
out = output[0] if isinstance(output, tuple) else output
cc.save_dump(os.path.join(dump_dir, "seanet-stage1.bin"), out[0].transpose(0, 1).contiguous())
seen_sn_s1["done"] = True
sn_layers[6].register_forward_hook(hook_sn_stage1)
seen_sn_s3 = {"done": False}
def hook_sn_stage3(module, args, output):
if seen_sn_s3["done"]:
return
out = output[0] if isinstance(output, tuple) else output
cc.save_dump(os.path.join(dump_dir, "seanet-stage3.bin"), out[0].transpose(0, 1).contiguous())
seen_sn_s3["done"] = True
sn_layers[12].register_forward_hook(hook_sn_stage3)
seen_mel = {"done": False}
def hook_spk_pre(module, args, kwargs):
if seen_mel["done"]:
@@ -274,7 +334,16 @@ def main():
ref_wav = ref_wav.astype(np.float32)
target_sr = model.speaker_encoder_sample_rate
if ref_sr != target_sr:
ref_wav = librosa.resample(y=ref_wav, orig_sr=int(ref_sr), target_sr=int(target_sr))
# Match C++ side audio_resample.h which is a torchaudio.functional.resample
# reimplementation. Using librosa.resample here would introduce a phase
# drift between the two waveforms that propagates through the SEANet
# stack and shows up as a measurable cossim drop on the codec encoder
# bisection stages.
import torchaudio
ref_wav = torchaudio.functional.resample(
torch.from_numpy(ref_wav.astype(np.float32)),
int(ref_sr), int(target_sr),
).numpy()
ref_sr = target_sr
print(f"[Python] RefWav: {ref_wav.shape[0]} samples {ref_sr} Hz {ref_wav.shape[0]/ref_sr:.2f}s")
@@ -299,6 +368,7 @@ def main():
HOP = 1920
aligned_T = (ref_wav.shape[0] // HOP) * HOP
ref_wav_aln = ref_wav[:aligned_T]
cc.save_dump(os.path.join(DUMP_PT, "audio-input.bin"), torch.from_numpy(ref_wav_aln.astype(np.float32)))
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()