diff --git a/README.md b/README.md index 0fef687..b2622c9 100644 --- a/README.md +++ b/README.md @@ -161,6 +161,82 @@ PY Model files are resolved with `huggingface-hub` by `QwenTTS.from_pretrained(...)` or passed directly to `QwenTTS(...)` as GGUF paths. +## Streaming packet sizes + +`QwenTTS.stream()` defaults to a **4-frame first packet** (320 ms of mono +24 kHz audio). Choose `first_chunk_frames=1`, `2`, `4`, or `8` for 80, 160, +320, or 640 ms of initial audio. A larger packet gives playback more audio +to start with, at the cost of waiting longer before the first yield. + +```python +for audio, sample_rate in tts.stream( + text="The sky is blue today.", + ref_spk_emb=spk, # Base model; use speaker= for CustomVoice + first_chunk_frames=4, + codec_chunk_sec=0.64, # Default later packets: 8 frames, independently sized +): + play_audio(audio, sample_rate) # Your playback/transport function +``` + +Later packets default to `codec_chunk_sec=0.64` (8 frames), matching the native +steady-state width without adding a larger batching delay after the first packet. +Explicit values round to the nearest 80 ms frame, with a one-frame minimum +(for example, 1.0 second rounds to 13 frames / 1.04 seconds). +The value must be finite and positive. Successful end-of-speech or the token +limit flushes any remaining audio as a short packet, including utterances +shorter than the requested first packet. Cancellation or errors discard the +unfinished packet; closing the iterator requests native cancellation. + +Packet assembly happens in Python, using the existing verified ABI v2 library. +The native decoder still emits its fixed 1→2→4→8-frame ramp, then 8-frame +chunks. Consequently, first packets of 1, 2, 4, and 8 frames become available +after native output has reached 1, 3, 7, and 15 frames respectively (or earlier +at end-of-speech). Packet boundaries preserve every PCM sample but do not +change native decode scheduling. Later packets may become available together +when a native callback spans several packet boundaries; this is not a timed +playback scheduler. `codec_left_context_sec` is ignored by the stateful native +stream. Buffered `synthesize()` retains its native codec chunking behavior. + +`last_stream_profile` keeps `first_callback_*` and `callback_count` for raw +native callbacks. `first_packet_ready_ms`, `first_packet_audio_s`, and +`packet_count` describe the assembled Python packets; `first_yield_ms` measures +delivery to the caller. This distinction includes the buffering cost in latency +measurements instead of treating the first native callback as audible output. + +Benchmark 1, 4, and 8 frames on local hardware (one warm-up, then three measured +runs per setting, with rotated order): + +```bash +python scripts/benchmark_first_chunk.py \ + --talker /path/to/qwen-talker-1.7b-base-Q8_0.gguf \ + --codec /path/to/qwen-tokenizer-12hz-Q8_0.gguf \ + --ref-spk /path/to/reference.spk \ + --require-metal --output /tmp/first-chunk-benchmark.json +``` + +For other model types, use `--speaker` or `--instruct` instead of `--ref-spk`. +Omit `--require-metal` on other backends. The report includes native first +callback time, first Python packet time/duration, and every packet's size. +The benchmark checks first/later packet sizes, finite non-silent audio, and the +final tail. Playback smoothness still depends on synthesis speed and the +player's buffering policy. + +Example measurement on an Apple M3 Pro with Metal, macOS 26.6.2, Python +3.12.13, the Q8_0 1.7B Base talker and Q8_0 codec above, and a cached speaker +embedding (three runs per setting after warm-up, medians): + +| First frames | Native first callback | First Python packet ready | First packet audio | +| --- | --- | --- | --- | +| 1 | 94.4 ms | 94.5 ms | 80 ms | +| 4 (default) | 84.2 ms | 369.8 ms | 320 ms | +| 8 | 85.7 ms | 826.4 ms | 640 ms | + +All runs used the default 8-frame later packets (`codec_chunk_sec=0.64`), preserved the +115-frame utterance, and flushed the final short tail. Four frames provide a +middle ground on this machine; these measurements do not guarantee gap-free +playback on other hardware. Raw native callbacks remain one frame initially; +the larger first packet is assembled by the binding. + ## Cached voice references qwentts.cpp ABI v2 can skip reference WAV encoding for Base voice cloning by diff --git a/scripts/benchmark_first_chunk.py b/scripts/benchmark_first_chunk.py new file mode 100644 index 0000000..bced75a --- /dev/null +++ b/scripts/benchmark_first_chunk.py @@ -0,0 +1,98 @@ +"""Compare first-packet latency and duration using local GGUF weights.""" +from __future__ import annotations + +import argparse +import json +import os +import platform +from pathlib import Path +import statistics + +import numpy as np + +from qwentts_cpp import QwenTTS, load_speaker_embedding + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--talker", type=Path, required=True) + parser.add_argument("--codec", type=Path, required=True) + parser.add_argument("--speaker") + parser.add_argument("--ref-spk", type=Path) + parser.add_argument("--instruct") + parser.add_argument("--require-metal", action="store_true") + parser.add_argument("--repeats", type=int, default=3) + parser.add_argument("--codec-chunk-sec", type=float, default=0.64) + parser.add_argument("--output", type=Path, default=Path("first-chunk-benchmark.json")) + parser.add_argument("--text", default=( + "This is a streaming speech test. The first packet should arrive promptly, " + "with enough audio to keep playback smooth while the next packet is generated." + )) + args = parser.parse_args() + if args.repeats < 1: + parser.error("--repeats must be positive") + if args.require_metal: + os.environ["GGML_BACKEND"] = "MTL0" + ref = load_speaker_embedding(args.ref_spk) if args.ref_spk else None + results = [] + with QwenTTS(args.talker, args.codec) as tts: + def run(first: int) -> dict: + packets = list(tts.stream( + text=args.text, speaker=args.speaker, instruct=args.instruct, + ref_spk_emb=ref, seed=42, max_new_tokens=256, + codec_chunk_sec=args.codec_chunk_sec, first_chunk_frames=first, + )) + assert len(packets) >= 3, "Use a longer utterance to measure steady-state packets" + assert all(p.size and np.isfinite(p).all() and sr == 24000 for p, sr in packets) + assert any(np.max(np.abs(p)) > 1e-5 for p, _ in packets), "Silent audio" + assert packets[0][0].size == first * 1920, "Incorrect first packet size" + profile = tts.last_stream_profile + assert all(p.size == profile["packet_frames"] * 1920 for p, _ in packets[1:-1]) + assert packets[-1][0].size <= profile["packet_frames"] * 1920 + return { + "first_chunk_frames": first, + "native_first_callback_ms": profile["first_callback_enter_ms"], + "native_first_callback_audio_ms": profile["first_callback_audio_s"] * 1000, + "first_packet_ready_ms": profile["first_packet_ready_ms"], + "first_yield_ms": profile["first_yield_ms"], + "first_packet_audio_ms": profile["first_packet_audio_s"] * 1000, + "packet_samples": [int(p.size) for p, _ in packets], + "total_ms": profile["consumer_done_ms"], + } + + # Warm the model/codec before measurements, then rotate the order to + # reduce systematic cache and temperature bias between settings. + print("Warming up (excluded from measurements)...", flush=True) + run(4) + for repeat in range(args.repeats): + order = (1, 4, 8) + shift = repeat % len(order) + for first in order[shift:] + order[:shift]: + result = run(first) + results.append(result) + print(f"frames={first}, run={repeat + 1}: native callback " + f"{result['native_first_callback_ms']:.1f} ms, first packet " + f"{result['first_packet_ready_ms']:.1f} ms, audio " + f"{result['first_packet_audio_ms']:.0f} ms", flush=True) + report = { + "platform": platform.platform(), "machine": platform.machine(), + "python": platform.python_version(), "native_version": tts.library.version(), + "backend": os.environ.get("GGML_BACKEND", "auto"), + "talker": args.talker.name, "codec": args.codec.name, + "codec_chunk_sec": args.codec_chunk_sec, "seed": 42, + "max_new_tokens": 256, "text": args.text, + "warmup_runs": 1, "repeats": args.repeats, "runs": results, + } + args.output.write_text(json.dumps(report, indent=2) + "\n") + print("\nMedians: frames | native callback ms | packet ready ms | packet audio ms") + for first in (1, 4, 8): + runs = [r for r in results if r["first_chunk_frames"] == first] + values = [statistics.median(r[key] for r in runs) for key in ( + "native_first_callback_ms", "first_packet_ready_ms", "first_packet_audio_ms", + )] + print(f"{first:>14} | {values[0]:>18.1f} | {values[1]:>15.1f} | {values[2]:>15.1f}") + print(f"Wrote {args.output}") + + +if __name__ == "__main__": + main() diff --git a/src/qwentts_cpp/_binding.py b/src/qwentts_cpp/_binding.py index 8c7dcd0..fe184f4 100644 --- a/src/qwentts_cpp/_binding.py +++ b/src/qwentts_cpp/_binding.py @@ -1,6 +1,8 @@ from __future__ import annotations import ctypes +import math +import numbers import os import queue import sys @@ -19,6 +21,34 @@ # writes a parameter struct; probing default_params itself is not memory-safe. QWENTTS_NATIVE_REVISION = "7df559a8ca25f66fee02970514ebe5f01dee9055" RVQ_CODE_BITS = 11 +CODEC_FRAME_SAMPLES = 1920 # Fixed 12.5 Hz codec at 24 kHz. + + +class _StreamPackets: + """Reframe owned native PCM without changing the native codec state.""" + + def __init__(self, first_frames: int, later_frames: int): + self.target = first_frames * CODEC_FRAME_SAMPLES + self.later_target = later_frames * CODEC_FRAME_SAMPLES + self.parts: list[np.ndarray] = [] + self.size = 0 + + def push(self, chunk: np.ndarray) -> Iterator[np.ndarray]: + offset = 0 + while offset < chunk.size: + take = min(self.target - self.size, chunk.size - offset) + self.parts.append(chunk[offset:offset + take]) + self.size += take + offset += take + if self.size == self.target: + yield self.flush() + self.target = self.later_target + + def flush(self) -> np.ndarray: + packet = self.parts[0] if len(self.parts) == 1 else np.concatenate(self.parts) + self.parts = [] + self.size = 0 + return packet class QwenStatus(IntEnum): @@ -821,12 +851,34 @@ def stream( subtalker_temperature: float | None = None, subtalker_top_k: int | None = None, subtalker_top_p: float | None = None, - codec_chunk_sec: float = 1.0, + codec_chunk_sec: float = 0.64, codec_left_context_sec: float = 2.0, + first_chunk_frames: int = 4, dump_dir: str | os.PathLike[str] | None = None, ) -> Iterator[Tuple[np.ndarray, int]]: + """Yield mono 24 kHz PCM with independent first and later packet sizes. + + The first packet covers 1, 2, 4 (default), or 8 codec frames (80 ms + each). Later packets default to 0.64 seconds (8 frames) and cover + codec_chunk_sec rounded to the nearest frame, at least one. + A short final packet flushes on successful EOS. + Python assembles packets from the native 1/2/4/8-frame callback ramp; + a packet can therefore wait for a native chunk crossing its boundary. + codec_left_context_sec is ignored by the stateful native stream. + """ + if (isinstance(first_chunk_frames, bool) + or not isinstance(first_chunk_frames, numbers.Integral) + or first_chunk_frames not in (1, 2, 4, 8)): + raise ValueError("first_chunk_frames must be one of 1, 2, 4, or 8") + if not math.isfinite(codec_chunk_sec) or codec_chunk_sec <= 0: + raise ValueError("codec_chunk_sec must be finite and positive") + later_frames = max(1, int(codec_chunk_sec * 12.5 + 0.5)) + packets = _StreamPackets(int(first_chunk_frames), later_frames) profile: dict[str, Any] = { "mode": "stream", + "first_chunk_frames": int(first_chunk_frames), + "packet_frames": later_frames, + "packet_count": 0, "codec_chunk_sec": float(codec_chunk_sec), "codec_left_context_sec": float(codec_left_context_sec), "callback_count": 0, @@ -846,6 +898,14 @@ def elapsed_ms() -> float: def cancel_cb(_user_data) -> bool: return cancel_event.is_set() + def emit_packet(packet: np.ndarray) -> None: + if profile["packet_count"] == 0: + profile["first_packet_ready_ms"] = elapsed_ms() + profile["first_packet_n_samples"] = int(packet.size) + profile["first_packet_audio_s"] = float(packet.size) / 24000.0 + profile["packet_count"] += 1 + chunks.put((packet, 24000)) + def on_chunk(samples, n_samples, _user_data) -> bool: if cancel_event.is_set(): return False @@ -863,7 +923,10 @@ def on_chunk(samples, n_samples, _user_data) -> bool: if is_first: profile["first_callback_copy_ms"] = copy_ms queue_start = time.perf_counter() - chunks.put((chunk.astype(np.float32, copy=False), 24000)) + for packet in packets.push(chunk): + if cancel_event.is_set(): + return False + emit_packet(packet) queue_ms = (time.perf_counter() - queue_start) * 1000 profile["callback_queue_ms_total"] += queue_ms if is_first: @@ -921,6 +984,8 @@ def producer() -> None: profile["native_return_ms"] = elapsed_ms() if rc != QwenStatus.OK and not cancel_event.is_set(): chunks.put(QwenTTSError(self.library.last_error() or f"qt_synthesize failed with status {rc}")) + elif rc == QwenStatus.OK and not cancel_event.is_set() and packets.size: + emit_packet(packets.flush()) except BaseException as exc: if not cancel_event.is_set(): chunks.put(exc) diff --git a/tests/test_stream.py b/tests/test_stream.py new file mode 100644 index 0000000..9ef694d --- /dev/null +++ b/tests/test_stream.py @@ -0,0 +1,160 @@ +"""Packet boundaries, native buffer ownership, and cooperative cancellation.""" +import ctypes +import threading +from types import SimpleNamespace + +import numpy as np +import pytest + +from qwentts_cpp import QwenTTS, QwenTTSError +from qwentts_cpp._binding import CODEC_FRAME_SAMPLES, QT_ABI_VERSION, QwenStatus + + +class NativeStream: + def __init__(self, frames=39, status=QwenStatus.OK, wait_for_cancel=False): + self.audio = np.arange(frames * CODEC_FRAME_SAMPLES, dtype=np.float32) + self.status = status + self.wait_for_cancel = wait_for_cancel + self.cancelled = threading.Event() + self.freed = threading.Event() + self.called = False + + def qt_tts_default_params(self, params): + params._obj.abi_version = QT_ABI_VERSION + + def qt_synthesize(self, ctx, params, audio): + self.called = True + params = params._obj + offset, width = 0, 1 + while offset < self.audio.size: + if params.cancel(None): + self.cancelled.set() + return QwenStatus.CANCELLED + end = min(offset + width * CODEC_FRAME_SAMPLES, self.audio.size) + buffer = self.audio[offset:end].copy() + proceed = params.on_chunk( + buffer.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), buffer.size, None, + ) + buffer.fill(-1) # Native callback memory is no longer valid. + if not proceed: + self.cancelled.set() + return QwenStatus.CANCELLED + offset = end + width = min(width * 2, 8) + if self.wait_for_cancel: + # The consumer receives a packet, closes the generator, and the + # native cancellation callback must observe that promptly. + for _ in range(500): + if params.cancel(None): + # Even if native delivers an in-flight chunk after close, + # the callback must reject it without queueing more PCM. + assert not params.on_chunk( + buffer.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), buffer.size, None, + ) + self.cancelled.set() + return QwenStatus.CANCELLED + self.cancelled.wait(0.01) + raise AssertionError("Stream close did not cancel native synthesis") + return self.status + + def qt_audio_free(self, audio): + self.freed.set() + + +def make_tts(native): + tts = QwenTTS.__new__(QwenTTS) + tts._ctx = 123 + tts._lock = threading.Lock() + tts.library = SimpleNamespace(_lib=native, last_error=lambda: "native failure") + return tts + + +@pytest.mark.parametrize("first", [1, 2, 4, 8]) +@pytest.mark.parametrize("seconds,later", [(0.01, 1), (0.32, 4), (0.64, 8), (1.0, 13)]) +def test_packet_sizes_and_audio_are_independent_of_native_ramp(first, seconds, later): + native = NativeStream() + tts = make_tts(native) + packets = list(tts.stream(text="test", first_chunk_frames=first, codec_chunk_sec=seconds)) + sizes = [packet.size // CODEC_FRAME_SAMPLES for packet, _ in packets] + full, tail = divmod(39 - first, later) + assert sizes == [first] + [later] * full + ([tail] if tail else []) + assert all(rate == 24000 for _, rate in packets) + np.testing.assert_array_equal(np.concatenate([p for p, _ in packets]), native.audio) + assert native.freed.is_set() + profile = tts.last_stream_profile + assert profile["first_callback_n_samples"] == CODEC_FRAME_SAMPLES + assert profile["first_packet_n_samples"] == first * CODEC_FRAME_SAMPLES + assert profile["first_callback_enter_ms"] <= profile["first_packet_ready_ms"] <= profile["first_yield_ms"] + assert profile["packet_count"] == len(packets) + + +def test_default_packet_sequence_is_four_then_eight_frames_and_short_tail(): + packets = list(make_tts(NativeStream()).stream(text="test")) + assert [packet.size for packet, _ in packets] == [ + frames * CODEC_FRAME_SAMPLES for frames in (4, 8, 8, 8, 8, 3) + ] + + +@pytest.mark.parametrize("frames", [0, 1, 3, 7]) +def test_short_utterance_flushes_at_successful_eos(frames): + native = NativeStream(frames=frames) + packets = list(make_tts(native).stream(text="test", first_chunk_frames=8)) + assert len(packets) == (1 if frames else 0) + if frames: + np.testing.assert_array_equal(packets[0][0], native.audio) + assert native.freed.is_set() + + +@pytest.mark.parametrize("value", [0, -1, 3, 16, 4.0, True, "4", None]) +def test_invalid_first_packet_fails_before_native_call(value): + native = NativeStream() + with pytest.raises(ValueError, match="first_chunk_frames"): + list(make_tts(native).stream(text="test", first_chunk_frames=value)) + assert not native.called + + +@pytest.mark.parametrize("seconds", [0, -1, float("inf"), float("nan")]) +def test_invalid_later_packet_fails_before_native_call(seconds): + native = NativeStream() + with pytest.raises(ValueError, match="codec_chunk_sec"): + list(make_tts(native).stream(text="test", codec_chunk_sec=seconds)) + assert not native.called + + +@pytest.mark.parametrize("status", [QwenStatus.GENERATE_FAILED, QwenStatus.CANCELLED]) +def test_failure_discards_unfinished_packet_and_frees_native_audio(status): + native = NativeStream(frames=3, status=status) + stream = make_tts(native).stream(text="test", first_chunk_frames=4) + with pytest.raises(QwenTTSError, match="native failure"): + next(stream) + assert native.freed.is_set() + + +def test_failure_after_first_packet_preserves_packet_then_raises(): + native = NativeStream(frames=7, status=QwenStatus.GENERATE_FAILED) + stream = make_tts(native).stream(text="test", first_chunk_frames=4) + np.testing.assert_array_equal(next(stream)[0], native.audio[:4 * CODEC_FRAME_SAMPLES]) + with pytest.raises(QwenTTSError, match="native failure"): + next(stream) + assert native.freed.is_set() + + +def test_repeated_stream_resets_packet_state(): + tts = make_tts(NativeStream(frames=3)) + for _ in range(2): + packets = list(tts.stream(text="test", first_chunk_frames=4)) + assert len(packets) == 1 + assert packets[0][0].size == 3 * CODEC_FRAME_SAMPLES + assert tts.last_stream_profile["packet_count"] == 1 + + +def test_close_cancels_native_generation_with_partial_packet_pending(): + native = NativeStream(frames=7, wait_for_cancel=True) + tts = make_tts(native) + stream = tts.stream(text="test", first_chunk_frames=4) + assert next(stream)[0].size == 4 * CODEC_FRAME_SAMPLES + stream.close() + assert native.cancelled.is_set() + assert native.freed.is_set() + assert tts.last_stream_profile["packet_count"] == 1 + assert not tts._lock.locked()