drewthomasson--ebook2audiobook
wehub-resource-sync
99 行
3.7 KiB
Python
99 行
3.7 KiB
Python
import sys
|
|
import json
|
|
from pathlib import Path
|
|
import shutil
|
|
|
|
# Make sure we can import from tools/Universal_TTS_Finetune
|
|
_DIR = Path(__file__).resolve().parent
|
|
sys.path.insert(0, str(_DIR))
|
|
|
|
from utils.piper_utils import export_piper_onnx
|
|
from utils.pipeline import _latest_matching_file, _finalize_training_artifacts
|
|
|
|
def main():
|
|
if len(sys.argv) < 2:
|
|
print("Usage: python export_checkpoint.py <run_dir_path>")
|
|
sys.exit(1)
|
|
|
|
run_dir = Path(sys.argv[1]).resolve()
|
|
if not run_dir.exists():
|
|
print(f"Error: Run directory {run_dir} does not exist.")
|
|
sys.exit(1)
|
|
|
|
# Check what kind of files are inside to auto-detect model type
|
|
# Piper uses Lightning (.ckpt), Coqui uses standard PyTorch (.pth)
|
|
has_ckpt = list(run_dir.glob("**/*.ckpt")) or list(run_dir.glob("*.ckpt"))
|
|
has_pth = list(run_dir.glob("**/*.pth")) or list(run_dir.glob("*.pth"))
|
|
|
|
# Try to infer model key from folder name: training_runs/<model_key>/<timestamp>
|
|
model_key = run_dir.parent.name
|
|
|
|
if has_ckpt and (not has_pth or model_key == "piper"):
|
|
print("Detected model family: Piper (ONNX format)")
|
|
|
|
preprocessed_dir = run_dir / "preprocessed"
|
|
config_path = preprocessed_dir / "config.json"
|
|
if not config_path.exists():
|
|
print(f"Error: config.json not found under {preprocessed_dir}")
|
|
sys.exit(1)
|
|
|
|
lightning_logs_dir = preprocessed_dir / "lightning_logs"
|
|
trained_ckpt = _latest_matching_file(lightning_logs_dir, ["**/*.ckpt", "*.ckpt"])
|
|
if not trained_ckpt:
|
|
trained_ckpt = _latest_matching_file(run_dir, ["**/*.ckpt", "*.ckpt"])
|
|
|
|
if not trained_ckpt:
|
|
print(f"Error: No .ckpt file found in {run_dir}")
|
|
sys.exit(1)
|
|
|
|
print(f"Using checkpoint: {trained_ckpt}")
|
|
|
|
ready_dir = run_dir / "ready"
|
|
ready_dir.mkdir(parents=True, exist_ok=True)
|
|
ready_onnx = ready_dir / "model.onnx"
|
|
|
|
print("Exporting to ONNX...")
|
|
export_piper_onnx(trained_ckpt, ready_onnx, config_path)
|
|
|
|
log_path = run_dir / "training.log"
|
|
artifacts = {
|
|
"model_key": "piper",
|
|
"model_label": "Piper TTS",
|
|
"family": "piper",
|
|
"training_root": str(run_dir),
|
|
"dataset_dir": "",
|
|
"checkpoint": str(ready_onnx),
|
|
"config": str(ready_onnx) + ".json",
|
|
"reference_wav": "",
|
|
"log_path": str(log_path) if log_path.exists() else "",
|
|
"unused_overrides": {},
|
|
}
|
|
|
|
artifacts_path = ready_dir / "artifacts.json"
|
|
artifacts_path.write_text(json.dumps(artifacts, indent=2), encoding="utf-8")
|
|
print("Successfully exported model and wrote artifacts.json!")
|
|
|
|
elif has_pth:
|
|
print(f"Detected model family: Coqui PTH (Model key inferred: {model_key})")
|
|
|
|
# Use pipeline's built-in finalizer for Coqui models (which handles XTTS optimizations too)
|
|
try:
|
|
artifacts = _finalize_training_artifacts(
|
|
spec_key=model_key,
|
|
training_root=run_dir,
|
|
dataset_dir=Path(""),
|
|
reference_wav=""
|
|
)
|
|
print("Successfully packaged Coqui checkpoint and wrote artifacts.json!")
|
|
except Exception as e:
|
|
print(f"Error finalizing Coqui artifacts: {e}")
|
|
sys.exit(1)
|
|
else:
|
|
print("Error: Could not find any .ckpt or .pth files to export.")
|
|
sys.exit(1)
|
|
|
|
print(f"You can now test this model in the GUI or via: python headless_cli.py synthesize --artifacts {run_dir}/ready --text '...' --model {model_key}")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|