diff --git a/README.md b/README.md index bb18f6e..5555d72 100644 --- a/README.md +++ b/README.md @@ -129,6 +129,27 @@ Then, to train the HiFiGAN: └── WavLM.py # wavlm modules (from original WavLM repo) ``` +## **Update for local system** + +1. Download the [prematch_g_02500000.pt](https://github.com/bshall/knn-vc/releases/download/v0.1/g_02500000.pt) model and move it to the "hifigan" folder +2. Download the [WavLM-Large.pt](https://github.com/bshall/knn-vc/releases/download/v0.1/WavLM-Large.pt) model and move it to the "wavlm" folder +3. (If the models are in a different folder) Make sure to make any changes in the hubconf_offline.py file as mentioned here: + + ```python + # ... + + local_file_path = cp / 'hifigan' / 'prematch_g_02500000.pt' # Path to your local HiFi-GAN model + + # ... + + local_file_path = Path(__file__).parent.absolute() / 'wavlm' / 'WavLM-Large.pt' # Path to your local WavLM model + + # ... + + ``` + +4. Run the knnvc_demo.py file (make sure to change the path for the source and reference audio files) + ## Acknowledgements diff --git a/hubconf_offline.py b/hubconf_offline.py new file mode 100644 index 0000000..67fe104 --- /dev/null +++ b/hubconf_offline.py @@ -0,0 +1,70 @@ +import torch +from torch import Tensor +import torch.nn as nn +import torch.nn.functional as F +import logging +import json +from pathlib import Path + +from wavlm.WavLM import WavLM, WavLMConfig +from hifigan.models import Generator as HiFiGAN +from hifigan.utils import AttrDict +from matcher import KNeighborsVC + + +def knn_vc(pretrained=True, progress=True, prematched=True, device='cuda') -> KNeighborsVC: + """ Load kNN-VC (WavLM encoder and HiFiGAN decoder). Optionally use vocoder trained on `prematched` data. """ + hifigan, hifigan_cfg = hifigan_wavlm(pretrained, progress, prematched, device) + wavlm = wavlm_large(pretrained, progress, device) + knnvc = KNeighborsVC(wavlm, hifigan, hifigan_cfg, device) + return knnvc + + +def hifigan_wavlm(pretrained=True, progress=True, prematched=True, device='cuda') -> HiFiGAN: + """ Load pretrained hifigan trained to vocode wavlm features. Optionally use weights trained on `prematched` data. """ + cp = Path(__file__).parent.absolute() + + # Use the local config file + with open(cp/'hifigan'/'config_v1_wavlm.json') as f: + data = f.read() + json_config = json.loads(data) + h = AttrDict(json_config) + device = torch.device(device) + + # Load HiFi-GAN model from local file + generator = HiFiGAN(h).to(device) + + if pretrained: + # Update this to load the local .pt file + local_file_path = cp / 'hifigan' / 'prematch_g_02500000.pt' # Path to your local HiFi-GAN model + state_dict_g = torch.load(local_file_path, map_location=device) # Load state dict from local file + generator.load_state_dict(state_dict_g['generator']) + + generator.eval() + generator.remove_weight_norm() + print(f"[HiFiGAN] Generator loaded with {sum([p.numel() for p in generator.parameters()]):,d} parameters.") + return generator, h + + +def wavlm_large(pretrained=True, progress=True, device='cuda') -> WavLM: + """Load the WavLM large checkpoint from the original paper. See https://github.com/microsoft/unilm/tree/master/wavlm for details. """ + if torch.cuda.is_available() == False: + if str(device) != 'cpu': + logging.warning(f"Overriding device {device} to cpu since no GPU is available.") + device = 'cpu' + + # Load WavLM model from local file + local_file_path = Path(__file__).parent.absolute() / 'wavlm' / 'WavLM-Large.pt' # Path to your local WavLM model + checkpoint = torch.load(local_file_path, map_location=device) # Load state dict from local file + + cfg = WavLMConfig(checkpoint['cfg']) + device = torch.device(device) + model = WavLM(cfg) + + if pretrained: + model.load_state_dict(checkpoint['model']) + + model = model.to(device) + model.eval() + print(f"WavLM-Large loaded with {sum([p.numel() for p in model.parameters()]):,d} parameters.") + return model diff --git a/knnvc_demo.py b/knnvc_demo.py new file mode 100644 index 0000000..c570928 --- /dev/null +++ b/knnvc_demo.py @@ -0,0 +1,25 @@ +import torch +import torchaudio +import numpy as np +from pathlib import Path +from hubconf_offline import knn_vc # Import the knn_vc function from hubconf.py + + +knnvc_model = knn_vc(pretrained=True, prematched=True, device='cpu') + +# path to 16kHz, single-channel, source waveform +src_wav_path = './src/src.wav' +# list of paths to all reference waveforms (each must be 16kHz, single-channel) from the target speaker +ref_wav_paths = ['./ref/ref1.wav', ] + +query_seq = knnvc_model.get_features(src_wav_path) +print(query_seq.shape) +matching_set = knnvc_model.get_matching_set(ref_wav_paths) +print(matching_set.shape) + +out_wav = knnvc_model.match(query_seq, matching_set, topk=4) +print(out_wav.shape) + +torchaudio.save('./out/out.wav', out_wav[None], 16000) + +