| from __future__ import annotations |
|
|
| """ |
| IndexTTS2 local audiotools compatibility layer. |
| |
| This file provides a minimal API subset used by IndexTTS2 and is based on |
| behavior from the upstream Descript projects: |
| - https://github.com/descriptinc/audiotools (tag: 0.7.4) |
| - https://github.com/descriptinc/descript-audio-codec |
| |
| Upstream package metadata credits authors Prem Seetharaman and Lucas Gestin. |
| Upstream audiotools is licensed under MIT (see upstream LICENSE file). |
| """ |
|
|
| from dataclasses import dataclass |
| from pathlib import Path |
| from types import SimpleNamespace |
| from typing import Optional, Union |
|
|
| import torch |
| import torch.nn.functional as F |
| import torchaudio |
| from torch import nn |
|
|
|
|
| _AUDIO_EXTENSIONS = { |
| ".wav", |
| ".flac", |
| ".mp3", |
| ".ogg", |
| ".m4a", |
| ".aac", |
| ".wma", |
| ".opus", |
| } |
|
|
|
|
| def find_audio(path: Union[str, Path]) -> list[Path]: |
| input_path = Path(path) |
| if input_path.is_file(): |
| return [input_path] if input_path.suffix.lower() in _AUDIO_EXTENSIONS else [] |
| if not input_path.exists(): |
| return [] |
| audio_files = [p for p in input_path.rglob("*") if p.is_file() and p.suffix.lower() in _AUDIO_EXTENSIONS] |
| return sorted(audio_files) |
|
|
|
|
| @dataclass |
| class STFTParams: |
| window_length: int |
| hop_length: int |
| match_stride: bool = False |
| window_type: Optional[str] = None |
|
|
|
|
| def _get_window(window_type: Optional[str], window_length: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor: |
| name = (window_type or "hann").lower() |
| if name in {"hann", "hann_window"}: |
| return torch.hann_window(window_length, device=device, dtype=dtype) |
| if name in {"hamming", "hamming_window"}: |
| return torch.hamming_window(window_length, device=device, dtype=dtype) |
| return torch.hann_window(window_length, device=device, dtype=dtype) |
|
|
|
|
| class AudioSignal: |
| def __init__( |
| self, |
| data: Union[str, Path, torch.Tensor], |
| sample_rate: Optional[int] = None, |
| stft_params: Optional[STFTParams] = None, |
| ): |
| if isinstance(data, (str, Path)): |
| waveform, sr = torchaudio.load(str(data)) |
| self.audio_data = waveform.unsqueeze(0) |
| self.sample_rate = int(sr) |
| else: |
| tensor = torch.as_tensor(data) |
| if tensor.ndim == 1: |
| tensor = tensor.unsqueeze(0).unsqueeze(0) |
| elif tensor.ndim == 2: |
| tensor = tensor.unsqueeze(0) |
| elif tensor.ndim != 3: |
| raise ValueError(f"Expected 1D/2D/3D audio tensor, got shape {tuple(tensor.shape)}") |
| if sample_rate is None: |
| raise ValueError("sample_rate is required when constructing AudioSignal from tensors") |
| self.audio_data = tensor |
| self.sample_rate = int(sample_rate) |
|
|
| if not torch.is_floating_point(self.audio_data): |
| self.audio_data = self.audio_data.float() |
|
|
| default_hop = max(1, int(self.sample_rate * 0.01)) |
| default_win = max(16, default_hop * 4) |
| self.stft_params = stft_params or STFTParams(window_length=default_win, hop_length=default_hop) |
| self.magnitude = None |
|
|
| @property |
| def device(self) -> torch.device: |
| return self.audio_data.device |
|
|
| @property |
| def shape(self): |
| return self.audio_data.shape |
|
|
| @property |
| def signal_length(self) -> int: |
| return int(self.audio_data.shape[-1]) |
|
|
| @property |
| def signal_duration(self) -> float: |
| return float(self.signal_length) / float(self.sample_rate) |
|
|
| def clone(self) -> "AudioSignal": |
| return AudioSignal(self.audio_data.clone(), self.sample_rate, stft_params=self.stft_params) |
|
|
| def __getitem__(self, item) -> "AudioSignal": |
| tensor = self.audio_data.__getitem__(item) |
| if tensor.ndim == 1: |
| tensor = tensor.unsqueeze(0).unsqueeze(0) |
| elif tensor.ndim == 2: |
| tensor = tensor.unsqueeze(1) |
| elif tensor.ndim != 3: |
| raise ValueError(f"Unsupported indexed shape {tuple(tensor.shape)}") |
| return AudioSignal(tensor, self.sample_rate, stft_params=self.stft_params) |
|
|
| def to(self, device: Union[str, torch.device]) -> "AudioSignal": |
| self.audio_data = self.audio_data.to(device) |
| return self |
|
|
| def zero_pad(self, left: int, right: int) -> "AudioSignal": |
| padded = F.pad(self.audio_data, (int(left), int(right))) |
| return AudioSignal(padded, self.sample_rate, stft_params=self.stft_params) |
|
|
| def resample(self, target_sample_rate: int) -> "AudioSignal": |
| target_sample_rate = int(target_sample_rate) |
| if target_sample_rate == self.sample_rate: |
| return self |
| batch, channels, samples = self.audio_data.shape |
| flat = self.audio_data.reshape(batch * channels, samples) |
| resampled = torchaudio.functional.resample(flat, self.sample_rate, target_sample_rate) |
| self.audio_data = resampled.reshape(batch, channels, -1) |
| self.sample_rate = target_sample_rate |
| self.magnitude = None |
| return self |
|
|
| def ffmpeg_resample(self, target_sample_rate: int) -> "AudioSignal": |
| return self.resample(target_sample_rate) |
|
|
| def loudness(self) -> torch.Tensor: |
| rms = torch.sqrt(torch.mean(self.audio_data.float() ** 2) + 1e-12) |
| db = 20.0 * torch.log10(rms.clamp_min(1e-7)) |
| return db.detach().cpu() |
|
|
| def ffmpeg_loudness(self) -> torch.Tensor: |
| return self.loudness() |
|
|
| def normalize(self, target_db: Union[float, torch.Tensor]) -> "AudioSignal": |
| target = float(torch.as_tensor(target_db).detach().cpu().item()) |
| current = float(self.loudness().item()) |
| gain = 10 ** ((target - current) / 20.0) |
| self.audio_data = self.audio_data * gain |
| self.magnitude = None |
| return self |
|
|
| def ensure_max_of_audio(self, max_value: float = 0.99) -> "AudioSignal": |
| peak = self.audio_data.abs().amax() |
| if torch.isfinite(peak) and peak.item() > max_value: |
| self.audio_data = self.audio_data * (max_value / peak) |
| return self |
|
|
| def stft( |
| self, |
| window_length: Optional[int] = None, |
| hop_length: Optional[int] = None, |
| window_type: Optional[str] = None, |
| ) -> torch.Tensor: |
| params = self.stft_params |
| win_length = int(window_length or params.window_length) |
| hop = int(hop_length or params.hop_length) |
| window = _get_window(window_type or params.window_type, win_length, self.device, self.audio_data.dtype) |
| batch, channels, samples = self.audio_data.shape |
| flat = self.audio_data.reshape(batch * channels, samples) |
| spec = torch.stft( |
| flat, |
| n_fft=win_length, |
| hop_length=hop, |
| win_length=win_length, |
| window=window, |
| center=True, |
| return_complex=True, |
| ) |
| spec = spec.reshape(batch, channels, spec.shape[-2], spec.shape[-1]) |
| self.magnitude = spec.abs() |
| return spec |
|
|
| def mel_spectrogram( |
| self, |
| n_mels: int, |
| mel_fmin: float = 0.0, |
| mel_fmax: Optional[float] = None, |
| window_length: Optional[int] = None, |
| hop_length: Optional[int] = None, |
| window_type: Optional[str] = None, |
| ) -> torch.Tensor: |
| params = self.stft_params |
| win_length = int(window_length or params.window_length) |
| hop = int(hop_length or params.hop_length) |
| window_fn = torch.hann_window |
| if (window_type or params.window_type or "").lower() in {"hamming", "hamming_window"}: |
| window_fn = torch.hamming_window |
| mel = torchaudio.transforms.MelSpectrogram( |
| sample_rate=self.sample_rate, |
| n_fft=win_length, |
| win_length=win_length, |
| hop_length=hop, |
| f_min=float(mel_fmin), |
| f_max=None if mel_fmax is None else float(mel_fmax), |
| n_mels=int(n_mels), |
| power=2.0, |
| center=True, |
| window_fn=window_fn, |
| ).to(self.device) |
| batch, channels, samples = self.audio_data.shape |
| flat = self.audio_data.reshape(batch * channels, samples) |
| mels = mel(flat) |
| return mels.reshape(batch, channels, mels.shape[-2], mels.shape[-1]) |
|
|
| def write(self, path: Union[str, Path]) -> Path: |
| out_path = Path(path) |
| out_path.parent.mkdir(parents=True, exist_ok=True) |
| waveform = self.audio_data[0].detach().cpu() |
| torchaudio.save(str(out_path), waveform, self.sample_rate) |
| return out_path |
|
|
| @classmethod |
| def load_from_file_with_ffmpeg(cls, path: Union[str, Path]) -> "AudioSignal": |
| return cls(path) |
|
|
|
|
| class BaseModel(nn.Module): |
| INTERN = [] |
| EXTERN = [] |
|
|
| @property |
| def device(self) -> torch.device: |
| try: |
| return next(self.parameters()).device |
| except StopIteration: |
| return torch.device("cpu") |
|
|
| @classmethod |
| def _extract_state_dict(cls, checkpoint): |
| if isinstance(checkpoint, dict): |
| for key in ("state_dict", "model", "generator", "weights"): |
| value = checkpoint.get(key) |
| if isinstance(value, dict): |
| return value |
| if checkpoint and all(torch.is_tensor(v) for v in checkpoint.values()): |
| return checkpoint |
| raise RuntimeError("Unsupported checkpoint format for BaseModel.load") |
|
|
| @classmethod |
| def _clean_state_dict(cls, state_dict): |
| cleaned = {} |
| for key, value in state_dict.items(): |
| if key.startswith("module."): |
| key = key[len("module.") :] |
| cleaned[key] = value |
| return cleaned |
|
|
| @classmethod |
| def load(cls, path: Union[str, Path]): |
| checkpoint = torch.load(path, map_location="cpu") |
| state_dict = cls._clean_state_dict(cls._extract_state_dict(checkpoint)) |
| model = cls() |
| model.load_state_dict(state_dict, strict=False) |
| model.eval() |
| return model |
|
|
|
|
| class Accelerator: |
| def __init__(self, *args, **kwargs): |
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
| def prepare(self, *objects): |
| if len(objects) == 1: |
| return objects[0] |
| return objects |
|
|
| @staticmethod |
| def unwrap_model(model): |
| return model |
|
|
|
|
| ml = SimpleNamespace(BaseModel=BaseModel, Accelerator=Accelerator) |
|
|
|
|
| __all__ = [ |
| "AudioSignal", |
| "STFTParams", |
| "BaseModel", |
| "Accelerator", |
| "ml", |
| "find_audio", |
| ] |
|
|
