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,48 +1,42 @@
 #!/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
+
+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()
@@ -52,39 +46,29 @@ def modify_model(state_dict: Dict[str,Any]):
             new_state_dict[name] = value
         return new_state_dict
 
-#sys.path.append("/content/acoustic-model/")
-sys.path.append("./acoustic-model/")
-from acoustic import AcousticModel
-
-hubert = torch.hub.load("bshall/hubert:main", "hubert_soft").cuda()
-
-#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) 
-acoustic_dict = torch.load(check_point_path) # this used ddp
+    acoustic_dict = torch.load(checkpoint_path) # this used ddp
     acoustic = AcousticModel().cuda()
     acoustic.eval()
 
-import importlib.util
+    new_state_dict = modify_model(state_dict=acoustic_dict)
-import sys
+    acoustic.load_state_dict(new_state_dict)
+    return acoustic
 
-#spec = importlib.util.spec_from_file_location("hifigan.utils", "/content/hifigan/hifigan/utils.py")
+# ===========================================
+
+def load_hifigan_model(hifigan_checkpoint_path=None):
+    # NB(bt, 2022-12-07): Not sure why this "spec loading" is needed.
     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()
@@ -92,48 +76,31 @@ if hifigan_checkpoint_path:
         hifigan.load_state_dict(checkpoint)
         hifigan.remove_weight_norm()
         hifigan.eval()
- 
     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
-import torch, torchaudio
-import sys  
-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("/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)
+    return source
-
 
 # ===========================================
 
+def do_voice_conversion(hubert, acoustic, hifigan, audio_source):
     #@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)
+        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)