From 7b9435c8860395c5ecabd2758eaab7f019985e2f Mon Sep 17 00:00:00 2001 From: Pascal Date: Sun, 10 May 2026 22:11:45 +0200 Subject: [PATCH] clone: mode B fix, librosa to torchaudio resample, plus SEANet bisection tooling --- src/pipeline-codec.cpp | 75 ++++++++++++++++++++++++++++++++++--- src/seanet-encoder.h | 43 +++++++++++++++++---- tests/debug-clone-cossim.py | 72 ++++++++++++++++++++++++++++++++++- 3 files changed, 176 insertions(+), 14 deletions(-) diff --git a/src/pipeline-codec.cpp b/src/pipeline-codec.cpp index be67ac5..078b6c2 100644 --- a/src/pipeline-codec.cpp +++ b/src/pipeline-codec.cpp @@ -265,7 +265,15 @@ std::vector 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 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 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 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 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. diff --git a/src/seanet-encoder.h b/src/seanet-encoder.h index 7a03fa7..e3e4e42 100644 --- a/src/seanet-encoder.h +++ b/src/seanet-encoder.h @@ -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); diff --git a/tests/debug-clone-cossim.py b/tests/debug-clone-cossim.py index ce901b8..2929aa8 100755 --- a/tests/debug-clone-cossim.py +++ b/tests/debug-clone-cossim.py @@ -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()