ad58f802f3
Co-authored-by: Davids048 <jundasu@ucsd.edu> Co-authored-by: Matthew Noto <99706358+RandNMR73@users.noreply.github.com>
62 lines
1.8 KiB
Python
62 lines
1.8 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Audio preprocessing helpers for LTX-2 training."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
import torchaudio
|
|
from torch import nn
|
|
|
|
|
|
class AudioProcessor(nn.Module):
|
|
"""Converts audio waveforms to log-mel spectrograms with resampling."""
|
|
|
|
def __init__(
|
|
self,
|
|
sample_rate: int,
|
|
mel_bins: int,
|
|
mel_hop_length: int,
|
|
n_fft: int,
|
|
) -> None:
|
|
super().__init__()
|
|
self.sample_rate = sample_rate
|
|
self.mel_transform = torchaudio.transforms.MelSpectrogram(
|
|
sample_rate=sample_rate,
|
|
n_fft=n_fft,
|
|
win_length=n_fft,
|
|
hop_length=mel_hop_length,
|
|
f_min=0.0,
|
|
f_max=sample_rate / 2.0,
|
|
n_mels=mel_bins,
|
|
window_fn=torch.hann_window,
|
|
center=True,
|
|
pad_mode="reflect",
|
|
power=1.0,
|
|
mel_scale="slaney",
|
|
norm="slaney",
|
|
)
|
|
|
|
def resample_waveform(
|
|
self,
|
|
waveform: torch.Tensor,
|
|
source_rate: int,
|
|
target_rate: int,
|
|
) -> torch.Tensor:
|
|
if source_rate == target_rate:
|
|
return waveform
|
|
resampled = torchaudio.functional.resample(
|
|
waveform, source_rate, target_rate)
|
|
return resampled.to(device=waveform.device, dtype=waveform.dtype)
|
|
|
|
def waveform_to_mel(
|
|
self,
|
|
waveform: torch.Tensor,
|
|
waveform_sample_rate: int,
|
|
) -> torch.Tensor:
|
|
waveform = self.resample_waveform(
|
|
waveform, waveform_sample_rate, self.sample_rate)
|
|
mel = self.mel_transform(waveform)
|
|
mel = torch.log(torch.clamp(mel, min=1e-5))
|
|
mel = mel.to(device=waveform.device, dtype=waveform.dtype)
|
|
return mel.permute(0, 1, 3, 2).contiguous()
|