Streamline Setup: Enable Offline Model Execution

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
 

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

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)
+
+