81 lines
2.7 KiB
Python
81 lines
2.7 KiB
Python
"""
|
|
Base abstractions for engine-specific training handlers.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from abc import ABC, abstractmethod
|
|
from typing import Any, Dict, Tuple
|
|
|
|
from utils.config_sanitizer import ConfigSanitizer
|
|
|
|
|
|
def unpack_tts_engine(tts_engine: Any) -> Tuple[str, Dict[str, Any]]:
|
|
"""
|
|
Normalize a TTS_ENGINE value into ``(engine_type, config_dict)``.
|
|
|
|
The suite currently uses two shapes:
|
|
- dict-based engine configs for most engines
|
|
- adapter instances for RVC
|
|
"""
|
|
if hasattr(tts_engine, "engine_type"):
|
|
engine_type = getattr(tts_engine, "engine_type", None)
|
|
config = getattr(tts_engine, "config", {}) or {}
|
|
if not engine_type:
|
|
raise ValueError("TTS engine adapter is missing engine_type")
|
|
if not isinstance(config, dict):
|
|
raise ValueError(f"{engine_type} engine config must be a dict")
|
|
return engine_type, ConfigSanitizer.sanitize(config)
|
|
|
|
if isinstance(tts_engine, dict):
|
|
engine_type = tts_engine.get("engine_type")
|
|
config = tts_engine.get("config", {}) or {}
|
|
if not config:
|
|
config = tts_engine
|
|
if not engine_type:
|
|
raise ValueError("TTS engine is missing engine_type")
|
|
if not isinstance(config, dict):
|
|
raise ValueError(f"{engine_type} engine config must be a dict")
|
|
return engine_type, ConfigSanitizer.sanitize(config)
|
|
|
|
raise ValueError("Invalid TTS_engine input")
|
|
|
|
|
|
class BaseTrainingHandler(ABC):
|
|
"""
|
|
Base class for engine-specific training backends.
|
|
"""
|
|
|
|
engine_type: str = ""
|
|
artifact_type: str = "model"
|
|
|
|
def ensure_engine_type(self, tts_engine: Any) -> Dict[str, Any]:
|
|
engine_type, config = unpack_tts_engine(tts_engine)
|
|
if engine_type != self.engine_type:
|
|
raise ValueError(
|
|
f"This training handler only supports '{self.engine_type}', got '{engine_type}'"
|
|
)
|
|
return config
|
|
|
|
@abstractmethod
|
|
def build_default_training_config(self, tts_engine: Any) -> Dict[str, Any]:
|
|
"""Return a default TRAINING_CONFIG payload for this engine."""
|
|
|
|
@abstractmethod
|
|
def prepare_dataset(self, tts_engine: Any, **kwargs) -> Dict[str, Any]:
|
|
"""Create a TRAINING_DATASET payload for this engine."""
|
|
|
|
@abstractmethod
|
|
def train(
|
|
self,
|
|
tts_engine: Any,
|
|
training_dataset: Dict[str, Any],
|
|
training_config: Dict[str, Any],
|
|
output_name: str = "",
|
|
resume: bool = False,
|
|
overwrite: bool = False,
|
|
continue_from: Any = None,
|
|
node_id: str = "",
|
|
) -> Dict[str, Any]:
|
|
"""Run the engine-specific training job and return TRAINING_ARTIFACTS."""
|