Files
qwentts.cpp/tests/debug-clone-cossim.py
T
Pascal 9a9a425260 tests: exercise the fused decode in every cossim harness
Each harness re-runs the C++ side with --codec-fused after the
buffered run and gates on codes-full.bin equality between the two:
the decode path cannot change the predictor, so any divergence is a
bug. The fused run's perf lines land in the log as [Perf Fused] next
to the buffered ones, so every grid cell carries the buffered vs
fused comparison per backend and quant.

The fused audio itself is not scored here: greedy amplitudes vanish
in the streamed PCM_16 and the fused FP envelope has its own
validation against the stream path.
2026-07-30 22:54:44 +02:00

492 lines
22 KiB
Python
Executable File

#!/usr/bin/env python3
"""Cossim debug : C++ qwen-tts vs Python Qwen3-TTS on the Voice Clone Mode B
(ICL) path.
Inputs (relative to CWD = tests/) :
../examples/prompt.txt target text fed to both pipelines
../examples/freeman.wav reference audio for cloning
../examples/freeman.txt transcript of the reference audio
Default mode is greedy on both sides, non_streaming_mode=False which is
the ICL branch upstream : text + codec streams are aligned to the codec
length, the shorter one padded with tts_pad / truncated as needed.
Cote Python the speaker embedding is captured directly via
model.extract_speaker_embedding, and the reference codec frames via
model.speech_tokenizer.encode. Both intermediates land as spk-emb.bin
and ref-codes.bin and are compared against the C++ side dumps emitted
by pipeline-tts.cpp when --ref-wav and --ref-text are set.
Dumps land in cpp/clone/ (C++) and python/clone/ (Python).
"""
import argparse
import os
import subprocess
import sys
import librosa
import numpy as np
import soundfile as sf
import torch
import cossim_common as cc
MODEL_T = "../models/qwen-talker-1.7b-base-{q}.gguf"
MODEL_CDC_T = "../models/qwen-tokenizer-12hz-{q}.gguf"
CKPT = "../checkpoints/Qwen3-TTS-12Hz-1.7B-Base"
DUMP_CPP = "cpp/clone"
DUMP_PT = "python/clone"
# 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. 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 + [
("MelHann", "mel-hann.bin"),
("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"),
("SpkFrontend", "spk-frontend.bin"),
("SpkBlock3", "spk-block3.bin"),
("SpkMFA", "spk-mfa.bin"),
("SpkASP", "spk-asp.bin"),
("SpeakerEmb", "spk-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)
# 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 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"]:
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")
ap.add_argument("--ref-wav", default="../examples/freeman.wav",
help="reference WAV path for voice cloning")
ap.add_argument("--ref-text", default="../examples/freeman.txt",
help="path to a UTF-8 file with the transcript of ref-wav")
ap.add_argument("--seed", type=int, default=42)
ap.add_argument("--lang", default="english")
ap.add_argument("--quant", default="F32",
help="GGUF quantization suffix (F32, BF16, Q8_0, Q4_K_M)")
ap.add_argument("--out-pt", default=os.path.join(DUMP_PT, "clone-python.wav"))
ap.add_argument("--out-cpp", default=os.path.join(DUMP_CPP, "clone-cpp.wav"))
ap.add_argument("--max-new-tokens", type=int, default=64)
ap.add_argument("--trace", action="store_true",
help="print per sample u and idx for the first 32 samples")
args = ap.parse_args()
cc.ensure_dir(DUMP_PT)
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, "r", encoding="utf-8") as f:
ref_text = f.read().strip()
print(f"[Input] Prompt: {len(text)} chars: {text[:60]}{'...' if len(text) > 60 else ''}")
print(f"[Input] RefAudio: {args.ref_wav}")
print(f"[Input] RefText: {len(ref_text)} chars: {ref_text[:60]}{'...' if len(ref_text) > 60 else ''}")
print(f"[Input] Lang: {args.lang} Seed: {args.seed} MaxNewTokens: {args.max_new_tokens}")
print(f"[Input] Mode: greedy ICL")
torch.manual_seed(args.seed)
np.random.seed(args.seed)
cc.set_trace(args.trace)
cc.register_qwen3_tts()
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"[Python] Device: {device}")
model = cc.AutoModel.from_pretrained(
CKPT,
device_map=device,
dtype=torch.float32,
attn_implementation="eager",
).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_wav, always_2d=False)
if ref_wav.ndim > 1:
ref_wav = ref_wav[:, 0]
ref_wav = ref_wav.astype(np.float32)
target_sr = model.speaker_encoder_sample_rate
if ref_sr != 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")
# 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}")
cc.save_dump(os.path.join(DUMP_PT, "spk-emb.bin"), spk_emb)
# 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. 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]
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()
print(f"[Python] RefCodes shape: {tuple(ref_code_kt.shape)} (K, T_codec)")
cc.save_dump_i32(os.path.join(DUMP_PT, "ref-codes.bin"), ref_code_kt)
# Tokenize the utterance and the reference text.
assistant_text = f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n"
inp_utt = processor(text=assistant_text, return_tensors="pt", padding=True)
input_ids = inp_utt["input_ids"].to(device)
if input_ids.dim() == 1:
input_ids = input_ids.unsqueeze(0)
print(f"[Python] InputIds shape: {tuple(input_ids.shape)}")
cc.save_dump_i32(os.path.join(DUMP_PT, "prompt-ids.bin"), input_ids[0])
ref_text_wrap = f"<|im_start|>assistant\n{ref_text}<|im_end|>\n"
inp_ref = processor(text=ref_text_wrap, return_tensors="pt", padding=True)
ref_ids = inp_ref["input_ids"].to(device)
if ref_ids.dim() == 1:
ref_ids = ref_ids.unsqueeze(0)
print(f"[Python] RefIds shape: {tuple(ref_ids.shape)}")
cc.save_dump_i32(os.path.join(DUMP_PT, "ref-ids.bin"), ref_ids[0])
cc.install_hooks(model, DUMP_PT)
# Custom subtalker_* kwargs are forwarded to talker.forward but not
# declared on GenerationMixin, so transformers 4.57 rejects them under
# the strict validator. Disable it on the talker only.
model.talker._validate_model_kwargs = lambda *a, **k: None
# voice_clone_prompt dict mirrors what _prompt_items_to_voice_clone_prompt
# builds for a single ICL prompt item : ref_code is the [T_codec, K]
# tensor, ref_spk_embedding is the [hidden] tensor, x_vector_only=False
# and icl_mode=True together select the mode B branch upstream.
voice_clone_prompt_dict = dict(
ref_code = [ref_code_pt],
ref_spk_embedding = [spk_emb],
x_vector_only_mode = [False],
icl_mode = [True],
)
talker_codes_list, _ = model.generate(
input_ids=[input_ids],
ref_ids=[ref_ids],
voice_clone_prompt=voice_clone_prompt_dict,
languages=[args.lang],
non_streaming_mode=False,
max_new_tokens=args.max_new_tokens,
**cc.GEN_KWARGS_GREEDY,
)
codes = talker_codes_list[0]
print(f"[Python] Codes shape: {tuple(codes.shape)} (T_frames, num_code_groups)")
cc.save_dump_i32(os.path.join(DUMP_PT, "codes-full.bin"), codes)
cc.save_dump_i32(os.path.join(DUMP_PT, "codes-step0.bin"), codes[0])
# The decode path prepends the reference codes and cuts the matching
# audio prefix afterwards, mirroring generate_voice_clone exactly so the
# produced WAV only covers the freshly generated portion.
cat_codes = torch.cat([ref_code_pt.to(codes.device), codes], dim=0)
wavs, fs = model.speech_tokenizer.decode([{"audio_codes": cat_codes}])
full_wav = np.asarray(wavs[0], dtype=np.float32)
ref_len = int(ref_code_pt.shape[0])
total_len = int(cat_codes.shape[0])
cut = int(ref_len / max(total_len, 1) * full_wav.shape[0])
audio_pt = full_wav[cut:]
sf.write(args.out_pt, audio_pt, fs, subtype="FLOAT")
cc.save_dump(os.path.join(DUMP_PT, "output-audio.bin"), audio_pt)
print(f"[Python] Audio: {audio_pt.shape[0]} samples {fs} Hz {audio_pt.shape[0]/fs:.2f}s -> {args.out_pt}")
if not os.path.isfile(cc.BIN):
print(f"[Cossim] FATAL: {cc.BIN} not found, build qwen-tts first")
sys.exit(1)
model_lm = MODEL_T.format(q=args.quant)
model_cdc = MODEL_CDC_T.format(q=args.quant)
for p in (model_lm, model_cdc):
if not os.path.isfile(p):
print(f"[Cossim] FATAL: GGUF not found: {p}")
sys.exit(1)
print(f"[Quant] {args.quant} -> {model_lm} + {model_cdc}")
del model
if device == "cuda":
torch.cuda.empty_cache()
cmd = [
cc.BIN,
"--model", model_lm,
"--codec", model_cdc,
"--seed", str(args.seed),
"--ref-wav", args.ref_wav,
"--ref-text", args.ref_text,
"--lang", args.lang,
"--max-new", str(args.max_new_tokens),
"--dump", DUMP_CPP,
"-o", args.out_cpp,
"--greedy",
]
print(f"[GGML] Cmd: {' '.join(cmd)}")
r = subprocess.run(cmd, input=text, text=True)
if r.returncode != 0:
sys.exit(r.returncode)
audio_cpp, sr = sf.read(args.out_cpp)
if audio_cpp.ndim > 1:
audio_cpp = audio_cpp[:, 0]
audio_cpp = audio_cpp.astype(np.float32)
print(f"[GGML] Audio: {audio_cpp.shape[0]} samples {sr} Hz {audio_cpp.shape[0]/sr:.2f}s -> {args.out_cpp}")
cc.compare_exact_i32("prompt-ids.bin", DUMP_CPP, DUMP_PT, "PromptIDs")
cc.compare_exact_i32("ref-codes.bin", DUMP_CPP, DUMP_PT, "RefCodes")
cc.compare_stages(STAGES_CLONE, DUMP_CPP, DUMP_PT)
cc.compare_exact_i32("codes-full.bin", DUMP_CPP, DUMP_PT, "CodesFull")
aa, ab = cc.pair("output-audio.bin", DUMP_CPP, DUMP_PT)
print(f"[Cossim] Audio cos: {cc.cos(aa, ab):.6f}")
n = min(aa.size, ab.size)
print(f"[Cossim] WAV stft_cos: {cc.stft_cos(aa.ravel()[:n], ab.ravel()[:n]):.6f} samples: {n}")
cc.fused_pass(cmd, text, DUMP_CPP)
if __name__ == "__main__":
main()