wfloat 2.0.0__py3-none-win_amd64.whl

wfloat/_model.py

@@ -0,0 +1,394 @@
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict, List, Mapping, Optional, Sequence
+
+from ._assets import fetch_model_assets
+from ._cache import (
+    CachedModelAssets,
+    cache_model_assets,
+)
+from ._constants import (
+    DEFAULT_MODEL_NAME,
+    DEFAULT_SILENCE_BETWEEN_SEGMENTS_SEC,
+    normalize_emotion,
+    normalize_intensity,
+    normalize_silence_padding_sec,
+    normalize_speed,
+    normalize_text,
+    normalize_voice_id,
+)
+from ._native import create_native_tts
+from ._results import Audio, GenerationResult, Timeline, TimelineChunk
+
+
+@dataclass
+class _PreparedChunk:
+    text: str
+    text_clean: str
+    highlight_start: int
+    highlight_end: int
+    sid: int
+    voice_id: Optional[object]
+    emotion: str
+    intensity: float
+    speed: float
+    silence_padding_sec: float
+    segment_index: Optional[int]
+
+
+class Model:
+    def __init__(
+        self,
+        model_name: str,
+        native_tts,
+        *,
+        cached_assets: Optional[CachedModelAssets] = None,
+    ) -> None:
+        self.model_name = model_name
+        self._native_tts = native_tts
+        self._cached_assets = cached_assets
+
+    @property
+    def sample_rate(self) -> int:
+        return int(self._native_tts.sample_rate)
+
+    @property
+    def num_speakers(self) -> int:
+        return int(self._native_tts.num_speakers)
+
+    def __repr__(self) -> str:
+        return "Model(model_name=%r, sample_rate=%r, num_speakers=%r)" % (
+            self.model_name,
+            self.sample_rate,
+            self.num_speakers,
+        )
+
+    def generate(
+        self,
+        *,
+        text: str,
+        voice_id=None,
+        emotion: Optional[str] = None,
+        intensity: Optional[float] = None,
+        speed: Optional[float] = None,
+        silence_padding_sec: Optional[float] = None,
+    ) -> GenerationResult:
+        normalized_text = normalize_text(text)
+        normalized_emotion = normalize_emotion(emotion)
+        normalized_intensity = normalize_intensity(intensity)
+        normalized_speed = normalize_speed(speed)
+        normalized_silence_padding_sec = normalize_silence_padding_sec(
+            silence_padding_sec
+        )
+        normalized_voice_id = voice_id
+        sid = normalize_voice_id(voice_id)
+
+        if hasattr(self._native_tts, "synthesize_result"):
+            return self._native_tts.synthesize_result(
+                model_id=self.model_name,
+                text=normalized_text,
+                voice=normalized_voice_id,
+                sid=sid,
+                emotion=normalized_emotion,
+                intensity=normalized_intensity,
+                speed=normalized_speed,
+                silence_padding_sec=normalized_silence_padding_sec,
+            )
+
+        prepared = self._native_tts.prepare_wfloat_text(
+            normalized_text,
+            normalized_emotion,
+            normalized_intensity,
+        )
+        prepared_chunks = self._build_prepared_chunks(
+            prepared.text,
+            prepared.text_clean,
+            sid=sid,
+            voice_id=voice_id,
+            emotion=normalized_emotion,
+            intensity=normalized_intensity,
+            speed=normalized_speed,
+            silence_padding_sec=normalized_silence_padding_sec,
+            segment_index=None,
+        )
+
+        result = self._synthesize_chunks(
+            prepared_chunks,
+            result_text=normalized_text,
+        )
+        return result
+
+    def synthesize(
+        self,
+        *,
+        text: str,
+        voice: Optional[object] = None,
+        speed: Optional[float] = None,
+        emotion: Optional[str] = None,
+        intensity: Optional[float] = None,
+        silence_padding_sec: Optional[float] = None,
+    ) -> GenerationResult:
+        return self.generate(
+            text=text,
+            voice_id=voice,
+            emotion=emotion,
+            intensity=intensity,
+            speed=speed,
+            silence_padding_sec=silence_padding_sec,
+        )
+
+    def generate_dialogue(
+        self,
+        *,
+        segments: Sequence[Mapping[str, Any]],
+        speed: Optional[float] = None,
+        silence_between_segments_sec: Optional[float] = None,
+    ) -> GenerationResult:
+        if not segments:
+            raise ValueError("segments is required.")
+
+        default_speed = normalize_speed(speed)
+        dialogue_silence_sec = normalize_silence_padding_sec(
+            silence_between_segments_sec,
+            default=DEFAULT_SILENCE_BETWEEN_SEGMENTS_SEC,
+        )
+
+        if hasattr(self._native_tts, "synthesize_dialogue_result"):
+            normalized_segments = []
+            for segment in segments:
+                if not isinstance(segment, Mapping):
+                    raise TypeError("Each segment must be a mapping.")
+
+                normalized_segments.append(
+                    self._normalize_dialogue_segment(
+                        segment,
+                        default_speed=default_speed,
+                    )
+                )
+
+            return self._native_tts.synthesize_dialogue_result(
+                model_id=self.model_name,
+                segments=normalized_segments,
+                silence_between_segments_sec=dialogue_silence_sec,
+            )
+
+        prepared_chunks = []
+        dialogue_texts = []
+
+        for segment_index, segment in enumerate(segments):
+            if not isinstance(segment, Mapping):
+                raise TypeError("Each segment must be a mapping.")
+
+            normalized_segment = self._normalize_dialogue_segment(
+                segment,
+                default_speed=default_speed,
+            )
+            dialogue_texts.append(normalized_segment["text"])
+
+            prepared = self._native_tts.prepare_wfloat_text(
+                normalized_segment["text"],
+                normalized_segment["emotion"],
+                normalized_segment["intensity"],
+            )
+            segment_chunks = self._build_prepared_chunks(
+                prepared.text,
+                prepared.text_clean,
+                sid=normalized_segment["sid"],
+                voice_id=normalized_segment["voice_id"],
+                emotion=normalized_segment["emotion"],
+                intensity=normalized_segment["intensity"],
+                speed=normalized_segment["speed"],
+                silence_padding_sec=normalized_segment["sentence_silence_padding_sec"],
+                segment_index=segment_index,
+            )
+            prepared_chunks.extend(segment_chunks)
+
+            if segment_index < len(segments) - 1 and segment_chunks:
+                prepared_chunks[-1].silence_padding_sec = dialogue_silence_sec
+
+        return self._synthesize_chunks(
+            prepared_chunks,
+            result_text=" ".join(dialogue_texts),
+        )
+
+    def synthesize_dialogue(
+        self,
+        *,
+        segments: Sequence[Mapping[str, Any]],
+        speed: Optional[float] = None,
+        silence_between_segments_sec: Optional[float] = None,
+    ) -> GenerationResult:
+        return self.generate_dialogue(
+            segments=segments,
+            speed=speed,
+            silence_between_segments_sec=silence_between_segments_sec,
+        )
+
+    def _normalize_dialogue_segment(
+        self,
+        segment: Mapping[str, Any],
+        *,
+        default_speed: float,
+    ) -> Dict[str, Any]:
+        text = normalize_text(segment.get("text"))  # type: ignore[arg-type]
+        voice_id = segment.get("voice", segment.get("voice_id"))
+        emotion = normalize_emotion(segment.get("emotion"))
+        intensity = normalize_intensity(segment.get("intensity"))
+        speed = normalize_speed(segment.get("speed"), default=default_speed)
+        sentence_silence_padding_sec = normalize_silence_padding_sec(
+            segment.get("sentence_silence_padding_sec")
+        )
+        sid = normalize_voice_id(voice_id)
+
+        return {
+            "text": text,
+            "voice_id": voice_id,
+            "emotion": emotion,
+            "intensity": intensity,
+            "speed": speed,
+            "sentence_silence_padding_sec": sentence_silence_padding_sec,
+            "sid": sid,
+        }
+
+    def _build_prepared_chunks(
+        self,
+        raw_text_chunks: Sequence[str],
+        clean_text_chunks: Sequence[str],
+        *,
+        sid: int,
+        voice_id,
+        emotion: str,
+        intensity: float,
+        speed: float,
+        silence_padding_sec: float,
+        segment_index: Optional[int],
+    ) -> List[_PreparedChunk]:
+        chunks = []
+        raw_cursor = 0
+        for index, raw_chunk_text in enumerate(raw_text_chunks):
+            highlight_start = raw_cursor
+            highlight_end = raw_cursor + len(raw_chunk_text)
+            raw_cursor = highlight_end
+
+            chunks.append(
+                _PreparedChunk(
+                    text=raw_chunk_text,
+                    text_clean=clean_text_chunks[index],
+                    highlight_start=highlight_start,
+                    highlight_end=highlight_end,
+                    sid=sid,
+                    voice_id=voice_id,
+                    emotion=emotion,
+                    intensity=intensity,
+                    speed=speed,
+                    silence_padding_sec=silence_padding_sec,
+                    segment_index=segment_index,
+                )
+            )
+        return chunks
+
+    def _synthesize_chunks(
+        self,
+        chunks: Sequence[_PreparedChunk],
+        *,
+        result_text: str,
+    ) -> GenerationResult:
+        if not chunks:
+            raise RuntimeError("Text preparation produced no synthesizeable chunks.")
+
+        all_samples = []
+        timeline_chunks = []
+        total_chunks = len(chunks)
+        cumulative_samples = 0
+        sample_rate = self.sample_rate
+
+        for index, chunk in enumerate(chunks):
+            generated_audio = self._native_tts.generate(
+                chunk.text_clean,
+                chunk.sid,
+                chunk.speed,
+            )
+            chunk_samples = [float(sample) for sample in generated_audio.samples]
+            chunk_sample_rate = int(generated_audio.sample_rate)
+            if chunk_sample_rate <= 0:
+                raise RuntimeError("Native generation returned an invalid sample rate.")
+
+            if index == 0:
+                sample_rate = chunk_sample_rate
+            elif chunk_sample_rate != sample_rate:
+                raise RuntimeError(
+                    "Native generation returned inconsistent sample rates across chunks."
+                )
+
+            start_sec = cumulative_samples / float(sample_rate)
+            all_samples.extend(chunk_samples)
+            cumulative_samples += len(chunk_samples)
+            end_sec = cumulative_samples / float(sample_rate)
+
+            timeline_chunks.append(
+                TimelineChunk(
+                    index=index,
+                    text=chunk.text,
+                    highlight_start=chunk.highlight_start,
+                    highlight_end=chunk.highlight_end,
+                    start_sec=start_sec,
+                    end_sec=end_sec,
+                    duration_sec=end_sec - start_sec,
+                    progress=float(index + 1) / float(total_chunks),
+                    voice_id=chunk.voice_id,
+                    sid=chunk.sid,
+                    emotion=chunk.emotion,
+                    intensity=chunk.intensity,
+                    speed=chunk.speed,
+                    segment_index=chunk.segment_index,
+                )
+            )
+
+            if index < total_chunks - 1 and chunk.silence_padding_sec > 0:
+                silence_samples = int(round(chunk.silence_padding_sec * sample_rate))
+                if silence_samples > 0:
+                    all_samples.extend([0.0] * silence_samples)
+                    cumulative_samples += silence_samples
+
+        audio = Audio(samples=all_samples, sample_rate=sample_rate)
+        timeline = Timeline(chunks=timeline_chunks, duration_sec=audio.duration_sec)
+        result = GenerationResult(
+            audio=audio,
+            timeline=timeline,
+            text=result_text,
+            model_name=self.model_name,
+        )
+
+        return result
+
+
+def load(
+    model_name: str = DEFAULT_MODEL_NAME,
+    *,
+    cache_dir=None,
+    force_download: bool = False,
+) -> Model:
+    normalized_model_name = normalize_text(model_name)
+    resolved_cache_dir = Path(cache_dir) if cache_dir is not None else None
+    assets = fetch_model_assets(normalized_model_name)
+    cached_assets = cache_model_assets(
+        normalized_model_name,
+        assets,
+        cache_dir=resolved_cache_dir,
+        force_download=force_download,
+    )
+    native_tts = create_native_tts(
+        normalized_model_name,
+        cached_assets.model_path,
+        cached_assets.tokens_path,
+        cached_assets.espeak_data_dir,
+    )
+    return Model(
+        normalized_model_name,
+        native_tts,
+        cached_assets=cached_assets,
+    )
+
+
+TtsModel = Model
+load_tts_model = load

wfloat/_native.py

@@ -0,0 +1,17 @@
+from pathlib import Path
+
+from ._core import create_core_tts
+
+
+def create_native_tts(
+    model_name: str,
+    model_path: Path,
+    tokens_path: Path,
+    espeak_data_dir: Path,
+):
+    return create_core_tts(
+        model_name=model_name,
+        model_path=model_path,
+        tokens_path=tokens_path,
+        espeak_data_dir=espeak_data_dir,
+    )

wfloat/_results.py

@@ -0,0 +1,163 @@
+import io
+import struct
+import wave
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Iterable, List, Optional
+
+
+def _coerce_samples(samples: Iterable[float]) -> List[float]:
+    return [float(sample) for sample in samples]
+
+
+def _sample_to_pcm16(sample: float) -> int:
+    clipped = max(-1.0, min(1.0, float(sample)))
+    if clipped <= -1.0:
+        return -32768
+    return int(round(clipped * 32767.0))
+
+
+def _write_wave_bytes(samples: List[float], sample_rate: int) -> bytes:
+    buffer = io.BytesIO()
+    with wave.open(buffer, "wb") as wav_file:
+        wav_file.setnchannels(1)
+        wav_file.setsampwidth(2)
+        wav_file.setframerate(sample_rate)
+        frame_bytes = bytearray()
+        for sample in samples:
+            frame_bytes.extend(struct.pack("<h", _sample_to_pcm16(sample)))
+        wav_file.writeframes(bytes(frame_bytes))
+    return buffer.getvalue()
+
+
+@dataclass
+class Audio:
+    samples: List[float]
+    sample_rate: int
+
+    def __post_init__(self) -> None:
+        self.samples = _coerce_samples(self.samples)
+        self.sample_rate = int(self.sample_rate)
+
+    @property
+    def duration_sec(self) -> float:
+        if self.sample_rate <= 0:
+            return 0.0
+        return len(self.samples) / float(self.sample_rate)
+
+    def save(self, path) -> None:
+        output_path = Path(path)
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+        output_path.write_bytes(self.wav_bytes())
+
+    def wav_bytes(self) -> bytes:
+        return _write_wave_bytes(self.samples, self.sample_rate)
+
+
+AudioResult = Audio
+
+
+@dataclass
+class TimelineChunk:
+    index: int
+    text: str
+    highlight_start: int
+    highlight_end: int
+    start_sec: float
+    end_sec: float
+    duration_sec: float
+    progress: float
+    voice_id: Optional[object]
+    sid: int
+    emotion: str
+    intensity: float
+    speed: float
+    segment_index: Optional[int] = None
+
+
+@dataclass
+class Timeline:
+    chunks: List[TimelineChunk]
+    duration_sec: float
+
+
+@dataclass
+class GenerationResult:
+    audio: Audio
+    timeline: Timeline
+    text: str
+    model_name: str
+
+    @property
+    def model_id(self) -> str:
+        return self.model_name
+
+    def __iter__(self):
+        yield self.audio
+        yield self.timeline
+
+
+TtsSynthesisResult = GenerationResult
+
+
+@dataclass
+class TranscriptionToken:
+    text: str
+    start_sec: float
+    duration_sec: float
+    confidence: float = 0.0
+
+
+@dataclass
+class TranscriptionSegment:
+    text: str
+    start_sec: float
+    duration_sec: float
+
+
+@dataclass
+class TranscriptionResult:
+    text: str
+    model_id: str
+    language: str = ""
+    emotion: str = ""
+    event: str = ""
+    json: str = ""
+    tokens: Optional[List[TranscriptionToken]] = None
+    segments: Optional[List[TranscriptionSegment]] = None
+
+
+@dataclass
+class StreamingTranscriptionResult:
+    text: str
+    model_id: str
+    is_endpoint: bool = False
+    json: str = ""
+
+
+@dataclass
+class LlmGenerationResult:
+    text: str
+    model_id: str
+    finish_reason: str = ""
+    json: str = ""
+    prompt_token_count: int = 0
+    completion_token_count: int = 0
+
+
+@dataclass
+class VadSegment:
+    start_sec: float
+    duration_sec: float
+    end_sec: float
+    start_sample: int
+    sample_count: int
+    sample_rate: int
+    audio: Audio
+
+
+@dataclass
+class VadDetectionResult:
+    model_id: str
+    segments: List[VadSegment]
+    speech_ratio: float

wfloat/_stt.py

@@ -0,0 +1,122 @@
+from __future__ import annotations
+
+import wave
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Iterable, Optional, Sequence
+
+from ._results import StreamingTranscriptionResult, TranscriptionResult
+
+
+def _coerce_audio_samples(audio: bytes | Sequence[float]) -> list[float]:
+    if isinstance(audio, (bytes, bytearray, memoryview)):
+        raise TypeError(
+            "Raw PCM bytes are not supported yet. Pass a float sequence or a WAV path."
+        )
+
+    return [float(sample) for sample in audio]
+
+
+def _load_wav_audio(path: Path) -> tuple[list[float], int]:
+    with wave.open(str(path), "rb") as wav_file:
+        channels = wav_file.getnchannels()
+        sample_width = wav_file.getsampwidth()
+        sample_rate = wav_file.getframerate()
+        frame_count = wav_file.getnframes()
+
+        if channels != 1:
+            raise ValueError("Only mono WAV files are supported.")
+        if sample_width != 2:
+            raise ValueError("Only 16-bit PCM WAV files are supported.")
+
+        frames = wav_file.readframes(frame_count)
+
+    samples: list[float] = []
+    for index in range(0, len(frames), 2):
+        sample = int.from_bytes(frames[index : index + 2], "little", signed=True)
+        samples.append(float(sample) / 32768.0)
+
+    return samples, sample_rate
+
+
+@dataclass
+class SttSession:
+    model_id: str
+    _native_session: object
+
+    def push(
+        self,
+        audio: Sequence[float],
+        *,
+        sample_rate: Optional[int] = None,
+    ) -> None:
+        if not hasattr(self._native_session, "push"):
+            raise RuntimeError("Native STT backend does not support push().")
+
+        self._native_session.push(audio, sample_rate=sample_rate)
+
+    def get_result(self) -> StreamingTranscriptionResult:
+        if not hasattr(self._native_session, "get_result"):
+            raise RuntimeError("Native STT backend does not support get_result().")
+
+        return self._native_session.get_result()
+
+    def finish(self) -> StreamingTranscriptionResult:
+        if not hasattr(self._native_session, "finish"):
+            raise RuntimeError("Native STT backend does not support finish().")
+
+        return self._native_session.finish()
+
+    def reset(self) -> None:
+        if not hasattr(self._native_session, "reset"):
+            raise RuntimeError("Native STT backend does not support reset().")
+
+        self._native_session.reset()
+
+    def close(self) -> None:
+        if hasattr(self._native_session, "close"):
+            self._native_session.close()
+
+
+@dataclass
+class SttModel:
+    model_id: str
+    _native_stt: object
+
+    def transcribe(
+        self,
+        *,
+        audio: bytes | Sequence[float] | str | Path,
+        sample_rate: Optional[int] = None,
+        language: Optional[str] = None,
+        task: Optional[str] = None,
+        hotwords: Optional[str] = None,
+    ) -> TranscriptionResult:
+        if isinstance(audio, (str, Path)):
+            samples, resolved_sample_rate = _load_wav_audio(Path(audio))
+        else:
+            if sample_rate is None or sample_rate <= 0:
+                raise ValueError("sample_rate is required for in-memory audio.")
+            samples = _coerce_audio_samples(audio)
+            resolved_sample_rate = int(sample_rate)
+
+        if not hasattr(self._native_stt, "transcribe_result"):
+            raise RuntimeError("Native STT backend does not support transcribe_result().")
+
+        return self._native_stt.transcribe_result(
+            model_id=self.model_id,
+            samples=samples,
+            sample_rate=resolved_sample_rate,
+            language=language,
+            task=task,
+            hotwords=hotwords,
+        )
+
+    def create_session(self) -> SttSession:
+        if not hasattr(self._native_stt, "create_session"):
+            raise RuntimeError("Native STT backend does not support create_session().")
+
+        return SttSession(
+            model_id=self.model_id,
+            _native_session=self._native_stt.create_session(),
+        )