| 1 | import argparse |
| 2 | import os |
| 3 | import shutil |
| 4 | from importlib.resources import files |
| 5 | |
| 6 | from cached_path import cached_path |
| 7 | |
| 8 | from f5_tts.model import CFM, DiT, Trainer, UNetT |
| 9 | from f5_tts.model.dataset import load_dataset |
| 10 | from f5_tts.model.utils import get_tokenizer |
| 11 | |
| 12 | |
| 13 | # -------------------------- Dataset Settings --------------------------- # |
| 14 | target_sample_rate = 24000 |
| 15 | n_mel_channels = 100 |
| 16 | hop_length = 256 |
| 17 | win_length = 1024 |
| 18 | n_fft = 1024 |
| 19 | mel_spec_type = "vocos" # 'vocos' or 'bigvgan' |
| 20 | |
| 21 | |
| 22 | # -------------------------- Argument Parsing --------------------------- # |
| 23 | def parse_args(): |
| 24 | parser = argparse.ArgumentParser(description="Train CFM Model") |
| 25 | |
| 26 | parser.add_argument( |
| 27 | "--exp_name", |
| 28 | type=str, |
| 29 | default="F5TTS_v1_Base", |
| 30 | choices=["F5TTS_v1_Base", "F5TTS_Base", "E2TTS_Base"], |
| 31 | help="Experiment name", |
| 32 | ) |
| 33 | parser.add_argument("--dataset_name", type=str, default="Emilia_ZH_EN", help="Name of the dataset to use") |
| 34 | parser.add_argument("--learning_rate", type=float, default=1e-5, help="Learning rate for training") |
| 35 | parser.add_argument("--batch_size_per_gpu", type=int, default=3200, help="Batch size per GPU") |
| 36 | parser.add_argument( |
| 37 | "--batch_size_type", type=str, default="frame", choices=["frame", "sample"], help="Batch size type" |
| 38 | ) |
| 39 | parser.add_argument("--max_samples", type=int, default=64, help="Max sequences per batch") |
| 40 | parser.add_argument("--grad_accumulation_steps", type=int, default=1, help="Gradient accumulation steps") |
| 41 | parser.add_argument("--max_grad_norm", type=float, default=1.0, help="Max gradient norm for clipping") |
| 42 | parser.add_argument("--epochs", type=int, default=100, help="Number of training epochs") |
| 43 | parser.add_argument("--num_warmup_updates", type=int, default=20000, help="Warmup updates") |
| 44 | parser.add_argument("--save_per_updates", type=int, default=50000, help="Save checkpoint every N updates") |
| 45 | parser.add_argument( |
| 46 | "--keep_last_n_checkpoints", |
| 47 | type=int, |
| 48 | default=-1, |
| 49 | help="-1 to keep all, 0 to not save intermediate, > 0 to keep last N checkpoints", |
| 50 | ) |
| 51 | parser.add_argument("--last_per_updates", type=int, default=5000, help="Save last checkpoint every N updates") |
| 52 | parser.add_argument("--finetune", action="store_true", help="Use Finetune") |
| 53 | parser.add_argument("--pretrain", type=str, default=None, help="the path to the checkpoint") |
| 54 | parser.add_argument( |
| 55 | "--tokenizer", type=str, default="pinyin", choices=["pinyin", "char", "custom"], help="Tokenizer type" |
| 56 | ) |
| 57 | parser.add_argument( |
| 58 | "--tokenizer_path", |
| 59 | type=str, |
| 60 | default=None, |
| 61 | help="Path to custom tokenizer vocab file (only used if tokenizer = 'custom')", |
| 62 | ) |
| 63 | parser.add_argument( |
| 64 | "--log_samples", |
| 65 | action="store_true", |
| 66 | help="Log inferenced samples per ckpt save updates", |
| 67 | ) |
| 68 | parser.add_argument("--logger", type=str, default=None, choices=[None, "wandb", "tensorboard"], help="logger") |
| 69 | parser.add_argument( |
| 70 | "--bnb_optimizer", |
| 71 | action="store_true", |
| 72 | help="Use 8-bit Adam optimizer from bitsandbytes", |
| 73 | ) |
| 74 | |
| 75 | return parser.parse_args() |
| 76 | |
| 77 | |
| 78 | # -------------------------- Training Settings -------------------------- # |
| 79 | |
| 80 | |
| 81 | def main(): |
| 82 | args = parse_args() |
| 83 | |
| 84 | checkpoint_path = str(files("f5_tts").joinpath(f"../../ckpts/{args.dataset_name}")) |
| 85 | |
| 86 | # Model parameters based on experiment name |
| 87 | |
| 88 | if args.exp_name == "F5TTS_v1_Base": |
| 89 | wandb_resume_id = None |
| 90 | model_cls = DiT |
| 91 | model_cfg = dict( |
| 92 | dim=1024, |
| 93 | depth=22, |
| 94 | heads=16, |
| 95 | ff_mult=2, |
| 96 | text_dim=512, |
| 97 | conv_layers=4, |
| 98 | ) |
| 99 | if args.finetune: |
| 100 | if args.pretrain is None: |
| 101 | ckpt_path = str(cached_path("hf://SWivid/F5-TTS/F5TTS_v1_Base/model_1250000.safetensors")) |
| 102 | else: |
| 103 | ckpt_path = args.pretrain |
| 104 | |
| 105 | elif args.exp_name == "F5TTS_Base": |
| 106 | wandb_resume_id = None |
| 107 | model_cls = DiT |
| 108 | model_cfg = dict( |
| 109 | dim=1024, |
| 110 | depth=22, |
| 111 | heads=16, |
| 112 | ff_mult=2, |
| 113 | text_dim=512, |
| 114 | text_mask_padding=False, |
| 115 | conv_layers=4, |
| 116 | pe_attn_head=1, |
| 117 | ) |
| 118 | if args.finetune: |
| 119 | if args.pretrain is None: |
| 120 | ckpt_path = str(cached_path("hf://SWivid/F5-TTS/F5TTS_Base/model_1200000.pt")) |
| 121 | else: |
| 122 | ckpt_path = args.pretrain |
| 123 | |
| 124 | elif args.exp_name == "E2TTS_Base": |
| 125 | wandb_resume_id = None |
| 126 | model_cls = UNetT |
| 127 | model_cfg = dict( |
| 128 | dim=1024, |
| 129 | depth=24, |
| 130 | heads=16, |
| 131 | ff_mult=4, |
| 132 | text_mask_padding=False, |
| 133 | pe_attn_head=1, |
| 134 | ) |
| 135 | if args.finetune: |
| 136 | if args.pretrain is None: |
| 137 | ckpt_path = str(cached_path("hf://SWivid/E2-TTS/E2TTS_Base/model_1200000.pt")) |
| 138 | else: |
| 139 | ckpt_path = args.pretrain |
| 140 | |
| 141 | if args.finetune: |
| 142 | if not os.path.isdir(checkpoint_path): |
| 143 | os.makedirs(checkpoint_path, exist_ok=True) |
| 144 | |
| 145 | file_checkpoint = os.path.basename(ckpt_path) |
| 146 | if not file_checkpoint.startswith("pretrained_"): # Change: Add 'pretrained_' prefix to copied model |
| 147 | file_checkpoint = "pretrained_" + file_checkpoint |
| 148 | file_checkpoint = os.path.join(checkpoint_path, file_checkpoint) |
| 149 | if not os.path.isfile(file_checkpoint): |
| 150 | shutil.copy2(ckpt_path, file_checkpoint) |
| 151 | print("copy checkpoint for finetune") |
| 152 | |
| 153 | # Use the tokenizer and tokenizer_path provided in the command line arguments |
| 154 | |
| 155 | tokenizer = args.tokenizer |
| 156 | if tokenizer == "custom": |
| 157 | if not args.tokenizer_path: |
| 158 | raise ValueError("Custom tokenizer selected, but no tokenizer_path provided.") |
| 159 | tokenizer_path = args.tokenizer_path |
| 160 | else: |
| 161 | tokenizer_path = args.dataset_name |
| 162 | |
| 163 | vocab_char_map, vocab_size = get_tokenizer(tokenizer_path, tokenizer) |
| 164 | |
| 165 | print("\nvocab : ", vocab_size) |
| 166 | print("\nvocoder : ", mel_spec_type) |
| 167 | |
| 168 | mel_spec_kwargs = dict( |
| 169 | n_fft=n_fft, |
| 170 | hop_length=hop_length, |
| 171 | win_length=win_length, |
| 172 | n_mel_channels=n_mel_channels, |
| 173 | target_sample_rate=target_sample_rate, |
| 174 | mel_spec_type=mel_spec_type, |
| 175 | ) |
| 176 | |
| 177 | model = CFM( |
| 178 | transformer=model_cls(**model_cfg, text_num_embeds=vocab_size, mel_dim=n_mel_channels), |
| 179 | mel_spec_kwargs=mel_spec_kwargs, |
| 180 | vocab_char_map=vocab_char_map, |
| 181 | ) |
| 182 | |
| 183 | trainer = Trainer( |
| 184 | model, |
| 185 | args.epochs, |
| 186 | args.learning_rate, |
| 187 | num_warmup_updates=args.num_warmup_updates, |
| 188 | save_per_updates=args.save_per_updates, |
| 189 | keep_last_n_checkpoints=args.keep_last_n_checkpoints, |
| 190 | checkpoint_path=checkpoint_path, |
| 191 | batch_size_per_gpu=args.batch_size_per_gpu, |
| 192 | batch_size_type=args.batch_size_type, |
| 193 | max_samples=args.max_samples, |
| 194 | grad_accumulation_steps=args.grad_accumulation_steps, |
| 195 | max_grad_norm=args.max_grad_norm, |
| 196 | logger=args.logger, |
| 197 | wandb_project=args.dataset_name, |
| 198 | wandb_run_name=args.exp_name, |
| 199 | wandb_resume_id=wandb_resume_id, |
| 200 | log_samples=args.log_samples, |
| 201 | last_per_updates=args.last_per_updates, |
| 202 | bnb_optimizer=args.bnb_optimizer, |
| 203 | ) |
| 204 | |
| 205 | train_dataset = load_dataset(args.dataset_name, tokenizer, mel_spec_kwargs=mel_spec_kwargs) |
| 206 | |
| 207 | trainer.train( |
| 208 | train_dataset, |
| 209 | resumable_with_seed=666, # seed for shuffling dataset |
| 210 | ) |
| 211 | |
| 212 | |
| 213 | if __name__ == "__main__": |
| 214 | main() |
| 215 |