Add scripts/browser_pipeline_sim.py

scripts/browser_pipeline_sim.py

@@ -0,0 +1,303 @@
+"""
+browser_pipeline_sim.py
+
+Python simulation of the browser worker.ts logic.
+Mirrors the EXACT token assembly, CFG loop, sampling, and decoding
+that the browser worker does — making it trivial to diff against worker.ts.
+
+If this produces correct audio, worker.ts with the same logic will too.
+If there's a bug in the browser, you can find it by diffing this script against worker.ts.
+
+Usage:
+    cd /workspaces/work
+    conda run -n chatterbox-onnx python browser_pipeline_sim.py [text] [ref_audio]
+
+Outputs:
+    _cmp/browser_sim_output.wav
+"""
+
+import os, sys, time, json, struct, math
+import numpy as np
+import onnxruntime as ort
+import soundfile as sf
+import librosa
+import requests
+from pathlib import Path
+from huggingface_hub import hf_hub_download
+
+sys.path.insert(0, "Chatterbox-Finnish")
+
+# ── Config — mirrors browser worker constants ─────────────────────────────────
+TEXT       = sys.argv[1] if len(sys.argv) > 1 else \
+             "Tervetuloa kokeilemaan hienoviritettyä suomenkielistä Chatterbox-puhesynteesiä."
+REF_AUDIO  = sys.argv[2] if len(sys.argv) > 2 else \
+             "Chatterbox-Finnish/samples/reference_finnish.wav"
+
+HF_BASE    = "onnx-community/chatterbox-multilingual-ONNX"
+HF_FI      = "RASMUS/Chatterbox-Finnish-ONNX"
+CACHE_DIR  = "_onnx_cache"
+OUT_DIR    = Path("_cmp"); OUT_DIR.mkdir(exist_ok=True)
+
+# Worker constants (from worker.ts)
+SOT            = 255    # [START] token
+EOT            = 0      # [STOP] token
+START_SPEECH   = 6561   # BOS speech token
+STOP_SPEECH    = 6562   # EOS speech token
+CFG_WEIGHT     = 0.5
+REP_PENALTY    = 1.2
+TEMPERATURE    = 0.8
+EXAGGERATION   = 0.6
+MIN_SPEECH_TOKENS = 40
+MAX_STEPS      = 1000
+SAMPLE_RATE    = 24000
+
+GROQ_KEY = os.environ.get("GROQ_API_KEY", "")
+
+# ── puncNorm — mirrors worker.ts puncNorm() ───────────────────────────────────
+def punc_norm(text: str) -> str:
+    """Mirrors puncNorm() in worker.ts"""
+    import re
+    t = text.strip()
+    t = re.sub(r'\s+', ' ', t)                         # normalize whitespace
+    t = t[0].upper() + t[1:] if t else t               # capitalize first letter
+    t = t.replace(' .', '.').replace(' ,', ',')        # remove space before punct
+    t = t.replace(' ?', '?').replace(' !', '!')
+    if t and t[-1] not in '.!?…':                      # ensure ending punctuation
+        t += '.'
+    return t
+
+
+# ── EnTokenizer (same as Chatterbox-Finnish/src/chatterbox_/models/tokenizers/tokenizer.py)
+def load_tokenizer(tokenizer_path: str):
+    from src.chatterbox_.models.tokenizers.tokenizer import EnTokenizer
+    return EnTokenizer(tokenizer_path)
+
+
+# ── puncNorm + tokenize → mirrors worker.ts tokenize() ───────────────────────
+def tokenize(tok, text: str) -> list[int]:
+    """
+    Mirrors browser worker tokenize():
+    1. puncNorm
+    2. replace spaces with [SPACE] token
+    3. encode grapheme by grapheme
+    4. wrap with [SOT, ..., EOT]
+    """
+    normed = punc_norm(text)
+    ids = tok.encode(normed)
+    return [SOT] + ids + [EOT]
+
+
+# ── Download helpers ───────────────────────────────────────────────────────────
+def dl(repo_id, filename):
+    return hf_hub_download(repo_id=repo_id, filename=filename,
+                           local_dir=CACHE_DIR, local_dir_use_symlinks=False)
+
+
+# ── KV cache helpers ───────────────────────────────────────────────────────────
+def empty_kv(n_layers=30, kv_dtype=np.float32):
+    return [(np.zeros((1, 16, 0, 64), dtype=kv_dtype),
+             np.zeros((1, 16, 0, 64), dtype=kv_dtype))
+            for _ in range(n_layers)]
+
+
+def make_kv_feeds(kv_cache):
+    feeds = {}
+    for i, (k, v) in enumerate(kv_cache):
+        feeds[f"past_key_values.{i}.key"]   = k
+        feeds[f"past_key_values.{i}.value"] = v
+    return feeds
+
+
+# ── Sampling helpers ───────────────────────────────────────────────────────────
+def apply_rep_penalty(logits, generated_set, penalty):
+    logits = logits.copy()
+    for tok in generated_set:
+        logits[tok] = logits[tok] / penalty if logits[tok] > 0 else logits[tok] * penalty
+    return logits
+
+
+def apply_min_p(logits, p=0.05):
+    """Mirrors applyMinP() in worker.ts"""
+    logits = logits.copy()
+    probs = np.exp(logits - logits.max())
+    probs /= probs.sum()
+    threshold = probs.max() * p
+    logits[probs < threshold] = -1e9
+    return logits
+
+
+def sample_with_temperature(logits, temperature):
+    """Mirrors sampleWithTemperature() in worker.ts"""
+    logits = (logits / temperature).astype(np.float64)
+    logits -= logits.max()
+    probs = np.exp(logits)
+    probs /= probs.sum()
+    return int(np.random.choice(len(probs), p=probs))
+
+
+# ── Main pipeline ─────────────────────────────────────────────────────────────
+def main():
+    providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+    opts = ort.SessionOptions()
+    opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+
+    # ── Load sessions ──
+    print("Loading ONNX sessions...")
+    sess_se = ort.InferenceSession(dl(HF_BASE, "onnx/speech_encoder.onnx"),  opts, providers)
+    sess_et = ort.InferenceSession(dl(HF_BASE, "onnx/embed_tokens.onnx"),    opts, providers)
+    sess_lm = ort.InferenceSession(dl(HF_FI,   "onnx/language_model.onnx"),  opts, providers)
+    sess_cd = ort.InferenceSession(dl(HF_BASE, "onnx/conditional_decoder.onnx"), opts, providers)
+
+    # ── Precomputed Finnish conditioning ──
+    cond_emb_raw = open(dl(HF_FI, "onnx/finnish_cond_emb.bin"), "rb").read()
+    cond_emb = np.frombuffer(cond_emb_raw, dtype=np.float32).reshape(1, 34, 1024)
+    print(f"  cond_emb: {cond_emb.shape}")
+
+    # ── KV dtype from model ──
+    kv_input_name = "past_key_values.0.key"
+    kv_dtype_str = next(
+        (inp.type for inp in sess_lm.get_inputs() if inp.name == kv_input_name), "tensor(float)"
+    )
+    kv_dtype = np.float16 if "float16" in kv_dtype_str else np.float32
+    print(f"  KV cache dtype: {kv_dtype}")
+
+    # ── Step 1: Reference audio → speaker embeddings ──
+    print(f"\nStep 1: speech_encoder  ({REF_AUDIO})")
+    ref_audio, ref_sr = librosa.load(REF_AUDIO, sr=None)
+    ref_16k = librosa.resample(ref_audio, orig_sr=ref_sr, target_sr=16000).astype(np.float32)
+    ref_input = ref_16k[np.newaxis, :]  # [1, T]
+
+    se_outs = sess_se.run(None, {"audio": ref_input})
+    speaker_emb   = se_outs[0]  # [1, 256]
+    prompt_tokens = se_outs[1]  # [1, N]
+    print(f"  speaker_emb: {speaker_emb.shape}")
+    print(f"  prompt_tokens: {prompt_tokens.shape}")
+    # Note: se_outs[2] is cond_emb from base model — we ignore it, use precomputed Finnish version
+
+    # ── Step 2: Tokenize text (mirrors worker.ts tokenize()) ──
+    print(f"\nStep 2: tokenize  '{TEXT[:60]}...'")
+    tok = load_tokenizer("Chatterbox-Finnish/pretrained_models/tokenizer.json")
+    token_ids = tokenize(tok, TEXT)
+    print(f"  token_ids ({len(token_ids)}): {token_ids[:6]}...{token_ids[-3:]}")
+    text_ids_np = np.array([token_ids], dtype=np.int64)  # [1, T]
+
+    # ── Step 3: Embed text tokens ──
+    print(f"\nStep 3: embed_tokens")
+    text_embeds = sess_et.run(None, {"input_ids": text_ids_np})[0]  # [1, T, 1024]
+    print(f"  text_embeds: {text_embeds.shape}")
+
+    # Embed BOS speech token
+    bos_ids = np.array([[START_SPEECH]], dtype=np.int64)
+    bos_emb = sess_et.run(None, {"input_ids": bos_ids})[0]  # [1, 1, 1024]
+
+    # ── Step 4: Build prefill ──
+    # Matches PyTorch: inputs_embeds = cat([cond_emb, text_emb, bos_emb])
+    prefill_cond   = np.concatenate([cond_emb, text_embeds, bos_emb], axis=1)
+    zeros_text     = np.zeros_like(text_embeds)
+    prefill_uncond = np.concatenate([cond_emb, zeros_text, bos_emb], axis=1)
+
+    mask_cond   = np.ones((1, prefill_cond.shape[1]),   dtype=np.int64)
+    mask_uncond = np.ones((1, prefill_uncond.shape[1]), dtype=np.int64)
+
+    kv_empty_layer = np.zeros((1, 16, 0, 64), dtype=kv_dtype)
+    kv_cond   = [(kv_empty_layer.copy(), kv_empty_layer.copy()) for _ in range(30)]
+    kv_uncond = [(kv_empty_layer.copy(), kv_empty_layer.copy()) for _ in range(30)]
+
+    def lm_step(embeds, mask, kv):
+        feeds = {"inputs_embeds": embeds, "attention_mask": mask}
+        feeds.update(make_kv_feeds(kv))
+        outs = sess_lm.run(None, feeds)
+        logits = outs[0]  # [1, seq, vocab]
+        new_kv = [(outs[1 + i*2], outs[1 + i*2 + 1]) for i in range(30)]
+        return logits, new_kv
+
+    # ── Step 5: Prefill both streams ──
+    print(f"\nStep 4: prefill  ({prefill_cond.shape[1]} tokens)")
+    t0 = time.time()
+    logits_c,  kv_cond   = lm_step(prefill_cond,   mask_cond,   kv_cond)
+    logits_uc, kv_uncond = lm_step(prefill_uncond, mask_uncond, kv_uncond)
+    print(f"  prefill done ({time.time()-t0:.1f}s)")
+
+    # ── Step 6: Autoregressive generation ──
+    print(f"\nStep 5: generate (max {MAX_STEPS} steps)")
+    generated = [START_SPEECH]
+    speech_tokens = []
+    t0 = time.time()
+
+    for step in range(MAX_STEPS):
+        last_ids = np.array([[generated[-1]]], dtype=np.int64)
+        last_emb = sess_et.run(None, {"input_ids": last_ids})[0]  # [1, 1, 1024]
+
+        seq_len_c  = kv_cond[0][0].shape[2] + 1
+        seq_len_uc = kv_uncond[0][0].shape[2] + 1
+        mask_c  = np.ones((1, seq_len_c),  dtype=np.int64)
+        mask_uc = np.ones((1, seq_len_uc), dtype=np.int64)
+
+        logits_c,  kv_cond   = lm_step(last_emb, mask_c,  kv_cond)
+        logits_uc, kv_uncond = lm_step(last_emb, mask_uc, kv_uncond)
+
+        # CFG: final = cond + cfg_weight * (cond - uncond)
+        lc  = logits_c[0, -1].astype(np.float32)
+        luc = logits_uc[0, -1].astype(np.float32)
+        final_logits = lc + CFG_WEIGHT * (lc - luc)
+
+        # Apply rep penalty + min_p + temperature sample
+        final_logits = apply_rep_penalty(final_logits, set(generated), REP_PENALTY)
+        final_logits = apply_min_p(final_logits, p=0.05)
+        token = sample_with_temperature(final_logits, TEMPERATURE)
+
+        if token == STOP_SPEECH and len(speech_tokens) >= MIN_SPEECH_TOKENS:
+            print(f"  EOS at step {step} ({len(speech_tokens)} speech tokens)")
+            break
+
+        generated.append(token)
+        if token < START_SPEECH:
+            speech_tokens.append(token)
+
+        if (step + 1) % 100 == 0:
+            elapsed = time.time() - t0
+            rate = (step + 1) / elapsed
+            print(f"  step {step+1:4d}: {len(speech_tokens):3d} speech tokens  ({rate:.1f} tok/s)")
+
+    elapsed = time.time() - t0
+    print(f"  generation done: {len(speech_tokens)} speech tokens in {elapsed:.1f}s")
+
+    # ── Step 7: Decode → waveform ──
+    print(f"\nStep 6: conditional_decoder")
+    speech_tok_arr = np.array([speech_tokens], dtype=np.int64)
+    cd_out = sess_cd.run(None, {
+        "speech_tokens":     speech_tok_arr,
+        "speaker_embeddings": speaker_emb,
+    })
+    wav = cd_out[0].squeeze().astype(np.float32)
+
+    # Normalize (mirrors browser worker)
+    peak = np.abs(wav).max()
+    if peak < 0.01:
+        print(f"  warning: very low amplitude (peak={peak:.4f}), auto-normalizing")
+        wav = wav * (0.9 / peak)
+    wav = np.clip(wav, -1.0, 1.0)
+
+    out_path = str(OUT_DIR / "browser_sim_output.wav")
+    sf.write(out_path, wav, SAMPLE_RATE)
+    print(f"\nSaved: {out_path}  ({len(wav)/SAMPLE_RATE:.2f}s, peak={np.abs(wav).max():.4f})")
+
+    # ── Transcribe ──
+    if GROQ_KEY:
+        print("\nTranscribing with Groq Whisper...")
+        with open(out_path, "rb") as f:
+            r = requests.post(
+                "https://api.groq.com/openai/v1/audio/transcriptions",
+                headers={"Authorization": f"Bearer {GROQ_KEY}"},
+                files={"file": (os.path.basename(out_path), f, "audio/wav")},
+                data={"model": "whisper-large-v3", "language": "fi", "response_format": "text"},
+            )
+        if r.ok:
+            print(f"\nTranscript:  '{r.text.strip()}'")
+            print(f"Target text: '{TEXT}'")
+        else:
+            print(f"  Groq error: {r.status_code} {r.text}")
+
+
+if __name__ == "__main__":
+    main()