diff --git a/.javis/memory.json b/.javis/memory.json new file mode 100644 index 0000000..da29c81 --- /dev/null +++ b/.javis/memory.json @@ -0,0 +1,10 @@ +{ + "entries": { + "spotify_pause_access_denied_unknown_speaker": { + "value": "spotify__pausePlayback mit deviceId 'a169794694f72b88e60439a45e2b506a30254979' faellt mit 'Access denied: unknown speaker' durch. Bei Spotify-Befehlern zuerst die Integration laden oder gueltige Geraete-IDs ueber getAvailableDevices ermitteln.", + "category": "spotify", + "created_at": "2026-10-10T10:43:30.486067+00:00", + "updated_at": "2026-10-10T10:43:30.486067+00:00" + } + } +} \ No newline at end of file diff --git a/.javis/memory_florian.json b/.javis/memory_florian.json new file mode 100644 index 0000000..1a2144d --- /dev/null +++ b/.javis/memory_florian.json @@ -0,0 +1,10 @@ +{ + "entries": { + "user_name": { + "value": "Der Nutzer heisst Florian.", + "category": "personal", + "created_at": "2026-10-10T10:40:14.388969+00:00", + "updated_at": "2026-10-10T10:40:14.388969+00:00" + } + } +} \ No newline at end of file diff --git a/.javis/quick_commands.json b/.javis/quick_commands.json index e2f8a11..a7d4b05 100644 --- a/.javis/quick_commands.json +++ b/.javis/quick_commands.json @@ -1,4 +1,22 @@ { "version": 1, - "commands": {} + "commands": { + "abfbc899-8ad1-436b-b386-152c80669dcd": { + "id": "abfbc899-8ad1-436b-b386-152c80669dcd", + "trigger": "küche ein", + "match_mode": "contains", + "actions": [ + { + "tool": "homey__control_zone", + "args": { + "capability": "onoff", + "value": true, + "zone": "Küche" + } + } + ], + "created_at": "2026-10-10T10:39:38.382223+00:00", + "updated_at": "2026-10-10T10:39:38.382223+00:00" + } + } } \ No newline at end of file diff --git a/.javis/slaves.json b/.javis/slaves.json new file mode 100644 index 0000000..5c71b98 --- /dev/null +++ b/.javis/slaves.json @@ -0,0 +1,13 @@ +{ + "devices": { + "a32f0630-04e5-4eda-9ae2-44fa4eeb6068": { + "name": "Tablet K\u00fcche", + "api_key_hash": "89db0c92e376e0160a756f8285d940ba8813d2867cd699badac4deb2b85b813b", + "room": "K\u00fcche", + "ip_address": "10.10.10.225", + "created_at": "2026-10-07T17:40:24+00:00", + "last_seen": "2026-10-10T10:44:31+00:00", + "last_seen_ip": "10.10.10.225" + } + } +} \ No newline at end of file diff --git a/.javis/voice_profiles/Florian.npy b/.javis/voice_profiles/Florian.npy new file mode 100644 index 0000000..b71f341 Binary files /dev/null and b/.javis/voice_profiles/Florian.npy differ diff --git a/.javis/voice_profiles/index.json b/.javis/voice_profiles/index.json new file mode 100644 index 0000000..c8cfb9e --- /dev/null +++ b/.javis/voice_profiles/index.json @@ -0,0 +1,9 @@ +[ + { + "name": "Florian", + "file": "Florian.npy", + "enrolled_at": "2026-10-07T21:32:39", + "clip_count": 13, + "last_seen": "2026-10-10T12:44:06" + } +] \ No newline at end of file diff --git a/.javis/voice_profiles/unidentified/index.json b/.javis/voice_profiles/unidentified/index.json new file mode 100644 index 0000000..0637a08 --- /dev/null +++ b/.javis/voice_profiles/unidentified/index.json @@ -0,0 +1 @@ +[] \ No newline at end of file diff --git a/README.md b/README.md index cf9c1c0..2d21860 100644 --- a/README.md +++ b/README.md @@ -89,6 +89,33 @@ Push-to-talk: `Ctrl+Space` in the window, `Ctrl+Alt+Space` system-wide (Windows) 3. Javis will VAD-split, transcribe, and build a clean dataset. 4. Test with zero-shot cloning (instant) or click **Fine-tune** for higher quality. +Fine-tuning runs in a console window of its own. Meanwhile Javis speaks with the +voice being trained, refreshed with the training progress every 15 minutes. Stop it +any time — Ctrl+C, closing the window, logging off or shutting down saves the +epoch so far, and the next run continues with its next batch. A crash loses at most +the last 10 minutes (`JAVIS_PIPER_PROGRESS_MINUTES`); training restarts by itself. + +## Speaker identification + +With `speaker_id.enabled`, Javis recognises who is talking (ECAPA embeddings) and manages +speakers in **Intelligence → Voice Profiles**: + +- **Enroll** — `+` records a few sentences and creates a profile. +- **Unidentified speakers** — utterances that matched no profile are kept (audio, + transcript, time, source device, closest profile). **▶ Play** listens to one; + **Sort…** adds it to an existing profile or creates a new one. Recordings that sound + alike are offered and pre-selected so you can sort several at once. ✕ / **Clear all** + delete them. +- **Enhance a voice** — 🎤 on a profile card opens a recorder: click **● Record sample**, + speak, review each sample (playback + how well it matches the profile; mismatches are + flagged), then **Add N samples** to blend them in. +- The "unknown speaker" pop-up (`speaker_id.auto_enroll`) can also play back what was said. + +Unknown-speaker clips are stored locally only, in `/unidentified/`. +They are capped by `speaker_id.max_unidentified` (default 50, oldest dropped first). Turn +the inbox off with `speaker_id.keep_unidentified: false`. Both settings are under +**Settings → Speaker Identification**. + ## MCP servers Edit `mcp_servers.json`. Schema matches VS Code's `mcp.json`, so you can copy/paste server entries. diff --git a/config.example.yaml b/config.example.yaml index b554a58..d13bfad 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -40,7 +40,7 @@ training: batch_size: 32 num_workers: 2 accelerator: "gpu" # gpu | cpu - precision: "16-mixed" # 32 | 16-mixed | bf16-mixed + precision: "bf16-mixed" # bf16-mixed (fastest) | 32 | 16-mixed language: "de" # espeak-ng voice code used for phonemisation (de, en-us, …) audio: diff --git a/docs/AUDIT.md b/docs/AUDIT.md index 562b907..d56381d 100644 --- a/docs/AUDIT.md +++ b/docs/AUDIT.md @@ -494,9 +494,12 @@ The overnight scheduler runs `run_pipeline()` via `asyncio.to_thread`; the threa on the training subprocess for hours. Cancelling the scheduler task (quit) released the training lock immediately while the thread kept training, and at interpreter exit Python joins executor threads — so Javis hung on quit until training finished. -**Fix:** cancellation calls `piper_train.request_stop()`, which terminates the child -process (`_run` now polls instead of blocking in `wait()`); the lock is only released -once the worker has actually finished. +**Fix:** cancellation calls `piper_train.request_stop()`: the training process finishes +its step, saves its progress — within the epoch too: the next run continues with the +next batch — and exits (it is killed if that takes longer than +`STOP_GRACE_SECONDS`), and the scheduler waits for that (`_run` polls instead of +blocking in `wait()`). The lock (`training/.training.lock`, held by `run_pipeline`) +names the training processes and is only released once they have exited. #### VOI-01 **Imported voice archives may carry pickle checkpoints** — *Medium (trust)* diff --git a/javis-slave/README.md b/javis-slave/README.md index b02e5a0..67f391a 100644 --- a/javis-slave/README.md +++ b/javis-slave/README.md @@ -37,6 +37,10 @@ screen (GUI mode) and played back through the device's speakers. If the master's turn fails it returns its spoken error phrase instead, so a screenless satellite still gets audible feedback; connection problems play a low falling tone. +5. **Conversation view** — The GUI lists the conversation with the newest + message on top; scroll down for older ones. The fullscreen button in the + conversation card (or `F11`) toggles window fullscreen; the orb and the + conversation stay visible. `Esc` leaves fullscreen. ## Setup diff --git a/javis-slave/src/javis_slave/main_window.py b/javis-slave/src/javis_slave/main_window.py index 9e81dbc..6d99e1b 100644 --- a/javis-slave/src/javis_slave/main_window.py +++ b/javis-slave/src/javis_slave/main_window.py @@ -132,6 +132,7 @@ QPushButton[variant="danger"] {{ color: {CL_ERROR_FG}; }} QPushButton[variant="danger"]:hover {{ background: {CL_ERROR_BG}; border-color: {CL_ERROR_BORDER}; }} QPushButton[variant="icon"] {{ padding: 0; min-height: 0; border-radius: 16px; }} QPushButton[variant="icon"]:checked {{ background: {CL_ERROR}; border-color: {CL_ERROR}; }} +QPushButton[variant="icon"][small="true"] {{ min-width: 46px; min-height: 46px; }} QScrollArea {{ background: transparent; border: 0; }} QScrollBar:vertical {{ background: transparent; width: 8px; margin: 4px 0; }} @@ -211,9 +212,28 @@ def _draw_tune(p: QtGui.QPainter, box: QtCore.QRectF, color: QtGui.QColor) -> No p.drawEllipse(QtCore.QPointF(x + knob * s, y + row * s), 0.09 * s, 0.09 * s) +def _draw_expand(p: QtGui.QPainter, box: QtCore.QRectF, color: QtGui.QColor, collapse: bool) -> None: + """Four corner brackets: pointing outwards to enter, inwards to leave fullscreen.""" + s, x, y = box.width(), box.x(), box.y() + pen = QtGui.QPen(color, max(1.5, s * 0.075)) + pen.setCapStyle(QtCore.Qt.RoundCap) + pen.setJoinStyle(QtCore.Qt.RoundJoin) + p.setPen(pen) + p.setBrush(QtCore.Qt.NoBrush) + leg = 0.26 * s + reach = 0.10 if collapse else 0.36 # corner distance from the centre + d = 1 if collapse else -1 # legs run away from / towards the centre + for sx in (-1, 1): + for sy in (-1, 1): + vx, vy = x + (0.5 + sx * reach) * s, y + (0.5 + sy * reach) * s + p.drawPolyline([QtCore.QPointF(vx + d * sx * leg, vy), + QtCore.QPointF(vx, vy), + QtCore.QPointF(vx, vy + d * sy * leg)]) + + @lru_cache(maxsize=None) def _icon(kind: str, color: str, size: int = 28) -> QtGui.QIcon: - """Render a crisp (2x) icon: 'mic', 'mic_off' or 'tune'.""" + """Render a crisp (2x) icon: 'mic', 'mic_off', 'tune', 'expand' or 'collapse'.""" pm = QtGui.QPixmap(size * 2, size * 2) pm.setDevicePixelRatio(2.0) pm.fill(QtCore.Qt.transparent) @@ -222,6 +242,8 @@ def _icon(kind: str, color: str, size: int = 28) -> QtGui.QIcon: box = QtCore.QRectF(0, 0, size, size) if kind == "tune": _draw_tune(p, box, QtGui.QColor(color)) + elif kind in ("expand", "collapse"): + _draw_expand(p, box, QtGui.QColor(color), collapse=(kind == "collapse")) else: _draw_mic(p, box, QtGui.QColor(color), slashed=(kind == "mic_off")) p.end() @@ -574,15 +596,18 @@ class MainScreen(QtWidgets.QWidget): Alexa-style UX features: - Mute toggle button on the orb - Connection health monitoring with retry countdown - - Transcript history (scrollable conversation log) + - Transcript history (scrollable conversation log, newest message on top) + - Fullscreen button in the conversation card (orb and chat stay visible) - Visual mute indicator overlay """ settingsRequested = QtCore.Signal() muteToggled = QtCore.Signal(bool) # new signal for mute toggle + fullscreenToggleRequested = QtCore.Signal() def __init__(self, parent=None): super().__init__(parent) + self._fullscreen = False self.setFocusPolicy(QtCore.Qt.ClickFocus) # tapping the background clears focus outer = QtWidgets.QVBoxLayout(self) outer.setContentsMargins(28, 20, 28, 24) @@ -654,10 +679,24 @@ class MainScreen(QtWidgets.QWidget): cv.setContentsMargins(24, 18, 14, 18) cv.setSpacing(10) - # Header label + # Header row: caption + fullscreen-chat toggle + header_row = QtWidgets.QHBoxLayout() + header_row.setContentsMargins(0, 0, 0, 0) history_header = QtWidgets.QLabel("Konversation") history_header.setProperty("role", "caption") - cv.addWidget(history_header) + header_row.addWidget(history_header) + header_row.addStretch(1) + self.expand_btn = QtWidgets.QPushButton() + self.expand_btn.setProperty("variant", "icon") + self.expand_btn.setProperty("small", True) # 48 px: the stylesheet's min-height would shrink it + self.expand_btn.setFixedSize(48, 48) + self.expand_btn.setIconSize(QtCore.QSize(26, 26)) + self.expand_btn.setCursor(QtCore.Qt.PointingHandCursor) + self.expand_btn.setFocusPolicy(QtCore.Qt.TabFocus) + self.expand_btn.clicked.connect(self.fullscreenToggleRequested.emit) + self._sync_expand_button() + header_row.addWidget(self.expand_btn) + cv.addLayout(header_row) # Transcript history (replaces single-turn display). # NOTE: this layout must belong to *exactly one* parent — we set it on @@ -708,6 +747,20 @@ class MainScreen(QtWidgets.QWidget): self.mute_btn.setToolTip(tip) self.mute_btn.setAccessibleName(tip) + def set_fullscreen(self, fullscreen: bool) -> None: + """Mirror the window's fullscreen state in the toggle button.""" + if fullscreen == self._fullscreen: + return + self._fullscreen = fullscreen + self._sync_expand_button() + + def _sync_expand_button(self) -> None: + fs = self._fullscreen + self.expand_btn.setIcon(_icon("collapse" if fs else "expand", CL_TEXT, 26)) + tip = "Vollbild beenden" if fs else "Vollbild" + self.expand_btn.setToolTip(tip) + self.expand_btn.setAccessibleName(tip) + def _bubble_max_width(self) -> int: return max(240, int(self.transcript_scroll.viewport().width() * 0.8)) @@ -767,7 +820,8 @@ class MainScreen(QtWidgets.QWidget): else: h.addWidget(bubble) h.addStretch(1) - self.transcript_layout.addWidget(row) + # Newest first: the latest bubble goes above the older ones. + self.transcript_layout.insertWidget(0, row) # ── slots (called via Qt signals from the engine bridge) ────────────── @@ -796,9 +850,8 @@ class MainScreen(QtWidgets.QWidget): def _ensure_transcript_spacer(self) -> None: """Keep exactly one expanding spacer, at the end of the transcript. - Bubbles are appended after the previous spacer, so stale spacers must - be removed from anywhere in the layout (they used to pile up between - turns and open large gaps). + It pushes short transcripts to the top; stray spacers (e.g. left from + the empty state) would open gaps between bubbles. """ for i in reversed(range(self.transcript_layout.count())): item = self.transcript_layout.itemAt(i) @@ -809,16 +862,16 @@ class MainScreen(QtWidgets.QWidget): self.transcript_layout.addSpacerItem(spacer) def _prune_transcript(self, max_bubbles: int = 40) -> None: - """Keep at most *max_bubbles* bubbles in the transcript layout (oldest first).""" + """Keep at most *max_bubbles* bubbles in the transcript layout (drops the oldest, at the bottom).""" rows = [self.transcript_layout.itemAt(i).widget() for i in range(self.transcript_layout.count())] rows = [row for row in rows if row is not None] - for row in rows[:max(0, len(rows) - max_bubbles)]: + for row in rows[max_bubbles:]: self.transcript_layout.removeWidget(row) row.deleteLater() def set_turn(self, user_text: str, reply_text: str) -> None: - """Add a new conversation turn to the transcript history (Alexa-style).""" + """Add a new conversation turn to the transcript history (newest on top).""" self._add_bubble(user_text, "user") # Assistant message bubble (if any) @@ -826,16 +879,16 @@ class MainScreen(QtWidgets.QWidget): self._add_bubble(reply_text, "assistant") self._prune_transcript() - # Spacer to push content to bottom (like Alexa's conversation history) + # Keeps short transcripts hugging the top edge self._ensure_transcript_spacer() - # Auto-scroll to bottom - QTimer.singleShot(50, self._scroll_transcript_to_bottom) + # Auto-scroll to top + QTimer.singleShot(50, self._scroll_transcript_to_top) - def _scroll_transcript_to_bottom(self) -> None: - """Auto-scroll transcript to show latest messages.""" + def _scroll_transcript_to_top(self) -> None: + """Auto-scroll transcript to show the latest messages (newest is on top).""" vbar = self.transcript_scroll.verticalScrollBar() - vbar.setValue(vbar.maximum()) + vbar.setValue(vbar.minimum()) def set_connection(self, ok: bool, retry_count: int = 0, backoff: float = 0.0) -> None: if ok: @@ -865,7 +918,7 @@ class MainScreen(QtWidgets.QWidget): self._add_bubble(f"⚠ {message}", "error") self._prune_transcript() self._ensure_transcript_spacer() - QTimer.singleShot(50, self._scroll_transcript_to_bottom) + QTimer.singleShot(50, self._scroll_transcript_to_top) class SlaveWindow(QtWidgets.QMainWindow): @@ -907,6 +960,7 @@ class SlaveWindow(QtWidgets.QMainWindow): self.engineError.connect(self.main_screen.show_error) self.engineConnection.connect(self.main_screen.set_connection) self.engineRoom.connect(lambda r: self.main_screen.set_room(r)) + self.main_screen.fullscreenToggleRequested.connect(self._toggle_fullscreen) QtGui.QShortcut(QtGui.QKeySequence("F11"), self, activated=self._toggle_fullscreen) # Kiosk: fullscreen with a hidden cursor (touch panels); Esc leaves fullscreen. @@ -919,6 +973,12 @@ class SlaveWindow(QtWidgets.QMainWindow): self.showFullScreen() QtWidgets.QApplication.setOverrideCursor(QtCore.Qt.CursorShape.BlankCursor) + def changeEvent(self, event): # noqa: N802 + super().changeEvent(event) + if event.type() == QtCore.QEvent.Type.WindowStateChange: + # Covers every path in (button, F11, kiosk, settings) and out (Esc). + self.main_screen.set_fullscreen(self.isFullScreen()) + def _leave_fullscreen(self) -> None: if self.isFullScreen(): self.showNormal() diff --git a/src/javis/agent/orchestrator.py b/src/javis/agent/orchestrator.py index 2b22d20..65f6278 100644 --- a/src/javis/agent/orchestrator.py +++ b/src/javis/agent/orchestrator.py @@ -2,9 +2,11 @@ from __future__ import annotations import asyncio +import contextlib import json import re import time +from collections.abc import AsyncIterator from dataclasses import dataclass from typing import Any, Optional @@ -16,9 +18,9 @@ from ..audio_io.vad import make_vad from ..audio_io.wake_word import WakeWordDetector from ..core.config import Config from ..core.events import EventBus +from ..core.hardware import GpuUnavailableError from ..core.logger import get_logger from ..core.state import AgentState -from ..core.hardware import GpuUnavailableError from ..core.tasks import cancel_and_wait, spawn from ..llm.conversation import Conversation from ..llm.hermes_client import HermesClient @@ -48,16 +50,17 @@ from .middleware import ( drain_background_tasks, ) from .prompts import load_system_prompt +from ..stt.unidentified import UnidentifiedStore from ..stt.voice_profiler import NullProfiler, SpeakerPrivileges, VoiceProfiler log = get_logger("agent") -SENTENCE_BOUNDARY = re.compile(r"(?<=[\.\!\?…])\s+|(?<=[\.\!\?…])$") - -# Detects whether the last sentence of a reply is a question. # Returned by Orchestrator._identify_speaker when speaker ID raised. _ID_FAILED = object() +SENTENCE_BOUNDARY = re.compile(r"(?<=[\.\!\?…])\s+|(?<=[\.\!\?…])$") + +# Detects whether the last sentence of a reply is a question. _FOLLOWUP_RE = re.compile(r'\?\s*$') @@ -353,6 +356,37 @@ class RemoteTurnResult: expects_followup: bool = False +@dataclass +class VoiceSample: + """A live voice sample recorded from the GUI to enhance a profile.""" + audio: np.ndarray + sample_rate: int + embedding: Optional[np.ndarray] = None + # Cosine similarity to the target profile (None if it does not exist yet). + score: Optional[float] = None + best_name: Optional[str] = None + best_score: Optional[float] = None + + @property + def duration_s(self) -> float: + return len(self.audio) / float(self.sample_rate) if self.sample_rate else 0.0 + + +@dataclass +class _PendingEnrollClip: + """Unknown-speaker clip buffered while the "who is this?" prompt is open.""" + audio: np.ndarray + sample_rate: int + # Task storing the same clip in the unidentified inbox (returns its id). + store_task: Optional["asyncio.Task[Optional[str]]"] = None + + def inbox_id(self) -> Optional[str]: + t = self.store_task + if t is None or not t.done() or t.cancelled() or t.exception() is not None: + return None + return t.result() + + class Orchestrator: """Long-lived voice agent. Owns all subsystems.""" @@ -441,8 +475,16 @@ class Orchestrator: self.profiler: VoiceProfiler | NullProfiler = ( VoiceProfiler(config) if config.speaker_id.enabled else NullProfiler() ) - self._enroll_clip_buffer: list[np.ndarray] = [] + self.unidentified = UnidentifiedStore( + config.paths.voice_profiles_dir / "unidentified", + config.speaker_id.max_unidentified, + ) + self._enroll_clip_buffer: list[_PendingEnrollClip] = [] self._enroll_ask_pending: bool = False + # Cleared while the GUI records voice samples so the wake-word loop + # (which reads the mic outside _turn_lock) stops competing for chunks. + self._mic_free = asyncio.Event() + self._mic_free.set() self._bg_quick_turn_count: int = 0 self._tool_cache: dict[str, tuple[str, float]] = {} @@ -560,8 +602,13 @@ class Orchestrator: # Disable any backend whose load() raised so later synthesis/prewarm # calls don't encounter RuntimeError("load() has not been called"). if self.tts is not None and getattr(self.tts, "_voice", True) is None: - log.warning("TTS model failed to load — speech output disabled") - self.tts = None + backend, self.tts = self.tts, None + if hasattr(backend, "model_available"): + log.warning("TTS model failed to load — speech output disabled until " + "a model is available") + spawn(self._revive_tts(backend), name="tts-revive") + else: + log.warning("TTS model failed to load — speech output disabled") await self.bus.publish("mcp_ready", self.mcp.status_snapshot()) ok = await self.llm.health() await self.bus.publish("llm_health", {"ok": ok, "url": self.config.llm.base_url, "model": self.config.llm.model}) @@ -854,8 +901,12 @@ class Orchestrator: await self.bus.publish( "speaker_id", {"name": None, "confidence": 0.0, "known": False} ) - if self.config.speaker_id.auto_enroll: - await self._auto_enroll_speaker(audio) + source = device_name or "satellite" + if room: + source = f"{source} ({room})" + await self._handle_unknown_speaker( + audio, sample_rate, text, source=source, + ) if not _privileges.can_respond: log.info( @@ -947,6 +998,10 @@ class Orchestrator: False when the listen loop exited without a detection (stale timeout, ``_stop`` set, or mic disconnected).""" assert self.wake is not None + # A GUI voice-sample recording owns the mic — wait until it is done. + mic_free: Optional[asyncio.Event] = getattr(self, "_mic_free", None) + if mic_free is not None and not mic_free.is_set(): + await mic_free.wait() await self.set_state(AgentState.IDLE) self.wake.reset() if self.mic is not None: @@ -954,6 +1009,8 @@ class Orchestrator: log.info(f"listening for wake word: '{self.config.wake_word.model}'") last_pub = 0.0 async for chunk in self.mic.chunks(stale_timeout=self._WAKE_STALE_S): + if mic_free is not None and not mic_free.is_set(): + return False # handed over to an exclusive recording score = self.wake.feed(chunk) # Everything heard before the wake word is background (music, TV, # this device's own speakers), so the recorder's voice gate is @@ -1213,6 +1270,14 @@ class Orchestrator: return True return None + async def _identify_speaker(self, audio: np.ndarray, sample_rate: int) -> Any: + """Run speaker ID off-loop; a failure must never abort the turn.""" + try: + return await asyncio.to_thread(self.profiler.identify, audio, sample_rate) + except Exception as exc: # noqa: BLE001 + log.exception(f"[profiler] speaker identification failed: {exc}") + return _ID_FAILED + def _unknown_privileges(self) -> SpeakerPrivileges: """Build a :class:`SpeakerPrivileges` from the unknown-speaker config fields.""" cfg = self.config.speaker_id @@ -1228,6 +1293,32 @@ class Orchestrator: denied_tools=denied, ) + async def _revive_tts(self, backend, interval: float = 30.0) -> None: + """Turn speech output on once the TTS model that failed to load is there. + + E.g. a voice still in its first training run: training exports the + voice's model as it goes (see ``javis.voices.piper_train``). + """ + while True: + try: + await asyncio.wait_for(self._stop.wait(), timeout=interval) + return + except asyncio.TimeoutError: + pass + try: + if not backend.model_available(): + continue + await asyncio.to_thread(backend.load) + except Exception as exc: # noqa: BLE001 - tried again once the model changes + log.warning(f"TTS model still failed to load: {exc}") + continue + if self._stop.is_set(): + return + self.tts = backend + log.info("TTS model loaded — speech output enabled") + spawn(self._prewarm_tts(), name="tts-prewarm") + return + async def _prewarm_tts(self) -> None: """Silently synthesize a short phrase to warm up TTS model buffers. @@ -1270,14 +1361,6 @@ class Orchestrator: # Reset to IDLE so the GUI doesn't stay stuck on "listening" # when the turn ends without any speech. await self.set_state(AgentState.IDLE) - async def _identify_speaker(self, audio: np.ndarray, sample_rate: int) -> Any: - """Run speaker ID off-loop; a failure must never abort the turn.""" - try: - return await asyncio.to_thread(self.profiler.identify, audio, sample_rate) - except Exception as exc: # noqa: BLE001 - log.exception(f"[profiler] speaker identification failed: {exc}") - return _ID_FAILED - return await self.set_state(AgentState.THINKING) @@ -1347,8 +1430,9 @@ class Orchestrator: await self.bus.publish( "speaker_id", {"name": None, "confidence": 0.0, "known": False} ) - if self.config.speaker_id.auto_enroll: - await self._auto_enroll_speaker(audio) + await self._handle_unknown_speaker( + audio, self.config.audio.sample_rate, tr.text or "", source="local", + ) # ── privilege: can_respond check ────────────────────────────────────── if not _privileges.can_respond: @@ -2007,8 +2091,7 @@ class Orchestrator: n = self.config.speaker_id.min_enroll_clips clips: list[np.ndarray] = [] - async with self._turn_lock: - await self.set_state(AgentState.ENROLLING) + async with self._exclusive_mic(): for i in range(n): await self.bus.publish("enroll_progress", { "name": name, "clip": i, "total": n, "recording": True, @@ -2016,7 +2099,6 @@ class Orchestrator: spawn(self._play_wake_chime()) audio = await record_utterance(self.config, self.mic, self.vad) if audio is None or len(audio) < 1600: - await self.set_state(AgentState.IDLE) await self.bus.publish("enroll_progress", { "name": name, "clip": i, "total": n, "error": "Kein Audio erkannt — Aufnahme abgebrochen.", @@ -2030,48 +2112,283 @@ class Orchestrator: sr = self.config.audio.sample_rate await asyncio.to_thread(self.profiler.enroll, name, clips, sr) await self.bus.publish("speaker_enrolled", {"name": name, "clip_count": len(clips)}) - await self.set_state(AgentState.IDLE) log.info(f"[profiler] enrollment complete for '{name}'") return True + @contextlib.asynccontextmanager + async def _exclusive_mic(self, wait_s: float = 0.0) -> AsyncIterator[MicStream]: + """Own the microphone for a GUI-initiated recording. + + Takes ``_turn_lock`` (waiting up to *wait_s*) and pauses the wake-word + loop via ``_mic_free`` so every chunk reaches the recorder. Raises + ``RuntimeError`` with a user-facing message when that is not possible. + """ + mic = self.mic + if mic is None: + raise RuntimeError("No microphone available — voice input is disabled.") + try: + await asyncio.wait_for(self._turn_lock.acquire(), timeout=max(0.01, wait_s)) + except asyncio.TimeoutError: + raise RuntimeError("Javis is busy right now — try again in a moment.") from None + self._mic_free.clear() + try: + await self.set_state(AgentState.ENROLLING) + mic.drain() + yield mic + finally: + self._mic_free.set() + try: + await self.set_state(AgentState.IDLE) + finally: + self._turn_lock.release() + + def _require_profiler(self) -> None: + if not self.config.speaker_id.enabled: + raise RuntimeError( + "Speaker identification is disabled (Settings → Speaker Identification).") + if not self.profiler.is_ready(): + raise RuntimeError("The speaker identification model is not loaded yet.") + + async def record_voice_sample( + self, name: Optional[str] = None, *, max_seconds: float = 15.0, + ) -> VoiceSample: + """Record one live utterance for enhancing *name*'s voice profile. + + Nothing is stored — pass the result to :meth:`add_voice_samples` once + the user has reviewed it. ``score`` is the cosine similarity to + *name*'s current profile (``None`` if it has none). Raises + ``RuntimeError`` with a user-facing message on failure. + """ + self._require_profiler() + sr = self.config.audio.sample_rate + async with self._exclusive_mic(wait_s=3.0) as mic: + # Chime first, then drop whatever the mic picked up of it. + await self._play_wake_chime() + mic.drain() + audio = await record_utterance( + self.config, mic, self.vad, + max_seconds=max_seconds, initial_grace_ms=5000, + ) + if audio is None or len(audio) < int(sr * 0.5): + raise RuntimeError("No speech detected — please try again.") + emb = await asyncio.to_thread(self.profiler.embed, audio, sr) + ranked = self.profiler.rank(emb) + score = next((s for n, s in ranked if n == name), None) if name else None + best_name, best_score = ranked[0] if ranked else (None, None) + return VoiceSample( + audio=audio, sample_rate=sr, embedding=emb, score=score, + best_name=best_name, best_score=best_score, + ) + + async def add_voice_samples(self, name: str, samples: list[VoiceSample]) -> int: + """Blend recorded *samples* into *name*'s profile (created if missing). + + Returns the number of samples used. + """ + name = (name or "").strip() + if not name: + raise ValueError("Profile name is empty.") + + def _apply() -> tuple[int, bool]: + embs = [] + for s in samples: + emb = s.embedding + if emb is None or not np.any(emb): + emb = self.profiler.embed(s.audio, s.sample_rate) + if np.any(emb): + embs.append(emb) + if not embs: + return 0, False + if self.profiler.has_profile(name): + used = sum(1 for e in embs if self.profiler.refine_embedding(name, e)) + return used, False + return (len(embs), True) if self.profiler.enroll_embeddings(name, embs) else (0, False) + + used, created = await asyncio.to_thread(_apply) + if not used: + raise RuntimeError("No usable voice samples — is the speaker model loaded?") + log.info(f"[profiler] added {used} live sample(s) to '{name}'") + await self.bus.publish("voice_profiles_changed", { + "name": name, "added": used, "created": created, + }) + return used + + # ---------------- unidentified speakers ---------------- + + async def _handle_unknown_speaker( + self, audio: np.ndarray, sample_rate: int, text: str, *, source: str, + ) -> None: + """Keep the clip in the unidentified inbox and/or ask for a name.""" + store_task: Optional["asyncio.Task[Optional[str]]"] = None + if self.config.speaker_id.keep_unidentified and (text or "").strip(): + store_task = spawn( + self._store_unidentified(audio, sample_rate, text, source), + name="unidentified-store", + ) + if self.config.speaker_id.auto_enroll: + await self._auto_enroll_speaker(audio, sample_rate, store_task) + + async def _store_unidentified( + self, audio: np.ndarray, sample_rate: int, text: str, source: str, + ) -> Optional[str]: + """Embed *audio* and add it to the inbox. Returns the clip id.""" + try: + emb = await asyncio.to_thread(self.profiler.embed, audio, sample_rate) + ranked = self.profiler.rank(emb) + best_name, best_score = ranked[0] if ranked else (None, None) + self.unidentified.max_clips = max(1, int(self.config.speaker_id.max_unidentified)) + clip_id = await asyncio.to_thread( + self.unidentified.add, audio, sample_rate, + embedding=emb, text=text, source=source, + best_name=best_name, best_score=best_score, + ) + except Exception as exc: # noqa: BLE001 + log.warning(f"[unidentified] could not store clip: {exc}") + return None + await self.bus.publish("unidentified_changed", {"added": clip_id}) + return clip_id + + async def assign_unidentified(self, clip_ids: list[str], name: str) -> dict: + """Sort inbox clips into *name*'s profile, creating it if needed. + + Returns ``{"name", "added", "created"}``. The clips are removed from + the inbox. Raises ``ValueError``/``RuntimeError`` with a user-facing + message on failure. + """ + name = (name or "").strip() + if not name: + raise ValueError("Profile name is empty.") + if not clip_ids: + raise ValueError("No clips selected.") + self._require_profiler() + + def _apply() -> tuple[list[str], bool]: + used: list[str] = [] + embs: list[np.ndarray] = [] + for cid in clip_ids: + emb = self.unidentified.load_embedding(cid) + if emb is None: + try: + audio, sr = self.unidentified.load_audio(cid) + except Exception as exc: # noqa: BLE001 + log.warning(f"[unidentified] clip {cid} unavailable: {exc}") + continue + emb = self.profiler.embed(audio, sr) + if np.any(emb): + embs.append(emb) + used.append(cid) + if not embs: + return [], False + if self.profiler.has_profile(name): + for e in embs: + self.profiler.refine_embedding(name, e) + created = False + else: + if not self.profiler.enroll_embeddings(name, embs): + return [], False + created = True + self.unidentified.remove(used) + return used, created + + used, created = await asyncio.to_thread(_apply) + if not used: + raise RuntimeError("None of the selected clips could be used.") + # Don't let a still-open "who is this?" prompt enroll them a second time. + done = set(used) + self._enroll_clip_buffer = [ + p for p in self._enroll_clip_buffer if p.inbox_id() not in done + ] + log.info( + f"[profiler] sorted {len(used)} unidentified clip(s) into " + f"{'new' if created else 'existing'} profile '{name}'" + ) + await self.bus.publish("voice_profiles_changed", { + "name": name, "added": len(used), "created": created, + }) + await self.bus.publish("unidentified_changed", {"removed": used}) + return {"name": name, "added": len(used), "created": created} + + async def delete_unidentified(self, clip_ids: Optional[list[str]] = None) -> int: + """Delete the given inbox clips, or all of them when *clip_ids* is None.""" + if clip_ids is None: + removed = await asyncio.to_thread(self.unidentified.clear) + else: + removed = await asyncio.to_thread(self.unidentified.remove, clip_ids) + if removed: + await self.bus.publish("unidentified_changed", {"removed": clip_ids}) + return removed + + def pending_enroll_clips(self) -> list[tuple[np.ndarray, int]]: + """Clips buffered for the open "who is this?" prompt (for playback).""" + return [(p.audio, p.sample_rate) for p in self._enroll_clip_buffer] + + def auto_enroll_dismiss(self) -> None: + """The user skipped the name prompt. + + Drops the buffered clips (they stay in the unidentified inbox when + ``keep_unidentified`` is on) and lets the prompt appear again later. + """ + self._enroll_clip_buffer.clear() + self._enroll_ask_pending = False + async def auto_enroll_confirm(self, key: str, name: str) -> bool: """Finalize auto-enrollment using all clips buffered since the dialog opened. Called by the GUI when the user submits their name. *key* is kept for API compatibility but is no longer used to look up audio — the full - clip buffer is consumed instead. + clip buffer is consumed instead. Clips that were also stored in the + unidentified inbox are removed from it. """ - clips = list(self._enroll_clip_buffer) + pending = list(self._enroll_clip_buffer) self._enroll_clip_buffer.clear() self._enroll_ask_pending = False - if not clips: + if not pending: log.warning("[profiler] auto_enroll_confirm: clip buffer is empty") return False - sr = self.config.audio.sample_rate - existing_names = {p["name"] for p in self.profiler.list_profiles()} - if name in existing_names: - # Merge all buffered clips into the existing profile. - for clip in clips: - await asyncio.to_thread(self.profiler.refine, name, clip, sr) - log.info( - f"[profiler] merged {len(clips)} clip(s) into existing profile '{name}'" - ) - else: + def _apply() -> int: + embs = [self.profiler.embed(p.audio, p.sample_rate) for p in pending] + if self.profiler.has_profile(name): + # Merge all buffered clips into the existing profile. + n = sum(1 for e in embs if self.profiler.refine_embedding(name, e)) + log.info(f"[profiler] merged {n} clip(s) into existing profile '{name}'") + return n # Build a new profile from all accumulated clips. - await asyncio.to_thread(self.profiler.enroll, name, clips, sr) - log.info(f"[profiler] enrolled '{name}' from {len(clips)} clip(s)") - await self.bus.publish("speaker_enrolled", {"name": name, "clip_count": len(clips)}) + return len(embs) if self.profiler.enroll_embeddings(name, embs) else 0 + + used = await asyncio.to_thread(_apply) + if not used: + return False + + inbox_ids: list[str] = [] + for p in pending: + if p.store_task is not None: + try: + cid = await p.store_task + except Exception: # noqa: BLE001 + cid = None + if cid: + inbox_ids.append(cid) + if inbox_ids: + await self.delete_unidentified(inbox_ids) + await self.bus.publish("speaker_enrolled", {"name": name, "clip_count": used}) return True - async def _auto_enroll_speaker(self, audio: np.ndarray) -> None: + async def _auto_enroll_speaker( + self, + audio: np.ndarray, + sample_rate: Optional[int] = None, + store_task: Optional["asyncio.Task[Optional[str]]"] = None, + ) -> None: """Buffer *audio* from an unknown speaker and ask for a name (once). Subsequent unknown utterances while the dialog is open are silently appended to the buffer so enrollment ultimately uses multiple clips. """ - self._enroll_clip_buffer.append(audio) + self._enroll_clip_buffer.append(_PendingEnrollClip( + audio, int(sample_rate or self.config.audio.sample_rate), store_task, + )) if not self._enroll_ask_pending: self._enroll_ask_pending = True # Use a stable key so the GUI can match confirm callbacks. diff --git a/src/javis/core/config.py b/src/javis/core/config.py index 08fd003..3c4a09c 100644 --- a/src/javis/core/config.py +++ b/src/javis/core/config.py @@ -253,6 +253,11 @@ class SpeakerIDConfig: # auto-refinement. Only high-confidence identifications (score >= # threshold + refine_margin) update the profile, keeping noise clips out. refine_margin: float = 0.10 + # Keep audio clips of unidentified speakers (WAV + embedding, stored under + # voice_profiles_dir/unidentified) so they can be played back and sorted + # into a profile from the GUI. Capped at max_unidentified (oldest dropped). + keep_unidentified: bool = True + max_unidentified: int = 50 # ── Unknown / guest speaker privileges ─────────────────────────────────── # Applied when the speaker cannot be matched to any known voice profile. # Controlled via the GUI "Settings → Speaker Identification" section. @@ -295,7 +300,9 @@ class TrainingConfig: batch_size: int = 32 num_workers: int = 2 accelerator: str = "gpu" # gpu | cpu - precision: str = "16-mixed" # 32 | 16-mixed | bf16-mixed + # bf16-mixed trains fastest and needs no loss scaling; GPUs without bf16 + # fall back to 32 at launch. + precision: str = "bf16-mixed" # 32 | 16-mixed | bf16-mixed language: str = "de" # espeak-ng voice code (de, en-us, ...) @@ -430,6 +437,8 @@ def validate(cfg: "Config") -> None: cfg.speaker_id.threshold, 0.0, 1.0, 0.65, "speaker_id.threshold") cfg.speaker_id.refine_margin = _clamp( cfg.speaker_id.refine_margin, 0.0, 0.5, 0.10, "speaker_id.refine_margin") + cfg.speaker_id.max_unidentified = _clamp_int( + cfg.speaker_id.max_unidentified, 1, 500, 50, "speaker_id.max_unidentified") cfg.wake_word.vad_gate = _clamp( cfg.wake_word.vad_gate, 0.0, 1.0, 0.0, "wake_word.vad_gate") cfg.wake_word.voice_gate_ratio = _clamp( @@ -512,8 +521,8 @@ def validate(cfg: "Config") -> None: if t.precision not in {"32", "16-mixed", "bf16-mixed"}: log.warning( f"config: training.precision='{t.precision}' is unknown " - "(32|16-mixed|bf16-mixed), using '16-mixed'") - t.precision = "16-mixed" + "(32|16-mixed|bf16-mixed), using 'bf16-mixed'") + t.precision = "bf16-mixed" if cfg.memory.provider not in {"ollama", "hermes"}: log.warning( diff --git a/src/javis/gui/main_window.py b/src/javis/gui/main_window.py index 2b365f2..5786cf4 100644 --- a/src/javis/gui/main_window.py +++ b/src/javis/gui/main_window.py @@ -167,6 +167,8 @@ def run_gui(config: Config) -> int: speakerEnrolled = QtCore.Signal(dict) speakerAskName = QtCore.Signal(dict) enrollProgress = QtCore.Signal(dict) + unidentifiedChanged = QtCore.Signal(dict) + voiceProfilesChanged = QtCore.Signal(dict) timerFired = QtCore.Signal(dict) timerUpdated = QtCore.Signal() @@ -192,6 +194,14 @@ def run_gui(config: Config) -> int: self._enroll_dlg: QtWidgets.QDialog | None = None # Non-modal "new speaker" dialogs: key -> QDialog (keeps Python ref alive). self._name_dialogs: dict[str, QtWidgets.QDialog] = {} + # Non-modal "sort clip" / "enhance voice" dialogs (refs kept alive). + self._sort_dialogs: dict[str, QtWidgets.QDialog] = {} + self._enhance_dialogs: dict[str, QtWidgets.QDialog] = {} + # Single voice-clip playback (a new clip stops the previous one). + self._play_task: asyncio.Task | None = None + self._play_key: str | None = None + self._play_btns: list = [] # every visible button showing "stop" for it + self._voice_refresh_pending = False # Last sent user text (for Ctrl+R retry) and assistant streaming state. self._last_sent_text: str = "" self._assistant_buf: str = "" @@ -282,6 +292,8 @@ def run_gui(config: Config) -> int: self.speakerEnrolled.connect(self._on_speaker_enrolled) self.speakerAskName.connect(self._show_ask_name_dialog) self.enrollProgress.connect(self._on_enroll_progress) + self.unidentifiedChanged.connect(self._schedule_voice_refresh) + self.voiceProfilesChanged.connect(self._schedule_voice_refresh) self.timerFired.connect(self._on_timer_fired) self.timerUpdated.connect(self._refresh_timers) @@ -1659,6 +1671,16 @@ def run_gui(config: Config) -> int: form.addRow(f"Auto-enroll{L}", _w_check("speaker_id.auto_enroll", config.speaker_id.auto_enroll, tip="Ask for name and enroll unknown speakers automatically")) + form.addRow(f"Keep unknown clips{L}", + _w_check("speaker_id.keep_unidentified", + config.speaker_id.keep_unidentified, + tip="Save clips of unrecognised speakers so you can play them " + "back and sort them into a profile (Intelligence → Voice " + "Profiles). Stored locally under voice_profiles/unidentified.")) + form.addRow(f"Max unknown clips{L}", + _w_spin("speaker_id.max_unidentified", + config.speaker_id.max_unidentified, 1, 500, 10, + tip="The oldest unidentified clips are deleted beyond this limit")) form.addRow(QtWidgets.QLabel("")) # spacer form.addRow(QtWidgets.QLabel("Unknown / guest speaker:")) form.addRow(f" Can respond{L}", @@ -1845,6 +1867,11 @@ def run_gui(config: Config) -> int: _mcp_upd = updates.get("mcp", {}) if "approval_timeout_seconds" in _mcp_upd: config.mcp.approval_timeout_seconds = int(_mcp_upd["approval_timeout_seconds"]) + _spk_upd = updates.get("speaker_id", {}) + if "keep_unidentified" in _spk_upd: + config.speaker_id.keep_unidentified = bool(_spk_upd["keep_unidentified"]) + if "max_unidentified" in _spk_upd: + config.speaker_id.max_unidentified = int(_spk_upd["max_unidentified"]) # Sync model combo if the Settings tab changed llm.model _settings_model = updates.get("llm", {}).get("model", "") @@ -2261,15 +2288,30 @@ def run_gui(config: Config) -> int: item.widget().deleteLater() orch = self.orch profiles: list[dict] = [] + clips: list[dict] = [] if orch is not None and orch.profiler is not None: try: profiles = orch.profiler.list_profiles() except Exception: log.debug("[voice] failed to list voice profiles") + if orch is not None and getattr(orch, "unidentified", None) is not None: + try: + clips = orch.unidentified.entries() + except Exception: + log.debug("[voice] failed to list unidentified clips") if self._voice_count_lbl is not None: - self._voice_count_lbl.setText( - f"{len(profiles)} profile{'s' if len(profiles) != 1 else ''}" - ) + text = f"{len(profiles)} profile{'s' if len(profiles) != 1 else ''}" + if clips: + text += f" · {len(clips)} unidentified" + self._voice_count_lbl.setText(text) + if clips: + self._voice_cards_layout.addWidget(self._make_voice_section_header( + f"Unidentified speakers ({len(clips)})", + action=("Clear all", self._clear_unidentified), + )) + for clip in clips: + self._voice_cards_layout.addWidget(self._make_unidentified_card(clip)) + self._voice_cards_layout.addWidget(self._make_voice_section_header("Profiles")) if not profiles: empty = QtWidgets.QLabel("No voice profiles yet.\nUse + to enroll a speaker.") empty.setAlignment(QtCore.Qt.AlignmentFlag.AlignCenter) @@ -2279,6 +2321,18 @@ def run_gui(config: Config) -> int: for profile in profiles: self._voice_cards_layout.addWidget(self._make_voice_card(profile)) + def _schedule_voice_refresh(self, *_args): + """Coalesce bursts of inbox/profile events into one panel rebuild.""" + if self._voice_refresh_pending: + return + self._voice_refresh_pending = True + + def _go(): + self._voice_refresh_pending = False + self._refresh_voice_cards() + + QtCore.QTimer.singleShot(150, _go) + def _make_voice_card(self, profile: dict) -> QtWidgets.QFrame: card = QtWidgets.QFrame() self._th(card, self._card_qss) @@ -2323,12 +2377,687 @@ def run_gui(config: Config) -> int: f"QPushButton:hover {{ color: {_CL_ACCENT}; background: #2a1810; }}")) priv_btn.clicked.connect(lambda _=False, n=name: self._edit_voice_privileges(n)) + enh_btn = QtWidgets.QPushButton("🎤") + enh_btn.setFixedSize(24, 24) + enh_btn.setToolTip(f"Enhance '{name}' with live voice samples") + enh_btn.setAccessibleName(f"Enhance voice of {name}") + self._th(enh_btn, lambda: ( + f"QPushButton {{ background: transparent; border: 0; color: {_CL_MUTED}; " + f"font-size: 10pt; border-radius: 4px; padding: 0; }}" + f"QPushButton:hover {{ color: {_CL_ACCENT}; background: {_CL_BTN_HOVER}; }}")) + enh_btn.clicked.connect(lambda _=False, n=name: self._enhance_voice_dialog(n)) + f_row.addWidget(date_lbl, 1) + f_row.addWidget(enh_btn) f_row.addWidget(priv_btn) f_row.addWidget(del_btn) v.addLayout(f_row) return card + # ─── unidentified speakers + live samples ─────────────────────────── + + def _make_voice_section_header(self, title: str, action=None) -> QtWidgets.QWidget: + w = QtWidgets.QWidget() + h = QtWidgets.QHBoxLayout(w) + h.setContentsMargins(2, 4, 0, 0) + h.setSpacing(6) + lbl = QtWidgets.QLabel(title.upper()) + self._th(lbl, lambda: ( + f"color: {_CL_MUTED}; font-size: 7.5pt; font-weight: 600; " + f"letter-spacing: 0.5px; background: transparent;")) + h.addWidget(lbl, 1) + if action is not None: + text, fn = action + btn = QtWidgets.QPushButton(text) + btn.setCursor(QtCore.Qt.CursorShape.PointingHandCursor) + self._th(btn, self._ghost_qss) + btn.clicked.connect(lambda _=False: fn()) + h.addWidget(btn) + return w + + def _small_btn_qss(self, accent: bool = False) -> str: + fg = _CL_ACCENT if accent else _CL_TEXT + return ( + f"QPushButton {{ background: {_CL_SURFACE2}; color: {fg}; " + f" border: 1px solid {_CL_BORDER2}; border-radius: 5px; " + f" padding: 2px 10px; font-size: 8.5pt; min-height: 20px; }}" + f"QPushButton:hover {{ background: {_CL_BTN_HOVER}; border-color: {_CL_STROKE}; }}" + f"QPushButton:pressed {{ background: {_CL_BTN_PRESSED}; }}" + f"QPushButton:focus {{ border-color: {_CL_FOCUS}; }}" + f"QPushButton:disabled {{ color: {_CL_DISABLED}; border-color: {_CL_BORDER}; }}" + ) + + @staticmethod + def _fmt_clip_time(stamp: str) -> str: + """'2026-05-01T14:03:22' → '14:03' today, else '01.05. 14:03'.""" + if len(stamp) < 16: + return stamp + if stamp[:10] == time.strftime("%Y-%m-%d"): + return stamp[11:16] + return f"{stamp[8:10]}.{stamp[5:7]}. {stamp[11:16]}" + + @staticmethod + def _fmt_clip_source(source: str) -> str: + return "this PC" if source in ("", "local") else source + + def _clip_matches(self, clip_id: str) -> list[tuple[str, float]]: + """Profiles ranked by similarity to an inbox clip (live, best first).""" + orch = self.orch + if orch is None: + return [] + try: + emb = orch.unidentified.load_embedding(clip_id) + return orch.profiler.rank(emb) if emb is not None else [] + except Exception: + return [] + + def _make_play_btn(self, idle: str = "▶ Play", busy: str = "■ Stop", + key: str | None = None) -> QtWidgets.QPushButton: + btn = QtWidgets.QPushButton(idle) + btn.setProperty("idleText", idle) + btn.setProperty("busyText", busy) + btn.setToolTip("Play / stop this recording") + btn.setCursor(QtCore.Qt.CursorShape.PointingHandCursor) + self._th(btn, self._small_btn_qss) + if key is not None and key == self._play_key: + btn.setText(busy) + self._play_btns.append(btn) + return btn + + @staticmethod + def _set_play_label(btn, playing: bool) -> None: + if btn is None: + return + try: + btn.setText(btn.property("busyText" if playing else "idleText") or "") + except RuntimeError: + pass # card was rebuilt while the clip played + + def _toggle_playback(self, key: str, loader, btn=None) -> None: + """Play the audio returned by *loader()*; clicking again stops it. + + Only one clip plays at a time — starting another stops the first. + """ + from ..audio_io.stream import play_pcm + + prev = self._play_task + if prev is not None and not prev.done(): + prev.cancel() + if self._play_key == key: + return + try: + audio, sr = loader() + except Exception as exc: + self.statusBar().showMessage(f"Cannot play recording: {exc}", 5000) + return + if audio is None or len(audio) == 0: + return + btns = [btn] if btn is not None else [] + self._play_key = key + self._play_btns = btns + self._set_play_label(btn, True) + task = spawn(play_pcm(audio, int(sr), config.audio.output_device), + name="voice-clip-playback") + self._play_task = task + + def _done(t): + if self._play_task is t: + self._play_task = None + self._play_key = None + self._play_btns = [] + for b in btns: + self._set_play_label(b, False) + if not t.cancelled() and t.exception() is not None: + self.statusBar().showMessage(f"Playback failed: {t.exception()}", 5000) + + task.add_done_callback(_done) + + def _play_inbox_clip(self, clip_id: str, btn=None) -> None: + orch = self.orch + if orch is None: + return + self._toggle_playback(f"inbox:{clip_id}", + lambda: orch.unidentified.load_audio(clip_id), btn) + + def _make_unidentified_card(self, clip: dict) -> QtWidgets.QFrame: + cid = str(clip.get("id", "")) + card = QtWidgets.QFrame() + self._th(card, self._card_qss) + v = QtWidgets.QVBoxLayout(card) + v.setContentsMargins(14, 10, 14, 9) + v.setSpacing(5) + + h = QtWidgets.QHBoxLayout() + h.setSpacing(6) + title = QtWidgets.QLabel("❔ Unknown speaker") + self._th(title, lambda: f"color: {_CL_TEXT}; font-weight: 600; font-size: 9.5pt;") + h.addWidget(title, 1) + dur = float(clip.get("duration_s") or 0.0) + dur_lbl = QtWidgets.QLabel(f" {dur:.1f} s ") + + def _dur_qss() -> str: + # Qt reads 8-digit hex as #AARRGGBB, so tint via rgba() instead. + c = QtGui.QColor(_CL_WARN_FG) + return (f"background: rgba({c.red()}, {c.green()}, {c.blue()}, 0.16); " + f"color: {_CL_WARN_FG}; font-size: 8pt; " + f"border-radius: 3px; padding: 1px 4px;") + + self._th(dur_lbl, _dur_qss) + h.addWidget(dur_lbl) + v.addLayout(h) + + text = str(clip.get("text") or "").strip() + if text: + shown = text if len(text) <= 140 else text[:137] + "…" + txt_lbl = QtWidgets.QLabel(f"“{shown}”") + txt_lbl.setWordWrap(True) + txt_lbl.setToolTip(text) + self._th(txt_lbl, lambda: f"color: {_CL_TEXT}; font-size: 9pt; font-style: italic;") + v.addWidget(txt_lbl) + + meta = [self._fmt_clip_time(str(clip.get("created_at") or "")), + self._fmt_clip_source(str(clip.get("source") or ""))] + matches = self._clip_matches(cid) + if matches: + meta.append(f"closest: {matches[0][0]} {max(0.0, matches[0][1]):.0%}") + meta_lbl = QtWidgets.QLabel(" · ".join(m for m in meta if m)) + meta_lbl.setWordWrap(True) + self._th(meta_lbl, lambda: f"color: {_CL_MUTED}; font-size: 8pt;") + v.addWidget(meta_lbl) + + row = QtWidgets.QHBoxLayout() + row.setSpacing(6) + play_btn = self._make_play_btn(key=f"inbox:{cid}") + play_btn.clicked.connect( + lambda _=False, c=cid, b=play_btn: self._play_inbox_clip(c, b)) + sort_btn = QtWidgets.QPushButton("Sort…") + sort_btn.setToolTip("Assign this voice to an existing profile or create a new one") + sort_btn.setCursor(QtCore.Qt.CursorShape.PointingHandCursor) + self._th(sort_btn, lambda: self._small_btn_qss(accent=True)) + sort_btn.clicked.connect(lambda _=False, c=cid: self._sort_unidentified_dialog(c)) + del_btn = QtWidgets.QPushButton("✕") + del_btn.setFixedSize(24, 24) + del_btn.setToolTip("Delete this recording") + del_btn.setAccessibleName("Delete recording") + self._th(del_btn, lambda: ( + f"QPushButton {{ background: transparent; border: 0; color: {_CL_MUTED}; " + f"font-size: 10pt; border-radius: 4px; padding: 0; }}" + f"QPushButton:hover {{ color: {_CL_ERROR}; background: {_CL_ERROR_BG}; }}")) + del_btn.clicked.connect(lambda _=False, c=cid: self._delete_unidentified([c])) + row.addWidget(play_btn) + row.addWidget(sort_btn) + row.addStretch(1) + row.addWidget(del_btn) + v.addLayout(row) + return card + + def _delete_unidentified(self, clip_ids: list[str] | None) -> None: + orch = self.orch + if orch is None: + return + key = self._play_key or "" + if (key.startswith("inbox:") and self._play_task is not None + and (clip_ids is None or key.removeprefix("inbox:") in clip_ids)): + self._play_task.cancel() + QtCore.QTimer.singleShot(0, lambda: spawn(orch.delete_unidentified(clip_ids))) + + def _clear_unidentified(self) -> None: + reply = QtWidgets.QMessageBox.question( + self, "Clear unidentified speakers", + "Delete all recordings of unidentified speakers?", + QtWidgets.QMessageBox.StandardButton.Yes | QtWidgets.QMessageBox.StandardButton.No, + ) + if reply == QtWidgets.QMessageBox.StandardButton.Yes: + self._delete_unidentified(None) + + def _sort_unidentified_dialog(self, clip_id: str) -> None: + """Non-modal dialog: assign an inbox clip (plus similar ones) to a profile.""" + import numpy as np + + orch = self.orch + if orch is None: + return + if clip_id in self._sort_dialogs: + self._sort_dialogs[clip_id].raise_() + self._sort_dialogs[clip_id].activateWindow() + return + clip = orch.unidentified.get(clip_id) + if clip is None: + self._refresh_voice_cards() + return + threshold = config.speaker_id.threshold + ranked = self._clip_matches(clip_id) + profile_names = [p["name"] for p in orch.profiler.list_profiles()] + # Best matches first, then any profile without a score. + ordered = [n for n, _ in ranked if n in profile_names] + ordered += [n for n in profile_names if n not in ordered] + + dlg = QtWidgets.QDialog(self) + dlg.setWindowTitle("Sort unidentified speaker") + dlg.setMinimumWidth(480) + dlg.setStyleSheet(self._dialog_qss()) + lay = QtWidgets.QVBoxLayout(dlg) + lay.setSpacing(12) + lay.setContentsMargins(20, 18, 20, 18) + + # ── the clip ───────────────────────────────────────────────────── + head = QtWidgets.QHBoxLayout() + info = QtWidgets.QLabel( + f"{self._fmt_clip_time(str(clip.get('created_at') or ''))}" + f" · {self._fmt_clip_source(str(clip.get('source') or ''))}" + f" · {float(clip.get('duration_s') or 0):.1f} s") + head.addWidget(info, 1) + play_btn = self._make_play_btn(key=f"inbox:{clip_id}") + play_btn.clicked.connect(lambda _=False: self._play_inbox_clip(clip_id, play_btn)) + head.addWidget(play_btn) + lay.addLayout(head) + text = str(clip.get("text") or "").strip() + if text: + t_lbl = QtWidgets.QLabel(f"“{text}”") + t_lbl.setWordWrap(True) + t_lbl.setStyleSheet(f"color: {_CL_TEXT}; font-style: italic;") + lay.addWidget(t_lbl) + if ranked: + m_lbl = QtWidgets.QLabel("Closest profiles: " + ", ".join( + f"{n} {max(0.0, s):.0%}" for n, s in ranked[:3]) + + f" (match needs ≥ {threshold:.0%})") + m_lbl.setWordWrap(True) + m_lbl.setStyleSheet(f"color: {_CL_MUTED}; font-size: 9pt;") + lay.addWidget(m_lbl) + + # ── target profile ─────────────────────────────────────────────── + radio_existing = QtWidgets.QRadioButton("Add to profile:") + combo = QtWidgets.QComboBox() + for nm in ordered: + combo.addItem(nm) + row_ex = QtWidgets.QHBoxLayout() + row_ex.setSpacing(8) + row_ex.addWidget(radio_existing) + row_ex.addWidget(combo, 1) + lay.addLayout(row_ex) + + radio_new = QtWidgets.QRadioButton("New profile:") + name_edit = QtWidgets.QLineEdit() + name_edit.setPlaceholderText("Name of the speaker…") + row_new = QtWidgets.QHBoxLayout() + row_new.setSpacing(8) + row_new.addWidget(radio_new) + row_new.addWidget(name_edit, 1) + lay.addLayout(row_new) + + hint_lbl = QtWidgets.QLabel("") + hint_lbl.setWordWrap(True) + hint_lbl.setStyleSheet(f"color: {_CL_WARN_FG}; font-size: 9pt;") + hint_lbl.hide() + lay.addWidget(hint_lbl) + + if not ordered: + radio_existing.setEnabled(False) + combo.setEnabled(False) + radio_new.setChecked(True) + elif ranked and ranked[0][1] >= threshold - 0.15: + radio_existing.setChecked(True) + else: + radio_new.setChecked(True) + + def _update_fields(): + use_new = radio_new.isChecked() + name_edit.setEnabled(use_new) + combo.setEnabled(not use_new and bool(ordered)) + nm = name_edit.text().strip() + if use_new and nm in profile_names: + hint_lbl.setText(f"A profile named '{nm}' exists — the clips will be " + "added to it.") + hint_lbl.show() + else: + hint_lbl.hide() + + radio_new.toggled.connect(lambda _=False: _update_fields()) + name_edit.textChanged.connect(lambda _=None: _update_fields()) + _update_fields() + + # ── other inbox clips that may be the same person ──────────────── + checks: list[tuple[str, QtWidgets.QCheckBox]] = [] + emb = orch.unidentified.load_embedding(clip_id) + others = [] + for e in orch.unidentified.entries(): + if e["id"] == clip_id: + continue + sim = None + o_emb = orch.unidentified.load_embedding(e["id"]) + if emb is not None and o_emb is not None and emb.shape == o_emb.shape: + denom = float(np.linalg.norm(emb) * np.linalg.norm(o_emb)) + sim = float(np.dot(emb, o_emb) / denom) if denom > 0 else None + others.append((e, sim)) + if others: + others.sort(key=lambda x: -(x[1] if x[1] is not None else -2.0)) + also_lbl = QtWidgets.QLabel( + "Also sort these recordings (pre-selected when they sound alike):") + also_lbl.setWordWrap(True) + lay.addWidget(also_lbl) + box = QtWidgets.QWidget() + box_lay = QtWidgets.QVBoxLayout(box) + box_lay.setContentsMargins(0, 0, 4, 0) + box_lay.setSpacing(4) + for e, sim in others: + r = QtWidgets.QHBoxLayout() + r.setSpacing(6) + o_text = str(e.get("text") or "").strip() + o_text = o_text if len(o_text) <= 48 else o_text[:45] + "…" + cb = QtWidgets.QCheckBox( + f"{self._fmt_clip_time(str(e.get('created_at') or ''))} “{o_text}”") + cb.setToolTip(str(e.get("text") or "")) + cb.setChecked(sim is not None and sim >= threshold) + # Long transcripts are clipped instead of forcing a side scrollbar. + cb.setSizePolicy(QtWidgets.QSizePolicy.Policy.Ignored, + QtWidgets.QSizePolicy.Policy.Fixed) + r.addWidget(cb, 1) + s_lbl = QtWidgets.QLabel("—" if sim is None else f"{max(0.0, sim):.0%}") + s_lbl.setToolTip("Similarity to the recording above") + s_lbl.setStyleSheet( + f"color: {_CL_SUCCESS if sim is not None and sim >= threshold else _CL_MUTED};" + " font-size: 9pt;") + r.addWidget(s_lbl) + p_btn = self._make_play_btn("▶", "■", key=f"inbox:{e['id']}") + p_btn.setFixedWidth(34) + p_btn.clicked.connect( + lambda _=False, c=e["id"], b=p_btn: self._play_inbox_clip(c, b)) + r.addWidget(p_btn) + box_lay.addLayout(r) + checks.append((e["id"], cb)) + scroll = QtWidgets.QScrollArea() + scroll.setWidgetResizable(True) + scroll.setFrameShape(QtWidgets.QFrame.Shape.NoFrame) + scroll.setHorizontalScrollBarPolicy(QtCore.Qt.ScrollBarPolicy.ScrollBarAlwaysOff) + scroll.setWidget(box) + scroll.setFixedHeight(min(200, box.sizeHint().height() + 2)) + lay.addWidget(scroll) + + err_lbl = QtWidgets.QLabel("") + err_lbl.setWordWrap(True) + err_lbl.setStyleSheet(f"color: {_CL_ERROR}; font-size: 9pt;") + err_lbl.hide() + lay.addWidget(err_lbl) + + btn_row = QtWidgets.QHBoxLayout() + cancel_btn = QtWidgets.QPushButton("Cancel") + cancel_btn.setStyleSheet(self._secondary_qss()) + assign_btn = QtWidgets.QPushButton("Assign") + assign_btn.setStyleSheet(self._primary_qss()) + assign_btn.setDefault(True) + btn_row.addStretch(1) + btn_row.addWidget(cancel_btn) + btn_row.addWidget(assign_btn) + lay.addLayout(btn_row) + + def _assign(): + nm = (name_edit.text() if radio_new.isChecked() else combo.currentText()).strip() + if not nm: + err_lbl.setText("⚠ Please enter a name.") + err_lbl.show() + name_edit.setFocus() + return + ids = [clip_id] + [c for c, cb in checks if cb.isChecked()] + assign_btn.setEnabled(False) + err_lbl.hide() + + async def _run(): + try: + res = await orch.assign_unidentified(ids, nm) + except Exception as exc: + try: + err_lbl.setText(f"⚠ {exc}") + err_lbl.show() + assign_btn.setEnabled(True) + except RuntimeError: + pass + return + try: + dlg.accept() + except RuntimeError: + pass # dialog already closed + verb = "Created" if res["created"] else "Updated" + n = res["added"] + self.statusBar().showMessage( + f"{verb} voice profile '{res['name']}' with " + f"{n} recording{'s' if n != 1 else ''}.", 6000) + + # Deferred: avoid re-entering qasync from inside the click handler. + QtCore.QTimer.singleShot(0, lambda: spawn(_run())) + + def _sort_finished(_r=0) -> None: + self._sort_dialogs.pop(clip_id, None) + dlg.deleteLater() + + assign_btn.clicked.connect(_assign) + cancel_btn.clicked.connect(dlg.reject) + name_edit.returnPressed.connect(_assign) + dlg.finished.connect(_sort_finished) + self._sort_dialogs[clip_id] = dlg + dlg.show() + if radio_new.isChecked(): + name_edit.setFocus() + + def _enhance_voice_dialog(self, name: str) -> None: + """Non-modal dialog: record live samples and blend them into *name*.""" + orch = self.orch + if orch is None: + return + if name in self._enhance_dialogs: + self._enhance_dialogs[name].raise_() + self._enhance_dialogs[name].activateWindow() + return + threshold = config.speaker_id.threshold + samples: list = [] + state: dict = {"task": None} + + class _EnhanceDialog(QtWidgets.QDialog): + """Asks before Esc / ✕ / Close throws away unsaved samples.""" + + def reject(self): + if samples: + ret = QtWidgets.QMessageBox.question( + self, "Discard samples", + f"Discard {len(samples)} recorded sample(s) without adding them?", + QtWidgets.QMessageBox.StandardButton.Yes + | QtWidgets.QMessageBox.StandardButton.No, + ) + if ret != QtWidgets.QMessageBox.StandardButton.Yes: + return + super().reject() + + def closeEvent(self, event): + event.ignore() + self.reject() + + dlg = _EnhanceDialog(self) + dlg.setWindowTitle(f"Enhance voice — {name}") + dlg.setMinimumWidth(460) + dlg.setStyleSheet(self._dialog_qss()) + lay = QtWidgets.QVBoxLayout(dlg) + lay.setSpacing(12) + lay.setContentsMargins(20, 18, 20, 18) + + intro = QtWidgets.QLabel( + f"Record a few sentences as {name}, speaking naturally from where " + "you usually talk to Javis. Recording stops after a short pause. Each " + "sample is compared with the current profile before you add it.") + intro.setWordWrap(True) + lay.addWidget(intro) + + rec_row = QtWidgets.QHBoxLayout() + rec_btn = QtWidgets.QPushButton("● Record sample") + rec_btn.setStyleSheet(self._primary_qss()) + rec_btn.setToolTip("Record one live voice sample from the microphone") + rec_row.addWidget(rec_btn) + rec_status = QtWidgets.QLabel("") + rec_status.setWordWrap(True) + rec_status.setStyleSheet(f"color: {_CL_MUTED}; font-size: 9pt;") + rec_row.addWidget(rec_status, 1) + lay.addLayout(rec_row) + + list_box = QtWidgets.QWidget() + list_lay = QtWidgets.QVBoxLayout(list_box) + list_lay.setContentsMargins(0, 0, 0, 0) + list_lay.setSpacing(4) + lay.addWidget(list_box) + + btn_row = QtWidgets.QHBoxLayout() + close_btn = QtWidgets.QPushButton("Close") + close_btn.setStyleSheet(self._secondary_qss()) + add_btn = QtWidgets.QPushButton("Add samples") + add_btn.setStyleSheet(self._primary_qss()) + btn_row.addStretch(1) + btn_row.addWidget(close_btn) + btn_row.addWidget(add_btn) + lay.addLayout(btn_row) + + def _set_status(text: str, color: str = _CL_MUTED) -> None: + try: + rec_status.setText(text) + rec_status.setStyleSheet(f"color: {color}; font-size: 9pt;") + except RuntimeError: + pass + + def _rebuild_rows() -> None: + while list_lay.count(): + item = list_lay.takeAt(0) + w = item.widget() + if w is not None: + w.deleteLater() + elif item.layout() is not None: + while item.layout().count(): + sub = item.layout().takeAt(0) + if sub.widget() is not None: + sub.widget().deleteLater() + if not samples: + empty = QtWidgets.QLabel("No samples recorded yet.") + empty.setStyleSheet(f"color: {_CL_MUTED}; font-size: 9pt;") + list_lay.addWidget(empty) + for i, s in enumerate(samples): + row = QtWidgets.QWidget() + r = QtWidgets.QHBoxLayout(row) + r.setContentsMargins(0, 0, 0, 0) + r.setSpacing(8) + r.addWidget(QtWidgets.QLabel(f"#{i + 1}")) + r.addWidget(QtWidgets.QLabel(f"{s.duration_s:.1f} s")) + if s.score is None: + txt, col = "new profile", _CL_MUTED + tip = "There is no stored profile to compare with yet." + elif s.score >= threshold: + txt, col = f"match {s.score:.0%}", _CL_SUCCESS + tip = f"Sounds like {name}." + else: + txt, col = f"match {max(0.0, s.score):.0%} — low", _CL_WARN_FG + tip = (f"This sample does not sound much like {name}; adding " + "it may weaken the profile.") + if (s.best_name and s.best_name != name and s.best_score is not None + and s.best_score > (s.score or 0.0) + 0.05): + txt += f" · sounds like {s.best_name}" + col = _CL_WARN_FG + m_lbl = QtWidgets.QLabel(txt) + m_lbl.setToolTip(tip) + m_lbl.setStyleSheet(f"color: {col}; font-size: 9pt;") + r.addWidget(m_lbl, 1) + p_btn = self._make_play_btn("▶", "■", key=f"sample:{name}:{id(s)}") + p_btn.setFixedWidth(34) + p_btn.clicked.connect( + lambda _=False, smp=s, b=p_btn: self._toggle_playback( + f"sample:{name}:{id(smp)}", + lambda: (smp.audio, smp.sample_rate), b)) + r.addWidget(p_btn) + x_btn = QtWidgets.QPushButton("✕") + x_btn.setFixedSize(26, 26) + x_btn.setToolTip("Discard this sample") + x_btn.setStyleSheet(self._icon_btn_qss()) + x_btn.clicked.connect(lambda _=False, smp=s: ( + samples.remove(smp), _rebuild_rows())) + r.addWidget(x_btn) + list_lay.addWidget(row) + n = len(samples) + add_btn.setText(f"Add {n} sample{'s' if n != 1 else ''}" if n else "Add samples") + add_btn.setEnabled(n > 0 and not (state["task"] and not state["task"].done())) + + def _record() -> None: + t = state["task"] + if t is not None and not t.done(): + return + rec_btn.setEnabled(False) + add_btn.setEnabled(False) + rec_btn.setText("● Listening…") + _set_status("Speak now…", _CL_ACCENT) + + async def _run(): + try: + sample = await orch.record_voice_sample(name) + except asyncio.CancelledError: + raise + except Exception as exc: + _set_status(f"⚠ {exc}", _CL_ERROR) + else: + samples.append(sample) + _set_status(f"✓ Sample {len(samples)} recorded " + f"({sample.duration_s:.1f} s).", _CL_SUCCESS) + finally: + try: + rec_btn.setEnabled(True) + rec_btn.setText("● Record sample") + state["task"] = None + _rebuild_rows() + except RuntimeError: + pass + + def _start(): + state["task"] = spawn(_run(), name="voice-sample-record") + + QtCore.QTimer.singleShot(0, _start) + + def _add() -> None: + if not samples: + return + add_btn.setEnabled(False) + rec_btn.setEnabled(False) + _set_status("Updating profile…") + + async def _run(): + try: + n = await orch.add_voice_samples(name, list(samples)) + except Exception as exc: + _set_status(f"⚠ {exc}", _CL_ERROR) + try: + add_btn.setEnabled(True) + rec_btn.setEnabled(True) + except RuntimeError: + pass + return + samples.clear() + try: + dlg.accept() + except RuntimeError: + pass # dialog already closed + self.statusBar().showMessage( + f"Added {n} live sample{'s' if n != 1 else ''} to '{name}'.", 6000) + + QtCore.QTimer.singleShot(0, lambda: spawn(_run())) + + def _finished(_r=0) -> None: + t = state["task"] + if t is not None and not t.done(): + t.cancel() # releases the microphone + if (self._play_task is not None + and (self._play_key or "").startswith(f"sample:{name}:")): + self._play_task.cancel() + self._enhance_dialogs.pop(name, None) + dlg.deleteLater() + + rec_btn.clicked.connect(_record) + add_btn.clicked.connect(_add) + close_btn.clicked.connect(dlg.reject) + dlg.finished.connect(_finished) + _rebuild_rows() + self._enhance_dialogs[name] = dlg + dlg.show() + def _add_voice_profile(self): self._enroll_speaker_dialog(None) @@ -2500,9 +3229,16 @@ def run_gui(config: Config) -> int: lay.setSpacing(14) lay.setContentsMargins(20, 20, 20, 20) - lay.addWidget(QtWidgets.QLabel( + head = QtWidgets.QHBoxLayout() + head.addWidget(QtWidgets.QLabel( "Javis detected an unknown speaker.\nHow would you like to continue?" - )) + ), 1) + play_btn = self._make_play_btn(key=f"ask:{key}") + play_btn.setToolTip("Listen to what the unknown speaker said") + play_btn.clicked.connect(lambda _=False: self._toggle_playback( + f"ask:{key}", self._pending_enroll_audio, play_btn)) + head.addWidget(play_btn, 0, QtCore.Qt.AlignmentFlag.AlignTop) + lay.addLayout(head) # ── New profile row ─────────────────────────────────────────────── radio_new = QtWidgets.QRadioButton("New profile:") @@ -2548,6 +3284,9 @@ def run_gui(config: Config) -> int: enroll_btn.setDefault(True) skip_btn = QtWidgets.QPushButton("Skip") skip_btn.setStyleSheet(self._secondary_qss()) + if config.speaker_id.keep_unidentified: + skip_btn.setToolTip("Decide later — the recording stays under " + "'Unidentified speakers' in Voice Profiles.") btn_row.addStretch(1) btn_row.addWidget(skip_btn) btn_row.addWidget(enroll_btn) @@ -2570,14 +3309,41 @@ def run_gui(config: Config) -> int: def _on_finished(): self._name_dialogs.pop(key, None) + if self._play_key == f"ask:{key}" and self._play_task is not None: + self._play_task.cancel() + + def _on_rejected(): + orch = self.orch + if orch is not None: + orch.auto_enroll_dismiss() enroll_btn.clicked.connect(_enroll) skip_btn.clicked.connect(dlg.reject) name_edit.returnPressed.connect(_enroll) + dlg.rejected.connect(_on_rejected) dlg.finished.connect(_on_finished) self._name_dialogs[key] = dlg dlg.show() # non-modal — does not block the asyncio event loop + def _pending_enroll_audio(self): + """All clips buffered for the open name prompt, joined with short gaps.""" + import numpy as np + + orch = self.orch + clips = orch.pending_enroll_clips() if orch is not None else [] + if not clips: + raise RuntimeError("the recording is no longer available") + sr = int(clips[0][1]) + gap = np.zeros(int(sr * 0.3), dtype=np.float32) + parts: list = [] + for audio, rate in clips: + if int(rate) != sr: + continue + if parts: + parts.append(gap) + parts.append(np.asarray(audio, dtype=np.float32).reshape(-1)) + return np.concatenate(parts), sr + def _on_speaker_identified(self, payload: dict): if self._current_speaker_lbl is None: @@ -3129,6 +3895,10 @@ def run_gui(config: Config) -> int: return if self._tray is not None: self._tray.hide() + # Lets a running transcode thread exit so interpreter shutdown + # doesn't block on it. + if self._ingest_stop_event is not None: + self._ingest_stop_event.set() event.accept() QtWidgets.QApplication.instance().quit() @@ -3329,10 +4099,15 @@ def run_gui(config: Config) -> int: elif phase == "stopped": at = p.get("at_index", 0); total = p.get("total", 1) pct = int(100 * at / max(1, total)) - self._task_finish(tid, ok=True, - msg=f"stopped at {at}/{total} — checkpoint saved") - self.voice_status.setText( - f"stopped at {at}/{total} ({pct}%) — checkpoint saved, resume next time") + if p.get("during") == "transcoding": + msg = "stopped while transcoding — nothing processed yet" + self._task_finish(tid, ok=True, msg=msg) + self.voice_status.setText(msg) + else: + self._task_finish(tid, ok=True, + msg=f"stopped at {at}/{total} — checkpoint saved") + self.voice_status.setText( + f"stopped at {at}/{total} ({pct}%) — checkpoint saved, resume next time") self.ingest_progress.hide() self.ingest_stop_btn.setEnabled(True) self.ingest_stop_btn.hide() @@ -3805,6 +4580,8 @@ def run_gui(config: Config) -> int: if not voice_name: self._piper_status_lbl.setText("Select a profile above.") return + from ..voices.piper_train import SNAPSHOT_MINUTES, describe_export, export_info + from ..voices.training_scheduler import training_in_progress voice_dir = config.paths.voices_dir / voice_name onnx = voice_dir / "piper_model.onnx" meta = voice_dir / "training" / "metadata.csv" @@ -3814,12 +4591,50 @@ def run_gui(config: Config) -> int: if onnx.exists(): size_mb = onnx.stat().st_size / (1024 * 1024) status = f"✓ Trained model ready ({size_mb:.1f} MB) · {n_clips} training utterances" + info = export_info(onnx) + if info: + status += f"\nModel trained up to {describe_export(info)}" + exported_at = info.get("exported_at") + if isinstance(exported_at, (int, float)): + status += time.strftime(" · exported %Y-%m-%d %H:%M", + time.localtime(exported_at)) else: status = f"Not trained yet · {n_clips} training utterances available" + try: + training = training_in_progress(voice_dir) + except Exception: # noqa: BLE001 - only a status line + training = False + if training: + status += (f"\nTraining in progress — the model is updated with its progress " + f"every {SNAPSHOT_MINUTES:g} minutes.") if base_ckpt: status += f"\nBase checkpoint: {base_ckpt}" self._piper_status_lbl.setText(status) + def _make_active_voice(self, name: str) -> str: + """Make *name* the voice Javis speaks with, right away and after restarts. + + Returns a note for the user if it can't take effect as is. + """ + config.tts.default_voice = name + try: + cfg_path = config.source_path + if cfg_path is None: + raise RuntimeError("No config.yaml loaded — restart with a config file.") + update_config_file(cfg_path, {"tts": {"default_voice": name}}) + except Exception as exc: + log.exception("voice activation config save failed") + QtWidgets.QMessageBox.warning( + self, "Config save failed", + f"Voice set in memory but could not persist to config.yaml:\n{exc}" + ) + self._populate_voice_combo() + if config.tts.backend != "piper": + return " Set tts.backend to piper and restart Javis to hear it." + if (config.tts.piper_model or "").strip(): + return " Note: tts.piper_model in config.yaml takes precedence; clear it." + return "" + def _train_piper_model(self) -> None: """Launch Piper VITS training for the selected voice profile.""" name = self._selected_voice_name() @@ -3896,6 +4711,9 @@ def run_gui(config: Config) -> int: self.voice_status.setText(f"Launching Piper training for '{name}'…") import sys as _sys + + from ..voices.piper_train import SNAPSHOT_MINUTES + onnx_exists = (voice_dir / "piper_model.onnx").exists() cmd = [ _sys.executable, "-m", "javis.voices.piper_train", "--voice", name, @@ -3906,13 +4724,30 @@ def run_gui(config: Config) -> int: try: from ..core.config import find_project_root _launch_in_console(cmd, str(find_project_root())) + except Exception as exc: + log.exception("piper_train launch failed") + QtWidgets.QMessageBox.critical(self, "Launch failed", str(exc)) + return + # Training updates the voice's model as it goes: speak with it to hear + # the progress. (Only with the Piper backend: for others the profile + # is the voice to clone.) + if config.tts.backend == "piper": + note = self._make_active_voice(name) + when = ("from the next reply on" if onnx_exists else + "once training has exported its first model") + self.voice_status.setText( + f"Training '{name}' launched in a new terminal. It is the active voice " + f"now: Javis speaks with the model as trained so far, {when}; the model " + f"is updated every {SNAPSHOT_MINUTES:g} minutes.{note}" + ) + else: self.voice_status.setText( f"Training '{name}' launched in a new terminal. " "This takes several hours. The model appears in the voice profile when done." ) - except Exception as exc: - log.exception("piper_train launch failed") - QtWidgets.QMessageBox.critical(self, "Launch failed", str(exc)) + # Once the run has started, the status shows it. + QtCore.QTimer.singleShot( + 15_000, self, lambda: self._refresh_piper_status(self._selected_voice_name())) def _filter_narrator_clips(self) -> None: """Score training clips by speaker similarity and drop non-narrator ones.""" @@ -4090,24 +4925,10 @@ def run_gui(config: Config) -> int: self._export_piper_snapshot() return - # Update the in-memory config. - config.tts.default_voice = name - try: - cfg_path = config.source_path - if cfg_path is None: - raise RuntimeError("No config.yaml loaded — restart with a config file.") - update_config_file(cfg_path, {"tts": {"default_voice": name}}) - except Exception as exc: - log.exception("voice activation config save failed") - QtWidgets.QMessageBox.warning( - self, "Config save failed", - f"Voice set in memory but could not persist to config.yaml:\n{exc}" - ) - - # Reload TTS backend live if the orchestrator is running. - self.voice_status.setText( - f"\u2713 Active voice set to '{name}' \u2014 restart Javis to apply." - ) + note = self._make_active_voice(name) + if config.tts.backend == "piper" and not note: + note = " Javis speaks with it from the next reply on." + self.voice_status.setText(f"\u2713 Active voice set to '{name}'.{note}") self._refresh_piper_status(name) # ── reference clip management ────────────────────────────────────────── @@ -5068,6 +5889,8 @@ def run_gui(config: Config) -> int: async def _on_speaker_enrolled_ev(p): window.speakerEnrolled.emit(p or {}) async def _on_speaker_ask_name_ev(p): window.speakerAskName.emit(p or {}) async def _on_enroll_progress_ev(p): window.enrollProgress.emit(p or {}) + async def _on_unidentified_ev(p): window.unidentifiedChanged.emit(p or {}) + async def _on_voice_profiles_ev(p): window.voiceProfilesChanged.emit(p or {}) async def _on_timer_fired_ev(p): window.timerFired.emit(p or {}) bus.subscribe("state", _on_state) @@ -5094,6 +5917,8 @@ def run_gui(config: Config) -> int: bus.subscribe("speaker_enrolled", _on_speaker_enrolled_ev) bus.subscribe("speaker_ask_name", _on_speaker_ask_name_ev) bus.subscribe("enroll_progress", _on_enroll_progress_ev) + bus.subscribe("unidentified_changed", _on_unidentified_ev) + bus.subscribe("voice_profiles_changed", _on_voice_profiles_ev) bus.subscribe("timer_fired", _on_timer_fired_ev) # MCP approval dialog diff --git a/src/javis/stt/unidentified.py b/src/javis/stt/unidentified.py new file mode 100644 index 0000000..6bb18a6 --- /dev/null +++ b/src/javis/stt/unidentified.py @@ -0,0 +1,186 @@ +"""Inbox of voice clips from speakers that speaker ID could not match. + +Each clip is kept so the user can listen to it in the GUI and sort it into an +existing voice profile or use it to create a new one:: + + ~/.javis/voice_profiles/unidentified/{id}.wav 16-bit PCM audio + ~/.javis/voice_profiles/unidentified/{id}.npy ECAPA embedding (optional) + ~/.javis/voice_profiles/unidentified/index.json metadata, oldest first + +The inbox is capped at ``max_clips``; the oldest clips are dropped first. +All methods are thread-safe (clips are added from worker threads). +""" +from __future__ import annotations + +import json +import threading +import time +import uuid +from collections.abc import Iterable +from pathlib import Path + +import numpy as np + +from ..core.logger import get_logger + +log = get_logger("unidentified") + +_INDEX_FILE = "index.json" + + +def _mono(audio: np.ndarray) -> np.ndarray: + a = np.asarray(audio, dtype=np.float32) + if a.ndim > 1: + a = a.mean(axis=1) + return a.reshape(-1) + + +class UnidentifiedStore: + """Bounded on-disk store of unidentified speaker clips.""" + + def __init__(self, directory: Path, max_clips: int = 50) -> None: + self.dir = Path(directory) + self.max_clips = max(1, int(max_clips)) + self._lock = threading.RLock() + self._index: list[dict] | None = None # lazily loaded + + # ── index I/O ───────────────────────────────────────────────────────────── + + def _entries(self) -> list[dict]: + if self._index is None: + idx = self.dir / _INDEX_FILE + data: list = [] + if idx.exists(): + try: + data = json.loads(idx.read_text(encoding="utf-8")) + except Exception as exc: # noqa: BLE001 + log.warning(f"[unidentified] index read failed: {exc}") + self._index = [ + e for e in data + if isinstance(e, dict) and e.get("id") and e.get("file") + and (self.dir / e["file"]).exists() + ] + return self._index + + def _save(self) -> None: + self.dir.mkdir(parents=True, exist_ok=True) + idx = self.dir / _INDEX_FILE + tmp = idx.with_suffix(".json.part") + try: + tmp.write_text( + json.dumps(self._entries(), ensure_ascii=False, indent=2), + encoding="utf-8", + ) + tmp.replace(idx) + except Exception as exc: # noqa: BLE001 + log.error(f"[unidentified] index write failed: {exc}") + + def _delete_files(self, entry: dict) -> None: + for key in ("file", "embedding"): + fname = entry.get(key) + if not fname: + continue + try: + (self.dir / fname).unlink(missing_ok=True) + except Exception as exc: # noqa: BLE001 + log.warning(f"[unidentified] could not delete {fname}: {exc}") + + # ── public API ──────────────────────────────────────────────────────────── + + def add( + self, + audio: np.ndarray, + sample_rate: int, + *, + embedding: np.ndarray | None = None, + text: str = "", + source: str = "", + best_name: str | None = None, + best_score: float | None = None, + ) -> str: + """Store a clip and return its id. Prunes the oldest beyond the cap.""" + import soundfile as sf + + mono = np.clip(_mono(audio), -1.0, 1.0) + clip_id = f"{time.strftime('%Y%m%d-%H%M%S')}-{uuid.uuid4().hex[:6]}" + with self._lock: + self.dir.mkdir(parents=True, exist_ok=True) + wav_name = f"{clip_id}.wav" + sf.write(str(self.dir / wav_name), mono, int(sample_rate), subtype="PCM_16") + emb_name = None + if embedding is not None and np.any(embedding): + emb_name = f"{clip_id}.npy" + np.save(str(self.dir / emb_name), + np.asarray(embedding, dtype=np.float32).reshape(-1)) + entries = self._entries() + entries.append({ + "id": clip_id, + "file": wav_name, + "embedding": emb_name, + "created_at": time.strftime("%Y-%m-%dT%H:%M:%S"), + "duration_s": round(len(mono) / float(sample_rate), 2), + "sample_rate": int(sample_rate), + "text": (text or "").strip(), + "source": source or "", + "best_name": best_name, + "best_score": None if best_score is None else round(float(best_score), 4), + }) + while len(entries) > self.max_clips: + self._delete_files(entries.pop(0)) + self._save() + log.info(f"[unidentified] stored clip {clip_id} ({len(mono) / sample_rate:.1f}s)") + return clip_id + + def entries(self) -> list[dict]: + """All clip metadata, newest first (copies — safe to keep).""" + with self._lock: + return [dict(e) for e in reversed(self._entries())] + + def get(self, clip_id: str) -> dict | None: + with self._lock: + entry = next((e for e in self._entries() if e["id"] == clip_id), None) + return dict(entry) if entry else None + + def load_audio(self, clip_id: str) -> tuple[np.ndarray, int]: + """Return ``(float32 mono audio, sample_rate)``; raises ``KeyError`` if unknown.""" + import soundfile as sf + + entry = self.get(clip_id) + if entry is None: + raise KeyError(clip_id) + audio, sr = sf.read(str(self.dir / entry["file"]), dtype="float32", always_2d=False) + return _mono(audio), int(sr) + + def load_embedding(self, clip_id: str) -> np.ndarray | None: + entry = self.get(clip_id) + if entry is None or not entry.get("embedding"): + return None + try: + return np.load(str(self.dir / entry["embedding"])).astype(np.float32) + except Exception as exc: # noqa: BLE001 + log.warning(f"[unidentified] embedding for {clip_id} unreadable: {exc}") + return None + + def remove(self, clip_ids: Iterable[str]) -> int: + """Delete the given clips; returns how many were removed.""" + wanted = set(clip_ids) + if not wanted: + return 0 + with self._lock: + entries = self._entries() + gone = [e for e in entries if e["id"] in wanted] + if not gone: + return 0 + for e in gone: + self._delete_files(e) + self._index = [e for e in entries if e["id"] not in wanted] + self._save() + return len(gone) + + def clear(self) -> int: + with self._lock: + return self.remove([e["id"] for e in self._entries()]) + + def __len__(self) -> int: + with self._lock: + return len(self._entries()) diff --git a/src/javis/stt/voice_profiler.py b/src/javis/stt/voice_profiler.py index b1abc74..c04ff04 100644 --- a/src/javis/stt/voice_profiler.py +++ b/src/javis/stt/voice_profiler.py @@ -47,6 +47,16 @@ def _l2(v: np.ndarray) -> np.ndarray: return v / n if n > 0 else v +def _as_vector(embedding) -> Optional[np.ndarray]: + """Flatten *embedding* to float32; ``None`` for missing or all-zero vectors.""" + if embedding is None: + return None + vec = np.asarray(embedding, dtype=np.float32).reshape(-1) + if vec.size == 0 or not np.any(vec): + return None + return vec + + # Minimum utterance duration for a reliable embedding. Shorter clips are # passed through but logged as a warning — ECAPA was trained on ≥ 1 s segments. _MIN_EMBED_DURATION_S = 0.5 @@ -125,6 +135,28 @@ class NullProfiler: def identify(self, audio_np: np.ndarray, sr: int) -> Optional[tuple[str, float]]: # noqa: ARG002 return None + def identify_embedding(self, embedding) -> Optional[tuple[str, float]]: # noqa: ARG002 + return None + + def is_ready(self) -> bool: + return False + + def has_profile(self, name: str) -> bool: # noqa: ARG002 + return False + + def rank(self, embedding) -> list[tuple[str, float]]: # noqa: ARG002 + return [] + + def enroll_embeddings(self, name: str, embeddings: list) -> bool: # noqa: ARG002 + log.warning("[profiler] NullProfiler — speechbrain unavailable, enroll skipped") + return False + + def refine(self, name: str, clip: np.ndarray, sr: int) -> None: # noqa: ARG002 + pass + + def refine_embedding(self, name: str, embedding) -> bool: # noqa: ARG002 + return False + def list_profiles(self) -> list[dict]: return [] @@ -155,8 +187,13 @@ class VoiceProfiler: self.config = config self._profiles_dir: Path = config.paths.voice_profiles_dir self._encoder = None # speechbrain EncoderClassifier or None + self._device = "cpu" # device the encoder runs on + self._encoder_lock = threading.Lock() # serialises the CPU fallback self._profiles: dict[str, np.ndarray] = {} # name → L2-normed embedding self._index: list[dict] = [] # metadata records + # identify/refine run on worker threads while the GUI enrolls/deletes + # on the loop thread — serialise access to _profiles/_index. + self._lock = threading.RLock() # ── lifecycle ───────────────────────────────────────────────────────────── @@ -164,9 +201,18 @@ class VoiceProfiler: """Load the ECAPA-TDNN encoder and all enrolled profiles from disk.""" self._profiles_dir.mkdir(parents=True, exist_ok=True) self._load_encoder() - self._load_index() + with self._lock: + self._load_index() log.info(f"[profiler] ready — {len(self._profiles)} profile(s) loaded") + def is_ready(self) -> bool: + """``True`` once the ECAPA encoder is loaded (embeddings are meaningful).""" + return self._encoder is not None + + def has_profile(self, name: str) -> bool: + with self._lock: + return name in self._profiles + def _load_encoder(self) -> None: try: import os @@ -187,8 +233,6 @@ class VoiceProfiler: category=UserWarning, ) # 3. speechbrain Pretrainer symlinks — harmless on Windows because - self._device = "cpu" # device the encoder runs on - self._encoder_lock = threading.Lock() # serialises the CPU fallback # model files are copied rather than symlinked. warnings.filterwarnings( "ignore", @@ -210,6 +254,33 @@ class VoiceProfiler: ) self._encoder = None + def _build_encoder(self, device: str): + from speechbrain.inference.classifiers import EncoderClassifier + + savedir = str(self.config.paths.models_dir / "ecapa-tdnn") + log.info(f"[profiler] loading ECAPA-TDNN on {device} → {savedir}") + return EncoderClassifier.from_hparams( + source=_MODEL_SOURCE, + savedir=savedir, + run_opts={"device": device}, + ) + + def _fallback_to_cpu(self, failed, exc: BaseException): + """Swap a broken CUDA encoder for a CPU one (once, thread-safe). + + CUDA errors here usually mean the context died (e.g. the NVIDIA driver + was updated while Javis was running), which never recovers in-process. + """ + with self._encoder_lock: + if self._encoder is failed: + log.warning( + f"[profiler] CUDA failed ({exc}) — switching speaker ID to CPU; " + "restart Javis to use the GPU again" + ) + self._encoder = self._build_encoder("cpu") + self._device = "cpu" + return self._encoder + @staticmethod def _try_cuda() -> bool: try: @@ -254,33 +325,6 @@ class VoiceProfiler: # ── embedding ───────────────────────────────────────────────────────────── - def _build_encoder(self, device: str): - from speechbrain.inference.classifiers import EncoderClassifier - - savedir = str(self.config.paths.models_dir / "ecapa-tdnn") - log.info(f"[profiler] loading ECAPA-TDNN on {device} → {savedir}") - return EncoderClassifier.from_hparams( - source=_MODEL_SOURCE, - savedir=savedir, - run_opts={"device": device}, - ) - - def _fallback_to_cpu(self, failed, exc: BaseException): - """Swap a broken CUDA encoder for a CPU one (once, thread-safe). - - CUDA errors here usually mean the context died (e.g. the NVIDIA driver - was updated while Javis was running), which never recovers in-process. - """ - with self._encoder_lock: - if self._encoder is failed: - log.warning( - f"[profiler] CUDA failed ({exc}) — switching speaker ID to CPU; " - "restart Javis to use the GPU again" - ) - self._encoder = self._build_encoder("cpu") - self._device = "cpu" - return self._encoder - def embed(self, audio_np: np.ndarray, sr: int) -> np.ndarray: """Return a L2-normalised ECAPA-TDNN speaker embedding for *audio_np*.""" if self._encoder is None: @@ -328,25 +372,58 @@ class VoiceProfiler: if self._encoder is None: log.warning("[profiler] enroll skipped — encoder not loaded") return - vecs = [self.embed(c, sr) for c in clips] - avg = _l2(np.mean(vecs, axis=0)) + self.enroll_embeddings(name, [self.embed(c, sr) for c in clips]) - # Build a filesystem-safe filename from the profile name. - safe = "".join(c if (c.isalnum() or c in "-_") else "_" for c in name) - npy_path = self._profiles_dir / f"{safe}.npy" - np.save(str(npy_path), avg) - self._profiles[name] = avg + def enroll_embeddings(self, name: str, embeddings: list[np.ndarray]) -> bool: + """Create (or replace) *name*'s profile from precomputed embeddings. - # Replace existing entry if the name already exists. - self._index = [e for e in self._index if e["name"] != name] - self._index.append({ - "name": name, - "file": f"{safe}.npy", - "enrolled_at": time.strftime("%Y-%m-%dT%H:%M:%S"), - "clip_count": len(clips), - }) - self._save_index() - log.info(f"[profiler] enrolled '{name}' from {len(clips)} clip(s)") + All-zero vectors (from an unloaded encoder) are ignored. Returns + ``False`` when no usable embedding remains. + """ + vecs = [v for v in (_as_vector(e) for e in embeddings) if v is not None] + if not vecs: + log.warning(f"[profiler] enroll '{name}' skipped — no usable embeddings") + return False + avg = _l2(np.mean(vecs, axis=0)).astype(np.float32) + + with self._lock: + self._profiles_dir.mkdir(parents=True, exist_ok=True) + old = next((e for e in self._index if e["name"] == name), None) + if old is not None: + fname = old["file"] + else: + # Filesystem-safe filename; never reuse another profile's file + # ("Anna B" and "Anna_B" sanitise to the same stem). + safe = "".join(c if (c.isalnum() or c in "-_") else "_" for c in name) + used = {e["file"] for e in self._index} + fname, i = f"{safe}.npy", 2 + while fname in used: + fname, i = f"{safe}_{i}.npy", i + 1 + np.save(str(self._profiles_dir / fname), avg) + self._profiles[name] = avg + + # Replace existing entry if the name already exists. + self._index = [e for e in self._index if e["name"] != name] + self._index.append({ + "name": name, + "file": fname, + "enrolled_at": time.strftime("%Y-%m-%dT%H:%M:%S"), + "clip_count": len(vecs), + }) + self._save_index() + log.info(f"[profiler] enrolled '{name}' from {len(vecs)} clip(s)") + return True + + def rank(self, embedding) -> list[tuple[str, float]]: + """Cosine similarity of *embedding* to every profile, best first.""" + vec = _as_vector(embedding) + if vec is None: + return [] + with self._lock: + items = list(self._profiles.items()) + scores = [(n, _cosine(vec, e)) for n, e in items if e.shape == vec.shape] + scores.sort(key=lambda x: -x[1]) + return scores def identify( self, audio_np: np.ndarray, sr: int @@ -358,20 +435,19 @@ class VoiceProfiler: """ if self._encoder is None or not self._profiles: return None - query = self.embed(audio_np, sr) - threshold = self.config.speaker_id.threshold + return self.identify_embedding(self.embed(audio_np, sr)) - scores: dict[str, float] = { - name: _cosine(query, emb) for name, emb in self._profiles.items() - } - best_name = max(scores, key=scores.__getitem__) - best_score = scores[best_name] + def identify_embedding(self, embedding) -> Optional[tuple[str, float]]: + """Like :meth:`identify` for an embedding computed by :meth:`embed`.""" + ranked = self.rank(embedding) + if not ranked: + return None + threshold = self.config.speaker_id.threshold + best_name, best_score = ranked[0] # Always log all candidate scores so identification failures are visible # in normal (INFO-level) logs without needing to enable DEBUG mode. - score_str = ", ".join( - f"{n}={s:.3f}" for n, s in sorted(scores.items(), key=lambda x: -x[1]) - ) + score_str = ", ".join(f"{n}={s:.3f}" for n, s in ranked) if best_score >= threshold: log.info(f"[profiler] identified '{best_name}' [{score_str}] threshold={threshold}") return best_name, float(best_score) @@ -382,7 +458,8 @@ class VoiceProfiler: def list_profiles(self) -> list[dict]: """Return a list of profile metadata dicts from the index.""" - return list(self._index) + with self._lock: + return list(self._index) def refine(self, name: str, clip: np.ndarray, sr: int) -> None: """Blend a new voice sample into *name*'s existing profile. @@ -393,36 +470,50 @@ class VoiceProfiler: """ if self._encoder is None: return - entry = next((e for e in self._index if e["name"] == name), None) - if entry is None or name not in self._profiles: + if not self.has_profile(name): log.warning(f"[profiler] refine: no profile found for '{name}'") return - new_emb = self.embed(clip, sr) - n = entry.get("clip_count", 1) - current = self._profiles[name] - # Running centroid: (current * n + new) / (n+1), then L2-normalise. - blended = _l2(current * n + new_emb) - npy_path = self._profiles_dir / entry["file"] - np.save(str(npy_path), blended) - self._profiles[name] = blended - entry["clip_count"] = n + 1 - entry["last_seen"] = time.strftime("%Y-%m-%dT%H:%M:%S") - self._save_index() + self.refine_embedding(name, self.embed(clip, sr)) + + def refine_embedding(self, name: str, embedding) -> bool: + """Blend a precomputed embedding into *name*'s centroid (see :meth:`refine`).""" + new_emb = _as_vector(embedding) + if new_emb is None: + return False + with self._lock: + entry = next((e for e in self._index if e["name"] == name), None) + if entry is None or name not in self._profiles: + log.warning(f"[profiler] refine: no profile found for '{name}'") + return False + n = entry.get("clip_count", 1) + current = self._profiles[name] + if current.shape != new_emb.shape: + log.warning(f"[profiler] refine '{name}': embedding shape mismatch") + return False + # Running centroid: (current * n + new) / (n+1), then L2-normalise. + blended = _l2(current * n + new_emb).astype(np.float32) + np.save(str(self._profiles_dir / entry["file"]), blended) + self._profiles[name] = blended + entry["clip_count"] = n + 1 + entry["last_seen"] = time.strftime("%Y-%m-%dT%H:%M:%S") + self._save_index() log.debug(f"[profiler] refined '{name}' — centroid now over {n + 1} clip(s)") + return True def delete_profile(self, name: str) -> bool: """Remove *name*'s embedding from disk and memory.""" - entry = next((e for e in self._index if e["name"] == name), None) - if entry is None: - return False - npy = self._profiles_dir / entry["file"] - try: - npy.unlink(missing_ok=True) - except Exception as exc: - log.warning(f"[profiler] delete {npy.name} failed: {exc}") - self._index = [e for e in self._index if e["name"] != name] - self._profiles.pop(name, None) - self._save_index() + with self._lock: + entry = next((e for e in self._index if e["name"] == name), None) + if entry is None: + return False + npy = self._profiles_dir / entry["file"] + try: + npy.unlink(missing_ok=True) + except Exception as exc: + log.warning(f"[profiler] delete {npy.name} failed: {exc}") + self._index = [e for e in self._index if e["name"] != name] + self._profiles.pop(name, None) + self._save_index() log.info(f"[profiler] deleted profile '{name}'") return True @@ -437,10 +528,11 @@ class VoiceProfiler: def set_privileges(self, name: str, privileges: SpeakerPrivileges) -> None: """Persist updated *privileges* for *name* to ``index.json``.""" - entry = next((e for e in self._index if e["name"] == name), None) - if entry is None: - log.warning(f"[profiler] set_privileges: no profile found for '{name}'") - return - entry["privileges"] = privileges.to_dict() - self._save_index() + with self._lock: + entry = next((e for e in self._index if e["name"] == name), None) + if entry is None: + log.warning(f"[profiler] set_privileges: no profile found for '{name}'") + return + entry["privileges"] = privileges.to_dict() + self._save_index() log.info(f"[profiler] privileges updated for '{name}'") diff --git a/src/javis/tts/piper_backend.py b/src/javis/tts/piper_backend.py index d747701..4800f80 100644 --- a/src/javis/tts/piper_backend.py +++ b/src/javis/tts/piper_backend.py @@ -12,6 +12,11 @@ The model path is resolved in this order: 1. ``config.tts.piper_model`` (explicit path in config.yaml) 2. ``{voices_dir}/{tts.default_voice}/piper_model.onnx`` (active voice dir) +The model is reloaded by itself when that changes: when training updates the +voice's model (it does every few minutes, see ``javis.voices.piper_train``) or +another voice is activated. Until the new model is loaded, the previous one +keeps speaking. + Inference is synchronous (ONNX runtime) so it runs in a thread-pool executor to avoid blocking the asyncio event loop. @@ -23,6 +28,7 @@ from __future__ import annotations import asyncio import re import threading +import time from pathlib import Path from typing import Optional @@ -59,42 +65,42 @@ def _clean(text: str) -> str: class PiperBackend(TTSBackend): """Piper ONNX TTS backend.""" + # Synthesis looks for a changed model at most this often. + CHECK_SECONDS = 2.0 + def __init__(self, config: Config) -> None: self.config = config self._voice: Optional[object] = None # piper.voice.PiperVoice self._sample_rate: int = 22050 - self._lock = threading.Lock() + self._lock = threading.Lock() # one synthesis at a time; guards the swap + self._reloading = threading.Lock() + self._loaded: tuple | None = None # _files_signature() of the loaded model + self._failed: tuple | None = None # ... of a model that failed to load + self._checked = 0.0 # ── TTSBackend interface ────────────────────────────────────────────────── def load(self) -> None: """Load the ONNX model. Called once at startup.""" - try: - from piper.voice import PiperVoice # type: ignore[import] - except ImportError as exc: - raise RuntimeError( - "piper-tts is not installed. Run: pip install piper-tts" - ) from exc - model_path = self._resolve_model_path() - config_path = Path(str(model_path) + ".json") - - log.info(f"Loading Piper model: {model_path}") - self._voice = PiperVoice.load( - str(model_path), - config_path=str(config_path) if config_path.exists() else None, - use_cuda=False, # ONNX runtime picks CUDA automatically when available - ) - self._sample_rate = ( - self._voice.config.sample_rate # type: ignore[union-attr] - if hasattr(self._voice, "config") - else 22050 - ) + signature = _files_signature(model_path) + try: + voice = self._load_voice(model_path) + except BaseException: + self._failed = signature + raise + with self._lock: + self._install(voice, signature) log.info( f"Piper model loaded — sample_rate={self._sample_rate} Hz, " - f"language={getattr(getattr(self._voice, 'config', None), 'espeak_voice', '?')}" + f"language={getattr(getattr(voice, 'config', None), 'espeak_voice', '?')}" ) + def model_available(self) -> bool: + """Whether there is a model for :meth:`load` (other than one it failed to load).""" + model_path = self._model_path() + return model_path is not None and _files_signature(model_path) != self._failed + async def synthesize( self, text: str, @@ -108,13 +114,17 @@ class PiperBackend(TTSBackend): sample_rate=self._sample_rate, ) - audio = await asyncio.to_thread(self._synthesize_sync, text) - return SynthesisResult(audio=audio, sample_rate=self._sample_rate) + audio, sample_rate = await asyncio.to_thread(self._synthesize_sync, text) + return SynthesisResult(audio=audio, sample_rate=sample_rate) # ── Internal ────────────────────────────────────────────────────────────── - def _synthesize_sync(self, text: str) -> np.ndarray: - """Run Piper inference synchronously (called from thread-pool).""" + def _synthesize_sync(self, text: str) -> tuple[np.ndarray, int]: + """Run Piper inference synchronously (called from thread-pool). + + Returns the audio and its sample rate. + """ + self._check_for_new_model() if self._voice is None: raise RuntimeError("PiperBackend.load() has not been called.") @@ -128,31 +138,103 @@ class PiperBackend(TTSBackend): chunks: list[np.ndarray] = [] with self._lock: + sample_rate = self._sample_rate for chunk in self._voice.synthesize(text, syn_config): arr = chunk.audio_float_array if arr is not None and len(arr) > 0: chunks.append(np.asarray(arr, dtype=np.float32)) if not chunks: - return np.zeros(0, dtype=np.float32) - return np.concatenate(chunks) + return np.zeros(0, dtype=np.float32), sample_rate + return np.concatenate(chunks), sample_rate - def _resolve_model_path(self) -> Path: - """Return the ONNX model path from config or voice profile dir.""" + @staticmethod + def _load_voice(model_path: Path): + try: + from piper.voice import PiperVoice # type: ignore[import] + except ImportError as exc: + raise RuntimeError( + "piper-tts is not installed. Run: pip install piper-tts" + ) from exc + + config_path = Path(str(model_path) + ".json") + log.info(f"Loading Piper model: {model_path}") + return PiperVoice.load( + str(model_path), + config_path=str(config_path) if config_path.exists() else None, + use_cuda=False, # ONNX runtime picks CUDA automatically when available + ) + + def _install(self, voice, signature: tuple | None) -> None: + """Make *voice* the one that speaks (call with ``_lock`` held).""" + self._voice = voice + self._sample_rate = ( + voice.config.sample_rate # type: ignore[union-attr] + if hasattr(voice, "config") + else 22050 + ) + self._loaded, self._failed = signature, None + + def _check_for_new_model(self) -> None: + """Start loading the model in the background if it changed since it was loaded.""" + if self._voice is None or time.monotonic() - self._checked < self.CHECK_SECONDS: + return + self._checked = time.monotonic() + model_path = self._model_path() + if model_path is None: + return + signature = _files_signature(model_path) + if signature is None or signature in (self._loaded, self._failed): + return + if not self._reloading.acquire(blocking=False): + return # being loaded already + try: + threading.Thread(target=self._reload, args=(model_path, signature), + name="piper-reload", daemon=True).start() + except BaseException: + self._reloading.release() + raise + + def _reload(self, model_path: Path, signature: tuple) -> None: + try: + voice = self._load_voice(model_path) + for _ in voice.synthesize("Bereit."): # warm up before it speaks + pass + except Exception as exc: # noqa: BLE001 - keep speaking with the loaded model + self._failed = signature + log.warning(f"Could not load the changed Piper model {model_path} ({exc}); " + "keeping the current one.") + else: + with self._lock: + self._install(voice, signature) + log.info(f"Piper model reloaded: {model_path}") + finally: + self._reloading.release() + + def _model_path(self, *, warn: bool = False) -> Path | None: + """The model to use (see the module docstring), or ``None`` if there is none.""" explicit = self.config.tts.piper_model.strip() if explicit: p = Path(explicit) if p.exists(): return p - log.warning(f"piper_model path not found: {p}") + if warn: + log.warning(f"piper_model path not found: {p}") # Fall back to active voice profile directory. + candidate = (self.config.paths.voices_dir / self.config.tts.default_voice + / "piper_model.onnx") + return candidate if candidate.exists() else None + + def _resolve_model_path(self) -> Path: + """Return the ONNX model path from config or voice profile dir.""" + model_path = self._model_path(warn=True) + if model_path is not None: + return model_path + voice_name = self.config.tts.default_voice voices_dir = self.config.paths.voices_dir candidate = voices_dir / voice_name / "piper_model.onnx" - if candidate.exists(): - return candidate - from ..voices.piper_train import has_checkpoints if has_checkpoints(voices_dir / voice_name): hint = (f"Training checkpoints exist — export them with: python -m " @@ -163,3 +245,16 @@ class PiperBackend(TTSBackend): f"Piper ONNX model not found. Expected at {candidate} " f"or set tts.piper_model in config.yaml. {hint}" ) + + +def _files_signature(model_path: Path) -> tuple | None: + """What tells a changed model apart: path, modification time and size of its files.""" + try: + st = model_path.stat() + except OSError: + return None + try: + config_mtime = Path(str(model_path) + ".json").stat().st_mtime_ns + except OSError: + config_mtime = None + return str(model_path), st.st_mtime_ns, st.st_size, config_mtime diff --git a/src/javis/voices/audiobook_ingest.py b/src/javis/voices/audiobook_ingest.py index ce1b3e1..d80018a 100644 --- a/src/javis/voices/audiobook_ingest.py +++ b/src/javis/voices/audiobook_ingest.py @@ -30,6 +30,7 @@ import json import os import shutil import sys +from collections.abc import Callable from dataclasses import dataclass, field, asdict from fractions import Fraction from pathlib import Path @@ -106,6 +107,11 @@ _SEGMENTS_FILE = "ingest_segments.npy" # Source audio is transcoded once to a flat PCM WAV (mono, TARGET_SR) so that # VAD and per-clip reads never decode compressed audio again. _CACHE_RAW = "ingest_audio_cache.i16" # raw int16 LE, mono, TARGET_SR Hz +# Free space that must remain on the cache drive; transcoding refuses to start +# or aborts mid-way below this. +_MIN_FREE_BYTES = 1024 * 1024 * 1024 +# Re-check free space every N 10-second blocks (~10 min of audio, ~26 MB cache). +_DISK_CHECK_EVERY_BLOCKS = 60 # Save checkpoint every N processed segments (not iterations — skipped ones # don't count) to limit how much work is lost on an unexpected stop. _CHECKPOINT_SAVE_INTERVAL = 20 @@ -328,55 +334,6 @@ def _is_clean_ref_clip(clip: np.ndarray, sr: int) -> bool: return True -def _load_audio(path: Path, target_sr: int = TARGET_SR) -> np.ndarray: - # soundfile 0.12+ ships libsndfile 1.2.2 which reads MP3 natively. - # OGG Vorbis and FLAC were always supported. - # M4A / AAC / OPUS / WMA are not supported by libsndfile — convert to - # MP3 or WAV first. - import soundfile as sf # type: ignore[import] - - info = sf.info(str(path)) - sr = info.samplerate - - # MP3 files with malformed ID3 tags (e.g. non-syncsafe CTOC frames) cause - # libsndfile to report INT64_MAX for the frame count. Calling sf.read() - # then tries to pre-allocate that many samples and raises a ValueError. - # Fall back to chunked streaming via sf.blocks() which reads until EOF - # without pre-allocating based on the (bogus) total-frame count. - # - # Memory note: we use a small blocksize (10 s) and resample + mono-mix - # each block immediately so that only the TARGET_SR data accumulates - # (roughly 2× smaller than keeping raw source-rate stereo chunks). - if info.frames >= sys.maxsize // 2: - log.warning(f"_load_audio: unknown frame count for {path.name} — using chunked read") - frac = Fraction(target_sr, sr).limit_denominator(1000) - resampled_chunks: list[np.ndarray] = [] - for block in sf.blocks(str(path), blocksize=sr * 10, dtype="float32", always_2d=True): - # Mono-mix immediately so we don't carry stereo data any longer than needed. - mono = block.mean(axis=1) if block.ndim > 1 else block.ravel() - chunk = ( - scipy.signal.resample_poly(mono, frac.numerator, frac.denominator).astype(np.float32) - if sr != target_sr else mono.astype(np.float32) - ) - resampled_chunks.append(chunk) - del block, mono # free source-rate data immediately - if not resampled_chunks: - raise ValueError(f"No audio data read from {path}") - audio: np.ndarray = np.concatenate(resampled_chunks) - del resampled_chunks - return audio # already mono, already at target_sr - else: - audio, _ = sf.read(str(path), dtype="float32", always_2d=True) - - if audio.ndim > 1: - audio = audio.mean(axis=1) - - if sr != target_sr: - frac = Fraction(target_sr, sr).limit_denominator(1000) - audio = scipy.signal.resample_poly(audio, frac.numerator, frac.denominator).astype(np.float32) - return audio - - def _segment(audio: np.ndarray, sr: int) -> list[tuple[int, int]]: """Return list of (start_sample, end_sample) speech segments via silero-vad.""" from silero_vad import load_silero_vad, get_speech_timestamps # heavy optional dep @@ -402,11 +359,21 @@ def _save_wav(path: Path, audio: np.ndarray, sr: int) -> None: sf.write(str(path), audio, sr, subtype="PCM_16") -def _transcode_to_raw(source: Path, target: Path, target_sr: int) -> None: +class TranscodeCancelled(Exception): + """Raised from the transcode thread when the caller requested a stop.""" + + +def _transcode_to_raw( + source: Path, + target: Path, + target_sr: int, + should_stop: Callable[[], bool] | None = None, +) -> None: """Stream *source* → raw int16 LE binary (mono, *target_sr* Hz). - Peak RAM: one 10-second block (~4 MiB). No file-size limit. - Writes atomically via a .part temp file. + Peak RAM: one 10-second block (~4 MiB). Writes atomically via a .part + temp file, re-checks free disk space while writing and raises + :class:`TranscodeCancelled` as soon as *should_stop* returns True. """ import soundfile as sf @@ -417,11 +384,11 @@ def _transcode_to_raw(source: Path, target: Path, target_sr: int) -> None: tmp = target.with_suffix(".part") # Pre-check disk space when duration is known (skip for malformed MP3s with - # bogus INT64_MAX frame counts). + # bogus INT64_MAX frame counts — the in-loop check below still applies). if info.frames < sys.maxsize // 2: estimated_bytes = int(info.frames / sr * target_sr) * 2 # int16 = 2 B free = shutil.disk_usage(str(target.parent)).free - if free < estimated_bytes + 64 * 1024 * 1024: + if free < estimated_bytes + _MIN_FREE_BYTES: raise RuntimeError( f"Not enough disk space to cache ‘{source.name}’: " f"need ~{estimated_bytes / 1_000_000_000:.1f} GB, " @@ -429,9 +396,18 @@ def _transcode_to_raw(source: Path, target: Path, target_sr: int) -> None: ) try: - with open(str(tmp), "wb") as out_f: - for block in sf.blocks(str(source), blocksize=sr * 10, - dtype="float32", always_2d=True): + # Read until SoundFile.read() returns 0 frames instead of using + # sf.blocks(): MP3s with malformed ID3 tags report INT64_MAX frames, + # and sf.blocks() ignores the real read count, so after EOF it keeps + # re-yielding its stale buffer forever (filled 800 GB once). + with sf.SoundFile(str(source)) as src, open(str(tmp), "wb") as out_f: + n_blocks = 0 + while True: + if should_stop is not None and should_stop(): + raise TranscodeCancelled(source.name) + block = src.read(sr * 10, dtype="float32", always_2d=True) + if len(block) == 0: + break mono = block.mean(axis=1) if block.ndim > 1 else block.ravel() if sr != target_sr: chunk = scipy.signal.resample_poly( @@ -442,6 +418,15 @@ def _transcode_to_raw(source: Path, target: Path, target_sr: int) -> None: i16 = np.clip(chunk * 32767.0, -32768, 32767).astype(np.int16) out_f.write(i16.tobytes()) del block, mono, chunk, i16 + n_blocks += 1 + if n_blocks % _DISK_CHECK_EVERY_BLOCKS == 0: + free = shutil.disk_usage(str(target.parent)).free + if free < _MIN_FREE_BYTES: + raise RuntimeError( + f"Aborted caching ‘{source.name}’: free space on " + f"{target.parent} dropped to {free / 1_000_000_000:.1f} GB " + f"after writing {out_f.tell() / 1_000_000_000:.1f} GB" + ) tmp.replace(target) except Exception: try: @@ -599,6 +584,14 @@ async def ingest_audiobook( # peak RAM per block, no file-size limit). VAD and clip reads operate on # this cache via numpy memmap — O(1) seek, ~90 MiB peak during segmentation. cache_raw = prof.dir / _CACHE_RAW + should_stop = stop_event.is_set if stop_event is not None else None + + async def _stopped_while_transcoding() -> IngestStats: + log.info(f"ingest stopped by user while transcoding {source.name}") + await _emit({"phase": "stopped", "during": "transcoding", + "at_index": 0, "total": 0}) + return stats + if cp is None: # Fresh run — remove any stale cache from a previous source. if cache_raw.exists(): @@ -608,7 +601,11 @@ async def ingest_audiobook( pass log.info(f"transcoding {source.name} → raw cache ...") await _emit({"phase": "loading", "source": source.name}) - await asyncio.to_thread(_transcode_to_raw, source, cache_raw, TARGET_SR) + try: + await asyncio.to_thread(_transcode_to_raw, source, cache_raw, TARGET_SR, + should_stop) + except TranscodeCancelled: + return await _stopped_while_transcoding() _cache_dur = cache_raw.stat().st_size / 2 / TARGET_SR # int16 = 2 bytes log.info(f"transcoded: {_cache_dur:.1f}s") await _emit({"phase": "loaded", "audio_s": _cache_dur}) @@ -616,7 +613,11 @@ async def ingest_audiobook( # Resuming but cache was deleted — retranscode. log.info(f"raw cache missing — retranscoding {source.name} ...") await _emit({"phase": "loading", "source": source.name}) - await asyncio.to_thread(_transcode_to_raw, source, cache_raw, TARGET_SR) + try: + await asyncio.to_thread(_transcode_to_raw, source, cache_raw, TARGET_SR, + should_stop) + except TranscodeCancelled: + return await _stopped_while_transcoding() log.info("retranscoded OK") if cp is not None: diff --git a/src/javis/voices/piper_compat.py b/src/javis/voices/piper_compat.py index d79c564..c136bca 100644 --- a/src/javis/voices/piper_compat.py +++ b/src/javis/voices/piper_compat.py @@ -9,18 +9,89 @@ ``onnxscript`` and rejects piper's ``dynamic_axes`` call. * A run resumed with ``--ckpt_path`` after a warm start would re-apply the base weights on top of the restored ones. +* piper's monotonic alignment search needs a Cython extension that the + piper-tts wheel doesn't ship (building it needs a C compiler), and a stale + build of it returned broken alignments — frames mapped to no phoneme or to + hundreds — so the KL loss was inf and no generator step succeeded, at any + precision. :func:`maximum_path` replaces it with a numba port of the same + algorithm (NumPy without numba). + +Training additions (``fit``): + +* ``bf16-mixed`` falls back to 32-bit on devices without native bf16. +* Training batches group clips of similar length (:class:`LengthBucketBatchSampler`). + Piper's random batches pad every clip to the longest of 32 random clips — + ~2.4× the real audio frames on a typical audiobook dataset. Padded lengths + are rounded up (:class:`RoundedCollate`), so cuDNN plans its convolutions + for a few dozen shapes instead of a new one nearly every step. Set + ``JAVIS_PIPER_BUCKETING=0`` to get piper's random batches back. +* cuDNN times its convolution algorithms during the first training step and + keeps the fastest (:class:`CudnnFirstStepTuning`). Set + ``JAVIS_PIPER_CUDNN_TUNING=0`` to skip that. +* :class:`NonFiniteGuard` skips optimizer steps with NaN/inf gradients (bf16 has + no GradScaler to do it) — on the GPU, through fused AdamW, so it never makes + the CPU wait — stops runs that keep diverging instead of letting them burn + GPU-hours without learning, and never saves NaN weights. +* Piper's training step read values back from the GPU ~460 times per step (a + Python loop over the batch to cut out training windows, debug and logging + reads, boolean masks in the duration predictor's spline flows), and each read + stalls the CPU until the GPU has caught up. :func:`_avoid_gpu_syncs` swaps in + equivalents that stay on the GPU, and pins batches so their upload no longer + blocks either. + +Reliability (``fit``): + +* Checkpoints are written beside their final name and renamed into place, so a + crash, power cut or full disk never leaves a truncated ``last.ckpt``; failed + saves fail the run instead of passing silently (:func:`_safe_checkpoint_files`). +* ``last.ckpt`` is refreshed every epoch, not only in epochs whose val_mel + improved, and a run no longer aborts when val_mos can't be computed + (:func:`_checkpoint_every_epoch`). +* :class:`ProgressSaver` saves the epoch so far to ``last.ckpt`` every + ``JAVIS_PIPER_PROGRESS_MINUTES`` (default 10; 0 turns that off), and before + the run stops: when the launcher creates its stop file or has gone away + (:class:`StopOnRequest`), or the console window is closed. The resumed run + continues with the epoch's next batch (:class:`ResumableDataLoader`). +* Preprocessed cache files are written atomically. An unreadable one is + deleted — the next run rebuilds it — and another utterance stands in for it + (:class:`RobustSubset`). +* Exit codes: 0 finished, :data:`~javis.voices.training_control.EXIT_STOPPED` + stopped after saving, :data:`~javis.voices.training_control.EXIT_DIVERGED` + NaN/inf (restarting would not help); anything else is a crash. Usage:: python -m javis.voices.piper_compat fit ... # = python -m piper.train fit ... python -m javis.voices.piper_compat export ... # = python -m piper.train.export_onnx ... + +``export --info-file X`` also writes which checkpoint (path, mtime) it exported, +with its epoch and step, to X (JSON), and refuses checkpoints with NaN/inf weights. """ +import contextlib +import dataclasses +import errno import functools import inspect +import io +import json +import math +import os +import random import sys +import time +import zlib +from collections.abc import Iterator, Sequence +import numpy as np +import torch +from lightning.pytorch.callbacks import Callback +from piper.train.vits import dataset as _piper_dataset from piper.train.vits import models as _models from pysilero_vad import SileroVoiceActivityDetector +from torch.nn import functional as F +from torch.utils.data import DataLoader, Sampler, Subset + +from javis.voices import training_control if not hasattr(SileroVoiceActivityDetector, "process_array"): SileroVoiceActivityDetector.process_array = SileroVoiceActivityDetector.process_samples @@ -64,13 +135,1329 @@ def _ignore_warmstart_on_resume() -> None: VitsModel.__init__ = _init +# ── Monotonic alignment search ────────────────────────────────────────────── + +_MAX_NEG = -1e9 + + +def _mas_path(path, value, t_y: int, t_x: int) -> None: + """Monotonic alignment of one item, as in VITS' ``core.pyx``. + + *value* (frames × phonemes) is overwritten with the DP table and ``path[y]`` + receives the phoneme of frame *y* (where ``core.pyx`` sets ``path[y, index] = 1``). + """ + index = t_x - 1 + for y in range(t_y): + for x in range(max(0, t_x + y - t_y), min(t_x, y + 1)): + v_cur = _MAX_NEG if x == y else value[y - 1, x] + if x == 0: + v_prev = 0.0 if y == 0 else _MAX_NEG + else: + v_prev = value[y - 1, x - 1] + value[y, x] += max(v_prev, v_cur) + for y in range(t_y - 1, -1, -1): + path[y] = index + if index != 0 and (index == y or value[y - 1, index] < value[y - 1, index - 1]): + index -= 1 + + +def _mas_numpy(paths, values, t_ys, t_xs) -> None: + """:func:`_mas_path` for a whole batch, vectorised over items and phonemes.""" + b, frames, phonemes = values.shape + x = np.arange(phonemes) + t_y, t_x = t_ys.astype(np.int64)[:, None], t_xs.astype(np.int64)[:, None] + neg = np.float32(_MAX_NEG) + for y in range(frames): + band = (x >= np.maximum(0, t_x + y - t_y)) & (x < np.minimum(t_x, y + 1)) + if y == 0: + best = np.where(x == 0, np.float32(0.0), neg) + else: + prev = values[:, y - 1] + v_cur = np.where(x == y, neg, prev) + v_prev = np.concatenate([np.full((b, 1), neg, np.float32), prev[:, :-1]], axis=1) + best = np.maximum(v_prev, v_cur) + values[:, y] += np.where(band, best, np.float32(0.0)) + rows = np.arange(b) + index = t_x[:, 0] - 1 + for y in range(frames - 1, -1, -1): + on = y < t_y[:, 0] + paths[on, y] = index[on] + if y == 0: + break + prev = values[:, y - 1] + lower = prev[rows, index] < prev[rows, np.maximum(index - 1, 0)] + index = index - (on & (index != 0) & ((index == y) | lower)) + + +_mas_kernel = None + + +def _batch_mas(): + """The batch kernel: numba-compiled (parallel over the batch) when available, else NumPy.""" + global _mas_kernel + if _mas_kernel is None: + try: + import numba + except ImportError: + _mas_kernel = _mas_numpy + else: + each = numba.njit(nogil=True)(_mas_path) + + @numba.njit(parallel=True, nogil=True) + def _mas_numba(paths, values, t_ys, t_xs): + for i in numba.prange(paths.shape[0]): + each(paths[i], values[i], t_ys[i], t_xs[i]) + + _mas_kernel = _mas_numba + return _mas_kernel + + +def maximum_path(neg_cent: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + """Drop-in for piper's ``monotonic_align.maximum_path``. + + *neg_cent* and *mask* are ``[batch, frames, phonemes]``; returns the 0/1 + alignment mapping every valid frame to one phoneme. Only the phoneme index + of each frame goes back to the device, not the dense matrix. + """ + values = neg_cent.detach().to("cpu", torch.float32, copy=True).numpy() # overwritten + m = mask.detach().float() + t_ys, t_xs = torch.stack([m.sum(1)[:, 0], m.sum(2)[:, 0]]).round().to("cpu", torch.int32).numpy() + phoneme = np.full(values.shape[:2], -1, dtype=np.int32) # padding frames stay -1 + _batch_mas()(phoneme, values, t_ys, t_xs) + ids = torch.arange(values.shape[2], dtype=torch.int32, device=neg_cent.device) + return (torch.from_numpy(phoneme).to(neg_cent.device).unsqueeze(2) == ids).to(neg_cent.dtype) + + +def _fix_monotonic_alignment() -> None: + try: + from piper.train.vits import monotonic_align + except ImportError: # extension not built: stand in for the whole module + import types + + from piper.train import vits + + monotonic_align = types.ModuleType("piper.train.vits.monotonic_align") + sys.modules[monotonic_align.__name__] = vits.monotonic_align = monotonic_align + monotonic_align.maximum_path = maximum_path + + +_fix_monotonic_alignment() + + +# ── Precision ─────────────────────────────────────────────────────────────── + +def _cuda_bf16_supported() -> bool: + if not torch.cuda.is_available(): + return False + try: + return torch.cuda.is_bf16_supported(including_emulation=False) + except TypeError: # torch < 2.3 has no including_emulation + return torch.cuda.is_bf16_supported() + + +def _resolve_precision(argv: list, *, bf16_supported: bool | None = None) -> None: + """Swap ``--trainer.precision bf16-*`` for 32-bit when bf16 isn't native here. + + Runs in the training process (which initialises CUDA anyway) rather than in + the launcher, so the launcher never holds a CUDA context of its own. + """ + for i, arg in enumerate(argv): + if arg == "--trainer.precision" and i + 1 < len(argv): + pos, prefix = i + 1, "" + elif arg.startswith("--trainer.precision="): + pos, prefix = i, "--trainer.precision=" + else: + continue + value = argv[pos][len(prefix):] + if not value.startswith("bf16"): + return + if bf16_supported is None: + on_cpu = any(a == "--trainer.accelerator=cpu" + or (a == "--trainer.accelerator" and j + 1 < len(argv) + and argv[j + 1] == "cpu") + for j, a in enumerate(argv)) + bf16_supported = not on_cpu and _cuda_bf16_supported() + if not bf16_supported: + print(f"NOTE: precision {value} needs a GPU with native bf16 (RTX 30xx or " + "newer); training in 32-bit instead.", flush=True) + argv[pos] = prefix + "32" + return + + +# ── Length-bucketed training batches ────────────────────────────────────────── + +class _EpochSetter: + """Lightning calls ``set_epoch`` on ``batch_sampler.sampler``, never on the batch sampler.""" + + def __init__(self, owner: "LengthBucketBatchSampler") -> None: + self._owner = owner + + def set_epoch(self, epoch: int) -> None: + self._owner.set_epoch(epoch) + + +class LengthBucketBatchSampler(Sampler): + """Batches of similar-length clips, reshuffled every epoch ("sortish" batching). + + Each epoch shuffles all clips, sorts pools of ``batch_size * pool_batches`` + clips by length, cuts them into batches and shuffles the batch order. Every + batch then holds clips of about the same length, so little padding is + computed, while batch composition and order still change every epoch. The + order depends only on (*seed*, epoch). + + :meth:`state_dict` records how far the current epoch got, so a run resumed + from a checkpoint saved mid-epoch trains the batches that epoch hadn't yet. + """ + + def __init__(self, lengths: Sequence[int], batch_size: int, *, drop_last: bool, + pool_batches: int = 50, seed: int = 0, longest_first: bool = True) -> None: + if batch_size < 1: + raise ValueError(f"batch_size must be >= 1, got {batch_size}") + self.lengths = list(lengths) + self.batch_size = batch_size + self.drop_last = drop_last + self.pool_size = batch_size * max(1, pool_batches) + self.seed = seed + self.epoch = 0 + # Fail fast: the first epoch of a run starts with its longest batch, so + # a batch size that doesn't fit runs out of memory within seconds + # instead of hours into the run. + self.longest_first = longest_first + self._first_epoch: int | None = None + self._order: tuple | None = None # (epoch, order of batches(epoch), fingerprint) + self._resume: tuple | None = None # (epoch, batches done, order) from a checkpoint + self.sampler = _EpochSetter(self) + + def set_epoch(self, epoch: int) -> None: + self.epoch = epoch + if self._resume is not None and self._resume[0] != epoch: + self._resume = None + + def __len__(self) -> int: + n = len(self.lengths) + return n // self.batch_size if self.drop_last else -(-n // self.batch_size) + + def batches(self, epoch: int) -> list: + """The batches (lists of dataset indices) for *epoch*, in training order.""" + rng = random.Random(self.seed * 1_000_003 + epoch) + order = list(range(len(self.lengths))) + rng.shuffle(order) + if self.drop_last: # drop random clips, not the longest ones of a pool + del order[len(order) - len(order) % self.batch_size:] + batches = [] + for start in range(0, len(order), self.pool_size): + pool = sorted(order[start:start + self.pool_size], key=self.lengths.__getitem__) + batches += [pool[i:i + self.batch_size] for i in range(0, len(pool), self.batch_size)] + rng.shuffle(batches) + return batches + + def __iter__(self) -> Iterator[list]: + # A DataLoader with workers calls __iter__ twice when it starts an epoch: + # both calls must return the same batches. + batches = self.batches(self.epoch) + if self._first_epoch is None: + self._first_epoch = self.epoch + order, done = list(range(len(batches))), 0 + if self._resume is not None and self._resume[0] == self.epoch: + done, order = self._resume[1], list(self._resume[2]) + if self.longest_first and self.epoch == self._first_epoch and done < len(order): + longest = max(range(done, len(order)), + key=lambda k: max(self.lengths[i] for i in batches[order[k]])) + order.insert(done, order.pop(longest)) + self._order = (self.epoch, order, _fingerprint(batches)) + return iter([batches[k] for k in order[done:]]) + + def state_dict(self, batches_done: int) -> dict: + """The current epoch's batch order, *batches_done* of which are trained.""" + if self._order is None: + return {} + epoch, order, fingerprint = self._order + return {"epoch": epoch, "batches_done": int(batches_done), "order": list(order), + "fingerprint": fingerprint} + + def load_state_dict(self, state: dict) -> None: + """Have the epoch saved in *state* continue where it was interrupted.""" + self._resume = None + try: + epoch, done = int(state["epoch"]), int(state["batches_done"]) + order = [int(k) for k in state["order"]] + fingerprint = int(state["fingerprint"]) + except (KeyError, TypeError, ValueError): + return # no epoch in progress (or a checkpoint of piper's random batches) + if not 0 < done < len(order): + return + batches = self.batches(epoch) + if sorted(order) != list(range(len(batches))) or fingerprint != _fingerprint(batches): + print("NOTE: the training data changed since the checkpoint was saved; the rest " + "of the interrupted epoch uses a new batch order.", flush=True) + return + # Lightning creates the first iterator of a resumed run before it calls + # set_epoch, so the sampler must already be in the interrupted epoch. + self.epoch = epoch + self._resume = (epoch, done, order) + + +def _fingerprint(batches: list) -> int: + return zlib.crc32(repr(batches).encode()) + + +class ResumableDataLoader(DataLoader): + """A DataLoader whose :class:`LengthBucketBatchSampler` position goes into checkpoints. + + Lightning saves and restores the state of loaders with ``state_dict`` and + ``load_state_dict`` methods; *batches_done* says how many batches of the + current epoch the trainer has completed. + """ + + def __init__(self, *args, batches_done=lambda: 0, **kwargs) -> None: + super().__init__(*args, **kwargs) + self.batches_done = batches_done + + def state_dict(self) -> dict: + return self.batch_sampler.state_dict(self.batches_done()) + + def load_state_dict(self, state: dict) -> None: + self.batch_sampler.load_state_dict(state) + + +def _batches_done(datamodule) -> int: + """Batches of the current epoch that *datamodule*'s trainer has completed.""" + try: + return int(datamodule.trainer.fit_loop.epoch_loop.batch_progress.current.ready) + except (AttributeError, TypeError, ValueError): + return 0 + + +def padding_factor(batches: Sequence[Sequence[int]], lengths: Sequence[int]) -> float: + """Padded frames computed per real frame when each batch pads to its longest clip.""" + real = sum(lengths[i] for b in batches for i in b) + padded = sum(len(b) * max(lengths[i] for i in b) for b in batches if b) + return padded / real if real else 1.0 + + +def _pad_to_multiple(t: torch.Tensor, multiple: int) -> torch.Tensor: + extra = -t.shape[-1] % multiple + return F.pad(t, (0, extra)) if extra else t + + +class RoundedCollate: + """piper's collate function, with each batch's padded length rounded up. + + cuDNN plans its convolutions per input shape, and planning a new shape costs + ~3 ms of CPU time per convolution layer. Nearly every piper batch has + lengths of its own; padded up to multiples of 32 spectrogram frames and 16 + phonemes, a dataset needs a few dozen shapes, each planned once. Training + steps got ~4% faster despite ~3% more padding. VITS masks everything past + each clip's length. + """ + + def __init__(self, collate, spectrogram_multiple: int = 32, phoneme_multiple: int = 16) -> None: + self.collate = collate + self.spectrogram_multiple = spectrogram_multiple + self.phoneme_multiple = phoneme_multiple + + def __call__(self, utterances): + batch = self.collate(utterances) + batch.spectrograms = _pad_to_multiple(batch.spectrograms, self.spectrogram_multiple) + batch.phoneme_ids = _pad_to_multiple(batch.phoneme_ids, self.phoneme_multiple) + return batch + + +def _clip_lengths(dataset) -> list: + """Length proxy per item: size of its cached float32 audio tensor (no loading).""" + indices = list(range(len(dataset))) + while isinstance(dataset, Subset): # random_split() wraps the dataset + indices = [dataset.indices[i] for i in indices] + dataset = dataset.dataset + lengths = [] + for i in indices: + try: + lengths.append(os.stat(dataset.utts[i].audio_norm_path).st_size) + except OSError: # missing cache file: fails at load time as it would unbucketed + lengths.append(0) + return lengths + + +def _bucketed_train_dataloader(datamodule): + from piper.train.vits.dataset import UtteranceCollate + + started = time.perf_counter() + dataset = datamodule.train_dataset + lengths = _clip_lengths(dataset) + seed = int(os.environ.get("PL_GLOBAL_SEED", "0")) + sampler = LengthBucketBatchSampler(lengths, datamodule.batch_size, + drop_last=len(dataset) > datamodule.batch_size, + seed=seed) + kwargs = { + "batch_sampler": sampler, + "collate_fn": RoundedCollate(UtteranceCollate(is_multispeaker=datamodule.num_speakers > 1, + segment_size=datamodule.segment_size)), + "num_workers": datamodule.num_workers, + "pin_memory": True, + } + if datamodule.num_workers > 0: + kwargs.update(persistent_workers=True, prefetch_factor=4) + loader = ResumableDataLoader(dataset, batches_done=functools.partial(_batches_done, datamodule), + **kwargs) + + order = list(range(len(lengths))) + random.Random(seed).shuffle(order) + unbucketed = [order[i:i + datamodule.batch_size] + for i in range(0, len(order), datamodule.batch_size)] + print(f"Length-bucketed batches: {len(sampler)} per epoch, padding " + f"{padding_factor(sampler.batches(0), lengths):.2f}x of the real audio " + f"(random batches: {padding_factor(unbucketed, lengths):.2f}x) " + f"[{time.perf_counter() - started:.1f} s]", flush=True) + return loader + + +# ── NaN/inf guard ───────────────────────────────────────────────────────────── + +class TrainingDivergedError(RuntimeError): + """Raised by :class:`NonFiniteGuard` when a run keeps producing NaN/inf.""" + + +def _total_norm(tensors: list) -> torch.Tensor: + get_total_norm = getattr(torch.nn.utils, "get_total_norm", None) + if get_total_norm is not None: # torch >= 2.6 + return get_total_norm(tensors) + return torch.linalg.vector_norm(torch.stack([torch.linalg.vector_norm(t) for t in tensors])) + + +_OPTIMIZER_NAMES = {0: ("generator", "g"), 1: ("discriminator", "d")} + + +@functools.cache +def _fused_adam_skips_on(device: torch.device) -> bool: + """Whether fused AdamW runs on *device* and leaves weights and state untouched + when its ``found_inf`` flag is set — the way GradScaler skips overflowed steps.""" + try: + param = torch.ones(1, device=device, requires_grad=True) + param.grad = torch.ones(1, device=device) + optimizer = torch.optim.AdamW([param], lr=1.0, fused=True) + optimizer.found_inf = torch.ones((), device=device) + optimizer.step() + return bool(param.detach().eq(1).all()) and float(optimizer.state[param]["step"]) == 0 + except Exception: # noqa: BLE001 - whatever fails, fused AdamW can't skip steps here + return False + + +def _use_fused_adam(optimizer) -> bool: + """Switch a torch Adam/AdamW to its fused kernel, which can skip a step on the GPU. + + Returns whether it did. Works on restored state too: foreach keeps each + parameter's ``step`` count on the CPU, fused on the parameter's device. + """ + if type(optimizer) not in (torch.optim.Adam, torch.optim.AdamW): + return False + groups = optimizer.param_groups + params = [p for group in groups for p in group["params"]] + if (not params or len({p.device for p in params}) > 1 + or not all(p.is_floating_point() for p in params) + or any(group.get("capturable") or group.get("differentiable") for group in groups) + or not _fused_adam_skips_on(params[0].device)): + return False + for group in groups: + group["fused"], group["foreach"] = True, False + for p in group["params"]: + state = optimizer.state.get(p) + if state and "step" in state: + state["step"] = torch.as_tensor(state["step"], dtype=torch.float32, device=p.device) + return True + + +def _remedy(trainer) -> str: + precision = str(getattr(trainer, "precision", "")) + if precision.startswith("16"): + return "Train with --precision bf16-mixed or 32, or check the data." + if precision.startswith("bf16"): + return "Check the data, or try --precision 32." + return "Check the data." + + +class NonFiniteGuard(Callback): + """Skip optimizer steps with NaN/inf gradients and stop runs that keep diverging. + + bf16 training has no GradScaler, so one bad batch would otherwise write NaN + into the weights. fp16's GradScaler does skip such steps, but carries on + even when *every* step is skipped and its scale has collapsed — the run then + keeps saving checkpoints while learning nothing. + + Deciding on the CPU whether to skip means waiting for the GPU to finish the + backward pass, twice per training step — ~6% of the training time. So + piper's AdamW runs fused and the GPU skips bad steps itself, as it does for + GradScaler; the counts are read back every few steps, early enough to stop at + exactly ``max_consecutive``. Other optimizers get the check on the CPU. + """ + + def __init__(self, max_consecutive: int = 100, min_fp16_scale: float = 1.0, + check_every: int = 25) -> None: + self.max_consecutive = max_consecutive + self.min_fp16_scale = min_fp16_scale + self.check_every = check_every + self.skipped: dict = {} # optimizer index -> skipped steps in total + self._streak: dict = {} # optimizer index -> consecutive skipped steps + self._fused: set = set() # indices of optimizers that skip steps on the GPU + self._counts: dict = {} # optimizer index -> on the GPU: [skipped in total, in a row] + self._until_check = 0 + + def on_train_start(self, trainer, pl_module) -> None: + # Lightning has restored the checkpoint's optimizer state, kernel choice included. + self._fused = {i for i, opt in enumerate(trainer.optimizers) if _use_fused_adam(opt)} + self._until_check = min(self.check_every, self.max_consecutive) + + def on_before_optimizer_step(self, trainer, pl_module, optimizer) -> None: + params = [p for group in optimizer.param_groups for p in group["params"] + if p.grad is not None] + if not params: + return + idx = next((i for i, opt in enumerate(trainer.optimizers) if opt is optimizer), -1) + finite = torch.isfinite(_total_norm([p.grad for p in params])) + if idx in self._fused: + bad = (~finite).float() + optimizer.found_inf = bad # 1: fused AdamW leaves weights and state alone + counts = self._counts.get(idx) + if counts is None: + counts = self._counts[idx] = bad.new_zeros(2) + counts[1].mul_(bad) # a good step ends the streak + counts.add_(bad) + return + if bool(finite): + self._streak[idx] = 0 + return + for p in params: + p.grad = None # optimizers leave parameters without a gradient untouched + self.skipped[idx] = self.skipped.get(idx, 0) + 1 + streak = self._streak[idx] = self._streak.get(idx, 0) + 1 + if streak == 1 or streak % 25 == 0 or streak >= self.max_consecutive: + self._report(trainer, idx, streak) + + def _report(self, trainer, idx: int, streak: int) -> None: + name = _OPTIMIZER_NAMES.get(idx, (f"optimizer {idx}",))[0] + print(f"WARNING: step {trainer.global_step}: skipped {name} steps with NaN/inf " + f"gradients ({self.skipped[idx]} in total, {streak} in a row)", + file=sys.stderr, flush=True) + if streak >= self.max_consecutive: + raise TrainingDivergedError( + f"{streak} consecutive {name} steps had NaN/inf gradients; training has " + f"diverged. {_remedy(trainer)}") + + def _read_back(self, trainer) -> None: + indices = sorted(self._counts) + rows = torch.stack([self._counts[i] for i in indices]).tolist() # the one GPU sync + longest = 0 + for idx, (skipped, streak) in zip(indices, rows): + skipped, streak = int(skipped), int(streak) + self._streak[idx] = streak + longest = max(longest, streak) + if skipped > self.skipped.get(idx, 0): + self.skipped[idx] = skipped + self._report(trainer, idx, streak) + # A streak grows by at most one per step: look again before it can reach the limit. + self._until_check = max(1, min(self.check_every, self.max_consecutive - longest)) + + def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx) -> None: + if self._counts: + self._until_check -= 1 + if self._until_check <= 0: + self._read_back(trainer) + if self.skipped: + pl_module.log_dict({f"skipped_steps_{_OPTIMIZER_NAMES.get(i, ('', str(i)))[1]}": + float(n) for i, n in self.skipped.items()}) + scaler = getattr(trainer.precision_plugin, "scaler", None) + if scaler is not None and scaler.is_enabled() and batch_idx % 10 == 0: + scale = scaler.get_scale() + if not (math.isfinite(scale) and scale >= self.min_fp16_scale): + raise TrainingDivergedError( + f"the fp16 loss scale collapsed to {scale:g}; every optimizer step is " + "being skipped. Train with --precision bf16-mixed or 32.") + + def on_train_end(self, trainer, pl_module) -> None: + if self._counts: + self._read_back(trainer) + + def on_save_checkpoint(self, trainer, pl_module, checkpoint) -> None: + params = [p.detach() for p in pl_module.parameters()] + if params and not bool(torch.isfinite(_total_norm(params))): + raise TrainingDivergedError("refusing to save a checkpoint with NaN/inf weights") + + +# ── Fewer GPU round trips ───────────────────────────────────────────────────── +# A training step is CPU-bound: Python queues ~17k kernels per step and the GPU +# idles over a third of the time. Piper's step also read values back from the +# GPU ~460 times; each read blocks the CPU until the GPU has caught up — e.g. +# right after the vocoder's forward and backward passes, the heaviest GPU work — +# instead of letting it queue the next kernels meanwhile. These drop-in +# equivalents compute the same values without reading anything back. + +def slice_segments(x: torch.Tensor, ids_str, segment_size: int = 4) -> torch.Tensor: + """``x[i, :, ids_str[i]:ids_str[i] + segment_size]`` for every item, as one gather. + + Piper's loop reads every start back from the GPU (3 syncs per item, 3 calls + per step). Negative starts count as 0, as in piper; a window running past + the end repeats the last frame where piper's loop would fail. + """ + width = min(segment_size, x.size(2)) + starts = torch.as_tensor(ids_str, device=x.device).long().clamp_min(0) + index = (starts.view(-1, 1) + torch.arange(width, device=x.device)).clamp_max(x.size(2) - 1) + return torch.gather(x, 2, index.unsqueeze(1).expand(-1, x.size(1), -1)) + + +def rand_slice_segments(x: torch.Tensor, x_lengths=None, segment_size: int = 4): + """Piper's random training windows, with the starts drawn on the device.""" + b, _, t = x.size() + if x_lengths is None: + x_lengths = t + ids_str = (torch.rand(b, device=x.device) * (x_lengths - segment_size + 1)).long() + return slice_segments(x, ids_str, segment_size), ids_str + + +def discriminator_loss(disc_real_outputs, disc_generated_outputs): + """Piper's discriminator loss, with the per-discriminator terms left on the device. + + Piper turned every term into a float (one sync each) for lists its training + step discards; they are detached tensors here. + """ + loss = 0 + r_losses, g_losses = [], [] + for dr, dg in zip(disc_real_outputs, disc_generated_outputs): + r_loss = torch.mean((1 - dr.float()) ** 2) + g_loss = torch.mean(dg.float() ** 2) + loss = loss + (r_loss + g_loss) + r_losses.append(r_loss.detach()) + g_losses.append(g_loss.detach()) + return loss, r_losses, g_losses + + +_hann_windows: dict = {} + + +def mel_spectrogram_torch(y: torch.Tensor, n_fft: int, num_mels: int, sampling_rate: int, + hop_size: int, win_size: int, fmin, fmax, center: bool = False): + """Piper's mel spectrogram without its debug check of the audio range. + + The check (``if torch.min(y) < -1``) blocked right after the vocoder's forward + pass, whose output ends in tanh, so it never fired. + """ + from piper.train.vits.mel_processing import spec_to_mel_torch + + key = (win_size, y.dtype, y.device) + window = _hann_windows.get(key) + if window is None: + window = _hann_windows[key] = torch.hann_window(win_size).to(y) + pad = (n_fft - hop_size) // 2 + y = F.pad(y.unsqueeze(1), (pad, pad), mode="reflect").squeeze(1) + spec = torch.stft(y, n_fft, hop_length=hop_size, win_length=win_size, window=window, + center=center, pad_mode="reflect", normalized=False, onesided=True, + return_complex=True) + spec = torch.sqrt(torch.view_as_real(spec).pow(2).sum(-1) + 1e-6) + return spec_to_mel_torch(spec, n_fft, num_mels, sampling_rate, fmin, fmax) + + +def _mask_free_spline(original): + """Piper's spline flow for training, with ``torch.where`` instead of boolean masks. + + Indexing with a mask reads the mask's size back from the GPU — 9 syncs per + flow in the forward pass and as many in the backward. Evaluating the spline + everywhere (inputs clamped into its interval) and keeping the identity + outside it gives the same values and gradients. Inference keeps piper's code. + """ + from piper.train.vits.transforms import DEFAULT_MIN_DERIVATIVE, rational_quadratic_spline + + @functools.wraps(original) + def spline(inputs, unnormalized_widths, unnormalized_heights, unnormalized_derivatives, + inverse=False, tails="linear", tail_bound=1.0, **kwargs): + if inverse or tails != "linear": + return original(inputs, unnormalized_widths, unnormalized_heights, + unnormalized_derivatives, inverse=inverse, tails=tails, + tail_bound=tail_bound, **kwargs) + # Linear tails: the derivative at both ends of the interval is fixed at 1. + min_derivative = kwargs.get("min_derivative", DEFAULT_MIN_DERIVATIVE) + derivatives = F.pad(unnormalized_derivatives, (1, 1), + value=math.log(math.exp(1 - min_derivative) - 1)) + inside = (inputs >= -tail_bound) & (inputs <= tail_bound) + outputs, logabsdet = rational_quadratic_spline( + inputs=inputs.clamp(-tail_bound, tail_bound), + unnormalized_widths=unnormalized_widths, + unnormalized_heights=unnormalized_heights, + unnormalized_derivatives=derivatives, + left=-tail_bound, right=tail_bound, bottom=-tail_bound, top=tail_bound, **kwargs) + return (torch.where(inside, outputs, inputs).to(inputs.dtype), + torch.where(inside, logabsdet, 0.0).to(inputs.dtype)) + + spline._javis_mask_free = True + return spline + + +def _flip_forward(self, x: torch.Tensor, *args, reverse: bool = False, **kwargs): + """Piper's ``Flip.forward``, with its zero log-determinant made on the device.""" + x = torch.flip(x, [1]) + if reverse: + return x + return x, x.new_zeros(x.size(0)) + + +def _pin_batch(self): + """``Batch.pin_memory``: lets the DataLoader's ``pin_memory=True`` pin piper's batches. + + Unpinned, every batch was uploaded synchronously at the start of its step. + """ + pinned = {} + for field in dataclasses.fields(self): + value = getattr(self, field.name) + if isinstance(value, torch.Tensor): + pinned[field.name] = value.pin_memory() + return dataclasses.replace(self, **pinned) + + +def _avoid_gpu_syncs() -> None: + from piper.train.vits import commons, losses, modules, transforms + from piper.train.vits import lightning as vits_lightning + from piper.train.vits.dataset import Batch + + # lightning.py imported these by name, so its references are replaced too. + commons.slice_segments = vits_lightning.slice_segments = slice_segments + commons.rand_slice_segments = rand_slice_segments + losses.discriminator_loss = vits_lightning.discriminator_loss = discriminator_loss + vits_lightning.mel_spectrogram_torch = mel_spectrogram_torch + spline = transforms.unconstrained_rational_quadratic_spline + if not getattr(spline, "_javis_mask_free", False): + transforms.unconstrained_rational_quadratic_spline = _mask_free_spline(spline) + modules.Flip.forward = _flip_forward + if not hasattr(Batch, "pin_memory"): + Batch.pin_memory = _pin_batch + + +# ── Convolution algorithms ──────────────────────────────────────────────────── + +class CudnnFirstStepTuning(Callback): + """Time cuDNN's convolution algorithms during the first training step only. + + In benchmark mode cuDNN times its candidate algorithms for each new input + shape and keeps the fastest. That pays off for the decoder and the + discriminators, which always see the same fixed-size training windows, but + left on it times every new batch length again: steps took ~8× as long. + Switched off after the first step, cuDNN keeps the timed choices and uses its + heuristics for shapes it hasn't seen: steps got ~7% faster for ~30 s of + timing at the start of a run. + """ + + def __init__(self) -> None: + self._restore: bool | None = None + + def on_train_start(self, trainer, pl_module) -> None: + self._restore = torch.backends.cudnn.benchmark # True if asked for (--trainer.benchmark) + torch.backends.cudnn.benchmark = True + + def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx) -> None: + if self._restore is not None: + torch.backends.cudnn.benchmark, self._restore = self._restore, None + + +# ── Files that survive a crash ──────────────────────────────────────────────── + +PARTIAL_SUFFIX = training_control.PARTIAL_SUFFIX + + +def _write_atomically(path, write, *, durable: bool) -> None: + """Write *path* by calling ``write(temporary path)`` and renaming the result into place. + + Readers — and the next run after a crash or power cut — see the complete + old file or the complete new one, never a truncated one. *durable* flushes + the data to disk before the rename (worth it for checkpoints, not for cache files). + """ + path = os.fspath(path) + partial = path + PARTIAL_SUFFIX + try: + write(partial) + if durable: + fd = os.open(partial, os.O_RDWR | getattr(os, "O_BINARY", 0)) + try: + os.fsync(fd) + finally: + os.close(fd) + training_control.replace_file(partial, path) + except BaseException: + with contextlib.suppress(OSError): + os.remove(partial) + raise + if durable and sys.platform != "win32": # make the rename itself durable + with contextlib.suppress(OSError): + fd = os.open(os.path.dirname(path) or ".", os.O_RDONLY) + try: + os.fsync(fd) + finally: + os.close(fd) + + +def remove_partial_files(directory) -> None: + """Delete ``*.partial`` files in *directory*: whatever a killed run was writing.""" + try: + with os.scandir(directory) as entries: + partial = [e.path for e in entries if e.name.endswith(PARTIAL_SUFFIX)] + except OSError: + return + for path in partial: + with contextlib.suppress(OSError): + os.remove(path) + + +def _safe_checkpoint_files() -> None: + """Write checkpoints beside their final name, then rename them into place. + + Lightning's "atomic" save buffers the whole checkpoint in memory (~1 GB + here), writes it to the system temp folder and moves it to the training + drive — a *copy* over the old file when the drives differ, so a crash or + full disk mid-copy truncated ``last.ckpt``. It also ignored a + ``PermissionError``: a failed save went unnoticed and the old checkpoint + stayed. Failed saves now fail the run, which the launcher restarts. + """ + from lightning.fabric.plugins.io import torch_io + from lightning.fabric.utilities.cloud_io import _is_local_file_protocol + + atomic_save = torch_io._atomic_save + if getattr(atomic_save, "_javis", False): + return + + @functools.wraps(atomic_save) + def _atomic_save(checkpoint, filepath) -> None: + if not _is_local_file_protocol(filepath): + return atomic_save(checkpoint, filepath) + _write_atomically(filepath, functools.partial(torch.save, checkpoint), durable=True) + + _atomic_save._javis = True + torch_io._atomic_save = _atomic_save + + remove = torch_io.TorchCheckpointIO.remove_checkpoint + + @functools.wraps(remove) + def _remove_checkpoint(self, path) -> None: + deadline = time.monotonic() + 10 + while True: + try: + return remove(self, path) + except PermissionError as exc: # open elsewhere for a moment + if time.monotonic() >= deadline: + print(f"WARNING: could not delete the old checkpoint {path}: {exc}", + file=sys.stderr, flush=True) + return + time.sleep(0.5) + + torch_io.TorchCheckpointIO.remove_checkpoint = _remove_checkpoint + + +_missing_monitors: set = set() + + +def _checkpoint_every_epoch() -> None: + """Keep ``last.ckpt`` current, and carry on when a monitored metric is missing. + + Lightning refreshes ``last.ckpt`` only when the same callback has just saved + one of its top-k checkpoints — with piper's callbacks, only in epochs whose + val_mel improved. Once val_mel flattens out, a crash or stop threw away + every epoch since the last improvement. + + piper's second callback keeps the best checkpoints by val_mos, which is never + logged when the MOS predictor couldn't be loaded: Lightning then aborted the + run at the end of its first epoch. + """ + from lightning.pytorch.callbacks import ModelCheckpoint + + if getattr(ModelCheckpoint.on_train_epoch_end, "_javis", False): + return + needed = ("_should_skip_saving_checkpoint", "_should_save_on_train_epoch_end", + "_save_last_checkpoint", "_monitor_candidates", "_save_topk_checkpoint") + if not all(hasattr(ModelCheckpoint, name) for name in needed): + print("WARNING: Lightning's ModelCheckpoint has changed; last.ckpt is only " + "refreshed in epochs whose val_mel improved.", file=sys.stderr, flush=True) + return + + def keeping_last_current(hook, at_epoch_end: bool): + @functools.wraps(hook) + def wrapper(self, trainer, pl_module) -> None: + due = (self.save_last is True and not self._should_skip_saving_checkpoint(trainer) + and self._should_save_on_train_epoch_end(trainer) == at_epoch_end) + hook(self, trainer, pl_module) + if due and self._last_global_step_saved != trainer.global_step: + self._save_last_checkpoint(trainer, self._monitor_candidates(trainer)) + + wrapper._javis = True + return wrapper + + ModelCheckpoint.on_train_epoch_end = keeping_last_current( + ModelCheckpoint.on_train_epoch_end, True) + ModelCheckpoint.on_validation_end = keeping_last_current( + ModelCheckpoint.on_validation_end, False) + + save_topk = ModelCheckpoint._save_topk_checkpoint + + @functools.wraps(save_topk) + def _save_topk_checkpoint(self, trainer, monitor_candidates) -> None: + if self.monitor is not None and self.save_top_k != 0 \ + and self.monitor not in monitor_candidates: + loops = getattr(getattr(trainer, "fit_loop", None), "epoch_loop", None) + validated = getattr(getattr(loops, "val_loop", None), "_has_run", False) + if validated and self.monitor not in _missing_monitors: + _missing_monitors.add(self.monitor) + print(f"WARNING: {self.monitor} was not logged, so no checkpoints are kept " + f"by {self.monitor}; the others (and last.ckpt) still are.", + file=sys.stderr, flush=True) + return + save_topk(self, trainer, monitor_candidates) + + ModelCheckpoint._save_topk_checkpoint = _save_topk_checkpoint + + +# ── Stopping on request ─────────────────────────────────────────────────────── + +def _say(message: str) -> None: + """Print to a console that may be gone: its window was closed.""" + try: + print(message, flush=True) + except (OSError, ValueError, SystemExit): # rich turns a broken pipe into SystemExit + pass + + +class StopTraining(Exception): + """Training stopped on request: after saving its progress, or before there was any.""" + + +class StopOnRequest(Callback): + """Stop when asked to, after saving the progress (:class:`ProgressSaver` does that mid-epoch). + + The launcher asks by creating the file named in ``JAVIS_PIPER_STOP_FILE``. + The run stops the same way when the launcher (``JAVIS_PIPER_LAUNCHER``) has + gone — killed, or Javis crashed — instead of training on unsupervised. + """ + + def __init__(self, stop_file: str | None = None, + launcher: tuple[int, int] | None = None) -> None: + self.stop_file = stop_file + self.launcher = launcher + self.reason: str | None = None + self._next_poll = 0.0 + + @classmethod + def from_environment(cls) -> "StopOnRequest | None": + stop_file = os.environ.get(training_control.STOP_FILE_ENV) or None + launcher = training_control.launcher_from_environment() + return cls(stop_file, launcher) if stop_file or launcher else None + + def requested(self, *, force: bool = False) -> bool: + """Whether to stop; looks at most once a second unless *force*d.""" + if self.reason is None and (force or time.monotonic() >= self._next_poll): + self._next_poll = time.monotonic() + 1.0 + if self.stop_file and os.path.exists(self.stop_file): + self.reason = "stop requested" + elif self.launcher and not training_control.process_alive(*self.launcher): + self.reason = "the process that started training has exited" + if self.reason: + _say(f"\n{self.reason[0].upper()}{self.reason[1:]}: saving the training " + "progress, then stopping…") + return self.reason is not None + + def raise_if_requested(self) -> None: + if self.requested(): + raise StopTraining(self.reason) + + def setup(self, trainer, pl_module, stage) -> None: + # Before the data loaders, the model and the sanity check start up. + if self.requested(force=True): + raise StopTraining(self.reason) + + def on_train_start(self, trainer, pl_module) -> None: + if self.requested(force=True): + raise StopTraining(self.reason) + + def on_train_epoch_end(self, trainer, pl_module) -> None: + # Asked during validation. Runs before the checkpoint callbacks: they + # still save this epoch. + if self.requested(force=True): + trainer.should_stop = True + + +def _progress_minutes() -> float: + """``JAVIS_PIPER_PROGRESS_MINUTES``: how often to save the epoch so far (0: never).""" + try: + minutes = float(os.environ.get("JAVIS_PIPER_PROGRESS_MINUTES", "10")) + except ValueError: + return 10.0 + return minutes if math.isfinite(minutes) and minutes > 0 else 0.0 + + +class ProgressSaver(Callback): + """Save the epoch so far to ``last.ckpt``: every *minutes*, and before stopping. + + An epoch of a large dataset takes a quarter of an hour or more, and a run + used to restart at the beginning of the epoch it was stopped or crashed in. + Now a stop (:class:`StopOnRequest`) or a closed console window + (:class:`~javis.voices.training_control.ConsoleClose`) saves after the + current step and ends the run; a crash loses at most *minutes* of training. + The resumed run continues with the next batch of that epoch + (:class:`ResumableDataLoader`). + + Never saves after the last batch of an epoch: that checkpoint would resume + into the epoch's end and skip its validation. The checkpoint callbacks save + there anyway, once validation has run. + """ + + def __init__(self, stopper: StopOnRequest | None = None, console=None, + minutes: float = 10.0) -> None: + self.stopper = stopper + self.console = console + self.interval = max(0.0, minutes) * 60 + self._last_save = time.monotonic() + self._warned = False + + def on_train_start(self, trainer, pl_module) -> None: + self._last_save = time.monotonic() + + def on_save_checkpoint(self, trainer, pl_module, checkpoint) -> None: + self._last_save = time.monotonic() + + def _closing(self) -> bool: + return self.console is not None and self.console.closing.is_set() + + def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx) -> None: + closing = self._closing() + stop = closing or (self.stopper is not None and self.stopper.requested()) + last = trainer.num_training_batches - 1 + if not math.isfinite(last) or batch_idx >= last: + if stop: + trainer.should_stop = True # validates and saves at the end of this epoch + return + age = time.monotonic() - self._last_save + # Before the last batch: validation and the end-of-epoch save take a while. + due = self.interval > 0 and (age >= self.interval + or (batch_idx == last - 1 and age >= self.interval / 2)) + if not (stop or due): + return + try: + saved = self._save(trainer) + finally: + # Also when closing began during a save that was merely due: that save + # holds all the progress, and there's no time for another step and save. + closing = self._closing() + if closing: + self.console.done() + stop = stop or closing + if saved and not closing: + _say(f"Saved the training progress: epoch {trainer.current_epoch}, " + f"batch {batch_idx + 1}/{last + 1}.") + if stop and saved: + raise StopTraining(self.console.reason if closing else self.stopper.reason) + if stop: + trainer.should_stop = True # ends the epoch early, validates and saves + + def _save(self, trainer) -> bool: + callback = next((cb for cb in trainer.checkpoint_callbacks + if getattr(cb, "save_last", False)), None) + if callback is None or not hasattr(callback, "_save_last_checkpoint"): + if not self._warned: + self._warned = True + _say("WARNING: no checkpoint callback keeps last.ckpt; progress within an " + "epoch is not saved.") + return False + # Lightning's own way to save last.ckpt, so the checkpoint callback's + # state stays consistent with the file. + callback._save_last_checkpoint(trainer, callback._monitor_candidates(trainer)) + return True + + +_OUT_OF_MEMORY = getattr(torch, "OutOfMemoryError", None) or torch.cuda.OutOfMemoryError + + +def _fit(main, stopper: StopOnRequest | None = None, console=None) -> int: + """Run piper's trainer and return the exit code for the launcher. + + 0: finished; ``EXIT_STOPPED``: stopped on request, after saving; + ``EXIT_DIVERGED``: restarting would not help. A crash raises (exit code 1). + """ + try: + main() + except StopTraining: + return training_control.EXIT_STOPPED + except TrainingDivergedError as exc: + print(f"\nERROR: {exc}", file=sys.stderr, flush=True) + return training_control.EXIT_DIVERGED + except _OUT_OF_MEMORY: + print("\nERROR: the GPU ran out of memory. If this keeps happening, train with a " + "smaller --batch-size, or close other programs that use the GPU.", + file=sys.stderr, flush=True) + raise + if (stopper is not None and stopper.reason) \ + or (console is not None and console.closing.is_set()): + return training_control.EXIT_STOPPED + return 0 + + +# ── Damaged cache files ─────────────────────────────────────────────────────── + +class CorruptCacheFile(RuntimeError): + """A preprocessed cache file was unreadable; it was deleted, to be rebuilt.""" + + +def _unreadable_content(exc: BaseException) -> bool: + """Whether *exc*, raised by ``torch.load``, means the file's content is damaged. + + Truncated or garbled files raise all sorts of errors — RuntimeError, + EOFError, UnpicklingError, UnicodeDecodeError — or, for a truncated zip, + ``OSError(EINVAL)``. Other OSErrors (missing, locked by a virus scanner) + and running out of memory say nothing about the file. + """ + if isinstance(exc, OSError): + return exc.errno == errno.EINVAL + if isinstance(exc, MemoryError) or not isinstance(exc, Exception): + return False + text = str(exc).lower() + return "out of memory" not in text and "not enough memory" not in text + + +_on_cache_write = None # called after each cache file prepare_data writes + + +class _CacheTorch: + """``torch``, as piper's dataset module sees it. + + Cache files are written atomically: a run killed while preprocessing used to + leave a truncated file that every later run tripped over, as ``prepare_data`` + only checks that a file exists. A damaged file is deleted when it fails to + load, so ``prepare_data`` rebuilds it. + """ + + def __getattr__(self, name): + return getattr(torch, name) + + @staticmethod + def save(obj, f, *args, **kwargs): + if not isinstance(f, (str, os.PathLike)): + return torch.save(obj, f, *args, **kwargs) + _write_atomically(f, lambda tmp: torch.save(obj, tmp, *args, **kwargs), durable=False) + if _on_cache_write is not None: + _on_cache_write() + + @staticmethod + def load(f, *args, **kwargs): + try: + return torch.load(f, *args, **kwargs) + except Exception as exc: + if not isinstance(f, (str, os.PathLike)) or not _unreadable_content(exc): + raise + with contextlib.suppress(OSError): + os.remove(f) + raise CorruptCacheFile(f"unreadable cache file {os.fspath(f)} (deleted; it is " + f"rebuilt when training starts next): {exc!r}") from exc + + +if not isinstance(_piper_dataset.torch, _CacheTorch): + # Module level: DataLoader workers re-import this module, not its __main__ block. + _piper_dataset.torch = _CacheTorch() + +_unreadable_utterances = 0 + + +def _report_unreadable(exc: BaseException) -> None: + global _unreadable_utterances + _unreadable_utterances += 1 + if _unreadable_utterances <= 10 or _unreadable_utterances % 100 == 0: + print(f"WARNING: used another utterance in place of one that can't be read " + f"({_unreadable_utterances} so far in this process): {exc}", + file=sys.stderr, flush=True) + + +class RobustSubset(Subset): + """A ``random_split`` part that substitutes a readable utterance for an unreadable one. + + A damaged or missing cache file would otherwise end the run at the batch + holding it — and every restarted run at the same place. + """ + + max_substitutes = 20 + + def __getitem__(self, idx): + if isinstance(idx, list): + return [self[i] for i in idx] + n = len(self.indices) + first = range(n)[idx] # IndexError past the end, which ends piper's loop over the test split + error = None + for k in range(min(n, self.max_substitutes + 1)): + try: + return self.dataset[self.indices[(first + k) % n]] + except (CorruptCacheFile, OSError) as exc: + error = error or exc + _report_unreadable(exc) + raise error + + def __getitems__(self, indices): + return [self[i] for i in indices] + + +def _robust_dataset(datamodule_class) -> None: + """Make piper's data module rebuild damaged cache files and work around unreadable ones.""" + if getattr(datamodule_class.setup, "_javis", False): + return + setup, prepare_data = datamodule_class.setup, datamodule_class.prepare_data + + @functools.wraps(prepare_data) + def _prepare_data(self): + remove_partial_files(self.cache_dir) + for attempt in range(1, 51): + try: + return prepare_data(self) + except CorruptCacheFile as exc: # deleted: the next pass rebuilds it + if attempt == 50: + raise + print(f"WARNING: {exc}", file=sys.stderr, flush=True) + + @functools.wraps(setup) + def _setup(self, stage): + setup(self, stage) + for name in ("train_dataset", "val_dataset", "test_dataset"): + part = getattr(self, name, None) + if type(part) is Subset: + setattr(self, name, RobustSubset(part.dataset, part.indices)) + + _setup._javis = True + datamodule_class.prepare_data = _prepare_data + datamodule_class.setup = _setup + + +def _install_training_aids(console=None) -> StopOnRequest | None: + """Patch piper's training; returns the :class:`StopOnRequest` callback, if a launcher set one up. + + *console*: a :class:`~javis.voices.training_control.ConsoleClose` whose + window closing saves the progress, like a stop. + """ + from piper.train.vits.dataset import VitsDataModule + from piper.train.vits.lightning import VitsModel + + global _on_cache_write + _safe_checkpoint_files() + _checkpoint_every_epoch() + _robust_dataset(VitsDataModule) + stopper = StopOnRequest.from_environment() + # Preprocessing a new dataset takes a while; stop between two of its files. + _on_cache_write = stopper.raise_if_requested if stopper is not None else None + saver = ProgressSaver(stopper, console, minutes=_progress_minutes()) + + configure_callbacks = VitsModel.configure_callbacks + tune_cudnn = os.environ.get("JAVIS_PIPER_CUDNN_TUNING", "1") != "0" + + @functools.wraps(configure_callbacks) + def _configure_callbacks(self): + callbacks = configure_callbacks(self) + if isinstance(callbacks, Callback): + callbacks = [callbacks] + return [*callbacks, NonFiniteGuard(), *([CudnnFirstStepTuning()] if tune_cudnn else []), + *([stopper] if stopper is not None else []), saver] + + VitsModel.configure_callbacks = _configure_callbacks + _avoid_gpu_syncs() + + bucketing = os.environ.get("JAVIS_PIPER_BUCKETING", "1") != "0" + train_dataloader = VitsDataModule.train_dataloader + + @functools.wraps(train_dataloader) + def _train_dataloader(self): + if bucketing: + try: + return _bucketed_train_dataloader(self) + except (AttributeError, ImportError, TypeError) as exc: # piper's data module changed + print(f"WARNING: length bucketing unavailable ({exc!r}); " + "using piper's random batches.", file=sys.stderr, flush=True) + loader = train_dataloader(self) + # Lightning restores loader states only into as many loaders with a + # state as saved them: checkpoints of bucketed runs must load here too. + loader.state_dict, loader.load_state_dict = dict, lambda state: None + return loader + + VitsDataModule.train_dataloader = _train_dataloader + return stopper + + +# ── ONNX export ─────────────────────────────────────────────────────────────── + +def _pop_option(argv: list, name: str) -> str | None: + """Remove ``name value`` or ``name=value`` from *argv*; returns the value.""" + for i, arg in enumerate(argv): + if arg == name and i + 1 < len(argv): + value = argv[i + 1] + del argv[i:i + 2] + return value + if arg.startswith(name + "="): + del argv[i] + return arg[len(name) + 1:] + return None + + +def _epoch_progress(checkpoint: dict) -> dict: + """How far into its epoch a checkpoint of bucketed batches was saved.""" + try: + loader = checkpoint["loops"]["fit_loop"]["state_dict"]["combined_loader"][0] + return {"batches_done": int(loader["batches_done"]), + "batches_per_epoch": len(loader["order"])} + except (KeyError, IndexError, TypeError, ValueError): + return {} + + +def _inspect_exported_checkpoint(info: dict) -> None: + """Record the training state that gets exported in *info*; refuse NaN/inf weights.""" + from lightning.pytorch.core import saving + + load = saving.pl_load + + @functools.wraps(load) + def _pl_load(path, *args, **kwargs): + source = path + if isinstance(path, (str, os.PathLike)) and os.path.isfile(path): + # Read it whole right away: while it is open, training can't replace it. + with open(path, "rb") as f: + info["checkpoint_mtime_ns"] = os.fstat(f.fileno()).st_mtime_ns + source = io.BytesIO(f.read()) + checkpoint = load(source, *args, **kwargs) + del source + for name, value in (checkpoint.get("state_dict") or {}).items(): + if torch.is_tensor(value) and value.is_floating_point() \ + and not bool(torch.isfinite(value).all()): + raise SystemExit(f"ERROR: {path} holds NaN/inf weights ({name}); " + "not exporting it.") + if isinstance(path, (str, os.PathLike)): + info["checkpoint"] = os.path.abspath(path) + info.update(epoch=checkpoint.get("epoch"), global_step=checkpoint.get("global_step"), + **_epoch_progress(checkpoint)) + return checkpoint + + saving.pl_load = _pl_load + + +def _export(argv: list) -> None: + """piper's ONNX export. ``--info-file X`` writes what was exported to X (JSON).""" + info_file = _pop_option(argv, "--info-file") + _legacy_onnx_export() + info: dict = {} + _inspect_exported_checkpoint(info) + from piper.train.export_onnx import main + main() + if info_file: + info["exported_at"] = time.time() + + def write(path) -> None: + with open(path, "w", encoding="utf-8") as f: + json.dump(info, f) + + _write_atomically(info_file, write, durable=False) + + if __name__ == "__main__": if len(sys.argv) > 1 and sys.argv[1] == "export": del sys.argv[1] - _legacy_onnx_export() - from piper.train.export_onnx import main + _export(sys.argv) else: if any(a == "--ckpt_path" or a.startswith("--ckpt_path=") for a in sys.argv): _ignore_warmstart_on_resume() + _resolve_precision(sys.argv) + _console = training_control.ConsoleClose() + _console.install() + _stopper = _install_training_aids(_console) from piper.train.__main__ import main - main() + sys.exit(_fit(main, _stopper, _console)) diff --git a/src/javis/voices/piper_train.py b/src/javis/voices/piper_train.py index 04dd44b..07be83e 100644 --- a/src/javis/voices/piper_train.py +++ b/src/javis/voices/piper_train.py @@ -12,9 +12,33 @@ Stage 2 — ONNX export: python -m javis.voices.piper_compat export --checkpoint ... --output-file ... The training module (``piper.train``) is bundled in the ``piper-tts`` package -installed via ``pip install piper-tts[train]`` plus the OHF-Voice repo at -``D:\\piper1-gpl`` (editable install provides the training source code and -the pure-Python monotonic_align fallback). +installed via ``pip install piper-tts[train]``. Javis runs it through +``javis.voices.piper_compat``, which patches it at runtime (version fixes, a +working alignment search, length-bucketed batches, fewer GPU round trips, a +guard that stops runs whose gradients go NaN, crash-safe checkpoint and cache +files). + +Long runs +--------- +* **Stopping.** Press Ctrl+C once (or create ``training/.stop_training``): the + training process saves its progress after its current step — the epoch so + far included — and exits; the next run continues with the next batch. + Closing the console window saves the progress too. Ctrl+C twice stops at + once, losing what was trained since the last save. +* **The voice follows training.** ``last.ckpt`` is refreshed every 10 minutes + (``JAVIS_PIPER_PROGRESS_MINUTES``), and every 15 minutes + (``--snapshot-minutes``) the voice's ``piper_model.onnx`` is re-exported from + it in the background, so Javis speaks with the voice as trained so far — it + picks up each export by itself. A stopped run exports its final state. +* **Crashes.** A crashed training process (CUDA error, out of memory, driver + reset...) is restarted from the newest checkpoint, unless it crashed twice + in a row without saving one. Checkpoints that can't be loaded or hold + NaN/inf are passed over. +* **One run per voice.** ``training/.training.lock`` keeps the GUI's training + console and the overnight schedule from training a voice at the same time; + it lapses as soon as its processes have exited. +* Windows doesn't go to sleep while training runs, and clicking into the + console window no longer pauses it. CLI usage:: @@ -22,15 +46,17 @@ CLI usage:: --voice \\ [--language de] \\ [--base-checkpoint /path/to/base.ckpt] \\ - [--batch-size 32] \\ [--num-workers 4] \\ [--max-epochs 6000] \\ + [--batch-size 32] \\ + [--num-workers 2] \\ + [--max-epochs 6000] \\ [--accelerator gpu|cpu] \\ - [--precision 32|16|bf16] \\ + [--precision bf16-mixed|32|16-mixed] \\ + [--snapshot-minutes 15] \\ [--config config.yaml] Prerequisites ------------- * ``pip install piper-tts[train]`` -* OHF-Voice/piper1-gpl cloned at ``D:\\piper1-gpl`` and installed editable * NVIDIA GPU with >=6 GB VRAM recommended; CPU training supported but slow Training time (RTX 3070, ~5 000 utterances) @@ -40,40 +66,110 @@ Training time (RTX 3070, ~5 000 utterances) """ from __future__ import annotations +import itertools +import json +import math import os import re import shutil import subprocess import sys import threading +import time from pathlib import Path +from .training_control import ( + EXIT_BUSY, + EXIT_DIVERGED, + EXIT_STOPPED, + LAUNCHER_ENV, + PARTIAL_SUFFIX, + STOP_FILE_ENV, + TrainingLock, + console_selection_cannot_pause, + keep_system_awake, + process_alive, + python_command, + replace_file, + this_process, +) + # Set by request_stop() (e.g. when Javis quits during scheduled training) so -# the in-process pipeline terminates its child instead of blocking exit. +# the in-process pipeline stops its child instead of blocking exit. _STOP = threading.Event() +# On a stop, the training process gets this long to save its progress — or, +# when the stop comes at the end of an epoch, to validate (~1 min with MOS +# scoring) and save — before it is killed. +STOP_GRACE_SECONDS = 180 +# A crashed training process is restarted after this pause (the GPU driver may +# be recovering from a reset)... +RESTART_DELAY_SECONDS = 30 +# ...unless this many attempts in a row crashed without saving a checkpoint. +MAX_FAILURES_WITHOUT_PROGRESS = 2 +STOP_FILE_NAME = ".stop_training" +# While training runs, the voice is re-exported from its newest progress this often. +SNAPSHOT_MINUTES = 15.0 + def request_stop() -> None: - """Ask a running :func:`run_pipeline` (in another thread) to abort.""" + """Ask a running :func:`run_pipeline` (in another thread) to save and stop.""" _STOP.set() - # ── Helpers ────────────────────────────────────────────────────────────────── +def _kill(proc: subprocess.Popen) -> None: + """End *proc* now. (Its DataLoader workers notice and exit by themselves.)""" + if proc.poll() is None: + proc.kill() + try: + proc.wait(timeout=30) + except subprocess.TimeoutExpired: + pass + + +def _begin_stop(proc: subprocess.Popen, stop_file: Path | None) -> float: + """Ask *proc* to stop: through *stop_file* if it has one, else by killing it.""" + if stop_file is not None: + try: + stop_file.touch() + except OSError as exc: + print(f"\nCould not create {stop_file} ({exc}); stopping training at once.", + file=sys.stderr, flush=True) + _kill(proc) + else: + print("\nStopping: training saves its progress after its current step (at the " + "end of an epoch it validates first, for a minute or two).\n" + "Press Ctrl+C again to stop at once and lose the progress since the " + "last save.", flush=True) + else: + _kill(proc) + return time.monotonic() + + def _run( cmd: list[str], *, cwd: Path | None = None, desc: str, env: dict | None = None, + stop_file: Path | None = None, + on_start=None, + on_poll=None, ) -> int: - """Run *cmd* as a subprocess, streaming output to stdout/stderr. + """Run *cmd* as a subprocess, streaming output to stdout/stderr; return its exit code. - Returns the process exit code. Handles Ctrl-C gracefully: forwards the - interrupt to the child process group and waits for it to finish before - returning exit code 1, suppressing the noisy BrokenPipeError tracebacks - that Windows multiprocessing workers emit when the pipe is torn down. + The child gets no Ctrl+C of its own (it runs in a new process group): this + process decides how it stops. With a *stop_file*, the child is a training + process that saves and exits (``EXIT_STOPPED``) once the file exists; Ctrl+C + and :func:`request_stop` create it. A second Ctrl+C, or + ``STOP_GRACE_SECONDS`` without exiting, kills the child. Other children + are killed right away. Returns ``EXIT_STOPPED`` for every stopped run + (``EXIT_DIVERGED`` if the run diverged meanwhile). + + *on_start* is called with the child's pid, and with ``None`` once it has exited. + *on_poll* is called about once a second while the child runs (until a stop). """ print(f"\n{'='*60}") print(f" {desc}") @@ -81,42 +177,75 @@ def _run( print(f"{'='*60}\n", flush=True) if _STOP.is_set(): - return 1 - # CREATE_NEW_PROCESS_GROUP lets us send CTRL_C_EVENT to only the child. + return EXIT_STOPPED + env = dict(os.environ if env is None else env) + if stop_file is not None: + stop_file.unlink(missing_ok=True) + launcher = this_process() + env[STOP_FILE_ENV] = str(stop_file) + env[LAUNCHER_ENV] = f"{launcher['pid']}:{launcher['start']}" + print("Press Ctrl+C to save the training progress and stop.\n", flush=True) flags = getattr(subprocess, "CREATE_NEW_PROCESS_GROUP", 0) proc = subprocess.Popen(cmd, cwd=cwd, env=env, creationflags=flags) + stopping_since: float | None = None try: + if on_start is not None: + on_start(proc.pid) while True: + try: # around all of it: a Ctrl+C anywhere in here stops gracefully + try: + proc.wait(timeout=1.0) + except subprocess.TimeoutExpired: + pass + if proc.poll() is not None: + break + if stopping_since is None: + if _STOP.is_set(): + stopping_since = _begin_stop(proc, stop_file) + elif on_poll is not None: + try: + on_poll() + except Exception as exc: # noqa: BLE001 - must not end the run + print(f"\nWARNING: {exc!r}", file=sys.stderr, flush=True) + on_poll = None + elif time.monotonic() - stopping_since > STOP_GRACE_SECONDS: + print("Training did not stop in time; stopping it at once.", + file=sys.stderr, flush=True) + _kill(proc) + except KeyboardInterrupt: + if stopping_since is None: + stopping_since = _begin_stop(proc, stop_file) + else: + print("\nStopping at once.", flush=True) + _kill(proc) + finally: + _kill(proc) + if stop_file is not None: try: - proc.wait(timeout=1.0) - break - except subprocess.TimeoutExpired: - if _STOP.is_set(): - print("\nStop requested — terminating training process…", flush=True) - proc.terminate() - try: - proc.wait(timeout=30) - except subprocess.TimeoutExpired: - proc.kill() - proc.wait() - return 1 - except KeyboardInterrupt: - print("\nInterrupt received — waiting for training process to stop…", - flush=True) - try: - import signal as _signal - proc.send_signal(_signal.CTRL_C_EVENT) # type: ignore[attr-defined] - except Exception: - proc.terminate() - try: - proc.wait(timeout=30) - except subprocess.TimeoutExpired: - proc.kill() - proc.wait() - return 1 + stop_file.unlink(missing_ok=True) + except OSError: + pass + if on_start is not None: + on_start(None) + if stopping_since is not None and proc.returncode != EXIT_DIVERGED: + return EXIT_STOPPED return proc.returncode +def _pause(seconds: float) -> bool: + """Wait *seconds*; ``False`` if a stop was requested (or Ctrl+C pressed) meanwhile.""" + deadline = time.monotonic() + seconds + try: + while not _STOP.is_set(): + left = deadline - time.monotonic() + if left <= 0: + return True + time.sleep(min(1.0, left)) + except KeyboardInterrupt: + pass + return False + + def _find_best_checkpoint(search_dir: Path) -> Path | None: """Return the best checkpoint under *search_dir* (recursive). @@ -146,11 +275,17 @@ def _find_best_checkpoint(search_dir: Path) -> Path | None: # Minimum free space required on the training drive before starting a run. -# Each ModelCheckpoint save is ~800 MB and Lightning's atomic save writes a -# full temp copy first (~2× headroom); 25 GB also leaves room for the -# per-version checkpoint retention (top-5 × 2 callbacks + last ≈ 8.75 GB). +# A checkpoint (~800 MB) is written next to the one it replaces before that +# one is removed; 25 GB also leaves room for the per-version checkpoint +# retention (top-5 × 2 callbacks + last ≈ 8.75 GB). MIN_FREE_BYTES = 25 * 1024**3 +# bf16 is the fastest precision measured for piper's VITS (on an RTX 3090 Ti +# fp16 took ~1.4× and TF32 "32" ~1.15× as long per step), and with fp32's +# range it needs no GradScaler. piper_compat falls back to 32-bit on GPUs +# that can't run it natively. +DEFAULT_PRECISION = "bf16-mixed" + # Model hyper-parameters that define the network shape (Piper "quality"). _ARCH_KEYS = ( "resblock", "resblock_kernel_sizes", "resblock_dilation_sizes", @@ -201,6 +336,123 @@ def _write_arch_config(arch: dict, dest: Path) -> Path: return dest +def _checkpoint_problem(ckpt: dict, precision: str) -> str | None: + """Why training can't continue from *ckpt*, or ``None`` if it looks healthy. + + A diverged run keeps writing checkpoints whose weights or optimizer state + are NaN/inf, or (fp16) whose GradScaler scale collapsed to 0 so every + later step is skipped. Resuming from those silently trains nothing. + """ + import math + + import torch + + def _non_finite(t) -> bool: + return (torch.is_tensor(t) and t.is_floating_point() + and not bool(torch.isfinite(t).all())) + + bad = [k for k, t in (ckpt.get("state_dict") or {}).items() if _non_finite(t)] + if bad: + parts = ", ".join(sorted({k.split(".", 1)[0] for k in bad})) + return f"{len(bad)} weight tensors contain NaN/inf (in {parts})" + for i, opt in enumerate(ckpt.get("optimizer_states") or []): + for state in (opt.get("state") or {}).values(): + if any(_non_finite(t) for t in state.values()): + return f"the state of optimizer {i} contains NaN/inf" + if precision.startswith("16"): # the loss scale only matters for fp16 + for key in ("MixedPrecision", "native_amp_scaling_state"): + scaler = ckpt.get(key) + if isinstance(scaler, dict) and scaler.get("scale") is not None: + scale = float(scaler["scale"]) + if not (math.isfinite(scale) and scale >= 1.0): + return (f"its fp16 loss scale collapsed to {scale:g}, so every " + "optimizer step would be skipped") + return None + + +def _usable_checkpoints(search_dir: Path, precision: str): + """Yield ``(path, problem, arch)`` for the checkpoints under *search_dir*. + + First the files that can't be loaded at all, then the others by training + progress (global step; ``last.ckpt`` first among equals, then the newest + file). *problem* says why a checkpoint can't be used (unreadable, + NaN/inf...), else it is ``None`` and *arch* holds the model's architecture. + Every checkpoint is unmapped again before it is yielded: Windows can't + replace a mapped file. + """ + import gc + + found, unreadable = [], [] + for path in search_dir.rglob("*.ckpt"): + try: + state = _load_checkpoint(path, mmap=True) + usable = isinstance(state, dict) and isinstance(state.get("state_dict"), dict) + step = int(state.get("global_step") or 0) if usable else 0 + mtime = path.stat().st_mtime + except FileNotFoundError: + continue # deleted meanwhile (a running training dropped it from its top 5) + except Exception as exc: # noqa: BLE001 - truncated or garbled: any error + unreadable.append((path, f"can't be loaded ({type(exc).__name__}: {exc})")) + continue + finally: + state = None + if usable: + found.append((step, path.name.startswith("last"), mtime, path)) + else: + unreadable.append((path, "not a Lightning training checkpoint")) + gc.collect() + for path, problem in unreadable: + yield path, problem, {} + for *_, path in sorted(found, key=lambda f: f[:3], reverse=True): + try: + state = _load_checkpoint(path, mmap=True) + problem, arch = _checkpoint_problem(state, precision), _arch_hparams(state) + except FileNotFoundError: + continue + except Exception as exc: # noqa: BLE001 + problem, arch = f"can't be loaded ({type(exc).__name__}: {exc})", {} + finally: + state = None + gc.collect() + yield path, problem, arch + + +def _newest_usable_checkpoint(search_dir: Path, precision: str): + """``(path, arch, problems)`` of the most advanced checkpoint training can continue from. + + *path* is ``None`` if there is none; *problems* describes the unreadable + checkpoints and each more advanced one that was passed over. + """ + problems = [] + for path, problem, arch in _usable_checkpoints(search_dir, precision): + if problem is None: + return path, arch, problems + problems.append(f"{path}: {problem}") + return None, {}, problems + + +def _newest_checkpoint_time(search_dir: Path) -> float: + times = [0.0] + for path in search_dir.rglob("*.ckpt"): + try: + times.append(path.stat().st_mtime) + except OSError: + pass + return max(times) + + +def _remove_partial_saves(lightning_dir: Path) -> None: + """Delete the remnants of checkpoint saves cut short by a crash or power cut. + + Call only while no training process runs (:func:`_run` returns once it has exited). + """ + for partial in lightning_dir.rglob("*" + PARTIAL_SUFFIX): + try: + partial.unlink() + except OSError: + pass + + def _write_portable_warmstart(src: Path, dest: Path) -> dict: """Save *src*'s weights in a form ``--model.warmstart_ckpt`` can load on this OS. @@ -284,6 +536,176 @@ def _prune_lightning_versions( print(f" Pruned {pruned_gb:.2f} GB total from {logs_dir}") +def _enough_disk_space(lightning_dir: Path) -> bool: + # Running out mid-save crashes the run (OSError: [Errno 28] No space left + # on device). Fail fast with a clear message instead. + free = shutil.disk_usage(lightning_dir).free + if free >= MIN_FREE_BYTES: + return True + print( + f"ERROR: insufficient disk space for training: " + f"{free / 1024**3:.1f} GB free on " + f"'{lightning_dir.drive or lightning_dir}', " + f"need ≥ {MIN_FREE_BYTES / 1024**3:.0f} GB. " + f"Free up space or move the voice's training directory to a " + f"larger drive, then retry.", + file=sys.stderr, + ) + return False + + +def _starting_point(training_dir: Path, base_checkpoint: Path | None, precision: str, + *, first: bool): + """Where training starts: ``(fit arguments, architecture, resume checkpoint)``. + + The voice's own most advanced usable checkpoint; else a warm start from + *base_checkpoint* or the best checkpoint in ``training/piper_base/``; else + scratch. ``None`` (after explaining why) if the voice has checkpoints but + none is usable. *first*: the first start of this run (prints more). + """ + lightning_dir = training_dir / "piper_lightning" + resume_ckpt, arch, problems = _newest_usable_checkpoint(lightning_dir, precision) + if resume_ckpt is None and problems: + print( + "ERROR: cannot continue training: none of this voice's checkpoints is usable.\n " + + "\n ".join(problems) + + "\nDelete or move away\n" + f" {lightning_dir}\n" + "and train again; the next run fine-tunes from the base model in\n" + f" {training_dir / 'piper_base'}", + file=sys.stderr, + ) + return None + for problem in problems: + print(f"WARNING: passing over {problem}", file=sys.stderr) + if resume_ckpt is not None: + # A base checkpoint only seeds the first run; afterwards training always + # continues from this voice's own newest checkpoint. + if first and base_checkpoint is not None: + print(f"Ignoring base checkpoint {base_checkpoint}: this voice already has its own " + f"checkpoints. Delete {lightning_dir} to start over from a base model.") + print(f"Resuming from checkpoint: {resume_ckpt}") + return ["--ckpt_path", str(resume_ckpt)], arch, resume_ckpt + + base_dir = training_dir / "piper_base" + warmstart_src = base_checkpoint + if warmstart_src is None and base_dir.is_dir(): + warmstart_src = _find_best_checkpoint(base_dir) + if warmstart_src is not None: + print(f"Fine-tuning from base checkpoint: {warmstart_src}") + warmstart_file = training_dir / "warmstart.pt" + arch = _write_portable_warmstart(warmstart_src, warmstart_file) + return ["--model.warmstart_ckpt", str(warmstart_file)], arch, None + if first: + print("No base checkpoint found — training from scratch.") + print( + "Tip: place a piper checkpoint (.ckpt) in\n" + f" {base_dir}\n" + "to fine-tune from a pre-trained voice (much faster convergence).\n" + "Pre-trained checkpoints: https://huggingface.co/datasets/rhasspy/piper-checkpoints" + ) + return [], {}, None + + +def _training_env() -> dict: + """The training process's environment: this one's, plus a few overrides.""" + # torch.compile is disabled: the only available backend on Windows without + # MSVC is aot_eager, which provides zero GPU kernel fusion (that requires + # inductor/Triton). aot_eager still builds a per-shape graph cache — with + # variable-length phoneme/audio sequences this cache grows unboundedly, + # filling VRAM after ~60 iterations and collapsing throughput. + # Pure eager mode is faster and holds VRAM steady. + # PYTORCH_NO_CUDA_GRAPH_CAPTURE: prevents CUDA graph capture from holding + # the GPU long enough to trigger Windows TDR (driver reset / PC crash). + env = {**os.environ, + "PYTORCH_NO_CUDA_GRAPH_CAPTURE": "1", + "TORCHDYNAMO_DISABLE": "1"} + # PYTORCH_CUDA_ALLOC_CONF: the epoch-1542 OOM failed on a tiny 58 MiB + # allocation while 549 MiB sat in reserved-but-unallocated segments + # (allocator fragmentation). expandable_segments lets the caching + # allocator grow segments in place instead of needing a contiguous block. + # PyTorch doesn't support it on Windows (it only warns there). + if sys.platform != "win32": + env["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" + return env + + +def _train(training_dir: Path, fit_args: list[str], *, lock: TrainingLock, + base_checkpoint: Path | None, precision: str, + snapshot_minutes: float = SNAPSHOT_MINUTES) -> int: + """Stage 1: run the training process; restart it after a crash. + + Every start continues from the most advanced usable checkpoint. Gives up + once ``MAX_FAILURES_WITHOUT_PROGRESS`` starts in a row crashed before + saving a checkpoint. A run that diverged (``EXIT_DIVERGED``) isn't + restarted: it would only diverge again. Meanwhile a :class:`SnapshotExporter` + keeps the voice's model close behind the training progress. + """ + lightning_dir = training_dir / "piper_lightning" + exporter = SnapshotExporter(training_dir.parent, snapshot_minutes) + failures, first = 0, True + try: + while True: + _remove_partial_saves(lightning_dir) + if not _enough_disk_space(lightning_dir): + return 1 + start = _starting_point(training_dir, base_checkpoint, precision, first=first) + if start is None: + return 1 + start_args, arch, resume_ckpt = start + first = False + + # ── Prune old Lightning versions (caps piper_lightning disk usage) ── + # After choosing the start so the version holding the resume checkpoint + # is protected from deletion. Piper keeps ~8.75 GB per version (top-5 × + # 2 callbacks + last), and every fit creates a new version_N even when + # resuming — without pruning this grows unbounded until the disk is full. + _prune_lightning_versions(lightning_dir, keep_ckpt=resume_ckpt) + + # The interpreter itself, not a virtual environment's python.exe + # launcher: that would end it at once when the console window closes, + # before it could save its progress. + python, python_env = python_command() + cmd: list[str] = [python, "-m", "javis.voices.piper_compat", "fit"] + if arch: + # Must come first: later CLI flags override config-file values. + cmd += ["--config", str(_write_arch_config(arch, training_dir / "model_arch.yaml"))] + cmd += fit_args + start_args + + saved_before = _newest_checkpoint_time(lightning_dir) + rc = _run(cmd, env={**_training_env(), **python_env}, + desc="Stage 1/2 — Phoneme preprocessing + VITS training", + stop_file=training_dir / STOP_FILE_NAME, on_start=lock.set_trainer, + on_poll=exporter.poll) + (training_dir / "warmstart.pt").unlink(missing_ok=True) # only read at fit start + if rc in (0, EXIT_STOPPED, EXIT_DIVERGED) or _STOP.is_set(): + break + failures = 0 if _newest_checkpoint_time(lightning_dir) > saved_before else failures + 1 + if failures >= MAX_FAILURES_WITHOUT_PROGRESS: + print(f"\nERROR: training failed (exit {rc}), {failures} times in a row before " + "it saved a checkpoint; not restarting it again.", file=sys.stderr) + return rc + print(f"\nTraining crashed (exit {rc}). Restarting it from its newest checkpoint in " + f"{RESTART_DELAY_SECONDS} s — press Ctrl+C to stop instead.", + file=sys.stderr, flush=True) + if not _pause(RESTART_DELAY_SECONDS): + rc = EXIT_STOPPED + break + finally: + exporter.stop() + if rc == EXIT_DIVERGED: + print(f"\nERROR: training diverged (exit {rc}; see above), so it was not restarted.", + file=sys.stderr) + return rc + if rc == EXIT_STOPPED or _STOP.is_set(): + print("\nTraining stopped; its progress is saved. Train again to continue where it " + "left off.") + return EXIT_STOPPED + if rc != 0: + print(f"\nERROR: training failed (exit {rc}).", file=sys.stderr) + return rc + + # ── Main pipeline ───────────────────────────────────────────────────────────── def run_pipeline( @@ -294,11 +716,12 @@ def run_pipeline( batch_size: int = 32, num_workers: int = 2, max_epochs: int = 6000, - precision: str = "16-mixed", + precision: str = DEFAULT_PRECISION, accelerator: str = "gpu", validation_split: float = 0.01, num_test_examples: int = 10, sample_rate: int = 22050, + snapshot_minutes: float = SNAPSHOT_MINUTES, ) -> int: """Run the full piper training pipeline and return the exit code. @@ -307,12 +730,24 @@ def run_pipeline( {voice_dir}/piper_model.onnx {voice_dir}/piper_model.onnx.json (voice config written during training) + {voice_dir}/piper_model.info.json (the checkpoint and epoch it is from) + + While training runs, the model there is refreshed from the training + progress every *snapshot_minutes* (0 = only at the end), and a run stopped + with Ctrl+C exports the progress it saved. If the voice already has checkpoints in ``training/piper_lightning/`` the - run resumes from the newest one. Otherwise it fine-tunes (warm start: + run continues from the most advanced one that loads and is healthy, + passing over corrupt ones (truncated, NaN weights, collapsed fp16 loss + scale); if none is, it refuses and explains how to start over. Otherwise + it fine-tunes (warm start: weights only, fresh optimizer) from *base_checkpoint* or the best checkpoint in ``training/piper_base/``, using that model's architecture (e.g. Piper "high" quality). Without either it trains from scratch. + + A crashed training process is restarted (see :func:`_train`). Returns + ``EXIT_STOPPED`` when stopped on request (Ctrl+C, :func:`request_stop`) + after saving, and ``EXIT_BUSY`` if this voice is being trained elsewhere. """ _STOP.clear() training_dir = voice_dir / "training" @@ -346,127 +781,61 @@ def run_pipeline( lightning_dir = training_dir / "piper_lightning" lightning_dir.mkdir(parents=True, exist_ok=True) - # ── Preflight: disk space ───────────────────────────────────────────── - # Lightning's atomic checkpoint save needs ~2× the checkpoint size - # (~800 MB) as headroom; running out mid-save crashes the whole run - # (OSError: [Errno 28] No space left on device) and leaves a partial - # temp file behind. Fail fast with a clear message instead. - free = shutil.disk_usage(lightning_dir).free - if free < MIN_FREE_BYTES: + lock = TrainingLock(voice_dir) + if not lock.acquire(): print( - f"ERROR: insufficient disk space for training: " - f"{free / 1024**3:.1f} GB free on " - f"'{lightning_dir.drive or lightning_dir}', " - f"need ≥ {MIN_FREE_BYTES / 1024**3:.0f} GB. " - f"Free up space or move the voice's training directory to a " - f"larger drive, then retry.", + f"ERROR: '{voice_dir.name}' is being trained already, in another training " + "console or by the training schedule. Stop that run first.\n" + f"(Its lock, {lock.path}, lapses as soon as that run has exited.)", file=sys.stderr, ) - return 1 + return EXIT_BUSY + try: + with keep_system_awake(): + fit_args = [ + "--data.voice_name", voice_dir.name, + "--data.csv_path", str(metadata_csv), + "--data.audio_dir", str(wavs_dir), + "--data.cache_dir", str(cache_dir), + "--data.config_path", str(config_json), + "--data.espeak_voice", language, + "--data.batch_size", str(batch_size), + "--data.num_workers", str(num_workers), + "--data.validation_split", str(validation_split), + "--data.num_test_examples", str(num_test_examples), + "--model.sample_rate", str(sample_rate), + "--trainer.max_epochs", str(max_epochs), + "--trainer.accelerator", accelerator, + "--trainer.devices", "1", + "--trainer.precision", precision, + "--trainer.default_root_dir", str(lightning_dir), + ] + rc = _train(training_dir, fit_args, lock=lock, + base_checkpoint=base_checkpoint, precision=precision, + snapshot_minutes=snapshot_minutes) + if rc == EXIT_STOPPED and not _STOP.is_set(): + _export_progress(voice_dir) # (not when Javis is quitting) + if rc != 0: + return rc - # ── Starting point ────────────────────────────────────────────────────────────────── - # A base checkpoint only seeds the first run; afterwards training always - # continues from this voice's own newest checkpoint. - resume_ckpt = _find_best_checkpoint(lightning_dir) - warmstart_src: Path | None = None - if resume_ckpt is None: - warmstart_src = base_checkpoint - base_dir = training_dir / "piper_base" - if warmstart_src is None and base_dir.is_dir(): - warmstart_src = _find_best_checkpoint(base_dir) - elif base_checkpoint is not None: - print(f"Ignoring base checkpoint {base_checkpoint}: this voice already has its own " - f"checkpoints. Delete {lightning_dir} to start over from a base model.") + # ── Stage 2: export ONNX ───────────────────────────────────────── + best_ckpt, _, problems = _newest_usable_checkpoint(lightning_dir, "32") + for problem in problems: + print(f"WARNING: passing over {problem}", file=sys.stderr) + if best_ckpt is None: + print( + "ERROR: no usable checkpoint found after training.\n" + f"Expected checkpoints under: {lightning_dir}", + file=sys.stderr, + ) + return 1 - warmstart_file = training_dir / "warmstart.pt" - start_args: list[str] = [] - if resume_ckpt is not None: - print(f"Resuming from checkpoint: {resume_ckpt}") - arch = _arch_hparams(_load_checkpoint(resume_ckpt, mmap=True)) - start_args = ["--ckpt_path", str(resume_ckpt)] - elif warmstart_src is not None: - print(f"Fine-tuning from base checkpoint: {warmstart_src}") - arch = _write_portable_warmstart(warmstart_src, warmstart_file) - start_args = ["--model.warmstart_ckpt", str(warmstart_file)] - else: - arch = {} - print("No base checkpoint found — training from scratch.") - print( - "Tip: place a piper checkpoint (.ckpt) in\n" - f" {training_dir / 'piper_base'}\n" - "to fine-tune from a pre-trained voice (much faster convergence).\n" - "Pre-trained checkpoints: https://huggingface.co/datasets/rhasspy/piper-checkpoints" - ) - - # ── Prune old Lightning versions (caps piper_lightning disk usage) ── - # After auto-detect so the version holding the resume checkpoint is - # protected from deletion. Piper keeps ~8.75 GB per version (top-5 × 2 - # callbacks + last), and every fit creates a new version_N even when - # resuming — without pruning this grows unbounded until the disk is full. - _prune_lightning_versions(lightning_dir, keep_ckpt=resume_ckpt) - - # ── Stage 1: train ──────────────────────────────────────────────────────── - train_cmd: list[str] = [sys.executable, "-m", "javis.voices.piper_compat", "fit"] - if arch: - # Must come first: later CLI flags override config-file values. - train_cmd += ["--config", str(_write_arch_config(arch, training_dir / "model_arch.yaml"))] - train_cmd += [ - "--data.voice_name", voice_dir.name, - "--data.csv_path", str(metadata_csv), - "--data.audio_dir", str(wavs_dir), - "--data.cache_dir", str(cache_dir), - "--data.config_path", str(config_json), - "--data.espeak_voice", language, - "--data.batch_size", str(batch_size), - "--data.num_workers", str(num_workers), - "--data.validation_split", str(validation_split), - "--data.num_test_examples", str(num_test_examples), - "--model.sample_rate", str(sample_rate), - "--trainer.max_epochs", str(max_epochs), - "--trainer.accelerator", accelerator, - "--trainer.devices", "1", - "--trainer.precision", precision, - "--trainer.default_root_dir", str(lightning_dir), - ] + start_args - - # Build subprocess environment: inherit everything, then overlay overrides. - # torch.compile is disabled: the only available backend on Windows without - # MSVC is aot_eager, which provides zero GPU kernel fusion (that requires - # inductor/Triton). aot_eager still builds a per-shape graph cache — with - # variable-length phoneme/audio sequences this cache grows unboundedly, - # filling VRAM after ~60 iterations and collapsing throughput. - # Pure eager mode is faster and holds VRAM steady. - # PYTORCH_NO_CUDA_GRAPH_CAPTURE: prevents CUDA graph capture from holding - # the GPU long enough to trigger Windows TDR (driver reset / PC crash). - # PYTORCH_CUDA_ALLOC_CONF: the epoch-1542 OOM failed on a tiny 58 MiB - # allocation while 549 MiB sat in reserved-but-unallocated segments - # (allocator fragmentation). expandable_segments lets the caching - # allocator grow segments in place instead of needing a contiguous block. - train_env = {**os.environ, - "PYTORCH_NO_CUDA_GRAPH_CAPTURE": "1", - "TORCHDYNAMO_DISABLE": "1", - "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True"} - - rc = _run(train_cmd, env=train_env, desc="Stage 1/2 — Phoneme preprocessing + VITS training") - warmstart_file.unlink(missing_ok=True) # only read at fit start - if rc != 0: - print(f"\nERROR: training failed (exit {rc}).", file=sys.stderr) - return rc - - # ── Stage 2: export ONNX ───────────────────────────────────────────────── - best_ckpt = _find_best_checkpoint(lightning_dir) - if best_ckpt is None: - print( - "ERROR: no checkpoint found after training.\n" - f"Expected checkpoints under: {lightning_dir}", - file=sys.stderr, - ) - return 1 - - onnx_out = voice_dir / "piper_model.onnx" - rc = _export_onnx(best_ckpt, onnx_out, voice_dir, desc="Stage 2/2 — ONNX export") - if rc != 0: - return rc + onnx_out = voice_dir / "piper_model.onnx" + rc = _export_onnx(best_ckpt, onnx_out, voice_dir, desc="Stage 2/2 — ONNX export") + if rc != 0: + return rc + finally: + lock.release() print(f"\nDone! Model saved to: {onnx_out}") print("\nTo use this voice in Javis, set in config.yaml:") @@ -514,40 +883,310 @@ def _rebuild_voice_config(lightning_dir: Path, dest: Path) -> bool: return True +def _voice_config_ok(path: Path) -> bool: + """Whether *path* holds a complete voice config (not one cut short by a crash).""" + import json + + try: + cfg = json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError): + return False + return isinstance(cfg, dict) and bool(cfg.get("phoneme_id_map")) + + def _export_onnx(ckpt: Path, onnx_out: Path, voice_dir: Path, *, desc: str) -> int: - """Export *ckpt* to *onnx_out* and put the voice config (``.onnx.json``) beside it.""" - rc = _run( - [ - sys.executable, "-m", "javis.voices.piper_compat", "export", - "--checkpoint", str(ckpt), - "--output-file", str(onnx_out), - ], - desc=desc, - ) - if rc != 0: - print(f"\nERROR: ONNX export failed (exit {rc}).", file=sys.stderr) - return rc + """Export *ckpt* to *onnx_out*, with the voice config (``.onnx.json``) and a + record of what was exported (``.info.json``) beside it. + + All of it is written under temporary names and renamed into place once + complete: a failed or interrupted export leaves the previous model as it was. + """ + _remove_stale_exports(onnx_out.parent) + files = _ExportFiles(onnx_out) + try: + cmd, env = _export_command(ckpt, files) + rc = _run(cmd, env=env, desc=desc) + if rc == EXIT_STOPPED: + print("\nONNX export stopped.", file=sys.stderr) + return rc + if rc != 0: + print(f"\nERROR: ONNX export failed (exit {rc}).", file=sys.stderr) + return rc + return _install_export(files, voice_dir, verbose=True) + finally: + files.discard() + + +_EXPORT_IDS = itertools.count(1) +# The temporary files of an export: ".partial--". +_EXPORT_PARTIAL = re.compile(re.escape(PARTIAL_SUFFIX) + r"-(\d+)-\d+\.\w+$") + + +class _ExportFiles: + """The files one ONNX export to *onnx_out* writes, and their temporary names. + + The temporary names are unique to the process and the export, so exports + running at the same time (in the background during training, from the + GUI...) can't write into each other's files. + """ + + def __init__(self, onnx_out: Path) -> None: + tag = f"{PARTIAL_SUFFIX}-{os.getpid()}-{next(_EXPORT_IDS)}" + self.onnx = onnx_out + self.config = onnx_out.with_suffix(".onnx.json") + self.info = export_info_path(onnx_out) + self.partial_onnx, self.partial_config, self.partial_info = ( + path.with_name(path.stem + tag + path.suffix) + for path in (self.onnx, self.config, self.info)) + + def discard(self) -> None: + """Delete the temporary files that are left.""" + for path in (self.partial_onnx, self.partial_config, self.partial_info): + try: + path.unlink(missing_ok=True) + except OSError: + pass + + +def _remove_stale_exports(directory: Path) -> None: + """Delete the temporary files of exports whose launcher has died (closed console...).""" + try: + entries = list(directory.iterdir()) + except OSError: + return + for path in entries: + match = _EXPORT_PARTIAL.search(path.name) + if match is None or int(match.group(1)) == os.getpid() or process_alive(int(match.group(1))): + continue + try: + path.unlink() + except OSError: + pass # still being written by an export that outlived its launcher + + +def _export_command(ckpt: Path, files: _ExportFiles) -> tuple[list[str], dict]: + """The command (and its environment) exporting *ckpt* into *files*' temporary names.""" + python, python_env = python_command() + cmd = [python, "-m", "javis.voices.piper_compat", "export", + "--checkpoint", str(ckpt), + "--output-file", str(files.partial_onnx), + "--info-file", str(files.partial_info)] + return cmd, {**os.environ, **python_env} + + +def _install_export(files: _ExportFiles, voice_dir: Path, *, verbose: bool = False) -> int: + """Put a finished export in place: the voice config, the model, then the record. + + In this order, whatever loads the model as soon as it changes (Javis does) + finds the matching config beside it. + """ + try: + complete = files.partial_onnx.stat().st_size > 0 + except OSError: + complete = False + if not complete: + print(f"\nERROR: ONNX export wrote no model ({files.partial_onnx}).", file=sys.stderr) + return 1 config_json = voice_dir / "training" / "piper_config.json" - dest_cfg = onnx_out.with_suffix(".onnx.json") - if config_json.exists(): - shutil.copy2(config_json, dest_cfg) - print(f"Config copied to: {dest_cfg}") - elif _rebuild_voice_config(voice_dir / "training" / "piper_lightning", dest_cfg): - print(f"{config_json.name} is missing — rebuilt the voice config from the " - f"training run: {dest_cfg}") - else: - print(f"ERROR: {config_json} is missing and could not be rebuilt; piper needs " - f"{dest_cfg.name} next to the model.", file=sys.stderr) + try: + if _voice_config_ok(config_json): + shutil.copyfile(config_json, files.partial_config) + if verbose: + print(f"Config copied to: {files.config}") + elif _rebuild_voice_config(voice_dir / "training" / "piper_lightning", + files.partial_config): + state = "damaged" if config_json.exists() else "missing" + print(f"{config_json.name} is {state} — rebuilt the voice config from the " + f"training run: {files.config}") + else: + print(f"ERROR: {config_json} is missing or damaged and could not be rebuilt; " + f"piper needs {files.config.name} next to the model.", file=sys.stderr) + return 1 + except OSError as exc: + print(f"ERROR: could not write {files.config.name} ({exc}).", file=sys.stderr) return 1 + + try: + replace_file(files.partial_config, files.config) + replace_file(files.partial_onnx, files.onnx) + except OSError as exc: + print(f"ERROR: could not replace {files.onnx.name} ({exc}). Close what is using " + "the voice and export again.", file=sys.stderr) + return 1 + try: + if files.partial_info.is_file(): + replace_file(files.partial_info, files.info) + else: + files.info.unlink(missing_ok=True) # it describes the model replaced + except OSError: + pass return 0 +def export_info_path(onnx: Path) -> Path: + """Where the record of what *onnx* was exported from is kept.""" + return onnx.with_suffix(".info.json") + + +def export_info(onnx: Path) -> dict | None: + """What the model *onnx* was exported from (written by each export), if recorded. + + Keys: ``checkpoint``, ``checkpoint_mtime_ns``, ``epoch``, ``global_step``, + ``exported_at`` (``time.time()``) and, for checkpoints saved within an + epoch, ``batches_done`` and ``batches_per_epoch``. + """ + try: + info = json.loads(export_info_path(onnx).read_text(encoding="utf-8")) + except (OSError, ValueError): + return None + return info if isinstance(info, dict) else None + + +def describe_export(info: dict | None) -> str: + """How far trained the exported model is, e.g. ``"epoch 3, batch 700/1456"``.""" + if not info or not isinstance(info.get("epoch"), int): + return "an unknown point of training" + epoch, done, total = info["epoch"], info.get("batches_done"), info.get("batches_per_epoch") + if isinstance(done, int) and isinstance(total, int) and 0 < done < total: + return f"epoch {epoch}, batch {done}/{total}" + if isinstance(done, int) and done == total: + return f"the end of epoch {epoch}" + return f"epoch {epoch}" + + +def _newest_progress(lightning_dir: Path) -> Path | None: + """The training's latest saved progress: the newest ``last*.ckpt``.""" + newest, newest_mtime = None, -1 + try: + for path in lightning_dir.rglob("last*.ckpt"): + try: + mtime = path.stat().st_mtime_ns + except OSError: + continue + if mtime > newest_mtime: + newest, newest_mtime = path, mtime + except OSError: + return None + return newest + + +def _checkpoint_signature(ckpt: Path) -> tuple | None: + try: + st = ckpt.stat() + except OSError: + return None + return os.path.normcase(os.path.abspath(ckpt)), st.st_mtime_ns, st.st_size + + +def _exported_from(ckpt: Path, onnx: Path) -> bool: + """Whether *onnx* was exported from *ckpt* as it is now.""" + info = export_info(onnx) + try: + return (onnx.is_file() + and os.path.normcase(os.path.abspath(info["checkpoint"])) + == os.path.normcase(os.path.abspath(ckpt)) + and info["checkpoint_mtime_ns"] == ckpt.stat().st_mtime_ns) + except (OSError, KeyError, TypeError, ValueError): + return False + + +class SnapshotExporter: + """Keeps the voice's model close behind the training, for Javis to speak with. + + :meth:`poll`, called while training runs, exports the training's newest + saved progress (``last.ckpt``) in the background — every *minutes* at most, + at below-normal priority — and puts the model in place once the export is + complete. Javis loads it by itself. A checkpoint whose export failed + isn't tried again. + """ + + def __init__(self, voice_dir: Path, minutes: float) -> None: + self.voice_dir = voice_dir + self.onnx = voice_dir / "piper_model.onnx" + self.lightning_dir = voice_dir / "training" / "piper_lightning" + self.log_path = voice_dir / "training" / "snapshot_export.log" + self.interval = minutes * 60 if math.isfinite(minutes) and minutes > 0 else 0.0 + self._next = time.monotonic() + self._proc: subprocess.Popen | None = None + self._job: tuple | None = None # (checkpoint signature, files) + self._failed: set = set() + + def poll(self) -> None: + """Install a finished export; start the next one when it's due.""" + if self._proc is not None: + if self._proc.poll() is None: + return + self._finish() + if not self.interval or time.monotonic() < self._next: + return + ckpt = _newest_progress(self.lightning_dir) + sig = None if ckpt is None else _checkpoint_signature(ckpt) + if sig is None or sig in self._failed or _exported_from(ckpt, self.onnx): + self._next = time.monotonic() + min(60.0, self.interval) # nothing new yet + return + self._next = time.monotonic() + self.interval + self._start(ckpt, sig) + + def stop(self) -> None: + """Abandon an export that is still running (it is out of date by now).""" + if self._proc is None: + return + _kill(self._proc) + _sig, files = self._job + self._proc = self._job = None + files.discard() + + def _start(self, ckpt: Path, sig: tuple) -> None: + _remove_stale_exports(self.voice_dir) + files = _ExportFiles(self.onnx) + cmd, env = _export_command(ckpt, files) + flags = (getattr(subprocess, "CREATE_NEW_PROCESS_GROUP", 0) + | getattr(subprocess, "BELOW_NORMAL_PRIORITY_CLASS", 0)) + try: + with open(self.log_path, "w", encoding="utf-8") as log: # the export has its own handle + self._proc = subprocess.Popen(cmd, env=env, stdin=subprocess.DEVNULL, stdout=log, + stderr=subprocess.STDOUT, creationflags=flags) + except BaseException: + files.discard() + raise + self._job = (sig, files) + + def _finish(self) -> None: + sig, files = self._job + rc = self._proc.returncode + self._proc = self._job = None + try: + if rc == 0 and _install_export(files, self.voice_dir) == 0: + print(f"\nVoice updated: it speaks as trained up to " + f"{describe_export(export_info(self.onnx))}.", flush=True) + else: + self._failed.add(sig) + print(f"\nWARNING: could not update the voice with the training progress " + f"(export exit {rc}); see {self.log_path}", file=sys.stderr, flush=True) + finally: + files.discard() + + +def _export_progress(voice_dir: Path) -> None: + """After a stop: let the voice speak as trained so far, unless it already does.""" + onnx = voice_dir / "piper_model.onnx" + ckpt = _newest_progress(voice_dir / "training" / "piper_lightning") + if ckpt is None or _exported_from(ckpt, onnx): + return + print("\nUpdating the voice with the saved training progress (about a minute; press " + "Ctrl+C to skip).", flush=True) + if _export_onnx(ckpt, onnx, voice_dir, desc="ONNX export of the training progress") == 0: + print(f"\nVoice updated: it speaks as trained up to " + f"{describe_export(export_info(onnx))}.") + + def export_snapshot(voice_dir: Path, *, output: Path | None = None) -> int: - """Export the best available training checkpoint to ONNX right now. + """Export the newest usable training checkpoint to ONNX right now. Does not start or resume training — only runs the ONNX export step on - whatever checkpoint exists. Useful for previewing a voice mid-training. + whatever checkpoint exists. Useful for previewing a voice mid-training + (which updates the voice in the background anyway, see :class:`SnapshotExporter`). Args: voice_dir: The voice profile directory (contains ``training/``). @@ -557,10 +1196,14 @@ def export_snapshot(voice_dir: Path, *, output: Path | None = None) -> int: Returns the subprocess exit code (0 = success). """ lightning_dir = voice_dir / "training" / "piper_lightning" - ckpt = _find_best_checkpoint(lightning_dir) + ckpt, problems = None, [] + if lightning_dir.is_dir(): + ckpt, _, problems = _newest_usable_checkpoint(lightning_dir, "32") + for problem in problems: + print(f"WARNING: passing over {problem}", file=sys.stderr) if ckpt is None: print( - f"ERROR: no checkpoint found under {lightning_dir}\n" + f"ERROR: no usable checkpoint found under {lightning_dir}\n" "Start training first with: python -m javis.voices.piper_train --voice ", file=sys.stderr, ) @@ -575,7 +1218,7 @@ def export_snapshot(voice_dir: Path, *, output: Path | None = None) -> int: return rc print(f"\nSnapshot ready: {onnx_out}") - print("(From checkpoint:", ckpt.name, ")") + print(f"(From checkpoint {ckpt.name}: {describe_export(export_info(onnx_out))})") return 0 @@ -615,10 +1258,17 @@ def _cli() -> int: p.add_argument("--accelerator", default="gpu", choices=["gpu", "cpu"], help="PyTorch Lightning accelerator (default: gpu)") - p.add_argument("--precision", default="16-mixed", + p.add_argument("--precision", default=DEFAULT_PRECISION, choices=["32", "16-mixed", "bf16-mixed"], - help="Training precision (default: 16-mixed; use 32 for debugging)") + help=f"Training precision (default: {DEFAULT_PRECISION}, the fastest; " + "falls back to 32 on GPUs without native bf16). 16-mixed works " + "but is ~1.4x slower.") p.add_argument("--config", type=Path, default=_project_root / "config.yaml") + p.add_argument("--snapshot-minutes", type=float, default=SNAPSHOT_MINUTES, + help="While training, update the voice's model from the training " + f"progress this often (default: {SNAPSHOT_MINUTES:g}; 0 = only " + "when training ends or is stopped). The export runs in the " + "background at low priority, on the CPU.") p.add_argument( "--export-snapshot", action="store_true", help="Export the best existing checkpoint to ONNX now (no training). " @@ -637,19 +1287,21 @@ def _cli() -> int: prof = mgr.get(args.voice) assert prof.dir is not None, "Voice profile has no directory" - if args.export_snapshot: - return export_snapshot(prof.dir, output=args.snapshot_output) + with console_selection_cannot_pause(): + if args.export_snapshot: + return export_snapshot(prof.dir, output=args.snapshot_output) - return run_pipeline( - prof.dir, - language=args.language, - base_checkpoint=args.base_checkpoint, - batch_size=args.batch_size, - num_workers=args.num_workers, - max_epochs=args.max_epochs, - accelerator=args.accelerator, - precision=args.precision, - ) + return run_pipeline( + prof.dir, + language=args.language, + base_checkpoint=args.base_checkpoint, + batch_size=args.batch_size, + num_workers=args.num_workers, + max_epochs=args.max_epochs, + accelerator=args.accelerator, + precision=args.precision, + snapshot_minutes=args.snapshot_minutes, + ) if __name__ == "__main__": diff --git a/src/javis/voices/training_control.py b/src/javis/voices/training_control.py new file mode 100644 index 0000000..55e909c --- /dev/null +++ b/src/javis/voices/training_control.py @@ -0,0 +1,495 @@ +"""Coordination between Javis and a running piper training process. + +Shared by the launcher (:mod:`javis.voices.piper_train`, the overnight +scheduler, the GUI) and the training process (:mod:`javis.voices.piper_compat` +and its data-loader workers), so it imports nothing but the standard library. + +* **Clean stops.** The launcher names a stop file in ``JAVIS_PIPER_STOP_FILE``; + once the file exists the trainer finishes its current step, saves its + progress — the epoch so far included — and exits with :data:`EXIT_STOPPED`. + It does the same if the launcher (``JAVIS_PIPER_LAUNCHER``) has gone away, + instead of training on unsupervised. +* **Closing the console window** (or logging off, shutting down) gives the + trainer the few seconds Windows allows to save its progress: + :class:`ConsoleClose`. :func:`python_command` starts it so that none of them + is lost: without the virtual environment's ``python.exe`` launcher, which + would end it at once, and without torch's Intel Fortran console handler. +* **One run per voice.** :class:`TrainingLock` records the processes of the run + holding it, so a crashed or killed run never blocks the next one. +* **Long runs on Windows.** :func:`keep_system_awake` and + :func:`console_selection_cannot_pause`. +""" +from __future__ import annotations + +import contextlib +import functools +import json +import os +import platform +import sys +import threading +import time +from pathlib import Path + +STOP_FILE_ENV = "JAVIS_PIPER_STOP_FILE" +LAUNCHER_ENV = "JAVIS_PIPER_LAUNCHER" # ":" of the supervising process + +# Exit codes of ``python -m javis.voices.piper_compat fit`` and of run_pipeline(). +EXIT_DIVERGED = 3 # training kept producing NaN/inf: restarting would not help +EXIT_STOPPED = 4 # stopped on request, after saving a checkpoint +EXIT_BUSY = 5 # run_pipeline(): another run of this voice holds the lock + +LOCK_FILE = ".training.lock" +_LEGACY_LOCK_SECONDS = 48 * 60 * 60 # locks without processes (older Javis versions) +_UNREADABLE_LOCK_SECONDS = 60 # a lock file that is being written + +PARTIAL_SUFFIX = ".partial" # a file being written; renamed into place once complete +# Windows: the target is open in another process for a moment (a snapshot export +# reading last.ckpt, a virus scanner): 32/33 sharing/lock violation, 1224 mapped. +_BUSY_WINERRORS = (32, 33, 1224) + + +def replace_file(src, dst, *, seconds: float = 60.0) -> None: + """``os.replace`` that gives other processes up to *seconds* to close *dst*.""" + delay, deadline = 0.1, time.monotonic() + seconds + while True: + try: + os.replace(src, dst) + return + except OSError as exc: + if not (isinstance(exc, PermissionError) + or getattr(exc, "winerror", None) in _BUSY_WINERRORS) \ + or time.monotonic() >= deadline: + raise + time.sleep(delay) + delay = min(delay * 2, 2.0) + + +# ── Processes ───────────────────────────────────────────────────────────────── + +@functools.cache +def _kernel32(): + import ctypes + from ctypes import wintypes + + # A private instance: argtypes set on ctypes.windll.kernel32 would leak to other code. + k = ctypes.WinDLL("kernel32", use_last_error=True) + k.OpenProcess.argtypes = (wintypes.DWORD, wintypes.BOOL, wintypes.DWORD) + k.OpenProcess.restype = wintypes.HANDLE + k.GetExitCodeProcess.argtypes = (wintypes.HANDLE, ctypes.POINTER(wintypes.DWORD)) + k.GetExitCodeProcess.restype = wintypes.BOOL + k.GetProcessTimes.argtypes = (wintypes.HANDLE,) + (ctypes.POINTER(wintypes.FILETIME),) * 4 + k.GetProcessTimes.restype = wintypes.BOOL + k.CloseHandle.argtypes = (wintypes.HANDLE,) + k.CloseHandle.restype = wintypes.BOOL + k.SetThreadExecutionState.argtypes = (wintypes.DWORD,) + k.SetThreadExecutionState.restype = wintypes.DWORD + k.GetStdHandle.argtypes = (wintypes.DWORD,) + k.GetStdHandle.restype = wintypes.HANDLE + k.GetConsoleMode.argtypes = (wintypes.HANDLE, ctypes.POINTER(wintypes.DWORD)) + k.GetConsoleMode.restype = wintypes.BOOL + k.SetConsoleMode.argtypes = (wintypes.HANDLE, wintypes.DWORD) + k.SetConsoleMode.restype = wintypes.BOOL + k.SetConsoleCtrlHandler.argtypes = (ctypes.c_void_p, wintypes.BOOL) + k.SetConsoleCtrlHandler.restype = wintypes.BOOL + k.SetProcessShutdownParameters.argtypes = (wintypes.DWORD, wintypes.DWORD) + k.SetProcessShutdownParameters.restype = wintypes.BOOL + k.GetModuleHandleW.argtypes = (wintypes.LPCWSTR,) + k.GetModuleHandleW.restype = wintypes.HMODULE + return k + + +@functools.cache +def _user32(): + """``(user32, WNDPROC, WNDCLASSW)``: what a hidden window takes.""" + import ctypes + from ctypes import wintypes + + wndproc = ctypes.WINFUNCTYPE(wintypes.LPARAM, wintypes.HWND, wintypes.UINT, + wintypes.WPARAM, wintypes.LPARAM) + + class WNDCLASSW(ctypes.Structure): + _fields_ = [("style", wintypes.UINT), ("lpfnWndProc", wndproc), + ("cbClsExtra", ctypes.c_int), ("cbWndExtra", ctypes.c_int), + ("hInstance", wintypes.HINSTANCE), ("hIcon", wintypes.HICON), + ("hCursor", wintypes.HANDLE), ("hbrBackground", wintypes.HBRUSH), + ("lpszMenuName", wintypes.LPCWSTR), ("lpszClassName", wintypes.LPCWSTR)] + + u = ctypes.WinDLL("user32", use_last_error=True) + u.RegisterClassW.argtypes = (ctypes.POINTER(WNDCLASSW),) + u.RegisterClassW.restype = wintypes.ATOM + u.UnregisterClassW.argtypes = (wintypes.LPCWSTR, wintypes.HINSTANCE) + u.UnregisterClassW.restype = wintypes.BOOL + u.CreateWindowExW.argtypes = (wintypes.DWORD, wintypes.LPCWSTR, wintypes.LPCWSTR, + wintypes.DWORD, ctypes.c_int, ctypes.c_int, ctypes.c_int, + ctypes.c_int, wintypes.HWND, wintypes.HMENU, + wintypes.HINSTANCE, wintypes.LPVOID) + u.CreateWindowExW.restype = wintypes.HWND + u.DefWindowProcW.argtypes = (wintypes.HWND, wintypes.UINT, wintypes.WPARAM, + wintypes.LPARAM) + u.DefWindowProcW.restype = wintypes.LPARAM + u.GetMessageW.argtypes = (ctypes.POINTER(wintypes.MSG), wintypes.HWND, wintypes.UINT, + wintypes.UINT) + u.GetMessageW.restype = wintypes.BOOL + u.DispatchMessageW.argtypes = (ctypes.POINTER(wintypes.MSG),) + u.DispatchMessageW.restype = wintypes.LPARAM + u.PostQuitMessage.argtypes = (ctypes.c_int,) + u.PostQuitMessage.restype = None + return u, wndproc, WNDCLASSW + + +def python_command() -> tuple[str, dict]: + """``(executable, extra environment)`` to start another Python of this environment. + + On Windows a virtual environment's ``python.exe`` is a launcher running the + real interpreter in a job that dies with it — and the launcher ends as soon + as its console window closes, so the interpreter couldn't save anything. + Like :mod:`multiprocessing`, start the real interpreter directly; + ``__PYVENV_LAUNCHER__`` tells it which environment it belongs to. + + Nor does it get the console handler of the Intel Fortran runtime in torch's + ``libiomp5md.dll``: on a closing window that aborts every process loading + torch ("forrtl: error (200)"), and slowly — Windows tells the processes of a + console one after the other, the newest first: data-loader workers and + exports before the trainer, which had seconds less to save its progress. + """ + if sys.platform != "win32": + return sys.executable, {} + env = {"FOR_DISABLE_CONSOLE_CTRL_HANDLER": "1"} # inherited by data-loader workers + base = getattr(sys, "_base_executable", "") or "" + if base and os.path.isfile(base) and os.path.normcase(os.path.abspath(base)) \ + != os.path.normcase(os.path.abspath(sys.executable)): + return base, {**env, "__PYVENV_LAUNCHER__": sys.executable} + return sys.executable, env + + +def process_start_time(pid: int) -> int | None: + """When process *pid* started, as an opaque number; ``None`` if it isn't running. + + ``0`` means it runs but its start time can't be read. + """ + if pid <= 0: + return None + if sys.platform != "win32": + try: + os.kill(pid, 0) # signal 0 only checks; on Windows os.kill would terminate + except PermissionError: + return 0 + except OSError: + return None + return 0 + + import ctypes + from ctypes import wintypes + + k = _kernel32() + handle = k.OpenProcess(0x1000, False, pid) # PROCESS_QUERY_LIMITED_INFORMATION + if not handle: + return 0 if ctypes.get_last_error() == 5 else None # 5: access denied, yet it exists + try: + code = wintypes.DWORD() + if k.GetExitCodeProcess(handle, ctypes.byref(code)) and code.value != 259: # STILL_ACTIVE + return None + created, exited, kernel, user = (wintypes.FILETIME() for _ in range(4)) + if not k.GetProcessTimes(handle, ctypes.byref(created), ctypes.byref(exited), + ctypes.byref(kernel), ctypes.byref(user)): + return 0 + return (created.dwHighDateTime << 32) | created.dwLowDateTime + finally: + k.CloseHandle(handle) + + +def process_alive(pid: int, start_time: int | None = None) -> bool: + """Whether *pid* is running and, given its *start_time*, is still the same process. + + Windows soon hands the ID of an exited process to a new one. + """ + current = process_start_time(pid) + if current is None: + return False + return not start_time or not current or current == start_time + + +def this_process() -> dict: + """``{"pid", "start"}`` of the calling process.""" + pid = os.getpid() + return {"pid": pid, "start": process_start_time(pid) or 0} + + +def launcher_from_environment() -> tuple[int, int] | None: + """``(pid, start time)`` of the launcher named in ``JAVIS_PIPER_LAUNCHER``, if any.""" + value = os.environ.get(LAUNCHER_ENV, "") + try: + pid, start = (int(part) for part in value.split(":")) + except ValueError: + return None + return pid, start + + +# ── One training run per voice ──────────────────────────────────────────────── + +def lock_path(voice_dir: Path) -> Path: + return Path(voice_dir) / "training" / LOCK_FILE + + +def _read_lock(path: Path) -> str | None: + """The lock file's text; ``None`` if there is none. Raises other ``OSError``s.""" + try: + return path.read_text(encoding="utf-8") + except FileNotFoundError: + return None + + +def _held(path: Path, text: str) -> bool: + try: + age = time.time() - path.stat().st_mtime + except FileNotFoundError: + return False + except OSError: + return True # can't tell: assume it is held + try: + info = json.loads(text) + except ValueError: + info = None + if not isinstance(info, dict): + return age < _UNREADABLE_LOCK_SECONDS + try: + processes = [(int(p.get("pid") or 0), int(p.get("start") or 0)) + for p in info.get("processes") or () if isinstance(p, dict)] + except (TypeError, ValueError): + processes = [] + if not processes or info.get("host", platform.node()) != platform.node(): + try: + started = float(info.get("started_epoch") or 0) or path.stat().st_mtime + except (OSError, TypeError, ValueError): + return True + return time.time() - started < _LEGACY_LOCK_SECONDS + return any(process_alive(pid, start) for pid, start in processes) + + +def _lock_held(path: Path) -> bool: + try: + text = _read_lock(path) + except OSError: + return True + return text is not None and _held(path, text) + + +def training_in_progress(voice_dir: Path) -> bool: + """Whether a training run of this voice is running (in any process).""" + return _lock_held(lock_path(voice_dir)) + + +class TrainingLock: + """``training/.training.lock``: one training run per voice at a time, across processes. + + Lists the launcher and, while one runs, its training process (which keeps + training for a minute or two after its launcher was killed, to save its + progress). A lock whose processes have all exited is stale and replaced. + """ + + def __init__(self, voice_dir: Path) -> None: + self.path = lock_path(voice_dir) + self._info: dict | None = None + + def acquire(self) -> bool: + """Take the lock; ``False`` if another run holds it.""" + info = {"host": platform.node(), "started_epoch": time.time(), + "processes": [this_process()]} + self.path.parent.mkdir(parents=True, exist_ok=True) + for _ in range(2): + try: + with open(self.path, "x", encoding="utf-8") as f: + json.dump(info, f) + except FileExistsError: + try: + text = _read_lock(self.path) + if text is not None: + if _held(self.path, text): + return False + # Left behind by a run that crashed or was killed — unless + # another launcher replaced it just now. + if _read_lock(self.path) == text: + self.path.unlink() + except FileNotFoundError: + pass + except OSError: + return False + continue + self._info = info + return True + return False + + def set_trainer(self, pid: int | None) -> None: + """Record the training process (or that none runs right now).""" + if self._info is None: + return + trainer = [{"pid": pid, "start": process_start_time(pid) or 0}] if pid else [] + self._info["processes"] = self._info["processes"][:1] + trainer + with contextlib.suppress(OSError): + self.path.write_text(json.dumps(self._info), encoding="utf-8") + + def release(self) -> None: + if self._info is None: + return + owner, self._info = self._info["processes"][0], None + try: + current = json.loads(self.path.read_text(encoding="utf-8")) + if (current.get("processes") or [{}])[0] == owner: + self.path.unlink() + except (OSError, ValueError, AttributeError): + pass + + +# ── Windows ─────────────────────────────────────────────────────────────────── + +@contextlib.contextmanager +def keep_system_awake(): + """Keep Windows from sleeping when idle (the display may still turn off).""" + if sys.platform != "win32": + yield + return + k = _kernel32() + es_continuous, es_system_required = 0x80000000, 0x00000001 + previous = k.SetThreadExecutionState(es_continuous | es_system_required) + try: + yield + finally: + k.SetThreadExecutionState(previous if previous & es_continuous else es_continuous) + + +@contextlib.contextmanager +def console_selection_cannot_pause(): + """Turn off QuickEdit in a classic console window while training runs. + + There, clicking into the window starts a selection that suspends every + process writing to it — training would stall until a key is pressed. + (Windows Terminal doesn't do that.) + """ + if sys.platform != "win32": + yield + return + import ctypes + from ctypes import wintypes + + k = _kernel32() + handle = k.GetStdHandle(wintypes.DWORD(-10).value) # STD_INPUT_HANDLE + mode = wintypes.DWORD() + quick_edit, extended_flags = 0x0040, 0x0080 + if not handle or not k.GetConsoleMode(handle, ctypes.byref(mode)) \ + or not mode.value & quick_edit: + yield + return + k.SetConsoleMode(handle, (mode.value | extended_flags) & ~quick_edit) + try: + yield + finally: + k.SetConsoleMode(handle, mode.value | extended_flags) + + +_CONSOLE_CLOSE_EVENTS = {2: "the console window was closed", 5: "the user is logging off", + 6: "the system is shutting down"} +_WM_DESTROY, _WM_QUERYENDSESSION, _WM_ENDSESSION = 0x0002, 0x0011, 0x0016 +_ENDSESSION_LOGOFF = 0x80000000 +# Ends before the other processes of the session — its data-loader workers +# among them: the top of the range Windows leaves to applications. +_SHUTDOWN_LEVEL = 0x3FF + + +class ConsoleClose: + """Lets the training process save its progress before Windows ends it. + + When its console window is closed (or the user logs off, the system shuts + down), Windows ends every process attached to it after a few seconds. + :meth:`install` handles those events: it sets :attr:`closing` — the training + loop checks it after every step — and holds the process until :meth:`done` + reports the progress saved, or *timeout* seconds have passed. + + Logging off and shutting down reach a console handler only in processes + that never loaded ``user32.dll`` — Python and torch load it. A hidden + window learns of them instead (``WM_ENDSESSION``). + """ + + def __init__(self, timeout: float = 30.0) -> None: + self.closing = threading.Event() + self.reason = "" + self.timeout = timeout + self._saved = threading.Event() + self._handler = None + self._window = None # the hidden window's handle + + def handle(self, event: int) -> bool: + """The console control handler; runs in a thread of its own.""" + reason = _CONSOLE_CLOSE_EVENTS.get(event) + if reason is None: + return False # Ctrl+C, Ctrl+Break: left to Python's own handler + self._close(reason) + return True + + def _close(self, reason: str) -> None: + self.reason = reason + self.closing.set() + self._saved.wait(self.timeout) + + def done(self) -> None: + """The progress is saved: Windows may end the process now.""" + self._saved.set() + + def install(self) -> None: + if sys.platform != "win32" or self._handler is not None: + return + import ctypes + from ctypes import wintypes + + handler = ctypes.WINFUNCTYPE(wintypes.BOOL, wintypes.DWORD)(self.handle) + if _kernel32().SetConsoleCtrlHandler(handler, True): + self._handler = handler # must outlive every call Windows may make + _kernel32().SetProcessShutdownParameters(_SHUTDOWN_LEVEL, 0) + self._watch_session_end() + + def _watch_session_end(self) -> None: + """Open the hidden window that learns of the user logging off, the system shutting down.""" + ready = threading.Event() + threading.Thread(target=self._session_window, args=(ready,), name="session end", + daemon=True).start() + ready.wait(10) + + def _session_window(self, ready: threading.Event) -> None: + import ctypes + from ctypes import wintypes + + try: + user32, wndproc, wndclass = _user32() + instance = _kernel32().GetModuleHandleW(None) + + def window_procedure(hwnd, message, wparam, lparam): + if message == _WM_QUERYENDSESSION: + return 1 # may end: the progress gets saved at WM_ENDSESSION + if message == _WM_ENDSESSION: + if wparam: # else cancelled + self._close("the user is logging off" if lparam & _ENDSESSION_LOGOFF + else "the system is shutting down") + return 0 + if message == _WM_DESTROY: + user32.PostQuitMessage(0) + return 0 + return user32.DefWindowProcW(hwnd, message, wparam, lparam) + + procedure = wndproc(window_procedure) # referenced while the window exists + name = f"JavisTrainingSession{os.getpid()}-{id(self):x}" + if not user32.RegisterClassW(ctypes.byref(wndclass( + lpfnWndProc=procedure, hInstance=instance, lpszClassName=name))): + return + # Top-level but never shown: message-only windows get no WM_ENDSESSION. + window = user32.CreateWindowExW(0, name, "Javis voice training", 0, 0, 0, 0, 0, + None, None, instance, None) + if not window: + user32.UnregisterClassW(name, instance) + return + self._window = window + finally: + ready.set() + message = wintypes.MSG() + while user32.GetMessageW(ctypes.byref(message), None, 0, 0) > 0: + user32.DispatchMessageW(ctypes.byref(message)) + user32.UnregisterClassW(name, instance) diff --git a/src/javis/voices/training_scheduler.py b/src/javis/voices/training_scheduler.py index 2207243..8f7b717 100644 --- a/src/javis/voices/training_scheduler.py +++ b/src/javis/voices/training_scheduler.py @@ -18,31 +18,29 @@ Design notes without a restart. * **Window guard.** Training will not *start* if fewer than ``min_remaining_minutes`` of the window are left. A run that has already - started is allowed to run to completion — PyTorch-Lightning cannot be - cleanly killed mid-epoch, and an unfinished run is safe: it leaves a - checkpoint that the next run resumes from. -* **Concurrency guard.** A lock file (``training/.training.lock``) prevents - the scheduler and the interactive GUI button from training the same voice - at the same time (they run in separate processes). + started keeps going; quitting Javis stops it cleanly — the training + process saves its progress after its current step (the epoch so far + included) first, so Javis's process can take a little while to exit — + and the next run resumes from where it stopped. While it runs, the + voice's model is updated with its progress (see ``piper_train``). +* **Concurrency guard.** ``run_pipeline`` holds a lock file + (``training/.training.lock``) naming its processes, so the scheduler and + the interactive GUI button (separate processes) never train the same voice + at the same time; a run that crashed or was killed leaves no lasting lock. """ from __future__ import annotations import asyncio -import json -import time from datetime import datetime, timedelta from pathlib import Path from typing import Optional from ..core.config import Config from ..core.logger import get_logger +from .training_control import EXIT_BUSY, EXIT_STOPPED, training_in_progress log = get_logger("voices.training") -# A lock is considered "live" for this long. Overnight runs take a few -# hours; 48 h comfortably bounds a stuck/crashed run so it self-expires. -_LOCK_STALE_SECONDS = 48 * 60 * 60 - # ── Pure helpers (unit-testable) ───────────────────────────────────────────── @@ -97,62 +95,10 @@ def window_end(now: datetime, start_hour: int, end_hour: int) -> datetime: return end -# ── Cross-process lock ──────────────────────────────────────────────────────── - -def _lock_path(voice_dir: Path) -> Path: - return voice_dir / "training" / ".training.lock" - - -def _lock_alive(voice_dir: Path) -> bool: - """True if a recent lock file indicates another training is in progress.""" - lp = _lock_path(voice_dir) - if not lp.exists(): - return False - try: - data = json.loads(lp.read_text(encoding="utf-8")) - ts = float(data.get("started_epoch", 0.0)) - except Exception: - try: - ts = lp.stat().st_mtime - except OSError: - return False - if not ts: - return False - return (time.time() - ts) < _LOCK_STALE_SECONDS - - -def _acquire_lock(voice_dir: Path) -> bool: - """Atomically claim the training lock. Returns False if already held.""" - lp = _lock_path(voice_dir) - lp.parent.mkdir(parents=True, exist_ok=True) - try: - # 'x' = exclusive create; fails if a live lock already exists. - with open(lp, "x", encoding="utf-8") as f: - f.write(json.dumps({"started_epoch": time.time()})) - return True - except FileExistsError: - return False - except OSError: - # On platforms where 'x' is flaky, fall back to the liveness check. - return not _lock_alive(voice_dir) - - -def _release_lock(voice_dir: Path) -> None: - try: - _lock_path(voice_dir).unlink(missing_ok=True) - except OSError: - pass - - -def training_in_progress(voice_dir: Path) -> bool: - """Public helper: is another training of this voice currently running?""" - return _lock_alive(voice_dir) - - # ── One run ─────────────────────────────────────────────────────────────────── async def _run_once(config: Config, voice_dir: Path, bus: Optional[object] = None) -> None: - """Execute a single scheduled training run (threaded) with locking.""" + """Execute a single scheduled training run (threaded).""" from .piper_train import run_pipeline # imported lazily (heavy deps) t = config.training @@ -166,14 +112,10 @@ async def _run_once(config: Config, voice_dir: Path, bus: Optional[object] = Non f"{t.min_remaining_minutes} min minimum); skipping start.") return - if _lock_alive(voice_dir): + if training_in_progress(voice_dir): log.info("Another training of this voice is already running; skipping.") return - if not _acquire_lock(voice_dir): - log.info("Could not acquire training lock (already held); skipping.") - return - voice = voice_dir.name n_clips = 0 meta = voice_dir / "training" / "metadata.csv" @@ -214,25 +156,30 @@ async def _run_once(config: Config, voice_dir: Path, bus: Optional[object] = Non try: rc = await asyncio.shield(work) except asyncio.CancelledError: - # Javis is quitting: stop the child process so the worker thread (and - # interpreter exit) doesn't block until training finishes. - from .piper_train import request_stop + # Javis is quitting: have the training process save a checkpoint and + # stop, so the worker thread (and interpreter exit) isn't kept waiting + # until training finishes. + from .piper_train import STOP_GRACE_SECONDS, request_stop request_stop() - await asyncio.wait({work}, timeout=60) + await asyncio.wait({work}, timeout=STOP_GRACE_SECONDS + 30) raise except Exception as exc: # noqa: BLE001 — training must never kill the loop log.exception(f"Scheduled training crashed: {exc}") rc = -1 - finally: - if work.done(): - _release_lock(voice_dir) - status = "completed" if rc == 0 else f"FAILED (exit {rc})" + if rc == 0: + status = "completed" + elif rc == EXIT_STOPPED: + status = "stopped (progress saved)" + elif rc == EXIT_BUSY: + status = "skipped (this voice is being trained already)" + else: + status = f"FAILED (exit {rc})" log.info(f"Scheduled training {status}: voice={voice!r}") if bus is not None: await _safe_publish( bus, "task_done", - {"id": "training", "ok": rc == 0, + {"id": "training", "ok": rc in (0, EXIT_STOPPED, EXIT_BUSY), "msg": f"Voice training {status}"}) diff --git a/tests/test_audiobook_ingest.py b/tests/test_audiobook_ingest.py new file mode 100644 index 0000000..390298d --- /dev/null +++ b/tests/test_audiobook_ingest.py @@ -0,0 +1,123 @@ +"""Audiobook ingest: raw-cache transcoding must terminate and clean up.""" +from __future__ import annotations + +import asyncio +import sys +import tempfile +import threading +import unittest +from collections import namedtuple +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import MagicMock, PropertyMock, patch + +import numpy as np + +SRC_SR = 44100 +SRC_SECONDS = 25 # two full 10 s blocks + one partial block +TARGET_SR = 22050 + +_DiskUsage = namedtuple("_DiskUsage", "total used free") + + +class TranscodeToRawTests(unittest.TestCase): + def setUp(self) -> None: + try: + import soundfile as sf + except ImportError: + self.skipTest("soundfile not installed") + from javis.voices import audiobook_ingest + self.sf = sf + self.ingest = audiobook_ingest + self._tmp = tempfile.TemporaryDirectory() + self.root = Path(self._tmp.name) + self.source = self.root / "book.wav" + t = np.arange(SRC_SR * SRC_SECONDS) / SRC_SR + tone = 0.3 * np.sin(2 * np.pi * 220 * t) + sf.write(str(self.source), np.stack([tone, tone], axis=1), SRC_SR, subtype="PCM_16") + self.target = self.root / "cache" / "ingest_audio_cache.i16" + + def tearDown(self) -> None: + self._tmp.cleanup() + + def _bogus_frame_count(self): + """Mimic an MP3 with a malformed ID3 tag: INT64_MAX frames, unseekable.""" + frames = patch.object(self.sf.SoundFile, "frames", new_callable=PropertyMock, + return_value=sys.maxsize) + seekable = patch.object(self.sf.SoundFile, "seekable", return_value=False) + return frames, seekable + + def _run_with_timeout(self, timeout: float = 60.0, **kw) -> BaseException | None: + errors: list[BaseException] = [] + + def run(): + try: + self.ingest._transcode_to_raw(self.source, self.target, TARGET_SR, **kw) + except BaseException as e: # noqa: BLE001 + errors.append(e) + + th = threading.Thread(target=run, daemon=True) + th.start() + th.join(timeout) + self.assertFalse(th.is_alive(), "transcode did not stop at EOF (runaway)") + return errors[0] if errors else None + + def test_unknown_frame_count_stops_at_eof(self) -> None: + frames, seekable = self._bogus_frame_count() + with frames, seekable: + self.assertGreaterEqual(self.sf.info(str(self.source)).frames, sys.maxsize // 2) + err = self._run_with_timeout() + self.assertIsNone(err) + self.assertEqual(self.target.stat().st_size, SRC_SECONDS * TARGET_SR * 2) + self.assertFalse(self.target.with_suffix(".part").exists()) + + def test_stop_request_aborts_and_removes_partial_file(self) -> None: + err = self._run_with_timeout(should_stop=lambda: True) + self.assertIsInstance(err, self.ingest.TranscodeCancelled) + self.assertFalse(self.target.exists()) + self.assertFalse(self.target.with_suffix(".part").exists()) + + def test_disk_running_full_aborts_and_removes_partial_file(self) -> None: + frames, seekable = self._bogus_frame_count() + full = patch.object(self.ingest.shutil, "disk_usage", + return_value=_DiskUsage(10**12, 10**12, 0)) + every_block = patch.object(self.ingest, "_DISK_CHECK_EVERY_BLOCKS", 1) + with frames, seekable, full, every_block: + err = self._run_with_timeout() + self.assertIsInstance(err, RuntimeError) + self.assertIn("free space", str(err)) + self.assertFalse(self.target.exists()) + self.assertFalse(self.target.with_suffix(".part").exists()) + + def test_insufficient_space_for_known_length_fails_fast(self) -> None: + with patch.object(self.ingest.shutil, "disk_usage", + return_value=_DiskUsage(10**12, 10**12, 0)): + err = self._run_with_timeout() + self.assertIsInstance(err, RuntimeError) + self.assertIn("Not enough disk space", str(err)) + self.assertFalse(self.target.with_suffix(".part").exists()) + + def test_stop_button_during_transcode_ends_ingest_as_stopped(self) -> None: + events: list[dict] = [] + + class Bus: + async def publish(self, topic, payload): + events.append(payload) + + mgr = MagicMock() + mgr.root = self.root + mgr.create.return_value = SimpleNamespace(dir=self.root / "voice") + stop = asyncio.Event() + stop.set() + with patch("javis.voices.manager.VoiceManager", return_value=mgr): + stats = asyncio.run(self.ingest.ingest_audiobook( + None, self.source, "voice", bus=Bus(), stop_event=stop, resume=False)) + self.assertEqual(stats.total_segments, 0) + self.assertEqual(events[-1]["phase"], "stopped") + self.assertEqual(events[-1]["during"], "transcoding") + self.assertFalse((self.root / "voice" / "ingest_audio_cache.i16").exists()) + self.assertFalse((self.root / "voice" / "ingest_audio_cache.part").exists()) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_slave_chat_ui.py b/tests/test_slave_chat_ui.py new file mode 100644 index 0000000..a423057 --- /dev/null +++ b/tests/test_slave_chat_ui.py @@ -0,0 +1,120 @@ +"""Slave GUI chat: newest-first transcript and fullscreen chat mode (offscreen Qt).""" +from __future__ import annotations + +import os +import sys +from pathlib import Path + +import pytest + +os.environ.setdefault("QT_QPA_PLATFORM", "offscreen") + +SLAVE_SRC = Path(__file__).resolve().parents[1] / "javis-slave" / "src" +if str(SLAVE_SRC) not in sys.path: + sys.path.insert(0, str(SLAVE_SRC)) + +QtWidgets = pytest.importorskip("PySide6.QtWidgets") +QtTest = pytest.importorskip("PySide6.QtTest") + +from javis_slave import main_window # noqa: E402 + + +@pytest.fixture(scope="module") +def qapp(): + return QtWidgets.QApplication.instance() or QtWidgets.QApplication([]) + + +@pytest.fixture +def screen(qapp): + scr = main_window.MainScreen() + scr.resize(900, 700) + scr.show() + yield scr + scr.close() + + +def _texts(scr) -> list[str]: + """Bubble texts from the top of the transcript to the bottom.""" + layout = scr.transcript_layout + out = [] + for i in range(layout.count()): + row = layout.itemAt(i).widget() + if row is not None: + out.append(row.findChild(QtWidgets.QLabel).text()) + return out + + +def test_newest_message_is_on_top(screen): + screen.set_turn("frage 1", "antwort 1") + screen.set_turn("frage 2", "antwort 2") + screen.show_error("kaputt") + assert _texts(screen) == [ + "⚠ kaputt", "antwort 2", "frage 2", "antwort 1", "frage 1"] + + +def test_prune_drops_oldest_at_the_bottom(screen): + for i in range(30): + screen.set_turn(f"u{i}", f"a{i}") + texts = _texts(screen) + assert len(texts) == 40 + assert texts[0] == "a29" + assert "u0" not in texts and "u29" in texts + + +def test_scrolls_to_top_on_new_message(screen): + for i in range(12): + screen.set_turn(f"u{i}", "x " * 200) + QtTest.QTest.qWait(150) # lets the layout settle (and the first auto-scroll fire) + vbar = screen.transcript_scroll.verticalScrollBar() + assert vbar.maximum() > 0 + vbar.setValue(vbar.maximum()) + screen.set_turn("neu", "neu") + QtTest.QTest.qWait(150) # the scroll is deferred by a 50 ms timer + assert vbar.value() == vbar.minimum() + + +def test_fullscreen_button_goes_fullscreen_and_keeps_orb_and_chat_visible(qapp): + win = main_window.SlaveWindow() + win.show_main() + win.show() + scr = win.main_screen + scr.set_turn("frage", "antwort") + try: + assert not win.isFullScreen() + assert scr.expand_btn.toolTip() == "Vollbild" + + scr.expand_btn.click() + QtTest.QTest.qWait(50) + assert win.isFullScreen() + assert scr.orb.isVisible() and scr.transcript_scroll.isVisible() + assert scr.expand_btn.toolTip() == "Vollbild beenden" + + scr.expand_btn.click() + QtTest.QTest.qWait(50) + assert not win.isFullScreen() + assert scr.orb.isVisible() and scr.transcript_scroll.isVisible() + assert scr.expand_btn.toolTip() == "Vollbild" + finally: + win.close() + + +def test_esc_and_f11_keep_button_in_sync(qapp): + win = main_window.SlaveWindow() + win.show_main() + win.show() + try: + win._toggle_fullscreen() # F11 + QtTest.QTest.qWait(50) + assert win.isFullScreen() and win.main_screen.expand_btn.toolTip() == "Vollbild beenden" + + win._leave_fullscreen() # Esc + QtTest.QTest.qWait(50) + assert not win.isFullScreen() and win.main_screen.expand_btn.toolTip() == "Vollbild" + finally: + win.close() + +def test_expand_icons_render(qapp): + for kind in ("expand", "collapse"): + img = main_window._icon(kind, "#cccccc", 26).pixmap(26, 26).toImage() + assert any(img.pixelColor(x, y).alpha() > 0 + for x in range(img.width()) for y in range(img.height())) diff --git a/tests/test_speaker_inbox.py b/tests/test_speaker_inbox.py new file mode 100644 index 0000000..474de71 --- /dev/null +++ b/tests/test_speaker_inbox.py @@ -0,0 +1,335 @@ +"""Unidentified-speaker inbox, sorting clips into profiles, and live samples.""" +from __future__ import annotations + +import asyncio +import tempfile +import unittest +from pathlib import Path +from unittest.mock import patch + +import numpy as np + +from javis.core.config import Config, PathsConfig +from javis.core.events import EventBus +from javis.core.state import AgentState +from javis.stt.unidentified import UnidentifiedStore +from javis.stt.voice_profiler import NullProfiler, VoiceProfiler + +SR = 16000 + + +def _clip(voice: int, seconds: float = 1.0) -> np.ndarray: + """Audio whose first sample encodes the (fake) speaker id.""" + a = np.full(int(SR * seconds), 0.2, dtype=np.float32) + a[0] = voice / 100.0 + return a + + +def _fake_embed(audio: np.ndarray, sr: int) -> np.ndarray: + voice = round(float(audio[0]) * 100) + v = np.full(192, 0.01, dtype=np.float32) + v[voice] = 1.0 + return v / np.linalg.norm(v) + + +def _profiler(base: Path) -> VoiceProfiler: + prof = VoiceProfiler(Config(paths=PathsConfig(base_dir=base))) + prof._encoder = object() # pretend the ECAPA model is loaded + prof.embed = _fake_embed # type: ignore[method-assign] + return prof + + +class UnidentifiedStoreTests(unittest.TestCase): + def test_add_list_load_remove_and_persist(self) -> None: + with tempfile.TemporaryDirectory() as d: + store = UnidentifiedStore(Path(d), max_clips=10) + emb = _fake_embed(_clip(3), SR) + a = store.add(_clip(3), SR, embedding=emb, text=" hallo ", source="local", + best_name="Bob", best_score=0.41) + b = store.add(_clip(4, 2.0), SR, text="zwei") + entries = store.entries() + self.assertEqual([e["id"] for e in entries], [b, a]) # newest first + self.assertEqual(entries[1]["text"], "hallo") + self.assertEqual(entries[1]["best_name"], "Bob") + self.assertAlmostEqual(entries[0]["duration_s"], 2.0) + + audio, sr = store.load_audio(a) + self.assertEqual(sr, SR) + self.assertEqual(len(audio), SR) + np.testing.assert_allclose(store.load_embedding(a), emb, rtol=1e-6) + self.assertIsNone(store.load_embedding(b)) + + # A fresh instance reads the same index from disk. + self.assertEqual(len(UnidentifiedStore(Path(d))), 2) + + self.assertEqual(store.remove([a, "nope"]), 1) + self.assertFalse((Path(d) / f"{a}.wav").exists()) + self.assertFalse((Path(d) / f"{a}.npy").exists()) + self.assertEqual(store.clear(), 1) + self.assertEqual(store.entries(), []) + with self.assertRaises(KeyError): + store.load_audio(a) + + def test_oldest_clips_are_pruned(self) -> None: + with tempfile.TemporaryDirectory() as d: + store = UnidentifiedStore(Path(d), max_clips=2) + ids = [store.add(_clip(i + 1), SR, text=str(i)) for i in range(3)] + self.assertEqual([e["id"] for e in store.entries()], [ids[2], ids[1]]) + self.assertFalse((Path(d) / f"{ids[0]}.wav").exists()) + + +class ProfilerEmbeddingTests(unittest.TestCase): + def test_enroll_rank_identify_refine(self) -> None: + with tempfile.TemporaryDirectory() as d: + prof = _profiler(Path(d)) + self.assertTrue(prof.enroll_embeddings("Alice", [_fake_embed(_clip(1), SR)] * 2)) + prof.enroll("Bob", [_clip(2)], SR) + ranked = prof.rank(_fake_embed(_clip(1), SR)) + self.assertEqual(ranked[0][0], "Alice") + self.assertGreater(ranked[0][1], 0.99) + self.assertEqual(prof.identify_embedding(_fake_embed(_clip(2), SR))[0], "Bob") + self.assertIsNone(prof.identify_embedding(_fake_embed(_clip(9), SR))) + self.assertEqual(prof.rank(np.zeros(192)), []) + + # Running centroid: a Bob sample pulls Alice's centroid towards Bob. + before = dict(prof.rank(_fake_embed(_clip(2), SR)))["Alice"] + self.assertTrue(prof.refine_embedding("Alice", _fake_embed(_clip(2), SR))) + after = dict(prof.rank(_fake_embed(_clip(2), SR)))["Alice"] + self.assertGreater(after, before) + alice = next(p for p in prof.list_profiles() if p["name"] == "Alice") + self.assertEqual(alice["clip_count"], 3) + + self.assertFalse(prof.refine_embedding("Nobody", _fake_embed(_clip(1), SR))) + self.assertFalse(prof.refine_embedding("Alice", np.ones(10))) # wrong shape + self.assertFalse(prof.enroll_embeddings("Zero", [np.zeros(192)])) + + def test_sanitised_name_collision_gets_its_own_file(self) -> None: + with tempfile.TemporaryDirectory() as d: + prof = _profiler(Path(d)) + prof.enroll_embeddings("Anna B", [_fake_embed(_clip(1), SR)]) + prof.enroll_embeddings("Anna_B", [_fake_embed(_clip(2), SR)]) + files = {p["name"]: p["file"] for p in prof.list_profiles()} + self.assertNotEqual(files["Anna B"], files["Anna_B"]) + prof2 = _profiler(Path(d)) + prof2._load_index() + self.assertEqual(prof2.rank(_fake_embed(_clip(1), SR))[0][0], "Anna B") + + def test_null_profiler_mirrors_api(self) -> None: + n = NullProfiler() + self.assertFalse(n.is_ready()) + self.assertFalse(n.has_profile("x")) + self.assertEqual(n.rank(np.ones(192)), []) + self.assertFalse(n.enroll_embeddings("x", [np.ones(192)])) + self.assertFalse(n.refine_embedding("x", np.ones(192))) + self.assertIsNone(n.identify_embedding(np.ones(192))) + + +class _FakeMic: + chunk_ms = 32 + + def __init__(self) -> None: + self.drains = 0 + + def drain(self) -> None: + self.drains += 1 + + +def _orchestrator(base: Path): + from javis.agent.orchestrator import Orchestrator + + orch = object.__new__(Orchestrator) + orch.config = Config(paths=PathsConfig(base_dir=base)) + orch.bus = EventBus() + orch.profiler = _profiler(base) + orch.unidentified = UnidentifiedStore(base / "voice_profiles" / "unidentified") + orch._enroll_clip_buffer = [] + orch._enroll_ask_pending = False + orch._turn_lock = asyncio.Lock() + orch._mic_free = asyncio.Event() + orch._mic_free.set() + orch.mic = _FakeMic() + orch.vad = None + orch.state = AgentState.IDLE + events: list[tuple[str, object]] = [] + + async def _record(name, payload=None): + events.append((name, payload)) + + for ev in ("unidentified_changed", "voice_profiles_changed", "speaker_enrolled", + "speaker_ask_name"): + orch.bus.subscribe(ev, lambda p, _n=ev: _record(_n, p)) + return orch, events + + +class AssignUnidentifiedTests(unittest.IsolatedAsyncioTestCase): + async def test_unknown_turn_is_stored_then_sorted_into_new_profile(self) -> None: + with tempfile.TemporaryDirectory() as d: + orch, events = _orchestrator(Path(d)) + orch.config.speaker_id.auto_enroll = False + await orch._handle_unknown_speaker(_clip(5), SR, "wer bin ich", source="Küche") + await orch._handle_unknown_speaker(_clip(5), SR, " ", source="local") # no text + await asyncio.sleep(0.2) + entries = orch.unidentified.entries() + self.assertEqual(len(entries), 1) + self.assertEqual(entries[0]["source"], "Küche") + self.assertIn(("unidentified_changed", {"added": entries[0]["id"]}), events) + + res = await orch.assign_unidentified([entries[0]["id"]], " Carla ") + self.assertEqual(res, {"name": "Carla", "added": 1, "created": True}) + self.assertTrue(orch.profiler.has_profile("Carla")) + self.assertEqual(orch.unidentified.entries(), []) + self.assertEqual(orch.profiler.identify_embedding(_fake_embed(_clip(5), SR))[0], + "Carla") + names = [e[0] for e in events] + self.assertIn("voice_profiles_changed", names) + + async def test_assign_to_existing_profile_refines_it(self) -> None: + with tempfile.TemporaryDirectory() as d: + orch, _ = _orchestrator(Path(d)) + orch.profiler.enroll("Dora", [_clip(6)], SR) + ids = [orch.unidentified.add(_clip(6), SR, text="a"), + orch.unidentified.add(_clip(6), SR, text="b")] # no stored embedding + res = await orch.assign_unidentified(ids, "Dora") + self.assertEqual(res["added"], 2) + self.assertFalse(res["created"]) + dora = next(p for p in orch.profiler.list_profiles() if p["name"] == "Dora") + self.assertEqual(dora["clip_count"], 3) + + async def test_assign_validates_input(self) -> None: + with tempfile.TemporaryDirectory() as d: + orch, _ = _orchestrator(Path(d)) + with self.assertRaises(ValueError): + await orch.assign_unidentified(["x"], " ") + with self.assertRaises(RuntimeError): + await orch.assign_unidentified(["missing"], "Eve") + orch.profiler._encoder = None + with self.assertRaises(RuntimeError): + await orch.assign_unidentified(["x"], "Eve") + + async def test_name_prompt_confirm_removes_inbox_copies(self) -> None: + with tempfile.TemporaryDirectory() as d: + orch, events = _orchestrator(Path(d)) + await orch._handle_unknown_speaker(_clip(7), SR, "hallo", source="local") + await orch._handle_unknown_speaker(_clip(7), SR, "noch mal", source="local") + self.assertEqual(len(orch.pending_enroll_clips()), 2) + self.assertEqual(sum(1 for e in events if e[0] == "speaker_ask_name"), 1) + self.assertTrue(await orch.auto_enroll_confirm("k", "Frank")) + self.assertEqual(orch.unidentified.entries(), []) + self.assertTrue(orch.profiler.has_profile("Frank")) + self.assertIn(("speaker_enrolled", {"name": "Frank", "clip_count": 2}), events) + + async def test_dismiss_keeps_inbox_and_allows_new_prompt(self) -> None: + with tempfile.TemporaryDirectory() as d: + orch, events = _orchestrator(Path(d)) + await orch._handle_unknown_speaker(_clip(8), SR, "hallo", source="local") + await asyncio.sleep(0.2) + orch.auto_enroll_dismiss() + self.assertEqual(orch.pending_enroll_clips(), []) + self.assertEqual(len(orch.unidentified.entries()), 1) + await orch._handle_unknown_speaker(_clip(8), SR, "wieder", source="local") + self.assertEqual(sum(1 for e in events if e[0] == "speaker_ask_name"), 2) + + +class LiveSampleTests(unittest.IsolatedAsyncioTestCase): + async def test_record_and_add_samples(self) -> None: + with tempfile.TemporaryDirectory() as d: + orch, events = _orchestrator(Path(d)) + orch.profiler.enroll("Gina", [_clip(9)], SR) + seen = {} + + async def fake_record(config, mic, vad, **kw): + # The wake loop must be paused and the turn lock held meanwhile. + seen["mic_free"] = orch._mic_free.is_set() + seen["locked"] = orch._turn_lock.locked() + seen["state"] = orch.state + return _clip(9) + + async def no_chime(): + return None + + orch._play_wake_chime = no_chime + with patch("javis.agent.orchestrator.record_utterance", fake_record): + sample = await orch.record_voice_sample("Gina") + self.assertEqual(seen, {"mic_free": False, "locked": True, + "state": AgentState.ENROLLING}) + self.assertTrue(orch._mic_free.is_set()) + self.assertFalse(orch._turn_lock.locked()) + self.assertEqual(orch.state, AgentState.IDLE) + self.assertGreater(sample.score, 0.99) + self.assertEqual(sample.best_name, "Gina") + self.assertAlmostEqual(sample.duration_s, 1.0) + + used = await orch.add_voice_samples("Gina", [sample, sample]) + self.assertEqual(used, 2) + gina = next(p for p in orch.profiler.list_profiles() if p["name"] == "Gina") + self.assertEqual(gina["clip_count"], 3) + self.assertIn(("voice_profiles_changed", + {"name": "Gina", "added": 2, "created": False}), events) + + async def test_record_errors_are_user_facing(self) -> None: + with tempfile.TemporaryDirectory() as d: + orch, _ = _orchestrator(Path(d)) + + async def no_speech(*a, **kw): + return None + + async def no_chime(): + return None + + orch._play_wake_chime = no_chime + with patch("javis.agent.orchestrator.record_utterance", no_speech), \ + self.assertRaisesRegex(RuntimeError, "No speech"): + await orch.record_voice_sample("Gina") + self.assertTrue(orch._mic_free.is_set()) + self.assertFalse(orch._turn_lock.locked()) + + async with orch._turn_lock: # a conversation turn is running + with self.assertRaisesRegex(RuntimeError, "busy"): + async with orch._exclusive_mic(wait_s=0.05): + pass + self.assertTrue(orch._mic_free.is_set()) + + orch.mic = None + with self.assertRaisesRegex(RuntimeError, "microphone"): + await orch.record_voice_sample("Gina") + + async def test_wake_loop_yields_mic_to_exclusive_recording(self) -> None: + from javis.agent.orchestrator import Orchestrator + + with tempfile.TemporaryDirectory() as d: + orch, _ = _orchestrator(Path(d)) + orch._stop = asyncio.Event() + fed: list[int] = [] + + class Wake: + def reset(self): + pass + + def feed(self, chunk): + fed.append(len(chunk)) + return 0.0 + + def triggered(self, score): + return False + + class Mic(_FakeMic): + async def chunks(self, stale_timeout=None): + while True: + await asyncio.sleep(0.005) + yield np.zeros(512, dtype=np.float32) + + orch.wake = Wake() + orch.mic = Mic() + orch._mic_free.clear() + task = asyncio.ensure_future(Orchestrator._wait_for_wake(orch)) + await asyncio.sleep(0.05) + self.assertEqual(fed, []) # waits while the mic is taken + orch._mic_free.set() + await asyncio.sleep(0.05) + self.assertTrue(fed) # resumed wake detection + orch._mic_free.clear() + self.assertFalse(await asyncio.wait_for(task, 1.0)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_voice_training.py b/tests/test_voice_training.py index d554d57..e061085 100644 --- a/tests/test_voice_training.py +++ b/tests/test_voice_training.py @@ -1,15 +1,33 @@ -"""Voice training pipeline: warm start/resume, piper compatibility fixes, data filters.""" +"""Voice training pipeline: warm start/resume, piper compatibility fixes, reliability, data filters.""" from __future__ import annotations +import asyncio +import contextlib +import errno +import functools +import io +import json +import math +import os +import platform +import re +import shutil +import subprocess import sys import tempfile +import threading +import time import unittest +import warnings +from datetime import timedelta from pathlib import Path from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import MagicMock, patch import numpy as np +from javis.voices.training_control import EXIT_BUSY, EXIT_DIVERGED, EXIT_STOPPED, process_alive + HIGH_ARCH = { "resblock": "1", "resblock_kernel_sizes": (3, 7, 11), "resblock_dilation_sizes": ((1, 3, 5), (1, 3, 5), (1, 3, 5)), @@ -25,6 +43,122 @@ def _voice_dir(root: Path) -> Path: return root / "v" +def _ckpt_dir(voice: Path, version: int = 0) -> Path: + return (voice / "training" / "piper_lightning" / "lightning_logs" / f"version_{version}" + / "checkpoints") + + +def _checkpoint(path: Path, *, step: int = 0, nan: bool = False) -> Path: + """A minimal Lightning checkpoint at *path* (with NaN weights if *nan*).""" + import torch + path.parent.mkdir(parents=True, exist_ok=True) + weights = torch.full((2,), float("nan")) if nan else torch.zeros(2) + torch.save({"state_dict": {"model_g.w": weights}, "global_step": step, + "hyper_parameters": HIGH_ARCH}, path) + return path + + +def _truncate(path: Path, size: int = 200) -> Path: + path.write_bytes(path.read_bytes()[:size]) + return path + + +def _exited_process() -> tuple[int, int]: + """``(pid, start time)`` of a process that has exited.""" + from javis.voices.training_control import process_start_time + proc = subprocess.Popen([sys.executable, "-c", "pass"]) + start = process_start_time(proc.pid) or 0 + proc.wait() + return proc.pid, start + + +def _captured(run): + """``(run(), stdout, stderr)``, with warnings suppressed.""" + out, err = io.StringIO(), io.StringIO() + with warnings.catch_warnings(), contextlib.redirect_stdout(out), \ + contextlib.redirect_stderr(err): + warnings.simplefilter("ignore") + result = run() + return result, out.getvalue(), err.getvalue() + + +def _toy_training(root: Path, *, epochs: int, callbacks=(), on_step=None, resumable=False): + """``(model, trainer)``: a tiny run with piper's checkpoint callbacks, saving to *root*. + + val_mel gets worse every epoch, so only epoch 0 makes the top 1; val_mos is + never logged (as when the MOS predictor can't be loaded). Four clips, one + per batch; ``model.items`` lists the clips trained. *resumable*: batched + like piper_compat's training, whose position within an epoch is saved. + """ + import lightning as L + import torch + from lightning.pytorch.callbacks import ModelCheckpoint + from torch.utils.data import DataLoader + + class Toy(L.LightningModule): + def __init__(self): + super().__init__() + self.layer = torch.nn.Linear(2, 1) + self.steps: list[tuple[int, int]] = [] # (epoch, batch) + self.items: list[int] = [] + + def training_step(self, batch, batch_idx): + self.steps.append((self.current_epoch, batch_idx)) + self.items += [int(clip) for clip in batch[:, 0]] + if on_step is not None: + on_step(self, batch_idx) + return self.layer(batch).pow(2).mean() + + def validation_step(self, batch, batch_idx): + self.log("val_mel", 1.0 + self.current_epoch) + + def configure_optimizers(self): + return torch.optim.SGD(self.parameters(), lr=0.01) + + def train_dataloader(self): + clips = torch.arange(4.0).unsqueeze(1).repeat(1, 2) # clip i: [i, i] + if not resumable: + return DataLoader(clips, batch_size=1) + from javis.voices import piper_compat as pc + return pc.ResumableDataLoader( + clips, batch_sampler=pc.LengthBucketBatchSampler([3, 1, 4, 2], 1, drop_last=False), + batches_done=functools.partial(pc._batches_done, self)) + + def val_dataloader(self): + return DataLoader(torch.ones(1, 2), batch_size=1) + + checkpoints = [ # as in piper.train.__main__ + ModelCheckpoint(dirpath=root, monitor="val_mel", mode="min", save_top_k=1, save_last=True, + filename="epoch={epoch}-val_mel={val_mel:.4f}", + auto_insert_metric_name=False), + ModelCheckpoint(dirpath=root, monitor="val_mos", mode="max", save_top_k=1, + save_last=False, filename="epoch={epoch}-val_mos={val_mos:.4f}", + auto_insert_metric_name=False), + ] + trainer = L.Trainer(accelerator="cpu", max_epochs=epochs, callbacks=[*checkpoints, *callbacks], + logger=False, num_sanity_val_steps=0, enable_progress_bar=False, + enable_model_summary=False) + return Toy(), trainer + + +@contextlib.contextmanager +def _checkpoint_patches(pc): + """piper_compat's checkpoint patches, undone afterwards: they patch Lightning itself.""" + from lightning.fabric.plugins.io import torch_io + from lightning.pytorch.callbacks import ModelCheckpoint + with contextlib.ExitStack() as stack: + for owner, name in ((torch_io, "_atomic_save"), + (torch_io.TorchCheckpointIO, "remove_checkpoint"), + (ModelCheckpoint, "on_train_epoch_end"), + (ModelCheckpoint, "on_validation_end"), + (ModelCheckpoint, "_save_topk_checkpoint")): + stack.enter_context(patch.object(owner, name, getattr(owner, name))) + stack.enter_context(patch.object(pc, "_missing_monitors", set())) + pc._safe_checkpoint_files() + pc._checkpoint_every_epoch() + yield + + class TrainingStartTests(unittest.TestCase): def setUp(self) -> None: try: @@ -34,16 +168,37 @@ class TrainingStartTests(unittest.TestCase): from javis.voices import piper_train self.pt = piper_train self.cmds: list[list[str]] = [] + # What the next training processes do: an exit code, or a callable that + # gets _run's keyword arguments and returns one. After those: crash (1). + self.outcomes: list = [] def fake_run(cmd, **kw): self.cmds.append(cmd) - return 1 # stop after stage 1 + if "export" in cmd: + Path(cmd[cmd.index("--output-file") + 1]).write_bytes(b"onnx") + if "--info-file" in cmd: # the record piper_compat's export writes + ckpt = Path(cmd[cmd.index("--checkpoint") + 1]) + Path(cmd[cmd.index("--info-file") + 1]).write_text(json.dumps({ + "checkpoint": os.path.abspath(ckpt), + "checkpoint_mtime_ns": ckpt.stat().st_mtime_ns, "epoch": 0}), + encoding="utf-8") + return 0 + outcome = self.outcomes.pop(0) if self.outcomes else 1 + return outcome(**kw) if callable(outcome) else outcome patches = [patch.object(piper_train, "_run", fake_run), - patch.object(piper_train, "MIN_FREE_BYTES", 0)] + patch.object(piper_train, "MIN_FREE_BYTES", 0), + patch.object(piper_train, "RESTART_DELAY_SECONDS", 0)] for p in patches: p.start() self.addCleanup(p.stop) + self.addCleanup(piper_train._STOP.clear) + + def _train(self, voice: Path, **kw) -> tuple[int, str, str]: + return _captured(lambda: self.pt.run_pipeline(voice, **kw)) + + def _fits(self) -> list[list[str]]: + return [cmd for cmd in self.cmds if "fit" in cmd] def _ckpt(self, path: Path, **hp) -> Path: import torch @@ -77,6 +232,227 @@ class TrainingStartTests(unittest.TestCase): cmd = self.cmds[0] self.assertEqual(cmd[cmd.index("--ckpt_path") + 1], str(last)) self.assertNotIn("--model.warmstart_ckpt", cmd) + self.assertEqual(cmd[cmd.index("--trainer.precision") + 1], "bf16-mixed") + + def test_refuses_to_resume_a_diverged_checkpoint(self) -> None: + import torch + with tempfile.TemporaryDirectory() as d: + voice = _voice_dir(Path(d)) + last = (voice / "training" / "piper_lightning" / "lightning_logs" / "version_0" + / "checkpoints" / "last.ckpt") + last.parent.mkdir(parents=True) + torch.save({"state_dict": {"model_g.w": torch.zeros(2), + "model_d.w": torch.full((2,), float("nan"))}, + "hyper_parameters": HIGH_ARCH}, last) + err = io.StringIO() + with contextlib.redirect_stderr(err), contextlib.redirect_stdout(io.StringIO()): + self.assertEqual(self.pt.run_pipeline(voice), 1) + self.assertEqual(self.cmds, []) # never started training + self.assertIn("model_d", err.getvalue()) + + def test_passes_over_damaged_checkpoints_to_the_newest_usable_one(self) -> None: + with tempfile.TemporaryDirectory() as d: + voice = _voice_dir(Path(d)) + ckpts = _ckpt_dir(voice) + usable = _checkpoint(ckpts / "epoch=1-step=100.ckpt", step=100) + _checkpoint(ckpts / "epoch=2-step=200.ckpt", step=200, nan=True) + _truncate(_checkpoint(ckpts / "last.ckpt", step=300)) # cut short by a power cut + self.outcomes = [EXIT_STOPPED] + _, _, err = self._train(voice) + cmd = self.cmds[0] + self.assertEqual(cmd[cmd.index("--ckpt_path") + 1], str(usable)) + self.assertIn("passing over", err) + self.assertIn("last.ckpt: can't be loaded", err) + self.assertIn("epoch=2-step=200.ckpt: 1 weight tensors contain NaN/inf", err) + + def test_crashed_training_restarts_from_its_newest_checkpoint(self) -> None: + with tempfile.TemporaryDirectory() as d: + voice = _voice_dir(Path(d)) + (voice / "training" / "piper_config.json").write_text( + json.dumps({"phoneme_id_map": {"a": [1]}}), encoding="utf-8") + + def crash_after_saving(step): + def training(**kw): + path = _checkpoint(_ckpt_dir(voice) / f"epoch={step}-step={step}.ckpt", + step=step) + later = time.time() + step # saved after everything before it + os.utime(path, (later, later)) + return 1 + return training + + self.outcomes = [crash_after_saving(1), crash_after_saving(2), 0] + rc, _, err = self._train(voice) + self.assertEqual(rc, 0) + fits = self._fits() + self.assertEqual(len(fits), 3) + self.assertNotIn("--ckpt_path", fits[0]) + for fit, step in zip(fits[1:], (1, 2)): + self.assertTrue(fit[fit.index("--ckpt_path") + 1].endswith( + f"epoch={step}-step={step}.ckpt")) + self.assertEqual(err.count("Training crashed (exit 1)"), 2) + export = self.cmds[-1] + self.assertTrue(export[export.index("--checkpoint") + 1].endswith("epoch=2-step=2.ckpt")) + self.assertEqual((voice / "piper_model.onnx").read_bytes(), b"onnx") + + def test_gives_up_on_crashes_that_save_nothing(self) -> None: + with tempfile.TemporaryDirectory() as d: + voice = _voice_dir(Path(d)) + self.outcomes = [7] * 5 + rc, _, err = self._train(voice) + self.assertEqual(rc, 7) + self.assertEqual(len(self.cmds), self.pt.MAX_FAILURES_WITHOUT_PROGRESS) # no export + self.assertIn("not restarting it again", err) + self.assertFalse((voice / "training" / ".training.lock").exists()) # released + + def test_diverged_or_stopped_training_is_not_restarted(self) -> None: + def stop_requested(**kw): + self.pt.request_stop() # e.g. Javis quits during a scheduled run + return 1 + + for outcome, expected in ((EXIT_DIVERGED, EXIT_DIVERGED), (EXIT_STOPPED, EXIT_STOPPED), + (stop_requested, EXIT_STOPPED)): + with self.subTest(outcome=outcome), tempfile.TemporaryDirectory() as d: + self.cmds.clear() + self.outcomes = [outcome, 0] + rc, _, _ = self._train(_voice_dir(Path(d))) + self.assertEqual(rc, expected) + self.assertEqual(len(self.cmds), 1) # neither restarted nor exported + + def test_a_stopped_run_lets_the_voice_speak_as_trained_so_far(self) -> None: + with tempfile.TemporaryDirectory() as d: + voice = _voice_dir(Path(d)) + (voice / "training" / "piper_config.json").write_text( + json.dumps({"phoneme_id_map": {"a": [1]}}), encoding="utf-8") + onnx, last = voice / "piper_model.onnx", _ckpt_dir(voice) / "last.ckpt" + + def stopped(**kw): # Ctrl+C: training saved its progress + _checkpoint(last, step=7) + return EXIT_STOPPED + + self.outcomes = [stopped] + rc, out, _ = self._train(voice) + self.assertEqual((rc, len(self.cmds)), (EXIT_STOPPED, 2)) + export = self.cmds[1] + self.assertIn("export", export) + self.assertEqual(export[export.index("--checkpoint") + 1], str(last)) + self.assertEqual(onnx.read_bytes(), b"onnx") + self.assertTrue(onnx.with_suffix(".onnx.json").is_file()) + self.assertIn("Voice updated: it speaks as trained up to epoch 0.", out) + + self.cmds.clear() # stopped again before saving anything new: nothing to export + self.outcomes = [EXIT_STOPPED] + rc, out, _ = self._train(voice) + self.assertEqual((rc, len(self.cmds)), (EXIT_STOPPED, 1)) + self.assertNotIn("Voice updated", out) + + def javis_quits(**kw): # no time for an export then + _checkpoint(last, step=8) + later = last.stat().st_mtime_ns + 10**9 + os.utime(last, ns=(later, later)) + self.pt.request_stop() + return EXIT_STOPPED + + self.cmds.clear() + self.outcomes = [javis_quits] + rc, out, _ = self._train(voice) + self.assertEqual((rc, len(self.cmds)), (EXIT_STOPPED, 1)) + self.assertFalse(self.pt._exported_from(last, onnx)) + + def test_training_keeps_the_voice_close_behind_its_progress(self) -> None: + polls = [] + + def training(rc): + def run(on_poll, **kw): + polls.append(on_poll) + return rc + return run + + with tempfile.TemporaryDirectory() as d, \ + patch.object(self.pt, "SnapshotExporter") as exporter_class: + voice = _voice_dir(Path(d)) + self.outcomes = [training(1), training(EXIT_STOPPED)] # crashed, restarted, stopped + rc, _, _ = self._train(voice, snapshot_minutes=5) + self.assertEqual(rc, EXIT_STOPPED) + exporter_class.assert_called_once_with(voice, 5) + exporter = exporter_class.return_value + self.assertEqual(polls, [exporter.poll] * 2) + exporter.stop.assert_called_once_with() # an export still running is abandoned + + def test_stop_during_the_restart_pause(self) -> None: + with tempfile.TemporaryDirectory() as d, \ + patch.object(self.pt, "_pause", return_value=False) as pause: + rc, out, _ = self._train(_voice_dir(Path(d))) + self.assertEqual(rc, EXIT_STOPPED) + pause.assert_called_once_with(0) + self.assertEqual(len(self.cmds), 1) + self.assertIn("Training stopped", out) + + def test_a_voice_is_trained_by_one_run_at_a_time(self) -> None: + from javis.voices.training_control import TrainingLock + with tempfile.TemporaryDirectory() as d: + voice = _voice_dir(Path(d)) + other_run = TrainingLock(voice) # held by a live process: this one + self.assertTrue(other_run.acquire()) + try: + rc, _, err = self._train(voice) + finally: + other_run.release() + self.assertEqual(rc, EXIT_BUSY) + self.assertEqual(self.cmds, []) + self.assertIn("being trained already", err) + + locks = [] + + def training(on_start, **kw): + on_start(os.getpid()) # stands in for the training process + locks.append(json.loads(other_run.path.read_text(encoding="utf-8"))) + return EXIT_STOPPED + + self.outcomes = [training] + self.assertEqual(self._train(voice)[0], EXIT_STOPPED) + self.assertEqual(len(locks[0]["processes"]), 2) # the launcher and its trainer + self.assertFalse(other_run.path.exists()) # released + + def test_leftovers_of_interrupted_saves_are_removed(self) -> None: + with tempfile.TemporaryDirectory() as d: + voice = _voice_dir(Path(d)) + last = _checkpoint(_ckpt_dir(voice) / "last.ckpt", step=5) + partial = last.with_name("last.ckpt.partial") + partial.write_bytes(b"cut short") + present_at_start = [] + + def killed_while_saving(**kw): + present_at_start.append(partial.exists()) + partial.write_bytes(b"cut short again") + return 1 + + def training(**kw): + present_at_start.append(partial.exists()) + return EXIT_STOPPED + + self.outcomes = [killed_while_saving, training] + self._train(voice) + self.assertEqual(present_at_start, [False, False]) # at the start and the restart + self.assertFalse(partial.exists()) + self.assertTrue(last.exists()) + + def test_checkpoint_problem_spots_divergence(self) -> None: + import torch + problem = self.pt._checkpoint_problem + healthy = {"state_dict": {"model_g.w": torch.zeros(2), "steps": torch.tensor(3)}, + "optimizer_states": [{"state": {0: {"exp_avg": torch.zeros(2)}}}], + "MixedPrecision": {"scale": 1024.0}} + self.assertIsNone(problem(healthy, "16-mixed")) + self.assertIn("model_g", problem( + {"state_dict": {"model_g.w": torch.tensor([1.0, float("inf")])}}, "32")) + self.assertIn("optimizer 1", problem( + {"state_dict": {}, "optimizer_states": [ + {"state": {}}, {"state": {0: {"exp_avg_sq": torch.tensor([float("nan")])}}}]}, + "bf16-mixed")) + collapsed = {"state_dict": {"model_g.w": torch.zeros(2)}, + "MixedPrecision": {"scale": torch.tensor(0.0)}} + self.assertIn("loss scale", problem(collapsed, "16-mixed")) + self.assertIsNone(problem(collapsed, "bf16-mixed")) # no scaler is used then def test_linux_checkpoint_with_posixpath_loads(self) -> None: import pathlib @@ -99,6 +475,905 @@ class TrainingStartTests(unittest.TestCase): self.assertEqual(str(ckpt["hyper_parameters"]["dataset"][0]), "/ssd/x") +class CheckpointPickerTests(unittest.TestCase): + """Training continues from the most advanced checkpoint that is usable.""" + + def setUp(self) -> None: + try: + import torch # noqa: F401 + except ImportError: + self.skipTest("torch not installed") + from javis.voices import piper_train + self.pick = piper_train._newest_usable_checkpoint + d = tempfile.TemporaryDirectory() + self.addCleanup(d.cleanup) + self.root = Path(d.name) + + def test_training_progress_beats_a_stale_last_checkpoint(self) -> None: + newest = _checkpoint(self.root / "version_0" / "epoch=9-step=900.ckpt", step=900) + # last.ckpt of a later run that crashed early: a newer file, but less trained + _checkpoint(self.root / "version_1" / "last.ckpt", step=500) + path, arch, problems = self.pick(self.root, "bf16-mixed") + self.assertEqual((path, problems), (newest, [])) + self.assertEqual(arch["resblock"], "1") + last = _checkpoint(self.root / "version_0" / "last.ckpt", step=900) + self.assertEqual(self.pick(self.root, "bf16-mixed")[0], last) # equal progress + + def test_unusable_checkpoints_are_passed_over_and_reported(self) -> None: + import torch + usable = _checkpoint(self.root / "a.ckpt", step=1) + _checkpoint(self.root / "b.ckpt", step=3, nan=True) + _truncate(_checkpoint(self.root / "c.ckpt", step=4)) + (self.root / "d.ckpt").write_bytes(b"") + torch.save({"weights": torch.zeros(1)}, self.root / "e.ckpt") # not from Lightning + path, _, problems = self.pick(self.root, "32") + self.assertEqual(path, usable) + self.assertEqual(len(problems), 4) + self.assertIn("b.ckpt: 1 weight tensors contain NaN/inf", problems[-1]) + usable.unlink() # not left mapped (Windows can't delete a mapped file) + path, _, problems = self.pick(self.root, "32") + self.assertIsNone(path) + self.assertEqual(len(problems), 4) + + +class TrainingControlTests(unittest.TestCase): + """training_control: one training run per voice, and its process helpers.""" + + def setUp(self) -> None: + from javis.voices import training_control + self.tc = training_control + d = tempfile.TemporaryDirectory() + self.addCleanup(d.cleanup) + self.tmp = Path(d.name) + self.voice = self.tmp / "voice" + self.lock_file = training_control.lock_path(self.voice) + + def _write_lock(self, content, *, age: float = 0.0) -> None: + self.lock_file.parent.mkdir(parents=True, exist_ok=True) + text = content if isinstance(content, str) else json.dumps(content) + self.lock_file.write_text(text, encoding="utf-8") + if age: + then = time.time() - age + os.utime(self.lock_file, (then, then)) + + def test_one_run_per_voice(self) -> None: + first, second = self.tc.TrainingLock(self.voice), self.tc.TrainingLock(self.voice) + self.assertTrue(first.acquire()) + self.assertTrue(self.tc.training_in_progress(self.voice)) + self.assertFalse(second.acquire()) + second.release() # not its lock + self.assertTrue(self.lock_file.exists()) + first.release() + self.assertFalse(self.lock_file.exists()) + self.assertFalse(self.tc.training_in_progress(self.voice)) + + def test_the_lock_of_a_run_that_has_exited_is_replaced(self) -> None: + pid, start = _exited_process() + self._write_lock({"host": platform.node(), "processes": [{"pid": pid, "start": start}]}) + self.assertFalse(self.tc.training_in_progress(self.voice)) + lock = self.tc.TrainingLock(self.voice) + self.assertTrue(lock.acquire()) + info = json.loads(self.lock_file.read_text(encoding="utf-8")) + self.assertEqual(info["processes"], [self.tc.this_process()]) + lock.release() + + def test_the_lock_lasts_while_the_training_process_saves(self) -> None: + pid, start = _exited_process() # the launcher was killed; its trainer still runs + self._write_lock({"host": platform.node(), "processes": [ + {"pid": pid, "start": start}, self.tc.this_process()]}) + self.assertTrue(self.tc.training_in_progress(self.voice)) + self.assertFalse(self.tc.TrainingLock(self.voice).acquire()) + + def test_set_trainer_and_release(self) -> None: + lock = self.tc.TrainingLock(self.voice) + self.assertTrue(lock.acquire()) + lock.set_trainer(os.getpid()) # stands in for the training process + info = json.loads(self.lock_file.read_text(encoding="utf-8")) + self.assertEqual([p["pid"] for p in info["processes"]], [os.getpid()] * 2) + lock.set_trainer(None) + info = json.loads(self.lock_file.read_text(encoding="utf-8")) + self.assertEqual(len(info["processes"]), 1) + # Taken over by another run meanwhile (this one's lock looked stale): left alone. + self._write_lock({"host": platform.node(), "processes": [{"pid": 4, "start": 1}]}) + lock.release() + self.assertTrue(self.lock_file.exists()) + + def test_locks_without_running_processes(self) -> None: + now = time.time() + cases = [ # (name, lock file content, its age in seconds, held?) + ("older Javis, recent", {"started_epoch": now - 3600}, 0, True), + ("older Javis, two days old", {"started_epoch": now - 49 * 3600}, 0, False), + ("another computer", {"host": "elsewhere", "started_epoch": now, + "processes": [{"pid": 4, "start": 1}]}, 0, True), + ("being written", "{", 0, True), + ("garbled", "{", 120, False), + ] + for name, content, age, held in cases: + with self.subTest(name): + self._write_lock(content, age=age) + self.assertEqual(self.tc.training_in_progress(self.voice), held) + self.assertEqual(self.tc.TrainingLock(self.voice).acquire(), not held) + self.lock_file.unlink() + + def test_process_alive(self) -> None: + me = self.tc.this_process() + self.assertTrue(self.tc.process_alive(me["pid"], me["start"])) + self.assertFalse(self.tc.process_alive(*_exited_process())) + if sys.platform == "win32": # IDs are soon reused there: the start time tells them apart + self.assertFalse(self.tc.process_alive(me["pid"], me["start"] + 1)) + + def test_launcher_from_environment(self) -> None: + for value, expected in (("12:345", (12, 345)), ("", None), ("12", None), ("a:b", None)): + with self.subTest(value), patch.dict(os.environ, {self.tc.LAUNCHER_ENV: value}): + self.assertEqual(self.tc.launcher_from_environment(), expected) + + def test_replace_file_waits_for_the_target_to_be_closed(self) -> None: + src, dst = self.tmp / "new", self.tmp / "current" + src.write_text("new") + dst.write_text("old") + real_replace, calls = os.replace, [] + + def busy_twice(a, b): + calls.append(a) + if len(calls) <= 2: + raise PermissionError(errno.EACCES, "in use") + real_replace(a, b) + + with patch.object(self.tc.os, "replace", busy_twice): + self.tc.replace_file(src, dst) + self.assertEqual((dst.read_text(), len(calls)), ("new", 3)) + src.write_text("newer") + with patch.object(self.tc.os, "replace", side_effect=PermissionError(errno.EACCES, "in use")), \ + self.assertRaises(PermissionError): + self.tc.replace_file(src, dst, seconds=0) + with self.assertRaises(FileNotFoundError): # other errors at once + self.tc.replace_file(self.tmp / "missing", dst) + + def test_windows_helpers_are_safe_anywhere(self) -> None: + with self.tc.keep_system_awake(), self.tc.console_selection_cannot_pause(): + pass + + def test_python_command_starts_this_environment(self) -> None: + exe, env = self.tc.python_command() + prefix = subprocess.run( + [exe, "-c", "import sys, javis, lightning; print(sys.prefix)"], + env={**os.environ, **env}, cwd=tempfile.gettempdir(), capture_output=True, + text=True, timeout=120, check=True).stdout.strip() + self.assertEqual(os.path.normcase(prefix), os.path.normcase(sys.prefix)) + if sys.platform == "win32": + # torch's Intel Fortran runtime must not abort it when its console window closes. + self.assertEqual(env.get("FOR_DISABLE_CONSOLE_CTRL_HANDLER"), "1") + if sys.platform == "win32" and sys.prefix != sys.base_prefix: + # Not the virtual environment's launcher: closing the console would end it at once. + self.assertEqual(os.path.normcase(exe), os.path.normcase(sys._base_executable)) + self.assertEqual(env.get("__PYVENV_LAUNCHER__"), sys.executable) + with patch.object(sys, "platform", "linux"): + self.assertEqual(self.tc.python_command(), (sys.executable, {})) + + def test_closing_console_holds_the_process_until_the_progress_is_saved(self) -> None: + console = self.tc.ConsoleClose(timeout=30) + self.assertEqual((console.handle(0), console.handle(1)), (False, False)) # Ctrl+C/Break + self.assertFalse(console.closing.is_set()) + handled = [] + handler = threading.Thread(target=lambda: handled.append(console.handle(5))) + handler.start() + self.assertTrue(console.closing.wait(10)) + self.assertEqual(console.reason, "the user is logging off") + handler.join(0.3) + self.assertTrue(handler.is_alive()) # Windows waits for the handler to return + console.done() + handler.join(10) + self.assertEqual(handled, [True]) + unsaved = self.tc.ConsoleClose(timeout=0.05) # nothing saves: let Windows go anyway + self.assertTrue(unsaved.handle(6)) + self.assertEqual(unsaved.reason, "the system is shutting down") + + @unittest.skipUnless(sys.platform == "win32", "a Windows console handler") + def test_console_close_handler_is_installed_once(self) -> None: + register, shutdown_order = MagicMock(return_value=0), MagicMock(return_value=1) + console = self.tc.ConsoleClose(timeout=0.05) + with patch.object(self.tc, "_kernel32", lambda: SimpleNamespace( + SetConsoleCtrlHandler=register, SetProcessShutdownParameters=shutdown_order)), \ + patch.object(console, "_watch_session_end") as watch_session_end: + console.install() # refused: tried again next time + register.return_value = 1 + console.install() + console.install() + self.assertEqual(register.call_count, 2) + watch_session_end.assert_called_once_with() + # Ends before the rest of the session: its data-loader workers are still there. + level = shutdown_order.call_args.args[0] + self.assertGreater(level, 0x280) # the default + self.assertLessEqual(level, 0x3FF) # the applications' range + handler, add = register.call_args.args + self.assertTrue(add) + self.assertEqual(handler(0), 0) # Ctrl+C: left to Python + self.assertFalse(console.closing.is_set()) + self.assertEqual(handler(2), 1) + self.assertEqual(console.reason, "the console window was closed") + with patch.object(sys, "platform", "linux"): + other = self.tc.ConsoleClose() + other.install() + self.assertIsNone(other._handler) + + @unittest.skipUnless(sys.platform == "win32", "a Windows window") + def test_logging_off_or_shutting_down_holds_the_process_until_the_progress_is_saved( + self) -> None: + import ctypes + from ctypes import wintypes + + user32 = ctypes.WinDLL("user32", use_last_error=True) + user32.SendMessageTimeoutW.argtypes = ( + wintypes.HWND, wintypes.UINT, wintypes.WPARAM, wintypes.LPARAM, wintypes.UINT, + wintypes.UINT, ctypes.POINTER(ctypes.c_size_t)) + user32.SendMessageTimeoutW.restype = wintypes.LPARAM + user32.PostMessageW.argtypes = (wintypes.HWND, wintypes.UINT, wintypes.WPARAM, + wintypes.LPARAM) + user32.GetAncestor.argtypes = (wintypes.HWND, wintypes.UINT) + user32.GetAncestor.restype = wintypes.HWND + user32.GetDesktopWindow.restype = wintypes.HWND + user32.IsWindowVisible.argtypes = (wintypes.HWND,) + query_end, end, logoff = 0x0011, 0x0016, 0x80000000 + + def session_window(console): + threads = set(threading.enumerate()) + console._watch_session_end() + self.assertTrue(console._window) + [thread] = set(threading.enumerate()) - threads + self.addCleanup(thread.join, 10) + self.addCleanup(user32.PostMessageW, console._window, 0x0010, 0, 0) # WM_CLOSE + # Top-level, as the session's messages need, but never shown. + self.assertEqual(user32.GetAncestor(console._window, 1), user32.GetDesktopWindow()) + self.assertFalse(user32.IsWindowVisible(console._window)) + return thread + + def send(console, message, wparam=0, lparam=0): + """As Windows sends it: from another thread, waiting for the answer.""" + result = ctypes.c_size_t() + if not user32.SendMessageTimeoutW(console._window, message, wparam, lparam, 0, + 20_000, ctypes.byref(result)): + raise ctypes.WinError(ctypes.get_last_error()) + return result.value + + console = self.tc.ConsoleClose(timeout=30) + thread = session_window(console) + self.assertTrue(thread.daemon) # never keeps the process + self.assertEqual(send(console, query_end, 0, logoff), 1) # may end + send(console, end, 0, logoff) # cancelled after all + self.assertFalse(console.closing.is_set()) + answers = [] + windows = threading.Thread(target=lambda: answers.append(send(console, end, 1, logoff))) + windows.start() + self.assertTrue(console.closing.wait(10)) + self.assertEqual(console.reason, "the user is logging off") + windows.join(0.3) + self.assertTrue(windows.is_alive()) # Windows waits for the window's answer + console.done() + windows.join(10) + self.assertEqual(answers, [0]) + + unsaved = self.tc.ConsoleClose(timeout=0.05) # nothing saves: let Windows go anyway + session_window(unsaved) + self.assertEqual(send(unsaved, end, 1, 0), 0) + self.assertEqual(unsaved.reason, "the system is shutting down") + + +class TrainingProcessStopTests(unittest.TestCase): + """piper_train._run asks training to save and stop; it kills it only if it must.""" + + # Stand-ins for the training process, run as ``python -c SCRIPT ``: + # once started, they write the pid of the interpreter itself (with a virtual + # environment, the process _run starts is a launcher that runs it). TRAINING + # checks what _run tells it, then stops like StopOnRequest once the stop file exists. + TRAINING = ( + "import os, sys, time\n" + "stop = os.environ['JAVIS_PIPER_STOP_FILE']\n" + "if os.path.exists(stop) or os.environ['JAVIS_PIPER_LAUNCHER'] != sys.argv[1]:\n" + " sys.exit(8)\n" + "open(sys.argv[2], 'w').write(str(os.getpid()))\n" + "deadline = time.monotonic() + 60\n" + "while not os.path.exists(stop) and time.monotonic() < deadline:\n" + " time.sleep(0.05)\n" + "sys.exit(4 if os.path.exists(stop) else 9)\n" + ) + IGNORES_STOP = "import os, sys, time; open(sys.argv[2], 'w').write(str(os.getpid())); time.sleep(60)" + + def setUp(self) -> None: + from javis.voices import piper_train + self.pt = piper_train + piper_train._STOP.clear() + self.addCleanup(piper_train._STOP.clear) + d = tempfile.TemporaryDirectory() + self.addCleanup(d.cleanup) + self.stop_file = Path(d.name) / ".stop_training" + self.pid_file = Path(d.name) / "pid" + + def _child_pid(self) -> int: + """The pid of the stand-in's interpreter, once it has started.""" + deadline = time.monotonic() + 30 + while time.monotonic() < deadline: + with contextlib.suppress(OSError, ValueError): + return int(self.pid_file.read_text(encoding="utf-8")) + time.sleep(0.05) + raise AssertionError("the training process did not start") + + def _assert_ended(self, pid: int) -> None: + deadline = time.monotonic() + 10 + while process_alive(pid) and time.monotonic() < deadline: + time.sleep(0.05) + self.assertFalse(process_alive(pid)) + + def _run(self, child: str, *, ctrl_c: int = 0, **kw): + """``(rc, child's exit code, pids passed to on_start, seconds)`` of running *child*. + + Ctrl+C is pressed *ctrl_c* times, at _run's polls once the child has started. + *kw* goes to _run; what it printed is in ``self.out`` and ``self.err``. + """ + from javis.voices.training_control import this_process + procs, pids, test = [], [], self + + class Popen(subprocess.Popen): + waits = 0 + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + procs.append(self) + + def wait(self, timeout=None): + if timeout == 1.0: # _run's polling + Popen.waits += 1 + if 2 <= Popen.waits < 2 + ctrl_c: + test._child_pid() + raise KeyboardInterrupt + return super().wait(timeout) + + me = this_process() + cmd = [sys.executable, "-c", child, f"{me['pid']}:{me['start']}", str(self.pid_file)] + started = time.monotonic() + with patch.object(self.pt.subprocess, "Popen", Popen): + rc, self.out, self.err = _captured(lambda: self.pt._run( + cmd, desc="test", stop_file=self.stop_file, on_start=pids.append, **kw)) + return rc, procs[0].returncode if procs else None, pids, time.monotonic() - started + + def _stop_after_start(self, seconds: float) -> None: + """request_stop() *seconds* after the child has started.""" + def stop(): + self._child_pid() + time.sleep(seconds) + self.pt.request_stop() + + thread = threading.Thread(target=stop, daemon=True) + thread.start() + self.addCleanup(thread.join, 60) + + def test_stop_request_lets_training_save_first(self) -> None: + self.stop_file.touch() # left over from an earlier run: removed before the start + self._stop_after_start(0.5) + rc, child_rc, pids, _ = self._run(self.TRAINING) + self.assertEqual((rc, child_rc), (EXIT_STOPPED, EXIT_STOPPED)) + self.assertEqual(pids, [pids[0], None]) + self.assertFalse(self.stop_file.exists()) + + def test_ctrl_c_stops_training_after_it_saved(self) -> None: + rc, child_rc, _, _ = self._run(self.TRAINING, ctrl_c=1) + self.assertEqual((rc, child_rc), (EXIT_STOPPED, EXIT_STOPPED)) + + def test_second_ctrl_c_stops_training_at_once(self) -> None: + rc, _, pids, seconds = self._run(self.IGNORES_STOP, ctrl_c=2) + self.assertEqual(rc, EXIT_STOPPED) + self.assertLess(seconds, 30) + self.assertFalse(process_alive(pids[0])) + self._assert_ended(self._child_pid()) # the interpreter too, not just its launcher + + def test_training_that_does_not_stop_is_killed_after_the_grace_period(self) -> None: + self._stop_after_start(0.5) + with patch.object(self.pt, "STOP_GRACE_SECONDS", 1): + rc, _, pids, seconds = self._run(self.IGNORES_STOP) + self.assertEqual(rc, EXIT_STOPPED) + self.assertLess(seconds, 30) + self.assertFalse(process_alive(pids[0])) + self._assert_ended(self._child_pid()) + + def test_nothing_starts_after_a_stop_request(self) -> None: + self.pt.request_stop() + self.assertEqual(self._run(self.TRAINING)[:3], (EXIT_STOPPED, None, [])) + + def test_on_poll_runs_while_training_runs_until_it_fails(self) -> None: + polls = [] + + def on_poll(): # e.g. SnapshotExporter.poll + polls.append(time.monotonic()) + if len(polls) == 2: + stop = threading.Timer(1.5, self.pt.request_stop) + stop.start() + self.addCleanup(stop.cancel) + raise OSError("disk full") + + rc, child_rc, _, _ = self._run(self.TRAINING, on_poll=on_poll) + self.assertEqual((rc, child_rc), (EXIT_STOPPED, EXIT_STOPPED)) # training went on + self.assertEqual(len(polls), 2) + self.assertIn("WARNING: OSError('disk full')", self.err) + + +class OnnxExportTests(unittest.TestCase): + """The voice model is replaced only by a complete export.""" + + def setUp(self) -> None: + from javis.voices import piper_train + self.pt = piper_train + d = tempfile.TemporaryDirectory() + self.addCleanup(d.cleanup) + self.voice = _voice_dir(Path(d.name)) + self.model = self.voice / "piper_model.onnx" + self.model_config = self.voice / "piper_model.onnx.json" + self.model.write_bytes(b"old model") + self.model_config.write_text("old config", encoding="utf-8") + self.config = self.voice / "training" / "piper_config.json" + self.config.write_text(json.dumps({"phoneme_id_map": {"a": [1]}}), encoding="utf-8") + + def _export(self, writes: bytes | None = b"new model", rc: int = 0) -> tuple[int, str]: + def fake_run(cmd, **kw): + if writes is not None: + Path(cmd[cmd.index("--output-file") + 1]).write_bytes(writes) + return rc + + with patch.object(self.pt, "_run", fake_run): + result, _, err = _captured(lambda: self.pt._export_onnx( + Path("last.ckpt"), self.model, self.voice, desc="test")) + return result, err + + def _leftovers(self) -> list[str]: + return sorted(p.name for p in self.voice.iterdir() if "partial" in p.name) + + def test_complete_export_replaces_the_model(self) -> None: + self.assertEqual(self._export()[0], 0) + self.assertEqual(self.model.read_bytes(), b"new model") + self.assertEqual(json.loads(self.model_config.read_text(encoding="utf-8")), + {"phoneme_id_map": {"a": [1]}}) + self.assertEqual(self._leftovers(), []) + + def test_failed_export_leaves_the_previous_model(self) -> None: + for writes, rc in ((b"half a mod", 1), (b"", 0), (None, 0), (b"stopped", EXIT_STOPPED)): + with self.subTest(writes=writes, rc=rc): + self.assertNotEqual(self._export(writes, rc)[0], 0) + self.assertEqual(self.model.read_bytes(), b"old model") + self.assertEqual(self.model_config.read_text(encoding="utf-8"), "old config") + self.assertEqual(self._leftovers(), []) + + def test_damaged_voice_config_is_rebuilt(self) -> None: + self.config.write_text('{"phoneme_id_map": {"a": [', encoding="utf-8") # cut short + + def rebuild(lightning_dir, dest): + dest.write_text(json.dumps({"phoneme_id_map": {"b": [2]}}), encoding="utf-8") + return True + + with patch.object(self.pt, "_rebuild_voice_config", side_effect=rebuild): + self.assertEqual(self._export()[0], 0) + self.assertIn('"b"', self.model_config.read_text(encoding="utf-8")) + with patch.object(self.pt, "_rebuild_voice_config", return_value=False): + rc, err = self._export(b"newer model") + self.assertEqual(rc, 1) + self.assertIn("could not be rebuilt", err) + self.assertEqual(self.model.read_bytes(), b"new model") + self.assertEqual(self._leftovers(), []) + + +class VoiceFollowsTrainingTests(unittest.TestCase): + """While training runs, the voice is re-exported from its progress in the background.""" + + def setUp(self) -> None: + from javis.voices import piper_train + self.pt = piper_train + d = tempfile.TemporaryDirectory() + self.addCleanup(d.cleanup) + self.voice = _voice_dir(Path(d.name)) + (self.voice / "training" / "piper_config.json").write_text( + json.dumps({"phoneme_id_map": {"a": [1]}}), encoding="utf-8") + self.onnx = self.voice / "piper_model.onnx" + self.last = _ckpt_dir(self.voice) / "last.ckpt" + self.last.parent.mkdir(parents=True) + self.last.write_bytes(b"progress") + self.exported: list[Path] = [] + self.rc, self.seconds = 0, 0 # how the next exports end, and how long they take + self.now = 1000.0 + for p in (patch.object(piper_train, "_export_command", self._export_command), + patch.object(piper_train, "time", SimpleNamespace(monotonic=lambda: self.now))): + p.start() + self.addCleanup(p.stop) + + def _export_command(self, ckpt: Path, files) -> tuple[list[str], dict]: + """Stands in for piper's export: writes what it would, takes self.seconds, exits self.rc.""" + self.exported.append(ckpt) + info = {"checkpoint": os.path.abspath(ckpt), "checkpoint_mtime_ns": ckpt.stat().st_mtime_ns, + "epoch": 2, "global_step": 21, "batches_done": 5, "batches_per_epoch": 8} + script = "\n".join([ + "import pathlib, sys, time", + f"pathlib.Path({str(files.partial_onnx)!r}).write_bytes(b'onnx')", + f"pathlib.Path({str(files.partial_info)!r}).write_text({json.dumps(info)!r})", + "print('exported', flush=True)", + f"time.sleep({self.seconds})", + f"sys.exit({self.rc})"]) + python, env = self.pt.python_command() + return [python, "-c", script], {**os.environ, **env} + + def _poll(self, exporter) -> tuple[str, str]: + """Poll *exporter* as training does, until the export it starts (if any) has ended.""" + def polls(): + exporter.poll() + deadline = time.monotonic() + 60 + while exporter._proc is not None and time.monotonic() < deadline: + time.sleep(0.05) + exporter.poll() + _, out, err = _captured(polls) + return out, err + + def _save_progress(self) -> None: + """Training saves its progress anew.""" + later = self.last.stat().st_mtime_ns + 10**9 + self.last.write_bytes(b"more progress") + os.utime(self.last, ns=(later, later)) + + def _leftovers(self) -> list[str]: + return sorted(p.name for p in self.voice.iterdir() if ".partial-" in p.name) + + def test_the_voice_is_updated_in_the_background(self) -> None: + exporter = self.pt.SnapshotExporter(self.voice, 15) + out, _ = self._poll(exporter) # the first export starts right away + self.assertEqual(self.exported, [self.last]) + self.assertIn("Voice updated: it speaks as trained up to epoch 2, batch 5/8.", out) + self.assertEqual(self.onnx.read_bytes(), b"onnx") + self.assertEqual(json.loads(self.onnx.with_suffix(".onnx.json").read_text(encoding="utf-8")), + {"phoneme_id_map": {"a": [1]}}) + self.assertEqual(self.pt.export_info(self.onnx)["batches_done"], 5) + self.assertEqual(self._leftovers(), []) + self.assertIn("exported", exporter.log_path.read_text(encoding="utf-8")) + + self._save_progress() + self.now += 15 * 60 - 1 # the next one, once the interval has passed + self._poll(exporter) + self.assertEqual(len(self.exported), 1) + self.now += 1 + self._poll(exporter) + self.assertEqual(len(self.exported), 2) + + self.now += 15 * 60 # due, but training saved nothing new meanwhile + out, _ = self._poll(exporter) + self.assertEqual((len(self.exported), out), (2, "")) + self._save_progress() + self.now += 59 # it looks for new progress every minute until then + self._poll(exporter) + self.assertEqual(len(self.exported), 2) + self.now += 1 + self._poll(exporter) + self.assertEqual(self.exported, [self.last] * 3) + + def test_a_failed_export_leaves_the_voice_as_it_was(self) -> None: + self.onnx.write_bytes(b"old") + self.rc = 3 + exporter = self.pt.SnapshotExporter(self.voice, 15) + out, err = self._poll(exporter) + self.assertEqual(self.onnx.read_bytes(), b"old") + self.assertNotIn("Voice updated", out) + self.assertIn("WARNING: could not update the voice with the training progress " + f"(export exit 3); see {exporter.log_path}", err) + self.assertEqual(self._leftovers(), []) + self.now += 15 * 60 # that progress isn't tried again... + self._poll(exporter) + self.assertEqual(len(self.exported), 1) + self.rc = 0 + self._save_progress() # ...the next is + self.now += 60 + self._poll(exporter) + self.assertEqual(len(self.exported), 2) + self.assertEqual(self.onnx.read_bytes(), b"onnx") + + def test_stopping_training_abandons_a_running_export(self) -> None: + self.seconds = 60 + exporter = self.pt.SnapshotExporter(self.voice, 15) + _captured(exporter.poll) + proc, deadline = exporter._proc, time.monotonic() + 60 + while "exported" not in exporter.log_path.read_text(encoding="utf-8") \ + and time.monotonic() < deadline: + time.sleep(0.05) + self.assertEqual(len(self._leftovers()), 2) # its model and record so far + started = time.monotonic() + exporter.stop() + self.assertLess(time.monotonic() - started, 30) + self.assertIsNotNone(proc.poll()) + self.assertEqual(self._leftovers(), []) + self.assertFalse(self.onnx.exists()) + exporter.stop() # nothing left to stop + + def test_no_snapshots_when_turned_off(self) -> None: + for minutes in (0, -5, float("nan"), float("inf")): + with self.subTest(minutes=minutes): + exporter = self.pt.SnapshotExporter(self.voice, minutes) + self.now += 10**6 + self._poll(exporter) + self.assertEqual(self.exported, []) + + def test_what_the_voice_was_exported_from(self) -> None: + for info, text in ( + (None, "an unknown point of training"), + ({"epoch": "3"}, "an unknown point of training"), + ({"epoch": 3}, "epoch 3"), + ({"epoch": 3, "batches_done": 700, "batches_per_epoch": 1456}, + "epoch 3, batch 700/1456"), + ({"epoch": 3, "batches_done": 1456, "batches_per_epoch": 1456}, + "the end of epoch 3"), + ({"epoch": 3, "batches_done": None, "batches_per_epoch": None}, "epoch 3")): + with self.subTest(info=info): + self.assertEqual(self.pt.describe_export(info), text) + + record = self.pt.export_info_path(self.onnx) + self.assertEqual(record.name, "piper_model.info.json") + self.assertIsNone(self.pt.export_info(self.onnx)) + for content in ("{", "[1]"): # cut short; not a record + record.write_text(content, encoding="utf-8") + self.assertIsNone(self.pt.export_info(self.onnx)) + + mtime = self.last.stat().st_mtime_ns + exported = {"checkpoint": str(self.last), "checkpoint_mtime_ns": mtime} + for info, current in ((exported, False), # the model is gone + (exported, True), + ({**exported, "checkpoint_mtime_ns": mtime - 1}, False), + ({**exported, "checkpoint": str(self.last.with_name("x.ckpt"))}, + False), + ({**exported, "checkpoint": None}, False), + ({"checkpoint": str(self.last)}, False)): + with self.subTest(info=info, model=self.onnx.exists()): + record.write_text(json.dumps(info), encoding="utf-8") + self.assertEqual(self.pt._exported_from(self.last, self.onnx), current) + self.onnx.write_bytes(b"onnx") + + def test_the_newest_saved_progress_is_exported(self) -> None: + lightning = self.voice / "training" / "piper_lightning" + newer = _ckpt_dir(self.voice, 1) / "last-v1.ckpt" + newer.parent.mkdir(parents=True) + others = [newer.with_name("epoch=3-step=9.ckpt"), newer.with_name("last.ckpt.partial")] + base = self.last.stat().st_mtime_ns + for i, path in enumerate([newer, *others], 1): # a top-k save; one being written + path.write_bytes(b"x") + os.utime(path, ns=(base + i * 10**9,) * 2) + self.assertEqual(self.pt._newest_progress(lightning), newer) + self.assertIsNone(self.pt._newest_progress(self.voice / "missing")) + + def test_leftovers_of_exports_whose_launcher_died_are_removed(self) -> None: + dead, live = 4001, 4002 + ours = self.pt._ExportFiles(self.onnx) + names = {ours.partial_onnx.name: True, + f"piper_model.partial-{dead}-1.onnx": False, + f"piper_model.onnx.partial-{dead}-1.json": False, + f"piper_model.info.partial-{dead}-1.json": False, + f"piper_model.partial-{live}-1.onnx": True, + "piper_model.onnx": True, "notes.partial.txt": True} + for name in names: + (self.voice / name).write_bytes(b"x") + with patch.object(self.pt, "process_alive", lambda pid, *_: pid == live): + self.pt._remove_stale_exports(self.voice) + self.assertEqual(sorted(p.name for p in self.voice.iterdir() if p.is_file()), + sorted(name for name, kept in names.items() if kept)) + + +class ScheduledTrainingTests(unittest.IsolatedAsyncioTestCase): + """The overnight scheduler reports how a run ended, and lets it save when Javis quits.""" + + async def asyncSetUp(self) -> None: + from javis.core.config import TrainingConfig + from javis.voices import piper_train, training_scheduler + self.ts, self.pt = training_scheduler, piper_train + d = tempfile.TemporaryDirectory() + self.addCleanup(d.cleanup) + self.voice = _voice_dir(Path(d.name)) + self.config = SimpleNamespace(training=TrainingConfig()) + self.events: list[tuple[str, dict]] = [] + self.bus = SimpleNamespace(publish=self._publish) + for p in (patch.object(training_scheduler, "window_end", + lambda now, *_: now + timedelta(hours=8)), + patch.object(training_scheduler, "training_in_progress", return_value=False)): + p.start() + self.addCleanup(p.stop) + self.addCleanup(piper_train._STOP.clear) + + async def _publish(self, event: str, payload: dict) -> None: + self.events.append((event, payload)) + + async def test_reports_how_the_run_ended(self) -> None: + for rc, ok, status in ((0, True, "completed"), + (EXIT_STOPPED, True, "stopped (progress saved)"), + (EXIT_BUSY, True, "being trained already"), + (1, False, "FAILED (exit 1)")): + with self.subTest(rc=rc), patch.object(self.pt, "run_pipeline", return_value=rc) as run: + self.events.clear() + await self.ts._run_once(self.config, self.voice, self.bus) + run.assert_called_once() + self.assertEqual([event for event, _ in self.events], ["task_start", "task_done"]) + self.assertEqual(self.events[-1][1]["ok"], ok) + self.assertIn(status, self.events[-1][1]["msg"]) + + async def test_quitting_javis_lets_training_save_and_stop(self) -> None: + started, finished = threading.Event(), threading.Event() + + def run_pipeline(*args, **kwargs): # stops when asked to, like the real one + started.set() + self.pt._STOP.wait(30) + finished.set() + return EXIT_STOPPED + + with patch.object(self.pt, "run_pipeline", run_pipeline): + run = asyncio.create_task(self.ts._run_once(self.config, self.voice, self.bus)) + self.assertTrue(await asyncio.to_thread(started.wait, 30)) + run.cancel() + with self.assertRaises(asyncio.CancelledError): + await run + self.assertTrue(self.pt._STOP.is_set()) + self.assertTrue(finished.is_set()) # waited for it to finish + + +class PiperVoiceReloadTests(unittest.TestCase): + """Javis speaks with the active voice's newest model, which training keeps updating.""" + + def setUp(self) -> None: + try: + import piper.config # noqa: F401 + except ImportError: + self.skipTest("piper-tts not installed") + from javis.core.config import Config, PathsConfig + from javis.tts import piper_backend + d = tempfile.TemporaryDirectory() + self.addCleanup(d.cleanup) + self.config = Config(paths=PathsConfig(base_dir=Path(d.name))) + self.config.tts.default_voice = "Rufus" + self.saves = 0 + self.attempts: list[bytes] = [] # the models loading was tried with + self.loaded: list[bytes] = [] # ...and those that loaded + self.broken: set[bytes] = set() + self.loading = threading.Event() # cleared: loading takes until it is set + self.loading.set() + patcher = patch.object(piper_backend.PiperBackend, "_load_voice", + staticmethod(self._load_voice)) + patcher.start() + self.addCleanup(patcher.stop) + self.backend = piper_backend.PiperBackend(self.config) + self.backend.CHECK_SECONDS = 0 + self.addCleanup(self._reloaded) + self.addCleanup(self.loading.set) # (cleanups run last first) + + def _export(self, voice: str, model: bytes) -> None: + """Training (or an export) writes *voice*'s model.""" + path = self.config.paths.voices_dir / voice / "piper_model.onnx" + path.parent.mkdir(parents=True, exist_ok=True) + self.saves += 1 + mtime = (1_700_000_000 + self.saves) * 10**9 + for file, content in ((Path(str(path) + ".json"), b"{}"), (path, model)): + file.write_bytes(content) + os.utime(file, ns=(mtime, mtime)) + + def _load_voice(self, model_path: Path): + self.assertTrue(self.loading.wait(10)) + model = Path(model_path).read_bytes() + self.attempts.append(model) + if model in self.broken: + raise RuntimeError(f"cannot load {model!r}") + self.loaded.append(model) + number = len(self.loaded) + + def synthesize(text, syn_config=None): + yield SimpleNamespace(audio_float_array=np.full(2, number, dtype=np.float32)) + + return SimpleNamespace(config=SimpleNamespace(sample_rate=16000 + number), + synthesize=synthesize) + + def _speaking(self) -> bytes: + """Synthesize something, as Javis does; returns the model that spoke.""" + audio, sample_rate = self.backend._synthesize_sync("Hallo") + number = int(audio[0]) + self.assertEqual(sample_rate, 16000 + number) + return self.loaded[number - 1] + + def _reloaded(self) -> None: + """Wait for the model being loaded in the background, if one is.""" + self.assertTrue(self.backend._reloading.acquire(timeout=10)) + self.backend._reloading.release() + + def test_the_model_training_updates_takes_over(self) -> None: + self._export("Rufus", b"epoch 1") + self.backend.load() + self.assertEqual(self._speaking(), b"epoch 1") + self.loading.clear() # loading the new model takes a while... + self._export("Rufus", b"epoch 2") + self.assertEqual(self._speaking(), b"epoch 1") # ...the old one speaks meanwhile + self.assertEqual(self._speaking(), b"epoch 1") + self.loading.set() + self._reloaded() + self.assertEqual(self._speaking(), b"epoch 2") + self.assertEqual(self.attempts, [b"epoch 1", b"epoch 2"]) # loaded once each + + def test_a_model_that_fails_to_load_leaves_the_voice_as_it_was(self) -> None: + self._export("Rufus", b"epoch 1") + self.backend.load() + self.broken.add(b"damaged") + self._export("Rufus", b"damaged") + self._speaking() + self._reloaded() + self.assertEqual(self._speaking(), b"epoch 1") + self._reloaded() + self.assertEqual(self.attempts, [b"epoch 1", b"damaged"]) # not tried again... + self._export("Rufus", b"epoch 2") + self._speaking() + self._reloaded() + self.assertEqual(self._speaking(), b"epoch 2") # ...but the next model is + + def test_activating_another_voice_switches_to_its_model(self) -> None: + self._export("Rufus", b"rufus") + self._export("Anna", b"anna") + self.backend.load() + self.config.tts.default_voice = "Anna" # e.g. the GUI activates the voice it trains + self._speaking() + self._reloaded() + self.assertEqual(self._speaking(), b"anna") + + def test_a_voice_whose_training_has_not_exported_a_model_yet(self) -> None: + self.assertFalse(self.backend.model_available()) + with self.assertRaisesRegex(FileNotFoundError, "Piper ONNX model not found"): + self.backend.load() + self.broken.add(b"damaged") + self._export("Rufus", b"damaged") + self.assertTrue(self.backend.model_available()) + with self.assertRaisesRegex(RuntimeError, "cannot load"): + self.backend.load() + self.assertFalse(self.backend.model_available()) # not tried again... + self._export("Rufus", b"epoch 1") + self.assertTrue(self.backend.model_available()) # ...until training exports anew + self.backend.load() + self.assertEqual(self._speaking(), b"epoch 1") + + +class TtsReviveTests(unittest.IsolatedAsyncioTestCase): + """Speech output comes on once the voice that had no usable model at startup has one.""" + + async def asyncSetUp(self) -> None: + from javis.agent import orchestrator + self.orch = object.__new__(orchestrator.Orchestrator) + self.orch._stop = asyncio.Event() + self.orch.tts = None + self.spawned: list[str] = [] + + def spawn(coro, name=None): + coro.close() + self.spawned.append(name) + + patcher = patch.object(orchestrator, "spawn", spawn) + patcher.start() + self.addCleanup(patcher.stop) + self.backend = MagicMock() + + async def _revive(self) -> None: + await asyncio.wait_for(self.orch._revive_tts(self.backend, interval=0.01), 10) + + async def test_speech_comes_on_once_the_model_loads(self) -> None: + self.backend.model_available.side_effect = [False, True, True] + self.backend.load.side_effect = [RuntimeError("damaged"), None] + await self._revive() + self.assertIs(self.orch.tts, self.backend) + self.assertEqual(self.backend.load.call_count, 2) + self.assertEqual(self.spawned, ["tts-prewarm"]) + + async def test_quitting_javis_ends_the_wait(self) -> None: + self.backend.model_available.return_value = False + revive = asyncio.create_task(self._revive()) + await asyncio.sleep(0.05) + self.orch._stop.set() + await revive + loop = asyncio.get_running_loop() + self.backend.model_available.return_value = True # quitting while the model loads + self.backend.load.side_effect = lambda: loop.call_soon_threadsafe(self.orch._stop.set) + self.orch._stop.clear() + await self._revive() + self.assertIsNone(self.orch.tts) + self.assertEqual(self.spawned, []) + + class PiperCompatTests(unittest.TestCase): def test_int_resblock_builds_high_quality_blocks(self) -> None: try: @@ -112,6 +1387,1432 @@ class PiperCompatTests(unittest.TestCase): self.assertTrue(hasattr(SileroVoiceActivityDetector, "process_array")) +def _piper_compat(): + try: + from javis.voices import piper_compat + except ImportError: + raise unittest.SkipTest("piper training extras not installed") + return piper_compat + + +def _pristine(name: str): + """A fresh copy of a ``piper.train.vits`` module, untouched by piper_compat's patches.""" + import importlib.util + + from piper.train import vits + spec = importlib.util.spec_from_file_location( + f"_pristine_vits_{name}", Path(vits.__file__).with_name(f"{name}.py")) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +class ExportInfoTests(unittest.TestCase): + """piper_compat's export records what it exported, and refuses NaN/inf weights.""" + + def setUp(self) -> None: + self.pc = _piper_compat() + from lightning.pytorch.core import saving + self.saving = saving + d = tempfile.TemporaryDirectory() + self.addCleanup(d.cleanup) + self.tmp = Path(d.name) + + def test_our_options_are_taken_out_of_piper_s(self) -> None: + for argv, value, rest in ( + (["x", "--info-file", "i.json", "--checkpoint", "c"], "i.json", + ["x", "--checkpoint", "c"]), + (["x", "--info-file=i.json"], "i.json", ["x"]), + (["x", "--info-file"], None, ["x", "--info-file"]), + (["x", "--info-files=i"], None, ["x", "--info-files=i"]), + (["x"], None, ["x"])): + with self.subTest(argv=argv): + argv = list(argv) + self.assertEqual(self.pc._pop_option(argv, "--info-file"), value) + self.assertEqual(argv, rest) + + def test_epoch_progress_of_damaged_records_is_unknown(self) -> None: + def saved(loaders): + return {"loops": {"fit_loop": {"state_dict": {"combined_loader": loaders}}}} + + self.assertEqual(self.pc._epoch_progress(saved([{"batches_done": 7, "order": [0] * 9}])), + {"batches_done": 7, "batches_per_epoch": 9}) + for checkpoint in ({}, {"loops": None}, saved(None), saved([]), saved([{}]), + saved([{"batches_done": "x", "order": [0]}]), + saved([{"batches_done": 1, "order": None}])): + with self.subTest(checkpoint=checkpoint): + self.assertEqual(self.pc._epoch_progress(checkpoint), {}) + + def test_the_checkpoint_read_is_recorded(self) -> None: + import torch + path, newer = self.tmp / "last.ckpt", self.tmp / "newer.ckpt" + path.write_bytes(b"saved progress") + newer.write_bytes(b"newer progress") + mtime = 1_700_000_000 * 10**9 + os.utime(path, ns=(mtime, mtime)) + read = [] + + def load(source, *args, **kwargs): + read.append((source.read(), args, kwargs)) + os.replace(newer, path) # training saves anew meanwhile: the file isn't held open + return {"state_dict": {"model_g.w": torch.zeros(2), "steps": torch.tensor([3])}, + "epoch": 3, "global_step": 700, + "loops": {"fit_loop": {"state_dict": {"combined_loader": [ + {"batches_done": 700, "order": list(range(1456))}]}}}} + + info: dict = {} + with patch.object(self.saving, "pl_load", load): + self.pc._inspect_exported_checkpoint(info) + self.assertEqual(self.saving.pl_load(str(path), map_location="cpu")["epoch"], 3) + self.assertEqual(read, [(b"saved progress", (), {"map_location": "cpu"})]) + self.assertEqual(info, {"checkpoint": os.path.abspath(path), "checkpoint_mtime_ns": mtime, + "epoch": 3, "global_step": 700, "batches_done": 700, + "batches_per_epoch": 1456}) + self.assertEqual(path.read_bytes(), b"newer progress") + + def test_nan_weights_are_not_exported(self) -> None: + import torch + path = self.tmp / "diverged.ckpt" + path.write_bytes(b"x") + diverged = {"state_dict": {"model_g.a": torch.zeros(2), + "model_g.b": torch.tensor([1.0, float("nan")])}} + with patch.object(self.saving, "pl_load", lambda source, *a, **kw: diverged): + self.pc._inspect_exported_checkpoint({}) + with self.assertRaisesRegex(SystemExit, + r"diverged\.ckpt holds NaN/inf weights \(model_g\.b\)"): + self.saving.pl_load(str(path)) + + def test_export_writes_its_record_when_asked_to(self) -> None: + from piper.train import export_onnx + ckpt = _checkpoint(self.tmp / "last.ckpt", step=9) + record = self.tmp / "piper_model.info.partial-1-1.json" + argvs = [] + + def main(): # piper's export loads the checkpoint + argvs.append(list(sys.argv)) + self.saving.pl_load(str(ckpt), map_location="cpu") + + for argv in (["piper_compat", "--checkpoint", str(ckpt)], + ["piper_compat", "--checkpoint", str(ckpt), "--info-file", str(record)]): + with self.subTest(argv=argv), \ + patch.object(self.saving, "pl_load", self.saving.pl_load), \ + patch.object(self.pc, "_legacy_onnx_export"), \ + patch.object(export_onnx, "main", main), patch.object(sys, "argv", list(argv)): + started = time.time() + self.pc._export(sys.argv) + self.assertEqual(argvs.pop(), ["piper_compat", "--checkpoint", str(ckpt)]) + self.assertEqual(record.exists(), "--info-file" in argv) + self.assertEqual(sorted(p.name for p in self.tmp.iterdir()), [ckpt.name, record.name]) + info = json.loads(record.read_text(encoding="utf-8")) + self.assertGreaterEqual(info.pop("exported_at"), started) + self.assertEqual(info, {"checkpoint": os.path.abspath(ckpt), + "checkpoint_mtime_ns": ckpt.stat().st_mtime_ns, + "epoch": None, "global_step": 9}) + + +class MonotonicAlignmentTests(unittest.TestCase): + """piper's alignment search runs on piper_compat's port, not a compiled extension.""" + + def setUp(self) -> None: + self.pc = _piper_compat() + self.kernels = {"numba": self.pc._batch_mas(), "numpy": self.pc._mas_numpy} + + @staticmethod + def _align(kernel, values, t_ys, t_xs) -> np.ndarray: + phoneme = np.full(values.shape[:2], -1, np.int32) + kernel(phoneme, values.astype(np.float32), np.asarray(t_ys, np.int32), + np.asarray(t_xs, np.int32)) + return phoneme + + def test_alignment_is_the_best_monotonic_path(self) -> None: + import itertools + rng = np.random.default_rng(0) + for t_y, t_x in [(1, 1), (5, 1), (4, 4), (6, 3), (7, 4), (8, 2)]: + for _ in range(5): + scores = rng.integers(-20, 20, (1, t_y + 2, t_x + 1)).astype(np.float32) # + padding + # every path from phoneme 0 to t_x - 1 that moves on by 0 or 1 phoneme per frame + best = max(scores[0, np.arange(t_y), np.cumsum((0,) + steps)].sum() + for steps in itertools.product((0, 1), repeat=t_y - 1) + if sum(steps) == t_x - 1) + for name, kernel in self.kernels.items(): + path = self._align(kernel, scores, [t_y], [t_x])[0] + with self.subTest(kernel=name, t_y=t_y, t_x=t_x): + self.assertEqual((path[0], path[t_y - 1]), (0, t_x - 1)) + self.assertTrue(set(np.diff(path[:t_y]).tolist()) <= {0, 1}) + self.assertTrue((path[t_y:] == -1).all()) + self.assertEqual(scores[0, np.arange(t_y), path[:t_y]].sum(), best) + + def test_kernels_agree_on_padded_batches(self) -> None: + rng = np.random.default_rng(1) + values = rng.normal(-200, 30, (6, 90, 25)).astype(np.float32) + t_ys, t_xs = [90, 61, 40, 25, 7, 1], [25, 20, 3, 25, 7, 1] + paths = {name: self._align(k, values, t_ys, t_xs) for name, k in self.kernels.items()} + np.testing.assert_array_equal(paths["numba"], paths["numpy"]) + for path, t_y, t_x in zip(paths["numba"], t_ys, t_xs): + self.assertEqual((path[0], path[t_y - 1]), (0, t_x - 1)) + self.assertTrue((path[t_y:] == -1).all()) + + def test_maximum_path_keeps_piper_contract(self) -> None: + import torch + from piper.train.vits import monotonic_align + neg_cent = torch.randn(3, 40, 12, generator=torch.Generator().manual_seed(2)) * 30 - 200 + t_ys, t_xs = torch.tensor([40, 31, 12]), torch.tensor([12, 9, 12]) + mask = ((torch.arange(40)[None, :, None] < t_ys[:, None, None]) + & (torch.arange(12)[None, None, :] < t_xs[:, None, None])).float() + before = neg_cent.clone() + attn = monotonic_align.maximum_path(neg_cent, mask) + self.assertIs(monotonic_align.maximum_path, self.pc.maximum_path) + self.assertEqual((attn.shape, attn.dtype), (neg_cent.shape, neg_cent.dtype)) + self.assertTrue(torch.equal(neg_cent, before)) # works on a copy + self.assertTrue(torch.equal(attn.sum(2), mask[:, :, 0])) # one phoneme per valid frame + self.assertTrue(torch.equal(attn * mask, attn)) + + def test_stands_in_for_a_missing_extension(self) -> None: + from piper.train import vits + name = "piper.train.vits.monotonic_align" + saved = vits.__dict__.pop("monotonic_align") + try: + with patch.dict(sys.modules, {name: None}): # its import now fails + self.pc._fix_monotonic_alignment() + from piper.train.vits import monotonic_align # as SynthesizerTrn.forward does + self.assertIs(monotonic_align.maximum_path, self.pc.maximum_path) + finally: + vits.monotonic_align = saved + + def test_piper_forward_maps_every_frame_to_one_phoneme(self) -> None: + import torch + from piper.train.vits.losses import kl_loss + from piper.train.vits.models import SynthesizerTrn + torch.manual_seed(0) + model = SynthesizerTrn( + n_vocab=12, spec_channels=17, segment_size=8, inter_channels=8, hidden_channels=8, + filter_channels=16, n_heads=2, n_layers=1, kernel_size=3, p_dropout=0.0, + resblock="2", resblock_kernel_sizes=(3,), resblock_dilation_sizes=((1, 3),), + upsample_rates=(2, 2), upsample_initial_channel=8, upsample_kernel_sizes=(4, 4), + use_sdp=False) + x_lengths, y_lengths = torch.tensor([9, 5]), torch.tensor([40, 23]) + with torch.no_grad(): + _, _, attn, _, _, y_mask, (_, z_p, m_p, logs_p, _, logs_q) = model( + torch.randint(1, 12, (2, 9)), x_lengths, torch.randn(2, 17, 40), y_lengths) + attn = attn.squeeze(1) + self.assertTrue(torch.equal(attn.sum(2), y_mask.squeeze(1))) + phoneme = attn.argmax(2) + for i, (t_x, t_y) in enumerate(zip(x_lengths.tolist(), y_lengths.tolist())): + self.assertEqual((phoneme[i, 0].item(), phoneme[i, t_y - 1].item()), (0, t_x - 1)) + self.assertTrue(set(torch.diff(phoneme[i, :t_y]).tolist()) <= {0, 1}) + self.assertTrue(torch.isfinite(kl_loss(z_p, logs_q, m_p, logs_p, y_mask))) + + +class LengthBucketTests(unittest.TestCase): + def setUp(self) -> None: + self.pc = _piper_compat() + rng = np.random.default_rng(1) + self.lengths = [int(x) for x in rng.integers(1, 1000, 1000)] + + def _sampler(self, **kw): + kw = {"drop_last": False, "seed": 7, **kw} + return self.pc.LengthBucketBatchSampler(self.lengths, kw.pop("batch_size", 8), **kw) + + def test_every_clip_once_per_epoch(self) -> None: + s = self._sampler(batch_size=7) + batches = list(s) + self.assertEqual(len(batches), len(s)) + self.assertEqual(sorted(i for b in batches for i in b), list(range(1000))) + self.assertTrue(all(1 <= len(b) <= 7 for b in batches)) + + def test_drop_last_keeps_full_batches(self) -> None: + s = self._sampler(batch_size=7, drop_last=True) + batches = list(s) + self.assertEqual(len(batches), len(s)) + self.assertEqual(len(s), 1000 // 7) + self.assertTrue(all(len(b) == 7 for b in batches)) + flat = [i for b in batches for i in b] + self.assertEqual(len(flat), len(set(flat))) + + def test_batches_pad_far_less_than_random_ones(self) -> None: + s = self._sampler() + rng = np.random.default_rng(2) + order = rng.permutation(1000).tolist() + unbucketed = [order[i:i + 8] for i in range(0, 1000, 8)] + self.assertLess(self.pc.padding_factor(s.batches(0), self.lengths), 1.1) + self.assertGreater(self.pc.padding_factor(unbucketed, self.lengths), 1.4) + + def test_order_is_reproducible_per_epoch_and_reshuffled_across_epochs(self) -> None: + self.assertEqual(self._sampler().batches(3), self._sampler().batches(3)) + self.assertNotEqual(self._sampler().batches(0), self._sampler().batches(1)) + self.assertNotEqual(self._sampler().batches(0), self._sampler(seed=8).batches(0)) + + def test_first_epoch_of_a_run_starts_with_the_longest_batch(self) -> None: + s = self._sampler() + s.set_epoch(4) + first, again = list(s), list(s) # a DataLoader with workers iterates twice + self.assertEqual(max(self.lengths[i] for i in first[0]), max(self.lengths)) + self.assertCountEqual(first, s.batches(4)) + self.assertEqual(again, first) + s.set_epoch(5) + self.assertEqual(list(s), s.batches(5)) # only the run's first epoch + + def test_longest_batch_first_with_dataloader_workers(self) -> None: + from torch.utils.data import DataLoader + loader = DataLoader(range(1000), batch_sampler=self._sampler(), num_workers=1) + first = next(iter(loader)).tolist() + self.assertEqual(max(self.lengths[i] for i in first), max(self.lengths)) + + def test_lightning_sets_the_epoch_through_the_dataloader(self) -> None: + from lightning.fabric.utilities.data import _set_sampler_epoch + from torch.utils.data import DataLoader + s = self._sampler(longest_first=False) + loader = DataLoader(range(1000), batch_sampler=s) + _set_sampler_epoch(loader, 5) + self.assertEqual(s.epoch, 5) + self.assertEqual([b.tolist() for b in loader], s.batches(5)) + + def _resumed(self, state, epoch: int, **kw): + """``(batches, what was printed)`` of a new run resuming from *state* in *epoch*.""" + sampler = self._sampler(**kw) + with contextlib.redirect_stdout(io.StringIO()) as out: + sampler.load_state_dict(state) + sampler.set_epoch(epoch) + return list(sampler), out.getvalue() + + def test_an_interrupted_epoch_resumes_with_the_batches_it_had_not_trained(self) -> None: + s = self._sampler() + self.assertEqual(s.state_dict(0), {}) # no epoch started + s.set_epoch(3) + first = list(s) # the run's first epoch: longest batch first + state = s.state_dict(40) + self.assertEqual((state["epoch"], state["batches_done"]), (3, 40)) + resumed = self._sampler() + with contextlib.redirect_stdout(io.StringIO()) as out: + resumed.load_state_dict(json.loads(json.dumps(state))) # (as saved and loaded) + rest = list(resumed) # as Lightning does: before it sets the epoch + resumed.set_epoch(3) + self.assertEqual(out.getvalue(), "") + self.assertEqual(list(resumed), rest) # a DataLoader with workers iterates twice + self.assertCountEqual(first[:40] + rest, s.batches(3)) # each batch once + self.assertEqual(max(self.lengths[i] for i in rest[0]), # longest first again + max(self.lengths[i] for b in first[40:] for i in b)) + resumed.set_epoch(4) + self.assertEqual(list(resumed), s.batches(4)) # later epochs are whole + later = self._sampler(longest_first=False) + later.set_epoch(4) + list(later) + rest, _ = self._resumed(later.state_dict(10), 4, longest_first=False) + self.assertEqual(rest, later.batches(4)[10:]) # in the order it would have had + + def test_the_loader_saves_and_restores_the_position_in_its_epoch(self) -> None: + from torch.utils.data import DataLoader + trained = {"batches": 0} + loader = self.pc.ResumableDataLoader(range(1000), batch_sampler=self._sampler(), + batches_done=lambda: trained["batches"]) + self.assertIsInstance(loader, DataLoader) + self.assertEqual(loader.state_dict(), {}) # (before its first epoch) + first = [b.tolist() for b in loader] + trained["batches"] = 100 + state = loader.state_dict() + self.assertEqual((state["epoch"], state["batches_done"]), (0, 100)) + resumed = self.pc.ResumableDataLoader(range(1000), batch_sampler=self._sampler()) + resumed.load_state_dict(state) + once, twice = [b.tolist() for b in resumed], [b.tolist() for b in resumed] + self.assertEqual(once, twice) + self.assertCountEqual(once, first[100:]) + + def test_saved_positions_that_cannot_resume_start_the_epoch_afresh(self) -> None: + s = self._sampler(longest_first=False) + s.set_epoch(2) + list(s) + cases = [ + ("piper's random batches", {}, 2), + ("nothing", None, 2), + ("damaged", {**s.state_dict(5), "order": "x"}, 2), + ("start of the epoch", s.state_dict(0), 2), + ("end of the epoch", s.state_dict(len(s)), 2), # validated and saved: the next one + ("another epoch", s.state_dict(5), 3), + ] + for name, state, epoch in cases: + with self.subTest(name): + self.assertEqual(self._resumed(state, epoch, longest_first=False), + (s.batches(epoch), "")) + + def test_changed_training_data_starts_the_interrupted_epoch_afresh(self) -> None: + s = self._sampler(longest_first=False) + list(s) + state = s.state_dict(5) + for change in ("a clip more", "a clip longer"): + with self.subTest(change): + self.lengths = list(s.lengths) + if change == "a clip more": + self.lengths.append(500) + else: # as many batches, but different ones + self.lengths[self.lengths.index(min(self.lengths))] = 5000 + fresh = self._sampler(longest_first=False).batches(0) + self.assertNotEqual(fresh, s.batches(0)) + batches, out = self._resumed(state, 0, longest_first=False) + self.assertIn("training data changed", out) + self.assertEqual(batches, fresh) + + def test_train_dataloader_buckets_piper_utterances(self) -> None: + import torch + from lightning.fabric.utilities.data import _set_sampler_epoch + from piper.train.vits.dataset import CachedUtterance, VitsDataset + from torch.utils.data import Subset + + samples = [1000, 20000, 1100, 19000, 5000, 5200, 300] + with tempfile.TemporaryDirectory() as d: + utts = [] + for i, n in enumerate(samples): + stem = Path(d) / f"u{i}" + torch.save(torch.tensor([1, 2, 3]), f"{stem}.phonemes.pt") + torch.save(torch.zeros(n), f"{stem}.audio.pt") + torch.save(torch.zeros(4, n // 256 + 1), f"{stem}.spec.pt") + utts.append(CachedUtterance(Path(f"{stem}.phonemes.pt"), Path(f"{stem}.audio.pt"), + Path(f"{stem}.spec.pt"), text="x")) + # random_split() hands piper a Subset (here nested); clip 6 is held out. + train = Subset(Subset(VitsDataset(utts), [6, 5, 4, 3, 2, 1, 0]), [1, 2, 3, 4, 5, 6]) + dm = SimpleNamespace(train_dataset=train, batch_size=2, num_workers=0, + num_speakers=1, segment_size=256) + with contextlib.redirect_stdout(io.StringIO()) as out: + loader = self.pc._bucketed_train_dataloader(dm) + self.assertIn("Length-bucketed batches: 3 per epoch", out.getvalue()) + self.assertEqual(len(loader), 3) + _set_sampler_epoch(loader, 2) + loaded = list(loader) + batches = [sorted(b.audio_lengths.tolist()) for b in loaded] + self.assertEqual(batches[0], [19000, 20000]) # longest first + self.assertCountEqual(batches, [[1000, 1100], [5000, 5200], [19000, 20000]]) + self.assertEqual({b.spectrograms.shape[-1] % 32 for b in loaded}, {0}) # rounded up + self.assertEqual({b.phoneme_ids.shape[-1] for b in loaded}, {16}) + + def test_collate_rounds_padded_lengths_up(self) -> None: + import pickle + + import torch + from piper.train.vits.dataset import UtteranceCollate, UtteranceTensors + utts = [UtteranceTensors(phoneme_ids=torch.arange(1, p + 1), spectrogram=torch.rand(4, t), + audio_norm=torch.rand(t * 256)) + for p, t in ((5, 40), (17, 33))] + plain = UtteranceCollate(is_multispeaker=False, segment_size=256) + # DataLoader workers get the collate function pickled. + rounded = pickle.loads(pickle.dumps(self.pc.RoundedCollate(plain))) + expected, batch = plain(utts), rounded(utts) + self.assertEqual(tuple(batch.spectrograms.shape), (2, 4, 64)) + self.assertEqual(tuple(batch.phoneme_ids.shape), (2, 32)) + torch.testing.assert_close(batch.spectrograms[..., :40], expected.spectrograms, rtol=0, atol=0) + torch.testing.assert_close(batch.phoneme_ids[:, :17], expected.phoneme_ids, rtol=0, atol=0) + self.assertFalse(batch.spectrograms[..., 40:].any()) + self.assertFalse(batch.phoneme_ids[:, 17:].any()) + for name in ("phoneme_lengths", "spectrogram_lengths", "audios", "audio_lengths"): + torch.testing.assert_close(getattr(batch, name), getattr(expected, name), rtol=0, atol=0) + exact = [UtteranceTensors(phoneme_ids=torch.arange(1, 17), spectrogram=torch.rand(4, 64), + audio_norm=torch.rand(64 * 256))] + self.assertEqual(rounded(exact).spectrograms.shape[-1], 64) + + def test_extra_padding_leaves_real_frames_alone(self) -> None: + import torch + from piper.train.vits.models import SynthesizerTrn + torch.manual_seed(0) + model = SynthesizerTrn( + n_vocab=12, spec_channels=17, segment_size=8, inter_channels=8, hidden_channels=8, + filter_channels=16, n_heads=2, n_layers=1, kernel_size=3, p_dropout=0.0, + resblock="2", resblock_kernel_sizes=(3,), resblock_dilation_sizes=((1, 3),), + upsample_rates=(2, 2), upsample_initial_channel=8, upsample_kernel_sizes=(4, 4), + use_sdp=False).eval() + x_lengths, y_lengths = torch.tensor([9, 5]), torch.tensor([40, 23]) + x = torch.randint(1, 12, (2, 9)) * (torch.arange(9) < x_lengths[:, None]) + y = torch.randn(2, 17, 40) * (torch.arange(40) < y_lengths[:, None, None]) + + def encode(x, y): # every part of piper's VITS that sees whole utterances + with torch.no_grad(): + h, m_p, logs_p, x_mask = model.enc_p(x, x_lengths) + _, m_q, logs_q, y_mask = model.enc_q(y, y_lengths) + return {"h": h, "m_p": m_p, "logs_p": logs_p, "logw": model.dp(h, x_mask), + "m_q": m_q, "logs_q": logs_q, "z_p": model.flow(m_q, y_mask)} + + padded = encode(self.pc._pad_to_multiple(x, 16), self.pc._pad_to_multiple(y, 32)) + for name, value in encode(x, y).items(): + with self.subTest(name): + self.assertGreater(padded[name].shape[-1], value.shape[-1]) + torch.testing.assert_close(padded[name][..., :value.shape[-1]], value) + + def test_training_aids_hook_into_piper(self) -> None: + from lightning.fabric.utilities.types import _Stateful + from lightning.pytorch.utilities.combined_loader import CombinedLoader + from piper.train.vits.dataset import VitsDataModule + from piper.train.vits.lightning import VitsModel + from torch.utils.data import DataLoader + tc = self.pc.training_control + console = SimpleNamespace(closing=threading.Event()) + for tuning, stop_file in (("1", None), ("0", "stop")): + with self.subTest(JAVIS_PIPER_CUDNN_TUNING=tuning, stop_file=stop_file), \ + patch.object(VitsModel, "configure_callbacks", VitsModel.configure_callbacks), \ + patch.object(VitsDataModule, "train_dataloader", + lambda dm: DataLoader(range(3))), \ + patch.object(self.pc, "_avoid_gpu_syncs") as avoid_gpu_syncs, \ + patch.object(self.pc, "_safe_checkpoint_files") as safe_checkpoint_files, \ + patch.object(self.pc, "_checkpoint_every_epoch") as checkpoint_every_epoch, \ + patch.object(self.pc, "_robust_dataset") as robust_dataset, \ + patch.object(self.pc, "_on_cache_write", None), \ + patch.dict(os.environ, {"JAVIS_PIPER_BUCKETING": "1", + "JAVIS_PIPER_CUDNN_TUNING": tuning, + "JAVIS_PIPER_PROGRESS_MINUTES": "4"}): + os.environ.pop(tc.LAUNCHER_ENV, None) + if stop_file: + os.environ[tc.STOP_FILE_ENV] = stop_file + else: + os.environ.pop(tc.STOP_FILE_ENV, None) + stopper = self.pc._install_training_aids(console) + for patched in (avoid_gpu_syncs, safe_checkpoint_files, checkpoint_every_epoch): + patched.assert_called_once_with() + robust_dataset.assert_called_once_with(VitsDataModule) + callbacks = VitsModel.configure_callbacks(SimpleNamespace()) + self.assertTrue(any(isinstance(c, self.pc.NonFiniteGuard) for c in callbacks)) + self.assertEqual(any(isinstance(c, self.pc.CudnnFirstStepTuning) for c in callbacks), + tuning == "1") + savers = [c for c in callbacks if isinstance(c, self.pc.ProgressSaver)] + self.assertEqual([(s.stopper, s.console, s.interval) for s in savers], + [(stopper, console, 240.0)]) + if stop_file: # started by the launcher: stops on request, even while preprocessing + self.assertIsInstance(stopper, self.pc.StopOnRequest) + self.assertIn(stopper, callbacks) + self.assertEqual(self.pc._on_cache_write, stopper.raise_if_requested) + else: + self.assertIsNone(stopper) + self.assertIsNone(self.pc._on_cache_write) + with contextlib.redirect_stderr(io.StringIO()) as err: # unusable for bucketing + loader = VitsDataModule.train_dataloader(SimpleNamespace()) + self.assertIn("length bucketing unavailable", err.getvalue()) + # piper's batches still load the checkpoints of bucketed ones + self.assertIsInstance(loader, _Stateful) + self.assertEqual(loader.state_dict(), {}) + CombinedLoader(loader)._load_state_dicts([{"epoch": 1, "batches_done": 5}]) + self.assertEqual(list(loader), [0, 1, 2]) + + +class PrecisionTests(unittest.TestCase): + def setUp(self) -> None: + self.pc = _piper_compat() + + def _resolve(self, argv, **kw): + with contextlib.redirect_stdout(io.StringIO()): + self.pc._resolve_precision(argv, **kw) + return argv + + def test_bf16_falls_back_to_32_without_native_support(self) -> None: + argv = ["fit", "--trainer.precision", "bf16-mixed", "--trainer.devices", "1"] + self.assertEqual(self._resolve(list(argv), bf16_supported=True), argv) + self.assertEqual(self._resolve(list(argv), bf16_supported=False)[2], "32") + self.assertEqual(self._resolve(["fit", "--trainer.precision=bf16-mixed"], + bf16_supported=False)[1], "--trainer.precision=32") + + def test_other_precisions_are_left_alone(self) -> None: + for value in ("32", "16-mixed"): + argv = ["fit", "--trainer.precision", value] + self.assertEqual(self._resolve(list(argv), bf16_supported=False), argv) + + def test_cpu_training_uses_32(self) -> None: + with patch.object(self.pc, "_cuda_bf16_supported", return_value=True): + argv = self._resolve(["fit", "--trainer.accelerator", "cpu", + "--trainer.precision", "bf16-mixed"]) + self.assertEqual(argv[-1], "32") + argv = self._resolve(["fit", "--trainer.accelerator", "gpu", + "--trainer.precision", "bf16-mixed"]) + self.assertEqual(argv[-1], "bf16-mixed") + + def test_config_defaults_to_bf16(self) -> None: + from javis.core.config import Config, validate + cfg = Config() + self.assertEqual(cfg.training.precision, "bf16-mixed") + cfg.training.precision = "fp8" + with self.assertLogs("javis.config", level="WARNING"): + validate(cfg) + self.assertEqual(cfg.training.precision, "bf16-mixed") + + +class NonFiniteGuardTests(unittest.TestCase): + """The guard inside a real Lightning loop with piper's GAN-style manual optimization.""" + + def setUp(self) -> None: + self.pc = _piper_compat() + + def _fit(self, nan_batches, guard, optimizer=None): + import lightning as L + import torch + from torch.utils.data import DataLoader + + make_optimizer = optimizer or torch.optim.AdamW + + class GanLike(L.LightningModule): + def __init__(self): + super().__init__() + self.automatic_optimization = False + self.g = torch.nn.Linear(2, 2) + self.d = torch.nn.Linear(2, 1) + self.batches = [] + + def training_step(self, batch, batch_idx): + self.batches.append(batch_idx) + opt_g, opt_d = self.optimizers() + loss_g = self.g(batch).sum() + if batch_idx in nan_batches: + loss_g = loss_g * float("inf") + opt_g.zero_grad() + self.manual_backward(loss_g) + opt_g.step() + opt_d.zero_grad() + self.manual_backward(self.d(batch).sum()) + opt_d.step() + + def configure_optimizers(self): + self.opts = [make_optimizer(self.g.parameters(), lr=0.1), + make_optimizer(self.d.parameters(), lr=0.1)] + return self.opts + + def train_dataloader(self): + return DataLoader(torch.ones(6, 2), batch_size=1) + + model = self.model = GanLike() + trainer = L.Trainer(accelerator="cpu", max_epochs=1, callbacks=[guard], logger=False, + enable_checkpointing=False, enable_progress_bar=False, + enable_model_summary=False) + with warnings.catch_warnings(), contextlib.redirect_stderr(io.StringIO()): + warnings.simplefilter("ignore") + trainer.fit(model) + return model + + def test_bad_steps_are_skipped_and_counted(self) -> None: + import torch + guard = self.pc.NonFiniteGuard(max_consecutive=3) + model = self._fit({1, 2, 4}, guard) + self.assertEqual(guard.skipped, {0: 3}) # generator only; streak was reset + self.assertTrue(all(bool(torch.isfinite(p).all()) for p in model.parameters())) + # piper's AdamW skips on the device: fused, its state untouched by bad steps + for opt, good_steps in zip(model.opts, (3, 6)): + self.assertTrue(all(group["fused"] for group in opt.param_groups)) + self.assertEqual({float(s["step"]) for s in opt.state.values()}, {good_steps}) + + def test_other_optimizers_are_checked_on_the_cpu(self) -> None: + import torch + guard = self.pc.NonFiniteGuard(max_consecutive=3) + model = self._fit({1, 2, 4}, guard, optimizer=torch.optim.SGD) + self.assertEqual(guard.skipped, {0: 3}) + self.assertTrue(all(bool(torch.isfinite(p).all()) for p in model.parameters())) + + def test_a_diverged_run_stops(self) -> None: + import torch + for optimizer in (torch.optim.AdamW, torch.optim.SGD): + with self.subTest(optimizer=optimizer.__name__): + guard = self.pc.NonFiniteGuard(max_consecutive=4, check_every=2) + with self.assertRaises(self.pc.TrainingDivergedError): + self._fit({1, 2, 3, 4, 5}, guard, optimizer) + self.assertEqual(self.model.batches, [0, 1, 2, 3, 4]) # at the 4th in a row + + def test_restored_adam_state_moves_to_the_fused_kernel(self) -> None: + import copy + + import torch + torch.manual_seed(0) + layer = torch.nn.Linear(3, 3) + twin = copy.deepcopy(layer) + + def train(model, optimizer, steps): + for _ in range(steps): + optimizer.zero_grad() + model(torch.ones(2, 3)).square().sum().backward() + optimizer.step() + + optimizer = torch.optim.AdamW(layer.parameters(), lr=0.01, foreach=True) + reference = torch.optim.AdamW(twin.parameters(), lr=0.01, foreach=True) + train(layer, optimizer, 2) + train(twin, reference, 2) + restored = torch.optim.AdamW(layer.parameters(), lr=0.01) # as Lightning resumes it + restored.load_state_dict(copy.deepcopy(optimizer.state_dict())) + self.assertTrue(self.pc._use_fused_adam(restored)) + self.assertTrue(all(g["fused"] and not g["foreach"] for g in restored.param_groups)) + train(layer, restored, 2) + train(twin, reference, 2) + for p, q in zip(layer.parameters(), twin.parameters()): + torch.testing.assert_close(p, q) + self.assertEqual({float(s["step"]) for s in restored.state.values()}, {4.0}) + resumed = torch.optim.AdamW(layer.parameters(), lr=0.01) # its checkpoint, resumed + resumed.load_state_dict(restored.state_dict()) + self.assertTrue(all(g["fused"] for g in resumed.param_groups)) + self.assertTrue(self.pc._use_fused_adam(resumed)) + self.assertFalse(self.pc._use_fused_adam(torch.optim.SGD(layer.parameters(), lr=0.1))) + + def test_skipping_on_the_gpu_never_waits_for_it(self) -> None: + import torch + if not torch.cuda.is_available(): + self.skipTest("needs CUDA") + layer = torch.nn.Linear(4, 4, device="cuda") + optimizer = torch.optim.AdamW(layer.parameters(), lr=0.1) + trainer = SimpleNamespace(optimizers=[optimizer], global_step=0) + guard = self.pc.NonFiniteGuard() + guard.on_train_start(trainer, None) + start = layer.weight.detach().clone() + x = torch.ones(2, 4, device="cuda") + weights = [] + for scale in (float("inf"), 1.0): + optimizer.zero_grad() + (layer(x).sum() * scale).backward() + torch.cuda.synchronize() + torch.cuda.set_sync_debug_mode("error") + try: + guard.on_before_optimizer_step(trainer, None, optimizer) + optimizer.step() + finally: + torch.cuda.set_sync_debug_mode("default") + weights.append(layer.weight.detach().clone()) + torch.testing.assert_close(weights[0], start, rtol=0, atol=0) # skipped + self.assertFalse(torch.equal(weights[1], start)) + self.assertEqual({float(s["step"]) for s in optimizer.state.values()}, {1.0}) + with contextlib.redirect_stderr(io.StringIO()): + guard.on_train_end(trainer, None) + self.assertEqual(guard.skipped, {0: 1}) + + def test_collapsed_fp16_scale_stops_the_run(self) -> None: + guard = self.pc.NonFiniteGuard() + module = SimpleNamespace(log_dict=lambda *a, **k: None) + for scale, fails in ((65536.0, False), (0.0, True), (float("inf"), True)): + scaler = SimpleNamespace(is_enabled=lambda: True, get_scale=lambda s=scale: s) + trainer = SimpleNamespace(precision_plugin=SimpleNamespace(scaler=scaler)) + if fails: + with self.assertRaises(self.pc.TrainingDivergedError): + guard.on_train_batch_end(trainer, module, None, None, 0) + else: + guard.on_train_batch_end(trainer, module, None, None, 0) + + def test_nan_weights_are_never_saved(self) -> None: + import torch + guard = self.pc.NonFiniteGuard() + layer = torch.nn.Linear(2, 2) + guard.on_save_checkpoint(None, layer, {}) + with torch.no_grad(): + layer.weight[0, 0] = float("nan") + with self.assertRaises(self.pc.TrainingDivergedError): + guard.on_save_checkpoint(None, layer, {}) + + +class GpuSyncTests(unittest.TestCase): + """piper_compat's sync-free replacements compute exactly what piper's code computes.""" + + def setUp(self) -> None: + self.pc = _piper_compat() + import torch + self.torch = torch + torch.manual_seed(0) + + def _run(self, fn, *inputs): + """``fn``'s outputs and the gradients of a fixed random projection of them.""" + torch = self.torch + leaves = [t.detach().clone().requires_grad_(t.is_floating_point()) for t in inputs] + outs = fn(*leaves) + outs = outs if isinstance(outs, tuple) else (outs,) + gen = torch.Generator().manual_seed(1) + sum((o * torch.randn(o.shape, generator=gen, dtype=o.dtype)).sum() for o in outs).backward() + return [o.detach() for o in outs] + [t.grad for t in leaves if t.requires_grad] + + def _assert_same(self, fn, reference, *inputs) -> None: + for mine, theirs in zip(self._run(fn, *inputs), self._run(reference, *inputs), strict=True): + self.torch.testing.assert_close(mine, theirs, rtol=0, atol=0) + + def test_slice_segments_matches_piper(self) -> None: + torch = self.torch + piper = _pristine("commons").slice_segments + cases = [(torch.randn(3, 4, 50, dtype=torch.float64), torch.tensor([0, 7, -2]), 16), + (torch.randn(2, 4, 10, dtype=torch.float64), torch.tensor([0, -1]), 16), + (torch.randn(2, 1, 300, dtype=torch.float64), torch.tensor([3, 5]) * 16, 64)] + for x, ids, size in cases: + self._assert_same(lambda t, ids=ids, size=size: self.pc.slice_segments(t, ids, size), + lambda t, ids=ids, size=size: piper(t, ids, size), x) + + def test_rand_slice_segments_matches_piper(self) -> None: + torch = self.torch + piper = _pristine("commons").rand_slice_segments + x, lengths = torch.randn(3, 4, 50), torch.tensor([50, 40, 10]) # 10: shorter than a window + starts = set() + for seed in range(10): + torch.manual_seed(seed) # on the CPU both draw the same starts + out, ids = self.pc.rand_slice_segments(x, lengths, 16) + torch.manual_seed(seed) + ref_out, ref_ids = piper(x, lengths, 16) + torch.testing.assert_close(ids, ref_ids, rtol=0, atol=0) + torch.testing.assert_close(out, ref_out, rtol=0, atol=0) + self.assertTrue(bool((ids.clamp_min(0) + 16 <= lengths.clamp_min(16)).all())) + starts.add(int(ids[0])) + self.assertGreater(len(starts), 1) + + def test_discriminator_loss_matches_piper(self) -> None: + torch = self.torch + piper = _pristine("losses").discriminator_loss + outputs = [torch.randn(2, 1, 30, dtype=torch.bfloat16), torch.randn(2, 5), + torch.randn(2, 1, 30, dtype=torch.bfloat16), torch.randn(2, 5)] + + def run(fn): + leaves = [t.clone().requires_grad_(True) for t in outputs] + loss, r_losses, g_losses = fn(leaves[:2], leaves[2:]) + loss.backward() + return loss.detach(), [float(v) for v in r_losses + g_losses], [t.grad for t in leaves] + + (loss, terms, grads), (ref_loss, ref_terms, ref_grads) = run(self.pc.discriminator_loss), run(piper) + torch.testing.assert_close(loss, ref_loss, rtol=0, atol=0) + self.assertEqual(terms, ref_terms) + for grad, ref_grad in zip(grads, ref_grads, strict=True): + torch.testing.assert_close(grad, ref_grad, rtol=0, atol=0) + + def test_spline_matches_piper(self) -> None: + torch = self.torch + reference, patched = _pristine("transforms"), _pristine("transforms") + patched.unconstrained_rational_quadratic_spline = self.pc._mask_free_spline( + patched.unconstrained_rational_quadratic_spline) + inputs = torch.linspace(-7, 7, 60, dtype=torch.float64).view(2, 1, 30) # partly outside ±5 + params = [torch.randn(2, 1, 30, n, dtype=torch.float64) for n in (10, 10, 9)] + for inverse, x in ((False, inputs), (True, inputs * 0.6)): # inverse: inference, piper's code + self._assert_same( + lambda *a, inverse=inverse: patched.piecewise_rational_quadratic_transform( + *a, inverse=inverse, tails="linear", tail_bound=5.0), + lambda *a, inverse=inverse: reference.piecewise_rational_quadratic_transform( + *a, inverse=inverse, tails="linear", tail_bound=5.0), + x, *params) + + def test_mel_spectrogram_matches_piper(self) -> None: + torch = self.torch + piper = _pristine("mel_processing").mel_spectrogram_torch + args = (1024, 80, 22050, 256, 1024, 0.0, None) + y = torch.rand(2, 8192) * 2 - 1 + self._assert_same(lambda t: self.pc.mel_spectrogram_torch(t, *args), + lambda t: piper(t, *args), y) + + def test_flip_matches_piper(self) -> None: + torch = self.torch + from piper.train.vits import modules + flip, x = modules.Flip(), torch.randn(2, 4, 7, dtype=torch.float64) + with patch.object(modules.Flip, "forward", self.pc._flip_forward): + flipped, logdet = flip(x) + restored = flip(flipped, reverse=True) + torch.testing.assert_close(flipped, torch.flip(x, [1]), rtol=0, atol=0) + torch.testing.assert_close(restored, x, rtol=0, atol=0) + torch.testing.assert_close(logdet, torch.zeros(2, dtype=torch.float64), rtol=0, atol=0) + + def test_batches_are_pinned(self) -> None: + import dataclasses + + from piper.train.vits.dataset import Batch + from torch.utils.data._utils.pin_memory import pin_memory + torch = self.torch + if not torch.cuda.is_available(): + self.skipTest("pinned memory needs CUDA") + batch = Batch(phoneme_ids=torch.ones(2, 3, dtype=torch.long), + phoneme_lengths=torch.tensor([3, 2]), spectrograms=torch.randn(2, 5, 4), + spectrogram_lengths=torch.tensor([4, 3]), audios=torch.randn(2, 1, 1024), + audio_lengths=torch.tensor([1024, 768]), speaker_ids=None) + with patch.object(Batch, "pin_memory", self.pc._pin_batch, create=True): + pinned = pin_memory(batch) # what the DataLoader's pinning thread calls + self.assertIsInstance(pinned, Batch) + self.assertIsNone(pinned.speaker_ids) + for field in dataclasses.fields(Batch): + value = getattr(batch, field.name) + if value is not None: + self.assertTrue(getattr(pinned, field.name).is_pinned(), field.name) + torch.testing.assert_close(getattr(pinned, field.name), value, rtol=0, atol=0) + + def test_training_math_never_waits_for_the_gpu(self) -> None: + torch = self.torch + if not torch.cuda.is_available(): + self.skipTest("needs CUDA") + cuda = torch.device("cuda") + spline = self.pc._mask_free_spline( + _pristine("transforms").unconstrained_rational_quadratic_spline) + z = torch.randn(4, 8, 100, device=cuda, requires_grad=True) + lengths = torch.tensor([100, 80, 60, 20], device=cuda) + audio = (torch.rand(4, 1, 100 * 256, device=cuda) * 2 - 1).requires_grad_() + flow_in = (torch.randn(4, 1, 50, device=cuda) * 4).requires_grad_() + flow_params = [torch.randn(4, 1, 50, n, device=cuda, requires_grad=True) for n in (10, 10, 9)] + mel_args = (1024, 80, 22050, 256, 1024, 0.0, None) + self.pc.mel_spectrogram_torch(audio[:, 0], *mel_args) # caches window and mel basis + torch.cuda.synchronize() + torch.cuda.set_sync_debug_mode("error") + try: + z_slice, ids = self.pc.rand_slice_segments(z, lengths, 32) + mel = self.pc.mel_spectrogram_torch( + self.pc.slice_segments(audio, ids * 256, 32 * 256).squeeze(1), *mel_args) + flow_out, flow_logdet = spline(flow_in, *flow_params, tails="linear", tail_bound=5.0) + flipped, flip_logdet = self.pc._flip_forward(None, z_slice) + d_loss, _, _ = self.pc.discriminator_loss([z_slice], [flipped]) + (mel.mean() + flow_out.sum() + flow_logdet.sum() + flip_logdet.sum() + d_loss).backward() + finally: + torch.cuda.set_sync_debug_mode("default") + torch.cuda.synchronize() + self.assertTrue(all(t.grad is not None for t in (z, audio, flow_in, *flow_params))) + + def test_installed_into_piper(self) -> None: + from piper.train.vits import commons, losses, modules, transforms + from piper.train.vits import lightning as vits_lightning + from piper.train.vits.dataset import Batch + targets = [(commons, "slice_segments"), (commons, "rand_slice_segments"), + (losses, "discriminator_loss"), (transforms, "unconstrained_rational_quadratic_spline"), + (vits_lightning, "slice_segments"), (vits_lightning, "discriminator_loss"), + (vits_lightning, "mel_spectrogram_torch"), (modules.Flip, "forward")] + had_pin_memory = "pin_memory" in vars(Batch) + + def drop_pin_memory(): + if not had_pin_memory and "pin_memory" in vars(Batch): + del Batch.pin_memory + + with contextlib.ExitStack() as restore: + for owner, name in targets: + restore.enter_context(patch.object(owner, name, getattr(owner, name))) + restore.callback(drop_pin_memory) + self.pc._avoid_gpu_syncs() + for module in (commons, vits_lightning): + self.assertIs(module.slice_segments, self.pc.slice_segments) + for module in (losses, vits_lightning): + self.assertIs(module.discriminator_loss, self.pc.discriminator_loss) + self.assertIs(commons.rand_slice_segments, self.pc.rand_slice_segments) + self.assertIs(vits_lightning.mel_spectrogram_torch, self.pc.mel_spectrogram_torch) + self.assertIs(modules.Flip.forward, self.pc._flip_forward) + self.assertTrue(hasattr(Batch, "pin_memory")) + spline = transforms.unconstrained_rational_quadratic_spline + self.assertTrue(spline._javis_mask_free) + self.pc._avoid_gpu_syncs() + self.assertIs(transforms.unconstrained_rational_quadratic_spline, spline) # not wrapped twice + + +class CudnnTuningTests(unittest.TestCase): + def test_cudnn_times_algorithms_during_the_first_step_only(self) -> None: + import lightning as L + import torch + from torch.utils.data import DataLoader + pc = _piper_compat() + self.addCleanup(setattr, torch.backends.cudnn, "benchmark", torch.backends.cudnn.benchmark) + seen = [] + + class Probe(L.LightningModule): + def __init__(self): + super().__init__() + self.w = torch.nn.Linear(1, 1) + + def training_step(self, batch, batch_idx): + seen.append(torch.backends.cudnn.benchmark) + return self.w(batch).sum() + + def configure_optimizers(self): + return torch.optim.SGD(self.parameters(), lr=0.1) + + def train_dataloader(self): + return DataLoader(torch.ones(3, 1), batch_size=1) + + for asked in (False, True): # --trainer.benchmark + seen.clear() + trainer = L.Trainer(accelerator="cpu", max_epochs=1, benchmark=asked, logger=False, + callbacks=[pc.CudnnFirstStepTuning()], enable_checkpointing=False, + enable_progress_bar=False, enable_model_summary=False) + with warnings.catch_warnings(), contextlib.redirect_stderr(io.StringIO()): + warnings.simplefilter("ignore") + trainer.fit(Probe()) + with self.subTest(asked=asked): + self.assertEqual(seen, [True, asked, asked]) + self.assertEqual(torch.backends.cudnn.benchmark, asked) + + +class CheckpointSavingTests(unittest.TestCase): + """Checkpoints are written whole, last.ckpt every epoch; a missing metric isn't fatal.""" + + def setUp(self) -> None: + self.pc = _piper_compat() + d = tempfile.TemporaryDirectory() + self.addCleanup(d.cleanup) + self.root = Path(d.name) + + def _files(self) -> list[str]: + return sorted(p.name for p in self.root.iterdir()) + + def test_last_checkpoint_is_saved_every_epoch(self) -> None: + import torch + with _checkpoint_patches(self.pc): + model, trainer = _toy_training(self.root, epochs=3) + _, _, err = _captured(lambda: trainer.fit(model)) + last = torch.load(self.root / "last.ckpt", map_location="cpu", weights_only=False) + self.assertEqual((last["epoch"], last["global_step"]), (2, 12)) + self.assertEqual(self._files(), ["epoch=0-val_mel=1.0000.ckpt", "last.ckpt"]) + self.assertEqual(err.count("val_mos was not logged"), 1) + + def test_a_failed_save_keeps_the_previous_checkpoint(self) -> None: + import torch + from lightning.fabric.plugins.io import torch_io + + class Unsaveable: + def __reduce__(self): + raise RuntimeError("disk full") + + path = self.root / "last.ckpt" + path.write_bytes(b"previous") + with _checkpoint_patches(self.pc): + with self.assertRaises(RuntimeError): + torch_io._atomic_save({"weights": Unsaveable()}, path) + self.assertEqual((path.read_bytes(), self._files()), (b"previous", ["last.ckpt"])) + with patch.object(self.pc.training_control, "replace_file", + side_effect=PermissionError(errno.EACCES, "in use")), \ + self.assertRaises(PermissionError): # Lightning's own save ignored this + torch_io._atomic_save({"step": 2}, path) + self.assertEqual((path.read_bytes(), self._files()), (b"previous", ["last.ckpt"])) + torch_io._atomic_save({"step": 3}, path) + self.assertEqual(torch.load(path, weights_only=False), {"step": 3}) + self.assertEqual(self._files(), ["last.ckpt"]) + + def test_an_old_checkpoint_open_elsewhere_is_deleted_once_closed(self) -> None: + from lightning.fabric.plugins.io import torch_io + attempts = [] + + def busy_once(checkpoint_io, path): + attempts.append(path) + if len(attempts) == 1: + raise PermissionError(errno.EACCES, "in use") + + with patch.object(torch_io.TorchCheckpointIO, "remove_checkpoint", busy_once), \ + _checkpoint_patches(self.pc): + torch_io.TorchCheckpointIO().remove_checkpoint("old.ckpt") + self.assertEqual(attempts, ["old.ckpt", "old.ckpt"]) + + +class StopOnRequestTests(unittest.TestCase): + """Training saves its progress and stops when asked to, or when its launcher or console goes.""" + + def setUp(self) -> None: + self.pc = _piper_compat() + d = tempfile.TemporaryDirectory() + self.addCleanup(d.cleanup) + self.stop_file = Path(d.name) / ".stop_training" + self.checkpoints = Path(d.name) / "checkpoints" + + def _fit(self, model, trainer, stopper, console=None) -> tuple[int, str]: + rc, out, _ = _captured(lambda: self.pc._fit(lambda: trainer.fit(model), stopper, console)) + return rc, out + + def _last(self) -> dict: + import torch + return torch.load(self.checkpoints / "last.ckpt", map_location="cpu", weights_only=False) + + def _resume(self, *, epochs: int = 2): + """The model of a run resumed from last.ckpt, once that has finished.""" + model, trainer = _toy_training(self.checkpoints, epochs=epochs, resumable=True) + _captured(lambda: trainer.fit(model, ckpt_path=self.checkpoints / "last.ckpt")) + return model + + def _stop_at(self, stopper, epoch: int, batch: int): + def on_step(model, batch_idx): + if (model.current_epoch, batch_idx) == (epoch, batch): + self.stop_file.touch() + stopper._next_poll = 0.0 # look now, not within the second + return on_step + + def test_stop_mid_epoch_saves_the_epoch_so_far_and_resumes_with_its_next_batch(self) -> None: + stopper = self.pc.StopOnRequest(str(self.stop_file)) + with _checkpoint_patches(self.pc): + model, trainer = _toy_training( + self.checkpoints, epochs=2, resumable=True, on_step=self._stop_at(stopper, 0, 1), + callbacks=[stopper, self.pc.ProgressSaver(stopper)]) + rc, out = self._fit(model, trainer, stopper) + self.assertEqual((rc, model.steps), (EXIT_STOPPED, [(0, 0), (0, 1)])) + self.assertIn("Stop requested", out) + self.assertIn("Saved the training progress: epoch 0, batch 2/4.", out) + last = self._last() + self.assertEqual((last["epoch"], last["global_step"]), (0, 2)) + self.assertEqual(self.pc._epoch_progress(last), + {"batches_done": 2, "batches_per_epoch": 4}) + self.assertEqual([p.name for p in self.checkpoints.iterdir()], ["last.ckpt"]) + resumed = self._resume() + self.assertEqual(resumed.steps, [(0, 2), (0, 3), (1, 0), (1, 1), (1, 2), (1, 3)]) + self.assertCountEqual(model.items + resumed.items[:2], range(4)) # each clip once + epoch_1 = self.pc.LengthBucketBatchSampler([3, 1, 4, 2], 1, drop_last=False).batches(1) + self.assertEqual(resumed.items[2:], [clip for batch in epoch_1 for clip in batch]) + self.assertTrue((self.checkpoints / "epoch=0-val_mel=1.0000.ckpt").exists()) # validated + + def test_stop_at_the_end_of_an_epoch_validates_and_saves_it(self) -> None: + from lightning.pytorch.callbacks import Callback + stop_file = self.stop_file + + class StopDuringValidation(Callback): + def on_validation_start(self, trainer, pl_module): + stop_file.touch() + + for when in ("at its last batch", "during its validation"): + with self.subTest(when), _checkpoint_patches(self.pc): + stop_file.unlink(missing_ok=True) + shutil.rmtree(self.checkpoints, ignore_errors=True) + stopper = self.pc.StopOnRequest(str(stop_file)) + extra = [StopDuringValidation()] if "validation" in when else [] + model, trainer = _toy_training( + self.checkpoints, epochs=3, resumable=True, + on_step=None if extra else self._stop_at(stopper, 0, 3), + callbacks=[stopper, self.pc.ProgressSaver(stopper), *extra]) + rc, out = self._fit(model, trainer, stopper) + self.assertEqual((rc, model.steps), (EXIT_STOPPED, [(0, i) for i in range(4)])) + self.assertNotIn("Saved the training progress", out) + last = self._last() + self.assertEqual((last["epoch"], last["global_step"]), (0, 4)) + self.assertEqual(self.pc._epoch_progress(last), + {"batches_done": 4, "batches_per_epoch": 4}) + self.assertTrue((self.checkpoints / "epoch=0-val_mel=1.0000.ckpt").exists()) + self.assertEqual(self._resume().steps, [(1, i) for i in range(4)]) + + def test_a_crash_loses_only_the_training_since_the_last_save(self) -> None: + class Clock: # every look at piper_compat's clock is an hour later + now = 0.0 + + def __getattr__(self, name): + return getattr(time, name) + + def monotonic(self): + Clock.now += 3600 + return Clock.now + + def on_step(model, batch_idx): + if (model.current_epoch, batch_idx) == (1, 2): + raise RuntimeError("crash") + + out = io.StringIO() + with _checkpoint_patches(self.pc), patch.object(self.pc, "time", Clock()): + model, trainer = _toy_training(self.checkpoints, epochs=2, resumable=True, + callbacks=[self.pc.ProgressSaver()], on_step=on_step) + with self.assertRaisesRegex(RuntimeError, "crash"), contextlib.redirect_stdout(out), \ + contextlib.redirect_stderr(io.StringIO()), warnings.catch_warnings(): + warnings.simplefilter("ignore") + trainer.fit(model) + # Saved after every step but an epoch's last: its checkpoint comes after validation. + self.assertEqual(re.findall(r"progress: epoch (\d), batch (\d)/4", out.getvalue()), + [("0", "1"), ("0", "2"), ("0", "3"), ("1", "1"), ("1", "2")]) + self.assertEqual(self._last()["global_step"], 6) + with _checkpoint_patches(self.pc): + resumed = self._resume() + self.assertEqual(resumed.steps, [(1, 2), (1, 3)]) + self.assertCountEqual(model.items[4:6] + resumed.items, range(4)) + + def test_closing_the_console_window_saves_before_windows_ends_training(self) -> None: + console = self.pc.training_control.ConsoleClose(timeout=30) + handled, handlers = [], [] + last = self.checkpoints / "last.ckpt" + + def close_window(): # what Windows does, in a thread of its own + handled.append(console.handle(2)) + handled.append(last.exists()) # Windows ends the process once the handler returns + + def on_step(model, batch_idx): + if batch_idx == 1: + handler = threading.Thread(target=close_window) + handler.start() + handlers.append(handler) + self.assertTrue(console.closing.wait(10)) + + with _checkpoint_patches(self.pc): + model, trainer = _toy_training(self.checkpoints, epochs=2, resumable=True, + callbacks=[self.pc.ProgressSaver(console=console)], + on_step=on_step) + rc, out = self._fit(model, trainer, None, console) + handlers[0].join(10) + self.assertEqual(handled, [True, True]) + self.assertEqual((rc, model.steps), (EXIT_STOPPED, [(0, 0), (0, 1)])) + self.assertNotIn("Saved the training progress", out) # the console is going + self.assertEqual(self._last()["global_step"], 2) + self.assertEqual(self._resume().steps[:2], [(0, 2), (0, 3)]) + + def test_closing_the_console_during_a_due_save_ends_training_after_that_save(self) -> None: + from lightning.pytorch.callbacks import Callback + console = self.pc.training_control.ConsoleClose(timeout=30) + handled, handlers, saves = [], [], [] + last = self.checkpoints / "last.ckpt" + + def close_window(): + handled.append(console.handle(2)) + handled.append(last.exists()) + + class CloseWhileSaving(Callback): + def on_save_checkpoint(self, trainer, pl_module, checkpoint): + saves.append(trainer.global_step) + if not handlers: + handlers.append(threading.Thread(target=close_window)) + handlers[0].start() + assert console.closing.wait(10) + + class Clock: # every look is an hour later: a save is due after every step + now = 0.0 + + def __getattr__(self, name): + return getattr(time, name) + + def monotonic(self): + Clock.now += 3600 + return Clock.now + + with _checkpoint_patches(self.pc), patch.object(self.pc, "time", Clock()): + model, trainer = _toy_training( + self.checkpoints, epochs=2, resumable=True, + callbacks=[self.pc.ProgressSaver(console=console), CloseWhileSaving()]) + rc, out = self._fit(model, trainer, None, console) + handlers[0].join(10) + self.assertEqual(handled, [True, True]) + # Not another step and save: Windows allows only seconds. + self.assertEqual((rc, model.steps, saves), (EXIT_STOPPED, [(0, 0)], [1])) + self.assertNotIn("Saved the training progress", out) + self.assertEqual(self._last()["global_step"], 1) + + def test_stop_requested_before_training_starts(self) -> None: + self.stop_file.touch() + stopper = self.pc.StopOnRequest(str(self.stop_file)) + with _checkpoint_patches(self.pc): + model, trainer = _toy_training(self.checkpoints, epochs=3, callbacks=[stopper]) + rc, _ = self._fit(model, trainer, stopper) + self.assertEqual((rc, model.steps), (EXIT_STOPPED, [])) + self.assertEqual(list(self.checkpoints.glob("*.ckpt")), []) + + def test_looks_for_the_stop_file_at_most_once_a_second(self) -> None: + stopper = self.pc.StopOnRequest(str(self.stop_file)) + self.assertFalse(stopper.requested()) + self.stop_file.touch() + self.assertFalse(stopper.requested()) # looked less than a second ago + with contextlib.redirect_stdout(io.StringIO()): + self.assertTrue(stopper.requested(force=True)) + self.assertEqual(stopper.reason, "stop requested") + + def test_training_stops_when_its_launcher_has_gone(self) -> None: + me = self.pc.training_control.this_process() + self.assertFalse(self.pc.StopOnRequest(launcher=(me["pid"], me["start"])).requested()) + stopper = self.pc.StopOnRequest(launcher=_exited_process()) + with contextlib.redirect_stdout(io.StringIO()) as out: + self.assertTrue(stopper.requested()) + with self.assertRaises(self.pc.StopTraining): + stopper.raise_if_requested() + self.assertIn("has exited", out.getvalue()) + + def test_from_environment(self) -> None: + tc = self.pc.training_control + with patch.dict(os.environ, {tc.STOP_FILE_ENV: "", tc.LAUNCHER_ENV: ""}): + self.assertIsNone(self.pc.StopOnRequest.from_environment()) + with patch.dict(os.environ, {tc.STOP_FILE_ENV: "stop", tc.LAUNCHER_ENV: "12:345"}): + stopper = self.pc.StopOnRequest.from_environment() + self.assertEqual((stopper.stop_file, stopper.launcher), ("stop", (12, 345))) + + def test_exit_codes(self) -> None: + def raises(exc): + def main(): + raise exc + return main + + stopped = self.pc.StopOnRequest() + stopped.reason = "stop requested" + for main, stopper, expected in ( + (lambda: None, None, 0), (lambda: None, self.pc.StopOnRequest(), 0), + (lambda: None, stopped, EXIT_STOPPED), + (raises(self.pc.StopTraining("stop requested")), None, EXIT_STOPPED), + (raises(self.pc.TrainingDivergedError("NaN")), None, EXIT_DIVERGED)): + with self.subTest(expected=expected, stopper=stopper): + self.assertEqual(_captured(functools.partial(self.pc._fit, main, stopper))[0], + expected) + with self.assertRaises(ValueError): # a crash: the launcher restarts training + self.pc._fit(raises(ValueError("bug"))) + err = io.StringIO() + with contextlib.redirect_stderr(err), self.assertRaises(self.pc._OUT_OF_MEMORY): + self.pc._fit(raises(self.pc._OUT_OF_MEMORY("CUDA out of memory"))) + self.assertIn("smaller --batch-size", err.getvalue()) + closed = SimpleNamespace(closing=threading.Event()) # ended at an epoch's end + closed.closing.set() + self.assertEqual(self.pc._fit(lambda: None, None, closed), EXIT_STOPPED) + + +class ProgressSaverTests(unittest.TestCase): + """When the progress within an epoch is saved.""" + + def setUp(self) -> None: + self.pc = _piper_compat() + + def _trainer(self, *, batches: float = 10, keeps_last: bool = True): + """``(trainer, its save of last.ckpt)``.""" + save = MagicMock() + callback = SimpleNamespace(save_last=keeps_last, _save_last_checkpoint=save, + _monitor_candidates=lambda trainer: {}) + return SimpleNamespace(num_training_batches=batches, current_epoch=3, should_stop=False, + checkpoint_callbacks=[callback]), save + + def test_every_few_minutes_and_before_an_epoch_ends(self) -> None: + saver = self.pc.ProgressSaver(minutes=10) + trainer, save = self._trainer() + batch_end = functools.partial(saver.on_train_batch_end, trainer, None, None, None) + saver._last_save = time.monotonic() - 400 + with contextlib.redirect_stdout(io.StringIO()) as out: + batch_end(5) + save.assert_not_called() + batch_end(8) # half due: validation and the epoch's checkpoint come next + self.assertEqual(save.call_count, 1) + saver._last_save = time.monotonic() - 601 + batch_end(9) # the epoch's last batch: never + self.assertEqual(save.call_count, 1) + batch_end(0) + self.assertEqual(save.call_count, 2) + self.assertIn("epoch 3, batch 9/10", out.getvalue()) + self.assertFalse(trainer.should_stop) + for minutes in (0, -5): + with self.subTest(minutes=minutes): + saver = self.pc.ProgressSaver(minutes=minutes) + saver._last_save = -math.inf + trainer, save = self._trainer() + saver.on_train_batch_end(trainer, None, None, None, 5) + save.assert_not_called() + + def test_stop_without_a_way_to_save_mid_epoch(self) -> None: + stopper = self.pc.StopOnRequest() + stopper.reason = "stop requested" + for batches, keeps_last in ((10, False), (math.inf, True)): # inf: an iterable dataset + with self.subTest(batches=batches, keeps_last=keeps_last): + saver = self.pc.ProgressSaver(stopper) + trainer, save = self._trainer(batches=batches, keeps_last=keeps_last) + with contextlib.redirect_stdout(io.StringIO()) as out: + saver.on_train_batch_end(trainer, None, None, None, 2) + saver.on_train_batch_end(trainer, None, None, None, 3) + save.assert_not_called() + self.assertTrue(trainer.should_stop) # validates and saves at the epoch's end + self.assertEqual(out.getvalue().count("progress within an epoch is not saved"), + 0 if keeps_last else 1) + + def test_a_closing_console_is_let_go_even_if_saving_fails(self) -> None: + console = SimpleNamespace(closing=threading.Event(), reason="closed", done=MagicMock()) + console.closing.set() + saver = self.pc.ProgressSaver(console=console) + trainer, save = self._trainer() + save.side_effect = OSError("disk full") + with self.assertRaises(OSError): + saver.on_train_batch_end(trainer, None, None, None, 2) + console.done.assert_called_once_with() + + def test_minutes_from_the_environment(self) -> None: + for value, expected in (("5", 5.0), ("0.5", 0.5), ("0", 0.0), ("-1", 0.0), + ("nan", 0.0), ("inf", 0.0), ("ten", 10.0), (None, 10.0)): + with self.subTest(value), patch.dict(os.environ): + os.environ.pop("JAVIS_PIPER_PROGRESS_MINUTES", None) + if value is not None: + os.environ["JAVIS_PIPER_PROGRESS_MINUTES"] = value + self.assertEqual(self.pc._progress_minutes(), expected) + + +class CacheRobustnessTests(unittest.TestCase): + """A damaged preprocessing cache file is rebuilt or passed over, never fatal.""" + + def setUp(self) -> None: + self.pc = _piper_compat() + d = tempfile.TemporaryDirectory() + self.addCleanup(d.cleanup) + self.root = Path(d.name) + + def test_piper_uses_the_safe_cache_functions(self) -> None: + import torch + from piper.train.vits import dataset + self.assertIsInstance(dataset.torch, self.pc._CacheTorch) + self.assertIs(dataset.torch.LongTensor, torch.LongTensor) + + def test_cache_files_are_written_whole(self) -> None: + import torch + cache, path = self.pc._CacheTorch(), self.root / "a.pt" + written = MagicMock() # stands in for StopOnRequest.raise_if_requested + with patch.object(self.pc, "_on_cache_write", written): + cache.save(torch.ones(2), path) + written.assert_called_once_with() + self.assertTrue(torch.equal(cache.load(path), torch.ones(2))) + + class Unsaveable: + def __reduce__(self): + raise RuntimeError("disk full") + + with self.assertRaises(RuntimeError): + cache.save([Unsaveable()], self.root / "b.pt") + self.assertEqual([p.name for p in self.root.iterdir()], ["a.pt"]) + + def test_a_damaged_cache_file_is_deleted_to_be_rebuilt(self) -> None: + import torch + cache, path = self.pc._CacheTorch(), self.root / "a.pt" + torch.save(torch.ones(1000), path) + _truncate(path) + with self.assertRaises(self.pc.CorruptCacheFile): + cache.load(path) + self.assertFalse(path.exists()) + with self.assertRaises(FileNotFoundError): # missing isn't damaged + cache.load(path) + + def test_what_counts_as_damaged_content(self) -> None: + import pickle + damaged = (RuntimeError("PytorchStreamReader failed reading zip archive"), EOFError(), + pickle.UnpicklingError("invalid load key"), + UnicodeDecodeError("utf-8", b"\xff", 0, 1, "invalid start byte"), + OSError(errno.EINVAL, "Invalid argument")) + not_damaged = (FileNotFoundError(errno.ENOENT, "missing"), + PermissionError(errno.EACCES, "in use"), MemoryError(), + RuntimeError("CUDA out of memory"), + RuntimeError("DefaultCPUAllocator: not enough memory"), KeyboardInterrupt()) + for exc in damaged + not_damaged: + with self.subTest(exc=exc): + self.assertEqual(self.pc._unreadable_content(exc), exc in damaged) + + def test_unreadable_utterances_are_replaced_by_readable_ones(self) -> None: + pc, unreadable = self.pc, {2, 3} + + class Utterances: + def __getitem__(self, i): + if i in unreadable: + raise pc.CorruptCacheFile(i) if i == 2 else FileNotFoundError(i) + return i + + part = pc.RobustSubset(Utterances(), [0, 1, 2, 3, 4]) + with patch.object(pc, "_unreadable_utterances", 0), \ + contextlib.redirect_stderr(io.StringIO()) as err: + self.assertEqual([part[i] for i in range(5)], [0, 1, 4, 4, 4]) + self.assertEqual(part.__getitems__([2, 0]), [4, 0]) + self.assertEqual(part[-1], 4) + with self.assertRaises(IndexError): # ends piper's loop over the test split + part[5] + unreadable.update(range(5)) + with self.assertRaises(pc.CorruptCacheFile): # nothing readable: the first error + part[2] + self.assertIn("used another utterance", err.getvalue()) + + def test_data_module_rebuilds_damaged_files(self) -> None: + from torch.utils.data import Subset + pc, root, attempts = self.pc, self.root, [] + (root / "x.pt.partial").write_bytes(b"cut short") # from a run killed while writing + + class DataModule: + cache_dir = root + damaged = 2 + + def prepare_data(self): + attempts.append(sorted(p.name for p in root.iterdir())) + if len(attempts) <= self.damaged: + raise pc.CorruptCacheFile(f"damaged file {len(attempts)}") + return "prepared" + + def setup(self, stage): + self.train_dataset = Subset(range(4), [0, 1]) + self.val_dataset = Subset(range(4), [2]) + self.test_dataset = None + + setup = DataModule.setup + pc._robust_dataset(DataModule) + pc._robust_dataset(DataModule) + self.assertIs(DataModule.setup.__wrapped__, setup) # patched once + dm = DataModule() + with contextlib.redirect_stderr(io.StringIO()) as err: + self.assertEqual(dm.prepare_data(), "prepared") + self.assertEqual(attempts, [[]] * 3) + self.assertEqual(err.getvalue().count("WARNING: damaged file"), 2) + dm.setup("fit") + self.assertIsInstance(dm.train_dataset, pc.RobustSubset) + self.assertIsInstance(dm.val_dataset, pc.RobustSubset) + self.assertEqual((dm.train_dataset.indices, dm.val_dataset.indices), ([0, 1], [2])) + self.assertIsNone(dm.test_dataset) + + attempts.clear() + DataModule.damaged = 100 + with contextlib.redirect_stderr(io.StringIO()), self.assertRaises(pc.CorruptCacheFile): + dm.prepare_data() + self.assertEqual(len(attempts), 50) + + class TranscriptGateTests(unittest.TestCase): def test_unreliable_transcripts_are_rejected(self) -> None: from javis.voices.audiobook_ingest import _transcript_ok