mirror of
https://github.com/SWivid/F5-TTS.git
synced 2026-01-08 03:12:37 -08:00
94 lines
4.0 KiB
Python
94 lines
4.0 KiB
Python
import argparse
|
|
from model import CFM, UNetT, DiT, MMDiT, Trainer
|
|
from model.utils import get_tokenizer
|
|
from model.dataset import load_dataset
|
|
|
|
# -------------------------- Dataset Settings --------------------------- #
|
|
target_sample_rate = 24000
|
|
n_mel_channels = 100
|
|
hop_length = 256
|
|
|
|
tokenizer = "pinyin" # 'pinyin', 'char', or 'custom'
|
|
tokenizer_path = None # if tokenizer = 'custom', define the path to the tokenizer you want to use (should be vocab.txt)
|
|
|
|
# -------------------------- Argument Parsing --------------------------- #
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(description='Train CFM Model')
|
|
|
|
parser.add_argument('--exp_name', type=str, default="F5TTS_Base", choices=["F5TTS_Base", "E2TTS_Base"],help='Experiment name')
|
|
parser.add_argument('--dataset_name', type=str, default="Emilia_ZH_EN", help='Name of the dataset to use')
|
|
parser.add_argument('--learning_rate', type=float, default=1e-4, help='Learning rate for training')
|
|
parser.add_argument('--batch_size_per_gpu', type=int, default=400, help='Batch size per GPU')
|
|
parser.add_argument('--batch_size_type', type=str, default="frame", choices=["frame", "sample"],help='Batch size type')
|
|
parser.add_argument('--max_samples', type=int, default=64, help='Max sequences per batch')
|
|
parser.add_argument('--grad_accumulation_steps', type=int, default=1,help='Gradient accumulation steps')
|
|
parser.add_argument('--max_grad_norm', type=float, default=1.0, help='Max gradient norm for clipping')
|
|
parser.add_argument('--epochs', type=int, default=11, help='Number of training epochs')
|
|
parser.add_argument('--num_warmup_updates', type=int, default=200, help='Warmup steps')
|
|
parser.add_argument('--save_per_updates', type=int, default=800, help='Save checkpoint every X steps')
|
|
parser.add_argument('--last_per_steps', type=int, default=400, help='Save last checkpoint every X steps')
|
|
|
|
return parser.parse_args()
|
|
|
|
# -------------------------- Training Settings -------------------------- #
|
|
|
|
def main():
|
|
args = parse_args()
|
|
|
|
# Model parameters based on experiment name
|
|
if args.exp_name == "F5TTS_Base":
|
|
wandb_resume_id = None
|
|
model_cls = DiT
|
|
model_cfg = dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4)
|
|
elif args.exp_name == "E2TTS_Base":
|
|
wandb_resume_id = None
|
|
model_cls = UNetT
|
|
model_cfg = dict(dim=1024, depth=24, heads=16, ff_mult=4)
|
|
|
|
# Use the dataset_name provided in the command line
|
|
tokenizer_path = args.dataset_name if tokenizer != "custom" else tokenizer_path
|
|
vocab_char_map, vocab_size = get_tokenizer(tokenizer_path, tokenizer)
|
|
|
|
mel_spec_kwargs = dict(
|
|
target_sample_rate=target_sample_rate,
|
|
n_mel_channels=n_mel_channels,
|
|
hop_length=hop_length,
|
|
)
|
|
|
|
e2tts = CFM(
|
|
transformer=model_cls(
|
|
**model_cfg,
|
|
text_num_embeds=vocab_size,
|
|
mel_dim=n_mel_channels
|
|
),
|
|
mel_spec_kwargs=mel_spec_kwargs,
|
|
vocab_char_map=vocab_char_map,
|
|
)
|
|
|
|
trainer = Trainer(
|
|
e2tts,
|
|
args.epochs,
|
|
args.learning_rate,
|
|
num_warmup_updates=args.num_warmup_updates,
|
|
save_per_updates=args.save_per_updates,
|
|
checkpoint_path=f'ckpts/{args.exp_name}',
|
|
batch_size=args.batch_size_per_gpu,
|
|
batch_size_type=args.batch_size_type,
|
|
max_samples=args.max_samples,
|
|
grad_accumulation_steps=args.grad_accumulation_steps,
|
|
max_grad_norm=args.max_grad_norm,
|
|
wandb_project="CFM-TTS",
|
|
wandb_run_name=args.exp_name,
|
|
wandb_resume_id=wandb_resume_id,
|
|
last_per_steps=args.last_per_steps,
|
|
)
|
|
|
|
train_dataset = load_dataset(args.dataset_name, tokenizer, mel_spec_kwargs=mel_spec_kwargs)
|
|
trainer.train(train_dataset,
|
|
resumable_with_seed=666 # seed for shuffling dataset
|
|
)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|