-
Notifications
You must be signed in to change notification settings - Fork 0
/
inference.py
35 lines (29 loc) · 1.01 KB
/
inference.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
import torch
import numpy as np
import wandb
from src.wavenet import WaveNet
from config import ModelConfig, MelSpectrogramConfig
from src.preprocessing import MelSpectrogram
import torchaudio
import warnings
warnings.filterwarnings('ignore')
import sys
model_config = ModelConfig()
device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")
model_config = ModelConfig()
model = WaveNet()
model.load_state_dict(torch.load(model_config.model_path, map_location=device))
model.to(device)
wav, sr = torchaudio.load("LJ045-0097.wav")
wav = wav.to(device)
featurizer = MelSpectrogram(MelSpectrogramConfig()).to(device)
model.eval()
with torch.no_grad():
mel = featurizer(wav)
y_pred = model.inference(mel).squeeze()
if model_config.wandb_log:
wandb.init(project="wavenet")
wandb.log({
"inference audio": [wandb.Audio(y_pred.cpu(), sample_rate=22050)],
"GT audio": [wandb.Audio(wav.squeeze().cpu(), sample_rate=22050)],
})