Add checkpoint cleaning option

src/piper/train/__main__.py

@@ -1,14 +1,51 @@
 import logging
+import sys
+from pathlib import Path
 
 import torch
 from lightning.pytorch.cli import LightningCLI
 
+from .checkpoint import clean_checkpoint
 from .vits.dataset import VitsDataModule
 from .vits.lightning import VitsModel
 
 _LOGGER = logging.getLogger(__package__)
 
 
+
+def _process_checkpoint_cleaning():
+    """Clean checkpoint and convert --ckpt_path to --model.pretrained_ckpt.
+
+    When --clean_checkpoint is specified, the checkpoint is treated as pretrained
+    weights only (not a training resume), allowing cross-architecture loading
+    such as single-speaker -> multi-speaker conversion.
+    """
+    sys.argv.remove("--clean_checkpoint")
+    try:
+        idx = sys.argv.index("--ckpt_path")
+        ckpt_path = Path(sys.argv[idx + 1])
+    except (ValueError, IndexError):
+        return
+    if not ckpt_path.exists():
+        return
+
+    num_speakers = 1
+    try:
+        ns_idx = sys.argv.index("--model.num_speakers")
+        num_speakers = int(sys.argv[ns_idx + 1])
+    except (ValueError, IndexError):
+        pass
+
+    cleaned_path = clean_checkpoint(ckpt_path, num_speakers)
+
+    # Replace --ckpt_path with --model.vocoder_warmstart_ckpt so Lightning loads
+    # weights only instead of resuming training state
+    sys.argv[idx] = "--model.vocoder_warmstart_ckpt"
+    sys.argv[idx + 1] = cleaned_path
+    sys.argv.append("--model.warmstart_pretrained")
+    sys.argv.append("true")
+
+
 class VitsLightningCLI(LightningCLI):
     def add_arguments_to_parser(self, parser):
         parser.link_arguments("data.batch_size", "model.batch_size")
@@ -23,9 +60,14 @@ def add_arguments_to_parser(self, parser):
 
 def main():
     logging.basicConfig(level=logging.INFO)
+
+    if "--clean_checkpoint" in sys.argv:
+        _process_checkpoint_cleaning()
+
     torch.backends.cuda.matmul.allow_tf32 = True
     torch.backends.cudnn.allow_tf32 = True
     torch.backends.cudnn.deterministic = False
+
     _cli = VitsLightningCLI(  # noqa: ignore=F841
         VitsModel, VitsDataModule, trainer_defaults={"max_epochs": -1}
     )

src/piper/train/checkpoint.py

@@ -0,0 +1,60 @@
+"""Checkpoint utilities for Piper training."""
+
+import inspect
+import logging
+import tempfile
+import torch
+
+_LOGGER = logging.getLogger(__name__)
+
+
+def clean_checkpoint(checkpoint_path, target_num_speakers):
+    """
+    Clean and prepare a checkpoint for training. 
+     - Removes invalid parameters from older training
+     - Updates single speaker checkpoints for multi speaker
+    """
+    from .vits.lightning import VitsModel
+
+    checkpoint = torch.load(checkpoint_path, weights_only=False, map_location="cpu")
+    if "hyper_parameters" not in checkpoint:
+        return checkpoint_path
+
+    ckpt_updated = False
+    hparams = checkpoint["hyper_parameters"]
+
+    # Remove invalid params
+    valid_params = set(inspect.signature(VitsModel.__init__).parameters.keys())
+    invalid_params = set(hparams.keys()) - valid_params
+
+    for param in invalid_params:
+        _LOGGER.info(f"Removing invalid parameter '{param}' from checkpoint")
+        del hparams[param]
+        ckpt_updated = True
+
+    # Prepare single-speaker checkpoint for multi-speaker training
+    if target_num_speakers > 1:
+        ckpt_num_speakers = hparams.get("num_speakers", 1)
+        has_speaker_emb = any(
+            k.startswith("model_g.emb_g.") for k in checkpoint.get("state_dict", {})
+        )
+
+        if ckpt_num_speakers == 1 and not has_speaker_emb:
+            _LOGGER.info(
+                f"Preparing single-speaker checkpoint for {target_num_speakers}-speaker training"
+            )
+            keys_to_keep = {"state_dict", "hyper_parameters", "pytorch-lightning_version"}
+            keys_to_remove = [k for k in checkpoint.keys() if k not in keys_to_keep]
+            for key in keys_to_remove:
+                del checkpoint[key]
+            ckpt_updated = True
+
+    if not ckpt_updated:
+        return checkpoint_path
+
+    with tempfile.NamedTemporaryFile(suffix=".ckpt", delete=False) as f:
+        torch.save(checkpoint, f.name)
+        temp_path = f.name
+
+    _LOGGER.info(f"Created cleaned checkpoint: {temp_path}")
+    return temp_path

src/piper/train/vits/lightning.py

@@ -72,6 +72,7 @@ def __init__(
         c_kl: float = 1.0,
         grad_clip: Optional[float] = None,
         vocoder_warmstart_ckpt: Optional[str] = None,
+        warmstart_pretrained: bool = False,
         # unused
         dataset: object = None,
         **kwargs,
@@ -116,6 +117,7 @@ def __init__(
         # Used to partially load the state dict from a checkpoint.
         # Only the text/phoneme agnostic portions are loaded.
         self._vocoder_warmstart_ckpt = vocoder_warmstart_ckpt
+        self._warmstart_pretrained = warmstart_pretrained
 
         # Set up models
         self.model_g = SynthesizerTrn(
@@ -315,7 +317,7 @@ def configure_optimizers(self):
         ]
 
         return optimizers, schedulers
-
+        
     def _warmstart_vocoder_from_ckpt(self, ckpt_path: str):
         ckpt = torch.load(ckpt_path, map_location=self.device, weights_only=False)
 
@@ -334,19 +336,25 @@ def _warmstart_vocoder_from_ckpt(self, ckpt_path: str):
             if not k.startswith(KEEP_PREFIXES):
                 continue
             if (k in new_sd) and (new_sd[k].shape == v.shape):
-                new_sd[k] = v
-                copied += 1
+                    new_sd[k] = v
+                    copied += 1
 
-        self.load_state_dict(new_sd, strict=False)
         _LOGGER.info(f"[warmstart] Copied {copied} vocoder parameters from {ckpt_path}")
-
+        return new_sd
+
     def on_fit_start(self):
         # Called once at the start of fit()
         if self._vocoder_warmstart_ckpt is None:
             return
-
-        # Make sure we’re on the correct device
-        self._warmstart_vocoder_from_ckpt(self._vocoder_warmstart_ckpt)
-
+
+        ckpt = torch.load(self._vocoder_warmstart_ckpt, map_location=self.device, weights_only=False)
+        if self._warmstart_pretrained:
+            new_sd = ckpt["state_dict"]
+            _LOGGER.info(f"Loaded pretrained weights from {self._vocoder_warmstart_ckpt}")
+        else:
+            # Make sure we’re on the correct device
+            new_sd = self._warmstart_vocoder_from_ckpt(self._vocoder_warmstart_ckpt)
+
+        self.load_state_dict(new_sd, strict=False)
         # Avoid re-running if Trainer restarts
         self._vocoder_warmstart_ckpt = None