wehub-resource-sync
411 lines
14 KiB
Python
411 lines
14 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""
|
|
STS (Speech-to-Speech) engine for oMLX.
|
|
|
|
This module provides an engine for audio processing (speech enhancement,
|
|
speech-to-speech conversion, source separation) using mlx-audio.
|
|
Unlike LLM engines, STS engines don't support streaming or chat completion.
|
|
mlx-audio is imported lazily inside start() to avoid module-level import errors
|
|
when mlx-audio is not installed.
|
|
|
|
Supported model families (mlx-audio >=0.4.0):
|
|
- DeepFilterNet: speech enhancement / noise removal
|
|
- MossFormer2: speech enhancement
|
|
- SAMAudio: text-guided sound/speech separation
|
|
- LFM2.5-Audio: multimodal speech-to-speech generation
|
|
"""
|
|
|
|
import asyncio
|
|
import gc
|
|
import logging
|
|
import os
|
|
import tempfile
|
|
from typing import Any, Dict, Optional
|
|
|
|
import mlx.core as mx
|
|
import numpy as np
|
|
|
|
from ..engine_core import get_mlx_executor
|
|
from .audio_utils import DEFAULT_SAMPLE_RATE as _DEFAULT_SAMPLE_RATE
|
|
from .audio_utils import audio_to_wav_bytes as _audio_to_wav_bytes
|
|
from .base import BaseNonStreamingEngine
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
# Maps config.json model_type / architecture values to STS engine families.
|
|
# This is the authoritative source — name-based matching is only a fallback.
|
|
_CONFIG_TYPE_TO_FAMILY: Dict[str, str] = {
|
|
# model_type values (from config.json or mlx-audio dir names)
|
|
"deepfilternet": "deepfilternet",
|
|
"mossformer2_se": "mossformer2",
|
|
"sam_audio": "sam_audio",
|
|
"lfm_audio": "lfm2",
|
|
"lfm2_audio": "lfm2",
|
|
"lfm2": "lfm2",
|
|
# architecture values
|
|
"DeepFilterNetModel": "deepfilternet",
|
|
"MossFormer2SEModel": "mossformer2",
|
|
"SAMAudio": "sam_audio",
|
|
"LFM2AudioModel": "lfm2",
|
|
}
|
|
|
|
|
|
def _detect_sts_family(model_name: str, config_model_type: str = "") -> str:
|
|
"""Detect STS model family, preferring config.json over name guessing.
|
|
|
|
Args:
|
|
model_name: HuggingFace repo name or local path
|
|
config_model_type: Raw model_type from config.json (passed via EngineEntry)
|
|
|
|
Returns one of: "deepfilternet", "mossformer2", "sam_audio", "lfm2", "generic"
|
|
"""
|
|
# 1) config.json model_type — most reliable signal
|
|
if config_model_type:
|
|
family = _CONFIG_TYPE_TO_FAMILY.get(config_model_type.lower())
|
|
if family:
|
|
return family
|
|
|
|
# 2) Try to read architectures from config.json on disk
|
|
config_path = os.path.join(model_name, "config.json")
|
|
if os.path.isfile(config_path):
|
|
try:
|
|
import json
|
|
with open(config_path) as f:
|
|
cfg = json.load(f)
|
|
for arch in cfg.get("architectures", []):
|
|
family = _CONFIG_TYPE_TO_FAMILY.get(arch)
|
|
if family:
|
|
return family
|
|
mt = cfg.get("model_type", "")
|
|
family = _CONFIG_TYPE_TO_FAMILY.get(mt.lower())
|
|
if family:
|
|
return family
|
|
except (OSError, ValueError):
|
|
pass
|
|
|
|
# 3) Fallback: name-based heuristic (tightened patterns)
|
|
name_lower = model_name.lower()
|
|
if "deepfilter" in name_lower:
|
|
return "deepfilternet"
|
|
if "mossformer" in name_lower:
|
|
return "mossformer2"
|
|
# Require "sam-audio" or "sam_audio" (not bare "sam")
|
|
if "sam-audio" in name_lower or "sam_audio" in name_lower:
|
|
return "sam_audio"
|
|
# Require "lfm2" or "lfm-audio" or "lfm_audio" (not bare "lfm")
|
|
if "lfm2" in name_lower or "lfm-audio" in name_lower or "lfm_audio" in name_lower:
|
|
return "lfm2"
|
|
return "generic"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Family-specific loaders
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _load_deepfilternet(model_name: str):
|
|
"""Load a DeepFilterNet model."""
|
|
from mlx_audio.sts.models.deepfilternet import DeepFilterNetModel
|
|
|
|
# DeepFilterNet-mlx uses subfolder='v3' by default; pass model_name_or_path
|
|
return DeepFilterNetModel.from_pretrained(model_name_or_path=model_name)
|
|
|
|
|
|
def _load_mossformer2(model_name: str):
|
|
"""Load a MossFormer2 SE model."""
|
|
from mlx_audio.sts.models.mossformer2_se import MossFormer2SEModel
|
|
|
|
return MossFormer2SEModel.from_pretrained(model_name)
|
|
|
|
|
|
def _load_sam_audio(model_name: str):
|
|
"""Load a SAMAudio model."""
|
|
from mlx_audio.sts.models.sam_audio import SAMAudio
|
|
|
|
return SAMAudio.from_pretrained(model_name)
|
|
|
|
|
|
def _load_lfm2(model_name: str):
|
|
"""Load LFM2.5-Audio model and processor."""
|
|
from mlx_audio.sts.models.lfm_audio import LFM2AudioModel, LFM2AudioProcessor
|
|
|
|
model = LFM2AudioModel.from_pretrained(model_name)
|
|
processor = LFM2AudioProcessor.from_pretrained(model_name)
|
|
return model, processor
|
|
|
|
|
|
_FAMILY_LOADERS = {
|
|
"deepfilternet": _load_deepfilternet,
|
|
"mossformer2": _load_mossformer2,
|
|
"sam_audio": _load_sam_audio,
|
|
"lfm2": _load_lfm2,
|
|
}
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Family-specific processors
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _process_deepfilternet(model, audio_path: str, **kwargs) -> bytes:
|
|
"""Enhance audio with DeepFilterNet (writes to temp file, returns WAV bytes)."""
|
|
fd, out_path = tempfile.mkstemp(suffix=".wav")
|
|
os.close(fd)
|
|
try:
|
|
model.enhance_file(str(audio_path), out_path)
|
|
with open(out_path, "rb") as f:
|
|
return f.read()
|
|
finally:
|
|
if os.path.exists(out_path):
|
|
os.unlink(out_path)
|
|
|
|
|
|
def _process_mossformer2(model, audio_path: str, **kwargs) -> bytes:
|
|
"""Enhance audio with MossFormer2 SE."""
|
|
enhanced = model.enhance(str(audio_path))
|
|
sample_rate = getattr(model.config, "sample_rate", 48000)
|
|
return _audio_to_wav_bytes(enhanced, int(sample_rate))
|
|
|
|
|
|
def _process_sam_audio(model, audio_path: str, **kwargs) -> bytes:
|
|
"""Separate audio with SAMAudio (text-guided)."""
|
|
descriptions = kwargs.get("descriptions", ["speech"])
|
|
result = model.separate(
|
|
audios=[str(audio_path)],
|
|
descriptions=descriptions,
|
|
)
|
|
# SeparationResult.target is List[mx.array] in batch mode
|
|
target = result.target[0] if isinstance(result.target, list) else result.target
|
|
sample_rate = getattr(
|
|
getattr(model, "config", None), "sample_rate", _DEFAULT_SAMPLE_RATE
|
|
)
|
|
return _audio_to_wav_bytes(target, int(sample_rate))
|
|
|
|
|
|
def _process_lfm2(model_and_processor, audio_path: str, **kwargs) -> bytes:
|
|
"""Run speech-to-speech generation with LFM2.5-Audio."""
|
|
from mlx_audio.sts.models.lfm_audio import ChatState, LFMModality
|
|
|
|
model, processor = model_and_processor
|
|
|
|
# Load audio
|
|
from mlx_audio import audio_io
|
|
|
|
audio_np, sr = audio_io.read(str(audio_path))
|
|
audio_mx = mx.array(audio_np.flatten(), dtype=mx.float32)
|
|
|
|
# Build chat state: user turn with audio, then start assistant turn
|
|
chat_state = ChatState(processor)
|
|
chat_state.new_turn("user")
|
|
chat_state.add_audio(audio_mx, sample_rate=sr)
|
|
chat_state.end_turn()
|
|
chat_state.new_turn("assistant")
|
|
|
|
# Generation parameters
|
|
max_new_tokens = kwargs.get("max_new_tokens", 512)
|
|
temperature = kwargs.get("temperature", 0.7)
|
|
audio_temperature = kwargs.get("audio_temperature", 0.8)
|
|
|
|
# Collect audio output frames (each is shape (num_codebooks,) or (1, num_codebooks))
|
|
audio_frames = []
|
|
for token, modality in model.generate_from_chat_state(
|
|
chat_state,
|
|
max_new_tokens=max_new_tokens,
|
|
temperature=temperature,
|
|
audio_temperature=audio_temperature,
|
|
):
|
|
if modality == LFMModality.AUDIO_OUT:
|
|
audio_frames.append(token)
|
|
|
|
if not audio_frames:
|
|
# No audio generated — return silence
|
|
return _audio_to_wav_bytes(np.zeros(1600, dtype=np.float32), _DEFAULT_SAMPLE_RATE)
|
|
|
|
# Stack frames: each frame is (num_codebooks,) -> stack to (T, num_codebooks)
|
|
# then transpose to (num_codebooks, T) for decode_audio
|
|
codes = mx.stack(audio_frames, axis=0) # (T, num_codebooks) or (T, 1, num_codebooks)
|
|
if codes.ndim == 3:
|
|
codes = codes.squeeze(1) # (T, num_codebooks)
|
|
codes = codes.transpose(1, 0) # (num_codebooks, T)
|
|
codes = codes[None, :, :] # (1, num_codebooks, T)
|
|
waveform = processor.decode_audio(codes)
|
|
|
|
# Output at 24kHz (LFM2 default)
|
|
sample_rate = getattr(
|
|
getattr(model, "config", None), "sample_rate", _DEFAULT_SAMPLE_RATE
|
|
)
|
|
return _audio_to_wav_bytes(waveform, int(sample_rate))
|
|
|
|
|
|
_FAMILY_PROCESSORS = {
|
|
"deepfilternet": _process_deepfilternet,
|
|
"mossformer2": _process_mossformer2,
|
|
"sam_audio": _process_sam_audio,
|
|
"lfm2": _process_lfm2,
|
|
}
|
|
|
|
|
|
class STSEngine(BaseNonStreamingEngine):
|
|
"""
|
|
Engine for speech-to-speech / audio processing (Speech-to-Speech).
|
|
|
|
This engine wraps mlx-audio STS models and provides async methods
|
|
for integration with the oMLX server. Each model family has its own
|
|
loading and processing logic since mlx-audio STS models do not share
|
|
a unified utils API (unlike STT/TTS).
|
|
|
|
Supported families:
|
|
- deepfilternet: speech enhancement / noise removal
|
|
- mossformer2: speech enhancement
|
|
- sam_audio: text-guided sound separation
|
|
- lfm2: multimodal speech-to-speech generation
|
|
"""
|
|
|
|
def __init__(self, model_name: str, config_model_type: str = "", **kwargs):
|
|
"""
|
|
Initialize the STS engine.
|
|
|
|
Args:
|
|
model_name: HuggingFace model name or local path
|
|
config_model_type: Raw model_type from config.json (passed by
|
|
EnginePool for reliable family detection)
|
|
**kwargs: Additional model-specific parameters
|
|
"""
|
|
super().__init__()
|
|
self._model_name = model_name
|
|
self._model = None # For lfm2, this is (model, processor) tuple
|
|
self._family = _detect_sts_family(model_name, config_model_type)
|
|
self._kwargs = kwargs
|
|
|
|
@property
|
|
def model_name(self) -> str:
|
|
"""Get the model name."""
|
|
return self._model_name
|
|
|
|
async def start(self) -> None:
|
|
"""Start the engine (load model if not loaded).
|
|
|
|
Model loading runs on the global MLX executor to avoid Metal
|
|
command buffer races with concurrent BatchGenerator steps.
|
|
mlx-audio is imported lazily inside the family-specific loaders
|
|
to avoid module-level errors when mlx-audio is not installed.
|
|
"""
|
|
if self._model is not None:
|
|
return
|
|
|
|
family = self._family
|
|
logger.info(f"Starting STS engine: {self._model_name} (family={family})")
|
|
|
|
loader = _FAMILY_LOADERS.get(family)
|
|
if loader is None:
|
|
raise ValueError(
|
|
f"Unsupported STS model family: {family!r}. "
|
|
f"Supported: {sorted(_FAMILY_LOADERS)}"
|
|
)
|
|
|
|
model_name = self._model_name
|
|
|
|
def _load_sync():
|
|
try:
|
|
return loader(model_name)
|
|
except ImportError as exc:
|
|
raise ImportError(
|
|
"mlx-audio is required for STS inference. "
|
|
"Install it with: pip install 'omlx[audio]'"
|
|
) from exc
|
|
|
|
loop = asyncio.get_running_loop()
|
|
self._model = await loop.run_in_executor(get_mlx_executor(), _load_sync)
|
|
logger.info(f"STS engine started: {self._model_name}")
|
|
|
|
async def stop(self) -> None:
|
|
"""Stop the engine and cleanup resources."""
|
|
if self._model is None:
|
|
return
|
|
|
|
logger.info(f"Stopping STS engine: {self._model_name}")
|
|
self._model = None
|
|
|
|
gc.collect()
|
|
loop = asyncio.get_running_loop()
|
|
await loop.run_in_executor(
|
|
get_mlx_executor(), lambda: (mx.synchronize(), mx.clear_cache())
|
|
)
|
|
logger.info(f"STS engine stopped: {self._model_name}")
|
|
|
|
async def process(self, audio_path: str, **kwargs) -> bytes:
|
|
"""
|
|
Process an audio file through the STS model.
|
|
|
|
For speech enhancement models (DeepFilterNet, MossFormer2), this
|
|
enhances / denoises the audio. For SAMAudio, this separates target
|
|
sounds. For LFM2, this runs speech-to-speech generation.
|
|
|
|
Args:
|
|
audio_path: Path to the audio file to process
|
|
**kwargs: Additional model-specific parameters:
|
|
- descriptions (list[str]): Target descriptions for SAMAudio
|
|
(default: ["speech"])
|
|
- max_new_tokens (int): Max tokens for LFM2 (default: 512)
|
|
- temperature (float): Sampling temperature for LFM2
|
|
- audio_temperature (float): Audio sampling temp for LFM2
|
|
|
|
Returns:
|
|
WAV-encoded bytes (RIFF header + 16-bit mono PCM) of processed audio
|
|
"""
|
|
if self._model is None:
|
|
raise RuntimeError("Engine not started. Call start() first.")
|
|
|
|
import time
|
|
|
|
file_size = os.path.getsize(audio_path) if os.path.exists(audio_path) else 0
|
|
logger.info(
|
|
"STS process: model=%s, family=%s, file=%s (%d bytes)",
|
|
self._model_name, self._family,
|
|
os.path.basename(audio_path), file_size,
|
|
)
|
|
|
|
family = self._family
|
|
processor_fn = _FAMILY_PROCESSORS.get(family)
|
|
if processor_fn is None:
|
|
raise ValueError(
|
|
f"Unsupported STS model family: {family!r}. "
|
|
f"Supported: {sorted(_FAMILY_PROCESSORS)}"
|
|
)
|
|
|
|
model = self._model
|
|
t0 = time.monotonic()
|
|
|
|
def _process_sync():
|
|
return processor_fn(model, str(audio_path), **kwargs)
|
|
|
|
activity_id = self._begin_activity(
|
|
"processing audio",
|
|
detail="Processing audio",
|
|
metadata={"file_size_bytes": file_size, "family": family},
|
|
)
|
|
try:
|
|
loop = asyncio.get_running_loop()
|
|
result = await loop.run_in_executor(get_mlx_executor(), _process_sync)
|
|
|
|
elapsed = time.monotonic() - t0
|
|
logger.info(
|
|
"STS process done: model=%s, %.2fs, %d bytes output",
|
|
self._model_name, elapsed, len(result),
|
|
)
|
|
return result
|
|
finally:
|
|
await self._finish_activity(activity_id)
|
|
|
|
def get_stats(self) -> Dict[str, Any]:
|
|
"""Get engine statistics."""
|
|
return {
|
|
"model_name": self._model_name,
|
|
"loaded": self._model is not None,
|
|
"family": self._family,
|
|
}
|
|
|
|
def __repr__(self) -> str:
|
|
s = "running" if self._model is not None else "stopped"
|
|
return f"<STSEngine model={self._model_name} family={self._family} {s}>"
|