Integrate YuE2 as an exclusive Athena mode

This commit is contained in:
Mikei386
2026-09-10 23:53:19 +02:00
parent 5c34afa7fa
commit f2052adb78
14 changed files with 278 additions and 15 deletions
@@ -30,6 +30,7 @@ TTS_LABEL_KEY = "com.mike-ai.tts-worker"
TTS_WORKER = os.environ.get("TTS_WORKER", "qwen3")
MUSIC_LABEL_KEY = "com.mike-ai.music-worker"
MUSIC_WORKER = os.environ.get("MUSIC_WORKER", "").strip()
YUE2_WORKER = os.environ.get("YUE2_WORKER", "").strip()
SEPARATOR_LABEL_KEY = "com.mike-ai.stem-separator"
SEPARATOR_WORKER = os.environ.get("SEPARATOR_WORKER", "").strip()
VOICE_LABEL_KEY = "com.mike-ai.voice-worker"
@@ -38,6 +39,8 @@ VOICE_CHANGE_LABEL_KEY = "com.mike-ai.voice-change-worker"
VOICE_CHANGE_WORKER = os.environ.get("VOICE_CHANGE_WORKER", "").strip()
APPLIO_LABEL_KEY = "com.mike-ai.applio-worker"
APPLIO_WORKER = os.environ.get("APPLIO_WORKER", "").strip()
TRELLIS_LABEL_KEY = "com.mike-ai.trellis-worker"
TRELLIS_WORKER = os.environ.get("TRELLIS_WORKER", "").strip()
LOCK = threading.Lock()
log = logging.getLogger("profile-controller")
@@ -117,6 +120,17 @@ def music_container() -> dict:
return matches[0]
def yue2_container() -> dict:
if not YUE2_WORKER:
raise RuntimeError("YuE2 worker is not configured")
matches = [item for item in labelled_containers(MUSIC_LABEL_KEY)
if item.get("Labels", {}).get(MUSIC_LABEL_KEY) == YUE2_WORKER]
if len(matches) != 1:
raise RuntimeError(
f"expected exactly one YuE2 worker {YUE2_WORKER!r}, found {len(matches)}")
return matches[0]
def separator_container() -> dict:
if not SEPARATOR_WORKER:
raise RuntimeError("stem separator is not configured")
@@ -161,11 +175,27 @@ def applio_container() -> dict:
return matches[0]
def trellis_container() -> dict:
if not TRELLIS_WORKER:
raise RuntimeError("TRELLIS worker is not configured")
matches = [item for item in labelled_containers(TRELLIS_LABEL_KEY)
if item.get("Labels", {}).get(TRELLIS_LABEL_KEY) == TRELLIS_WORKER]
if len(matches) != 1:
raise RuntimeError(
f"expected exactly one TRELLIS worker {TRELLIS_WORKER!r}, found {len(matches)}")
return matches[0]
def stop_music_if_configured() -> None:
if MUSIC_WORKER:
stop_container(music_container(), timeout=30)
def stop_yue2_if_configured() -> None:
if YUE2_WORKER:
stop_container(yue2_container(), timeout=30)
def stop_separator_if_configured() -> None:
if SEPARATOR_WORKER:
stop_container(separator_container(), timeout=30)
@@ -186,6 +216,11 @@ def stop_applio_if_configured() -> None:
stop_container(applio_container(), timeout=30)
def stop_trellis_if_configured() -> None:
if TRELLIS_WORKER:
stop_container(trellis_container(), timeout=30)
def stop_voice_tools(except_kind: str | None = None) -> None:
if except_kind != "voice":
stop_voice_if_configured()
@@ -233,8 +268,10 @@ def set_image_worker(running: bool, kind: str = IMAGE_WORKER) -> dict:
# Qwen3-TTS. TTS is unavailable during this exclusive GPU phase.
stop_container(tts_container(), timeout=30)
stop_music_if_configured()
stop_yue2_if_configured()
stop_separator_if_configured()
stop_voice_tools()
stop_trellis_if_configured()
for other in image_containers():
if other["Id"] != item["Id"]:
stop_container(other, timeout=20)
@@ -262,6 +299,8 @@ def set_music_worker(running: bool) -> dict:
stop_container(tts_container(), timeout=30)
stop_separator_if_configured()
stop_voice_tools()
stop_yue2_if_configured()
stop_trellis_if_configured()
start_container(item)
else:
stop_container(item, timeout=30)
@@ -269,6 +308,27 @@ def set_music_worker(running: bool) -> dict:
"state": "running" if running else "stopped"}
def set_yue2_worker(running: bool) -> dict:
"""Start YuE2 exclusively, or stop it before another mode is loaded."""
with LOCK:
item = yue2_container()
if running:
for profile_item in containers().values():
stop_container(profile_item)
for worker in image_containers():
stop_container(worker, timeout=20)
stop_container(tts_container(), timeout=30)
stop_music_if_configured()
stop_separator_if_configured()
stop_voice_tools()
stop_trellis_if_configured()
start_container(item)
else:
stop_container(item, timeout=30)
return {"yue2_worker": YUE2_WORKER,
"state": "running" if running else "stopped"}
def set_separator_worker(running: bool) -> dict:
"""Start vocal separation exclusively, or stop it before LLM restoration."""
with LOCK:
@@ -280,7 +340,9 @@ def set_separator_worker(running: bool) -> dict:
stop_container(worker, timeout=20)
stop_container(tts_container(), timeout=30)
stop_music_if_configured()
stop_yue2_if_configured()
stop_voice_tools()
stop_trellis_if_configured()
start_container(item)
else:
stop_container(item, timeout=30)
@@ -299,8 +361,10 @@ def set_voice_worker(running: bool) -> dict:
stop_container(worker, timeout=20)
stop_container(tts_container(), timeout=30)
stop_music_if_configured()
stop_yue2_if_configured()
stop_separator_if_configured()
stop_voice_tools("voice")
stop_trellis_if_configured()
start_container(item)
else:
stop_container(item, timeout=30)
@@ -319,8 +383,10 @@ def set_voice_change_worker(running: bool) -> dict:
stop_container(worker, timeout=20)
stop_container(tts_container(), timeout=30)
stop_music_if_configured()
stop_yue2_if_configured()
stop_separator_if_configured()
stop_voice_tools("voicechange")
stop_trellis_if_configured()
start_container(item)
else:
stop_container(item, timeout=30)
@@ -339,8 +405,10 @@ def set_applio_worker(running: bool) -> dict:
stop_container(worker, timeout=20)
stop_container(tts_container(), timeout=30)
stop_music_if_configured()
stop_yue2_if_configured()
stop_separator_if_configured()
stop_voice_tools("applio")
stop_trellis_if_configured()
start_container(item)
else:
stop_container(item, timeout=30)
@@ -348,6 +416,27 @@ def set_applio_worker(running: bool) -> dict:
"state": "running" if running else "stopped"}
def set_trellis_worker(running: bool) -> dict:
"""Start TRELLIS.2 exclusively, or stop it before another mode is loaded."""
with LOCK:
item = trellis_container()
if running:
for profile_item in containers().values():
stop_container(profile_item)
for worker in image_containers():
stop_container(worker, timeout=20)
stop_container(tts_container(), timeout=30)
stop_music_if_configured()
stop_yue2_if_configured()
stop_separator_if_configured()
stop_voice_tools()
start_container(item)
else:
stop_container(item, timeout=30)
return {"trellis_worker": TRELLIS_WORKER,
"state": "running" if running else "stopped"}
def active_profile(items: dict[str, dict] | None = None) -> str | None:
items = items or containers()
active = [name for name, item in items.items() if item.get("State") == "running"]
@@ -364,8 +453,10 @@ def activate(profile: str) -> dict:
for worker in image_containers():
stop_container(worker)
stop_music_if_configured()
stop_yue2_if_configured()
stop_separator_if_configured()
stop_voice_tools()
stop_trellis_if_configured()
start_container(tts_container())
items = containers()
missing = [name for name in ALLOWED if name not in items]
@@ -437,6 +528,13 @@ class Handler(BaseHTTPRequestHandler):
"unhealthy" if "(unhealthy)" in music_status else
"starting" if music.get("State") == "running" else
"stopped")
yue2 = yue2_container() if YUE2_WORKER else {}
yue2_status = yue2.get("Status", "")
yue2_health = ("disabled" if not YUE2_WORKER else
"healthy" if "(healthy)" in yue2_status else
"unhealthy" if "(unhealthy)" in yue2_status else
"starting" if yue2.get("State") == "running" else
"stopped")
separator = separator_container() if SEPARATOR_WORKER else {}
separator_status = separator.get("Status", "")
separator_health = ("disabled" if not SEPARATOR_WORKER else
@@ -465,9 +563,18 @@ class Handler(BaseHTTPRequestHandler):
"unhealthy" if "(unhealthy)" in applio_status else
"starting" if applio.get("State") == "running" else
"stopped")
trellis = trellis_container() if TRELLIS_WORKER else {}
trellis_status = trellis.get("Status", "")
trellis_health = ("disabled" if not TRELLIS_WORKER else
"healthy" if "(healthy)" in trellis_status else
"unhealthy" if "(unhealthy)" in trellis_status else
"starting" if trellis.get("State") == "running" else
"stopped")
self.reply(200, {"active_profile": active_profile(items),
"music_worker": music.get("State", "disabled"),
"music_health": music_health,
"yue2_worker": yue2.get("State", "disabled"),
"yue2_health": yue2_health,
"separator_worker": separator.get("State", "disabled"),
"separator_health": separator_health,
"voice_worker": voice.get("State", "disabled"),
@@ -476,6 +583,8 @@ class Handler(BaseHTTPRequestHandler):
"voice_change_health": voice_change_health,
"applio_worker": applio.get("State", "disabled"),
"applio_health": applio_health,
"trellis_worker": trellis.get("State", "disabled"),
"trellis_health": trellis_health,
"profiles": {name: items.get(name, {}).get(
"State", "missing") for name in ALLOWED}})
except Exception as exc:
@@ -500,6 +609,13 @@ class Handler(BaseHTTPRequestHandler):
log.exception("music worker transition failed")
self.reply(503, {"error": str(exc)})
return
if self.path in {"/workers/yue2/start", "/workers/yue2/stop"}:
try:
self.reply(200, set_yue2_worker(self.path.endswith("/start")))
except Exception as exc:
log.exception("YuE2 worker transition failed")
self.reply(503, {"error": str(exc)})
return
if self.path in {"/workers/separator/start", "/workers/separator/stop"}:
try:
self.reply(200, set_separator_worker(self.path.endswith("/start")))
@@ -528,6 +644,13 @@ class Handler(BaseHTTPRequestHandler):
log.exception("Applio worker transition failed")
self.reply(503, {"error": str(exc)})
return
if self.path in {"/workers/trellis/start", "/workers/trellis/stop"}:
try:
self.reply(200, set_trellis_worker(self.path.endswith("/start")))
except Exception as exc:
log.exception("TRELLIS worker transition failed")
self.reply(503, {"error": str(exc)})
return
worker_paths = {
"/workers/image/start": (IMAGE_WORKER, True),
"/workers/image/stop": (IMAGE_WORKER, False),
@@ -84,6 +84,8 @@ start_proxy 8008 voice-studio:8008
start_proxy 8009 xvc-studio:8009
start_proxy 8011 applio-studio:6969
start_proxy 8012 mikes-applio-ui:8012
start_proxy 8013 trellis-studio:8080
start_proxy 8014 yue2-studio:8014
start_proxy 8202 mcp-athena-operator:8000
start_proxy 9443 portainer:9443