support finetune_cli hydra and fix some minor bugs

src/f5_tts/config/E2TTS_Base_finetune.yaml

@@ -0,0 +1,46 @@
+hydra:
+  run:
+    dir: ckpts/finetune_${model.name}_${model.mel_spec.mel_spec_type}_${model.tokenizer}_${datasets.name}/${now:%Y-%m-%d}/${now:%H-%M-%S}
+    
+datasets:
+  name: Emilia_ZH_EN # dataset name
+  batch_size_per_gpu: 38400  # 8 GPUs, 8 * 38400 = 307200
+  batch_size_type: frame # "frame" or "sample"
+  max_samples: 64  # max sequences per batch if use frame-wise batch_size. we set 32 for small models, 64 for base models
+  num_workers: 16 # number of workers
+
+optim:
+  epochs: 15 # max epochs
+  learning_rate: 7.5e-5 # learning rate
+  num_warmup_updates: 20000  # warmup steps
+  grad_accumulation_steps: 1  # note: updates = steps / grad_accumulation_steps
+  max_grad_norm: 1.0 # gradient clipping
+  bnb_optimizer: False # use bnb optimizer or not
+
+model:
+  name: F5TTS_Base # model name
+  tokenizer: pinyin # tokenizer type
+  tokenizer_path: None  # if tokenizer = 'custom', define the path to the tokenizer you want to use (should be vocab.txt)
+  arch:
+    dim: 1024 # model dim
+    depth: 22 # model depth
+    heads: 16 # model heads
+    ff_mult: 2 # feedforward expansion
+    text_dim: 512 # text encoder dim
+    conv_layers: 4 # convolution layers
+  mel_spec:
+    target_sample_rate: 24000 # target sample rate
+    n_mel_channels: 100 # mel channel
+    hop_length: 256 # hop length
+    win_length: 1024 # window length
+    n_fft: 1024 # fft length
+    mel_spec_type: vocos  # 'vocos' or 'bigvgan'
+    is_local_vocoder: False # use local vocoder or not
+    local_vocoder_path: None # local vocoder path
+
+ckpts:
+  logger: wandb # wandb | tensorboard | None
+  save_per_updates: 50000 # save checkpoint per steps
+  last_per_steps: 5000 # save last checkpoint per steps
+  pretain_ckpt_path: ckpts/E2TTS_Base/model_1200000.pt
+  save_dir: ckpts/${model.name}_${model.mel_spec.mel_spec_type}_${model.tokenizer}_${datasets.name}/${now:%Y-%m-%d}/${now:%H-%M-%S}

src/f5_tts/config/E2TTS_Base_train.yaml

@@ -7,6 +7,7 @@ datasets:
   batch_size_per_gpu: 38400  # 8 GPUs, 8 * 38400 = 307200
   batch_size_type: frame # "frame" or "sample"
   max_samples: 64  # max sequences per batch if use frame-wise batch_size. we set 32 for small models, 64 for base models
+  num_workers: 16 # number of workers
 
 optim:
   epochs: 15 # max epochs
@@ -14,6 +15,7 @@ optim:
   num_warmup_updates: 20000  # warmup steps
   grad_accumulation_steps: 1  # note: updates = steps / grad_accumulation_steps
   max_grad_norm: 1.0 # gradient clipping
+  bnb_optimizer: False # use bnb optimizer or not
 
 model:
   name: E2TTS_Base # model name
@@ -35,6 +37,7 @@ model:
     local_vocoder_path: None # path to local vocoder
 
 ckpts:
+  logger: wandb # wandb | tensorboard | None
   save_per_updates: 50000 # save checkpoint per steps
   last_per_steps: 5000 # save last checkpoint per steps
   save_dir: ckpts/${model.name}_${model.mel_spec.mel_spec_type}_${model.tokenizer}_${datasets.name}/${now:%Y-%m-%d}/${now:%H-%M-%S}

src/f5_tts/config/E2TTS_Small_train.yaml

@@ -7,6 +7,7 @@ datasets:
   batch_size_per_gpu: 38400  # 8 GPUs, 8 * 38400 = 307200
   batch_size_type: frame # "frame" or "sample"
   max_samples: 64  # max sequences per batch if use frame-wise batch_size. we set 32 for small models, 64 for base models
+  num_workers: 16 # number of workers
 
 optim:
   epochs: 15
@@ -14,6 +15,7 @@ optim:
   num_warmup_updates: 20000  # warmup steps
   grad_accumulation_steps: 1  # note: updates = steps / grad_accumulation_steps
   max_grad_norm: 1.0
+  bnb_optimizer: False  
 
 model:
   name: E2TTS_Small
@@ -35,6 +37,7 @@ model:
     local_vocoder_path: None
 
 ckpts:
+  logger: wandb # wandb | tensorboard | None
   save_per_updates: 50000 # save checkpoint per steps
   last_per_steps: 5000 # save last checkpoint per steps
   save_dir: ckpts/${model.name}_${model.mel_spec.mel_spec_type}_${model.tokenizer}_${datasets.name}/${now:%Y-%m-%d}/${now:%H-%M-%S}

src/f5_tts/config/F5TTS_Base_finetune.yaml

@@ -0,0 +1,46 @@
+hydra:
+  run:
+    dir: ckpts/finetune_${model.name}_${model.mel_spec.mel_spec_type}_${model.tokenizer}_${datasets.name}/${now:%Y-%m-%d}/${now:%H-%M-%S}
+    
+datasets:
+  name: Emilia_ZH_EN # dataset name
+  batch_size_per_gpu: 38400  # 8 GPUs, 8 * 38400 = 307200
+  batch_size_type: frame # "frame" or "sample"
+  max_samples: 64  # max sequences per batch if use frame-wise batch_size. we set 32 for small models, 64 for base models
+  num_workers: 16 # number of workers
+
+optim:
+  epochs: 15 # max epochs
+  learning_rate: 7.5e-5 # learning rate
+  num_warmup_updates: 20000  # warmup steps
+  grad_accumulation_steps: 1  # note: updates = steps / grad_accumulation_steps
+  max_grad_norm: 1.0 # gradient clipping
+  bnb_optimizer: False # use bnb optimizer or not
+
+model:
+  name: F5TTS_Base # model name
+  tokenizer: pinyin # tokenizer type
+  tokenizer_path: None  # if tokenizer = 'custom', define the path to the tokenizer you want to use (should be vocab.txt)
+  arch:
+    dim: 1024 # model dim
+    depth: 22 # model depth
+    heads: 16 # model heads
+    ff_mult: 2 # feedforward expansion
+    text_dim: 512 # text encoder dim
+    conv_layers: 4 # convolution layers
+  mel_spec:
+    target_sample_rate: 24000 # target sample rate
+    n_mel_channels: 100 # mel channel
+    hop_length: 256 # hop length
+    win_length: 1024 # window length
+    n_fft: 1024 # fft length
+    mel_spec_type: vocos  # 'vocos' or 'bigvgan'
+    is_local_vocoder: False # use local vocoder or not
+    local_vocoder_path: None # local vocoder path
+
+ckpts:
+  logger: wandb # wandb | tensorboard | None
+  save_per_updates: 50000 # save checkpoint per steps
+  last_per_steps: 5000 # save last checkpoint per steps
+  pretain_ckpt_path: ckpts/F5TTS_Base/model_1200000.pt
+  save_dir: ckpts/${model.name}_${model.mel_spec.mel_spec_type}_${model.tokenizer}_${datasets.name}/${now:%Y-%m-%d}/${now:%H-%M-%S}

src/f5_tts/config/F5TTS_Base_train.yaml

@@ -7,6 +7,7 @@ datasets:
   batch_size_per_gpu: 38400  # 8 GPUs, 8 * 38400 = 307200
   batch_size_type: frame # "frame" or "sample"
   max_samples: 64  # max sequences per batch if use frame-wise batch_size. we set 32 for small models, 64 for base models
+  num_workers: 16 # number of workers
 
 optim:
   epochs: 15 # max epochs
@@ -14,6 +15,7 @@ optim:
   num_warmup_updates: 20000  # warmup steps
   grad_accumulation_steps: 1  # note: updates = steps / grad_accumulation_steps
   max_grad_norm: 1.0 # gradient clipping
+  bnb_optimizer: False # use bnb optimizer or not
 
 model:
   name: F5TTS_Base # model name
@@ -37,6 +39,7 @@ model:
     local_vocoder_path: None # local vocoder path
 
 ckpts:
+  logger: wandb # wandb | tensorboard | None
   save_per_updates: 50000 # save checkpoint per steps
   last_per_steps: 5000 # save last checkpoint per steps
   save_dir: ckpts/${model.name}_${model.mel_spec.mel_spec_type}_${model.tokenizer}_${datasets.name}/${now:%Y-%m-%d}/${now:%H-%M-%S}

src/f5_tts/config/F5TTS_Small_train.yaml

@@ -7,6 +7,7 @@ datasets:
   batch_size_per_gpu: 38400  # 8 GPUs, 8 * 38400 = 307200
   batch_size_type: frame # "frame" or "sample"
   max_samples: 64  # max sequences per batch if use frame-wise batch_size. we set 32 for small models, 64 for base models
+  num_workers: 16 # number of workers
 
 optim:
   epochs: 15
@@ -14,6 +15,7 @@ optim:
   num_warmup_updates: 20000  # warmup steps
   grad_accumulation_steps: 1  # note: updates = steps / grad_accumulation_steps
   max_grad_norm: 1.0
+  bnb_optimizer: False
 
 model:
   name: F5TTS_Small
@@ -37,6 +39,7 @@ model:
     local_vocoder_path: None
 
 ckpts:
+  logger: wandb # wandb | tensorboard | None
   save_per_updates: 50000 # save checkpoint per steps
   last_per_steps: 5000 # save last checkpoint per steps
   save_dir: ckpts/${model.name}_${model.mel_spec.mel_spec_type}_${model.tokenizer}_${datasets.name}/${now:%Y-%m-%d}/${now:%H-%M-%S}

src/f5_tts/model/trainer.py

@@ -91,7 +91,7 @@ class Trainer:
         elif self.logger == "tensorboard":
             from torch.utils.tensorboard import SummaryWriter
 
-            self.writer = SummaryWriter(log_dir=f"runs/{wandb_run_name}")
+            self.writer = SummaryWriter(log_dir=f"{checkpoint_path}/runs/{wandb_run_name}")
 
         self.model = model
 

src/f5_tts/model/utils.py

@@ -113,7 +113,7 @@ def get_tokenizer(dataset_name, tokenizer: str = "pinyin"):
         with open(tokenizer_path, "r", encoding="utf-8") as f:
             vocab_char_map = {}
             for i, char in enumerate(f):
-                vocab_char_map[char[:-1]] = i
+                vocab_char_map[char.strip()] = i  # ignore \n
         vocab_size = len(vocab_char_map)
         assert vocab_char_map[" "] == 0, "make sure space is of idx 0 in vocab.txt, cuz 0 is used for unknown char"
 
@@ -125,7 +125,7 @@ def get_tokenizer(dataset_name, tokenizer: str = "pinyin"):
         with open(dataset_name, "r", encoding="utf-8") as f:
             vocab_char_map = {}
             for i, char in enumerate(f):
-                vocab_char_map[char[:-1]] = i
+                vocab_char_map[char.strip()] = i
         vocab_size = len(vocab_char_map)
 
     return vocab_char_map, vocab_size

src/f5_tts/train/finetune_cli.py

@@ -1,6 +1,6 @@
-import argparse
 import os
 import shutil
+import hydra
 
 from cached_path import cached_path
 from f5_tts.model import CFM, UNetT, DiT, Trainer
@@ -9,163 +9,76 @@ from f5_tts.model.dataset import load_dataset
 from importlib.resources import files
 
 
-# -------------------------- Dataset Settings --------------------------- #
-target_sample_rate = 24000
-n_mel_channels = 100
-hop_length = 256
-win_length = 1024
-n_fft = 1024
-mel_spec_type = "vocos"  # 'vocos' or 'bigvgan'
+@hydra.main(config_path=os.path.join("..", "configs"), config_name=None)
+def main(cfg):
+    tokenizer = cfg.model.tokenizer
+    mel_spec_type = cfg.model.mel_spec.mel_spec_type
+    exp_name = f"finetune_{cfg.model.name}_{mel_spec_type}_{cfg.model.tokenizer}_{cfg.datasets.name}"
 
+    # set text tokenizer
+    if tokenizer != "custom":
+        tokenizer_path = cfg.datasets.name
+    else:
+        tokenizer_path = cfg.model.tokenizer_path
+    vocab_char_map, vocab_size = get_tokenizer(tokenizer_path, tokenizer)
 
-# -------------------------- Argument Parsing --------------------------- #
-def parse_args():
-    # batch_size_per_gpu = 1000 settting for gpu 8GB
-    # batch_size_per_gpu = 1600 settting for gpu 12GB
-    # batch_size_per_gpu = 2000 settting for gpu 16GB
-    # batch_size_per_gpu = 3200 settting for gpu 24GB
-
-    # num_warmup_updates = 300 for 5000 sample about 10 hours
-
-    # change save_per_updates , last_per_steps change this value what you need  ,
-
-    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-5, help="Learning rate for training")
-    parser.add_argument("--batch_size_per_gpu", type=int, default=3200, 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=100, help="Number of training epochs")
-    parser.add_argument("--num_warmup_updates", type=int, default=300, help="Warmup steps")
-    parser.add_argument("--save_per_updates", type=int, default=10000, help="Save checkpoint every X steps")
-    parser.add_argument("--last_per_steps", type=int, default=50000, help="Save last checkpoint every X steps")
-    parser.add_argument("--finetune", type=bool, default=True, help="Use Finetune")
-    parser.add_argument("--pretrain", type=str, default=None, help="the path to the checkpoint")
-    parser.add_argument(
-        "--tokenizer", type=str, default="pinyin", choices=["pinyin", "char", "custom"], help="Tokenizer type"
-    )
-    parser.add_argument(
-        "--tokenizer_path",
-        type=str,
-        default=None,
-        help="Path to custom tokenizer vocab file (only used if tokenizer = 'custom')",
-    )
-    parser.add_argument(
-        "--log_samples",
-        type=bool,
-        default=False,
-        help="Log inferenced samples per ckpt save steps",
-    )
-    parser.add_argument("--logger", type=str, default=None, choices=["wandb", "tensorboard"], help="logger")
-    parser.add_argument(
-        "--bnb_optimizer",
-        type=bool,
-        default=False,
-        help="Use 8-bit Adam optimizer from bitsandbytes",
-    )
-
-    return parser.parse_args()
-
-
-# -------------------------- Training Settings -------------------------- #
-
-
-def main():
-    args = parse_args()
-
-    checkpoint_path = str(files("f5_tts").joinpath(f"../../ckpts/{args.dataset_name}"))
+    print("\nvocab : ", vocab_size)
+    print("\nvocoder : ", mel_spec_type)
 
     # Model parameters based on experiment name
-    if args.exp_name == "F5TTS_Base":
-        wandb_resume_id = None
+    if "F5TTS" in cfg.model.name:
         model_cls = DiT
-        model_cfg = dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4)
-        if args.finetune:
-            if args.pretrain is None:
-                ckpt_path = str(cached_path("hf://SWivid/F5-TTS/F5TTS_Base/model_1200000.pt"))
-            else:
-                ckpt_path = args.pretrain
-    elif args.exp_name == "E2TTS_Base":
-        wandb_resume_id = None
+        ckpt_path = cfg.ckpts.pretain_ckpt_path or str(cached_path("hf://SWivid/F5-TTS/F5TTS_Base/model_1200000.pt"))
+    elif "E2TTS" in cfg.model.name:
         model_cls = UNetT
-        model_cfg = dict(dim=1024, depth=24, heads=16, ff_mult=4)
-        if args.finetune:
-            if args.pretrain is None:
-                ckpt_path = str(cached_path("hf://SWivid/E2-TTS/E2TTS_Base/model_1200000.pt"))
-            else:
-                ckpt_path = args.pretrain
-
-    if args.finetune:
-        if not os.path.isdir(checkpoint_path):
-            os.makedirs(checkpoint_path, exist_ok=True)
+        ckpt_path = cfg.ckpts.pretain_ckpt_path or str(cached_path("hf://SWivid/F5-TTS/E2TTS_Base/model_1200000.pt"))
+    wandb_resume_id = None
 
-        file_checkpoint = os.path.join(checkpoint_path, os.path.basename(ckpt_path))
-        if not os.path.isfile(file_checkpoint):
-            shutil.copy2(ckpt_path, file_checkpoint)
-            print("copy checkpoint for finetune")
+    checkpoint_path = str(files("f5_tts").joinpath(f"../../{cfg.ckpts.save_dir}"))
 
-    # Use the tokenizer and tokenizer_path provided in the command line arguments
-    tokenizer = args.tokenizer
-    if tokenizer == "custom":
-        if not args.tokenizer_path:
-            raise ValueError("Custom tokenizer selected, but no tokenizer_path provided.")
-        tokenizer_path = args.tokenizer_path
-    else:
-        tokenizer_path = args.dataset_name
+    if not os.path.isdir(checkpoint_path):
+        os.makedirs(checkpoint_path, exist_ok=True)
 
-    vocab_char_map, vocab_size = get_tokenizer(tokenizer_path, tokenizer)
-
-    print("\nvocab : ", vocab_size)
-    print("\nvocoder : ", mel_spec_type)
-
-    mel_spec_kwargs = dict(
-        n_fft=n_fft,
-        hop_length=hop_length,
-        win_length=win_length,
-        n_mel_channels=n_mel_channels,
-        target_sample_rate=target_sample_rate,
-        mel_spec_type=mel_spec_type,
-    )
+    file_checkpoint = os.path.join(checkpoint_path, os.path.basename(ckpt_path))
+    if not os.path.isfile(file_checkpoint):
+        shutil.copy2(ckpt_path, file_checkpoint)
+        print("copy checkpoint for finetune")
 
     model = CFM(
-        transformer=model_cls(**model_cfg, text_num_embeds=vocab_size, mel_dim=n_mel_channels),
-        mel_spec_kwargs=mel_spec_kwargs,
+        transformer=model_cls(**cfg.model.arch, text_num_embeds=vocab_size, mel_dim=cfg.model.mel_spec.n_mel_channels),
+        mel_spec_kwargs=cfg.model.mel_spec,
         vocab_char_map=vocab_char_map,
     )
 
     trainer = Trainer(
         model,
-        args.epochs,
-        args.learning_rate,
-        num_warmup_updates=args.num_warmup_updates,
-        save_per_updates=args.save_per_updates,
+        epochs=cfg.optim.epochs,
+        learning_rate=cfg.optim.learning_rate,
+        num_warmup_updates=cfg.optim.num_warmup_updates,
+        save_per_updates=cfg.ckpts.save_per_updates,
         checkpoint_path=checkpoint_path,
-        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,
-        logger=args.logger,
-        wandb_project=args.dataset_name,
-        wandb_run_name=args.exp_name,
+        batch_size=cfg.datasets.batch_size_per_gpu,
+        batch_size_type=cfg.datasets.batch_size_type,
+        max_samples=cfg.datasets.max_samples,
+        grad_accumulation_steps=cfg.optim.grad_accumulation_steps,
+        max_grad_norm=cfg.optim.max_grad_norm,
+        logger=cfg.ckpts.logger,
+        wandb_project=cfg.datasets.name,
+        wandb_run_name=exp_name,
         wandb_resume_id=wandb_resume_id,
-        log_samples=args.log_samples,
-        last_per_steps=args.last_per_steps,
-        bnb_optimizer=args.bnb_optimizer,
+        log_samples=True,
+        last_per_steps=cfg.ckpts.last_per_steps,
+        bnb_optimizer=cfg.optim.bnb_optimizer,
+        mel_spec_type=mel_spec_type,
+        is_local_vocoder=cfg.model.mel_spec.is_local_vocoder,
+        local_vocoder_path=cfg.model.mel_spec.local_vocoder_path,
     )
 
-    train_dataset = load_dataset(args.dataset_name, tokenizer, mel_spec_kwargs=mel_spec_kwargs)
+    train_dataset = load_dataset(cfg.datasets.name, tokenizer, mel_spec_kwargs=cfg.model.mel_spec)
 
     trainer.train(
         train_dataset,
+        num_workers=cfg.datasets.num_workers,
         resumable_with_seed=666,  # seed for shuffling dataset
     )
 

src/f5_tts/train/train.py

@@ -48,11 +48,13 @@ def main(cfg):
         max_samples=cfg.datasets.max_samples,
         grad_accumulation_steps=cfg.optim.grad_accumulation_steps,
         max_grad_norm=cfg.optim.max_grad_norm,
+        logger=cfg.ckpts.logger,
         wandb_project="CFM-TTS",
         wandb_run_name=exp_name,
         wandb_resume_id=wandb_resume_id,
         last_per_steps=cfg.ckpts.last_per_steps,
         log_samples=True,
+        bnb_optimizer=cfg.optim.bnb_optimizer,
         mel_spec_type=mel_spec_type,
         is_local_vocoder=cfg.model.mel_spec.is_local_vocoder,
         local_vocoder_path=cfg.model.mel_spec.local_vocoder_path,
@@ -61,6 +63,7 @@ def main(cfg):
     train_dataset = load_dataset(cfg.datasets.name, tokenizer, mel_spec_kwargs=cfg.model.mel_spec)
     trainer.train(
         train_dataset,
+        num_workers=cfg.datasets.num_workers,
         resumable_with_seed=666,  # seed for shuffling dataset
     )