prabaerode/zero-shot-tts

1import argparse
2from model import CFM, UNetT, DiT, Trainer
3from model.utils import get_tokenizer
4from model.dataset import load_dataset
5from cached_path import cached_path
6import shutil
7import os
8
9# -------------------------- Dataset Settings --------------------------- #
10target_sample_rate = 24000
11n_mel_channels = 100
12hop_length = 256
13
14
15# -------------------------- Argument Parsing --------------------------- #
16def parse_args():
17    parser = argparse.ArgumentParser(description="Train CFM Model")
18
19    parser.add_argument(
20        "--exp_name", type=str, default="F5TTS_Base", choices=["F5TTS_Base", "E2TTS_Base"], help="Experiment name"
21    )
22    parser.add_argument("--dataset_name", type=str, default="Emilia_ZH_EN", help="Name of the dataset to use")
23    parser.add_argument("--learning_rate", type=float, default=1e-4, help="Learning rate for training")
24    parser.add_argument("--batch_size_per_gpu", type=int, default=256, help="Batch size per GPU")
25    parser.add_argument(
26        "--batch_size_type", type=str, default="frame", choices=["frame", "sample"], help="Batch size type"
27    )
28    parser.add_argument("--max_samples", type=int, default=16, help="Max sequences per batch")
29    parser.add_argument("--grad_accumulation_steps", type=int, default=1, help="Gradient accumulation steps")
30    parser.add_argument("--max_grad_norm", type=float, default=1.0, help="Max gradient norm for clipping")
31    parser.add_argument("--epochs", type=int, default=10, help="Number of training epochs")
32    parser.add_argument("--num_warmup_updates", type=int, default=5, help="Warmup steps")
33    parser.add_argument("--save_per_updates", type=int, default=10, help="Save checkpoint every X steps")
34    parser.add_argument("--last_per_steps", type=int, default=10, help="Save last checkpoint every X steps")
35    parser.add_argument("--finetune", type=bool, default=True, help="Use Finetune")
36
37    parser.add_argument(
38        "--tokenizer", type=str, default="pinyin", choices=["pinyin", "char", "custom"], help="Tokenizer type"
39    )
40    parser.add_argument(
41        "--tokenizer_path",
42        type=str,
43        default=None,
44        help="Path to custom tokenizer vocab file (only used if tokenizer = 'custom')",
45    )
46
47    return parser.parse_args()
48
49
50# -------------------------- Training Settings -------------------------- #
51
52
53def main():
54    args = parse_args()
55
56    # Model parameters based on experiment name
57    if args.exp_name == "F5TTS_Base":
58        wandb_resume_id = None
59        model_cls = DiT
60        model_cfg = dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4)
61        if args.finetune:
62            ckpt_path = str(cached_path("hf://SWivid/F5-TTS/F5TTS_Base/model_1200000.pt"))
63    elif args.exp_name == "E2TTS_Base":
64        wandb_resume_id = None
65        model_cls = UNetT
66        model_cfg = dict(dim=1024, depth=24, heads=16, ff_mult=4)
67        if args.finetune:
68            ckpt_path = str(cached_path("hf://SWivid/E2-TTS/E2TTS_Base/model_1200000.pt"))
69
70    if args.finetune:
71        path_ckpt = os.path.join("ckpts", args.dataset_name)
72        if not os.path.isdir(path_ckpt):
73            os.makedirs(path_ckpt, exist_ok=True)
74            shutil.copy2(ckpt_path, os.path.join(path_ckpt, os.path.basename(ckpt_path)))
75
76    checkpoint_path = os.path.join("ckpts", args.dataset_name)
77
78    # Use the tokenizer and tokenizer_path provided in the command line arguments
79    tokenizer = args.tokenizer
80    if tokenizer == "custom":
81        if not args.tokenizer_path:
82            raise ValueError("Custom tokenizer selected, but no tokenizer_path provided.")
83        tokenizer_path = args.tokenizer_path
84    else:
85        tokenizer_path = args.dataset_name
86
87    vocab_char_map, vocab_size = get_tokenizer(tokenizer_path, tokenizer)
88
89    mel_spec_kwargs = dict(
90        target_sample_rate=target_sample_rate,
91        n_mel_channels=n_mel_channels,
92        hop_length=hop_length,
93    )
94
95    e2tts = CFM(
96        transformer=model_cls(**model_cfg, text_num_embeds=vocab_size, mel_dim=n_mel_channels),
97        mel_spec_kwargs=mel_spec_kwargs,
98        vocab_char_map=vocab_char_map,
99    )
100
101    trainer = Trainer(
102        e2tts,
103        args.epochs,
104        args.learning_rate,
105        num_warmup_updates=args.num_warmup_updates,
106        save_per_updates=args.save_per_updates,
107        checkpoint_path=checkpoint_path,
108        batch_size=args.batch_size_per_gpu,
109        batch_size_type=args.batch_size_type,
110        max_samples=args.max_samples,
111        grad_accumulation_steps=args.grad_accumulation_steps,
112        max_grad_norm=args.max_grad_norm,
113        wandb_project="CFM-TTS",
114        wandb_run_name=args.exp_name,
115        wandb_resume_id=wandb_resume_id,
116        last_per_steps=args.last_per_steps,
117    )
118
119    train_dataset = load_dataset(args.dataset_name, tokenizer, mel_spec_kwargs=mel_spec_kwargs)
120    trainer.train(
121        train_dataset,
122        resumable_with_seed=666,  # seed for shuffling dataset
123    )
124
125
126if __name__ == "__main__":
127    main()
128