From 538eddd74d57ac5a8fea0a8ac0d5288e20678ccb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Zhengyuan=20Su=20=28=E8=8B=8F=E6=94=BF=E6=B8=8A=29?= Date: Wed, 16 Sep 2026 20:49:04 +0800 Subject: [PATCH] [Perf][CosyVoice3] Bounded-window incremental HiFT vocoder streaming (#7521) Signed-off-by: Zhengyuan Su --- .../bench_cosyvoice3_concurrency_matrix.py | 517 ++++++++++++++++++ pyproject.toml | 2 - requirements/common.txt | 5 + .../cosyvoice3/test_cosyvoice3_components.py | 61 ++- .../test_cosyvoice3_incremental_hift.py | 415 ++++++++++++++ .../minicpmo_4_5/test_cuda_graph_wrapper.py | 2 + .../cosyvoice3/code2wav_core/hifigan.py | 170 +++++- .../models/cosyvoice3/cosyvoice3_code2wav.py | 136 +++-- 8 files changed, 1224 insertions(+), 84 deletions(-) create mode 100644 benchmarks/tts/bench_cosyvoice3_concurrency_matrix.py create mode 100644 tests/model_executor/models/cosyvoice3/test_cosyvoice3_incremental_hift.py diff --git a/benchmarks/tts/bench_cosyvoice3_concurrency_matrix.py b/benchmarks/tts/bench_cosyvoice3_concurrency_matrix.py new file mode 100644 index 000000000..fb45b1743 --- /dev/null +++ b/benchmarks/tts/bench_cosyvoice3_concurrency_matrix.py @@ -0,0 +1,517 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""CosyVoice3 streaming e2e benchmark: concurrency x utterance-length matrix. + +Launches the server once per (concurrency, step-size) combo, streams TTS +requests cycling through several utterance lengths, and reports +TTFA/E2E/RTF/per-chunk latency (mean, stdev, range) per cell. "Step size" is +``codec_chunk_frames``, the token-hop length the first emitted chunk waits +for; later chunks wait for a doubling hop capped at ``4 * codec_chunk_frames`` +(see ``_token_hop_lengths``). Per-chunk latency is reconstructed from that +same ramp rather than from raw transport read boundaries, which don't align +to model chunks (see ``one_request``'s docstring). + +Running with no flags reproduces the default matrix. Per-request correctness +is checked: non-silent, long enough, and (with a fixed ``--seed``) +byte-identical PCM across requests for the same text. +""" + +import argparse +import json +import os +import signal +import statistics +import subprocess +import sys +import threading +import time +from pathlib import Path + +import httpx +import numpy as np + +REPO = Path(__file__).resolve().parents[2] +DEFAULT_MODEL = "FunAudioLLM/Fun-CosyVoice3-0.5B-2512" + +DEFAULT_REF_AUDIO = "https://raw.githubusercontent.com/FunAudioLLM/CosyVoice/main/asset/zero_shot_prompt.wav" +DEFAULT_REF_TEXT = "希望你以后能够做的比我还好呦。" + +DEFAULT_TEXTS = [ + "Thank you so much for the birthday gift you sent me.", + "The weather is lovely today, with warm sunshine and a gentle breeze. " + "We decided to take a walk in the park to enjoy the spring flowers and the fresh air.", + "The weather is lovely today, with warm sunshine and a gentle breeze. We decided to take a walk in the " + "park to enjoy the spring flowers and the fresh air. Many people were out walking, children were " + "running and playing on the grass, and older folks sat on benches chatting with one another. We " + "strolled slowly along the path and came across a cherry tree in full bloom, its pink petals drifting " + "on the wind. It was so beautiful that we stopped to take some photos to remember the moment. Then we " + "continued on and reached a small lake, its water so clear you could see straight to the bottom, with " + "the blue sky and white clouds reflected on its surface. A few ducks paddled back and forth across the " + "water, occasionally letting out a soft quack. We sat on a bench by the lake for a while, enjoying the " + "peaceful afternoon. Time passed quickly, and before we knew it evening had arrived, with the setting " + "sun painting the sky in shades of gold. Reluctant to leave, we finally headed home, already planning " + "our next visit. On the way back we talked about everything we had seen and heard that day, feeling " + "wonderfully happy. It was truly an unforgettable weekend.", +] + +# PCM16 mono @ 24 kHz (CosyVoice3 output rate) +SAMPLE_RATE = 24000 + +# Well below any real speech's RMS; catches all-zero/near-silent PCM that byte +# count and duration checks alone pass through. +MIN_RMS = 50.0 + +# Overlay: only codec_chunk_frames is overridden; must live next to cosyvoice3.yaml (relative base_config) +DEPLOY_TEMPLATE = """# Auto-generated by bench_cosyvoice3_concurrency_matrix.py -- do not edit. +base_config: cosyvoice3.yaml + +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + codec_streaming: true + connector_get_sleep_s: 0.01 + connector_get_max_wait_first_chunk: 3000 + connector_get_max_wait: 300 + codec_chunk_frames: {chunk_frames} + codec_max_chunk_frames: {max_chunk_frames} + codec_stream_scale_factor: {stream_scale_factor} + codec_left_context_frames: 25 + codec_vocab_size: 6561 +""" + +# CosyVoice3Config constants (transformers_utils/configs/cosyvoice3.py): tokens +# are consumed at token_mel_ratio mel frames each, and HiFT emits +# samples_per_mel PCM16 samples per mel frame. Pinned here (not left to the +# server's defaults) since one_request's byte-boundary reconstruction below +# depends on knowing them exactly. +TOKEN_MEL_RATIO = 2 +SAMPLES_PER_MEL = 480 +STREAM_SCALE_FACTOR = 2 + + +def _stats(values: list[float]) -> dict[str, float]: + """Mean/median/stdev/p95/min/max/n for a list of per-request measurements.""" + n = len(values) + if n == 0: + return { + "n": 0, + "mean": float("nan"), + "median": float("nan"), + "stdev": float("nan"), + "p95": float("nan"), + "min": float("nan"), + "max": float("nan"), + } + return { + "n": n, + "mean": statistics.mean(values), + "median": statistics.median(values), + "stdev": statistics.stdev(values) if n >= 2 else 0.0, + "p95": statistics.quantiles(values, n=20)[18] if n >= 2 else values[0], + "min": min(values), + "max": max(values), + } + + +def vllm_bin() -> str: + for c in (REPO / ".venv/bin/vllm", Path(sys.executable).parent / "vllm"): + if c.is_file(): + return str(c) + return "vllm" + + +def wait_ready(url: str, proc: subprocess.Popen, log: Path, timeout: float) -> None: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if proc.poll() is not None: + tail = log.read_text(errors="replace").splitlines()[-40:] if log.exists() else [] + raise RuntimeError("server exited during startup:\n " + "\n ".join(tail)) + try: + if httpx.get(f"{url}/health", timeout=2.0).status_code == 200: + return + except Exception: + pass + time.sleep(2.0) + raise TimeoutError(f"server not ready within {timeout}s (log: {log})") + + +def _token_hop_lengths(chunk_frames: int, max_chunk_frames: int): + """The token-hop-length ramp talker2code2wav_async_chunk uses to decide + when to fire the next flow-solve/emission: chunk_frames, then doubling + (STREAM_SCALE_FACTOR) each step, capped at max_chunk_frames. This -- not + codec_chunk_frames alone -- is what actually paces emitted audio chunks: + with the defaults here, hops are chunk_frames, 2x, 4x, 4x, 4x, ... + + The very first emission is consistently one token short of chunk_frames + (confirmed against real speech.audio.delta event sizes) even though the + ramp's own progression is computed from the nominal chunk_frames, not the + actual first hop -- so only the first yield gets the -1 adjustment. + """ + yield chunk_frames - 1 + hop = chunk_frames + while True: + hop = min(max_chunk_frames, max(chunk_frames, hop * STREAM_SCALE_FACTOR)) + yield hop + + +def one_request(client, url, model, text, ref_audio, ref_text, chunk_frames, seed=None): + """Run one streaming request, returning TTFA, E2E, per-chunk gaps, and audio bytes. + + ``response_format=pcm`` + ``stream_format=audio`` streams raw bytes with no + framing, so a client read boundary alone is a transport artifact, not a + model chunk boundary. Reconstructed instead by replicating the server's + own token-hop-length ramp (see ``_token_hop_lengths``) to compute each + real chunk's expected byte size ahead of time, and slicing the + accumulated byte stream at those offsets -- independent of how the + transport happens to fragment the underlying reads. The final chunk + (finalize, which releases extra held-back frames) is whatever remains + when the stream closes, not a predicted hop. + """ + payload = { + "model": model, + "input": text, + "response_format": "pcm", + "stream": True, + "stream_format": "audio", + "ref_audio": ref_audio, + "ref_text": ref_text, + } + if seed is not None: + payload["seed"] = seed + t0 = time.perf_counter() + ttfa = None + nbytes = 0 + audio_bytes = bytearray() + chunk_gaps_ms = [] # inter-chunk arrival gaps (per model chunk), by step index + prev_t = None + step_idx = 0 + + bytes_per_token = TOKEN_MEL_RATIO * SAMPLES_PER_MEL * 2 + hop_gen = _token_hop_lengths(chunk_frames, 4 * chunk_frames) + next_boundary = next(hop_gen) * bytes_per_token + last_crossed_bytes = 0 + received = 0 + last_now = t0 + + with client.stream("POST", f"{url}/v1/audio/speech", json=payload, headers={"Authorization": "Bearer EMPTY"}) as r: + if r.status_code != 200: + raise RuntimeError(f"{r.status_code}: {r.read()[:300]!r}") + for chunk in r.iter_bytes(): + if not chunk: + continue + now = time.perf_counter() + if ttfa is None: + ttfa = now - t0 + last_now = now + received += len(chunk) + nbytes += len(chunk) + audio_bytes += chunk + while received >= next_boundary: + if prev_t is not None: + chunk_gaps_ms.append((step_idx, (now - prev_t) * 1e3)) + step_idx += 1 + prev_t = now + last_crossed_bytes = next_boundary + next_boundary += next(hop_gen) * bytes_per_token + + if received > last_crossed_bytes: + # Finalize chunk: releases extra held-back frames, so it doesn't + # match a predicted hop size -- its arrival is the stream closing. + if prev_t is not None: + chunk_gaps_ms.append((step_idx, (last_now - prev_t) * 1e3)) + step_idx += 1 + + e2e = time.perf_counter() - t0 + audio_s = nbytes / 2 / SAMPLE_RATE + return { + "ttfa_s": ttfa, + "e2e_s": e2e, + "audio_s": audio_s, + "rtf": e2e / audio_s if audio_s > 0 else float("nan"), + "chunk_gaps_ms": chunk_gaps_ms, + "nbytes": nbytes, + "audio_bytes": audio_bytes, + } + + +def check_correctness(r: dict, min_audio_s: float) -> list[str]: + """Return a list of correctness violations (empty = pass).""" + problems = [] + if r["nbytes"] <= 0: + problems.append("empty audio (0 bytes)") + if r["audio_s"] < min_audio_s: + problems.append(f"audio too short ({r['audio_s']:.2f}s < {min_audio_s:.2f}s)") + if r["nbytes"] > 0: + samples = np.frombuffer(bytes(r["audio_bytes"]), dtype=np.int16).astype(np.float64) + rms = float(np.sqrt(np.mean(samples**2))) + if rms < MIN_RMS: + problems.append(f"near-silent audio (rms={rms:.1f} < {MIN_RMS})") + return problems + + +def drive(url, args, n, conc, label, results, lock, idx, chunk_frames): + """Run ``n`` requests at concurrency ``conc``, cycling through ``args.texts``.""" + texts = args.texts + + def worker(): + with httpx.Client(timeout=600.0) as c: + while True: + with lock: + i = idx[0] + idx[0] += 1 + if i >= n: + return + text = texts[i % len(texts)] + try: + r = one_request( + c, url, args.model, text, args.ref_audio, args.ref_text, chunk_frames, seed=args.seed + ) + r["text"] = text + with lock: + results.append(r) + except Exception as e: + print(f" [{label}] request {i} failed: {e}", flush=True) + + ts = [threading.Thread(target=worker) for _ in range(conc)] + t0 = time.perf_counter() + for t in ts: + t.start() + for t in ts: + t.join() + return time.perf_counter() - t0 + + +def run_combo(args, conc: int, chunk_frames: int, out: Path) -> dict: + tag = f"c{conc}_chunk{chunk_frames}" + run_dir = out / tag + run_dir.mkdir(parents=True, exist_ok=True) + log = run_dir / "server.log" + + deploy = REPO / "vllm_omni/deploy" / f"_autogen_steps_{tag}.yaml" + deploy.write_text( + DEPLOY_TEMPLATE.format( + chunk_frames=chunk_frames, + max_chunk_frames=4 * chunk_frames, + stream_scale_factor=STREAM_SCALE_FACTOR, + ) + ) + + cmd = [ + vllm_bin(), + "serve", + args.model, + "--omni", + "--deploy-config", + str(deploy), + "--host", + "127.0.0.1", + "--port", + str(args.port), + "--allowed-local-media-path", + str(REPO), + ] + env = dict(os.environ) + if args.cuda_visible_devices: + env["CUDA_VISIBLE_DEVICES"] = args.cuda_visible_devices + if args.trt is not None: + env["COSYVOICE3_TRT"] = args.trt + + print(f"\n=== concurrency={conc} step_size={chunk_frames} ===", flush=True) + print(f" launching server (log: {log})", flush=True) + with open(log, "w") as lf: + proc = subprocess.Popen(cmd, env=env, stdout=lf, stderr=subprocess.STDOUT, start_new_session=True) + url = f"http://127.0.0.1:{args.port}" + try: + wait_ready(url, proc, log, args.server_timeout) + print(" server ready", flush=True) + + results: list[dict] = [] + lock = threading.Lock() + idx = [0] + if args.warmup: + print(f" warmup ({args.warmup} req)", flush=True) + drive(url, args, args.warmup, min(conc, args.warmup), "warmup", results, lock, idx, chunk_frames) + results.clear() + idx[0] = 0 + + print(f" running {args.num_prompts} req @ concurrency {conc}", flush=True) + wall = drive(url, args, args.num_prompts, conc, "bench", results, lock, idx, chunk_frames) + + ok = [r for r in results if r.get("ttfa_s") is not None] + problems = [] + for r in ok: + problems += check_correctness(r, args.min_audio_s) + # With a fixed seed, same-text requests must produce byte-identical PCM (per text group) + if args.seed is not None: + by_text: dict[str, bytes] = {} + for r in ok: + t = r.get("text") + ref = by_text.get(t) + if ref is None: + by_text[t] = r.get("audio_bytes") + elif r.get("audio_bytes") != ref: + problems.append(f"byte-identity violation for text {t!r}: seeded requests produced different PCM") + n_ok = len(ok) + total_audio_s = sum(r["audio_s"] for r in ok) + all_gaps = [g for r in ok for g in r["chunk_gaps_ms"]] # (step_index, gap_ms) per request + gap_vals = [g for _, g in all_gaps] + + # bucket gaps by chunk position within a request, to see any O(n) slope + by_idx: dict[int, list[float]] = {} + for step_index, g in all_gaps: + by_idx.setdefault(step_index, []).append(g) + step_latency_by_idx = {i: statistics.median(vals) for i, vals in sorted(by_idx.items())} + + agg = { + "concurrency": conc, + "step_size": chunk_frames, + "ok": n_ok, + "errors": args.num_prompts - n_ok, + "correctness_problems": len(problems), + "correctness_sample": problems[:5], + "total_audio_s": total_audio_s, + "throughput_audio_s_per_s": total_audio_s / wall if wall > 0 else float("nan"), + "ttfa_s": _stats([r["ttfa_s"] for r in ok]), + "e2e_s": _stats([r["e2e_s"] for r in ok]), + "rtf": _stats([r["rtf"] for r in ok]), + "step_latency_ms": _stats(gap_vals), + "step_latency_ms_by_idx": step_latency_by_idx, + "num_steps": len(all_gaps), + # per-utterance-length breakdown; n per text = num_prompts / len(texts) + "by_text": { + t: { + "audio_s": _stats([r["audio_s"] for r in ok if r.get("text") == t]), + "e2e_s": _stats([r["e2e_s"] for r in ok if r.get("text") == t]), + "rtf": _stats([r["rtf"] for r in ok if r.get("text") == t]), + } + for t in sorted({r.get("text") for r in ok}) + }, + } + print(json.dumps(agg, indent=2, ensure_ascii=False), flush=True) + return agg + finally: + print(" shutting down server", flush=True) + try: + os.killpg(os.getpgid(proc.pid), signal.SIGTERM) + except Exception: + proc.terminate() + try: + proc.wait(timeout=60) + except subprocess.TimeoutExpired: + try: + os.killpg(os.getpgid(proc.pid), signal.SIGKILL) + except Exception: + proc.kill() + proc.wait(timeout=20) + deploy.unlink(missing_ok=True) + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--model", default=DEFAULT_MODEL) + p.add_argument( + "--text", + dest="texts", + nargs="+", + default=DEFAULT_TEXTS, + help="One or more utterance texts. With several, requests cycle through " + "them so a single server launch covers all utterance lengths; results " + "are bucketed per text.", + ) + p.add_argument("--ref-audio", default=DEFAULT_REF_AUDIO) + p.add_argument("--ref-text", default=DEFAULT_REF_TEXT) + p.add_argument("--concurrency", type=int, nargs="+", default=[1, 8]) + p.add_argument( + "--step-size", + type=int, + nargs="+", + default=[15], + help="codec_chunk_frames values to sweep (the streaming step size).", + ) + p.add_argument("--num-prompts", type=int, default=30) + p.add_argument("--warmup", type=int, default=2) + p.add_argument("--port", type=int, default=8091) + p.add_argument( + "--server-timeout", + type=float, + default=1800.0, + help="Startup budget; first run also downloads weights and builds the TRT engine.", + ) + p.add_argument("--output-dir", default="./perf_steps") + p.add_argument("--cuda-visible-devices", default="0") + p.add_argument( + "--trt", + choices=["0", "1"], + default=None, + help="Pin COSYVOICE3_TRT. 0=force torch estimator, 1=allow TensorRT. Default: unset.", + ) + p.add_argument( + "--min-audio-s", + type=float, + default=1.0, + help="Minimum expected audio seconds per request for the correctness check.", + ) + p.add_argument( + "--seed", + type=int, + default=42, + help="Fixed seed for reproducible (byte-identical) audio. Pass --seed '' " + "equivalent (any non-int) to disable -- there is no unset sentinel here " + "by design, since the default sweep is meant to be exactly reproducible.", + ) + return p.parse_args() + + +def main() -> None: + args = parse_args() + out = Path(args.output_dir).resolve() + out.mkdir(parents=True, exist_ok=True) + + all_results = [] + failed = False + for conc in args.concurrency: + for chunk_frames in args.step_size: + try: + result = run_combo(args, conc, chunk_frames, out) + all_results.append(result) + if result["errors"] > 0 or result["correctness_problems"] > 0: + failed = True + except Exception as e: + print(f" combo c={conc} chunk={chunk_frames} FAILED: {e}", flush=True) + failed = True + + out_path = out / "steps_results.json" + with open(out_path, "w", encoding="utf-8") as f: + json.dump({"model": args.model, "texts": args.texts, "runs": all_results}, f, indent=2, ensure_ascii=False) + print(f"\nResults written to {out_path}", flush=True) + + print("\n=== SUMMARY ===") + header = ( + f"{'conc':>4} {'step':>4} {'ok':>4} {'corr':>4} " + f"{'TTFA med(ms)':>12} {'step med(ms)':>12} {'step p95(ms)':>12} {'RTF mean+-sd':>14}" + ) + print(header) + print("-" * len(header)) + for r in all_results: + print( + f"{r['concurrency']:>4} {r['step_size']:>4} {r['ok']:>4} {r['correctness_problems']:>4} " + f"{r['ttfa_s']['median'] * 1e3:>12.1f} " + f"{r['step_latency_ms']['median']:>12.1f} {r['step_latency_ms']['p95']:>12.1f} " + f"{r['rtf']['mean']:>6.3f}+-{r['rtf']['stdev']:>5.3f}" + ) + for t, stats in r.get("by_text", {}).items(): + print( + f" text {t[:24]!r}: e2e {stats['e2e_s']['mean'] * 1e3:8.1f}" + f"+-{stats['e2e_s']['stdev'] * 1e3:6.1f} ms (n={stats['e2e_s']['n']}, " + f"audio {stats['audio_s']['mean']:.1f}s)" + ) + + if failed: + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index 434c4305c..c463d3917 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -66,14 +66,12 @@ dev = [ "jiwer==4.0.0", "zhon==2.1.1", "zhconv==1.4.3", - "scipy==1.18.0", "funasr==1.3.29", "lpips==0.1.4", # for quantization fp8 "orjson==3.11.9", # to make v1/embeddings API fast "voxcpm==2.0.3", # for VoxCPM2 path-locating in tests "torchdiffeq==0.2.5", # for covo_audio Test "FlagEmbedding==1.4.0", # for hunyuan_image3 tests - "s3tokenizer==0.3.0", # for tts test # Provides top-level cosyvoice2; do not co-install stepaudio2-minicpmo, # whose top-level stepaudio2 package conflicts with this distribution. "step-audio2==1.0.0", # Step-Audio2 / MiniCPM-o 4.5 Token2Wav (flashcosyvoice / hyperpyyaml) diff --git a/requirements/common.txt b/requirements/common.txt index 855fb6b0a..d7b845e32 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -41,3 +41,8 @@ pydantic-settings>=2.0.0 gguf>=0.10.0 # Required for running the Cosmos3 model with guardrails enabled (default behavior). cosmos-guardrail>=0.3.1 +# Required for CosyVoice3 (get_window/resample_poly). >=1.18.0 requires +# Python >=3.12, breaking base installs on the project's supported 3.10/3.11; +# these are long-stable APIs with no version-specific need for 1.18. +scipy>=1.11.0 +s3tokenizer==0.3.0 diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py index 7dfbe6e74..74fc5059b 100644 --- a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py +++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py @@ -15,6 +15,7 @@ from vllm_omni.diffusion.config import set_current_diffusion_config from vllm_omni.diffusion.data import AttentionConfig from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import ( CausalHiFTGenerator, + HiFTGenerator, ) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -76,6 +77,39 @@ def test_causal_hift_stft_moves_window_to_input_device(causal_hift): assert causal_hift.stft_window.device == waveform.device +class _MinimalF0Predictor(nn.Module): + def __init__(self) -> None: + super().__init__() + self.proj = nn.Conv1d(80, 1, 1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.proj(x).squeeze(1).abs() + + +@pytest.mark.parametrize("sampling_rate", [22050, 24000]) +def test_noncausal_hift_inference_runs_with_real_sinegen(sampling_rate): + """Base (non-causal) HiFTGenerator with a real SineGen/SineGen2 (not a + test double) must still run: SourceModuleHnNSF/SineGen/SineGen2/_f02sine + all branch on `self.causal` to return the pre-streaming (shorter) tuple + contract for non-causal callers, and any one of them regressing back to + always returning the streaming-length tuple breaks this immediately. + """ + hift = HiFTGenerator( + base_channels=32, + sampling_rate=sampling_rate, + upsample_rates=[8, 5, 3], + upsample_kernel_sizes=[16, 11, 7], + source_resblock_kernel_sizes=[7, 7, 11], + source_resblock_dilation_sizes=[[1, 3, 5]] * 3, + f0_predictor=_MinimalF0Predictor(), + ).eval() + + speech, _ = hift.inference(torch.randn(1, 80, 20)) + + assert speech.shape[0] == 1 + assert torch.isfinite(speech).all() + + def test_causal_hift_routes_npu_transforms_to_cpu_fallback(causal_hift, monkeypatch): class NpuTensor: device = SimpleNamespace(type="npu") @@ -83,16 +117,17 @@ def test_causal_hift_routes_npu_transforms_to_cpu_fallback(causal_hift, monkeypa tensor = NpuTensor() stft_calls = [] istft_calls = [] - monkeypatch.setattr( - causal_hift, - "_stft_on_cpu", - lambda value: stft_calls.append(value) or ("real", "imag"), - ) - monkeypatch.setattr( - causal_hift, - "_istft_on_cpu", - lambda magnitude, phase: istft_calls.append((magnitude, phase)) or "waveform", - ) + + def fake_stft(value): + stft_calls.append(value) + return ("real", "imag") + + def fake_istft(magnitude, phase): + istft_calls.append((magnitude, phase)) + return "waveform" + + monkeypatch.setattr(causal_hift, "_stft_on_cpu", fake_stft) + monkeypatch.setattr(causal_hift, "_istft_on_cpu", fake_istft) assert causal_hift._stft(tensor) == ("real", "imag") assert causal_hift._istft(tensor, tensor) == "waveform" @@ -627,7 +662,11 @@ def test_code2wav_forward_finalizes_hift_tail(): def inference(self, speech_feat, finalize=True): self.finalize_calls.append(bool(finalize)) - return torch.zeros((speech_feat.shape[0], 1, speech_feat.shape[-1]), dtype=speech_feat.dtype), None + return ( + torch.zeros((speech_feat.shape[0], 1, speech_feat.shape[-1]), dtype=speech_feat.dtype), + None, + None, + ) model = object.__new__(CosyVoice3Code2Wav) nn.Module.__init__(model) diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_incremental_hift.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_incremental_hift.py new file mode 100644 index 000000000..1e8eff599 --- /dev/null +++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_incremental_hift.py @@ -0,0 +1,415 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Correctness + boundedness tests for the incremental (windowed) HiFT path. + +The streaming vocoder runs HiFT over a bounded mel window instead of the full +cumulative history, carrying harmonic phase and noise-buffer offset across +chunks. + +* **Correctness** — windowed output must match the full cumulative re-run. + With full history the match is bit-exact (proves the phase/noise carry is + right). With a bounded window, float64 still agrees to ~1e-16 (the window + is algebraically exact); float32 differs by ~2e-7 absolute (reassociation + rounding, not truncation), so tests assert a tight atol/rtol rather than + PCM16 byte-identity, since a sub-LSB perturbation can flip a boundary sample. +* **Boundedness** — the per-chunk window must not grow with mel history. + +Tests are parametrized across REAL_CONFIG (24000 Hz -> SineGen2, what +CosyVoice3 ships) and LEGACY_CONFIG (22050 Hz -> SineGen) so a class-selection +regression in either path is caught. +""" + +import pytest +import torch +import torch.nn as nn + +from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import ( + CausalConvRNNF0Predictor, + CausalHiFTGenerator, + SineGen, + SineGen2, + _wrapped_slice, +) +from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3_code2wav import CosyVoice3Code2Wav + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +CHUNK_LEN = 24 +TOTAL_MEL = 96 +SPM = 480 # samples per mel frame for the small test HiFT (8*5*3*4), matching the real model + +LONG_TOTAL_MEL = 400 # long enough that the 64-frame window actually truncates +PCM16_LSB = 1.0 / 32767.0 # smallest representable PCM16 difference + +# Worst observed deviation is 1.97e-7 (~2x float32 eps); atol=1e-6 keeps 5x headroom. +ATOL = 1e-6 +RTOL = 1e-5 + +# What CosyVoice3 ships (-> SineGen2). nb_harmonics=8 is required, not a free +# choice: SineGen2's causal buffers are hardcoded to dim 9 (= harmonic_num+1). +REAL_CONFIG = {"sampling_rate": 24000, "nb_harmonics": 8} +LEGACY_CONFIG = {"sampling_rate": 22050, "nb_harmonics": 4} # -> SineGen, still shipped code + +CONFIGS = [ + pytest.param(REAL_CONFIG, id="sinegen2_real_config"), + pytest.param(LEGACY_CONFIG, id="sinegen1_legacy_config"), +] + + +def _make_hift(config: dict = REAL_CONFIG) -> CausalHiFTGenerator: + """Small CausalHiFTGenerator; defaults to the config CosyVoice3 ships.""" + torch.manual_seed(0) + f0_predictor = CausalConvRNNF0Predictor(num_class=1, in_channels=80, cond_channels=16) + return CausalHiFTGenerator( + in_channels=80, + base_channels=32, + nb_harmonics=config["nb_harmonics"], + sampling_rate=config["sampling_rate"], + upsample_rates=[8, 5, 3], + upsample_kernel_sizes=[16, 11, 7], + source_resblock_kernel_sizes=[7, 7, 11], + source_resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5], [1, 3, 5]], + resblock_kernel_sizes=[3, 7, 11], + resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5], [1, 3, 5]], + f0_predictor=f0_predictor, + ).eval() + + +def test_real_config_resolves_to_sinegen2(): + """Guard against silently testing the wrong SineGen class again.""" + hift = _make_hift(REAL_CONFIG) + assert isinstance(hift.m_source.l_sin_gen, SineGen2) + + hift_legacy = _make_hift(LEGACY_CONFIG) + assert isinstance(hift_legacy.m_source.l_sin_gen, SineGen) + + +def _make_model(hift: CausalHiFTGenerator, window_len: int) -> CosyVoice3Code2Wav: + """Build a CosyVoice3Code2Wav shell wired to the given HiFT.""" + model = object.__new__(CosyVoice3Code2Wav) + nn.Module.__init__(model) + model.hift = hift + model._hift_window_len = window_len + return model + + +def _chunks(total_mel: int = TOTAL_MEL, chunk_len: int = CHUNK_LEN) -> list[torch.Tensor]: + torch.manual_seed(1) + mel = torch.randn(1, 80, total_mel) + return [mel[:, :, i : i + chunk_len] for i in range(0, total_mel, chunk_len)] + + +def _variable_chunks(total_mel: int, first_len: int, step_len: int) -> list[torch.Tensor]: + """First chunk sized like a real prompt prefill, then small fixed-size steps.""" + torch.manual_seed(1) + mel = torch.randn(1, 80, total_mel) + chunks = [mel[:, :, :first_len]] + pos = first_len + while pos < total_mel: + chunks.append(mel[:, :, pos : pos + step_len]) + pos += step_len + return chunks + + +def _full_reference(model: CosyVoice3Code2Wav, chunks: list[torch.Tensor]) -> torch.Tensor: + """Re-run HiFT over the full cumulative mel each chunk (pre-windowing behavior). + + Uses non-finalize inference per chunk, matching the old streaming vocoder; a + single finalize pass would emit the look-right tail streaming holds back. + """ + torch.manual_seed(0) + emitted = [] + cache: dict[str, torch.Tensor] | None = None + for chunk in chunks: + cached = None if cache is None else cache.get("mel") + if cached is not None and cached.numel() > 0: + tts_mel = torch.cat([cached, chunk], dim=-1) + else: + tts_mel = chunk + speech, _, _ = model.hift.inference(speech_feat=tts_mel, finalize=False) + speech = speech.reshape(speech.shape[0], -1) + offset = 0 if cache is None else cache.get("speech_offset") + emitted.append(speech[:, offset:]) + cache = {"mel": tts_mel.detach().cpu(), "speech_offset": int(speech.shape[-1])} + return torch.cat(emitted, dim=-1) + + +def _pcm16(x: torch.Tensor) -> torch.Tensor: + """Quantize float audio to PCM16, the format the server streams to clients.""" + return (x.clamp(-1, 1) * 32767).round().to(torch.int16) + + +def _incremental(model: CosyVoice3Code2Wav, chunks: list[torch.Tensor]) -> torch.Tensor: + """Run the new windowed path via _stream_hift_from_feat.""" + torch.manual_seed(0) + emitted = [] + cache: dict[str, torch.Tensor] | None = None + for chunk in chunks: + speech, cache = model._stream_hift_from_feat(chunk, cache_state=cache, finalize=False) + emitted.append(speech.reshape(speech.shape[0], -1)) + return torch.cat(emitted, dim=-1) + + +def _rel_divergence(a: torch.Tensor, b: torch.Tensor) -> float: + return ((a - b).abs().mean() / a.abs().mean().clamp_min(1e-6)).item() + + +@pytest.mark.parametrize("config", CONFIGS) +def test_incremental_hift_matches_full_reference_exactly_with_full_history(config): + """With a window covering the full history, the windowed path must match the + full cumulative re-run exactly. This proves the phase/noise carry is right.""" + hift = _make_hift(config) + model = _make_model(hift, window_len=TOTAL_MEL) # window >= history + chunks = _chunks() + + full = _full_reference(model, chunks) + incr = _incremental(model, chunks) + + assert full.shape == incr.shape, f"{full.shape} vs {incr.shape}" + assert _rel_divergence(full, incr) < 1e-6 + + +@pytest.mark.parametrize("config", CONFIGS) +def test_incremental_hift_bounded_window_is_close(config): + """With a bounded window, the windowed path stays close to the full re-run; + only the truncated deep history diverges, and it is small.""" + hift = _make_hift(config) + model = _make_model(hift, window_len=48) + chunks = _chunks() + + full = _full_reference(model, chunks) + incr = _incremental(model, chunks) + + assert full.shape == incr.shape + torch.testing.assert_close(incr, full, atol=ATOL, rtol=RTOL) + + +@pytest.mark.parametrize("config", CONFIGS) +def test_incremental_hift_matches_streaming_reference_within_tolerance(config): + """Windowed output matches the full-cumulative reference within float32 rounding. + + The load-bearing correctness test: 400 mel frames far exceeds the 64-frame + window, so truncation genuinely exercises the phase/noise carry. Float64 + agrees to ~1e-16, confirming the residual here is rounding, not truncation. + """ + hift = _make_hift(config) + model = _make_model(hift, window_len=64) # the real _hift_window_len + chunks = _chunks(total_mel=LONG_TOTAL_MEL) + + trim = int(hift.f0_predictor.condnet[0].causal_padding) + assert LONG_TOTAL_MEL > 64 + trim + CHUNK_LEN # window must actually truncate + + full = _full_reference(model, chunks) + incr = _incremental(model, chunks) + + assert full.shape == incr.shape, f"{full.shape} vs {incr.shape}" + torch.testing.assert_close(incr, full, atol=ATOL, rtol=RTOL) + + max_dev = (full - incr).abs().max().item() # stays far below one PCM16 step + assert max_dev < PCM16_LSB / 10, f"max deviation {max_dev:.3e} vs PCM16 LSB {PCM16_LSB:.3e}" + + +def test_incremental_hift_window_is_bounded(): + """Per-chunk window must not grow with cumulative mel history.""" + hift = _make_hift() + model = _make_model(hift, window_len=32) + chunks = _chunks(total_mel=200, chunk_len=24) + + trim = int(hift.f0_predictor.condnet[0].causal_padding) + f0_margin = int(hift.f0_predictor.left_context_frames) + + cache = None + window_sizes = [] + for chunk in chunks: + _, cache = model._stream_hift_from_feat(chunk, cache_state=cache, finalize=False) + window_sizes.append(int(cache["mel"].shape[-1])) + + assert max(window_sizes) <= 32 + trim + f0_margin + 24 # bounded regardless of utterance length + steady = max(window_sizes) + assert window_sizes.count(steady) >= 2 # steady-state value repeats once history exceeds window + + +def test_incremental_hift_finalize_releases_tail(): + """Finalize must emit the released look-right tail without crashing.""" + hift = _make_hift() + model = _make_model(hift, window_len=32) + chunks = _chunks(total_mel=72, chunk_len=24) + + cache = None + emitted = [] + for i, chunk in enumerate(chunks): + finalize = i == len(chunks) - 1 + speech, cache = model._stream_hift_from_feat(chunk, cache_state=cache, finalize=finalize) + emitted.append(speech.reshape(speech.shape[0], -1)) + if finalize: + assert cache is None + + full = torch.cat(emitted, dim=-1) + assert full.shape[-1] > 0 + assert torch.isfinite(full).all() + + +def test_incremental_hift_emits_full_audio_length(): + """Streamed output covers the full mel history at samples-per-mel resolution. + + Guards against emission arithmetic dropping an upsample stage, which + silently truncates audio (e.g. 1/3 of expected length with [8,5,3]). + """ + hift = _make_hift() + model = _make_model(hift, window_len=32) + chunks = _chunks(total_mel=72, chunk_len=24) + + cache = None + emitted = [] + for i, chunk in enumerate(chunks): + finalize = i == len(chunks) - 1 + speech, cache = model._stream_hift_from_feat(chunk, cache_state=cache, finalize=finalize) + emitted.append(speech.reshape(speech.shape[0], -1)) + if finalize: + assert cache is None + + full = torch.cat(emitted, dim=-1) + # total_mel * SPM, +slack for finalize's look-right tail, -slack for a dropped upsample stage + total_mel = sum(c.shape[-1] for c in chunks) + assert full.shape[-1] > total_mel * SPM * 0.8, f"{full.shape[-1]} vs {total_mel * SPM}" + assert full.shape[-1] < total_mel * SPM * 1.5, f"{full.shape[-1]} vs {total_mel * SPM}" + + +@pytest.mark.parametrize("config", CONFIGS) +def test_finalize_matches_single_pass_length_exactly(config): + """Streamed total (chunked, last=finalize) must exactly match a single + non-streaming finalize=True call over the whole mel -- not just approximately. + + The finalize release must include `trim` (the F0 predictor's own held-back + frames), or the last 3 mel frames (1,440 samples / 60ms at 24kHz) of every + stream's audio go missing. + """ + hift = _make_hift(config) + chunks = _chunks(total_mel=200, chunk_len=24) + + torch.manual_seed(0) + full_mel = torch.cat(chunks, dim=-1) + ground_truth, _, _ = hift.inference(speech_feat=full_mel, finalize=True) + ground_truth = ground_truth.reshape(ground_truth.shape[0], -1) + + hift2 = _make_hift(config) + model = _make_model(hift2, window_len=64) + torch.manual_seed(0) + cache = None + emitted = [] + for i, chunk in enumerate(chunks): + finalize = i == len(chunks) - 1 + speech, cache = model._stream_hift_from_feat(chunk, cache_state=cache, finalize=finalize) + emitted.append(speech.reshape(speech.shape[0], -1)) + streamed = torch.cat(emitted, dim=-1) + + assert streamed.shape[-1] == ground_truth.shape[-1], ( + f"{streamed.shape[-1]} vs {ground_truth.shape[-1]} (diff={ground_truth.shape[-1] - streamed.shape[-1]} samples)" + ) + torch.testing.assert_close(streamed, ground_truth, atol=ATOL, rtol=RTOL) + + +@pytest.mark.parametrize("config", CONFIGS) +def test_incremental_hift_matches_reference_for_voiced_f0(config): + """Forces F0 above nsf_voiced_threshold (untrained weights never trigger voiced, + so no other test here does) and checks windowed output against the full- + cumulative reference, which exercises condnet's left-causal convs and + SineGen/SineGen2's phase integration under genuinely voiced harmonics. + + Tolerance is looser than ATOL/RTOL: summing this sustained constant-F0 phase + in per-window chunks vs. one recompute reassociates float32 adds differently; + real speech's varying/unvoiced F0 keeps this far smaller in practice. + """ + hift = _make_hift(config) + with torch.no_grad(): + hift.f0_predictor.classifier.bias.fill_(150.0) # clearly voiced (threshold is 10) + chunks = _chunks(total_mel=200, chunk_len=24) + + full = _full_reference(_make_model(hift, window_len=400), chunks) + incr = _incremental(_make_model(hift, window_len=64), chunks) + + assert full.shape == incr.shape + torch.testing.assert_close(incr, full, atol=2e-4, rtol=2e-4) + + +@pytest.mark.parametrize("config", CONFIGS) +@pytest.mark.parametrize("first_len,step_len", [(10, 6), (10, 3), (20, 3), (10, 1)]) +def test_incremental_hift_matches_reference_for_voiced_f0_small_chunks(config, first_len, step_len): + """Same as test_incremental_hift_matches_reference_for_voiced_f0, but with small, + sub-receptive-field chunk sizes (down to 1 mel frame) instead of the uniform + 24-frame chunks used elsewhere. cache_state must carry the F0 predictor's own + left receptive field across calls regardless of how small a single chunk is, + or the next call's margin comes up short and substitutes zero-padding for real + history -- corrupting SineGen's phase integration the same way as an absent + margin entirely. + """ + hift = _make_hift(config) + with torch.no_grad(): + hift.f0_predictor.classifier.bias.fill_(150.0) # clearly voiced (threshold is 10) + chunks = _variable_chunks(total_mel=200, first_len=first_len, step_len=step_len) + + full = _full_reference(_make_model(hift, window_len=400), chunks) + incr = _incremental(_make_model(hift, window_len=64), chunks) + + assert full.shape == incr.shape + torch.testing.assert_close(incr, full, atol=2e-4, rtol=2e-4) + + +def test_wrapped_slice_wraps_position_deterministically(): + """SineGen/SineGen2/SourceModuleHnNSF's noise buffers are a fixed 300s at + 24kHz; a stream running past that must keep repeating the realization + (position-deterministic) rather than the slice coming back short and + failing to broadcast against the rest of that call's tensors. + """ + buf = torch.arange(10).reshape(1, 10, 1) + + assert _wrapped_slice(buf, 2, 3).flatten().tolist() == [2, 3, 4] + assert _wrapped_slice(buf, 8, 4).flatten().tolist() == [8, 9, 0, 1] + assert _wrapped_slice(buf, 12, 3).flatten().tolist() == [2, 3, 4] # offset past the end + assert _wrapped_slice(buf, 5, 13).flatten().tolist() == [5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6, 7] + + +def test_incremental_hift_empty_middle_chunk_emits_nothing(): + """An empty middle chunk (reachable from forward_streaming_batch when a + row's post-trim feat comes out empty near the lookahead boundary) must + emit zero samples and leave cache_state's phase_acc unchanged, not + re-emit the re-decoded overlap window as if it were new audio. + """ + hift = _make_hift(REAL_CONFIG) + model = _make_model(hift, window_len=64) + chunks = _chunks(total_mel=96, chunk_len=24) + empty = chunks[0][:, :, :0] + + cache = None + for chunk in chunks[:2]: + _, cache = model._stream_hift_from_feat(chunk, cache_state=cache, finalize=False) + assert cache is not None + phase_acc_before = cache["phase_acc"] + + speech, cache = model._stream_hift_from_feat(empty, cache_state=cache, finalize=False) + + assert speech.shape[-1] == 0 + if phase_acc_before is None: + assert cache["phase_acc"] is None + else: + torch.testing.assert_close(cache["phase_acc"], phase_acc_before) + + +@pytest.mark.parametrize("config", CONFIGS) +@pytest.mark.parametrize("mel_frames", [1, 2, 3, 4]) +def test_incremental_hift_short_finalize_chunk_does_not_crash(config, mel_frames): + """A one-token (or otherwise sub-trim) utterance that finalizes on its very + first chunk must not crash. On finalize the F0 predictor consumes the whole + window (it does not hold back `trim` frames), so the phase-carry index must + not be inflated by `trim` -- it would point past the end of a short window. + """ + hift = _make_hift(config) + model = _make_model(hift, window_len=32) + chunk = torch.randn(1, 80, mel_frames) + + speech, cache = model._stream_hift_from_feat(chunk, cache_state=None, finalize=True) + + assert cache is None + assert speech.shape[0] == 1 + assert speech.shape[-1] > 0 + assert torch.isfinite(speech).all() diff --git a/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py b/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py index 959156b8f..555e72e46 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py +++ b/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py @@ -177,6 +177,8 @@ class _MiniDiT(nn.Module): cnn_cache_buffer: torch.Tensor | None = None, att_cache_buffer: torch.Tensor | None = None, ) -> torch.Tensor: + assert cnn_cache_buffer is not None + assert att_cache_buffer is not None x = x.transpose(1, 2) x = self.in_proj(x) for b_idx in range(len(self.blocks)): diff --git a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py index 7cbd31e9c..2b8f3e61e 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py +++ b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py @@ -98,6 +98,46 @@ class ResBlock(torch.nn.Module): remove_weight_norm(self.convs2[idx]) +def _carry_phase_at_boundary( + cum_total: torch.Tensor, + phase_acc: torch.Tensor | None, + next_overlap: int, + *, + dim: int, + trim: int = 0, +) -> torch.Tensor | None: + """Accumulated phase (mod 1) carried into the next streaming window. + + Shared by SineGen and SineGen2 so harmonic phase stays continuous across + chunk boundaries when HiFT runs over a bounded mel window. ``cum_total`` + must already include ``phase_acc``; falls back to it unchanged if + ``next_overlap`` (+ ``trim``) covers the whole window. + """ + length = cum_total.shape[dim] + if next_overlap >= length + trim: + return phase_acc + idx = length + trim - next_overlap - 1 + return cum_total.narrow(dim, idx, 1) % 1 + + +def _wrapped_slice(buf: torch.Tensor, offset: int, length: int, dim: int = 1) -> torch.Tensor: + """``length`` elements of ``buf`` along ``dim`` starting at ``offset``, + wrapping (mod ``buf.shape[dim]``) instead of running past the end and + coming back short. Shared by SineGen/SineGen2/SourceModuleHnNSF's + position-indexed noise buffers, now that streams can outlast their 300s. + """ + n = buf.shape[dim] + start = offset % n + parts = [] + remaining = length + while remaining > 0: + take = min(remaining, n - start) + parts.append(buf.narrow(dim, start, take)) + remaining -= take + start = 0 + return parts[0] if len(parts) == 1 else torch.cat(parts, dim=dim) + + class SineGen(torch.nn.Module): """Definition of sine generator SineGen(samp_rate, harmonic_num = 0, @@ -114,13 +154,24 @@ class SineGen(torch.nn.Module): segment is always sin(np.pi) or cos(0) """ - def __init__(self, samp_rate, harmonic_num=0, sine_amp=0.1, noise_std=0.003, voiced_threshold=0): + def __init__(self, samp_rate, harmonic_num=0, sine_amp=0.1, noise_std=0.003, voiced_threshold=0, causal=False): super().__init__() self.sine_amp = sine_amp self.noise_std = noise_std self.harmonic_num = harmonic_num self.sampling_rate = samp_rate self.voiced_threshold = voiced_threshold + self.causal = causal + # Fixed per model instance like SineGen2's rand_ini, not redrawn per + # call, or each harmonic jumps to a new random phase every window. + self.phase_vec = Uniform(low=-np.pi, high=np.pi).sample(sample_shape=(1, harmonic_num + 1, 1)) + self.phase_vec[:, 0, :] = 0 + if causal: + self.register_buffer( + "noise_buf", + torch.randn(1, 300 * 24000, harmonic_num + 1), + persistent=False, + ) def _f02uv(self, f0): # generate uv signal @@ -128,22 +179,32 @@ class SineGen(torch.nn.Module): return uv @torch.no_grad() - def forward(self, f0): + def forward(self, f0, phase_acc=None, next_overlap=0, trim=0, noise_offset=0): """sine_tensor, uv = forward(f0) input F0: tensor(batchsize=1, dim=1, length) f0 for unvoiced steps should be 0 output sine_tensor: tensor(batchsize=1, length, dim) output uv: tensor(batchsize=1, length, 1) + + ``next_overlap`` and ``trim`` are sample-indexed here (``f0`` arrives + post-``f0_upsamp``), unlike SineGen2 which downsamples to mel rate + before its own cumsum -- callers must scale both by samples-per-mel. """ f0 = f0.transpose(1, 2) F_mat = torch.zeros((f0.size(0), self.harmonic_num + 1, f0.size(-1))).to(f0.device) for i in range(self.harmonic_num + 1): F_mat[:, i : i + 1, :] = f0 * (i + 1) / self.sampling_rate - theta_mat = 2 * np.pi * (torch.cumsum(F_mat, dim=-1) % 1) - u_dist = Uniform(low=-np.pi, high=np.pi) - phase_vec = u_dist.sample(sample_shape=(f0.size(0), self.harmonic_num + 1, 1)).to(F_mat.device) - phase_vec[:, 0, :] = 0 + cum = torch.cumsum(F_mat, dim=-1) + cum_total = cum if phase_acc is None else phase_acc + cum + theta_mat = 2 * np.pi * (cum_total % 1) + new_phase_acc = _carry_phase_at_boundary(cum_total, phase_acc, next_overlap, dim=-1, trim=trim) + if self.causal: + phase_vec = self.phase_vec.to(F_mat.device) + else: + u_dist = Uniform(low=-np.pi, high=np.pi) + phase_vec = u_dist.sample(sample_shape=(f0.size(0), self.harmonic_num + 1, 1)).to(F_mat.device) + phase_vec[:, 0, :] = 0 # generate sine waveforms sine_waves = self.sine_amp * torch.sin(theta_mat + phase_vec) @@ -155,11 +216,16 @@ class SineGen(torch.nn.Module): # std = self.sine_amp/3 -> max value ~ self.sine_amp # . for voiced regions is self.noise_std noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3 - noise = noise_amp * torch.randn_like(sine_waves) + if self.training is False and self.causal is True: + noise = noise_amp * _wrapped_slice(self.noise_buf, noise_offset, sine_waves.shape[-1]).transpose(1, 2) + else: + noise = noise_amp * torch.randn_like(sine_waves) # first: set the unvoiced part to 0 by uv # then: additive noise sine_waves = sine_waves * uv + noise + if self.causal: + return sine_waves.transpose(1, 2), uv.transpose(1, 2), noise, new_phase_acc return sine_waves.transpose(1, 2), uv.transpose(1, 2), noise @@ -215,7 +281,7 @@ class SineGen2(torch.nn.Module): uv = (f0 > self.voiced_threshold).type(torch.float32) return uv - def _f02sine(self, f0_values): + def _f02sine(self, f0_values, phase_acc=None, next_overlap=0, trim=0): """f0_values: (batchsize, length, dim) where dim indicates fundamental tone and overtones """ @@ -223,9 +289,12 @@ class SineGen2(torch.nn.Module): # because 2 * np.pi * n doesn't affect phase rad_values = (f0_values / self.sampling_rate) % 1 - # initial phase noise (no noise for fundamental component) + # initial phase noise (no noise for fundamental component). Only applied + # on the genuine first chunk of a stream (see rationale above the class); + # phase_acc is None exactly when there is no carried-forward phase yet. if self.training is False and self.causal is True: - rad_values[:, 0, :] = rad_values[:, 0, :] + self.rand_ini.to(rad_values.device) + if phase_acc is None: + rad_values[:, 0, :] = rad_values[:, 0, :] + self.rand_ini.to(rad_values.device) else: rand_ini = torch.rand(f0_values.shape[0], f0_values.shape[2], device=f0_values.device) rand_ini[:, 0] = 0 @@ -237,7 +306,10 @@ class SineGen2(torch.nn.Module): rad_values.transpose(1, 2), scale_factor=1 / self.upsample_scale, mode="linear" ).transpose(1, 2) - phase = torch.cumsum(rad_values, dim=1) * 2 * np.pi + cum = torch.cumsum(rad_values, dim=1) + phase = cum if phase_acc is None else cum + phase_acc + new_phase_acc = _carry_phase_at_boundary(phase, phase_acc, next_overlap, dim=1, trim=trim) + phase = phase * 2 * np.pi phase = torch.nn.functional.interpolate( phase.transpose(1, 2) * self.upsample_scale, scale_factor=self.upsample_scale, @@ -272,9 +344,10 @@ class SineGen2(torch.nn.Module): # get the sines sines = torch.cos(i_phase * 2 * np.pi) - return sines + new_phase_acc = None + return (sines, new_phase_acc) if self.causal else sines - def forward(self, f0): + def forward(self, f0, phase_acc=None, next_overlap=0, trim=0, noise_offset=0): """sine_tensor, uv = forward(f0) input F0: tensor(batchsize=1, length, dim=1) f0 for unvoiced steps should be 0 @@ -284,8 +357,12 @@ class SineGen2(torch.nn.Module): # fundamental component fn = torch.multiply(f0, self.harmonic_ids) - # generate sine waveforms - sine_waves = self._f02sine(fn) * self.sine_amp + if self.causal: + sine_waves, new_phase_acc = self._f02sine(fn, phase_acc, next_overlap, trim) + else: + sine_waves = self._f02sine(fn) + new_phase_acc = None + sine_waves = sine_waves * self.sine_amp # generate uv signal uv = self._f02uv(f0) @@ -295,14 +372,15 @@ class SineGen2(torch.nn.Module): # . for voiced regions is self.noise_std noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3 if self.training is False and self.causal is True: - noise = noise_amp * self.sine_waves[:, : sine_waves.shape[1]].to(sine_waves.device) + # indexed by absolute sample position, not window-relative position + noise = noise_amp * _wrapped_slice(self.sine_waves, noise_offset, sine_waves.shape[1]).to(sine_waves.device) else: noise = noise_amp * torch.randn_like(sine_waves) # first: set the unvoiced part to 0 by uv # then: additive noise sine_waves = sine_waves * uv + noise - return sine_waves, uv, noise + return (sine_waves, uv, noise, new_phase_acc) if self.causal else (sine_waves, uv, noise) class SourceModuleHnNSF(torch.nn.Module): @@ -338,10 +416,13 @@ class SourceModuleHnNSF(torch.nn.Module): self.sine_amp = sine_amp self.noise_std = add_noise_std + self.upsample_scale = upsample_scale # to produce sine waveforms if sinegen_type == "1": - self.l_sin_gen = SineGen(sampling_rate, harmonic_num, sine_amp, add_noise_std, voiced_threshold) + self.l_sin_gen = SineGen( + sampling_rate, harmonic_num, sine_amp, add_noise_std, voiced_threshold, causal=causal + ) else: self.l_sin_gen = SineGen2( sampling_rate, upsample_scale, harmonic_num, sine_amp, add_noise_std, voiced_threshold, causal=causal @@ -354,24 +435,38 @@ class SourceModuleHnNSF(torch.nn.Module): if causal is True: self.uv = torch.rand(1, 300 * 24000, 1) - def forward(self, x): + def forward(self, x, phase_acc=None, uv_offset=0, next_overlap=0, trim=0): """ Sine_source, noise_source = SourceModuleHnNSF(F0_sampled) F0_sampled (batchsize, length, 1) Sine_source (batchsize, length, 1) noise_source (batchsize, length 1) """ - # source for harmonic branch with torch.no_grad(): - sine_wavs, uv, _ = self.l_sin_gen(x) + if not self.causal: + sine_wavs, uv, _ = self.l_sin_gen(x) + new_phase_acc = None + elif isinstance(self.l_sin_gen, SineGen): + # next_overlap/trim arrive in mel frames; SineGen needs samples. + sine_wavs, uv, _, new_phase_acc = self.l_sin_gen( + x, + phase_acc, + next_overlap * self.upsample_scale, + trim * self.upsample_scale, + noise_offset=uv_offset, + ) + else: + sine_wavs, uv, _, new_phase_acc = self.l_sin_gen( + x, phase_acc, next_overlap, trim, noise_offset=uv_offset + ) sine_merge = self.l_tanh(self.l_linear(sine_wavs)) # source for noise branch, in the same shape as uv if self.training is False and self.causal is True: - noise = self.uv[:, : uv.shape[1]] * self.sine_amp / 3 + noise = _wrapped_slice(self.uv, uv_offset, uv.shape[1]) * self.sine_amp / 3 else: noise = torch.randn_like(uv) * self.sine_amp / 3 - return sine_merge, noise, uv + return sine_merge, noise, uv, new_phase_acc class HiFTGenerator(nn.Module): @@ -606,7 +701,7 @@ class HiFTGenerator(nn.Module): f0 = self.f0_predictor(speech_feat) # f0->source s = self.f0_upsamp(f0[:, None]).transpose(1, 2) # bs,n,t - s, _, _ = self.m_source(s) + s, _, _, _ = self.m_source(s) s = s.transpose(1, 2) # mel+source->speech generated_speech = self.decode(x=speech_feat, s=s) @@ -622,7 +717,7 @@ class HiFTGenerator(nn.Module): f0 = self.f0_predictor(speech_feat) # f0->source s = self.f0_upsamp(f0[:, None]).transpose(1, 2) # bs,n,t - s, _, _ = self.m_source(s) + s, _, _, _ = self.m_source(s) s = s.transpose(1, 2) # use cache_source to avoid glitch if cache_source.shape[2] != 0: @@ -785,14 +880,28 @@ class CausalHiFTGenerator(HiFTGenerator): return x @torch.inference_mode() - def inference(self, speech_feat: torch.Tensor, finalize: bool = True) -> torch.Tensor: + def inference( + self, + speech_feat: torch.Tensor, + finalize: bool = True, + phase_acc: torch.Tensor | None = None, + uv_offset: int = 0, + next_overlap: int = 0, + trim: int = 0, + f0_margin: int = 0, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: # mel->f0 NOTE f0_predictor precision is crucial for causal inference, move # self.f0_predictor to cpu if necessary self.f0_predictor.to("cpu") f0 = self.f0_predictor(speech_feat.cpu(), finalize=finalize).to(speech_feat) + if f0_margin > 0: + f0 = f0[:, f0_margin:] + speech_feat = speech_feat[:, :, f0_margin:] # f0->source s = self.f0_upsamp(f0[:, None]).transpose(1, 2) # bs,n,t - s, _, _ = self.m_source(s) + s, _, _, new_phase_acc = self.m_source( + s, phase_acc=phase_acc, uv_offset=uv_offset, next_overlap=next_overlap, trim=trim + ) s = s.transpose(1, 2) if finalize is True: generated_speech = self.decode(x=speech_feat, s=s, finalize=finalize) @@ -800,7 +909,7 @@ class CausalHiFTGenerator(HiFTGenerator): generated_speech = self.decode( x=speech_feat[:, :, : -self.f0_predictor.condnet[0].causal_padding], s=s, finalize=finalize ) - return generated_speech, s + return generated_speech, s, new_phase_acc class CausalConv1dUpsample(torch.nn.Conv1d): @@ -954,6 +1063,11 @@ class CausalConvRNNF0Predictor(nn.Module): nn.ELU(), ) self.classifier = nn.Linear(in_features=cond_channels, out_features=self.num_class) + self.left_context_frames = sum( + layer.causal_padding + for layer in self.condnet + if isinstance(layer, CausalConv1d) and layer.causal_type == "left" + ) def forward(self, x: torch.Tensor, finalize: bool = True) -> torch.Tensor: if finalize is True: diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py index ef706b6c8..1713467c8 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py @@ -9,10 +9,8 @@ This module contains the code2wav (token-to-waveform) stage which uses: 3. HiFiGAN vocoder for waveform synthesis """ -from __future__ import annotations - from collections import Counter -from typing import cast +from typing import TypedDict, cast import numpy as np import torch @@ -39,6 +37,21 @@ from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config logger = init_logger(__name__) +class StreamingFlowItem(TypedDict, total=False): + """One entry of the batched-streaming item list passed to forward_streaming_batch.""" + + index: int + req_id: str | None + stream_finished: bool + token: torch.Tensor + prompt_token: torch.Tensor + prompt_feat: torch.Tensor + embedding: torch.Tensor + cache_state: dict[str, torch.Tensor] | None + token_offset_tokens: int + finalize: bool + + class CosyVoice3Code2Wav(nn.Module): """CosyVoice3 Code2Wav stage for token-to-waveform conversion. @@ -119,6 +132,10 @@ class CosyVoice3Code2Wav(nn.Module): self.mel_cache_len = 20 self.source_cache_len = int(self.mel_cache_len * 256) self.speech_window = np.hamming(2 * self.source_cache_len) + # Must cover decode()'s own causal receptive field, not just the F0 + # margin; window_len=48 already passes + # test_incremental_hift_bounded_window_is_close, so 64 has headroom. + self._hift_window_len = 64 @property def input_frame_rate(self) -> int: @@ -226,40 +243,79 @@ class CosyVoice3Code2Wav(nn.Module): chunk_mel = feat.to(device=hift_weight.device, dtype=hift_weight.dtype) cached_mel = None if not cache_state else cache_state.get("mel") - speech_offset_obj = None if not cache_state else cache_state.get("speech_offset") - try: - speech_offset = int(speech_offset_obj) if speech_offset_obj is not None else 0 - except (TypeError, ValueError): - speech_offset = 0 - + cached_offset = 0 if not cache_state else int(cache_state.get("mel_offset", 0)) + window_len = self._hift_window_len + samples_per_mel = int(np.prod(self.hift.upsample_rates) * self.hift.istft_params["hop_len"]) + trim = int(self.hift.f0_predictor.condnet[0].causal_padding) + f0_margin_frames = int(self.hift.f0_predictor.left_context_frames) + f0_margin = 0 if isinstance(cached_mel, torch.Tensor) and cached_mel.numel() > 0: cached_mel = cached_mel.to(device=chunk_mel.device, dtype=chunk_mel.dtype) - tts_mel = torch.cat([cached_mel, chunk_mel], dim=-1) if chunk_mel.numel() > 0 else cached_mel + total_hist = cached_mel.shape[-1] + overlap = min(window_len + trim, total_hist) + window_mel = torch.cat([cached_mel[..., -overlap:], chunk_mel], dim=-1) + window_offset = cached_offset + total_hist - overlap + uv_offset = window_offset * samples_per_mel + f0_margin = min(f0_margin_frames, total_hist - overlap) + f0_input_mel = ( + torch.cat( + [cached_mel[..., total_hist - overlap - f0_margin : total_hist - overlap], window_mel], dim=-1 + ) + if f0_margin > 0 + else window_mel + ) else: - tts_mel = chunk_mel + window_mel = chunk_mel + f0_input_mel = chunk_mel + window_offset = 0 + uv_offset = 0 - if tts_mel.shape[-1] == 0: + phase_acc = None if not cache_state else cache_state.get("phase_acc") + if phase_acc is not None: + phase_acc = phase_acc.to(device=chunk_mel.device, dtype=chunk_mel.dtype) + + if window_mel.shape[-1] == 0: tts_speech = torch.zeros((chunk_mel.shape[0], 1, 0), device=chunk_mel.device, dtype=chunk_mel.dtype) + new_phase_acc = phase_acc else: - tts_speech, _ = self.hift.inference(speech_feat=tts_mel, finalize=finalize) + next_overlap = min(window_len + trim, window_mel.shape[-1]) + # On finalize the F0 predictor consumes the whole window (it does not + # hold back `trim` frames), so trim must not inflate the phase-carry + # index -- it would point past the end of a short window. + carry_trim = 0 if finalize else trim + tts_speech, _, new_phase_acc = self.hift.inference( + speech_feat=f0_input_mel, + finalize=finalize, + phase_acc=phase_acc, + uv_offset=uv_offset, + next_overlap=next_overlap, + trim=carry_trim, + f0_margin=f0_margin, + ) tts_speech = tts_speech.reshape(tts_speech.shape[0], -1) - speech_offset = max(0, min(speech_offset, int(tts_speech.shape[-1]))) - emitted_speech = tts_speech[:, speech_offset:] + new_samples = chunk_mel.shape[-1] * samples_per_mel + if finalize: + frames_withheld_per_chunk = trim + self.hift.conv_pre_look_right + 1 + released = frames_withheld_per_chunk * samples_per_mel + emitted_speech = tts_speech[:, -(new_samples + released) :] + else: + emitted_speech = tts_speech[:, -new_samples:] if new_samples > 0 else tts_speech[:, :0] if finalize: return emitted_speech.reshape(emitted_speech.shape[0], 1, -1), None new_state = { - "mel": tts_mel.detach().cpu().contiguous(), - "speech_offset": int(tts_speech.shape[-1]), + "mel": f0_input_mel.detach().cpu().contiguous(), + "mel_offset": window_offset - f0_margin, + "phase_acc": new_phase_acc.detach().cpu().contiguous() if new_phase_acc is not None else None, } return emitted_speech.reshape(emitted_speech.shape[0], 1, -1), new_state @torch.inference_mode() def forward_streaming_batch( self, - items: list[dict[str, object]], + items: list[StreamingFlowItem], *, n_timesteps: int = 10, ) -> list[tuple[torch.Tensor, dict[str, torch.Tensor] | None]]: @@ -270,20 +326,16 @@ class CosyVoice3Code2Wav(nn.Module): group and passed to the flow as per-row token lengths. """ results: list[tuple[torch.Tensor, dict[str, torch.Tensor] | None] | None] = [None] * len(items) - groups: dict[tuple[int, int, int, bool], list[tuple[int, dict[str, object]]]] = {} + groups: dict[tuple[int, int, int, bool], list[tuple[int, StreamingFlowItem]]] = {} for index, item in enumerate(items): - token = item["token"] - prompt_token = item["prompt_token"] - prompt_feat = item["prompt_feat"] - embedding = item["embedding"] - assert isinstance(token, torch.Tensor) - assert isinstance(prompt_token, torch.Tensor) - assert isinstance(prompt_feat, torch.Tensor) - assert isinstance(embedding, torch.Tensor) + assert isinstance(item["token"], torch.Tensor) + assert isinstance(item["prompt_token"], torch.Tensor) + assert isinstance(item["prompt_feat"], torch.Tensor) + assert isinstance(item["embedding"], torch.Tensor) key = ( - int(prompt_token.shape[1]), - int(prompt_feat.shape[1]), - int(embedding.shape[1]), + int(item["prompt_token"].shape[1]), + int(item["prompt_feat"].shape[1]), + int(item["embedding"].shape[1]), bool(item.get("finalize", False)), ) groups.setdefault(key, []).append((index, item)) @@ -305,11 +357,11 @@ class CosyVoice3Code2Wav(nn.Module): if len(group) == 1: index, item = group[0] result = self.forward_streaming( - token=item["token"], # type: ignore[arg-type] - prompt_token=item["prompt_token"], # type: ignore[arg-type] - prompt_feat=item["prompt_feat"], # type: ignore[arg-type] - embedding=item["embedding"], # type: ignore[arg-type] - cache_state=item.get("cache_state"), # type: ignore[arg-type] + token=item["token"], + prompt_token=item["prompt_token"], + prompt_feat=item["prompt_feat"], + embedding=item["embedding"], + cache_state=item.get("cache_state"), n_timesteps=n_timesteps, token_offset_tokens=int(item.get("token_offset_tokens", 0)), finalize=bool(item.get("finalize", False)), @@ -318,15 +370,13 @@ class CosyVoice3Code2Wav(nn.Module): continue token_tensors = [item["token"] for _, item in group] - assert all(isinstance(token, torch.Tensor) for token in token_tensors) token_lens = torch.tensor( - [int(token.shape[1]) for token in token_tensors], # type: ignore[union-attr] + [int(token.shape[1]) for token in token_tensors], dtype=torch.int32, ) max_token_len = int(token_lens.max().item()) padded_tokens = [] for token in token_tensors: - assert isinstance(token, torch.Tensor) if int(token.shape[1]) == max_token_len: padded_tokens.append(token) else: @@ -337,9 +387,9 @@ class CosyVoice3Code2Wav(nn.Module): ) padded_tokens.append(torch.cat([token, pad], dim=1)) tokens = torch.cat(padded_tokens, dim=0) - prompt_tokens = torch.cat([item["prompt_token"] for _, item in group], dim=0) # type: ignore[list-item] - prompt_feats = torch.cat([item["prompt_feat"] for _, item in group], dim=0) # type: ignore[list-item] - embeddings = torch.cat([item["embedding"] for _, item in group], dim=0) # type: ignore[list-item] + prompt_tokens = torch.cat([item["prompt_token"] for _, item in group], dim=0) + prompt_feats = torch.cat([item["prompt_feat"] for _, item in group], dim=0) + embeddings = torch.cat([item["embedding"] for _, item in group], dim=0) prompt_token_lens = torch.full((len(group),), prompt_tokens.shape[1], dtype=torch.int32) prompt_feat_lens = torch.full((len(group),), prompt_feats.shape[1], dtype=torch.int32) finalize = bool(group[0][1].get("finalize", False)) @@ -370,7 +420,7 @@ class CosyVoice3Code2Wav(nn.Module): row_feat = row_feat[:, :, trim_mel:] results[index] = self._stream_hift_from_feat( row_feat, - cache_state=item.get("cache_state"), # type: ignore[arg-type] + cache_state=item.get("cache_state"), finalize=finalize, ) @@ -443,7 +493,7 @@ class CosyVoice3Code2Wav(nn.Module): dtype=tts_mel.dtype, ) else: - tts_speech, _ = self.hift.inference(speech_feat=tts_mel, finalize=True) + tts_speech, _, _ = self.hift.inference(speech_feat=tts_mel, finalize=True) return tts_speech