# -*- coding: utf-8 -*- """SoftVCTraceAndOnnx.ipynb Automatically generated by Colaboratory. Original file is located at https://colab.research.google.com/drive/1g0YsQQY8fzXHmZUWAMzQO60c6OUlyKBl """ !pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu116 !pip3 install sklearn !pip3 install soundfile !pip3 install matplotlib !git clone https://github.com/bshall/hubert !git clone https://github.com/bshall/acoustic-model !git clone https://github.com/bshall/hifigan !nvidia-smi from google.colab import drive drive.mount('/content/drive') # Torch Script Inference Hubert import copy from typing import Optional, Tuple import random from sklearn.cluster import KMeans import torch import torch.nn as nn import torch.nn.functional as F from torch.nn.modules.utils import consume_prefix_in_state_dict_if_present import torch.nn.utils as utils from torch.nn.utils.weight_norm import WeightNorm from typing import Dict, Any from collections import OrderedDict def remove_weight_norm(module): module_list = [mod for mod in module.children()] if len(module_list) == 0: for k, hook in module._forward_pre_hooks.items(): if isinstance(hook, WeightNorm): hook.remove(module) del module._forward_pre_hooks[k] else: for mod in module_list: remove_weight_norm(mod) def modify_model(state_dict:Dict[str,Any]): new_state_dict = OrderedDict() for k in state_dict['acoustic-model'].keys(): value = state_dict['acoustic-model'][k] name = k[7:] # remove `module.` because they used ddp training to save new_state_dict[name] = value return new_state_dict # @torch.jit.export # def generate(self, x: torch.Tensor) -> torch.Tensor: # x = self.encoder(x) # return self.decoder.generate(x) class AcousticModelScripted(nn.Module): def __init__(self, discrete: bool = False, upsample: bool = True): super().__init__() self.encoder = Encoder(discrete, upsample) self.decoder = Decoder() def forward(self, x: torch.Tensor, mels: torch.Tensor) -> torch.Tensor: x = self.encoder(x) return self.decoder(x, mels) @torch.jit.export def generate(self, x: torch.Tensor) -> torch.Tensor: x = self.encoder(x) return self.decoder.generate(x) class Encoder(nn.Module): def __init__(self, discrete: bool = False, upsample: bool = True): super().__init__() self.embedding = nn.Embedding(100 + 1, 256) if discrete else None self.prenet = PreNet(256, 256, 256) self.convs = nn.Sequential( nn.Conv1d(256, 512, 5, 1, 2), nn.ReLU(), nn.InstanceNorm1d(512), nn.ConvTranspose1d(512, 512, 4, 2, 1) if upsample else nn.Identity(), nn.Conv1d(512, 512, 5, 1, 2), nn.ReLU(), nn.InstanceNorm1d(512), nn.Conv1d(512, 512, 5, 1, 2), nn.ReLU(), nn.InstanceNorm1d(512), ) def forward(self, x: torch.Tensor) -> torch.Tensor: if self.embedding is not None: x = self.embedding(x) x = self.prenet(x) x = self.convs(x.transpose(1, 2)) return x.transpose(1, 2) class Decoder(nn.Module): def __init__(self): super().__init__() self.prenet = PreNet(128, 256, 256) self.lstm1 = nn.LSTM(512 + 256, 768, batch_first=True) self.lstm2 = nn.LSTM(768, 768, batch_first=True) self.lstm3 = nn.LSTM(768, 768, batch_first=True) self.proj = nn.Linear(768, 128, bias=False) def forward(self, x: torch.Tensor, mels: torch.Tensor) -> torch.Tensor: mels = self.prenet(mels) x, _ = self.lstm1(torch.cat((x, mels), dim=-1)) res = x x, _ = self.lstm2(x) x = res + x res = x x, _ = self.lstm3(x) x = res + x return self.proj(x) @torch.inference_mode() def generate(self, xs: torch.Tensor) -> torch.Tensor: m = torch.zeros(xs.size(0), 128, device=xs.device) h1 = torch.zeros(1, xs.size(0), 768, device=xs.device) c1 = torch.zeros(1, xs.size(0), 768, device=xs.device) h2 = torch.zeros(1, xs.size(0), 768, device=xs.device) c2 = torch.zeros(1, xs.size(0), 768, device=xs.device) h3 = torch.zeros(1, xs.size(0), 768, device=xs.device) c3 = torch.zeros(1, xs.size(0), 768, device=xs.device) mel = [] for x in torch.unbind(xs, dim=1): m = self.prenet(m) x = torch.cat((x, m), dim=1).unsqueeze(1) x1, (h1, c1) = self.lstm1(x, (h1, c1)) x2, (h2, c2) = self.lstm2(x1, (h2, c2)) x = x1 + x2 x3, (h3, c3) = self.lstm3(x, (h3, c3)) x = x + x3 m = self.proj(x).squeeze(1) mel.append(m) return torch.stack(mel, dim=1) class PreNet(nn.Module): def __init__( self, input_size: int, hidden_size: int, output_size: int, dropout: float = 0.5, ): super().__init__() self.net = nn.Sequential( nn.Linear(input_size, hidden_size), nn.ReLU(), nn.Dropout(dropout), nn.Linear(hidden_size, output_size), nn.ReLU(), nn.Dropout(dropout), ) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.net(x) import IPython.display as display import torch, torchaudio import torch.nn as nn import sys import soundfile as sf from time import time import torch.jit sys.path.insert(0, './hubert/hubert') from model import HubertSoft sys.path.insert(0, './acoustic-model/') from acoustic.model import AcousticModel #@title select input wav file_name = "/content/drive/MyDrive/VCModels/Source/brandon1.wav"#@param { type: "string"} source, sr = torchaudio.load(f"{file_name}") source = torchaudio.functional.resample(source, sr, 16000) source = source.unsqueeze(0).cuda() print(f"Wav Shape and Type {source.shape} {source.dtype}") if source.shape[1] == 2: # remove channel source = torch.mean(source, dim=0).unsqueeze(0) print(f"Removing Channel to mono Wav Shape and Type {source.shape} {source.dtype}") display.Audio(source.squeeze().cpu(), rate=16000) import datetime hubert_soft = torch.hub.load("bshall/hubert:main", "hubert_soft").cuda() hubert_soft.eval() script_acoustic = True #@param {type: "boolean"} print(script_acoustic) # Load the acoustic model (either hubert_soft or hubert_discrete) input_acoustic_model_file = "/content/drive/MyDrive/VCModels/Padme/model-24000.pt" #@param {type: "string"} acoustic_dict = torch.load(input_acoustic_model_file) # this used ddp ## REPLACE THIS! <-- ! if script_acoustic == True: acoustic = AcousticModelScripted().cuda() else: acoustic = AcousticModel().cuda() acoustic.eval() new_state_dict = modify_model(state_dict=acoustic_dict) acoustic.load_state_dict(new_state_dict) import importlib.util import sys spec = importlib.util.spec_from_file_location("hifigan.utils", "/content/hifigan/hifigan/utils.py") u = importlib.util.module_from_spec(spec) sys.modules["hifigan.utils"] = u spec.loader.exec_module(u) spec = importlib.util.spec_from_file_location("hifigan.generator", "/content/hifigan/hifigan/generator.py") g = importlib.util.module_from_spec(spec) sys.modules["hifigan.generator"] = g spec.loader.exec_module(g) new_state_dict = modify_model(state_dict=acoustic_dict) acoustic.load_state_dict(new_state_dict) input_hifigan_checkpoint_path = "/content/drive/MyDrive/VCModels/Hifigan/Padme.pt" #@param {type: "string"} if input_hifigan_checkpoint_path: checkpoint = torch.load(input_hifigan_checkpoint_path)["generator"]["model"] hifigan = g.HifiganGenerator() consume_prefix_in_state_dict_if_present(checkpoint, "module.") hifigan.load_state_dict(checkpoint) hifigan.remove_weight_norm() else: hifigan = torch.hub.load("bshall/hifigan:main", "hifigan_hubert_soft") hifigan.cuda() hifigan.eval() #utils.remove_weight_norm(hubert_soft.positional_embedding.conv) # Convert to the target speaker units = None mel = None unit_start_time = 0.0 unit_end_time = 0.0 acoustic_start_time = 0.0 acoustic_end_time = 0.0 hifi_start_time = 0.0 hifi_end_time = 0.0 with torch.inference_mode(): unit_start_time = datetime.datetime.now() units = hubert_soft.units(source) print(f"Units and Type {units.shape} {units.dtype}") unit_end_time = datetime.datetime.now() hubert_soft_script = torch.jit.trace_module(hubert_soft,{"units":source}) torch.jit.save(hubert_soft_script, 'hubert_soft.jit') print(f"Unit Generation Time {(unit_end_time - unit_start_time).total_seconds() * 1000}") print("Hubert Soft Exported") acoustic_start_time = datetime.datetime.now() mel = acoustic.generate(units).transpose(1, 2) print(f"Mel and Type {mel.shape} {mel.dtype}") acoustic_end_time = datetime.datetime.now() print(f"Acoustic Time {(acoustic_end_time - acoustic_start_time).total_seconds() * 1000}") if script_acoustic == True: output_acoustic_scripted = '/content/drive/MyDrive/VCModels/FinishedModels/Padme.recast' #@param {type: "string"} print("Exporting Scripted") acoustic_model_script = torch.jit.script(acoustic) torch.jit.save(acoustic_model_script, output_acoustic_scripted) else: print("Exporting Traced") acoustic_model_script = torch.jit.trace_module(acoustic,{"generate":units}) output_acoustic_traced = '' #@param {type: "string"} torch.jit.save(acoustic_model_script, output_acoustic_traced) print("Acoustic Model Exported") hifi_start_time = datetime.datetime.now() target = hifigan(mel) print(f"Wavform and Type {target.shape} {target.dtype}") hifi_end_time = datetime.datetime.now() print(f"MEL {target.shape}") print(f"Hifi Gan Time {(hifi_end_time - hifi_start_time).total_seconds() * 1000}") hifi_gan_script = torch.jit.trace_module(hifigan,{"forward":mel}) output_hifigan_jit = '/content/drive/MyDrive/VCModels/FinishedModels/Padme-hifigan-jit.pt' #@param {type: "string"} torch.jit.save(hifi_gan_script, output_hifigan_jit) print("Hifi Gan Model Exported") import datetime jit_unit_start_time = 0.0 jit_unit_end_time = 0.0 jit_unit_final = 0.0 jit_acoustic_start_time = 0.0 jit_acoustic_end_time = 0.0 jit_acoustic_final = 0.0 jit_hifi_start_time = 0.0 jit_hifi_end_time = 0.0 jit_hifi_final = 0.0 jit_model = torch.jit.load('hubert_soft.jit').cuda() jit_model.eval() with torch.no_grad(): jit_unit_start_time = datetime.datetime.now() jit_tensor = jit_model.units(source) jit_unit_end_time = datetime.datetime.now() jit_unit_final = jit_unit_end_time - jit_unit_start_time print(jit_unit_final.total_seconds() * 1000) print(jit_tensor.shape) assert torch.allclose(jit_tensor,units) print(script_acoustic) if script_acoustic == True: acoustic_jit_model = torch.jit.load(output_acoustic_scripted).cuda() print("Using Scripted Model") else: acoustic_jit_model = torch.jit.load(output_acoustic_traced).cuda() acoustic_jit_model.eval() with torch.no_grad(): if script_acoustic == True: print("Using Scripted Model") jit_acoustic_start_time = datetime.datetime.now() jit_mel = acoustic_jit_model.generate(jit_tensor) else: jit_acoustic_start_time = datetime.datetime.now() jit_mel = acoustic_jit_model.generate(jit_tensor) jit_acoustic_end_time = datetime.datetime.now() jit_acoustic_final = jit_acoustic_end_time - jit_acoustic_start_time print(jit_acoustic_final.total_seconds() * 1000) jit_mel = jit_mel.transpose(1,2) print(jit_mel.shape) assert torch.allclose(mel,jit_mel) hifi_gan = torch.jit.load(output_hifigan_jit).cuda() hifi_gan.eval() with torch.no_grad(): jit_hifi_start_time = datetime.datetime.now() jit_target = hifi_gan(jit_mel) jit_hifi_end_time = datetime.datetime.now() jit_hifi_final = jit_hifi_end_time - jit_hifi_start_time print(jit_hifi_final.total_seconds() * 1000) assert torch.allclose(target,jit_target) sf.write(f'./sample_jit.wav', jit_target.squeeze().detach().cpu().numpy(), samplerate=16000) sf.write(f'./sample.wav', target.squeeze().cpu().detach().numpy(), samplerate=16000) #uploaded = files.upload() import datetime jit_unit_start_time = 0.0 jit_unit_end_time = 0.0 jit_unit_final = 0.0 jit_acoustic_start_time = 0.0 jit_acoustic_end_time = 0.0 jit_acoustic_final = 0.0 jit_hifi_start_time = 0.0 jit_hifi_end_time = 0.0 jit_hifi_final = 0.0 source, sr = torchaudio.load(f"{file_name}") source = torchaudio.functional.resample(source, sr, 16000) source = source.unsqueeze(0).cuda() print(f"Wav Shape and Type {source.shape} {source.dtype}") if source.shape[1] == 2: # remove channel source = torch.mean(source, dim=1).unsqueeze(0) print(f"Removing Channel to mono Wav Shape and Type {source.shape} {source.dtype}") display.Audio(source.squeeze().cpu(), rate=16000) jit_model = torch.jit.load('hubert_soft.jit').cuda() jit_model.eval() if script_acoustic == True: acoustic_jit_model = torch.jit.load(output_acoustic_scripted).cuda() print("Loading Scripted Model") else: acoustic_jit_model = torch.jit.load(output_acoustic_traced).cuda() print("Loading Tracing Model") acoustic_jit_model.eval() hifi_gan = torch.jit.load(output_hifigan_jit).cuda() hifi_gan.eval() print("Loaded Models") # Warm up the model to get the real time for i in range(0,5): with torch.no_grad(): jit_unit_start_time = datetime.datetime.now() jit_tensor = jit_model.units(source) jit_unit_end_time = datetime.datetime.now() if script_acoustic == True: print("Using Scripted Model TT2") jit_acoustic_start_time = datetime.datetime.now() jit_mel = acoustic_jit_model.generate(jit_tensor) else: print("Using Traced Model TT2") jit_acoustic_start_time = datetime.datetime.now() jit_mel = acoustic_jit_model.generate(jit_tensor) jit_mel = jit_mel.transpose(1, 2) jit_acoustic_end_time = datetime.datetime.now() jit_hifi_start_time = datetime.datetime.now() jit_target = hifi_gan(jit_mel) jit_hifi_end_time = datetime.datetime.now() jit_unit_final = jit_unit_end_time - jit_unit_start_time jit_acoustic_final = jit_acoustic_end_time - jit_acoustic_start_time jit_hifi_final = jit_hifi_end_time - jit_hifi_start_time print(f"16 Khz Wav Shape {source.shape}") print(f"Jit Tensor Shape for Unit Inputs: {jit_tensor.shape}") print(f"Hubert Runtime : {jit_unit_final.total_seconds() * 1000}") print(f"Tactron 2 Runtime Mel Output Shape: {jit_acoustic_final.total_seconds() * 1000}") print(f"Jit Tensor Shape for Mel output: {jit_mel.shape}") print(f"Hifi Runtime Mel Output Shape: {jit_hifi_final.total_seconds() * 1000}") print(f"Jit Tensor Shape for Mel output: {jit_target.shape}") sf.write(f'./{file_name}_jit.wav', jit_target.squeeze().detach().cpu().numpy(), samplerate=16000) display.Audio(jit_target.squeeze().cpu(), rate=16000) # silero vad import datetime source, sr = torchaudio.load(f"{file_name}") source = torchaudio.functional.resample(source, sr, 16000) print(source.shape) input_vad_file = "/content/drive/MyDrive/tts/softvc-checkpoints/silero_vad.jit" #@param {type: "string"} class VADIterator: # You will only need to verify the 100 ms chunks to see this checks a minute of audio # sound_sample_size run vad in 30 ms, 60 ms or 100 ms def __init__(self,threshold: float = 0.5, sampling_rate: int = 16000,sound_sample_size:float = 100): self.silero_jit = torch.jit.load(input_vad_file).cuda() self.silero_jit.eval() self.sampling_rate = sampling_rate self.sound_sample_size = sound_sample_size jit_silero_start_time = 0.0 jit_silero_end_time = 0.0 jit_silero_final = 0.0 def operating_frame_size(self,sampling_rate:int,ms_of_audio:float): return int((ms_of_audio * sampling_rate) / 1000) # number of frames def check_voice_activity(self,source_sample) -> float: probs_over_time = [] # formula we know that we can get 1600 operating frames / 16000 sr * 1000 = 100 ms of audio from 1600 frames. # this gives you the operating frames given the audio size you want. op_frame_size = self.operating_frame_size(sampling_rate=self.sampling_rate,ms_of_audio=self.sound_sample_size) print(op_frame_size) iters = int(source_sample.shape[-1] / op_frame_size) # get the chunks #print(iters) for i in range(0,iters): # get sound_sample_size ms chunks #print(window_size) chunk = source_sample[i * op_frame_size:(i + 1) * op_frame_size] with torch.no_grad():# slide window across using VAD 100ms jit_silero_start_time = datetime.datetime.now() try: print(f"Current {i}") prob = self.silero_jit(chunk.cuda(),sr=self.sampling_rate).item() # input shape is [1*] just an array of samples. 1 = false jit_silero_end_time = datetime.datetime.now() print(f"Probability: {prob}") jit_silero_final = jit_silero_end_time - jit_silero_start_time #print(f"Silero VAD Time: {jit_silero_final.total_seconds() * 1000}") probs_over_time.append(prob) except Exception as e: print(e) return probs_over_time iterator = VADIterator(sound_sample_size=100) source_sample = source_sample=source.squeeze().squeeze() print(source_sample.shape) probs = iterator.check_voice_activity(source_sample=source_sample) #print(probs)