| from tinygrad import nn | |
| class TinySileroVAD: | |
| def __init__(self): | |
| """ | |
| from tinygrad.nn.state import safe_load, load_state_dict | |
| tiny_model = TinySileroVAD() | |
| state_dict = safe_load('data/silero_vad_16k.safetensors') | |
| load_state_dict(tiny_model, state_dict) | |
| """ | |
| self.n_fft = 256 | |
| self.stride = 128 | |
| self.pad = 64 | |
| self.cutoff = int(self.n_fft // 2) + 1 | |
| self.stft_conv = nn.Conv1d(1, 258, kernel_size=256, stride=self.stride, padding=0, bias=False) | |
| self.conv1 = nn.Conv1d(129, 128, kernel_size=3, stride=1, padding=1) | |
| self.conv2 = nn.Conv1d(128, 64, kernel_size=3, stride=2, padding=1) | |
| self.conv3 = nn.Conv1d(64, 64, kernel_size=3, stride=2, padding=1) | |
| self.conv4 = nn.Conv1d(64, 128, kernel_size=3, stride=1, padding=1) | |
| self.lstm_cell = nn.LSTMCell(128, 128) | |
| self.final_conv = nn.Conv1d(128, 1, 1) | |
| def __call__(self, x, state=None): | |
| """ | |
| # full audio example: | |
| import torch | |
| from tinygrad import Tensor | |
| wav = read_audio(audio_path, sampling_rate=16000).unsqueeze(0) | |
| num_samples = 512 | |
| context_size = 64 | |
| context = Tensor(np.zeros((1, context_size))).float() | |
| outs = [] | |
| state = None | |
| if wav.shape[1] % num_samples: | |
| pad_num = num_samples - (wav.shape[1] % num_samples) | |
| wav = torch.nn.functional.pad(wav, (0, pad_num), 'constant', value=0.0) | |
| wav = torch.nn.functional.pad(wav, (context_size, 0)) | |
| wav = Tensor(wav.numpy()).float() | |
| for i in tqdm(range(context_size, wav.shape[1], num_samples)): | |
| wavs_batch = wav[:, i-context_size:i+num_samples] | |
| out_chunk, state = tiny_model(wavs_batch, state) | |
| #outs.append(out_chunk.numpy()) | |
| outs.append(out_chunk) | |
| predict = outs[0].cat(*outs[1:], dim=1).numpy() | |
| """ | |
| if state is not None: | |
| state = (state[0], state[1]) | |
| x = x.pad((0, self.pad), "reflect").unsqueeze(1) | |
| x = self.stft_conv(x) | |
| x = (x[:, :self.cutoff, :]**2 + x[:, self.cutoff:, :]**2).sqrt() | |
| x = self.conv1(x).relu() | |
| x = self.conv2(x).relu() | |
| x = self.conv3(x).relu() | |
| x = self.conv4(x).relu().squeeze(-1) | |
| h, c = self.lstm_cell(x, state) | |
| x = h.unsqueeze(-1) | |
| state = h.stack(c, dim=0) | |
| x = x.relu() | |
| x = self.final_conv(x).sigmoid() | |
| x = x.squeeze(1).mean(axis=1).unsqueeze(1) | |
| return x, state | |