505 lines
17 KiB
Python
505 lines
17 KiB
Python
import base64
|
|
import io
|
|
import json
|
|
import os
|
|
import struct
|
|
import sys
|
|
import threading
|
|
import types
|
|
import urllib.request
|
|
import wave
|
|
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
|
from pathlib import Path
|
|
from urllib.parse import parse_qs, urlsplit
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[2]
|
|
if str(REPO_ROOT) not in sys.path:
|
|
sys.path.insert(0, str(REPO_ROOT))
|
|
|
|
from utils.audio_cpp.client import (
|
|
AudioCppClient,
|
|
AudioCppHTTPError,
|
|
AudioCppProtocolError,
|
|
)
|
|
from utils.audio_cpp.catalog import load_catalog
|
|
from utils.audio_cpp.discovery import package_install_path
|
|
from utils.audio_cpp.process import AudioCppServerProcess, normalize_audio_cpp_task
|
|
from utils.audio_cpp.settings import AudioCppSettings
|
|
from utils.audio_cpp import resolver as audio_cpp_resolver
|
|
from utils.audio_cpp.session import (
|
|
audio_cpp_session_statuses,
|
|
close_all_audio_cpp_sessions,
|
|
get_audio_cpp_session,
|
|
)
|
|
|
|
|
|
def _wav_bytes(sample_rate=16000, channels=1, frames=32):
|
|
buffer = io.BytesIO()
|
|
with wave.open(buffer, "wb") as wav_file:
|
|
wav_file.setnchannels(channels)
|
|
wav_file.setsampwidth(2)
|
|
wav_file.setframerate(sample_rate)
|
|
samples = []
|
|
for index in range(frames):
|
|
value = int(16000 * ((index % 4) - 1.5) / 1.5)
|
|
samples.extend([value] * channels)
|
|
wav_file.writeframes(struct.pack(f"<{len(samples)}h", *samples))
|
|
return buffer.getvalue()
|
|
|
|
|
|
class _FakeAudioCppHandler(BaseHTTPRequestHandler):
|
|
server_version = "FakeAudioCpp/0.5.1"
|
|
|
|
def log_message(self, format, *args):
|
|
return None
|
|
|
|
def _json(self, payload, status=200):
|
|
body = json.dumps(payload).encode("utf-8")
|
|
self.send_response(status)
|
|
self.send_header("Content-Type", "application/json")
|
|
self.send_header("Content-Length", str(len(body)))
|
|
self.end_headers()
|
|
self.wfile.write(body)
|
|
|
|
def do_GET(self):
|
|
parsed = urlsplit(self.path)
|
|
if parsed.path == "/health":
|
|
self._json({"status": "ok", "models": 1, "features": ["unload_models"]})
|
|
elif parsed.path == "/v1/models":
|
|
self.server.model_queries += 1
|
|
self._json({
|
|
"object": "list",
|
|
"data": [{
|
|
"id": self.server.model_id,
|
|
"object": "model",
|
|
"family": "pocket_tts",
|
|
"task": "tts",
|
|
"mode": "offline",
|
|
}],
|
|
})
|
|
elif parsed.path == "/v1/audio/voices":
|
|
self.server.last_voice_query = parse_qs(parsed.query)
|
|
self._json({"voices": ["alba", "cosette"]})
|
|
else:
|
|
self._json({"error": {"message": "not found", "type": "not_found"}}, 404)
|
|
|
|
def do_POST(self):
|
|
length = int(self.headers.get("Content-Length", "0"))
|
|
payload = json.loads(self.rfile.read(length).decode("utf-8"))
|
|
self.server.requests.append(payload)
|
|
request = payload.get("request", {})
|
|
text = request.get("text")
|
|
if text == "http-error":
|
|
self._json(
|
|
{"error": {"message": "model is busy", "type": "server_busy"}},
|
|
503,
|
|
)
|
|
return
|
|
|
|
encoded = base64.b64encode(self.server.wav_bytes).decode("ascii")
|
|
if text == "named-only":
|
|
self._json({
|
|
"named_audio_outputs": [{
|
|
"id": "speech",
|
|
"audio": encoded,
|
|
"sample_rate": 16000,
|
|
"channels": 1,
|
|
}],
|
|
"timing": {},
|
|
})
|
|
elif text == "ambiguous":
|
|
self._json({
|
|
"named_audio_outputs": [
|
|
{"id": "left", "audio": encoded},
|
|
{"id": "right", "audio": encoded},
|
|
]
|
|
})
|
|
elif text == "transcript-only":
|
|
self._json({
|
|
"text": "hello world",
|
|
"language": "en",
|
|
"words": [
|
|
{"word": "hello", "start_sample": 0, "end_sample": 8000},
|
|
{"word": "world", "start_sample": 8000, "end_sample": 16000},
|
|
],
|
|
"timing": {"wall_ms": 1.0},
|
|
})
|
|
else:
|
|
self._json({
|
|
"audio": encoded,
|
|
"sample_rate": 16000,
|
|
"channels": 1,
|
|
"timing": {"wall_ms": 1.0},
|
|
})
|
|
|
|
|
|
@pytest.fixture
|
|
def fake_audio_cpp_server():
|
|
server = ThreadingHTTPServer(("127.0.0.1", 0), _FakeAudioCppHandler)
|
|
server.model_id = "pocket"
|
|
server.wav_bytes = _wav_bytes()
|
|
server.requests = []
|
|
server.last_voice_query = None
|
|
server.model_queries = 0
|
|
thread = threading.Thread(target=server.serve_forever, daemon=True)
|
|
thread.start()
|
|
try:
|
|
yield server, f"http://127.0.0.1:{server.server_port}"
|
|
finally:
|
|
server.shutdown()
|
|
server.server_close()
|
|
thread.join(timeout=2)
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def clean_audio_cpp_sessions():
|
|
close_all_audio_cpp_sessions()
|
|
yield
|
|
close_all_audio_cpp_sessions()
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_client_health_models_voices_and_primary_audio(fake_audio_cpp_server):
|
|
server, url = fake_audio_cpp_server
|
|
client = AudioCppClient(url)
|
|
|
|
assert client.health()["status"] == "ok"
|
|
assert client.models()[0]["id"] == "pocket"
|
|
assert client.voices("pocket") == ["alba", "cosette"]
|
|
assert server.last_voice_query == {"model": ["pocket"]}
|
|
assert client.supports_feature("unload_models") is True
|
|
|
|
result = client.run_task("pocket", {"text": "hello", "seed": "42"})
|
|
assert result.sample_rate == 16000
|
|
assert result.channels == 1
|
|
assert result.waveform.shape == (1, 32)
|
|
assert result.waveform.dtype == torch.float32
|
|
assert result.waveform.device.type == "cpu"
|
|
assert server.requests[-1] == {
|
|
"model": "pocket",
|
|
"request": {"text": "hello", "seed": "42"},
|
|
}
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_client_selects_sole_named_audio_and_rejects_ambiguous_output(fake_audio_cpp_server):
|
|
_, url = fake_audio_cpp_server
|
|
client = AudioCppClient(url)
|
|
|
|
result = client.run_task("pocket", {"text": "named-only"})
|
|
assert result.sample_rate == 16000
|
|
assert list(result.named_audio) == ["speech"]
|
|
assert result.waveform.data_ptr() == result.named_audio["speech"].waveform.data_ptr()
|
|
|
|
with pytest.raises(AudioCppProtocolError, match="multiple named audio"):
|
|
client.run_task("pocket", {"text": "ambiguous"})
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_client_accepts_structured_transcript_without_audio(fake_audio_cpp_server):
|
|
_, url = fake_audio_cpp_server
|
|
result = AudioCppClient(url).run_task("asr", {"text": "transcript-only"})
|
|
|
|
assert result.waveform is None
|
|
assert result.sample_rate is None
|
|
assert result.raw["text"] == "hello world"
|
|
assert len(result.raw["words"]) == 2
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_client_surfaces_structured_http_error(fake_audio_cpp_server):
|
|
_, url = fake_audio_cpp_server
|
|
client = AudioCppClient(url)
|
|
|
|
with pytest.raises(AudioCppHTTPError) as captured:
|
|
client.run_task("pocket", {"text": "http-error"})
|
|
|
|
assert captured.value.status == 503
|
|
assert captured.value.error_type == "server_busy"
|
|
assert "model is busy" in str(captured.value)
|
|
|
|
|
|
def _write_fake_server_script(path: Path) -> Path:
|
|
script = r'''
|
|
import argparse
|
|
import base64
|
|
import io
|
|
import json
|
|
import struct
|
|
import wave
|
|
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
|
from urllib.parse import urlsplit
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--config", required=True)
|
|
args = parser.parse_args()
|
|
with open(args.config, "r", encoding="utf-8") as handle:
|
|
config = json.load(handle)
|
|
model_id = config["models"][0]["id"]
|
|
|
|
buffer = io.BytesIO()
|
|
with wave.open(buffer, "wb") as wav_file:
|
|
wav_file.setnchannels(1)
|
|
wav_file.setsampwidth(2)
|
|
wav_file.setframerate(22050)
|
|
wav_file.writeframes(struct.pack("<16h", *range(16)))
|
|
encoded = base64.b64encode(buffer.getvalue()).decode("ascii")
|
|
|
|
class Handler(BaseHTTPRequestHandler):
|
|
def log_message(self, format, *args):
|
|
return None
|
|
def send_json(self, payload):
|
|
body = json.dumps(payload).encode("utf-8")
|
|
self.send_response(200)
|
|
self.send_header("Content-Type", "application/json")
|
|
self.send_header("Content-Length", str(len(body)))
|
|
self.end_headers()
|
|
self.wfile.write(body)
|
|
def do_GET(self):
|
|
path = urlsplit(self.path).path
|
|
if path == "/health":
|
|
self.send_json({"status": "ok", "models": 1})
|
|
elif path == "/v1/models":
|
|
self.send_json({"object": "list", "data": [{"id": model_id}]})
|
|
elif path == "/v1/audio/voices":
|
|
self.send_json({"voices": ["managed"]})
|
|
else:
|
|
self.send_json({})
|
|
def do_POST(self):
|
|
length = int(self.headers.get("Content-Length", "0"))
|
|
json.loads(self.rfile.read(length).decode("utf-8"))
|
|
self.send_json({"audio": encoded, "sample_rate": 22050, "channels": 1})
|
|
|
|
server = ThreadingHTTPServer((config["host"], config["port"]), Handler)
|
|
print("fake audio.cpp ready", flush=True)
|
|
server.serve_forever()
|
|
'''
|
|
path.write_text(script, encoding="utf-8")
|
|
return path
|
|
|
|
|
|
def _owned_config(tmp_path: Path, script: Path):
|
|
model_path = tmp_path / "model"
|
|
model_path.mkdir(exist_ok=True)
|
|
(model_path / "weights.gguf").write_bytes(b"audio-cpp-test-model")
|
|
return {
|
|
"connection_mode": "existing_binary",
|
|
"binary_path": sys.executable,
|
|
"binary_args": [str(script)],
|
|
"model_path": str(model_path),
|
|
"model_id": "managed-pocket",
|
|
"family": "pocket_tts",
|
|
"package_id": "pocket_tts_english_q8_0",
|
|
"task": "tts",
|
|
"backend": "cpu",
|
|
"startup_timeout": 5.0,
|
|
"connect_timeout": 1.0,
|
|
"request_timeout": 5.0,
|
|
"stop_timeout": 2.0,
|
|
}
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_owned_process_writes_secure_config_and_stops_exact_child(tmp_path):
|
|
assert normalize_audio_cpp_task("clone") == "clon"
|
|
assert normalize_audio_cpp_task("voice design") == "vdes"
|
|
script = _write_fake_server_script(tmp_path / "fake_audio_cpp_server.py")
|
|
runtime = AudioCppServerProcess(_owned_config(tmp_path, script)).start()
|
|
process = runtime.process
|
|
config_path = runtime.config_path
|
|
assert process is not None and process.poll() is None
|
|
assert config_path is not None and config_path.exists()
|
|
|
|
payload = json.loads(config_path.read_text(encoding="utf-8"))
|
|
assert payload["host"] == "127.0.0.1"
|
|
assert payload["cors_origins"] == ""
|
|
assert payload["log_request_body"] is False
|
|
assert payload["model_spec_override"].endswith("utils\\audio_cpp\\model_specs") or payload[
|
|
"model_spec_override"
|
|
].endswith("utils/audio_cpp/model_specs")
|
|
assert payload["models"][0]["task"] == "tts"
|
|
assert Path(payload["models"][0]["path"]).is_absolute()
|
|
assert runtime.client.voices("managed-pocket") == ["managed"]
|
|
|
|
runtime.close()
|
|
assert process.poll() is not None
|
|
assert not config_path.exists()
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_external_session_is_keyed_and_never_terminates_server(fake_audio_cpp_server):
|
|
server, url = fake_audio_cpp_server
|
|
config = {
|
|
"connection_mode": "existing_server",
|
|
"server_url": url,
|
|
"model_id": "pocket",
|
|
}
|
|
first = get_audio_cpp_session(config)
|
|
second = get_audio_cpp_session(dict(config))
|
|
assert first is second
|
|
assert first.owned is False
|
|
assert first.model_id == "pocket"
|
|
assert first.model_metadata["family"] == "pocket_tts"
|
|
assert first.task == "tts"
|
|
assert first.voices() == ["alba", "cosette"]
|
|
assert server.model_queries == 1
|
|
|
|
first.close()
|
|
with urllib.request.urlopen(f"{url}/health", timeout=1) as response:
|
|
assert json.load(response)["status"] == "ok"
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_owned_session_restarts_and_reregisters_after_exact_child_exit(tmp_path, monkeypatch):
|
|
script = _write_fake_server_script(tmp_path / "fake_audio_cpp_server.py")
|
|
|
|
fake_management = types.ModuleType("comfy.model_management")
|
|
fake_management.current_loaded_models = []
|
|
fake_management.cleanup_models = lambda: None
|
|
|
|
class LoadedModel:
|
|
def __init__(self, model):
|
|
self.model = model
|
|
|
|
fake_management.LoadedModel = LoadedModel
|
|
monkeypatch.setitem(sys.modules, "comfy.model_management", fake_management)
|
|
comfy_module = sys.modules.get("comfy")
|
|
if comfy_module is not None:
|
|
monkeypatch.setattr(comfy_module, "model_management", fake_management, raising=False)
|
|
|
|
session = get_audio_cpp_session(_owned_config(tmp_path, script))
|
|
assert session.proxy.model_size() >= len(b"audio-cpp-test-model")
|
|
first = session.run({"text": "first"})
|
|
first_runtime = session.process
|
|
first_process = first_runtime.process
|
|
assert first.sample_rate == 22050
|
|
assert len(fake_management.current_loaded_models) == 1
|
|
|
|
first_process.kill()
|
|
first_process.wait(timeout=2)
|
|
second = session.run({"text": "second"})
|
|
second_runtime = session.process
|
|
second_process = second_runtime.process
|
|
assert second.sample_rate == 22050
|
|
assert second_runtime is not first_runtime
|
|
assert second_process is not first_process
|
|
assert len(fake_management.current_loaded_models) == 1
|
|
|
|
session.restart_owned_runtime()
|
|
reset_process = session.process.process
|
|
assert second_process.poll() is not None
|
|
assert reset_process is not second_process
|
|
assert session.run({"text": "after explicit reset"}).sample_rate == 22050
|
|
assert len(fake_management.current_loaded_models) == 1
|
|
|
|
assert session.proxy.partially_unload("cpu", 1) == 0
|
|
tracked_model = fake_management.current_loaded_models[0]
|
|
session.proxy.unpatch_model("cpu")
|
|
assert reset_process.poll() is not None
|
|
assert fake_management.current_loaded_models == [tracked_model]
|
|
fake_management.current_loaded_models.pop(0)
|
|
|
|
session.run({"text": "third"})
|
|
third_process = session.process.process
|
|
assert third_process.poll() is None
|
|
assert len(fake_management.current_loaded_models) == 1
|
|
session.close()
|
|
assert third_process.poll() is not None
|
|
assert len(fake_management.current_loaded_models) == 0
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_session_makes_voice_reference_path_absolute(fake_audio_cpp_server, tmp_path, monkeypatch):
|
|
server, url = fake_audio_cpp_server
|
|
monkeypatch.chdir(tmp_path)
|
|
reference = tmp_path / "voice.wav"
|
|
reference.write_bytes(_wav_bytes())
|
|
session = get_audio_cpp_session({
|
|
"connection_mode": "existing_server",
|
|
"server_url": url,
|
|
"model_id": "pocket",
|
|
})
|
|
|
|
session.run({"text": "absolute path", "voice_ref": "voice.wav"})
|
|
|
|
sent_path = server.requests[-1]["request"]["voice_ref"]
|
|
assert Path(sent_path).is_absolute()
|
|
assert Path(sent_path) == reference
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_session_resolver_discovers_existing_model_root(tmp_path, monkeypatch):
|
|
script = _write_fake_server_script(tmp_path / "fake_audio_cpp_server.py")
|
|
external_root = tmp_path / "existing-models"
|
|
managed_root = tmp_path / "suite-managed-models"
|
|
package = load_catalog().package("pocket_tts_english_q8_0")
|
|
installed = package_install_path(package, external_root)
|
|
for relative_path in package.local_files:
|
|
target = installed / relative_path
|
|
target.parent.mkdir(parents=True, exist_ok=True)
|
|
target.write_bytes(b"installed")
|
|
|
|
monkeypatch.setattr(
|
|
audio_cpp_resolver,
|
|
"load_settings",
|
|
lambda: AudioCppSettings(
|
|
connection_mode="managed",
|
|
model_roots=(str(external_root),),
|
|
managed_model_root=str(managed_root),
|
|
runtime_backend="cpu",
|
|
),
|
|
)
|
|
session = get_audio_cpp_session({
|
|
"connection_mode": "managed",
|
|
"family": "pocket_tts",
|
|
"package_id": package.id,
|
|
"task": "tts",
|
|
"backend": "cpu",
|
|
"binary_path": sys.executable,
|
|
"binary_args": [str(script)],
|
|
"startup_timeout": 5.0,
|
|
"connect_timeout": 1.0,
|
|
"request_timeout": 5.0,
|
|
})
|
|
|
|
result = session.run({"text": "resolved"})
|
|
|
|
assert result.sample_rate == 22050
|
|
assert session.config["model_path"] == str(installed.resolve())
|
|
assert not managed_root.exists()
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_external_session_selects_the_servers_sole_model(fake_audio_cpp_server):
|
|
server, url = fake_audio_cpp_server
|
|
|
|
session = get_audio_cpp_session({
|
|
"connection_mode": "external_server",
|
|
"server_url": url,
|
|
"model_id": "",
|
|
})
|
|
reused = get_audio_cpp_session({
|
|
"connection_mode": "external_server",
|
|
"server_url": url,
|
|
"model_id": "",
|
|
})
|
|
|
|
assert reused is session
|
|
assert session.model_id == "pocket"
|
|
assert session.model_metadata["family"] == "pocket_tts"
|
|
assert session.task == "tts"
|
|
assert server.model_queries == 1
|
|
|
|
before = next(item for item in audio_cpp_session_statuses() if item["model_id"] == "pocket")
|
|
assert before["state"] == "server_ready"
|
|
assert before["owned"] is False
|
|
assert before["endpoint"] == url
|
|
|
|
session.run({"text": "status"})
|
|
after = next(item for item in audio_cpp_session_statuses() if item["model_id"] == "pocket")
|
|
assert after["state"] == "model_ready"
|