Add model exporting and CLI, update docs and testing data

README.md

@@ -1,4 +1,4 @@
-<img src="./vall-e.png" width="450px"></img>
+<img src="./vall-e.png" width="550px"></img>
 
 # VALL-E
 
@@ -6,11 +6,11 @@ An unofficial PyTorch implementation of [VALL-E](https://valle-demo.github.io/),
 
 [!["Buy Me A Coffee"](https://www.buymeacoffee.com/assets/img/custom_images/orange_img.png)](https://www.buymeacoffee.com/enhuiz)
 
-## Install
+## Get Started
 
 ### Requirements
 
-Since the trainer is based on [DeepSpeed](https://github.com/microsoft/DeepSpeed.git), you will need to have a GPU that DeepSpeed has developed and tested against, as well as a CUDA or ROCm compiler pre-installed to install this package.
+Since the trainer is based on [DeepSpeed](https://github.com/microsoft/DeepSpeed#requirements), you will need to have a GPU that DeepSpeed has developed and tested against, as well as a CUDA or ROCm compiler pre-installed to install this package.
 
 ### Install
 
@@ -18,7 +18,7 @@ Since the trainer is based on [DeepSpeed](https://github.com/microsoft/DeepSpeed
 pip install git+https://github.com/enhuiz/vall-e
 ```
 
-### Clone
+Or you may clone by:
 
 ```
 git clone --recurse-submodules https://github.com/enhuiz/vall-e.git
@@ -28,6 +28,8 @@ Note that the code is only tested under `Python 3.10.7`.
 
 ## Usage
 
+### Training
+
 1. Put your data into a folder, e.g. `data/your_data`. Audio files should be named with the suffix `.wav` and text files with `.normalized.txt`.
 
 2. Quantize the data:
@@ -50,6 +52,24 @@ python -m vall_e.emb.g2p data/your_data
 python -m vall_e.train yaml=config/your_data/ar_or_nar.yml
 ```
 
+You may quit your training any time by just typing `quit` in your CLI. The latest checkpoint will be automatically saved.
+
+6. Export trained models:
+
+Both trained models need to be exported to a certain path. To export either of them, run:
+
+```
+python -m vall_e.export zoo/ar_or_nar.pt yaml=config/your_data/ar_or_nar.yml
+```
+
+This will export the latest checkpoint.
+
+### Synthesis
+
+```
+python -m vall_e <text> <ref_path> <out_path> --ar-ckpt zoo/ar.pt --nar-ckpt zoo/nar.pt
+```
+
 ## TODO
 
 - [x] AR model for the first quantizer
@@ -59,6 +79,7 @@ python -m vall_e.train yaml=config/your_data/ar_or_nar.yml
 - [x] Implement AdaLN for NAR model.
 - [x] Sample-wise quantization level sampling for NAR training.
 - [ ] Pre-trained checkpoint and demos on LibriTTS
+- [x] CLI synthesis interface
 
 ## Notice
 

config/test/ar.yml

@@ -2,3 +2,7 @@ data_dirs: [data/test]
 
 model: ar-quarter
 batch_size: 1
+eval_batch_size: 1
+save_ckpt_every: 500
+eval_every: 500
+max_iter: 1000

config/test/nar.yml

@@ -2,3 +2,7 @@ data_dirs: [data/test]
 
 model: nar-quarter
 batch_size: 1
+eval_batch_size: 1
+save_ckpt_every: 500
+eval_every: 500
+max_iter: 1000

data/test/test2.phn.txt

@@ -0,0 +1 @@
+test.phn.txt

data/test/test2.qnt.pt

@@ -0,0 +1 @@
+test.qnt.pt

vall_e/__main__.py

@@ -1,14 +1,43 @@
 import argparse
 from pathlib import Path
 
+import torch
+from einops import rearrange
+
+from .emb import g2p, qnt
+from .utils import to_device
+
 
 def main():
     parser = argparse.ArgumentParser("VALL-E TTS")
     parser.add_argument("text")
-    parser.add_argument("output")
-    parser.add_argument("--reference", type=Path)
+    parser.add_argument("reference", type=Path)
+    parser.add_argument("out_path", type=Path)
+    parser.add_argument("--ar-ckpt", type=Path, default="zoo/ar.pt")
+    parser.add_argument("--nar-ckpt", type=Path, default="zoo/nar.pt")
+    parser.add_argument("--device", default="cuda")
     args = parser.parse_args()
 
+    ar = torch.load(args.ar_ckpt).to(args.device)
+    nar = torch.load(args.nar_ckpt).to(args.device)
+
+    symmap = ar.phone_symmap
+
+    proms = qnt.encode_from_file(args.reference)
+    proms = rearrange(proms, "1 l t -> t l")
+
+    phns = torch.tensor([symmap[p] for p in g2p.encode(args.text)])
+
+    proms = to_device(proms, args.device)
+    phns = to_device(phns, args.device)
+
+    resp_list = ar(text_list=[phns], proms_list=[proms])
+    resps_list = [r.unsqueeze(-1) for r in resp_list]
+
+    resps_list = nar(text_list=[phns], proms_list=[proms], resps_list=resps_list)
+    qnt.decode_to_file(resps=resps_list[0], path=args.out_path)
+    print(args.out_path, "saved.")
+
 
 if __name__ == "__main__":
     main()

vall_e/data.py

@@ -220,7 +220,7 @@ def _load_train_val_paths():
     val_paths = []
 
     for data_dir in cfg.data_dirs:
-        paths.extend(tqdm(data_dir.rglob("**/*.qnt.pt")))
+        paths.extend(tqdm(data_dir.rglob("*.qnt.pt")))
 
     if len(paths) == 0:
         raise RuntimeError(f"Failed to find any .qnt.pt file in {cfg.data_dirs}.")
@@ -244,7 +244,7 @@ def _load_train_val_paths():
 def _load_test_paths():
     test_paths = []
     for data_dir in cfg.test_data_dirs:
-        test_paths.extend(data_dir.rglob("**/*.asr.txt"))
+        test_paths.extend(data_dir.rglob("*.phn.txt"))
     test_paths = sorted(test_paths)
     return test_paths
 

vall_e/emb/qnt.py

@@ -52,7 +52,7 @@ def _replace_file_extension(path, suffix):
 
 
 @torch.inference_mode()
-def encode(wav, sr, device="cuda"):
+def encode(wav: Tensor, sr: int, device="cuda"):
     """
     Args:
         wav: (t)
@@ -67,6 +67,13 @@ def encode(wav, sr, device="cuda"):
     return qnt
 
 
+def encode_from_file(path, device="cuda"):
+    wav, sr = torchaudio.load(str(path))
+    if wav.shape[0] == 2:
+        wav = wav[:1]
+    return encode(wav, sr, device)
+
+
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument("folder", type=Path)
@@ -80,10 +87,7 @@ def main():
         out_path = _replace_file_extension(path, ".qnt.pt")
         if out_path.exists():
             continue
-        wav, sr = torchaudio.load(path)
-        if wav.shape[0] == 2:
-            wav = wav[:1]
-        qnt = encode(wav, sr)
+        qnt = encode_from_file(path)
         torch.save(qnt.cpu(), out_path)
 
 

vall_e/export.py

@@ -0,0 +1,25 @@
+import argparse
+
+import torch
+
+from .data import VALLEDatset, create_train_val_dataloader
+from .train import load_engines
+
+
+def main():
+    parser = argparse.ArgumentParser("Save trained model to path.")
+    parser.add_argument("path")
+    args = parser.parse_args()
+
+    engine = load_engines()
+    model = engine["model"].module.cpu()
+    train_dl, *_ = create_train_val_dataloader()
+    assert isinstance(train_dl.dataset, VALLEDatset)
+    model.phone_symmap = train_dl.dataset.phone_symmap
+    model.spkr_symmap = train_dl.dataset.spkr_symmap
+    torch.save(model, args.path)
+    print(args.path, "saved.")
+
+
+if __name__ == "__main__":
+    main()