Refactor train CLI into core, with CLI and API shims

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,
+    )

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

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:
 

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)

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",
+]