Upload scripts/compare_onnx_vs_pytorch_parity.py with huggingface_hub

scripts/compare_onnx_vs_pytorch_parity.py

@@ -0,0 +1,485 @@
+"""
+compare_onnx_vs_pytorch_parity.py
+
+Parity-first Finnish Chatterbox comparison harness.
+
+Purpose:
+- Keep `Chatterbox-Finnish/inference_example.py` as the behavioral source of truth.
+- Compare that PyTorch path against an explicit ONNX path that mirrors the future
+  browser worker logic as closely as possible.
+- Avoid mixing debug shortcuts into the default comparison path.
+
+Notes:
+- This script requires CUDA for ONNX Runtime. If CUDA is unavailable, it fails
+  loudly instead of silently falling back to CPU.
+- The ONNX path still uses precomputed Finnish conditioning because the Finnish
+  `cond_enc` path is not yet fully packaged for the browser runtime.
+"""
+
+import os
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+import requests
+import soundfile as sf
+
+
+CONFIG = {
+    "MODE": "parity",
+    "TEXT": "Tervetuloa kokeilemaan hienoviritettyä suomenkielistä Chatterbox-puhesynteesiä.",
+    "REFERENCE_AUDIO": "Chatterbox-Finnish/samples/reference_finnish.wav",
+    "FINETUNED_WEIGHTS": "Chatterbox-Finnish/models/best_finnish_multilingual_cp986.safetensors",
+    "PRETRAINED_DIR": "Chatterbox-Finnish/pretrained_models",
+    "OUT_DIR": "_cmp",
+    "ONNX_CACHE_DIR": "_onnx_cache",
+    "HF_BASE_REPO": "onnx-community/chatterbox-multilingual-ONNX",
+    "HF_FINNISH_REPO": "RASMUS/Chatterbox-Finnish-ONNX",
+    "SEED": 42,
+    "REPETITION_PENALTY": 1.2,
+    "TEMPERATURE": 0.8,
+    "EXAGGERATION": 0.6,
+    "CFG_WEIGHT": 0.5,
+    "MIN_P": 0.05,
+    "MIN_SPEECH_TOKENS": 40,
+    "MAX_GENERATION_STEPS": 800,
+    "RUN_TRANSCRIPTION": True,
+    "RUN_ANALYZE_AUDIO": False,
+}
+
+
+OUT_DIR = Path(CONFIG["OUT_DIR"])
+OUT_DIR.mkdir(exist_ok=True)
+CACHE_DIR = Path(CONFIG["ONNX_CACHE_DIR"])
+CACHE_DIR.mkdir(exist_ok=True)
+
+START_SPEECH = 6561
+STOP_SPEECH = 6562
+SOT_TEXT = 255
+EOT_TEXT = 0
+
+
+def seed_everything(seed: int) -> None:
+    import torch
+
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def hf_download(repo_id: str, filename: str) -> str:
+    from huggingface_hub import hf_hub_download
+
+    return hf_hub_download(
+        repo_id=repo_id,
+        filename=filename,
+        local_dir=str(CACHE_DIR),
+        local_dir_use_symlinks=False,
+    )
+
+
+def require_cuda_providers():
+    import onnxruntime as ort
+
+    providers = ort.get_available_providers()
+    if "CUDAExecutionProvider" not in providers:
+        raise RuntimeError(
+            "CUDAExecutionProvider not available. "
+            "Set LD_LIBRARY_PATH to the conda env cuDNN path before running."
+        )
+    return ["CUDAExecutionProvider", "CPUExecutionProvider"]
+
+
+def save_wav(arr: np.ndarray, path: str, sr: int) -> None:
+    sf.write(path, arr, sr)
+    duration = len(arr) / sr
+    peak = float(np.abs(arr).max())
+    print(f"  saved -> {path}  ({duration:.2f}s, peak={peak:.4f})")
+
+
+def transcribe(wav_path: str, lang: str = "fi") -> str:
+    groq_key = os.environ.get("GROQ_API_KEY", "")
+    if not groq_key:
+        return "(no GROQ_API_KEY)"
+
+    with open(wav_path, "rb") as handle:
+        response = requests.post(
+            "https://api.groq.com/openai/v1/audio/transcriptions",
+            headers={"Authorization": f"Bearer {groq_key}"},
+            files={"file": (os.path.basename(wav_path), handle, "audio/wav")},
+            data={"model": "whisper-large-v3", "language": lang, "response_format": "text"},
+            timeout=300,
+        )
+    response.raise_for_status()
+    return response.text.strip()
+
+
+def apply_rep_penalty(logits: np.ndarray, generated: list[int], penalty: float) -> np.ndarray:
+    updated = logits.copy()
+    for token in set(generated):
+        if updated[token] > 0:
+            updated[token] /= penalty
+        else:
+            updated[token] *= penalty
+    return updated
+
+
+def apply_min_p(logits: np.ndarray, min_p: float) -> np.ndarray:
+    updated = logits.copy()
+    probs = np.exp(updated - updated.max())
+    probs /= probs.sum()
+    updated[probs < probs.max() * min_p] = -1e9
+    return updated
+
+
+def sample_with_temperature(logits: np.ndarray, temperature: float) -> int:
+    scaled = logits / temperature
+    scaled -= scaled.max()
+    probs = np.exp(scaled)
+    probs /= probs.sum()
+    return int(np.random.choice(len(probs), p=probs))
+
+
+def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
+    flat_a = np.asarray(a).reshape(-1)
+    flat_b = np.asarray(b).reshape(-1)
+    denom = np.linalg.norm(flat_a) * np.linalg.norm(flat_b) + 1e-12
+    return float(np.dot(flat_a, flat_b) / denom)
+
+
+def run_pytorch() -> str:
+    print("\n" + "=" * 64)
+    print("1. PYTORCH INFERENCE")
+    print("=" * 64)
+
+    import torch
+    from safetensors.torch import load_file
+
+    sys.path.insert(0, "Chatterbox-Finnish")
+    from src.chatterbox_.tts import ChatterboxTTS
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    engine = ChatterboxTTS.from_local(CONFIG["PRETRAINED_DIR"], device=device)
+    checkpoint = load_file(CONFIG["FINETUNED_WEIGHTS"])
+    t3_state = {k[3:] if k.startswith("t3.") else k: v for k, v in checkpoint.items()}
+    engine.t3.load_state_dict(t3_state, strict=False)
+
+    start = time.time()
+    wav = engine.generate(
+        text=CONFIG["TEXT"],
+        audio_prompt_path=CONFIG["REFERENCE_AUDIO"],
+        repetition_penalty=CONFIG["REPETITION_PENALTY"],
+        temperature=CONFIG["TEMPERATURE"],
+        exaggeration=CONFIG["EXAGGERATION"],
+        cfg_weight=CONFIG["CFG_WEIGHT"],
+        min_p=CONFIG["MIN_P"],
+    )
+    elapsed = time.time() - start
+
+    array = wav.squeeze().cpu().numpy()
+    output_path = str(OUT_DIR / "pytorch_output.wav")
+    save_wav(array, output_path, engine.sr)
+    print(f"  inference time: {elapsed:.1f}s")
+    return output_path
+
+
+def run_onnx() -> str:
+    print("\n" + "=" * 64)
+    print("2. ONNX INFERENCE")
+    print("=" * 64)
+
+    import librosa
+    import onnxruntime as ort
+
+    providers = require_cuda_providers()
+    options = ort.SessionOptions()
+    options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+
+    speech_encoder_path = hf_download(CONFIG["HF_BASE_REPO"], "onnx/speech_encoder.onnx")
+    hf_download(CONFIG["HF_BASE_REPO"], "onnx/speech_encoder.onnx_data")
+    embed_tokens_path = hf_download(CONFIG["HF_BASE_REPO"], "onnx/embed_tokens.onnx")
+    hf_download(CONFIG["HF_BASE_REPO"], "onnx/embed_tokens.onnx_data")
+    conditional_decoder_path = hf_download(CONFIG["HF_BASE_REPO"], "onnx/conditional_decoder.onnx")
+    hf_download(CONFIG["HF_BASE_REPO"], "onnx/conditional_decoder.onnx_data")
+    language_model_path = hf_download(CONFIG["HF_FINNISH_REPO"], "onnx/language_model.onnx")
+    hf_download(CONFIG["HF_FINNISH_REPO"], "onnx/language_model.onnx_data")
+
+    print("  loading sessions...")
+    sess_se = ort.InferenceSession(speech_encoder_path, sess_options=options, providers=providers)
+    sess_et = ort.InferenceSession(embed_tokens_path, sess_options=options, providers=providers)
+    sess_lm = ort.InferenceSession(language_model_path, sess_options=options, providers=providers)
+    sess_cd = ort.InferenceSession(conditional_decoder_path, sess_options=options, providers=providers)
+
+    cond_emb_path = hf_download(CONFIG["HF_FINNISH_REPO"], "onnx/finnish_cond_emb.bin")
+    with open(cond_emb_path, "rb") as handle:
+        cond_emb = np.frombuffer(handle.read(), dtype=np.float32).reshape(1, 34, 1024)
+    print(f"  cond_emb: {cond_emb.shape}")
+
+    sys.path.insert(0, "Chatterbox-Finnish")
+    from src.chatterbox_.models.tokenizers.tokenizer import EnTokenizer
+    from src.chatterbox_.tts import punc_norm
+
+    tokenizer = EnTokenizer(os.path.join(CONFIG["PRETRAINED_DIR"], "tokenizer.json"))
+    normalized_text = punc_norm(CONFIG["TEXT"])
+    token_ids = tokenizer.encode(normalized_text)
+    text_ids = np.array([[SOT_TEXT] + token_ids + [EOT_TEXT]], dtype=np.int64)
+    print(f"  text tokens: {text_ids.shape}")
+
+    # Use 24kHz reference audio for speech_encoder, matching the official ONNX pipeline.
+    ref_24k, _ = librosa.load(CONFIG["REFERENCE_AUDIO"], sr=24000)
+    se_out = sess_se.run(None, {"audio_values": ref_24k[np.newaxis, :].astype(np.float32)})
+    prompt_tokens = se_out[1]
+    speaker_embeddings = se_out[2]
+    speaker_features = se_out[3]
+    print(f"  prompt_tokens: {prompt_tokens.shape}")
+
+    exaggeration = np.array([CONFIG["EXAGGERATION"]], dtype=np.float32)
+    text_pos = np.arange(text_ids.shape[1], dtype=np.int64)[np.newaxis, :]
+    text_embeds = sess_et.run(
+        None,
+        {"input_ids": text_ids, "position_ids": text_pos, "exaggeration": exaggeration},
+    )[0]
+
+    bos_emb = sess_et.run(
+        None,
+        {
+            "input_ids": np.array([[START_SPEECH]], dtype=np.int64),
+            "position_ids": np.array([[0]], dtype=np.int64),
+            "exaggeration": exaggeration,
+        },
+    )[0]
+
+    prefill_cond = np.concatenate([cond_emb, text_embeds, bos_emb], axis=1)
+    prefill_uncond = np.concatenate([cond_emb, np.zeros_like(text_embeds), bos_emb], axis=1)
+
+    kv_meta = next(inp for inp in sess_lm.get_inputs() if inp.name == "past_key_values.0.key")
+    kv_dtype = np.float16 if "float16" in kv_meta.type else np.float32
+    kv_empty = np.zeros((1, 16, 0, 64), dtype=kv_dtype)
+    layer_count = 30
+    kv_cond = [(kv_empty.copy(), kv_empty.copy()) for _ in range(layer_count)]
+    kv_uncond = [(kv_empty.copy(), kv_empty.copy()) for _ in range(layer_count)]
+
+    def make_kv_feeds(kv_cache):
+        feeds = {}
+        for layer_index in range(layer_count):
+            feeds[f"past_key_values.{layer_index}.key"] = kv_cache[layer_index][0]
+            feeds[f"past_key_values.{layer_index}.value"] = kv_cache[layer_index][1]
+        return feeds
+
+    def lm_step(inputs_embeds: np.ndarray, attention_mask: np.ndarray, kv_cache):
+        feeds = {"inputs_embeds": inputs_embeds, "attention_mask": attention_mask}
+        feeds.update(make_kv_feeds(kv_cache))
+        outputs = sess_lm.run(None, feeds)
+        logits = outputs[0]
+        next_kv = [(outputs[1 + 2 * idx], outputs[2 + 2 * idx]) for idx in range(layer_count)]
+        return logits, next_kv
+
+    mask_cond = np.ones((1, prefill_cond.shape[1]), dtype=np.int64)
+    mask_uncond = np.ones((1, prefill_uncond.shape[1]), dtype=np.int64)
+    logits_cond, kv_cond = lm_step(prefill_cond, mask_cond, kv_cond)
+    logits_uncond, kv_uncond = lm_step(prefill_uncond, mask_uncond, kv_uncond)
+
+    generated = [START_SPEECH]
+    speech_tokens = []
+
+    def choose_next_token(logits_c: np.ndarray, logits_u: np.ndarray) -> int:
+        cond_last = logits_c[0, -1].astype(np.float32)
+        uncond_last = logits_u[0, -1].astype(np.float32)
+        merged = cond_last + CONFIG["CFG_WEIGHT"] * (cond_last - uncond_last)
+        merged = apply_rep_penalty(merged, generated, CONFIG["REPETITION_PENALTY"])
+        merged = apply_min_p(merged, CONFIG["MIN_P"])
+        return sample_with_temperature(merged, CONFIG["TEMPERATURE"])
+
+    first_token = choose_next_token(logits_cond, logits_uncond)
+    generated.append(first_token)
+    if first_token < START_SPEECH:
+        speech_tokens.append(first_token)
+
+    print("  generating...")
+    generation_start = time.time()
+    for step in range(1, CONFIG["MAX_GENERATION_STEPS"]):
+        step_emb = sess_et.run(
+            None,
+            {
+                "input_ids": np.array([[generated[-1]]], dtype=np.int64),
+                "position_ids": np.array([[step]], dtype=np.int64),
+                "exaggeration": exaggeration,
+            },
+        )[0]
+
+        step_mask_cond = np.ones((1, kv_cond[0][0].shape[2] + 1), dtype=np.int64)
+        step_mask_uncond = np.ones((1, kv_uncond[0][0].shape[2] + 1), dtype=np.int64)
+
+        logits_cond, kv_cond = lm_step(step_emb, step_mask_cond, kv_cond)
+        logits_uncond, kv_uncond = lm_step(step_emb, step_mask_uncond, kv_uncond)
+        next_token = choose_next_token(logits_cond, logits_uncond)
+
+        if next_token == STOP_SPEECH and len(speech_tokens) >= CONFIG["MIN_SPEECH_TOKENS"]:
+            print(f"  EOS at step {step} ({len(speech_tokens)} speech tokens)")
+            break
+
+        generated.append(next_token)
+        if next_token < START_SPEECH:
+            speech_tokens.append(next_token)
+
+        if (step + 1) % 100 == 0:
+            elapsed = time.time() - generation_start
+            rate = (step + 1) / elapsed
+            print(f"  step {step + 1}: {len(speech_tokens)} speech tokens ({rate:.1f} tok/s)")
+
+    print(f"  generation time: {time.time() - generation_start:.1f}s")
+
+    generated_arr = np.array([speech_tokens], dtype=np.int64)
+    decoder_tokens = np.concatenate([prompt_tokens, generated_arr], axis=1)
+    print(f"  decoder input: {decoder_tokens.shape}")
+    wav = sess_cd.run(
+        None,
+        {
+            "speech_tokens": decoder_tokens,
+            "speaker_embeddings": speaker_embeddings,
+            "speaker_features": speaker_features,
+        },
+    )[0].squeeze().astype(np.float32)
+
+    peak = float(np.abs(wav).max())
+    if peak < 0.01 and peak > 0:
+        wav = wav * (0.9 / peak)
+    wav = np.clip(wav, -1.0, 1.0)
+
+    output_path = str(OUT_DIR / "onnx_output_parity.wav")
+    save_wav(wav, output_path, 24000)
+    return output_path
+
+
+def compare_outputs(pytorch_wav: str, onnx_wav: str) -> None:
+    print("\n" + "=" * 64)
+    print("3. OUTPUT COMPARISON")
+    print("=" * 64)
+
+    if CONFIG["RUN_TRANSCRIPTION"]:
+        print(f"  ref text: {CONFIG['TEXT']}")
+        print(f"  PyTorch: {transcribe(pytorch_wav)}")
+        print(f"  ONNX:    {transcribe(onnx_wav)}")
+
+    if CONFIG["RUN_ANALYZE_AUDIO"]:
+        import subprocess
+
+        subprocess.run(
+            [
+                sys.executable,
+                "analyze_audio.py",
+                pytorch_wav,
+                onnx_wav,
+                "--label-a",
+                "PyTorch",
+                "--label-b",
+                "ONNX",
+                "--ref-text",
+                CONFIG["TEXT"],
+                "--lang",
+                "fi",
+            ],
+            check=False,
+        )
+
+
+def run_debug() -> None:
+    print("\n" + "=" * 64)
+    print("4. COMPONENT DEBUG")
+    print("=" * 64)
+
+    import librosa
+    import onnxruntime as ort
+    import torch
+    from safetensors.torch import load_file
+
+    providers = require_cuda_providers()
+    options = ort.SessionOptions()
+    options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+
+    sys.path.insert(0, "Chatterbox-Finnish")
+    from src.chatterbox_.models.tokenizers.tokenizer import EnTokenizer
+    from src.chatterbox_.tts import ChatterboxTTS, punc_norm
+
+    checkpoint = load_file(CONFIG["FINETUNED_WEIGHTS"])
+    t3_state = {k[3:] if k.startswith("t3.") else k: v for k, v in checkpoint.items()}
+
+    engine = ChatterboxTTS.from_local(CONFIG["PRETRAINED_DIR"], device="cuda")
+    engine.t3.load_state_dict(t3_state, strict=False)
+
+    speech_encoder_path = hf_download(CONFIG["HF_BASE_REPO"], "onnx/speech_encoder.onnx")
+    hf_download(CONFIG["HF_BASE_REPO"], "onnx/speech_encoder.onnx_data")
+    embed_tokens_path = hf_download(CONFIG["HF_BASE_REPO"], "onnx/embed_tokens.onnx")
+    hf_download(CONFIG["HF_BASE_REPO"], "onnx/embed_tokens.onnx_data")
+
+    sess_se = ort.InferenceSession(speech_encoder_path, sess_options=options, providers=providers)
+    sess_et = ort.InferenceSession(embed_tokens_path, sess_options=options, providers=providers)
+
+    ref_24k, _ = librosa.load(CONFIG["REFERENCE_AUDIO"], sr=24000)
+    se_out = sess_se.run(None, {"audio_values": ref_24k[np.newaxis, :].astype(np.float32)})
+    onnx_prompt_tokens = se_out[1]
+    onnx_speaker_embeddings = se_out[2]
+    onnx_speaker_features = se_out[3]
+
+    with torch.no_grad():
+        engine.prepare_conditionals(CONFIG["REFERENCE_AUDIO"], exaggeration=CONFIG["EXAGGERATION"])
+        pytorch_ref = engine.conds.gen
+        pt_prompt_tokens = pytorch_ref["prompt_token"].cpu().numpy()
+        pt_speaker_features = pytorch_ref["prompt_feat"].cpu().numpy()
+        pt_speaker_embeddings = pytorch_ref["embedding"].cpu().numpy()
+
+    print("  speech_encoder vs PyTorch conditionals")
+    print(f"    prompt_tokens exact match: {np.array_equal(onnx_prompt_tokens, pt_prompt_tokens)}")
+    print(f"    speaker_embeddings cosine: {cosine_similarity(onnx_speaker_embeddings, pt_speaker_embeddings):.6f}")
+    print(f"    onnx speaker_features: {onnx_speaker_features.shape}")
+    print(f"    pytorch speaker_features: {pt_speaker_features.shape}")
+    if onnx_speaker_features.shape == pt_speaker_features.shape:
+        max_diff = float(np.abs(onnx_speaker_features - pt_speaker_features).max())
+        print(f"    speaker_features max diff: {max_diff:.6f}")
+
+    tokenizer = EnTokenizer(os.path.join(CONFIG["PRETRAINED_DIR"], "tokenizer.json"))
+    normalized_text = punc_norm(CONFIG["TEXT"])
+    onnx_ids = [SOT_TEXT] + tokenizer.encode(normalized_text) + [EOT_TEXT]
+
+    with torch.no_grad():
+        pt_ids = engine.tokenizer.text_to_tokens(normalized_text)[0].tolist()
+        pt_ids = [engine.t3.hp.start_text_token] + pt_ids + [engine.t3.hp.stop_text_token]
+
+    print("  text tokenization")
+    print(f"    exact match: {onnx_ids == pt_ids}")
+    print(f"    onnx ids head: {onnx_ids[:8]}")
+    print(f"    pytorch ids head: {pt_ids[:8]}")
+
+    exaggeration = np.array([CONFIG["EXAGGERATION"]], dtype=np.float32)
+    onnx_ids_arr = np.array([onnx_ids], dtype=np.int64)
+    pos_ids = np.arange(len(onnx_ids), dtype=np.int64)[np.newaxis, :]
+    onnx_embeds = sess_et.run(
+        None,
+        {"input_ids": onnx_ids_arr, "position_ids": pos_ids, "exaggeration": exaggeration},
+    )[0]
+
+    with torch.no_grad():
+        pt_embeds = engine.t3.text_emb(torch.tensor(onnx_ids_arr, device="cuda")).cpu().numpy()
+
+    print("  embed_tokens")
+    print(f"    cosine: {cosine_similarity(onnx_embeds, pt_embeds):.6f}")
+    print(f"    max diff: {float(np.abs(onnx_embeds - pt_embeds).max()):.6f}")
+
+
+def main() -> None:
+    seed_everything(CONFIG["SEED"])
+    if CONFIG["MODE"] == "debug":
+        run_debug()
+        print("\nDone. Debug checks completed.")
+        return
+
+    pytorch_wav = run_pytorch()
+    onnx_wav = run_onnx()
+    compare_outputs(pytorch_wav, onnx_wav)
+    print(f"\nDone. Outputs saved in {OUT_DIR}/")
+
+
+if __name__ == "__main__":
+    main()