diff --git a/README.md b/README.md
index 02843ee..b3ba375 100644
--- a/README.md
+++ b/README.md
@@ -1,4 +1,4 @@
-
+
# VALL-E
@@ -6,11 +6,11 @@ An unofficial PyTorch implementation of [VALL-E](https://valle-demo.github.io/),
[](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 --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
diff --git a/config/test/ar.yml b/config/test/ar.yml
index 7dcf809..46d208e 100644
--- a/config/test/ar.yml
+++ b/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
diff --git a/config/test/nar.yml b/config/test/nar.yml
index 80775dd..2fc372b 100644
--- a/config/test/nar.yml
+++ b/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
diff --git a/data/test/test.ar.recon.wav b/data/test/test.ar.recon.wav
deleted file mode 100644
index 5ee5fe2..0000000
Binary files a/data/test/test.ar.recon.wav and /dev/null differ
diff --git a/data/test/test.nar.1-1.wav b/data/test/test.nar.1-1.wav
deleted file mode 100644
index 3043e17..0000000
Binary files a/data/test/test.nar.1-1.wav and /dev/null differ
diff --git a/data/test/test.nar.1-2.wav b/data/test/test.nar.1-2.wav
deleted file mode 100644
index e870081..0000000
Binary files a/data/test/test.nar.1-2.wav and /dev/null differ
diff --git a/data/test/test.nar.1-3.wav b/data/test/test.nar.1-3.wav
deleted file mode 100644
index 9d7f137..0000000
Binary files a/data/test/test.nar.1-3.wav and /dev/null differ
diff --git a/data/test/test.nar.1-4.wav b/data/test/test.nar.1-4.wav
deleted file mode 100644
index 9d7f137..0000000
Binary files a/data/test/test.nar.1-4.wav and /dev/null differ
diff --git a/data/test/test.nar.1-5.wav b/data/test/test.nar.1-5.wav
deleted file mode 100644
index 9d7f137..0000000
Binary files a/data/test/test.nar.1-5.wav and /dev/null differ
diff --git a/data/test/test.nar.1-6.wav b/data/test/test.nar.1-6.wav
deleted file mode 100644
index 9d7f137..0000000
Binary files a/data/test/test.nar.1-6.wav and /dev/null differ
diff --git a/data/test/test.nar.1-7.wav b/data/test/test.nar.1-7.wav
deleted file mode 100644
index 9d7f137..0000000
Binary files a/data/test/test.nar.1-7.wav and /dev/null differ
diff --git a/data/test/test.nar.init.wav b/data/test/test.nar.init.wav
deleted file mode 100644
index f303877..0000000
Binary files a/data/test/test.nar.init.wav and /dev/null differ
diff --git a/data/test/test2.phn.txt b/data/test/test2.phn.txt
new file mode 120000
index 0000000..dc8968d
--- /dev/null
+++ b/data/test/test2.phn.txt
@@ -0,0 +1 @@
+test.phn.txt
\ No newline at end of file
diff --git a/data/test/test2.qnt.pt b/data/test/test2.qnt.pt
new file mode 120000
index 0000000..90f9762
--- /dev/null
+++ b/data/test/test2.qnt.pt
@@ -0,0 +1 @@
+test.qnt.pt
\ No newline at end of file
diff --git a/vall_e/__main__.py b/vall_e/__main__.py
index 9f416d0..f20df51 100644
--- a/vall_e/__main__.py
+++ b/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()
diff --git a/vall_e/data.py b/vall_e/data.py
index 2da64ff..b92435f 100644
--- a/vall_e/data.py
+++ b/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
diff --git a/vall_e/emb/qnt.py b/vall_e/emb/qnt.py
index a5a4052..e0efd72 100644
--- a/vall_e/emb/qnt.py
+++ b/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)
diff --git a/vall_e/export.py b/vall_e/export.py
new file mode 100644
index 0000000..a79cc40
--- /dev/null
+++ b/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()