first commit
This commit is contained in:
@@ -0,0 +1,462 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import multiprocessing
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
import traceback
|
||||
import uuid
|
||||
from multiprocessing.connection import Connection
|
||||
from typing import Any, Callable
|
||||
|
||||
|
||||
MAX_NEW_TOKENS = 384
|
||||
DEFAULT_SYNTHESIS_TIMEOUT_SECONDS = 120.0
|
||||
DEFAULT_STARTUP_TIMEOUT_SECONDS = 480.0
|
||||
# 启动失败后再次拉起 worker 的最小间隔：避免"启动超时→立即重启→再超时"的死循环
|
||||
# 在系统高负载时持续加载 torch/CUDA，进一步加剧卡顿。
|
||||
STARTUP_FAILURE_BACKOFF_SECONDS = 60.0
|
||||
DEFAULT_CPU_THREADS = 4
|
||||
DEFAULT_CPU_AFFINITY_COUNT = 8
|
||||
DEFAULT_PROCESS_PRIORITY = "below_normal"
|
||||
|
||||
|
||||
class FasterQwenWorkerError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
class FasterQwenWorkerTimeout(FasterQwenWorkerError):
|
||||
pass
|
||||
|
||||
|
||||
def _generation_kwargs(settings: dict[str, Any], text: str) -> dict[str, Any]:
|
||||
return {
|
||||
"text": text,
|
||||
"language": str(settings.get("language") or "Chinese"),
|
||||
"non_streaming_mode": bool(settings.get("non_streaming_mode", True)),
|
||||
"max_new_tokens": MAX_NEW_TOKENS,
|
||||
}
|
||||
|
||||
|
||||
def _bounded_int(value: Any, default: int, *, minimum: int, maximum: int) -> int:
|
||||
try:
|
||||
parsed = int(value)
|
||||
except (TypeError, ValueError):
|
||||
parsed = default
|
||||
return max(minimum, min(maximum, parsed))
|
||||
|
||||
|
||||
def _configure_worker_environment(settings: dict[str, Any]) -> int:
|
||||
cpu_threads = _bounded_int(
|
||||
settings.get("cpu_threads"),
|
||||
DEFAULT_CPU_THREADS,
|
||||
minimum=1,
|
||||
maximum=8,
|
||||
)
|
||||
thread_value = str(cpu_threads)
|
||||
for name in (
|
||||
"OMP_NUM_THREADS",
|
||||
"MKL_NUM_THREADS",
|
||||
"OPENBLAS_NUM_THREADS",
|
||||
"NUMEXPR_NUM_THREADS",
|
||||
"VECLIB_MAXIMUM_THREADS",
|
||||
"BLIS_NUM_THREADS",
|
||||
):
|
||||
os.environ[name] = thread_value
|
||||
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
||||
return cpu_threads
|
||||
|
||||
|
||||
def _configure_torch_threads(torch_module: Any, cpu_threads: int) -> None:
|
||||
torch_module.set_num_threads(cpu_threads)
|
||||
try:
|
||||
torch_module.set_num_interop_threads(1)
|
||||
except RuntimeError:
|
||||
# PyTorch only allows setting interop threads before parallel work starts.
|
||||
pass
|
||||
|
||||
|
||||
def _apply_worker_process_limits(settings: dict[str, Any]) -> dict[str, Any]:
|
||||
cpu_count = max(1, int(os.cpu_count() or 1))
|
||||
affinity_count = _bounded_int(
|
||||
settings.get("cpu_affinity_count"),
|
||||
DEFAULT_CPU_AFFINITY_COUNT,
|
||||
minimum=0,
|
||||
maximum=min(cpu_count, 63),
|
||||
)
|
||||
priority = str(settings.get("process_priority") or DEFAULT_PROCESS_PRIORITY).strip().lower()
|
||||
applied_affinity = 0
|
||||
applied_priority = "default"
|
||||
|
||||
if os.name == "nt":
|
||||
try:
|
||||
import ctypes
|
||||
from ctypes import wintypes
|
||||
|
||||
kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
|
||||
kernel32.GetCurrentProcess.restype = wintypes.HANDLE
|
||||
kernel32.SetPriorityClass.argtypes = [wintypes.HANDLE, wintypes.DWORD]
|
||||
kernel32.SetPriorityClass.restype = wintypes.BOOL
|
||||
kernel32.SetProcessAffinityMask.argtypes = [wintypes.HANDLE, ctypes.c_size_t]
|
||||
kernel32.SetProcessAffinityMask.restype = wintypes.BOOL
|
||||
process_handle = kernel32.GetCurrentProcess()
|
||||
priority_classes = {
|
||||
"idle": 0x00000040,
|
||||
"below_normal": 0x00004000,
|
||||
"normal": 0x00000020,
|
||||
}
|
||||
priority_class = priority_classes.get(priority, priority_classes[DEFAULT_PROCESS_PRIORITY])
|
||||
if kernel32.SetPriorityClass(process_handle, priority_class):
|
||||
applied_priority = priority if priority in priority_classes else DEFAULT_PROCESS_PRIORITY
|
||||
if affinity_count > 0:
|
||||
affinity_mask = (1 << affinity_count) - 1
|
||||
if kernel32.SetProcessAffinityMask(process_handle, affinity_mask):
|
||||
applied_affinity = affinity_count
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return {
|
||||
"cpu_affinity_count": applied_affinity,
|
||||
"process_priority": applied_priority,
|
||||
}
|
||||
|
||||
|
||||
def _load_runtime(settings: dict[str, Any]) -> dict[str, Any]:
|
||||
device = str(settings.get("device") or "cuda")
|
||||
if device == "cpu" and "CUDA_VISIBLE_DEVICES" not in os.environ:
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = ""
|
||||
|
||||
cpu_threads = _configure_worker_environment(settings)
|
||||
import torch
|
||||
_configure_torch_threads(torch, cpu_threads)
|
||||
from faster_qwen3_tts import FasterQwen3TTS
|
||||
|
||||
if device == "cpu":
|
||||
torch.cuda.is_available = lambda: False
|
||||
|
||||
load_kwargs: dict[str, Any] = {}
|
||||
if device == "cpu":
|
||||
load_kwargs["device"] = "cpu"
|
||||
model = FasterQwen3TTS.from_pretrained(
|
||||
str(settings.get("model_name_or_path") or "Qwen/Qwen3-TTS-12Hz-0.6B-Base"),
|
||||
**load_kwargs,
|
||||
)
|
||||
|
||||
# voice_clone_prompt 不再预计算：当前 faster_qwen3_tts 的 FasterQwen3TTS 没有
|
||||
# create_voice_clone_prompt 方法，预计算只会失败。改为每次合成时在
|
||||
# _synthesize_wav 里直接传 ref_audio/ref_text/xvec_only 参数。
|
||||
voice_clone_prompt = None
|
||||
|
||||
return {
|
||||
"model": model,
|
||||
"torch": torch,
|
||||
"voice_clone_prompt": voice_clone_prompt,
|
||||
"settings": settings,
|
||||
"cpu_threads": cpu_threads,
|
||||
}
|
||||
|
||||
|
||||
def _synthesize_wav(runtime: dict[str, Any], text: str) -> bytes:
|
||||
import soundfile as sf
|
||||
|
||||
model = runtime["model"]
|
||||
torch = runtime["torch"]
|
||||
settings = runtime["settings"]
|
||||
safe_text = str(text or "").strip()[:80] or "欢迎来到直播间。"
|
||||
kwargs = _generation_kwargs(settings, safe_text)
|
||||
voice_clone_prompt = runtime.get("voice_clone_prompt")
|
||||
if voice_clone_prompt is not None:
|
||||
kwargs["voice_clone_prompt"] = voice_clone_prompt
|
||||
else:
|
||||
ref_audio = str(settings.get("ref_audio") or "")
|
||||
if not ref_audio:
|
||||
raise RuntimeError("Faster-Qwen3-TTS requires ref_audio")
|
||||
kwargs.update({
|
||||
"ref_audio": ref_audio,
|
||||
"ref_text": str(settings.get("ref_text") or "") or None,
|
||||
"xvec_only": bool(settings.get("xvec_only", True)),
|
||||
"append_silence": bool(settings.get("append_silence", True)),
|
||||
})
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
torch.backends.cudnn.enabled = False
|
||||
with torch.inference_mode():
|
||||
wavs, sample_rate = model.generate_voice_clone(**kwargs)
|
||||
|
||||
output = io.BytesIO()
|
||||
audio = wavs[0]
|
||||
if isinstance(audio, torch.Tensor):
|
||||
audio = audio.cpu().numpy()
|
||||
sf.write(output, audio, sample_rate, format="WAV")
|
||||
return output.getvalue()
|
||||
|
||||
|
||||
def faster_qwen_worker_main(connection: Connection, settings: dict[str, Any]) -> None:
|
||||
try:
|
||||
cpu_threads = _configure_worker_environment(settings)
|
||||
process_limits = _apply_worker_process_limits(settings)
|
||||
load_started = time.monotonic()
|
||||
runtime = _load_runtime(settings)
|
||||
load_ms = int((time.monotonic() - load_started) * 1000)
|
||||
|
||||
warmup_started = time.monotonic()
|
||||
_synthesize_wav(runtime, "系统启动")
|
||||
warmup_ms = int((time.monotonic() - warmup_started) * 1000)
|
||||
connection.send({
|
||||
"type": "ready",
|
||||
"pid": os.getpid(),
|
||||
"load_ms": load_ms,
|
||||
"warmup_ms": warmup_ms,
|
||||
"max_new_tokens": MAX_NEW_TOKENS,
|
||||
"cpu_threads": cpu_threads,
|
||||
**process_limits,
|
||||
})
|
||||
except BaseException as exc:
|
||||
try:
|
||||
connection.send({
|
||||
"type": "startup_error",
|
||||
"error_type": type(exc).__name__,
|
||||
"error": str(exc),
|
||||
"traceback": traceback.format_exc(),
|
||||
})
|
||||
except Exception:
|
||||
pass
|
||||
connection.close()
|
||||
return
|
||||
|
||||
while True:
|
||||
try:
|
||||
message = connection.recv()
|
||||
except (EOFError, OSError):
|
||||
break
|
||||
if not isinstance(message, dict):
|
||||
continue
|
||||
command = str(message.get("command") or "")
|
||||
if command == "stop":
|
||||
break
|
||||
if command != "synthesize":
|
||||
continue
|
||||
|
||||
request_id = str(message.get("request_id") or "")
|
||||
started = time.monotonic()
|
||||
try:
|
||||
audio = _synthesize_wav(runtime, str(message.get("text") or ""))
|
||||
connection.send({
|
||||
"type": "result",
|
||||
"request_id": request_id,
|
||||
"audio": audio,
|
||||
"duration_ms": int((time.monotonic() - started) * 1000),
|
||||
"bytes": len(audio),
|
||||
"max_new_tokens": MAX_NEW_TOKENS,
|
||||
})
|
||||
except BaseException as exc:
|
||||
try:
|
||||
connection.send({
|
||||
"type": "error",
|
||||
"request_id": request_id,
|
||||
"error_type": type(exc).__name__,
|
||||
"error": str(exc),
|
||||
"traceback": traceback.format_exc(),
|
||||
})
|
||||
except Exception:
|
||||
break
|
||||
connection.close()
|
||||
|
||||
|
||||
class FasterQwenWorkerClient:
|
||||
def __init__(
|
||||
self,
|
||||
settings: dict[str, Any],
|
||||
logger=None,
|
||||
*,
|
||||
synthesis_timeout_seconds: float = DEFAULT_SYNTHESIS_TIMEOUT_SECONDS,
|
||||
startup_timeout_seconds: float = DEFAULT_STARTUP_TIMEOUT_SECONDS,
|
||||
context=None,
|
||||
process_target: Callable[..., None] | None = None,
|
||||
):
|
||||
self.settings = dict(settings)
|
||||
self.logger = logger
|
||||
self.synthesis_timeout_seconds = max(1.0, float(synthesis_timeout_seconds))
|
||||
self.startup_timeout_seconds = max(10.0, float(startup_timeout_seconds))
|
||||
self._context = context or multiprocessing.get_context("spawn")
|
||||
self._process_target = process_target or faster_qwen_worker_main
|
||||
self._lock = threading.RLock()
|
||||
self._process = None
|
||||
self._connection = None
|
||||
self._worker_pid = 0
|
||||
self._next_start_after = 0.0
|
||||
|
||||
@property
|
||||
def worker_pid(self) -> int:
|
||||
return int(self._worker_pid or 0)
|
||||
|
||||
def _log(self, level: str, message: str, *args) -> None:
|
||||
if self.logger:
|
||||
getattr(self.logger, level)(message, *args)
|
||||
|
||||
def _is_alive_locked(self) -> bool:
|
||||
return bool(self._process is not None and self._process.is_alive())
|
||||
|
||||
def ensure_ready(self) -> dict[str, Any]:
|
||||
with self._lock:
|
||||
if self._is_alive_locked() and self._connection is not None:
|
||||
return {"pid": self.worker_pid, "reused": True}
|
||||
return self._start_worker_locked()
|
||||
|
||||
def _start_worker_locked(self) -> dict[str, Any]:
|
||||
now = time.monotonic()
|
||||
if now < self._next_start_after:
|
||||
wait = int(self._next_start_after - now)
|
||||
raise FasterQwenWorkerError(
|
||||
f"TTS worker 启动退避中，距上次启动失败不足 {int(STARTUP_FAILURE_BACKOFF_SECONDS)} 秒，"
|
||||
f"约 {wait} 秒后可重试"
|
||||
)
|
||||
self._terminate_worker_locked("replace_stale_worker", graceful=False)
|
||||
parent_connection, child_connection = self._context.Pipe(duplex=True)
|
||||
process = self._context.Process(
|
||||
target=self._process_target,
|
||||
args=(child_connection, self.settings),
|
||||
name="FasterQwen3TTSWorker",
|
||||
daemon=True,
|
||||
)
|
||||
process.start()
|
||||
try:
|
||||
child_connection.close()
|
||||
except Exception:
|
||||
pass
|
||||
self._process = process
|
||||
self._connection = parent_connection
|
||||
self._worker_pid = int(getattr(process, "pid", 0) or 0)
|
||||
self._log("info", "[FasterQwenTTS] worker 已启动, pid=%s，正在加载和预热", self.worker_pid)
|
||||
|
||||
if not parent_connection.poll(self.startup_timeout_seconds):
|
||||
self._next_start_after = time.monotonic() + STARTUP_FAILURE_BACKOFF_SECONDS
|
||||
self._terminate_worker_locked("startup_timeout", graceful=False)
|
||||
raise FasterQwenWorkerError(
|
||||
f"Faster-Qwen3-TTS worker startup exceeded {self.startup_timeout_seconds:.0f}s"
|
||||
)
|
||||
try:
|
||||
message = parent_connection.recv()
|
||||
except (EOFError, OSError) as exc:
|
||||
self._next_start_after = time.monotonic() + STARTUP_FAILURE_BACKOFF_SECONDS
|
||||
self._terminate_worker_locked("startup_connection_closed", graceful=False)
|
||||
raise FasterQwenWorkerError("Faster-Qwen3-TTS worker exited during startup") from exc
|
||||
if not isinstance(message, dict) or message.get("type") != "ready":
|
||||
if isinstance(message, dict):
|
||||
error = str(message.get("error") or message.get("error_type") or "unknown startup error")
|
||||
else:
|
||||
error = "invalid startup response"
|
||||
self._next_start_after = time.monotonic() + STARTUP_FAILURE_BACKOFF_SECONDS
|
||||
self._terminate_worker_locked("startup_error", graceful=False)
|
||||
raise FasterQwenWorkerError(error)
|
||||
self._next_start_after = 0.0
|
||||
self._log(
|
||||
"info",
|
||||
"[FasterQwenTTS] worker 预热完成, pid=%s load=%sms warmup=%sms max_new_tokens=%s",
|
||||
self.worker_pid,
|
||||
message.get("load_ms"),
|
||||
message.get("warmup_ms"),
|
||||
message.get("max_new_tokens"),
|
||||
)
|
||||
self._log(
|
||||
"info",
|
||||
"[FasterQwenTTS] worker 资源限制: cpu_threads=%s affinity=%s priority=%s",
|
||||
message.get("cpu_threads"),
|
||||
message.get("cpu_affinity_count"),
|
||||
message.get("process_priority"),
|
||||
)
|
||||
return message
|
||||
|
||||
def synthesize(self, text: str) -> tuple[bytes, dict[str, Any]]:
|
||||
with self._lock:
|
||||
self.ensure_ready()
|
||||
request_id = uuid.uuid4().hex
|
||||
connection = self._connection
|
||||
try:
|
||||
connection.send({
|
||||
"command": "synthesize",
|
||||
"request_id": request_id,
|
||||
"text": str(text or ""),
|
||||
})
|
||||
except (BrokenPipeError, EOFError, OSError) as exc:
|
||||
self._restart_after_failure_locked("send_failed")
|
||||
raise FasterQwenWorkerError("Faster-Qwen3-TTS worker connection failed") from exc
|
||||
|
||||
if not connection.poll(self.synthesis_timeout_seconds):
|
||||
self._log(
|
||||
"error",
|
||||
"[FasterQwenTTS] 单次合成超过 %.0f 秒，强制终止 worker pid=%s",
|
||||
self.synthesis_timeout_seconds,
|
||||
self.worker_pid,
|
||||
)
|
||||
restart_error = self._restart_after_failure_locked("synthesis_timeout")
|
||||
suffix = f"; restart failed: {restart_error}" if restart_error else ""
|
||||
raise FasterQwenWorkerTimeout(
|
||||
f"Faster-Qwen3-TTS synthesis exceeded {self.synthesis_timeout_seconds:.0f}s{suffix}"
|
||||
)
|
||||
try:
|
||||
message = connection.recv()
|
||||
except (EOFError, OSError) as exc:
|
||||
self._restart_after_failure_locked("worker_exited")
|
||||
raise FasterQwenWorkerError("Faster-Qwen3-TTS worker exited during synthesis") from exc
|
||||
|
||||
if not isinstance(message, dict) or message.get("request_id") != request_id:
|
||||
self._restart_after_failure_locked("invalid_response")
|
||||
raise FasterQwenWorkerError("Faster-Qwen3-TTS worker returned an invalid response")
|
||||
if message.get("type") == "error":
|
||||
error = str(message.get("error") or message.get("error_type") or "synthesis failed")
|
||||
lowered = error.lower()
|
||||
if "cuda" in lowered or "out of memory" in lowered or "device-side" in lowered:
|
||||
self._restart_after_failure_locked("cuda_error")
|
||||
raise FasterQwenWorkerError(error)
|
||||
if message.get("type") != "result" or not isinstance(message.get("audio"), bytes):
|
||||
self._restart_after_failure_locked("invalid_result")
|
||||
raise FasterQwenWorkerError("Faster-Qwen3-TTS worker returned no audio")
|
||||
return message["audio"], message
|
||||
|
||||
def _restart_after_failure_locked(self, reason: str) -> str:
|
||||
self._terminate_worker_locked(reason, graceful=False)
|
||||
try:
|
||||
self._start_worker_locked()
|
||||
return ""
|
||||
except Exception as exc:
|
||||
self._log("error", "[FasterQwenTTS] worker 自动重启失败: %s", exc)
|
||||
return str(exc)
|
||||
|
||||
def _terminate_worker_locked(self, reason: str, *, graceful: bool) -> None:
|
||||
process = self._process
|
||||
connection = self._connection
|
||||
self._process = None
|
||||
self._connection = None
|
||||
self._worker_pid = 0
|
||||
if process is None:
|
||||
if connection is not None:
|
||||
try:
|
||||
connection.close()
|
||||
except Exception:
|
||||
pass
|
||||
return
|
||||
|
||||
if graceful and process.is_alive() and connection is not None:
|
||||
try:
|
||||
connection.send({"command": "stop"})
|
||||
process.join(timeout=3.0)
|
||||
except Exception:
|
||||
pass
|
||||
if process.is_alive():
|
||||
self._log("warning", "[FasterQwenTTS] 终止 worker, reason=%s pid=%s", reason, process.pid)
|
||||
process.terminate()
|
||||
process.join(timeout=10.0)
|
||||
if process.is_alive():
|
||||
process.kill()
|
||||
process.join(timeout=5.0)
|
||||
if connection is not None:
|
||||
try:
|
||||
connection.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def close(self) -> None:
|
||||
with self._lock:
|
||||
self._terminate_worker_locked("shutdown", graceful=True)
|
||||
Reference in New Issue
Block a user