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,157 +1,128 @@
 #!/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]):
-    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
-
-#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()
+#sys.path.insert(0, './acoustic-model/acoustic/')
+#from model import AcousticModel
 
-#check_point_path = '/content/drive/MyDrive/tts/softvc-checkpoints/tay/acoustic/taylor_swift-acoustic.pt' #@param {type:"string"}
+sys.path.append("./hifigan/")
-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 = AcousticModel().cuda()
-acoustic.eval()
-
-import importlib.util
-import sys
-
-#spec = importlib.util.spec_from_file_location("hifigan.utils", "/content/hifigan/hifigan/utils.py")
-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:
-  hifigan = torch.hub.load("bshall/hifigan:main", "hifigan_hubert_soft")
-hifigan.cuda()
-
-
-# ===========================================
-
-
-#@title This code selects the source file
-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)
-
 
 # ===========================================
 
-#@title Run the inference and display the result
+parser = argparse.ArgumentParser()
-# Convert to the target speaker
+parser.add_argument('--acoustic_model_filename', type=str, required=True)
-with torch.inference_mode():
+parser.add_argument('--hifigan_model_filename', type=str, required=True)
-    # Extract speech units
+parser.add_argument('--source_audio_filename', type=str, required=True)
-    start = time()
+parser.add_argument('--output_audio_filename', type=str, required=True)
-    units = hubert.units(source)
+args = parser.parse_args()
-    end = time()
-    # Generate target spectrogram
-    print(f"time for feature embedding:{end-start}")
 
-    start = time()
+print('\n'.join(f'[arg] {k} = {v}' for k, v in vars(args).items()))
-    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}")
-
-display.Audio(target.squeeze().cpu(), rate=16000)
 
 # ===========================================
 
-#@title Save the output to file
+def load_acoustic_model(checkpoint_path):
-filename = "output.wav" #@param{type: "string"}
+
-sf.write(filename, target.squeeze().cpu(), samplerate=16000)
+    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
+
+    # Load the acoustic model (either hubert_soft or hubert_discrete) 
+    acoustic_dict = torch.load(checkpoint_path) # this used ddp
+    acoustic = AcousticModel().cuda()
+    acoustic.eval()
+
+    new_state_dict = modify_model(state_dict=acoustic_dict)
+    acoustic.load_state_dict(new_state_dict)
+    return acoustic
+
+# ===========================================
+
+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", "./hifigan/hifigan/generator.py")
+    g = importlib.util.module_from_spec(spec)
+    sys.modules["hifigan.generator"] = g
+    spec.loader.exec_module(g)
+
+    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()
+    return hifigan
+
+# ===========================================
+
+
+def load_wav_data(file_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()
+    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(audio_source)
+        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}")
+        return target
+# ===========================================
+
+hubert = torch.hub.load("bshall/hubert:main", "hubert_soft").cuda()
+
+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)