| 1 | # training script. |
| 2 | |
| 3 | import os |
| 4 | from importlib.resources import files |
| 5 | |
| 6 | import hydra |
| 7 | from omegaconf import OmegaConf |
| 8 | |
| 9 | from f5_tts.model import CFM, Trainer |
| 10 | from f5_tts.model.dataset import load_dataset |
| 11 | from f5_tts.model.utils import get_tokenizer |
| 12 | |
| 13 | |
| 14 | os.chdir(str(files("f5_tts").joinpath("../.."))) # change working directory to root of project (local editable) |
| 15 | |
| 16 | |
| 17 | @hydra.main(version_base="1.3", config_path=str(files("f5_tts").joinpath("configs")), config_name=None) |
| 18 | def main(model_cfg): |
| 19 | model_cls = hydra.utils.get_class(f"f5_tts.model.{model_cfg.model.backbone}") |
| 20 | model_arc = model_cfg.model.arch |
| 21 | tokenizer = model_cfg.model.tokenizer |
| 22 | mel_spec_type = model_cfg.model.mel_spec.mel_spec_type |
| 23 | |
| 24 | wandb_project = model_cfg.ckpts.get("wandb_project", "CFM-TTS") |
| 25 | wandb_run_name = model_cfg.ckpts.get( |
| 26 | "wandb_run_name", |
| 27 | f"{model_cfg.model.name}_{mel_spec_type}_{model_cfg.model.tokenizer}_{model_cfg.datasets.name}", |
| 28 | ) |
| 29 | wandb_resume_id = model_cfg.ckpts.get("wandb_resume_id", None) |
| 30 | |
| 31 | # set text tokenizer |
| 32 | if tokenizer != "custom": |
| 33 | tokenizer_path = model_cfg.datasets.name |
| 34 | else: |
| 35 | tokenizer_path = model_cfg.model.tokenizer_path |
| 36 | vocab_char_map, vocab_size = get_tokenizer(tokenizer_path, tokenizer) |
| 37 | |
| 38 | # set model |
| 39 | model = CFM( |
| 40 | transformer=model_cls(**model_arc, text_num_embeds=vocab_size, mel_dim=model_cfg.model.mel_spec.n_mel_channels), |
| 41 | mel_spec_kwargs=model_cfg.model.mel_spec, |
| 42 | vocab_char_map=vocab_char_map, |
| 43 | ) |
| 44 | |
| 45 | # init trainer |
| 46 | trainer = Trainer( |
| 47 | model, |
| 48 | epochs=model_cfg.optim.epochs, |
| 49 | learning_rate=model_cfg.optim.learning_rate, |
| 50 | num_warmup_updates=model_cfg.optim.num_warmup_updates, |
| 51 | save_per_updates=model_cfg.ckpts.save_per_updates, |
| 52 | keep_last_n_checkpoints=model_cfg.ckpts.keep_last_n_checkpoints, |
| 53 | checkpoint_path=str(files("f5_tts").joinpath(f"../../{model_cfg.ckpts.save_dir}")), |
| 54 | batch_size_per_gpu=model_cfg.datasets.batch_size_per_gpu, |
| 55 | batch_size_type=model_cfg.datasets.batch_size_type, |
| 56 | max_samples=model_cfg.datasets.max_samples, |
| 57 | grad_accumulation_steps=model_cfg.optim.grad_accumulation_steps, |
| 58 | max_grad_norm=model_cfg.optim.max_grad_norm, |
| 59 | logger=model_cfg.ckpts.logger, |
| 60 | wandb_project=wandb_project, |
| 61 | wandb_run_name=wandb_run_name, |
| 62 | wandb_resume_id=wandb_resume_id, |
| 63 | last_per_updates=model_cfg.ckpts.last_per_updates, |
| 64 | log_samples=model_cfg.ckpts.log_samples, |
| 65 | bnb_optimizer=model_cfg.optim.bnb_optimizer, |
| 66 | mel_spec_type=mel_spec_type, |
| 67 | is_local_vocoder=model_cfg.model.vocoder.is_local, |
| 68 | local_vocoder_path=model_cfg.model.vocoder.local_path, |
| 69 | model_cfg_dict=OmegaConf.to_container(model_cfg, resolve=True), |
| 70 | ) |
| 71 | |
| 72 | train_dataset = load_dataset(model_cfg.datasets.name, tokenizer, mel_spec_kwargs=model_cfg.model.mel_spec) |
| 73 | trainer.train( |
| 74 | train_dataset, |
| 75 | num_workers=model_cfg.datasets.num_workers, |
| 76 | resumable_with_seed=666, # seed for shuffling dataset |
| 77 | ) |
| 78 | |
| 79 | |
| 80 | if __name__ == "__main__": |
| 81 | main() |
| 82 |