diff --git a/README.md b/README.md index 77c5504..9f5d509 100644 --- a/README.md +++ b/README.md @@ -46,16 +46,16 @@ variants. Use them when the PyPI CUDA 12.8 wheel does not match the runtime or GPU target, for example DGX Spark / GB10 with CUDA 13: ```bash -pip install "qwentts-cpp-python==0.2.0+cpu" \ +pip install "qwentts-cpp-python==0.3.0+cpu" \ -f https://huggingface.co/datasets/andito/qwentts-cpp-python-wheels/resolve/main/whl/cpu.html -pip install "qwentts-cpp-python==0.2.0+cu124" \ +pip install "qwentts-cpp-python==0.3.0+cu124" \ -f https://huggingface.co/datasets/andito/qwentts-cpp-python-wheels/resolve/main/whl/cu124.html -pip install "qwentts-cpp-python==0.2.0+cu128" \ +pip install "qwentts-cpp-python==0.3.0+cu128" \ -f https://huggingface.co/datasets/andito/qwentts-cpp-python-wheels/resolve/main/whl/cu128.html -pip install "qwentts-cpp-python==0.2.0+cu130" \ +pip install "qwentts-cpp-python==0.3.0+cu130" \ -f https://huggingface.co/datasets/andito/qwentts-cpp-python-wheels/resolve/main/whl/cu130.html ``` @@ -95,6 +95,19 @@ passing precomputed latents: - `.spk`: raw float32 speaker embedding from `qwen-codec --talker` - `.rvq`: packed 11-bit reference codec stream from `qwen-codec` +The wrapper can create those files in-process from decoded mono float32 audio at +24 kHz: + +```python +from qwentts_cpp import QwenTTS + +tts = QwenTTS.from_pretrained("Qwen/Qwen3-TTS-12Hz-1.7B-Base", quant="Q4_K_M") + +# ref_audio_24k is a 1-D numpy float32 array, already resampled to 24 kHz. +voice_ref = tts.extract_voice_ref(ref_audio_24k) +voice_ref.save("reference.spk", "reference.rvq") +``` + ```python from qwentts_cpp import QwenTTS, load_speaker_embedding diff --git a/pyproject.toml b/pyproject.toml index 30c83e8..2ecebc2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "qwentts-cpp-python" -version = "0.2.0" +version = "0.3.0" description = "Python ctypes bindings and wheel packaging for qwentts.cpp" readme = "README.md" license = "MIT" diff --git a/src/qwentts_cpp/__init__.py b/src/qwentts_cpp/__init__.py index 0b573a6..71bd553 100644 --- a/src/qwentts_cpp/__init__.py +++ b/src/qwentts_cpp/__init__.py @@ -6,12 +6,17 @@ QwenStatus, QwenTTS, QwenTTSError, + VoiceRef, load_rvq_codes, load_speaker_embedding, + load_voice_ref, + save_rvq_codes, + save_speaker_embedding, + save_voice_ref, ) from .models import GGUF_REPO, resolve_gguf_paths -__version__ = "0.2.0" +__version__ = "0.3.0" __all__ = [ "GGUF_REPO", @@ -22,7 +27,12 @@ "QwenStatus", "QwenTTS", "QwenTTSError", + "VoiceRef", "load_rvq_codes", "load_speaker_embedding", + "load_voice_ref", "resolve_gguf_paths", + "save_rvq_codes", + "save_speaker_embedding", + "save_voice_ref", ] diff --git a/src/qwentts_cpp/_binding.py b/src/qwentts_cpp/_binding.py index ad59f52..16dd1e2 100644 --- a/src/qwentts_cpp/_binding.py +++ b/src/qwentts_cpp/_binding.py @@ -6,6 +6,7 @@ import sys import threading import time +from dataclasses import dataclass from enum import IntEnum from pathlib import Path from typing import Any, Iterator, Sequence, Tuple @@ -97,12 +98,74 @@ class QtTTSParams(ctypes.Structure): ] +class QtVoiceRef(ctypes.Structure): + _fields_ = [ + ("ref_spk_emb", ctypes.POINTER(ctypes.c_float)), + ("ref_spk_dim", ctypes.c_int), + ("ref_codes", ctypes.POINTER(ctypes.c_int32)), + ("ref_T", ctypes.c_int), + ("num_codebooks", ctypes.c_int), + ] + + +@dataclass(frozen=True) +class VoiceRef: + """Reusable Base voice-clone conditioning extracted by qwentts.cpp.""" + + ref_spk_emb: np.ndarray + ref_codes: np.ndarray + + def __post_init__(self) -> None: + object.__setattr__(self, "ref_spk_emb", _prepare_speaker_embedding(self.ref_spk_emb)) + object.__setattr__(self, "ref_codes", _prepare_rvq_matrix(self.ref_codes)) + + @property + def num_codebooks(self) -> int: + return int(self.ref_codes.shape[0]) + + @property + def ref_T(self) -> int: + return int(self.ref_codes.shape[1]) + + def save( + self, + spk_path: str | os.PathLike[str], + rvq_path: str | os.PathLike[str], + *, + code_bits: int = RVQ_CODE_BITS, + ) -> tuple[Path, Path]: + """Write this reference as qwentts.cpp-compatible `.spk` and `.rvq` files.""" + return save_voice_ref(self, spk_path, rvq_path, code_bits=code_bits) + + +def _prepare_speaker_embedding(embedding: np.ndarray) -> np.ndarray: + spk = np.ascontiguousarray(embedding, dtype=np.float32).reshape(-1) + if spk.size == 0: + raise ValueError("Speaker embedding must not be empty") + return spk + + +def _prepare_rvq_matrix(codes: np.ndarray) -> np.ndarray: + rvq = np.asarray(codes, dtype=np.int32) + if rvq.ndim != 2 or rvq.shape[0] <= 0 or rvq.shape[1] <= 0: + raise ValueError("RVQ codes must have shape [num_codebooks, T] with positive dimensions") + return np.ascontiguousarray(rvq, dtype=np.int32) + + def load_speaker_embedding(path: str | os.PathLike[str]) -> np.ndarray: """Load a qwentts.cpp `.spk` file as a contiguous float32 vector.""" data = np.fromfile(path, dtype=np.float32) if data.size == 0: raise ValueError(f"Speaker embedding file is empty: {path}") - return np.ascontiguousarray(data, dtype=np.float32) + return _prepare_speaker_embedding(data) + + +def save_speaker_embedding(path: str | os.PathLike[str], embedding: np.ndarray) -> Path: + """Write a qwentts.cpp `.spk` file containing raw float32 speaker values.""" + output = Path(path) + output.parent.mkdir(parents=True, exist_ok=True) + _prepare_speaker_embedding(embedding).tofile(output) + return output def load_rvq_codes( @@ -131,6 +194,47 @@ def load_rvq_codes( return codes.reshape(int(num_codebooks), n_codes // int(num_codebooks)) +def save_rvq_codes( + path: str | os.PathLike[str], + codes: np.ndarray, + *, + code_bits: int = RVQ_CODE_BITS, +) -> Path: + """Write qwentts.cpp `.rvq` packed 11-bit reference codec codes.""" + output = Path(path) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_bytes(_pack_rvq_codes(_prepare_rvq_matrix(codes).reshape(-1), int(code_bits))) + return output + + +def load_voice_ref( + spk_path: str | os.PathLike[str], + rvq_path: str | os.PathLike[str], + num_codebooks: int, + *, + code_bits: int = RVQ_CODE_BITS, +) -> VoiceRef: + """Load reusable Base voice-clone conditioning from `.spk` and `.rvq` files.""" + return VoiceRef( + ref_spk_emb=load_speaker_embedding(spk_path), + ref_codes=load_rvq_codes(rvq_path, num_codebooks, code_bits=code_bits), + ) + + +def save_voice_ref( + voice_ref: VoiceRef, + spk_path: str | os.PathLike[str], + rvq_path: str | os.PathLike[str], + *, + code_bits: int = RVQ_CODE_BITS, +) -> tuple[Path, Path]: + """Write a reusable Base voice-clone reference to `.spk` and `.rvq` files.""" + ref = VoiceRef(voice_ref.ref_spk_emb, voice_ref.ref_codes) + spk = save_speaker_embedding(spk_path, ref.ref_spk_emb) + rvq = save_rvq_codes(rvq_path, ref.ref_codes, code_bits=code_bits) + return spk, rvq + + def _unpack_rvq_codes(packed: np.ndarray, n_codes: int, code_bits: int) -> np.ndarray: mask = (1 << code_bits) - 1 out = np.empty(n_codes, dtype=np.int32) @@ -149,6 +253,37 @@ def _unpack_rvq_codes(packed: np.ndarray, n_codes: int, code_bits: int) -> np.nd return out +def _pack_rvq_codes(codes: np.ndarray, code_bits: int) -> bytes: + if code_bits <= 0 or code_bits >= 32: + raise ValueError(f"code_bits must be in [1, 31], got {code_bits}") + flat = np.asarray(codes, dtype=np.int64).reshape(-1) + if flat.size == 0: + raise ValueError("RVQ codes must not be empty") + + max_code = (1 << code_bits) - 1 + invalid = (flat < 0) | (flat > max_code) + if bool(np.any(invalid)): + bad = int(flat[np.nonzero(invalid)[0][0]]) + raise ValueError(f"RVQ code {bad} is outside the {code_bits}-bit range [0, {max_code}]") + + total_bits = int(flat.size) * int(code_bits) + out = bytearray((total_bits + 7) // 8) + acc = 0 + bits_in_acc = 0 + out_pos = 0 + for code in flat.tolist(): + acc |= int(code) << bits_in_acc + bits_in_acc += int(code_bits) + while bits_in_acc >= 8: + out[out_pos] = acc & 0xFF + out_pos += 1 + acc >>= 8 + bits_in_acc -= 8 + if bits_in_acc > 0: + out[out_pos] = acc & 0xFF + return bytes(out) + + def _as_utf8(value: str | os.PathLike[str] | None, keepalive: list[object]) -> bytes | None: if value is None: return None @@ -224,6 +359,8 @@ def __init__(self, library_path: str | os.PathLike[str] | None = None): self._has_qt_num_codebooks = False self._has_qt_n_speakers = False self._has_qt_speaker_name = False + self._has_qt_extract_voice_ref = False + self._has_qt_voice_ref_free = False self._lib = self._load_cdll(self.path) self._bind() @@ -289,6 +426,23 @@ def _bind(self) -> None: self._has_qt_speaker_name = True except AttributeError: self._has_qt_speaker_name = False + try: + lib.qt_extract_voice_ref.argtypes = [ + ctypes.c_void_p, + ctypes.POINTER(ctypes.c_float), + ctypes.c_int, + ctypes.POINTER(QtVoiceRef), + ] + lib.qt_extract_voice_ref.restype = ctypes.c_int + self._has_qt_extract_voice_ref = True + except AttributeError: + self._has_qt_extract_voice_ref = False + try: + lib.qt_voice_ref_free.argtypes = [ctypes.POINTER(QtVoiceRef)] + lib.qt_voice_ref_free.restype = None + self._has_qt_voice_ref_free = True + except AttributeError: + self._has_qt_voice_ref_free = False def version(self) -> str: return self._lib.qt_version().decode("utf-8", errors="replace") @@ -333,6 +487,7 @@ def __init__( self._lock = threading.Lock() self.last_synthesize_profile: dict[str, Any] | None = None self.last_stream_profile: dict[str, Any] | None = None + self.last_extract_voice_ref_profile: dict[str, Any] | None = None self._init(talker_path, codec_path, use_fa=use_fa, clamp_fp16=clamp_fp16) @classmethod @@ -430,6 +585,79 @@ def speaker_names(self) -> list[str]: def load_rvq_codes(self, path: str | os.PathLike[str], *, code_bits: int = RVQ_CODE_BITS) -> np.ndarray: return load_rvq_codes(path, self.num_codebooks(), code_bits=code_bits) + def load_voice_ref( + self, + spk_path: str | os.PathLike[str], + rvq_path: str | os.PathLike[str], + *, + code_bits: int = RVQ_CODE_BITS, + ) -> VoiceRef: + return load_voice_ref(spk_path, rvq_path, self.num_codebooks(), code_bits=code_bits) + + def extract_voice_ref(self, ref_audio_24k: np.ndarray) -> VoiceRef: + """Extract reusable Base voice-clone conditioning from 24 kHz mono audio.""" + if not (self.library._has_qt_extract_voice_ref and self.library._has_qt_voice_ref_free): + raise QwenTTSError("qt_extract_voice_ref is unavailable; voice reference extraction requires qwentts.cpp ABI v2") + + profile: dict[str, Any] = {} + start = time.perf_counter() + audio = np.ascontiguousarray(ref_audio_24k, dtype=np.float32).reshape(-1) + if audio.size == 0: + raise ValueError("ref_audio_24k must not be empty") + profile["audio_prepare_ms"] = (time.perf_counter() - start) * 1000 + profile["ref_n_samples"] = int(audio.size) + + out = QtVoiceRef() + lock_start = time.perf_counter() + with self._lock: + profile["lock_wait_ms"] = (time.perf_counter() - lock_start) * 1000 + native_start = time.perf_counter() + rc = self.library._lib.qt_extract_voice_ref( + self._require_ctx(), + audio.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), + int(audio.size), + ctypes.byref(out), + ) + profile["native_extract_ms"] = (time.perf_counter() - native_start) * 1000 + + try: + if rc != QwenStatus.OK: + raise QwenTTSError(self.library.last_error() or f"qt_extract_voice_ref failed with status {rc}") + if not out.ref_spk_emb or out.ref_spk_dim <= 0: + raise QwenTTSError("qt_extract_voice_ref returned an empty speaker embedding") + if not out.ref_codes or out.num_codebooks <= 0 or out.ref_T <= 0: + raise QwenTTSError("qt_extract_voice_ref returned empty RVQ codes") + + copy_start = time.perf_counter() + spk = np.ctypeslib.as_array(out.ref_spk_emb, shape=(int(out.ref_spk_dim),)).copy() + codes = np.ctypeslib.as_array( + out.ref_codes, + shape=(int(out.num_codebooks) * int(out.ref_T),), + ).copy() + codes = codes.reshape(int(out.num_codebooks), int(out.ref_T)) + profile["copy_ms"] = (time.perf_counter() - copy_start) * 1000 + profile["ref_spk_dim"] = int(out.ref_spk_dim) + profile["num_codebooks"] = int(out.num_codebooks) + profile["ref_T"] = int(out.ref_T) + profile["total_ms"] = (time.perf_counter() - start) * 1000 + self.last_extract_voice_ref_profile = profile + return VoiceRef(ref_spk_emb=spk, ref_codes=codes) + finally: + self.library._lib.qt_voice_ref_free(ctypes.byref(out)) + + def save_voice_ref( + self, + ref_audio_24k: np.ndarray, + spk_path: str | os.PathLike[str], + rvq_path: str | os.PathLike[str], + *, + code_bits: int = RVQ_CODE_BITS, + ) -> VoiceRef: + """Extract and save reusable Base voice-clone conditioning from reference audio.""" + voice_ref = self.extract_voice_ref(ref_audio_24k) + voice_ref.save(spk_path, rvq_path, code_bits=code_bits) + return voice_ref + def set_log_callback(self, callback) -> None: self.library.set_log_callback(callback) diff --git a/tests/test_binding.py b/tests/test_binding.py index faa6049..ba86fe1 100644 --- a/tests/test_binding.py +++ b/tests/test_binding.py @@ -1,12 +1,25 @@ from __future__ import annotations +import ctypes import os +import threading import numpy as np import pytest -from qwentts_cpp import LibraryNotFoundError, QwenLibrary, load_rvq_codes, load_speaker_embedding -from qwentts_cpp._binding import QtTTSParams +from qwentts_cpp import ( + LibraryNotFoundError, + QwenLibrary, + QwenTTS, + VoiceRef, + load_rvq_codes, + load_speaker_embedding, + load_voice_ref, + save_rvq_codes, + save_speaker_embedding, + save_voice_ref, +) +from qwentts_cpp._binding import QtTTSParams, QtVoiceRef def _pack_rvq_codes(codes, code_bits=11): @@ -53,6 +66,16 @@ def test_tts_params_contains_abi_v2_latent_tail_fields(): ] +def test_voice_ref_struct_matches_abi(): + assert [name for name, _ctype in QtVoiceRef._fields_] == [ + "ref_spk_emb", + "ref_spk_dim", + "ref_codes", + "ref_T", + "num_codebooks", + ] + + def test_load_speaker_embedding_reads_raw_float32(tmp_path): path = tmp_path / "speaker.spk" expected = np.array([0.25, -0.5, 1.0], dtype=np.float32) @@ -65,6 +88,16 @@ def test_load_speaker_embedding_reads_raw_float32(tmp_path): np.testing.assert_array_equal(loaded, expected) +def test_save_speaker_embedding_writes_raw_float32(tmp_path): + path = tmp_path / "nested" / "speaker.spk" + expected = np.array([0.25, -0.5, 1.0], dtype=np.float32) + + saved = save_speaker_embedding(path, expected) + + assert saved == path + np.testing.assert_array_equal(load_speaker_embedding(path), expected) + + def test_load_rvq_codes_unpacks_lsb_first_matrix(tmp_path): path = tmp_path / "reference.rvq" expected = np.array( @@ -84,9 +117,124 @@ def test_load_rvq_codes_unpacks_lsb_first_matrix(tmp_path): np.testing.assert_array_equal(loaded, expected) +def test_save_rvq_codes_packs_lsb_first_matrix(tmp_path): + path = tmp_path / "nested" / "reference.rvq" + expected = np.array( + [ + [1, 2, 3], + [2047, 17, 42], + [0, 999, 123], + [456, 789, 1024], + ], + dtype=np.int32, + ) + + saved = save_rvq_codes(path, expected) + + assert saved == path + assert path.read_bytes() == _pack_rvq_codes(expected.reshape(-1).tolist()) + np.testing.assert_array_equal(load_rvq_codes(path, num_codebooks=expected.shape[0]), expected) + + def test_load_rvq_codes_rejects_wrong_codebook_count(tmp_path): path = tmp_path / "reference.rvq" path.write_bytes(_pack_rvq_codes([1, 2, 3, 4])) with pytest.raises(ValueError, match="num_codebooks"): load_rvq_codes(path, num_codebooks=3) + + +def test_save_rvq_codes_rejects_out_of_range_codes(tmp_path): + with pytest.raises(ValueError, match="outside"): + save_rvq_codes(tmp_path / "bad.rvq", np.array([[0, 2048]], dtype=np.int32)) + + +def test_voice_ref_save_and_load_round_trips_files(tmp_path): + spk = np.array([0.25, -0.5, 1.0], dtype=np.float32) + codes = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int32) + ref = VoiceRef(spk, codes) + + spk_path, rvq_path = save_voice_ref(ref, tmp_path / "voice.spk", tmp_path / "voice.rvq") + loaded = load_voice_ref(spk_path, rvq_path, num_codebooks=2) + + assert ref.num_codebooks == 2 + assert ref.ref_T == 3 + np.testing.assert_array_equal(loaded.ref_spk_emb, spk) + np.testing.assert_array_equal(loaded.ref_codes, codes) + + +def test_voice_ref_instance_save_round_trips_files(tmp_path): + ref = VoiceRef( + np.array([0.5, 1.5], dtype=np.float32), + np.array([[7, 8], [9, 10]], dtype=np.int32), + ) + + ref.save(tmp_path / "speaker.spk", tmp_path / "reference.rvq") + + loaded = load_voice_ref(tmp_path / "speaker.spk", tmp_path / "reference.rvq", num_codebooks=2) + np.testing.assert_array_equal(loaded.ref_spk_emb, ref.ref_spk_emb) + np.testing.assert_array_equal(loaded.ref_codes, ref.ref_codes) + + +class _FakeExtractLib: + def __init__(self): + self.free_calls = 0 + self.spk_buf = None + self.codes_buf = None + self.ctx = None + self.ref_audio = None + + def qt_extract_voice_ref(self, ctx, audio_ptr, n_samples, out_ptr): + self.ctx = ctx + self.ref_audio = np.ctypeslib.as_array(audio_ptr, shape=(n_samples,)).copy() + self.spk_buf = (ctypes.c_float * 3)(0.25, -0.5, 1.0) + self.codes_buf = (ctypes.c_int32 * 6)(1, 2, 3, 4, 5, 6) + out = out_ptr._obj + out.ref_spk_emb = ctypes.cast(self.spk_buf, ctypes.POINTER(ctypes.c_float)) + out.ref_spk_dim = 3 + out.ref_codes = ctypes.cast(self.codes_buf, ctypes.POINTER(ctypes.c_int32)) + out.ref_T = 3 + out.num_codebooks = 2 + return 0 + + def qt_voice_ref_free(self, out_ptr): + self.free_calls += 1 + out = out_ptr._obj + out.ref_spk_emb = ctypes.POINTER(ctypes.c_float)() + out.ref_spk_dim = 0 + out.ref_codes = ctypes.POINTER(ctypes.c_int32)() + out.ref_T = 0 + out.num_codebooks = 0 + + +class _FakeLibrary: + def __init__(self, lib): + self._lib = lib + self._has_qt_extract_voice_ref = True + self._has_qt_voice_ref_free = True + + def last_error(self): + return "fake error" + + +def test_extract_voice_ref_copies_native_buffers_before_free(): + fake_lib = _FakeExtractLib() + tts = QwenTTS.__new__(QwenTTS) + tts.library = _FakeLibrary(fake_lib) + tts._ctx = 123 + tts._lock = threading.Lock() + tts.last_extract_voice_ref_profile = None + + ref = tts.extract_voice_ref(np.array([0.0, 0.5, -0.5], dtype=np.float64)) + + assert fake_lib.ctx == 123 + assert fake_lib.free_calls == 1 + np.testing.assert_array_equal(fake_lib.ref_audio, np.array([0.0, 0.5, -0.5], dtype=np.float32)) + np.testing.assert_array_equal(ref.ref_spk_emb, np.array([0.25, -0.5, 1.0], dtype=np.float32)) + np.testing.assert_array_equal(ref.ref_codes, np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int32)) + + fake_lib.spk_buf[0] = 99.0 + fake_lib.codes_buf[0] = 99 + assert ref.ref_spk_emb[0] == np.float32(0.25) + assert ref.ref_codes[0, 0] == 1 + assert tts.last_extract_voice_ref_profile["ref_spk_dim"] == 3