Fabio
c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
81 lines
2 KiB
Python
81 lines
2 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import argparse
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import librosa
|
|
import numpy as np
|
|
import torch
|
|
import torchaudio
|
|
|
|
|
|
def summarize(values: torch.Tensor) -> dict:
|
|
flat = values.detach().cpu().float().contiguous().view(-1)
|
|
return {
|
|
"frames": int(values.shape[0]),
|
|
"dims": int(values.shape[1]),
|
|
"sum": float(flat.double().sum()),
|
|
"first": flat[:32].tolist(),
|
|
}
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("reference")
|
|
args = parser.parse_args()
|
|
|
|
waveform, sample_rate = torchaudio.load(args.reference)
|
|
waveform = waveform.mean(dim=0, keepdim=True)
|
|
|
|
prompt = torchaudio.functional.resample(
|
|
waveform, sample_rate, 24000
|
|
).squeeze(0)
|
|
window = torch.hann_window(1920, periodic=True)
|
|
padded = torch.nn.functional.pad(
|
|
prompt.unsqueeze(0), (720, 720), mode="reflect"
|
|
).squeeze(0)
|
|
stft = torch.stft(
|
|
padded,
|
|
n_fft=1920,
|
|
hop_length=480,
|
|
win_length=1920,
|
|
window=window,
|
|
center=False,
|
|
return_complex=True,
|
|
)
|
|
magnitude = torch.sqrt(stft.real.square() + stft.imag.square() + 1e-9)
|
|
mel_basis = torch.from_numpy(
|
|
librosa.filters.mel(
|
|
sr=24000,
|
|
n_fft=1920,
|
|
n_mels=80,
|
|
fmin=0,
|
|
fmax=8000,
|
|
htk=False,
|
|
norm="slaney",
|
|
)
|
|
).float()
|
|
mel = torch.log(torch.clamp(mel_basis @ magnitude, min=1e-5)).transpose(0, 1)
|
|
|
|
campplus = torchaudio.functional.resample(
|
|
waveform, sample_rate, 16000
|
|
)
|
|
fbank = torchaudio.compliance.kaldi.fbank(
|
|
campplus,
|
|
num_mel_bins=80,
|
|
dither=0.0,
|
|
sample_frequency=16000,
|
|
window_type="povey",
|
|
frame_length=25.0,
|
|
frame_shift=10.0,
|
|
snip_edges=True,
|
|
)
|
|
fbank = fbank - fbank.mean(dim=0, keepdim=True)
|
|
|
|
print(json.dumps({"mel": summarize(mel), "fbank": summarize(fbank)}))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|