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/), [!["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 --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()