627 lines
23 KiB
Python
627 lines
23 KiB
Python
"""Keyed audio.cpp sessions and ComfyUI lifecycle integration."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import atexit
|
|
import ipaddress
|
|
import json
|
|
import os
|
|
import re
|
|
import sys
|
|
import threading
|
|
import time
|
|
import urllib.parse
|
|
import uuid
|
|
from pathlib import Path
|
|
from typing import Any, Dict, Mapping, Optional
|
|
|
|
from .client import (
|
|
AudioCppClient,
|
|
AudioCppConnectionError,
|
|
AudioCppTaskResult,
|
|
AudioCppTimeoutError,
|
|
)
|
|
from .lifecycle import AudioCppRuntimeProxy
|
|
from .process import AudioCppServerProcess, normalize_audio_cpp_task
|
|
from .resolver import resolve_audio_cpp_config
|
|
|
|
|
|
def _warn(message: str, exc: Optional[BaseException] = None) -> None:
|
|
"""Emit diagnostics without assuming a UTF-8 Windows console."""
|
|
text = f"WARNING: {message}"
|
|
if exc is not None:
|
|
text += f": {exc}"
|
|
encoding = getattr(sys.stderr, "encoding", None) or "ascii"
|
|
try:
|
|
text = text.encode(encoding, errors="replace").decode(encoding, errors="replace")
|
|
except LookupError:
|
|
text = text.encode("ascii", errors="replace").decode("ascii")
|
|
print(text, file=sys.stderr)
|
|
|
|
|
|
def _flatten_config(config: Mapping[str, Any]) -> Dict[str, Any]:
|
|
if not isinstance(config, Mapping):
|
|
raise TypeError("audio.cpp config must be a mapping")
|
|
nested = config.get("config")
|
|
flattened: Dict[str, Any] = dict(nested) if isinstance(nested, Mapping) else {}
|
|
flattened.update({key: value for key, value in config.items() if key != "config"})
|
|
return flattened
|
|
|
|
|
|
def _first(config: Mapping[str, Any], *keys: str, default: Any = None) -> Any:
|
|
for key in keys:
|
|
if key in config and config[key] is not None:
|
|
return config[key]
|
|
return default
|
|
|
|
|
|
def _connection_mode(config: Mapping[str, Any]) -> str:
|
|
raw = str(_first(config, "connection_mode", "source", default="auto"))
|
|
normalized = re.sub(r"[^a-z0-9]+", "_", raw.strip().lower()).strip("_")
|
|
external_aliases = {
|
|
"existing_server",
|
|
"external_server",
|
|
"server",
|
|
"remote_server",
|
|
"connect",
|
|
}
|
|
owned_aliases = {
|
|
"owned_process",
|
|
"existing_binary",
|
|
"managed_binary",
|
|
"managed",
|
|
"binary",
|
|
"local_binary",
|
|
}
|
|
if normalized in external_aliases:
|
|
return "external_server"
|
|
if normalized in owned_aliases:
|
|
return "owned_process"
|
|
if normalized not in {"", "auto"}:
|
|
raise ValueError(f"Unsupported audio.cpp connection mode: {raw}")
|
|
endpoint = _first(config, "server_url", "endpoint", "base_url")
|
|
return "external_server" if endpoint else "owned_process"
|
|
|
|
|
|
def _is_loopback_url(url: str) -> bool:
|
|
parsed = urllib.parse.urlsplit(url)
|
|
host = (parsed.hostname or "").strip().lower()
|
|
if host == "localhost":
|
|
return True
|
|
try:
|
|
return ipaddress.ip_address(host).is_loopback
|
|
except ValueError:
|
|
return False
|
|
|
|
|
|
def _canonical_url(url: Any) -> str:
|
|
parsed = urllib.parse.urlsplit(str(url or "").strip())
|
|
if parsed.scheme not in {"http", "https"} or not parsed.netloc:
|
|
raise ValueError(f"Invalid audio.cpp external server URL: {url!r}")
|
|
if parsed.query or parsed.fragment:
|
|
raise ValueError("audio.cpp external server URL must not contain a query or fragment")
|
|
return urllib.parse.urlunsplit(
|
|
(parsed.scheme, parsed.netloc, parsed.path.rstrip("/"), "", "")
|
|
)
|
|
|
|
|
|
def _safe_model_id(config: Mapping[str, Any]) -> str:
|
|
raw = _first(
|
|
config,
|
|
"model_id",
|
|
"server_model_id",
|
|
"package_id",
|
|
default=_first(config, "family", "model_family", default="audio-cpp-model"),
|
|
)
|
|
value = str(raw or "audio-cpp-model").strip()
|
|
cleaned = "".join(char if char.isalnum() or char in "._-" else "-" for char in value)
|
|
return cleaned.strip("-.") or "audio-cpp-model"
|
|
|
|
|
|
def _jsonable(value: Any) -> Any:
|
|
if isinstance(value, Mapping):
|
|
return {str(key): _jsonable(item) for key, item in sorted(value.items(), key=lambda item: str(item[0]))}
|
|
if isinstance(value, (list, tuple)):
|
|
return [_jsonable(item) for item in value]
|
|
if isinstance(value, Path):
|
|
return str(value.expanduser().resolve())
|
|
if isinstance(value, (str, int, float, bool)) or value is None:
|
|
return value
|
|
return repr(value)
|
|
|
|
|
|
def _session_key(config: Mapping[str, Any], mode: str, model_id: str, endpoint: str = "") -> str:
|
|
if mode == "external_server":
|
|
identity = {
|
|
"mode": mode,
|
|
"endpoint": endpoint,
|
|
"model_id": model_id,
|
|
}
|
|
else:
|
|
identity_keys = (
|
|
"binary_path",
|
|
"server_binary",
|
|
"executable_path",
|
|
"audio_cpp_binary",
|
|
"binary_args",
|
|
"model_path",
|
|
"package_path",
|
|
"gguf_path",
|
|
"family",
|
|
"model_family",
|
|
"task",
|
|
"backend",
|
|
"device",
|
|
"device_index",
|
|
"threads",
|
|
"load_options",
|
|
"session_options",
|
|
"default_request_options",
|
|
"config_id",
|
|
"weight_id",
|
|
"model_spec_override",
|
|
"show_server_console",
|
|
)
|
|
identity = {"mode": mode, "model_id": model_id}
|
|
for key in identity_keys:
|
|
if key in config:
|
|
identity[key] = config[key]
|
|
return json.dumps(_jsonable(identity), ensure_ascii=True, sort_keys=True, separators=(",", ":"))
|
|
|
|
|
|
def _normalize_request_paths(request: Mapping[str, Any]) -> Dict[str, Any]:
|
|
normalized = dict(request)
|
|
path_fields = {
|
|
"voice_ref",
|
|
"audio_path",
|
|
"source_audio",
|
|
"target_voice",
|
|
"prosody_ref",
|
|
"style_ref",
|
|
}
|
|
for key in path_fields:
|
|
value = normalized.get(key)
|
|
if isinstance(value, os.PathLike) or (isinstance(value, str) and value.strip()):
|
|
normalized[key] = str(
|
|
Path(os.path.expandvars(os.path.expanduser(str(value)))).resolve()
|
|
)
|
|
return normalized
|
|
|
|
|
|
class AudioCppSession:
|
|
"""Persistent external connection or restartable suite-owned audio.cpp server."""
|
|
|
|
def __init__(
|
|
self,
|
|
config: Mapping[str, Any],
|
|
*,
|
|
owned: bool,
|
|
model_id: str,
|
|
client: Optional[AudioCppClient] = None,
|
|
endpoint: str = "",
|
|
model_metadata: Optional[Mapping[str, Any]] = None,
|
|
) -> None:
|
|
self.config = dict(config)
|
|
self.owned = bool(owned)
|
|
self.model_id = str(model_id)
|
|
self.endpoint = endpoint
|
|
self.model_metadata = dict(model_metadata or {})
|
|
self.family = str(
|
|
_first(
|
|
self.model_metadata,
|
|
"family",
|
|
default=_first(self.config, "family", "model_family", default=""),
|
|
)
|
|
or ""
|
|
)
|
|
self.task = normalize_audio_cpp_task(
|
|
_first(
|
|
self.model_metadata,
|
|
"task",
|
|
default=_first(self.config, "task", default="tts"),
|
|
)
|
|
)
|
|
self.unload_models_supported = False
|
|
self._client = client
|
|
self._process: Optional[AudioCppServerProcess] = None
|
|
self._lock = threading.RLock()
|
|
self._closed = False
|
|
self._model_ready_reported = False
|
|
self.ui_session_id = uuid.uuid4().hex
|
|
self._proxy = AudioCppRuntimeProxy(self) if self.owned else None
|
|
|
|
@property
|
|
def process(self) -> Optional[AudioCppServerProcess]:
|
|
return self._process
|
|
|
|
@property
|
|
def running(self) -> bool:
|
|
if not self.owned:
|
|
return not self._closed
|
|
return self._process is not None and self._process.running
|
|
|
|
@property
|
|
def proxy(self) -> Optional[AudioCppRuntimeProxy]:
|
|
return self._proxy
|
|
|
|
def _probe_opt_in_features(self) -> None:
|
|
if not bool(self.config.get("probe_unload_models", False)) or self._client is None:
|
|
return
|
|
try:
|
|
self.unload_models_supported = self._client.supports_feature("unload_models")
|
|
except Exception as exc:
|
|
_warn("audio.cpp unload_models feature probe failed", exc)
|
|
self.unload_models_supported = False
|
|
|
|
def _start_owned_runtime(self) -> None:
|
|
if not self.owned:
|
|
return
|
|
with _RUNTIME_START_LOCK:
|
|
if self._process is not None and self._process.running:
|
|
return
|
|
_stop_conflicting_audio_cpp_sessions(self)
|
|
if str(self.config.get("backend", "cuda")).lower() in {"cuda", "hip", "auto"}:
|
|
_clear_conflicting_suite_tts_models()
|
|
process_config = dict(self.config)
|
|
process_config["model_id"] = self.model_id
|
|
family = self.family or str(self.config.get("family", "unknown"))
|
|
backend = str(self.config.get("backend", "auto"))
|
|
print(
|
|
f"🚀 audio.cpp: Starting {backend} server for {family} "
|
|
f"('{self.model_id}')..."
|
|
)
|
|
started = time.monotonic()
|
|
process = AudioCppServerProcess(process_config)
|
|
process.start()
|
|
self._process = process
|
|
self._client = process.client
|
|
self.endpoint = process.base_url
|
|
self._probe_opt_in_features()
|
|
if self._proxy is not None:
|
|
self._proxy.register()
|
|
print(
|
|
f"✅ audio.cpp: Server ready at {self.endpoint} "
|
|
f"({time.monotonic() - started:.2f}s); model will load on first request"
|
|
)
|
|
|
|
def _ensure_client(self) -> AudioCppClient:
|
|
if self._closed:
|
|
raise RuntimeError("audio.cpp session is closed")
|
|
if self.owned:
|
|
if self._process is None or not self._process.running or self._client is None:
|
|
if self._proxy is not None:
|
|
self._proxy.unregister()
|
|
if self._process is not None:
|
|
self._process.close()
|
|
self._process = None
|
|
self._client = None
|
|
self._start_owned_runtime()
|
|
if self._client is None:
|
|
raise RuntimeError("audio.cpp session has no HTTP client")
|
|
return self._client
|
|
|
|
def _restart_after_transport_failure(self) -> AudioCppClient:
|
|
if not self.owned:
|
|
raise RuntimeError("Cannot restart an external audio.cpp server")
|
|
self._stop_owned_runtime()
|
|
self._start_owned_runtime()
|
|
if self._client is None:
|
|
raise RuntimeError("audio.cpp owned runtime restart did not create a client")
|
|
return self._client
|
|
|
|
def restart_owned_runtime(self) -> None:
|
|
"""Recreate an owned server when an upstream session cannot be reused safely."""
|
|
if not self.owned:
|
|
raise RuntimeError("Cannot restart an external audio.cpp server")
|
|
with self._lock:
|
|
self._stop_owned_runtime()
|
|
self._start_owned_runtime()
|
|
|
|
def run(self, request: Mapping[str, Any]) -> AudioCppTaskResult:
|
|
if not isinstance(request, Mapping):
|
|
raise TypeError("audio.cpp task request must be a mapping")
|
|
normalized_request = _normalize_request_paths(request)
|
|
timeout = float(_first(self.config, "request_timeout", "request_timeout_seconds", default=600.0))
|
|
with self._lock:
|
|
client = self._ensure_client()
|
|
first_request = not self._model_ready_reported
|
|
started = time.monotonic()
|
|
if first_request:
|
|
action = "Loading model" if self.owned else "Sending first request to model"
|
|
print(f"⏳ audio.cpp: {action} '{self.model_id}'...")
|
|
try:
|
|
result = client.run_task(self.model_id, normalized_request, timeout=timeout)
|
|
except AudioCppConnectionError:
|
|
if not self.owned:
|
|
raise
|
|
client = self._restart_after_transport_failure()
|
|
result = client.run_task(self.model_id, normalized_request, timeout=timeout)
|
|
except AudioCppTimeoutError:
|
|
# A live server may still be executing after the client times out.
|
|
# Only restart when the exact child has actually exited.
|
|
if not self.owned or (self._process is not None and self._process.running):
|
|
raise
|
|
client = self._restart_after_transport_failure()
|
|
result = client.run_task(self.model_id, normalized_request, timeout=timeout)
|
|
if first_request:
|
|
self._model_ready_reported = True
|
|
print(
|
|
f"✅ audio.cpp: Model '{self.model_id}' loaded; first {self.task} request completed "
|
|
f"in {time.monotonic() - started:.2f}s"
|
|
)
|
|
return result
|
|
|
|
def voices(self) -> list[str]:
|
|
timeout = float(_first(self.config, "connect_timeout", "connect_timeout_seconds", default=5.0))
|
|
with self._lock:
|
|
client = self._ensure_client()
|
|
try:
|
|
return client.voices(self.model_id, timeout=timeout)
|
|
except AudioCppConnectionError:
|
|
if not self.owned:
|
|
raise
|
|
return self._restart_after_transport_failure().voices(
|
|
self.model_id, timeout=timeout
|
|
)
|
|
|
|
def _stop_owned_runtime(self, *, unregister: bool = True) -> None:
|
|
if not self.owned:
|
|
return
|
|
with self._lock:
|
|
process = self._process
|
|
self._process = None
|
|
self._client = None
|
|
self._model_ready_reported = False
|
|
if unregister and self._proxy is not None:
|
|
self._proxy.unregister()
|
|
if process is not None:
|
|
process.close()
|
|
|
|
def close(self) -> None:
|
|
with self._lock:
|
|
if self._closed:
|
|
return
|
|
self._closed = True
|
|
if self.owned:
|
|
self._stop_owned_runtime()
|
|
else:
|
|
# External servers are never unloaded, reconfigured, or terminated.
|
|
self._client = None
|
|
|
|
|
|
_SESSIONS: Dict[str, AudioCppSession] = {}
|
|
_SESSIONS_LOCK = threading.RLock()
|
|
_RUNTIME_START_LOCK = threading.RLock()
|
|
|
|
|
|
def _stop_conflicting_audio_cpp_sessions(active: AudioCppSession) -> None:
|
|
if str(active.config.get("backend", "cuda")).lower() not in {"cuda", "hip", "auto"}:
|
|
return
|
|
with _SESSIONS_LOCK:
|
|
conflicts = [
|
|
session
|
|
for session in _SESSIONS.values()
|
|
if session is not active
|
|
and session.owned
|
|
and session.running
|
|
and str(session.config.get("backend", "cuda")).lower()
|
|
in {"cuda", "hip", "auto"}
|
|
]
|
|
for session in conflicts:
|
|
session._stop_owned_runtime()
|
|
|
|
|
|
def _clear_conflicting_suite_tts_models() -> None:
|
|
"""Clear known suite-managed TTS resources only when their modules are already live."""
|
|
interface_module = sys.modules.get("utils.models.unified_model_interface")
|
|
interface = getattr(interface_module, "unified_model_interface", None)
|
|
if interface is not None:
|
|
try:
|
|
isolated = getattr(interface, "_isolated_model_cache", None)
|
|
remover = getattr(interface, "_remove_isolated_model", None)
|
|
if isinstance(isolated, dict) and callable(remover):
|
|
for cache_key in list(isolated):
|
|
if "_tts_" in cache_key:
|
|
remover(cache_key)
|
|
except Exception as exc:
|
|
_warn("Could not clear a conflicting isolated TTS runtime", exc)
|
|
|
|
wrapper_module = sys.modules.get("utils.models.comfyui_model_wrapper")
|
|
manager = getattr(wrapper_module, "tts_model_manager", None)
|
|
cache = getattr(manager, "_model_cache", None)
|
|
remover = getattr(manager, "remove_model", None)
|
|
if isinstance(cache, dict) and callable(remover):
|
|
try:
|
|
for cache_key, wrapper in list(cache.items()):
|
|
model_info = getattr(wrapper, "model_info", None)
|
|
if getattr(model_info, "model_type", None) == "tts":
|
|
remover(cache_key)
|
|
except Exception as exc:
|
|
_warn("Could not clear a conflicting embedded TTS model", exc)
|
|
|
|
|
|
def _external_client(
|
|
config: Mapping[str, Any],
|
|
) -> tuple[AudioCppClient, str, str, Dict[str, Any]]:
|
|
endpoint = _canonical_url(_first(config, "server_url", "endpoint", "base_url"))
|
|
if not _is_loopback_url(endpoint) and not bool(config.get("allow_remote_server", False)):
|
|
raise ValueError(
|
|
"External audio.cpp servers must use loopback by default. "
|
|
"Set allow_remote_server only when transport security and path access are understood."
|
|
)
|
|
client = AudioCppClient(
|
|
endpoint,
|
|
connect_timeout=float(
|
|
_first(config, "connect_timeout", "connect_timeout_seconds", default=5.0)
|
|
),
|
|
request_timeout=float(
|
|
_first(config, "request_timeout", "request_timeout_seconds", default=600.0)
|
|
),
|
|
)
|
|
models = client.models()
|
|
requested_id = _first(config, "model_id", "server_model_id")
|
|
available_models = [dict(item) for item in models if item.get("id") is not None]
|
|
available_ids = [str(item["id"]) for item in available_models]
|
|
if requested_id is not None:
|
|
model_id = str(requested_id)
|
|
if model_id not in available_ids:
|
|
raise ValueError(
|
|
f"External audio.cpp server does not expose model '{model_id}'. "
|
|
f"Available: {', '.join(available_ids) or '(none)'}"
|
|
)
|
|
elif len(available_ids) == 1:
|
|
model_id = available_ids[0]
|
|
elif not available_ids:
|
|
raise ValueError("External audio.cpp server does not expose any configured models")
|
|
else:
|
|
raise ValueError(
|
|
"External audio.cpp server exposes multiple models; select model_id explicitly"
|
|
)
|
|
selected_metadata = next(
|
|
(item for item in available_models if str(item.get("id")) == model_id),
|
|
{"id": model_id},
|
|
)
|
|
return client, endpoint, model_id, selected_metadata
|
|
|
|
|
|
def get_audio_cpp_session(config: Mapping[str, Any]) -> AudioCppSession:
|
|
"""Return a keyed persistent audio.cpp session for a loose engine config dict."""
|
|
flattened = resolve_audio_cpp_config(_flatten_config(config))
|
|
mode = _connection_mode(flattened)
|
|
if mode == "external_server":
|
|
endpoint_hint = _canonical_url(
|
|
_first(flattened, "server_url", "endpoint", "base_url")
|
|
)
|
|
model_hint = _first(flattened, "model_id", "server_model_id")
|
|
auto_key = None
|
|
if isinstance(model_hint, str) and model_hint.strip():
|
|
hinted_key = _session_key(flattened, mode, model_hint.strip(), endpoint_hint)
|
|
with _SESSIONS_LOCK:
|
|
existing = _SESSIONS.get(hinted_key)
|
|
if existing is not None and not existing._closed:
|
|
return existing
|
|
else:
|
|
# An omitted model id means "select the server's sole model". Cache
|
|
# that resolution per endpoint so every text chunk does not repeat
|
|
# /v1/models before reaching the already-persistent session.
|
|
auto_key = _session_key(flattened, mode, "", endpoint_hint)
|
|
with _SESSIONS_LOCK:
|
|
existing = _SESSIONS.get(auto_key)
|
|
if existing is not None and not existing._closed:
|
|
return existing
|
|
client, endpoint, model_id, model_metadata = _external_client(flattened)
|
|
key = _session_key(flattened, mode, model_id, endpoint)
|
|
with _SESSIONS_LOCK:
|
|
existing = _SESSIONS.get(key)
|
|
if existing is not None and not existing._closed:
|
|
if auto_key is not None:
|
|
_SESSIONS[auto_key] = existing
|
|
return existing
|
|
session = AudioCppSession(
|
|
flattened,
|
|
owned=False,
|
|
model_id=model_id,
|
|
client=client,
|
|
endpoint=endpoint,
|
|
model_metadata=model_metadata,
|
|
)
|
|
session._probe_opt_in_features()
|
|
_SESSIONS[key] = session
|
|
if auto_key is not None:
|
|
_SESSIONS[auto_key] = session
|
|
return session
|
|
|
|
model_id = _safe_model_id(flattened)
|
|
key = _session_key(flattened, mode, model_id)
|
|
with _SESSIONS_LOCK:
|
|
existing = _SESSIONS.get(key)
|
|
if existing is not None and not existing._closed:
|
|
return existing
|
|
session = AudioCppSession(flattened, owned=True, model_id=model_id)
|
|
_SESSIONS[key] = session
|
|
return session
|
|
|
|
|
|
def close_all_audio_cpp_sessions() -> None:
|
|
with _SESSIONS_LOCK:
|
|
sessions = list({id(session): session for session in _SESSIONS.values()}.values())
|
|
_SESSIONS.clear()
|
|
for session in sessions:
|
|
try:
|
|
session.close()
|
|
except Exception as exc:
|
|
_warn("Failed to close audio.cpp session during shutdown", exc)
|
|
|
|
|
|
def audio_cpp_session_statuses() -> list[Dict[str, Any]]:
|
|
"""Return a path-free, side-effect-free snapshot for the frontend indicator."""
|
|
with _SESSIONS_LOCK:
|
|
sessions = list({id(session): session for session in _SESSIONS.values()}.values())
|
|
statuses = []
|
|
for session in sessions:
|
|
if session._closed:
|
|
continue
|
|
if session.owned:
|
|
if session.running:
|
|
state = "model_ready" if session._model_ready_reported else "server_ready"
|
|
else:
|
|
state = "configured"
|
|
else:
|
|
state = "model_ready" if session._model_ready_reported else "server_ready"
|
|
statuses.append(
|
|
{
|
|
"session_id": session.ui_session_id,
|
|
"state": state,
|
|
"owned": session.owned,
|
|
"family": session.family,
|
|
"model_id": session.model_id,
|
|
"endpoint": session.endpoint if not session.owned else "",
|
|
"pid": session.process.process.pid if session.owned and session.running else None,
|
|
**_process_memory_status(
|
|
session.process.process.pid if session.owned and session.running else None
|
|
),
|
|
}
|
|
)
|
|
return statuses
|
|
|
|
|
|
def _process_memory_status(pid: Optional[int]) -> Dict[str, Optional[int]]:
|
|
if not pid:
|
|
return {"working_set_bytes": None, "private_bytes": None}
|
|
try:
|
|
import psutil
|
|
|
|
info = psutil.Process(pid).memory_info()
|
|
return {
|
|
"working_set_bytes": int(info.rss),
|
|
"private_bytes": int(getattr(info, "private", info.vms)),
|
|
}
|
|
except Exception:
|
|
return {"working_set_bytes": None, "private_bytes": None}
|
|
|
|
|
|
def stop_owned_audio_cpp_session(session_id: str) -> bool:
|
|
"""Stop, but retain, an exact Suite-owned session for lazy restart."""
|
|
with _SESSIONS_LOCK:
|
|
sessions = list({id(session): session for session in _SESSIONS.values()}.values())
|
|
session = next((item for item in sessions if item.ui_session_id == session_id), None)
|
|
if session is None:
|
|
return False
|
|
if not session.owned:
|
|
raise PermissionError("External audio.cpp servers cannot be stopped by the Suite")
|
|
session._stop_owned_runtime()
|
|
return True
|
|
|
|
|
|
atexit.register(close_all_audio_cpp_sessions)
|
|
|
|
|
|
__all__ = [
|
|
"AudioCppRuntimeProxy",
|
|
"AudioCppSession",
|
|
"audio_cpp_session_statuses",
|
|
"close_all_audio_cpp_sessions",
|
|
"get_audio_cpp_session",
|
|
"stop_owned_audio_cpp_session",
|
|
]
|