179 lines
7.6 KiB
Python
179 lines
7.6 KiB
Python
"""Voice training pipeline: warm start/resume, piper compatibility fixes, data filters."""
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest.mock import patch
|
|
|
|
import numpy as np
|
|
|
|
HIGH_ARCH = {
|
|
"resblock": "1", "resblock_kernel_sizes": (3, 7, 11),
|
|
"resblock_dilation_sizes": ((1, 3, 5), (1, 3, 5), (1, 3, 5)),
|
|
"upsample_rates": (8, 8, 2, 2), "upsample_initial_channel": 512,
|
|
"upsample_kernel_sizes": (16, 16, 4, 4),
|
|
}
|
|
|
|
|
|
def _voice_dir(root: Path) -> Path:
|
|
training = root / "v" / "training"
|
|
(training / "wavs").mkdir(parents=True)
|
|
(training / "metadata.csv").write_text("clip_00000|hallo|hallo\n", encoding="utf-8")
|
|
return root / "v"
|
|
|
|
|
|
class TrainingStartTests(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
try:
|
|
import torch # noqa: F401
|
|
except ImportError:
|
|
self.skipTest("torch not installed")
|
|
from javis.voices import piper_train
|
|
self.pt = piper_train
|
|
self.cmds: list[list[str]] = []
|
|
|
|
def fake_run(cmd, **kw):
|
|
self.cmds.append(cmd)
|
|
return 1 # stop after stage 1
|
|
|
|
patches = [patch.object(piper_train, "_run", fake_run),
|
|
patch.object(piper_train, "MIN_FREE_BYTES", 0)]
|
|
for p in patches:
|
|
p.start()
|
|
self.addCleanup(p.stop)
|
|
|
|
def _ckpt(self, path: Path, **hp) -> Path:
|
|
import torch
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
torch.save({"state_dict": {"w": torch.zeros(1)}, "hyper_parameters": hp}, path)
|
|
return path
|
|
|
|
def test_first_run_warm_starts_from_base_with_its_architecture(self) -> None:
|
|
import yaml
|
|
with tempfile.TemporaryDirectory() as d:
|
|
voice = _voice_dir(Path(d))
|
|
self._ckpt(voice / "training" / "piper_base" / "epoch=2665-step=1.ckpt",
|
|
dataset=[Path("x")], **HIGH_ARCH)
|
|
self.pt.run_pipeline(voice)
|
|
cmd = self.cmds[0]
|
|
self.assertIn("javis.voices.piper_compat", cmd)
|
|
self.assertIn("--model.warmstart_ckpt", cmd)
|
|
self.assertNotIn("--ckpt_path", cmd)
|
|
arch = yaml.safe_load(Path(cmd[cmd.index("--config") + 1]).read_text(encoding="utf-8"))
|
|
self.assertEqual(arch["model"]["resblock"], "1")
|
|
self.assertEqual(arch["model"]["upsample_rates"], [8, 8, 2, 2])
|
|
self.assertFalse((voice / "training" / "warmstart.pt").exists()) # removed after fit
|
|
|
|
def test_later_runs_resume_own_checkpoint_and_ignore_base(self) -> None:
|
|
with tempfile.TemporaryDirectory() as d:
|
|
voice = _voice_dir(Path(d))
|
|
base = self._ckpt(voice / "training" / "piper_base" / "epoch=1.ckpt", **HIGH_ARCH)
|
|
last = self._ckpt(voice / "training" / "piper_lightning" / "lightning_logs"
|
|
/ "version_0" / "checkpoints" / "last.ckpt", **HIGH_ARCH)
|
|
self.pt.run_pipeline(voice, base_checkpoint=base)
|
|
cmd = self.cmds[0]
|
|
self.assertEqual(cmd[cmd.index("--ckpt_path") + 1], str(last))
|
|
self.assertNotIn("--model.warmstart_ckpt", cmd)
|
|
|
|
def test_linux_checkpoint_with_posixpath_loads(self) -> None:
|
|
import pathlib
|
|
import torch
|
|
|
|
class FakePosixPath: # pickled as pathlib.PosixPath, like a Linux checkpoint
|
|
def __init__(self, p):
|
|
self.p = p
|
|
|
|
def __reduce__(self):
|
|
return (FakePosixPath, (self.p,))
|
|
|
|
FakePosixPath.__module__ = "pathlib"
|
|
FakePosixPath.__name__ = FakePosixPath.__qualname__ = "PosixPath"
|
|
with tempfile.TemporaryDirectory() as d:
|
|
p = Path(d) / "linux.ckpt"
|
|
with patch.object(pathlib, "PosixPath", FakePosixPath):
|
|
torch.save({"hyper_parameters": {"dataset": [FakePosixPath("/ssd/x")]}}, p)
|
|
ckpt = self.pt._load_checkpoint(p)
|
|
self.assertEqual(str(ckpt["hyper_parameters"]["dataset"][0]), "/ssd/x")
|
|
|
|
|
|
class PiperCompatTests(unittest.TestCase):
|
|
def test_int_resblock_builds_high_quality_blocks(self) -> None:
|
|
try:
|
|
import javis.voices.piper_compat # noqa: F401 — applies the fixes
|
|
from piper.train.vits import models, modules
|
|
from pysilero_vad import SileroVoiceActivityDetector
|
|
except ImportError:
|
|
self.skipTest("piper training extras not installed")
|
|
g = models.Generator(192, 1, (3, 7, 11), ((1, 3, 5),) * 3, (8, 8, 2, 2), 64, (16, 16, 4, 4))
|
|
self.assertIsInstance(g.resblocks[0], modules.ResBlock1)
|
|
self.assertTrue(hasattr(SileroVoiceActivityDetector, "process_array"))
|
|
|
|
|
|
class TranscriptGateTests(unittest.TestCase):
|
|
def test_unreliable_transcripts_are_rejected(self) -> None:
|
|
from javis.voices.audiobook_ingest import _transcript_ok
|
|
|
|
def tr(**kw):
|
|
return SimpleNamespace(**{"avg_logprob": -0.2, "no_speech_prob": 0.01,
|
|
"compression_ratio": 1.4, **kw})
|
|
self.assertTrue(_transcript_ok(tr()))
|
|
self.assertTrue(_transcript_ok(tr(avg_logprob=None, no_speech_prob=None,
|
|
compression_ratio=None)))
|
|
self.assertFalse(_transcript_ok(tr(avg_logprob=-1.3)))
|
|
self.assertFalse(_transcript_ok(tr(no_speech_prob=0.9)))
|
|
self.assertFalse(_transcript_ok(tr(compression_ratio=3.1)))
|
|
|
|
|
|
class NarratorFilterTests(unittest.TestCase):
|
|
def test_dominant_speaker_scores_high(self) -> None:
|
|
from javis.voices.dataset_clean import narrator_scores
|
|
|
|
rng = np.random.default_rng(0)
|
|
a, b = np.eye(8)[0], np.eye(8)[1]
|
|
embs = np.array([v + rng.normal(0, 0.1, 8) for v in [a] * 70 + [b] * 30])
|
|
embs /= np.linalg.norm(embs, axis=1, keepdims=True)
|
|
s = narrator_scores(embs)
|
|
self.assertGreater(s[:70].min(), 0.8)
|
|
self.assertLess(s[70:].max(), 0.4)
|
|
|
|
def test_clean_dataset_filters_and_keeps_new_imports(self) -> None:
|
|
import soundfile as sf
|
|
from javis.voices.dataset_clean import clean_dataset
|
|
|
|
def embed(audio, sr):
|
|
# The first sample marks the speaker: 0.5 = narrator, else character voice.
|
|
return np.eye(4)[0] if audio[0] > 0.4 else np.eye(4)[1]
|
|
|
|
with tempfile.TemporaryDirectory() as d:
|
|
training = Path(d)
|
|
(training / "wavs").mkdir()
|
|
|
|
def add(stem: str, narrator: bool) -> str:
|
|
sf.write(str(training / "wavs" / f"{stem}.wav"),
|
|
np.full(100, 0.5 if narrator else 0.1, dtype=np.float32), 22050)
|
|
return f"{stem}|text {stem}|text {stem}\n"
|
|
|
|
lines = [add(f"clip_{i:05d}", narrator=i % 4 != 0) for i in range(8)]
|
|
(training / "metadata.csv").write_text("".join(lines), encoding="utf-8")
|
|
|
|
report = clean_dataset(training, embed, dry_run=True)
|
|
self.assertEqual((report["kept"], report["dropped"]), (6, 2))
|
|
self.assertEqual(len((training / "metadata.csv").read_text().splitlines()), 8)
|
|
|
|
clean_dataset(training, embed)
|
|
self.assertEqual(len((training / "metadata.csv").read_text().splitlines()), 6)
|
|
self.assertEqual(len((training / "metadata.unfiltered.csv").read_text().splitlines()), 8)
|
|
|
|
# A later import appends to metadata.csv; re-filtering must keep it.
|
|
with open(training / "metadata.csv", "a", encoding="utf-8") as f:
|
|
f.write(add("clip_00100", narrator=True))
|
|
report = clean_dataset(training, embed)
|
|
self.assertEqual((report["clips"], report["kept"]), (9, 7))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(unittest.main())
|