Spaces:
Paused
Paused
| """Audio loading and resampling utilities. WAV is handled by the stdlib; | |
| other formats fall back to `soundfile`.""" | |
| import wave | |
| from pathlib import Path | |
| from typing import IO | |
| import numpy as np | |
| import torch | |
| from muscriptor.utils.resample import resample_frac | |
| def _read_wav_file(source) -> tuple[torch.Tensor, int]: | |
| """Load a PCM WAV file using the stdlib `wave` module. | |
| `source` may be a filesystem path or a binary file-like object. | |
| Returns: | |
| (wav, sr) where wav has shape [C, T] and is float32 in [-1, 1]. | |
| """ | |
| if hasattr(source, "read"): | |
| opened = wave.open(source, "rb") | |
| else: | |
| opened = wave.open(str(source), "rb") | |
| with opened as wf: | |
| n_channels = wf.getnchannels() | |
| sr = wf.getframerate() | |
| sampwidth = wf.getsampwidth() | |
| n_frames = wf.getnframes() | |
| raw = wf.readframes(n_frames) | |
| if sampwidth == 1: | |
| data = np.frombuffer(raw, dtype=np.uint8).astype(np.float32) | |
| data = (data - 128.0) / 128.0 | |
| elif sampwidth == 2: | |
| data = np.frombuffer(raw, dtype=np.int16).astype(np.float32) / 32768.0 | |
| elif sampwidth == 3: | |
| bytes_ = np.frombuffer(raw, dtype=np.uint8).reshape(-1, 3) | |
| as_int32 = ( | |
| bytes_[:, 0].astype(np.int32) | |
| | (bytes_[:, 1].astype(np.int32) << 8) | |
| | (bytes_[:, 2].astype(np.int32) << 16) | |
| ) | |
| as_int32 = np.where(as_int32 >= (1 << 23), as_int32 - (1 << 24), as_int32) | |
| data = as_int32.astype(np.float32) / float(1 << 23) | |
| elif sampwidth == 4: | |
| data = np.frombuffer(raw, dtype=np.int32).astype(np.float32) / float(1 << 31) | |
| else: | |
| raise ValueError(f"Unsupported WAV sample width: {sampwidth} bytes") | |
| data = data.reshape(-1, n_channels) | |
| return torch.from_numpy(np.ascontiguousarray(data.T)), sr | |
| def _read_non_wav_file(source: str | Path | IO[bytes]) -> tuple[torch.Tensor, int]: | |
| """Load a non-WAV audio file using `soundfile`. | |
| `source` may be a filesystem path or a binary file-like object (e.g. an | |
| ``io.BytesIO`` of an uploaded file), since libsndfile reads either. | |
| Returns: | |
| (wav, sr) where wav has shape [C, T] and is float32 in [-1, 1]. | |
| """ | |
| try: | |
| import soundfile as sf | |
| except ImportError as e: | |
| raise ImportError( | |
| "soundfile is required to read non-WAV audio files. " | |
| "Install with: `pip install soundfile` or `uvx --with soundfile`" | |
| ) from e | |
| target = str(source) if isinstance(source, (str, Path)) else source | |
| data, sample_rate = sf.read(target, dtype="float32") | |
| if data.ndim == 1: | |
| data = data[:, None] | |
| wav = torch.from_numpy(np.ascontiguousarray(data.T)) | |
| return wav, sample_rate | |
| def resample( | |
| waveform: torch.Tensor, | |
| orig_freq: int, | |
| new_freq: int, | |
| ) -> torch.Tensor: | |
| """Sinc resampler via julius `resample_frac`. Operates along the last dim.""" | |
| if orig_freq == new_freq: | |
| return waveform | |
| return resample_frac(waveform, int(orig_freq), int(new_freq)) | |
| def load_audio(path: str | Path, target_sr: int = 16000) -> torch.Tensor: | |
| """Load an audio file and return a mono float32 tensor at target_sr. | |
| PCM WAV files are read with the stdlib `wave` module. Other formats (mp3, | |
| flac, ogg, m4a, …) are decoded via `soundfile`. Dispatch is by content, not | |
| file extension, so misnamed files (e.g. an MP3 upload saved as .wav) still | |
| load. | |
| Returns: | |
| Tensor of shape [1, T] at target_sr. | |
| """ | |
| filepath = Path(path) | |
| try: | |
| wav, sr = _read_wav_file(str(filepath)) | |
| except (wave.Error, EOFError): | |
| wav, sr = _read_non_wav_file(str(filepath)) | |
| if wav.shape[0] > 1: | |
| wav = wav.mean(dim=0, keepdim=True) | |
| if sr != target_sr: | |
| wav = resample(wav, sr, target_sr) | |
| return wav | |