Checkpoint Parameter Removal Support

docs/TRAINING.md

@@ -72,7 +72,42 @@ where:
 
 Using `--ckpt_path` is recommended since it will speed up training a lot, even if the checkpoint is from a different language. Only `medium` quality checkpoints are supported without [tweaking other settings][audio-config].
 
-Run `python3 -m piper.train fit --help` for many more options.
+Run `python3 -m piper.train fit --help` to view all options. Alternatively, they are provided here:
+
+```txt
+usage: __main__.py [options] fit [-h] [-c CONFIG] [--print_config[=flags]] [--seed_everything SEED_EVERYTHING] [--trainer CONFIG] [--trainer.accelerator.help CLASS_PATH_OR_NAME]
+
+      [--trainer.accelerator ACCELERATOR] [--trainer.strategy.help CLASS_PATH_OR_NAME] [--trainer.strategy STRATEGY] [--trainer.devices DEVICES]
+      [--trainer.num_nodes NUM_NODES] [--trainer.precision PRECISION] [--trainer.logger.help CLASS_PATH_OR_NAME] [--trainer.logger LOGGER]
+      [--trainer.callbacks.help CLASS_PATH_OR_NAME] [--trainer.callbacks CALLBACKS] [--trainer.fast_dev_run FAST_DEV_RUN] [--trainer.max_epochs MAX_EPOCHS]
+      [--trainer.min_epochs MIN_EPOCHS] [--trainer.max_steps MAX_STEPS] [--trainer.min_steps MIN_STEPS] [--trainer.max_time MAX_TIME]
+      [--trainer.limit_train_batches LIMIT_TRAIN_BATCHES] [--trainer.limit_val_batches LIMIT_VAL_BATCHES] [--trainer.limit_test_batches LIMIT_TEST_BATCHES]
+      [--trainer.limit_predict_batches LIMIT_PREDICT_BATCHES] [--trainer.overfit_batches OVERFIT_BATCHES] [--trainer.val_check_interval VAL_CHECK_INTERVAL]
+      [--trainer.check_val_every_n_epoch CHECK_VAL_EVERY_N_EPOCH] [--trainer.num_sanity_val_steps NUM_SANITY_VAL_STEPS]
+      [--trainer.log_every_n_steps LOG_EVERY_N_STEPS] [--trainer.enable_checkpointing {true,false,null}] [--trainer.enable_progress_bar {true,false,null}]
+      [--trainer.enable_model_summary {true,false,null}] [--trainer.accumulate_grad_batches ACCUMULATE_GRAD_BATCHES]
+      [--trainer.gradient_clip_val GRADIENT_CLIP_VAL] [--trainer.gradient_clip_algorithm GRADIENT_CLIP_ALGORITHM] [--trainer.deterministic DETERMINISTIC]
+      [--trainer.benchmark {true,false,null}] [--trainer.inference_mode {true,false}] [--trainer.use_distributed_sampler {true,false}]
+      [--trainer.profiler.help CLASS_PATH_OR_NAME] [--trainer.profiler PROFILER] [--trainer.detect_anomaly {true,false}] [--trainer.barebones {true,false}]
+      [--trainer.plugins.help CLASS_PATH_OR_NAME] [--trainer.plugins PLUGINS] [--trainer.sync_batchnorm {true,false}]
+      [--trainer.reload_dataloaders_every_n_epochs RELOAD_DATALOADERS_EVERY_N_EPOCHS] [--trainer.default_root_dir DEFAULT_ROOT_DIR]
+      [--trainer.model_registry MODEL_REGISTRY] [--model CONFIG] [--model.sample_rate SAMPLE_RATE] [--model.num_speakers NUM_SPEAKERS] [--model.resblock RESBLOCK]
+      [--model.resblock_kernel_sizes RESBLOCK_KERNEL_SIZES] [--model.resblock_dilation_sizes RESBLOCK_DILATION_SIZES] [--model.upsample_rates UPSAMPLE_RATES]
+      [--model.upsample_initial_channel UPSAMPLE_INITIAL_CHANNEL] [--model.upsample_kernel_sizes UPSAMPLE_KERNEL_SIZES] [--model.filter_length FILTER_LENGTH]
+      [--model.hop_length HOP_LENGTH] [--model.win_length WIN_LENGTH] [--model.mel_channels MEL_CHANNELS] [--model.mel_fmin MEL_FMIN] [--model.mel_fmax MEL_FMAX]
+      [--model.inter_channels INTER_CHANNELS] [--model.hidden_channels HIDDEN_CHANNELS] [--model.filter_channels FILTER_CHANNELS] [--model.n_heads N_HEADS]
+      [--model.n_layers N_LAYERS] [--model.kernel_size KERNEL_SIZE] [--model.p_dropout P_DROPOUT] [--model.n_layers_q N_LAYERS_Q]
+      [--model.use_spectral_norm {true,false}] [--model.gin_channels GIN_CHANNELS] [--model.use_sdp {true,false}] [--model.segment_size SEGMENT_SIZE]
+      [--model.learning_rate LEARNING_RATE] [--model.learning_rate_d LEARNING_RATE_D] [--model.betas [ITEM,...]] [--model.betas_d [ITEM,...]] [--model.eps EPS]
+      [--model.lr_decay LR_DECAY] [--model.lr_decay_d LR_DECAY_D] [--model.init_lr_ratio INIT_LR_RATIO] [--model.warmup_epochs WARMUP_EPOCHS] [--model.c_mel C_MEL]
+      [--model.c_kl C_KL] [--model.grad_clip GRAD_CLIP] [--model.dataset.help [CLASS_PATH_OR_NAME]] [--model.dataset DATASET] [--data CONFIG]
+      --data.csv_path CSV_PATH --data.cache_dir CACHE_DIR --data.espeak_voice ESPEAK_VOICE --data.config_path CONFIG_PATH --data.voice_name VOICE_NAME
+      [--data.audio_dir AUDIO_DIR] [--data.alignments_dir ALIGNMENTS_DIR] [--data.num_symbols NUM_SYMBOLS] [--data.batch_size BATCH_SIZE]
+      [--data.validation_split VALIDATION_SPLIT] [--data.num_test_examples NUM_TEST_EXAMPLES] [--data.num_workers NUM_WORKERS] [--data.trim_silence {true,false}]
+      [--data.keep_seconds_before_silence KEEP_SECONDS_BEFORE_SILENCE] [--data.keep_seconds_after_silence KEEP_SECONDS_AFTER_SILENCE]
+      [--optimizer.help [CLASS_PATH_OR_NAME]] [--optimizer CONFIG | CLASS_PATH_OR_NAME | .INIT_ARG_NAME VALUE] [--lr_scheduler.help CLASS_PATH_OR_NAME]
+      [--lr_scheduler CONFIG | CLASS_PATH_OR_NAME | .INIT_ARG_NAME VALUE] [--ckpt_path CKPT_PATH]
+```
 
 ## Exporting
 

src/piper/train/__main__.py

@@ -1,5 +1,9 @@
 import logging
 
+import inspect
+import tempfile
+import sys
+
 import torch
 from lightning.pytorch.cli import LightningCLI
 
@@ -21,11 +25,58 @@ def add_arguments_to_parser(self, parser):
         parser.link_arguments("model.segment_size", "data.segment_size")
 
 
+def clean_checkpoint(checkpoint_path):
+    checkpoint = torch.load(checkpoint_path, weights_only=False, map_location="cpu")
+
+    if "hyper_parameters" not in checkpoint:
+        return checkpoint_path
+
+    init_signature = inspect.signature(VitsModel.__init__)
+    valid_params = set(init_signature.parameters.keys())
+    checkpoint_params = set(checkpoint["hyper_parameters"].keys())
+    invalid_params = checkpoint_params - valid_params
+
+    if not invalid_params:
+        return checkpoint_path
+
+    for param in invalid_params:
+        _LOGGER.info(f"Removing invalid parameter '{param}' from checkpoint")
+        del checkpoint["hyper_parameters"][param]
+
+    temp_file = tempfile.NamedTemporaryFile(suffix=".ckpt", delete=False)
+    torch.save(checkpoint, temp_file.name)
+    temp_file.close()
+
+    _LOGGER.info(f"Created cleaned checkpoint: {temp_file.name}")
+    return temp_file.name
+
 def main():
+
     logging.basicConfig(level=logging.INFO)
     torch.backends.cuda.matmul.allow_tf32 = True
     torch.backends.cudnn.allow_tf32 = True
     torch.backends.cudnn.deterministic = False
+
+    ckpt_path = None
+    if '--ckpt_path' in sys.argv:
+        try:
+            ckpt_idx = sys.argv.index('--ckpt_path')
+            if ckpt_idx + 1 < len(sys.argv):
+                ckpt_path = sys.argv[ckpt_idx + 1]
+        except (IndexError, ValueError):
+            pass
+
+    if ckpt_path:
+        cleaned_ckpt_path = clean_checkpoint(ckpt_path)
+        # issue a ckpt path replacement to fixed ckpt path
+        if '--ckpt_path' in sys.argv:
+            ckpt_idx = sys.argv.index('--ckpt_path')
+            if ckpt_idx + 1 < len(sys.argv):
+                sys.argv[ckpt_idx + 1] = cleaned_ckpt_path
+
+    else:
+        _LOGGER.info("No checkpoint path provided; skipping checkpoint cleaning.")
+
     _cli = VitsLightningCLI(  # noqa: ignore=F841
         VitsModel, VitsDataModule, trainer_defaults={"max_epochs": -1}
     )