initial release
This commit is contained in:
civ
Executable
+408
@@ -0,0 +1,408 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Cossim debug : C++ omnivoice-tts vs Python OmniVoice on voice design.
|
||||
|
||||
Inputs (relative to CWD = tests/) :
|
||||
../examples/prompt.txt target text fed to both pipelines
|
||||
|
||||
Both sides run with :
|
||||
instruct=male, language=English, seed=42, F32 weights, no pre or post
|
||||
process. Defaults match : num_step=32, guidance_scale=2.0, t_shift=0.1,
|
||||
layer_penalty_factor=5.0, position_temperature=5.0, class_temperature=0.0.
|
||||
|
||||
Dumps land in cpp/ (C++) and python/ (Python). The script compares each
|
||||
matching .bin pair via cosine similarity over the f32 payload, plus exact
|
||||
match rate for tensors that originated as int (mg-tokens). All paths are
|
||||
relative, no absolute paths anywhere.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import struct
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
# Strict F32 matmul on both sides. NVIDIA_TF32_OVERRIDE=0 forces full FP32
|
||||
# mantissa in cuBLAS for both PyTorch and the C++ child via inheritance.
|
||||
# Must be set BEFORE torch imports so the cuBLAS handle reads it on init.
|
||||
os.environ["NVIDIA_TF32_OVERRIDE"] = "0"
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
import torch
|
||||
|
||||
# Belt and suspenders : disable PyTorch's own TF32 toggles too. Some code
|
||||
# paths bypass NVIDIA_TF32_OVERRIDE through cudnn or torch internal flags.
|
||||
torch.backends.cuda.matmul.allow_tf32 = False
|
||||
torch.backends.cudnn.allow_tf32 = False
|
||||
torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
|
||||
torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
|
||||
torch.set_float32_matmul_precision("highest")
|
||||
|
||||
from omnivoice import OmniVoice
|
||||
from omnivoice.utils.common import fix_random_seed
|
||||
|
||||
BIN = "../build/omnivoice-tts"
|
||||
MODEL_LM_T = "../models/omnivoice-base-{q}.gguf"
|
||||
MODEL_CDC_T = "../models/omnivoice-tokenizer-{q}.gguf"
|
||||
CKPT = "../checkpoints/OmniVoice"
|
||||
DUMP_CPP = "cpp"
|
||||
DUMP_PT = "python"
|
||||
|
||||
def ensure_dir(path):
|
||||
os.makedirs(path, exist_ok=True)
|
||||
|
||||
def save_dump(path, data):
|
||||
"""Write a tensor in the C++ debug.h format :
|
||||
[ndim:i32] [shape:i32 x ndim] [data:f32 x numel]
|
||||
"""
|
||||
if isinstance(data, torch.Tensor):
|
||||
data = data.detach().to(torch.float32).cpu().numpy()
|
||||
data = np.ascontiguousarray(data.astype(np.float32))
|
||||
shape = data.shape
|
||||
with open(path, "wb") as f:
|
||||
f.write(struct.pack("i", len(shape)))
|
||||
for s in shape:
|
||||
f.write(struct.pack("i", s))
|
||||
f.write(data.tobytes())
|
||||
|
||||
def load_dump(path):
|
||||
"""Inverse of save_dump : returns (data:f32 numpy, shape:tuple)."""
|
||||
raw = np.fromfile(path, dtype=np.uint8)
|
||||
ndim = int(np.frombuffer(raw[0:4], dtype=np.int32)[0])
|
||||
shape = tuple(int(x) for x in np.frombuffer(raw[4:4 + 4 * ndim], dtype=np.int32))
|
||||
body = np.frombuffer(raw[4 + 4 * ndim:], dtype=np.float32)
|
||||
return body.reshape(shape), shape
|
||||
|
||||
def cos(a, b):
|
||||
a = a.astype(np.float64).ravel()
|
||||
b = b.astype(np.float64).ravel()
|
||||
n = min(len(a), len(b))
|
||||
a, b = a[:n], b[:n]
|
||||
d = float(np.linalg.norm(a) * np.linalg.norm(b))
|
||||
return float(np.dot(a, b) / d) if d > 1e-10 else 0.0
|
||||
|
||||
def stft_cos(a, b, win=2048, hop=512):
|
||||
# STFT magnitude cosine. Drops phase, so a constant time shift between
|
||||
# the two waveforms does not collapse the score. The plain cos() on
|
||||
# raw samples falls to ~0 the moment chunks land a few samples apart.
|
||||
a = a.astype(np.float64).ravel()
|
||||
b = b.astype(np.float64).ravel()
|
||||
n = min(len(a), len(b))
|
||||
a, b = a[:n], b[:n]
|
||||
window = np.hanning(win)
|
||||
frames = (n - win) // hop + 1
|
||||
if frames <= 0:
|
||||
return 0.0
|
||||
sa = np.zeros((frames, win // 2 + 1))
|
||||
sb = np.zeros((frames, win // 2 + 1))
|
||||
for i in range(frames):
|
||||
s = i * hop
|
||||
sa[i] = np.abs(np.fft.rfft(a[s:s + win] * window))
|
||||
sb[i] = np.abs(np.fft.rfft(b[s:s + win] * window))
|
||||
return cos(sa.ravel(), sb.ravel())
|
||||
|
||||
def install_hooks(model, dump_dir):
|
||||
# First call to _prepare_embed_inputs at step 0 returns the input embedding
|
||||
# right before layer 0 of the LLM, mirroring the C++ inputs_embeds dump.
|
||||
# Also captures the raw input_ids row k=0 for cond and uncond, so any
|
||||
# token-level divergence localizes upstream of the embed lookup.
|
||||
seen_embed = {"done": False}
|
||||
orig_prepare = model._prepare_embed_inputs
|
||||
def hooked_prepare(input_ids, audio_mask):
|
||||
out = orig_prepare(input_ids, audio_mask)
|
||||
if not seen_embed["done"] and out.dim() == 3 and out.shape[0] >= 2:
|
||||
cond = out[0].detach().to(torch.float32).cpu().numpy()
|
||||
uncond = out[1].detach().to(torch.float32).cpu().numpy()
|
||||
save_dump(os.path.join(dump_dir, "lm-hidden-step0-cond-embed.bin"), cond)
|
||||
save_dump(os.path.join(dump_dir, "lm-hidden-step0-uncond-embed.bin"), uncond)
|
||||
# Style and text tokens duplicate across all K codebooks so k=0
|
||||
# carries the full sequence for diagnostic. Cast to f32 keeps the
|
||||
# debug.h binary format identical on both sides.
|
||||
cond_ids = input_ids[0, 0, :].detach().to(torch.float32).cpu().numpy()
|
||||
uncond_ids = input_ids[1, 0, :].detach().to(torch.float32).cpu().numpy()
|
||||
save_dump(os.path.join(dump_dir, "prompt-cond-ids.bin"), cond_ids)
|
||||
save_dump(os.path.join(dump_dir, "prompt-uncond-ids.bin"), uncond_ids)
|
||||
seen_embed["done"] = True
|
||||
return out
|
||||
model._prepare_embed_inputs = hooked_prepare
|
||||
|
||||
# Bisection : dump cond and uncond hidden states after a few layers so a
|
||||
# mismatch can be localized within the 28 layer Qwen3 stack.
|
||||
bisect_layers = [0, 6, 13, 20]
|
||||
seen_layers = {idx: False for idx in bisect_layers}
|
||||
def make_layer_hook(layer_idx):
|
||||
def hook(module, inputs, output):
|
||||
if seen_layers[layer_idx]:
|
||||
return
|
||||
h = output[0] if isinstance(output, tuple) else output
|
||||
if h.dim() == 3 and h.shape[0] >= 2:
|
||||
cond = h[0].detach().to(torch.float32).cpu().numpy()
|
||||
uncond = h[1].detach().to(torch.float32).cpu().numpy()
|
||||
save_dump(os.path.join(dump_dir, f"lm-hidden-step0-cond-l{layer_idx}.bin"), cond)
|
||||
save_dump(os.path.join(dump_dir, f"lm-hidden-step0-uncond-l{layer_idx}.bin"), uncond)
|
||||
seen_layers[layer_idx] = True
|
||||
return hook
|
||||
for layer_idx in bisect_layers:
|
||||
model.llm.layers[layer_idx].register_forward_hook(make_layer_hook(layer_idx))
|
||||
|
||||
# First call to audio_heads corresponds to step 0 of the MaskGIT loop.
|
||||
# The input to audio_heads is the final hidden state, shape [B, S, D],
|
||||
# mirroring what the C++ side reads back via dump_hidden_dir before the
|
||||
# lm_head matmul. We dump cond (b=0) and uncond (b=1) separately.
|
||||
seen_hidden = {"done": False}
|
||||
def pre_audio_heads(module, inputs):
|
||||
if not seen_hidden["done"]:
|
||||
h = inputs[0]
|
||||
if h.dim() == 3 and h.shape[0] >= 2:
|
||||
cond = h[0].detach().to(torch.float32).cpu().numpy()
|
||||
uncond = h[1].detach().to(torch.float32).cpu().numpy()
|
||||
save_dump(os.path.join(dump_dir, "lm-hidden-step0-cond.bin"), cond)
|
||||
save_dump(os.path.join(dump_dir, "lm-hidden-step0-uncond.bin"), uncond)
|
||||
seen_hidden["done"] = True
|
||||
model.audio_heads.register_forward_pre_hook(pre_audio_heads)
|
||||
|
||||
# First call to _predict_tokens_with_scoring corresponds to step 0 of the
|
||||
# MaskGIT loop. Capture cond and uncond logits in [K, T, V] layout (squeeze
|
||||
# the batch axis to match the C++ dump shape). Replicate the predict math
|
||||
# locally so log_probs / pred_tokens / scores can all be dumped from a
|
||||
# single source of truth on the Python side.
|
||||
seen = {"step0": False, "mg_tokens": False, "audio": False}
|
||||
orig_pred = model._predict_tokens_with_scoring
|
||||
def hooked_pred(c_logits, u_logits, gen_config):
|
||||
pred_tokens, scores = orig_pred(c_logits, u_logits, gen_config)
|
||||
if not seen["step0"]:
|
||||
if gen_config.guidance_scale != 0:
|
||||
c_lp = torch.nn.functional.log_softmax(c_logits, dim=-1)
|
||||
u_lp = torch.nn.functional.log_softmax(u_logits, dim=-1)
|
||||
log_probs = torch.log_softmax(
|
||||
c_lp + gen_config.guidance_scale * (c_lp - u_lp), dim=-1)
|
||||
else:
|
||||
log_probs = torch.nn.functional.log_softmax(c_logits, dim=-1)
|
||||
log_probs = log_probs.clone()
|
||||
log_probs[..., model.config.audio_mask_id] = float("-inf")
|
||||
|
||||
c = c_logits.detach().to(torch.float32).cpu().numpy()
|
||||
u = u_logits.detach().to(torch.float32).cpu().numpy()
|
||||
if c.ndim == 4:
|
||||
c = c[0]
|
||||
if u.ndim == 4:
|
||||
u = u[0]
|
||||
save_dump(os.path.join(dump_dir, "lm-logits-step0-cond.bin"), c)
|
||||
save_dump(os.path.join(dump_dir, "lm-logits-step0-uncond.bin"), u)
|
||||
|
||||
lp_arr = log_probs.detach().to(torch.float32).cpu().numpy()
|
||||
if lp_arr.ndim == 4:
|
||||
lp_arr = lp_arr[0]
|
||||
save_dump(os.path.join(dump_dir, "mg-log-probs-step0.bin"), lp_arr)
|
||||
|
||||
pt_arr = pred_tokens.detach().to(torch.float32).cpu().numpy()
|
||||
sc_arr = scores.detach().to(torch.float32).cpu().numpy()
|
||||
if pt_arr.ndim == 3:
|
||||
pt_arr = pt_arr[0]
|
||||
if sc_arr.ndim == 3:
|
||||
sc_arr = sc_arr[0]
|
||||
save_dump(os.path.join(dump_dir, "mg-pred-tokens-step0.bin"), pt_arr)
|
||||
save_dump(os.path.join(dump_dir, "mg-scores-step0.bin"), sc_arr)
|
||||
seen["step0"] = True
|
||||
return pred_tokens, scores
|
||||
model._predict_tokens_with_scoring = hooked_pred
|
||||
|
||||
orig_generate = model._generate_iterative
|
||||
def hooked_generate(task, gen_config):
|
||||
out = orig_generate(task, gen_config)
|
||||
# out is a list of (K, T_i) long tensors, one per batch item.
|
||||
# In chunked mode this hook fires once per chunk; only the first
|
||||
# call mirrors the chunk 0 dump on the C++ side.
|
||||
if not seen["mg_tokens"]:
|
||||
save_dump(os.path.join(dump_dir, "mg-tokens.bin"), out[0])
|
||||
seen["mg_tokens"] = True
|
||||
return out
|
||||
model._generate_iterative = hooked_generate
|
||||
|
||||
orig_decode = model.audio_tokenizer.decode
|
||||
def hooked_decode(*args, **kwargs):
|
||||
out = orig_decode(*args, **kwargs)
|
||||
# The audio tokenizer returns either a tensor or a wrapper holding
|
||||
# audio_values shape [B, C, N]. Unwrap and dump the first item mono.
|
||||
# Same first-chunk guard as the mg-tokens hook above.
|
||||
if seen["audio"]:
|
||||
return out
|
||||
wav = getattr(out, "audio_values", out)
|
||||
if isinstance(wav, torch.Tensor):
|
||||
arr = wav.detach().to(torch.float32).cpu().numpy()
|
||||
else:
|
||||
arr = np.asarray(wav, dtype=np.float32)
|
||||
if arr.ndim == 3:
|
||||
arr = arr[0, 0]
|
||||
elif arr.ndim == 2:
|
||||
arr = arr[0]
|
||||
save_dump(os.path.join(dump_dir, "output-audio.bin"), arr)
|
||||
seen["audio"] = True
|
||||
return out
|
||||
model.audio_tokenizer.decode = hooked_decode
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--prompt", default="../examples/prompt.txt")
|
||||
ap.add_argument("--seed", type=int, default=42)
|
||||
ap.add_argument("--instruct", default="male, young adult, moderate pitch")
|
||||
ap.add_argument("--lang", default="English")
|
||||
ap.add_argument("--duration", type=float, default=None)
|
||||
ap.add_argument("--quant", default="F32",
|
||||
help="quantization suffix for GGUF (default: F32, e.g. BF16, Q8_0, Q4_K_M)")
|
||||
ap.add_argument("--out-cpp", default="cpp/tts-cpp.wav")
|
||||
ap.add_argument("--out-pt", default="python/tts-python.wav")
|
||||
args = ap.parse_args()
|
||||
|
||||
model_lm = MODEL_LM_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"[Error] GGUF not found: {p}")
|
||||
sys.exit(1)
|
||||
print(f"[Quant] {args.quant} -> {model_lm} + {model_cdc}")
|
||||
|
||||
ensure_dir(DUMP_CPP)
|
||||
ensure_dir(DUMP_PT)
|
||||
os.makedirs(os.path.dirname(args.out_cpp) or ".", exist_ok=True)
|
||||
|
||||
with open(args.prompt, "r", encoding="utf-8") as f:
|
||||
text = f.read().strip()
|
||||
print(f"[Input] Prompt: {len(text)} chars: {text[:60]}{'...' if len(text) > 60 else ''}")
|
||||
print(f"[Input] Instruct: {args.instruct}")
|
||||
print(f"[Input] Language: {args.lang}")
|
||||
print(f"[Input] Seed: {args.seed}")
|
||||
|
||||
# Python reference path : F32, voice design male, no pre or post process.
|
||||
fix_random_seed(args.seed)
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
model = OmniVoice.from_pretrained(
|
||||
CKPT,
|
||||
torch_dtype=torch.float32,
|
||||
attn_implementation="eager",
|
||||
).to(device).eval()
|
||||
install_hooks(model, DUMP_PT)
|
||||
gen_kwargs = dict(
|
||||
text=text,
|
||||
language=args.lang,
|
||||
instruct=args.instruct,
|
||||
duration=args.duration,
|
||||
)
|
||||
audios = model.generate(**gen_kwargs)
|
||||
audio_pt = np.asarray(audios[0], dtype=np.float32)
|
||||
sf.write(args.out_pt, audio_pt, 24000, subtype="FLOAT")
|
||||
print(f"[Python] Audio: {audio_pt.shape[0]} samples {audio_pt.shape[0] / 24000:.2f}s -> {args.out_pt}")
|
||||
|
||||
# Free the GPU before launching the C++ binary so it has room to load
|
||||
# the F32 GGUFs without fighting for VRAM.
|
||||
del model
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
# C++ path : same text, same instruct, same seed, F32 GGUF weights,
|
||||
# dumps under cpp/.
|
||||
cmd = [
|
||||
BIN,
|
||||
"--model", model_lm,
|
||||
"--codec", model_cdc,
|
||||
"--seed", str(args.seed),
|
||||
"--instruct", args.instruct,
|
||||
"--lang", args.lang,
|
||||
"--format", "wav32",
|
||||
"--dump", DUMP_CPP,
|
||||
"--no-fa",
|
||||
"-o", args.out_cpp,
|
||||
]
|
||||
if args.duration:
|
||||
cmd += ["--duration", str(args.duration)]
|
||||
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}")
|
||||
|
||||
# Cossim in pipeline order: prompt-ids -> embed -> l0 -> l6 -> l13 -> l20
|
||||
# -> final -> logits -> tokens -> audio. Cond and uncond on the same line
|
||||
# so a drift localizes immediately to the originating stage.
|
||||
def pair(name):
|
||||
a, _ = load_dump(os.path.join(DUMP_CPP, name))
|
||||
b, _ = load_dump(os.path.join(DUMP_PT, name))
|
||||
return a, b
|
||||
|
||||
def ids_exact(a, b):
|
||||
n = min(a.size, b.size)
|
||||
ai = a.astype(np.int64).ravel()[:n]
|
||||
bi = b.astype(np.int64).ravel()[:n]
|
||||
diffs = np.where(ai != bi)[0]
|
||||
return 100.0 * float(np.mean(ai == bi)), diffs, ai, bi
|
||||
|
||||
ca, cb = pair("prompt-cond-ids.bin")
|
||||
ua, ub = pair("prompt-uncond-ids.bin")
|
||||
cm, cd, cai, cbi = ids_exact(ca, cb)
|
||||
um, ud, uai, ubi = ids_exact(ua, ub)
|
||||
print(f"[Cossim] PromptIDs cond exact: {cm:.2f}% uncond exact: {um:.2f}%")
|
||||
for s in cd[:20]:
|
||||
print(f"[Cossim] PromptIDs cond diff at s={s}: ggml={cai[s]} python={cbi[s]}")
|
||||
for s in ud[:20]:
|
||||
print(f"[Cossim] PromptIDs uncond diff at s={s}: ggml={uai[s]} python={ubi[s]}")
|
||||
|
||||
stages = [
|
||||
("Embed", "lm-hidden-step0-{}-embed.bin"),
|
||||
("L0", "lm-hidden-step0-{}-l0.bin"),
|
||||
("L6", "lm-hidden-step0-{}-l6.bin"),
|
||||
("L13", "lm-hidden-step0-{}-l13.bin"),
|
||||
("L20", "lm-hidden-step0-{}-l20.bin"),
|
||||
("Final", "lm-hidden-step0-{}.bin"),
|
||||
("Logits", "lm-logits-step0-{}.bin"),
|
||||
]
|
||||
def metric(a, b):
|
||||
n = min(a.size, b.size)
|
||||
af = a.astype(np.float64).ravel()[:n]
|
||||
bf = b.astype(np.float64).ravel()[:n]
|
||||
d = np.abs(af - bf)
|
||||
nrm_a = float(np.linalg.norm(af))
|
||||
nrm_b = float(np.linalg.norm(bf))
|
||||
c = float(np.dot(af, bf) / (nrm_a * nrm_b)) if nrm_a > 1e-10 and nrm_b > 1e-10 else 0.0
|
||||
return c, float(d.max()), float(d.mean())
|
||||
|
||||
for label, fmt in stages:
|
||||
ca, cb = pair(fmt.format("cond"))
|
||||
ua, ub = pair(fmt.format("uncond"))
|
||||
cc, cmax, cmean = metric(ca, cb)
|
||||
uc, umax, umean = metric(ua, ub)
|
||||
print(f"[Cossim] {label} cond cos: {cc:.6f} max: {cmax:.4e} mean: {cmean:.4e} uncond cos: {uc:.6f} max: {umax:.4e} mean: {umean:.4e}")
|
||||
|
||||
pa, pb = pair("mg-pred-tokens-step0.bin")
|
||||
n = min(pa.size, pb.size)
|
||||
ai = pa.astype(np.int64).ravel()[:n]
|
||||
bi = pb.astype(np.int64).ravel()[:n]
|
||||
diffs = np.where(ai != bi)[0]
|
||||
print(f"[Cossim] Step0Tokens exact: {100.0 * float((ai == bi).mean()):.2f}% diffs: {diffs.size}")
|
||||
|
||||
sa, sb = pair("mg-scores-step0.bin")
|
||||
sd = np.abs(sa - sb)
|
||||
print(f"[Cossim] Step0Scores cos: {cos(sa, sb):.6f} max_abs_diff: {sd.max():.6f} mean_abs_diff: {sd.mean():.6f}")
|
||||
|
||||
la, lb = pair("mg-log-probs-step0.bin")
|
||||
finite = np.isfinite(la) & np.isfinite(lb)
|
||||
laf = la[finite]; lbf = lb[finite]
|
||||
ld = np.abs(laf - lbf)
|
||||
print(f"[Cossim] Step0LogProbs cos: {cos(laf, lbf):.6f} max_abs_diff: {ld.max():.6f} mean_abs_diff: {ld.mean():.6f}")
|
||||
|
||||
ta, tb = pair("mg-tokens.bin")
|
||||
n = min(ta.size, tb.size)
|
||||
ai = ta.astype(np.int64).ravel()[:n]
|
||||
bi = tb.astype(np.int64).ravel()[:n]
|
||||
print(f"[Cossim] Tokens: {cos(ta, tb):.6f} exact: {100.0 * float(np.mean(ai == bi)):.2f}%")
|
||||
|
||||
aa, ab = pair("output-audio.bin")
|
||||
print(f"[Cossim] Audio: {cos(aa, ab):.6f}")
|
||||
|
||||
n = min(audio_cpp.size, audio_pt.size)
|
||||
print(f"[Cossim] WAV stft_cos: {stft_cos(audio_cpp[:n], audio_pt[:n]):.6f} samples: {n}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user