diff --git a/styletts2/api.py b/styletts2/api.py index 4f614cbb..b8d8053a 100644 --- a/styletts2/api.py +++ b/styletts2/api.py @@ -47,3 +47,37 @@ def preprocess( overwrite=overwrite, debug=debug, ) + + +def train( + config: "StyleTTS2Config", + config_file: str | Path | None = None, + mode: str | core.TrainingMode = "first", + precision: str = "32", + accelerator: str = "auto", + devices: str | int = "auto", + nodes: int = 1, + strategy: str = "ddp", +): + """Train an end-to-end (StyleTTS2) model + + Args: + config (StyleTTS2Config): your StyleTTS2 configuration + config_file (str | Path | None): for logging purposes only -- if provided, the logs will include a copy of this file + mode ("first" | "second" | "finetune"): Training mode: "first" (acoustic pre-training with TMS), "second" (joint diffusion+adversarial), or "finetune" + precision (str): Floating-point precision passed to Lightning Trainer (e.g., "32", "16-mixed", "bf16-mixed") + accelerator (str): PyTorch Lightning Accelerator to use: https://pytorch-lightning.readthedocs.io/en/stable/extensions/accelerator.html + devices ("auto" | str | int): the number of GPUs to use on each node as a str or int; use "auto" to let pytoch-lightning decide + nodes (int): the number of nodes to use + strategy (str): the strategy for data parallelization: https://pytorch-lightning.readthedocs.io/en/stable/accelerators/gpu_intermediate.html" + """ + core.train( + config=config, + config_file=config_file, + mode=core.TrainingMode(mode), + precision=precision, + accelerator=accelerator, + devices=str(devices), + nodes=nodes, + strategy=strategy, + ) diff --git a/styletts2/cli/fetch_pretrained.py b/styletts2/cli/fetch_pretrained.py index 27ae978d..ef3f42f2 100644 --- a/styletts2/cli/fetch_pretrained.py +++ b/styletts2/cli/fetch_pretrained.py @@ -12,7 +12,7 @@ def fetch_pretrained( file_okay=True, ), ): - """Download StyleTTS2 pretrained model weights from HuggingFace. + """Download StyleTTS2 pretrained model weights from HuggingFace Run this command on a node with internet access before submitting a GPU training job. The files are stored in the HuggingFace hub cache and are diff --git a/styletts2/cli/synthesize.py b/styletts2/cli/synthesize.py index 35cef778..92a8e3c9 100644 --- a/styletts2/cli/synthesize.py +++ b/styletts2/cli/synthesize.py @@ -19,7 +19,7 @@ def load_styletts2_model(model_path: Path, device): - """Load a StyleTTS2 model and mel transform from a checkpoint.""" + """Load a StyleTTS2 model and mel transform from a checkpoint""" import torch from ..lightning import StyleTTS2 @@ -43,7 +43,7 @@ def load_reference_style( reference_path: Path, device, ): - """Load a reference audio file and return a pre-computed style encoding. + """Load a reference audio file and return a pre-computed style encoding Runs ``_load_reference_mel`` then ``model._encode_reference``, returning ``ref_s`` of shape ``[1, 256]`` on ``device``. Call this at startup to @@ -73,7 +73,7 @@ def synthesize_one( language: str | None = None, text_representation: DatasetTextRepresentation | None = None, ): - """Synthesize a single utterance and return a float32 numpy waveform. + """Synthesize a single utterance and return a float32 numpy waveform Works only with stage-2 (or finetune) checkpoints that include the diffusion sampler. Stage-1 checkpoints will raise an AttributeError @@ -227,7 +227,7 @@ def synthesize( help="Blend weight for prosody style (0 = pure reference, 1 = pure diffusion).", ), ): - """Synthesize audio from text using a trained StyleTTS2 model. + """Synthesize audio from text using a trained StyleTTS2 model Examples: diff --git a/styletts2/cli/train.py b/styletts2/cli/train.py index 0486afb9..f449ed96 100644 --- a/styletts2/cli/train.py +++ b/styletts2/cli/train.py @@ -1,31 +1,23 @@ -import os -import shutil -import sys -from enum import Enum -from pathlib import Path from typing import Annotated import typer from everyvoice.base_cli.interfaces import train_base_command_interface from merge_args import merge_args - -class Mode(str, Enum): - first = "first" - second = "second" - finetune = "finetune" +from .. import core +from ..core.train import TrainingMode @merge_args(train_base_command_interface) def train( mode: Annotated[ - Mode, + TrainingMode, typer.Option( "-m", "--mode", help="Training mode: 'first' (acoustic pre-training with TMA), 'second' (joint diffusion+adversarial), or 'finetune'.", ), - ] = Mode.first, + ] = TrainingMode.first, precision: Annotated[ str, typer.Option( @@ -40,142 +32,8 @@ def train( **styletts2 train config/everyvoice-text-to-wav.yaml --mode first** """ - from everyvoice.utils import spinner - - with spinner(): - import torch - - if not torch.cuda.is_available(): - # device="cuda" is assumed in multiple places, so let's just tell the user up front - # It's also pointless to try on CPU if it takes around a week on GPU... - sys.exit( - "ERROR: StyleTTS2 training requires a GPU with the cuda accellerator" - ) - - import lightning as L - from everyvoice.utils import update_config_from_cli_args - from lightning.pytorch.callbacks import LearningRateMonitor, ModelCheckpoint - from lightning.pytorch.loggers import TensorBoardLogger - from lightning.pytorch.strategies import DDPStrategy - - from ..ev_config import StyleTTS2Config - from ..ev_config.translation import to_native_config - from ..lightning import StyleTTS2, StyleTTS2DataModule - - config_file: Path = kwargs["config_file"] - config_args: list[str] = kwargs.get("config_args", []) - - ev_config = StyleTTS2Config.load_config_from_path(config_file) - update_config_from_cli_args(config_args, ev_config) - - config = to_native_config(ev_config) - - from everyvoice.text.lookups import lookuptables_from_config - - lang2id, _ = lookuptables_from_config(ev_config) - if ev_config.model.multilingual and not lang2id: - sys.exit( - "ERROR: model.multilingual is true but no 'language' values were " - "found in your training/validation filelists." - ) - config["lang2id"] = lang2id - - tr = ev_config.training - max_epochs = ( - tr.epochs_1st - if mode == Mode.first - else tr.epochs_2nd if mode == Mode.second else tr.max_epochs - ) - - log_dir = config["log_dir"] - os.makedirs(log_dir, exist_ok=True) - shutil.copy(str(config_file), os.path.join(log_dir, config_file.name)) - - # Stage 1 and stage 2 share the same log_dir; give each mode its own - # tensorboard sub-run so stage 2's restarted step/epoch counters don't - # get drawn as a continuation of stage 1's curve. - mode_sub_dirs = { - Mode.first: "stage-1", - Mode.second: "stage-2", - Mode.finetune: "finetune", - } - tb_logger = TensorBoardLogger( - save_dir=tr.logger.save_dir, - name=tr.logger.name, - version=tr.logger.version, - sub_dir=mode_sub_dirs[mode], - ) - - ckpt_dir = os.path.join(log_dir, "checkpoints") - ckpt_filename = f"epoch_{mode.value}_" + "{epoch:05d}" - # Always keep the last checkpoint regardless of performance. - last_ckpt_callback = ModelCheckpoint( - dirpath=ckpt_dir, - filename=ckpt_filename, - save_top_k=1, - save_last=True, - every_n_train_steps=tr.ckpt_steps, - every_n_epochs=tr.ckpt_epochs, - enable_version_counter=True, - save_on_train_epoch_end=True, - ) - # Stage 1 and stage 2 share the same checkpoints dir; name the "last" - # checkpoint per-mode so stage 2 doesn't replace stage 1's checkpoint. - mode_ckpt_names = { - Mode.first: "stage-1-last", - Mode.second: "stage-2-last", - Mode.finetune: "finetune-last", - } - last_ckpt_callback.CHECKPOINT_NAME_LAST = mode_ckpt_names[mode] - # Keep only the top-k checkpoints ranked by val/mel (lower is better). - monitored_ckpt_callback = ModelCheckpoint( - dirpath=ckpt_dir, - filename=ckpt_filename, - monitor="val/mel", - mode="min", - save_top_k=tr.save_top_k_ckpts, - every_n_train_steps=tr.ckpt_steps, - every_n_epochs=tr.ckpt_epochs, - enable_version_counter=False, + ev_config = core.load_config( + config_file=kwargs["config_file"], # don't pop this one, train() needs it + config_args=kwargs.pop("config_args"), ) - lr_monitor = LearningRateMonitor(logging_interval="step") - - devices = kwargs.get("devices", "auto") - strategy = kwargs.get("strategy", "ddp") - - # GAN training uses separate discriminator/generator backward passes, - # so find_unused_parameters=True is required for DDP correctness. - try: - n_devices = int(devices) - multi_gpu = n_devices > 1 - except (TypeError, ValueError): - multi_gpu = devices not in ("auto", "1", 1) - - if strategy == "ddp" or (multi_gpu and strategy == "auto"): - resolved_strategy = DDPStrategy(find_unused_parameters=True) - else: - resolved_strategy = strategy - - trainer = L.Trainer( - max_epochs=max_epochs, - devices=devices, - num_nodes=kwargs.get("nodes", 1), - accelerator=kwargs.get("accelerator", "auto"), - strategy=resolved_strategy, - precision=precision, - logger=tb_logger, - callbacks=[monitored_ckpt_callback, last_ckpt_callback, lr_monitor], - log_every_n_steps=config.get("log_interval", 10), - enable_progress_bar=True, - ) - - datamodule = StyleTTS2DataModule(config, load_for_everyvoice=True) - model = StyleTTS2(config, mode=mode.value) - - resume_ckpt = ( - str(tr.finetune_checkpoint) - if tr.finetune_checkpoint and os.path.exists(tr.finetune_checkpoint) - else None - ) - - trainer.fit(model, datamodule=datamodule, ckpt_path=resume_ckpt) + core.train(config=ev_config, mode=mode, precision=precision, **kwargs) diff --git a/styletts2/core/__init__.py b/styletts2/core/__init__.py index a0e51f15..c8fabbe7 100644 --- a/styletts2/core/__init__.py +++ b/styletts2/core/__init__.py @@ -1,3 +1,10 @@ from .preprocess import PREPROCESS_CATEGORIES, load_config, preprocess +from .train import TrainingMode, train -__all__ = ["load_config", "preprocess", "PREPROCESS_CATEGORIES"] +__all__ = [ + "load_config", + "preprocess", + "PREPROCESS_CATEGORIES", + "train", + "TrainingMode", +] diff --git a/styletts2/core/train.py b/styletts2/core/train.py new file mode 100644 index 00000000..4a8607e2 --- /dev/null +++ b/styletts2/core/train.py @@ -0,0 +1,168 @@ +import os +import shutil +import sys +from enum import Enum +from pathlib import Path +from typing import TYPE_CHECKING + +import yaml + +if TYPE_CHECKING: + from ..ev_config import StyleTTS2Config + + +class TrainingMode(str, Enum): + first = "first" + second = "second" + finetune = "finetune" + + +def train( + config: "StyleTTS2Config", + config_file: str | Path | None, + mode: TrainingMode, + precision: str, + accelerator: str, + devices: str, + nodes: int, + strategy: str, +): + from everyvoice.utils import spinner + + with spinner(): + import torch + + if not torch.cuda.is_available(): + # device="cuda" is assumed in multiple places, so let's just tell the user up front + # It's also pointless to try on CPU if it takes around a week on GPU... + sys.exit( + "ERROR: StyleTTS2 training requires a GPU with the cuda accellerator" + ) + + import lightning as L + from everyvoice.text.lookups import lookuptables_from_config + from lightning.pytorch.callbacks import LearningRateMonitor, ModelCheckpoint + from lightning.pytorch.loggers import TensorBoardLogger + from lightning.pytorch.strategies import DDPStrategy + + from ..ev_config.translation import to_native_config + from ..lightning import StyleTTS2, StyleTTS2DataModule + + ev_config = config + native_config = to_native_config(ev_config) + + lang2id, _ = lookuptables_from_config(ev_config) + if ev_config.model.multilingual and not lang2id: + sys.exit( + "ERROR: model.multilingual is true but no 'language' values were " + "found in your training/validation filelists." + ) + native_config["lang2id"] = lang2id + + tr = ev_config.training + max_epochs = ( + tr.epochs_1st + if mode == TrainingMode.first + else tr.epochs_2nd if mode == TrainingMode.second else tr.max_epochs + ) + + log_dir = native_config["log_dir"] + os.makedirs(log_dir, exist_ok=True) + if config_file is not None: + config_file_s = str(config_file) + shutil.copy(config_file_s, os.path.join(log_dir, config_file_s)) + # Log the full in-memory config object as really used, including any + # --config-args overrides on the CLI or programmatic changes to config + with open( + os.path.join(log_dir, "effective-text-to-wav-config.yaml"), + mode="w", + encoding="utf-8", + ) as f: + config_dict = ev_config.model_dump() + yaml.dump(config_dict, f, default_flow_style=None, allow_unicode=True) + + # Stage 1 and stage 2 share the same log_dir; give each mode its own + # tensorboard sub-run so stage 2's restarted step/epoch counters don't + # get drawn as a continuation of stage 1's curve. + mode_sub_dirs = { + TrainingMode.first: "stage-1", + TrainingMode.second: "stage-2", + TrainingMode.finetune: "finetune", + } + tb_logger = TensorBoardLogger( + save_dir=tr.logger.save_dir, + name=tr.logger.name, + version=tr.logger.version, + sub_dir=mode_sub_dirs[mode], + ) + + ckpt_dir = os.path.join(log_dir, "checkpoints") + ckpt_filename = f"epoch_{mode.value}_" + "{epoch:05d}" + # Always keep the last checkpoint regardless of performance. + last_ckpt_callback = ModelCheckpoint( + dirpath=ckpt_dir, + filename=ckpt_filename, + save_top_k=1, + save_last=True, + every_n_train_steps=tr.ckpt_steps, + every_n_epochs=tr.ckpt_epochs, + enable_version_counter=True, + save_on_train_epoch_end=True, + ) + # Stage 1 and stage 2 share the same checkpoints dir; name the "last" + # checkpoint per-mode so stage 2 doesn't replace stage 1's checkpoint. + mode_ckpt_names = { + TrainingMode.first: "stage-1-last", + TrainingMode.second: "stage-2-last", + TrainingMode.finetune: "finetune-last", + } + last_ckpt_callback.CHECKPOINT_NAME_LAST = mode_ckpt_names[mode] + # Keep only the top-k checkpoints ranked by val/mel (lower is better). + monitored_ckpt_callback = ModelCheckpoint( + dirpath=ckpt_dir, + filename=ckpt_filename, + monitor="val/mel", + mode="min", + save_top_k=tr.save_top_k_ckpts, + every_n_train_steps=tr.ckpt_steps, + every_n_epochs=tr.ckpt_epochs, + enable_version_counter=False, + ) + lr_monitor = LearningRateMonitor(logging_interval="step") + + # GAN training uses separate discriminator/generator backward passes, + # so find_unused_parameters=True is required for DDP correctness. + try: + n_devices = int(devices) + multi_gpu = n_devices > 1 + except (TypeError, ValueError): + multi_gpu = devices not in ("auto", "1", 1) + + if strategy == "ddp" or (multi_gpu and strategy == "auto"): + resolved_strategy = DDPStrategy(find_unused_parameters=True) + else: + resolved_strategy = strategy + + trainer = L.Trainer( + max_epochs=max_epochs, + devices=devices, + num_nodes=nodes, + accelerator=accelerator, + strategy=resolved_strategy, + precision=precision, + logger=tb_logger, + callbacks=[monitored_ckpt_callback, last_ckpt_callback, lr_monitor], + log_every_n_steps=native_config.get("log_interval", 10), + enable_progress_bar=True, + ) + + datamodule = StyleTTS2DataModule(native_config, load_for_everyvoice=True) + model = StyleTTS2(native_config, mode=mode.value) + + resume_ckpt = ( + str(tr.finetune_checkpoint) + if tr.finetune_checkpoint and os.path.exists(tr.finetune_checkpoint) + else None + ) + + trainer.fit(model, datamodule=datamodule, ckpt_path=resume_ckpt)