SoftVC training notebook; training patch

2026-10-16 00:13:07 +00:00 · 2022-12-07 21:41:47 -05:00
parent 4837def3c2
commit 6281ebb454
2 changed files with 349 additions and 0 deletions
@@ -0,0 +1,69 @@
+--- /content/acoustic-model/train.py	2022-09-03 22:51:08.632610312 -0700
+++ train(3).py	2022-09-03 23:15:29.700072538 -0700
+@@ -17,7 +17,7 @@
+ from acoustic.utils import Metric, save_checkpoint, load_checkpoint, plot_spectrogram
+ 
+ 
+-logging.basicConfig(level=logging.DEBUG)
+logging.basicConfig(level=logging.INFO)
+ logger = logging.getLogger(__name__)
+ 
+ ########################################################################################
+@@ -28,10 +28,10 @@
+ LEARNING_RATE = 4e-4
+ BETAS = (0.8, 0.99)
+ WEIGHT_DECAY = 1e-5
+-STEPS = 80000
+STEPS = 20000
+ LOG_INTERVAL = 5
+-VALIDATION_INTERVAL = 1000
+-CHECKPOINT_INTERVAL = 1000
+VALIDATION_INTERVAL = 2000
+CHECKPOINT_INTERVAL = 2000
+ BACKEND = "nccl"
+ INIT_METHOD = "tcp://localhost:54321"
+ 
+@@ -200,6 +200,19 @@
+             # Start validation loop
+             # --------------------------------------------------------------------------#
+ 
+
+
+            if global_step % CHECKPOINT_INTERVAL == 0:
+              save_checkpoint(
+                checkpoint_dir=args.checkpoint_dir,
+                acoustic=acoustic,
+                optimizer=optimizer,
+                step=global_step,
+                loss=validation_loss.value,
+                best=9999,
+                logger=logger,
+            )
+
+             if global_step % VALIDATION_INTERVAL == 0:
+                 acoustic.eval()
+                 validation_loss.reset()
+@@ -242,22 +255,7 @@
+                         f"valid -- epoch: {epoch}, loss: {validation_loss.value:.4f}"
+                     )
+ 
+-                new_best = best_loss > validation_loss.value
+-                if new_best or global_step % CHECKPOINT_INTERVAL:
+-                    if new_best:
+-                        logger.info("-------- new best model found!")
+-                        best_loss = validation_loss.value
+-
+-                    if rank == 0:
+-                        save_checkpoint(
+-                            checkpoint_dir=args.checkpoint_dir,
+-                            acoustic=acoustic,
+-                            optimizer=optimizer,
+-                            step=global_step,
+-                            loss=validation_loss.value,
+-                            best=new_best,
+-                            logger=logger,
+-                        )
+              
+ 
+             # -----------------------------------------------------------------------------#
+             # End validation loop
@@ -0,0 +1,280 @@
+# -*- coding: utf-8 -*-
+"""softvc_pynb.ipynb
+
+Automatically generated by Colaboratory.
+
+Original file is located at
+    https://colab.research.google.com/drive/1FzpL3Ze3X9ESt_nB4w9rTgnfGJxunha0
+"""
+
+!nvidia-smi
+
+#@title Download code
+!git clone https://github.com/bshall/acoustic-model
+!git clone https://github.com/bshall/hubert
+!git clone https://github.com/bshall/hifigan
+
+#@title Download and apply patches
+!pip install gdown
+import gdown
+gdown.download(id="1tcaBRORA627ZGoyMqPEGbmW5Hr2PBDXK", output="/content/acoustic-model-train.patch")
+!patch -p2 < /content/acoustic-model-train.patch
+
+#@title Create directories and setup the colab. Grab the models that are needed for this.
+from pathlib import Path
+import os
+import requests
+
+
+def download_url(url:str,path:Path):
+    url_file_name = url.split("/")[-1]
+    response = requests.get(url)
+    with open(path / Path(f"{url_file_name}"), 'wb') as f:
+        f.write(response.content)
+
+dataset = Path("./dataset")
+
+if dataset.exists() == False:
+    os.mkdir(dataset)
+
+wavs = Path("./dataset/wavs")
+if wavs.exists() == False:
+    os.mkdir(wavs)
+
+mels = Path("./dataset/mels")
+if mels.exists() == False:
+    os.mkdir(mels)
+
+train = Path("./dataset/wavs/train")
+if train.exists() == False:
+    os.mkdir(train)
+
+dev = Path("./dataset/wavs/dev")
+if dev.exists() == False:
+    os.mkdir(dev)
+
+
+check_point = Path("./checkpoint/")
+if check_point.exists() == False:
+    os.mkdir(check_point)
+
+batch_convert_wavs = Path("./batch_convert/")
+if check_point.exists() == False:
+    os.mkdir(batch_convert_wavs)
+
+
+download_url("https://github.com/bshall/acoustic-model/releases/download/v0.1/hubert-soft-0321fd7e.pt",check_point)
+download_url("https://github.com/bshall/acoustic-model/releases/download/v0.1/hubert-discrete-d49e1c77.pt",check_point)
+
+print("Models and Directories Setup")
+
+#@title Mount Google Drive
+from google.colab import drive
+drive.mount('/content/drive')
+
+#@title Unzip dataset
+#@markdown the inside of the dataset zip should be structured like this:
+#@markdown ```
+#@markdown /wavs/
+#@markdown ------/train
+#@markdown -------------/1.wav
+#@markdown -------------/2.wav
+#@markdown -------------/etc.wav
+#@markdown ------/dev
+#@markdown ------------/3.wav
+#@markdown ------------/4.wav
+#@markdown ------------/etc.wav
+#@markdown ```
+#@markdown - dev is the validation dataset and should be roughly 1/10th of the training data
+#@markdown - the wav names don't matter
+#@markdown - the training data should be at least half an hour long for decent results
+    
+import zipfile
+path_to_zipfile = "/content/drive/MyDrive/tts/datasets/bshall-JGSpeech-1.7.zip"#@param { type: "string"}
+with zipfile.ZipFile(path_to_zipfile, 'r') as zip_ref:
+    zip_ref.extractall("/content/dataset")
+
+!ls /content/dataset/wavs/dev | wc -l
+!ls /content/dataset/wavs/train | wc -l
+!du -sh /content/dataset/wavs/dev
+!du -sh /content/dataset/wavs/train
+
+#@title Convert to mels
+!python ./acoustic-model/mels.py ./dataset/wavs ./dataset/mels/ 
+!python ./hubert/encode.py --extension .wav soft ./dataset/wavs ./dataset/soft
+
+
+
+#@title Start training
+#@markdown - Edit hyperparameters in /content/acoustic-model/train.py before you run this cell
+#@markdown - to resume from a previous session, change resume_path to your previous checkpoint
+#@markdown - If this cell hangs with no output, restart the runtime
+
+import os, sys, argparse
+from pathlib import Path
+import torch
+import torch.distributed as dist
+
+sys.path.append("/content/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)
+
+"""# Inference"""
+
+#@title Load the models 
+#@markdown leave hifigan blank to use defaults
+import sys
+import torch
+from typing import Dict, Any
+from collections import OrderedDict
+from torch.nn.modules.utils import consume_prefix_in_state_dict_if_present
+
+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/")
+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"}
+
+# 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")
+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")
+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"}
+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"}
+
+#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
+# Convert to the target speaker
+with torch.inference_mode():
+    # Extract speech units
+    start = time()
+    units = hubert.units(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}")
+
+display.Audio(target.squeeze().cpu(), rate=16000)
+
+#@title Save the output to file
+filename = "bruce.wav" #@param{type: "string"}
+sf.write(filename, target.squeeze().cpu(), samplerate=16000)
+
+#@title Batch convert
+def batch_convert_wavs(folder:Path):
+    # Load the content encoder (either hubert_soft or hubert_discrete)
+    hubert = torch.hub.load("bshall/hubert:main", "hubert_soft").cuda()
+
+    # 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()
+
+    new_state_dict = modify_model(state_dict=acoustic_dict)
+    acoustic.load_state_dict(new_state_dict)
+
+    # Load the vocoder (either hifigan_hubert_soft or hifigan_hubert_discrete)
+    hifigan = torch.hub.load("bshall/hifigan:main", "hifigan_hubert_soft").cuda()
+
+    
+    # Convert to the target speaker
+    with torch.inference_mode():
+        # Extract speech units
+        units = hubert.units(source)
+        # Generate target spectrogram
+        mel = acoustic.generate(units).transpose(1, 2)
+        # Generate audio waveform
+        target = hifigan(mel)
+
+    # TODO later.