Copy the files in the new repo for public release

pocket_tts/data/audio.py

@@ -0,0 +1,126 @@
+"""
+Audio IO methods are defined in this module (info, read, write),
+We rely on av library for faster read when possible, otherwise on torchaudio.
+"""
+
+import logging
+import os
+import sys
+import wave
+from contextlib import nullcontext
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+from beartype.typing import Iterator
+
+logger = logging.getLogger(__name__)
+
+FIRST_CHUNK_LENGTH_SECONDS = float(os.environ.get("FIRST_CHUNK_LENGTH_SECONDS", "0"))
+
+
+def audio_read(filepath: str | Path) -> tuple[torch.Tensor, int]:
+    """Read audio using Python's wave module."""
+    with wave.open(str(filepath), "rb") as wav_file:
+        sample_rate = wav_file.getframerate()
+
+        # Read all audio data as 16-bit signed integers
+        raw_data = wav_file.readframes(-1)
+        samples = np.frombuffer(raw_data, dtype=np.int16).astype(np.float32) / 32768.0
+
+        # Return as mono tensor (channels, samples)
+        wav = torch.from_numpy(samples.reshape(1, -1))
+        return wav, sample_rate
+
+
+class StreamingWAVWriter:
+    """WAV writer using Python's standard library wave module."""
+
+    def __init__(self, output_stream, sample_rate: int):
+        self.output_stream = output_stream
+        self.sample_rate = sample_rate
+        self.wave_writer = None
+        self.first_chunk_buffer = []
+
+    def write_header(self, sample_rate: int):
+        """Initialize WAV writer with header."""
+        # For stdout streaming, we need to handle the unseekable stream case
+        # The wave module supports unseekable streams since Python 3.4
+        self.wave_writer = wave.open(self.output_stream, "wb")
+        self.wave_writer.setnchannels(1)  # Mono
+        self.wave_writer.setsampwidth(2)  # 16-bit
+        self.wave_writer.setframerate(sample_rate)
+        self.wave_writer.setnframes(1_000_000_000)
+
+    def write_pcm_data(self, audio_chunk: torch.Tensor):
+        """Write PCM data using wave module."""
+        # Convert to int16 PCM bytes
+        chunk_int16 = (audio_chunk.clamp(-1, 1) * 32767).short()
+        chunk_bytes = chunk_int16.detach().cpu().numpy().tobytes()
+
+        if self.first_chunk_buffer is not None:
+            self.first_chunk_buffer.append(chunk_bytes)
+            total_length = sum(len(c) for c in self.first_chunk_buffer)
+            target_length = (
+                int(self.sample_rate * FIRST_CHUNK_LENGTH_SECONDS) * 2
+            )  # 2 bytes per sample
+            if total_length < target_length:
+                return
+            self._flush()
+            return
+
+        # Use writeframesraw to avoid frame count validation for streaming
+        self.wave_writer.writeframesraw(chunk_bytes)
+
+    def _flush(self):
+        if self.first_chunk_buffer is not None:
+            self.wave_writer.writeframesraw(b"".join(self.first_chunk_buffer))
+            self.first_chunk_buffer = None
+
+    def finalize(self):
+        """Close the wave writer."""
+        self._flush()
+
+        # Let's add 200ms of silence to ensure proper playback
+        silence_duration_sec = 0.2
+        num_silence_samples = int(self.sample_rate * silence_duration_sec)
+
+        self.wave_writer.writeframesraw(bytes(num_silence_samples * 2))
+
+        if self.wave_writer:
+            # do not update the header for unseekable streams
+            self.wave_writer._patchheader = lambda: None
+            self.wave_writer.close()
+
+
+def is_file_like(obj):
+    """Check if object has basic file-like methods."""
+    return all(hasattr(obj, attr) for attr in ["write", "close"])
+
+
+def stream_audio_chunks(
+    path: str | Path | None | Any, audio_chunks: Iterator[torch.Tensor], sample_rate: int
+):
+    """Stream audio chunks to a WAV file or stdout, optionally playing them."""
+    if path == "-":
+        f = sys.stdout.buffer
+    elif path is None:
+        f = nullcontext()
+    elif is_file_like(path):
+        f = path
+    else:
+        f = open(path, "wb")
+
+    with f:
+        if path is not None:
+            writer = StreamingWAVWriter(f, sample_rate)
+            writer.write_header(sample_rate)
+
+        for chunk in audio_chunks:
+            # Then write to file
+            if path is not None:
+                writer.write_pcm_data(chunk)
+
+        if path is not None:
+            writer.finalize()