Integrate YuE2 as an exclusive Athena mode
This commit is contained in:
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user