mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-13 23:53:02 +00:00
Paul Sajna
f86b44d8fc
* 48KHz softvc * add some notes * almost forgot the inference script * fix readme link * requirements.txt
122 lines
4.0 KiB
Python
122 lines
4.0 KiB
Python
#@title Load the models
|
|
#@markdown leave hifigan blank to use defaults
|
|
import sys
|
|
import torch, torchaudio
|
|
from typing import Dict, Any
|
|
from collections import OrderedDict
|
|
from torch.nn.modules.utils import consume_prefix_in_state_dict_if_present
|
|
|
|
# Inputs
|
|
check_point_path = './acoustic-model/exp/glados/model-35000.pt' #@param {type:"string"}
|
|
hifigan_checkpoint_path = "./hifigan/exp/vctk/model-156000.pt" #@param {type: "string"}
|
|
file_name = "/mnt/8terra/tts/voice-conversion-sources/ieatmeat.wav" #@param {type: "string"}
|
|
|
|
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
|
|
|
|
def repeat_expand_2d(content: torch.Tensor, target_len: int) -> torch.Tensor:
|
|
# content : [h, t]
|
|
src_len = content.shape[1]
|
|
if target_len < src_len:
|
|
return content[:, target_len]
|
|
else:
|
|
return torch.nn.functional.interpolate(
|
|
content.unsqueeze(0), size=target_len, mode="nearest"
|
|
)
|
|
|
|
|
|
sys.path.append("./acoustic-model/")
|
|
from acoustic import AcousticModel
|
|
|
|
hubert = torch.hub.load("bshall/hubert:main", "hubert_soft").cuda()
|
|
|
|
|
|
# Load the acoustic model (either hubert_soft or hubert_discrete)
|
|
acoustic_dict = torch.load(check_point_path) # this used ddp
|
|
acoustic = AcousticModel().cuda()
|
|
acoustic.eval()
|
|
|
|
import importlib.util
|
|
import sys
|
|
|
|
spec = importlib.util.spec_from_file_location("hifigan.utils", "./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", "./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)
|
|
if hifigan_checkpoint_path:
|
|
checkpoint = torch.load(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()
|
|
hifigan.eval()
|
|
else:
|
|
hifigan = torch.hub.load("bshall/hifigan:main", "hifigan_hubert_soft")
|
|
hifigan.cuda()
|
|
|
|
#@title This code selects the source file
|
|
#import IPython.display as display
|
|
sys.path.insert(0, './acoustic-model/acoustic/')
|
|
from model import AcousticModel
|
|
import soundfile as sf
|
|
#from google.colab import files
|
|
from time import time
|
|
sys.path.append("./hifigan/")
|
|
from hifigan.generator import HifiganGenerator
|
|
from hifigan.generator import consume_prefix_in_state_dict_if_present
|
|
#uploaded = files.upload()
|
|
|
|
#for name in uploaded.keys():
|
|
#file_name = name
|
|
source, sr = torchaudio.load(f"{file_name}")
|
|
source = torchaudio.functional.resample(source, sr, 16000)
|
|
source = torch.mean(source, dim=0).unsqueeze(0)
|
|
source = source.unsqueeze(0).cuda()
|
|
#display.Audio(source.squeeze().cpu(), rate=16000)
|
|
|
|
#@title Run the inference and display the result
|
|
# Convert to the target speaker
|
|
with torch.inference_mode():
|
|
# Extract speech units
|
|
start = time()
|
|
units = hubert.units(source)
|
|
print(units.shape)
|
|
#units = repeat_expand_2d(units, target_len)
|
|
units = torch.nn.functional.interpolate(
|
|
units.transpose(1,2), size=units.shape[1]*3, mode="nearest"
|
|
).transpose(1,2)
|
|
|
|
print(units.shape)
|
|
#units = units.unsqueeze(0)
|
|
#units = units.transpose(1,2)
|
|
|
|
end = time()
|
|
# Generate target spectrogram
|
|
print(f"time for feature embedding:{end-start}")
|
|
|
|
start = time()
|
|
mel = acoustic.generate(units).transpose(1,2)
|
|
end = time()
|
|
print(f"time for inference:{end-start}")
|
|
# Generate audio waveform
|
|
start = time()
|
|
target = hifigan(mel)
|
|
end = time()
|
|
print(f"time for wavform {end-start}")
|
|
print(target.squeeze().cpu().unsqueeze(0).shape)
|
|
|
|
torchaudio.save("output.wav", target.squeeze().cpu().unsqueeze(0), 48000)
|