Inference cleanup, gitignore additions, etc.

.gitignore

@@ -20,6 +20,7 @@ __pycache__/
 models/Wav2Lip/python/
 models/tacotron2/python/
 tts/python/
+voice_conversion/softvc/python/
 
 # Rust build
 target/
@@ -34,3 +35,5 @@ tts/hifigan/hifisr
 # Other garbage to prune/remove
 tts/hifigan/LJSpeech-1.1/
 
+# Output files
+output.wav

start_tts_server.sh

@@ -1,6 +1,6 @@
 #!/bin/bash
+# This is a helpful script for local development using Docker images.
 
 cd models/tts
 source python/bin/activate
 ./vocodes_server.py
-

voice_conversion/softvc/.gitignore

@@ -1,2 +1,10 @@
 # Python Virtual Env
 python/
+
+# Python object cache
+__pycache__/
+
+# Vim cruft
+*.swo
+*.swp
+*~

voice_conversion/softvc/cli_inference.py

@@ -1,50 +1,44 @@
 #!/usr/bin/env python
 
-# import os, sys, argparse
+import IPython.display as display
-# from pathlib import Path
+import argparse
-# import torch
+import importlib.util
-# import torch.distributed as dist
+import soundfile as sf
-# 
+import sys
-# sys.path.append("./acoustic-model/")
-# from train import train
-# 
-# checkpoint_save_path = "/content/drive/MyDrive/tts/softvc-checkpoints/jeff_v2/acoustic" #@param {type: "string"}
-# _checkpoint_save_path = Path(checkpoint_save_path)
-# if _checkpoint_save_path.exists() == False:
-#     os.mkdir(_checkpoint_save_path)
-# 
-# resume_path = "/content/checkpoint/hubert-soft-0321fd7e.pt" #@param {type: "string"}
-# 
-# world_size = torch.cuda.device_count()
-# 
-# 
-# class DictObj:
-#   def __init__(self, in_dict:dict):
-#     assert isinstance(in_dict, dict)
-#     for key, val in in_dict.items():
-#       if isinstance(val, (list, tuple)):
-#         setattr(self, key, [DictObj(x) if isinstance(x, dict) else x for x in val])
-#       else:
-#         setattr(self, key, DictObj(val) if isinstance(val, dict) else val)
-# 
-# args=DictObj({"checkpoint_dir": _checkpoint_save_path, "dataset_dir": Path("/content/dataset"), "resume":Path(resume_path), "discrete": False})
-# try:
-#   dist.destroy_process_group()
-# except:
-#   pass
-# train(0, world_size, args)
-# 
-
-
-#@title Load the models 
-#@markdown leave hifigan blank to use defaults
 import sys  
 import torch
-from typing import Dict, Any
+import torchaudio
 from collections import OrderedDict
+from time import time
 from torch.nn.modules.utils import consume_prefix_in_state_dict_if_present
+from typing import Dict, Any
 
-def modify_model(state_dict: Dict[str,Any]):
+sys.path.append("./acoustic-model/")
+from acoustic import AcousticModel
+
+#sys.path.insert(0, './acoustic-model/acoustic/')
+#from model import AcousticModel
+
+sys.path.append("./hifigan/")
+from hifigan.generator import HifiganGenerator
+from hifigan.generator import consume_prefix_in_state_dict_if_present
+
+# ===========================================
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--acoustic_model_filename', type=str, required=True)
+parser.add_argument('--hifigan_model_filename', type=str, required=True)
+parser.add_argument('--source_audio_filename', type=str, required=True)
+parser.add_argument('--output_audio_filename', type=str, required=True)
+args = parser.parse_args()
+
+print('\n'.join(f'[arg] {k} = {v}' for k, v in vars(args).items()))
+
+# ===========================================
+
+def load_acoustic_model(checkpoint_path):
+
+    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]
@@ -52,88 +46,61 @@ def modify_model(state_dict: Dict[str,Any]):
             new_state_dict[name] = value
         return new_state_dict
 
-#sys.path.append("/content/acoustic-model/")
+    # Load the acoustic model (either hubert_soft or hubert_discrete) 
-sys.path.append("./acoustic-model/")
+    acoustic_dict = torch.load(checkpoint_path) # this used ddp
-from acoustic import AcousticModel
+    acoustic = AcousticModel().cuda()
+    acoustic.eval()
 
-hubert = torch.hub.load("bshall/hubert:main", "hubert_soft").cuda()
+    new_state_dict = modify_model(state_dict=acoustic_dict)
+    acoustic.load_state_dict(new_state_dict)
+    return acoustic
 
-#check_point_path = '/content/drive/MyDrive/tts/softvc-checkpoints/tay/acoustic/taylor_swift-acoustic.pt' #@param {type:"string"}
+# ===========================================
-check_point_path = '/home/bt/models/voice-conversion/trump-acoustic.pt' #@param {type:"string"}
 
-# Load the acoustic model (either hubert_soft or hubert_discrete) 
+def load_hifigan_model(hifigan_checkpoint_path=None):
-acoustic_dict = torch.load(check_point_path) # this used ddp
+    # NB(bt, 2022-12-07): Not sure why this "spec loading" is needed.
-acoustic = AcousticModel().cuda()
+    spec = importlib.util.spec_from_file_location("hifigan.utils", "./hifigan/hifigan/utils.py")
-acoustic.eval()
+    u = importlib.util.module_from_spec(spec)
+    sys.modules["hifigan.utils"] = u
+    spec.loader.exec_module(u)
 
-import importlib.util
+    spec = importlib.util.spec_from_file_location("hifigan.generator", "./hifigan/hifigan/generator.py")
-import sys
+    g = importlib.util.module_from_spec(spec)
+    sys.modules["hifigan.generator"] = g
+    spec.loader.exec_module(g)
 
-#spec = importlib.util.spec_from_file_location("hifigan.utils", "/content/hifigan/hifigan/utils.py")
+    if hifigan_checkpoint_path:
-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", "/content/hifigan/hifigan/generator.py")
-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)
-#hifigan_checkpoint_path = "/content/drive/MyDrive/tts/softvc-checkpoints/tay/hifigan/taylor_swift-hifigan.pt" #@param {type: "string"}
-hifigan_checkpoint_path = "/home/bt/models/voice-conversion/trump-hifigan.pt" #@param {type: "string"}
-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:
-else:
         hifigan = torch.hub.load("bshall/hifigan:main", "hifigan_hubert_soft")
-hifigan.cuda()
 
+    hifigan.cuda()
+    return hifigan
 
 # ===========================================
 
 
-#@title This code selects the source file
+def load_wav_data(file_name):
-import IPython.display as display
+    source, sr = torchaudio.load(f"{file_name}")
-import torch, torchaudio
+    source = torchaudio.functional.resample(source, sr, 16000)
-import sys  
+    source = torch.mean(source, dim=0).unsqueeze(0)
-sys.path.insert(0, './acoustic-model/acoustic/')
+    source = source.unsqueeze(0).cuda()
-from model import AcousticModel
+    return source
-import soundfile as sf
-#from google.colab import files
-from time import time
-sys.path.append("/content/hifigan/")
-from hifigan.generator import HifiganGenerator
-from hifigan.generator import consume_prefix_in_state_dict_if_present
-#uploaded = files.upload()
-#file_name = "/content/drive/MyDrive/20220914110418-9dd00c7845-godzilla_mdxKaraoke_vocals.[mvsep.com].wav" #@param {type: "string"}
-file_name = "./samples/source/easy/brandon1.wav"
-
-#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
+def do_voice_conversion(hubert, acoustic, hifigan, audio_source):
-# Convert to the target speaker
+    #@title Run the inference and display the result
-with torch.inference_mode():
+    # Convert to the target speaker
+    with torch.inference_mode():
         # Extract speech units
         start = time()
-    units = hubert.units(source)
+        units = hubert.units(audio_source)
         end = time()
         # Generate target spectrogram
         print(f"time for feature embedding:{end-start}")
@@ -147,11 +114,15 @@ with torch.inference_mode():
         target = hifigan(mel)
         end = time()
         print(f"time for wavform {end-start}")
-
+        return target
-display.Audio(target.squeeze().cpu(), rate=16000)
-
 # ===========================================
 
-#@title Save the output to file
+hubert = torch.hub.load("bshall/hubert:main", "hubert_soft").cuda()
-filename = "output.wav" #@param{type: "string"}
+
-sf.write(filename, target.squeeze().cpu(), samplerate=16000)
+acoustic = load_acoustic_model(args.acoustic_model_filename)
+hifigan = load_hifigan_model(args.hifigan_model_filename)
+
+source_audio = load_wav_data(args.source_audio_filename)
+target_audio = do_voice_conversion(hubert, acoustic, hifigan, source_audio)
+
+sf.write(args.output_audio_filename, target_audio.squeeze().cpu(), samplerate=16000)