mirror of
https://github.com/calesthio/OpenMontage.git
synced 2026-08-25 17:42:50 +08:00
Merge remote-tracking branch 'origin/main' into codex/repair-pr-482
# Conflicts: # docs/PROVIDERS.md
This commit is contained in:
calesthio
@@ -17,13 +17,14 @@ from tools.base_tool import (
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
from tools.audio.comfyui_music import ComfyUIMusic
|
||||
from tools.graphics.comfyui_image import ComfyUIImage
|
||||
from tools.graphics.image_selector import ImageSelector
|
||||
from tools.tool_registry import ToolRegistry
|
||||
from tools.video.video_selector import VideoSelector
|
||||
from tools.video.comfyui_video import ComfyUIVideo
|
||||
|
||||
TOOLS = [ComfyUIImage, ComfyUIVideo]
|
||||
TOOLS = [ComfyUIImage, ComfyUIVideo, ComfyUIMusic]
|
||||
WORKFLOW_DIR = Path(__file__).resolve().parent.parent.parent / "tools" / "_comfyui" / "workflows"
|
||||
PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent
|
||||
|
||||
@@ -136,6 +137,7 @@ EXPECTED_WORKFLOWS = [
|
||||
"flux2-txt2img.json",
|
||||
"wan22-i2v-4step.json",
|
||||
"wan22-t2v-4step.json",
|
||||
"ace-step-1-t2a.json",
|
||||
]
|
||||
|
||||
|
||||
@@ -282,6 +284,73 @@ class TestClientHelpers:
|
||||
"folder_type": "temp",
|
||||
}
|
||||
|
||||
def test_poll_timeout_carries_prompt_id_for_recovery(self, monkeypatch):
|
||||
from tools._comfyui.client import ComfyUIClient, ComfyUIError
|
||||
|
||||
client = ComfyUIClient("http://comfy.test")
|
||||
monkeypatch.setattr(
|
||||
"tools._comfyui.client.requests.get",
|
||||
lambda *a, **k: type("R", (), {
|
||||
"raise_for_status": lambda self: None,
|
||||
"json": lambda self: {},
|
||||
})(),
|
||||
)
|
||||
monkeypatch.setattr("tools._comfyui.client.time.sleep", lambda s: None)
|
||||
|
||||
with pytest.raises(ComfyUIError) as excinfo:
|
||||
client.poll("prompt-timeout-1", timeout=0, interval=0)
|
||||
|
||||
assert excinfo.value.prompt_id == "prompt-timeout-1"
|
||||
assert "prompt-timeout-1" in str(excinfo.value)
|
||||
|
||||
def test_generate_resume_prompt_id_skips_resubmit(self, monkeypatch, tmp_path):
|
||||
from tools._comfyui.client import ComfyUIClient
|
||||
import sys
|
||||
|
||||
client = ComfyUIClient("http://comfy.test")
|
||||
|
||||
def fail_submit(workflow):
|
||||
raise AssertionError("submit() should not be called when resuming")
|
||||
|
||||
monkeypatch.setattr(client, "submit", fail_submit)
|
||||
_install_fake_websocket(monkeypatch, frames=[])
|
||||
monkeypatch.setattr(
|
||||
sys.modules["websocket"],
|
||||
"create_connection",
|
||||
lambda *a, **k: (_ for _ in ()).throw(
|
||||
AssertionError("resumed jobs must use history polling")
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(client, "poll", lambda prompt_id, **kwargs: {
|
||||
"outputs": {"9": {"images": [{
|
||||
"filename": "resumed.png", "subfolder": "", "type": "output",
|
||||
}]}}
|
||||
})
|
||||
monkeypatch.setattr(client, "download", lambda filename, subfolder, dest, folder_type="output": Path(dest))
|
||||
|
||||
paths = client.generate(
|
||||
{"9": {"inputs": {}}}, "9", tmp_path / "out.png",
|
||||
resume_prompt_id="already-running-id",
|
||||
)
|
||||
assert paths == [tmp_path / "out.png"]
|
||||
|
||||
def test_generate_reads_audio_key_from_savaudio_node(self, monkeypatch, tmp_path):
|
||||
"""The native SaveAudio node writes outputs under "audio", not
|
||||
"images"/"gifs" -- comfyui_music depends on this being handled."""
|
||||
from tools._comfyui.client import ComfyUIClient
|
||||
|
||||
client = ComfyUIClient("http://comfy.test")
|
||||
monkeypatch.setattr(client, "submit", lambda workflow: "p1")
|
||||
monkeypatch.setattr(client, "poll", lambda prompt_id, **kwargs: {
|
||||
"outputs": {"9": {"audio": [{
|
||||
"filename": "track.flac", "subfolder": "", "type": "output",
|
||||
}]}}
|
||||
})
|
||||
monkeypatch.setattr(client, "download", lambda filename, subfolder, dest, folder_type="output": Path(dest))
|
||||
|
||||
paths = client.generate({"9": {"inputs": {}}}, "9", tmp_path / "out.flac")
|
||||
assert paths == [tmp_path / "out.flac"]
|
||||
|
||||
def test_is_default_url_when_env_not_set(self, monkeypatch):
|
||||
from tools._comfyui.client import ComfyUIClient
|
||||
monkeypatch.delenv("COMFYUI_SERVER_URL", raising=False)
|
||||
@@ -310,6 +379,290 @@ class TestClientHelpers:
|
||||
assert "myhost:9999" in msg
|
||||
assert "COMFYUI_SERVER_URL" not in msg
|
||||
|
||||
def test_submit_includes_client_id_for_websocket_targeting(self, monkeypatch):
|
||||
from tools._comfyui.client import ComfyUIClient
|
||||
|
||||
client = ComfyUIClient("http://comfy.test")
|
||||
seen = {}
|
||||
|
||||
def fake_post(url, json=None, timeout=None):
|
||||
seen.update(json)
|
||||
return type("R", (), {
|
||||
"raise_for_status": lambda self: None,
|
||||
"json": lambda self: {"prompt_id": "abc"},
|
||||
})()
|
||||
|
||||
monkeypatch.setattr("tools._comfyui.client.requests.post", fake_post)
|
||||
client.submit({"1": {"inputs": {}}})
|
||||
assert seen["client_id"] == client.client_id
|
||||
|
||||
|
||||
class TestMultiServer:
|
||||
|
||||
def test_capability_env_var_takes_priority_over_shared(self, monkeypatch):
|
||||
from tools._comfyui.client import ComfyUIClient
|
||||
|
||||
monkeypatch.setenv("COMFYUI_SERVER_URL", "http://shared:8188")
|
||||
monkeypatch.setenv("COMFYUI_VIDEO_SERVER_URL", "http://video-gpu:8188")
|
||||
client = ComfyUIClient(capability="video")
|
||||
assert client.server_url == "http://video-gpu:8188"
|
||||
|
||||
def test_falls_back_to_shared_when_capability_var_unset(self, monkeypatch):
|
||||
from tools._comfyui.client import ComfyUIClient
|
||||
|
||||
monkeypatch.setenv("COMFYUI_SERVER_URL", "http://shared:8188")
|
||||
monkeypatch.delenv("COMFYUI_IMAGE_SERVER_URL", raising=False)
|
||||
client = ComfyUIClient(capability="image")
|
||||
assert client.server_url == "http://shared:8188"
|
||||
|
||||
def test_other_capability_env_var_does_not_leak_across_tools(self, monkeypatch):
|
||||
from tools._comfyui.client import ComfyUIClient
|
||||
|
||||
monkeypatch.delenv("COMFYUI_SERVER_URL", raising=False)
|
||||
monkeypatch.setenv("COMFYUI_IMAGE_SERVER_URL", "http://image-gpu:8188")
|
||||
monkeypatch.delenv("COMFYUI_VIDEO_SERVER_URL", raising=False)
|
||||
video_client = ComfyUIClient(capability="video")
|
||||
assert video_client.server_url == "http://localhost:8188"
|
||||
|
||||
def test_explicit_server_url_wins_over_capability_env_var(self, monkeypatch):
|
||||
from tools._comfyui.client import ComfyUIClient
|
||||
|
||||
monkeypatch.setenv("COMFYUI_VIDEO_SERVER_URL", "http://video-gpu:8188")
|
||||
client = ComfyUIClient("http://explicit:1234", capability="video")
|
||||
assert client.server_url == "http://explicit:1234"
|
||||
|
||||
def test_no_capability_behaves_as_before(self, monkeypatch):
|
||||
from tools._comfyui.client import ComfyUIClient
|
||||
|
||||
monkeypatch.setenv("COMFYUI_SERVER_URL", "http://shared:8188")
|
||||
client = ComfyUIClient()
|
||||
assert client.server_url == "http://shared:8188"
|
||||
assert client.is_default_url is False
|
||||
|
||||
def test_is_default_url_true_only_when_both_vars_unset(self, monkeypatch):
|
||||
from tools._comfyui.client import ComfyUIClient
|
||||
|
||||
monkeypatch.delenv("COMFYUI_SERVER_URL", raising=False)
|
||||
monkeypatch.delenv("COMFYUI_IMAGE_SERVER_URL", raising=False)
|
||||
client = ComfyUIClient(capability="image")
|
||||
assert client.is_default_url is True
|
||||
|
||||
monkeypatch.setenv("COMFYUI_IMAGE_SERVER_URL", "http://image-gpu:8188")
|
||||
client2 = ComfyUIClient(capability="image")
|
||||
assert client2.is_default_url is False
|
||||
|
||||
def test_unavailable_reason_mentions_capability_and_shared_var(self, monkeypatch):
|
||||
from tools._comfyui.client import ComfyUIClient
|
||||
|
||||
monkeypatch.delenv("COMFYUI_SERVER_URL", raising=False)
|
||||
monkeypatch.delenv("COMFYUI_VIDEO_SERVER_URL", raising=False)
|
||||
client = ComfyUIClient(capability="video")
|
||||
msg = client.unavailable_reason()
|
||||
assert "COMFYUI_VIDEO_SERVER_URL" in msg
|
||||
assert "COMFYUI_SERVER_URL" in msg
|
||||
|
||||
def test_image_and_video_tools_use_independent_servers(self, monkeypatch):
|
||||
from tools.graphics.comfyui_image import ComfyUIImage
|
||||
from tools.video.comfyui_video import ComfyUIVideo
|
||||
|
||||
monkeypatch.delenv("COMFYUI_SERVER_URL", raising=False)
|
||||
monkeypatch.setenv("COMFYUI_IMAGE_SERVER_URL", "http://image-gpu:8188")
|
||||
monkeypatch.setenv("COMFYUI_VIDEO_SERVER_URL", "http://video-gpu:8188")
|
||||
|
||||
image_tool = ComfyUIImage()
|
||||
video_tool = ComfyUIVideo()
|
||||
|
||||
assert image_tool._client.server_url == "http://image-gpu:8188"
|
||||
assert video_tool._client.server_url == "http://video-gpu:8188"
|
||||
|
||||
|
||||
class _FakeWSTimeout(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class _FakeWSConn:
|
||||
def __init__(self, frames):
|
||||
self._frames = list(frames)
|
||||
|
||||
def settimeout(self, value):
|
||||
pass
|
||||
|
||||
def recv(self):
|
||||
if not self._frames:
|
||||
raise _FakeWSTimeout()
|
||||
return self._frames.pop(0)
|
||||
|
||||
def close(self):
|
||||
pass
|
||||
|
||||
|
||||
def _install_fake_websocket(monkeypatch, frames):
|
||||
"""Inject a fake `websocket` module so wait_ws() runs without the real
|
||||
optional websocket-client dependency installed."""
|
||||
import sys
|
||||
import types
|
||||
|
||||
fake_module = types.SimpleNamespace(
|
||||
WebSocketTimeoutException=_FakeWSTimeout,
|
||||
create_connection=lambda url, timeout=10: _FakeWSConn(frames),
|
||||
)
|
||||
monkeypatch.setitem(sys.modules, "websocket", fake_module)
|
||||
|
||||
|
||||
class TestWebsocketWait:
|
||||
|
||||
def test_wait_ws_returns_job_completed_before_connection(self, monkeypatch):
|
||||
from tools._comfyui.client import ComfyUIClient
|
||||
import sys
|
||||
|
||||
client = ComfyUIClient("http://comfy.test")
|
||||
_install_fake_websocket(monkeypatch, frames=[])
|
||||
monkeypatch.setattr(
|
||||
sys.modules["websocket"],
|
||||
"create_connection",
|
||||
lambda *a, **k: (_ for _ in ()).throw(
|
||||
AssertionError("completed history must avoid websocket connection")
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"tools._comfyui.client.requests.get",
|
||||
lambda *a, **k: type("R", (), {
|
||||
"raise_for_status": lambda self: None,
|
||||
"json": lambda self: {"done": {"outputs": {"9": {}}}},
|
||||
})(),
|
||||
)
|
||||
|
||||
assert client.wait_ws("done", timeout=5) == {"outputs": {"9": {}}}
|
||||
|
||||
def test_wait_ws_completes_on_executing_none_node(self, monkeypatch, tmp_path):
|
||||
from tools._comfyui.client import ComfyUIClient
|
||||
|
||||
client = ComfyUIClient("http://comfy.test")
|
||||
progress_events = []
|
||||
frames = [
|
||||
json.dumps({"type": "progress", "data": {
|
||||
"value": 2, "max": 20, "prompt_id": "p1",
|
||||
}}),
|
||||
json.dumps({"type": "executing", "data": {
|
||||
"node": None, "prompt_id": "p1",
|
||||
}}),
|
||||
]
|
||||
_install_fake_websocket(monkeypatch, frames)
|
||||
history_calls = iter(({}, {}, {"p1": {"outputs": {"9": {}}}}))
|
||||
monkeypatch.setattr(
|
||||
"tools._comfyui.client.requests.get",
|
||||
lambda *a, **k: type("R", (), {
|
||||
"raise_for_status": lambda self: None,
|
||||
"json": lambda self: next(history_calls),
|
||||
})(),
|
||||
)
|
||||
|
||||
entry = client.wait_ws("p1", timeout=5, on_progress=progress_events.append)
|
||||
|
||||
assert entry == {"outputs": {"9": {}}}
|
||||
assert progress_events == [{"value": 2, "max": 20, "prompt_id": "p1"}]
|
||||
|
||||
def test_wait_ws_execution_error_raises_with_prompt_id(self, monkeypatch):
|
||||
from tools._comfyui.client import ComfyUIClient, ComfyUIError
|
||||
|
||||
client = ComfyUIClient("http://comfy.test")
|
||||
frames = [
|
||||
json.dumps({"type": "execution_error", "data": {
|
||||
"prompt_id": "p2", "exception_message": "boom",
|
||||
}}),
|
||||
]
|
||||
_install_fake_websocket(monkeypatch, frames)
|
||||
|
||||
with pytest.raises(ComfyUIError) as excinfo:
|
||||
client.wait_ws("p2", timeout=5)
|
||||
|
||||
assert excinfo.value.prompt_id == "p2"
|
||||
|
||||
def test_wait_ws_ignores_other_prompts_on_shared_connection(self, monkeypatch):
|
||||
from tools._comfyui.client import ComfyUIClient
|
||||
|
||||
client = ComfyUIClient("http://comfy.test")
|
||||
frames = [
|
||||
# Another job's event on the same client_id -- must not trigger completion.
|
||||
json.dumps({"type": "executing", "data": {
|
||||
"node": None, "prompt_id": "someone-elses-job",
|
||||
}}),
|
||||
json.dumps({"type": "executing", "data": {
|
||||
"node": None, "prompt_id": "p3",
|
||||
}}),
|
||||
]
|
||||
_install_fake_websocket(monkeypatch, frames)
|
||||
monkeypatch.setattr(
|
||||
"tools._comfyui.client.requests.get",
|
||||
lambda *a, **k: type("R", (), {
|
||||
"raise_for_status": lambda self: None,
|
||||
"json": lambda self: {"p3": {"outputs": {}}},
|
||||
})(),
|
||||
)
|
||||
|
||||
entry = client.wait_ws("p3", timeout=5)
|
||||
assert entry == {"outputs": {}}
|
||||
|
||||
def test_wait_ws_timeout_raises_comfyuierror_with_prompt_id(self, monkeypatch):
|
||||
from tools._comfyui.client import ComfyUIClient, ComfyUIError
|
||||
|
||||
client = ComfyUIClient("http://comfy.test")
|
||||
_install_fake_websocket(monkeypatch, frames=[]) # recv() always times out
|
||||
|
||||
with pytest.raises(ComfyUIError) as excinfo:
|
||||
client.wait_ws("p4", timeout=0)
|
||||
|
||||
assert excinfo.value.prompt_id == "p4"
|
||||
|
||||
def test_wait_falls_back_to_poll_when_websocket_unavailable(self, monkeypatch):
|
||||
"""No websocket-client installed (or any transport failure) must
|
||||
silently fall back to REST polling, not blow up the whole call."""
|
||||
from tools._comfyui.client import ComfyUIClient
|
||||
import sys
|
||||
|
||||
client = ComfyUIClient("http://comfy.test")
|
||||
monkeypatch.delitem(sys.modules, "websocket", raising=False)
|
||||
monkeypatch.setattr(
|
||||
"builtins.__import__",
|
||||
_raise_on_websocket_import(__import__),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
client, "poll", lambda prompt_id, **kwargs: {"outputs": {"used": "poll"}}
|
||||
)
|
||||
|
||||
entry = client._wait("p5", timeout=5, interval=5)
|
||||
assert entry == {"outputs": {"used": "poll"}}
|
||||
|
||||
def test_wait_does_not_swallow_genuine_comfyuierror_from_websocket(self, monkeypatch):
|
||||
"""A real execution error detected over the websocket must propagate,
|
||||
not be masked by a fallback-to-poll retry."""
|
||||
from tools._comfyui.client import ComfyUIClient, ComfyUIError
|
||||
|
||||
client = ComfyUIClient("http://comfy.test")
|
||||
frames = [
|
||||
json.dumps({"type": "execution_error", "data": {
|
||||
"prompt_id": "p6", "exception_message": "bad node",
|
||||
}}),
|
||||
]
|
||||
_install_fake_websocket(monkeypatch, frames)
|
||||
|
||||
def fail_poll(prompt_id, **kwargs):
|
||||
raise AssertionError("poll() should not be called after a real ws error")
|
||||
|
||||
monkeypatch.setattr(client, "poll", fail_poll)
|
||||
|
||||
with pytest.raises(ComfyUIError) as excinfo:
|
||||
client._wait("p6", timeout=5, interval=5)
|
||||
assert excinfo.value.prompt_id == "p6"
|
||||
|
||||
|
||||
def _raise_on_websocket_import(real_import):
|
||||
def _import(name, *args, **kwargs):
|
||||
if name == "websocket":
|
||||
raise ImportError("no module named websocket")
|
||||
return real_import(name, *args, **kwargs)
|
||||
return _import
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Model discovery (offline, no server needed)
|
||||
@@ -332,6 +685,11 @@ class TestModelRequirements:
|
||||
assert len(_REQUIRED_MODELS_T2V) > 0
|
||||
assert any("t2v" in m.lower() for m in _REQUIRED_MODELS_T2V)
|
||||
|
||||
def test_music_tool_has_required_models(self):
|
||||
from tools.audio.comfyui_music import _REQUIRED_MODELS
|
||||
assert len(_REQUIRED_MODELS) > 0
|
||||
assert any("ace_step" in m.lower() for m in _REQUIRED_MODELS)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Custom workflow contract and provenance
|
||||
@@ -418,6 +776,77 @@ class TestCustomWorkflowContract:
|
||||
assert provenance["model_stack"] == [{"role": "lora", "name": "style.safetensors"}]
|
||||
assert provenance["model_stack_source"] == "caller_supplied"
|
||||
|
||||
def test_video_timeout_surfaces_resumable_prompt_id(self, tmp_path):
|
||||
from tools._comfyui.client import ComfyUIError
|
||||
|
||||
tool = ComfyUIVideo()
|
||||
tool._client.is_available = lambda: True
|
||||
|
||||
def fake_generate(workflow, output_node, dest, **kwargs):
|
||||
raise ComfyUIError("Prompt timed-out-id did not complete within 5s", prompt_id="timed-out-id")
|
||||
|
||||
tool._client.generate = fake_generate
|
||||
|
||||
result = tool.execute({
|
||||
"prompt": "test",
|
||||
"workflow_json": json.dumps({"42": {"inputs": {}}}),
|
||||
"output_node": "42",
|
||||
"output_path": str(tmp_path / "video.mp4"),
|
||||
"timeout_seconds": 5,
|
||||
})
|
||||
|
||||
assert result.success is False
|
||||
assert result.data["prompt_id"] == "timed-out-id"
|
||||
assert "resume_prompt_id" in result.error
|
||||
assert "timed-out-id" in result.error
|
||||
|
||||
def test_video_passes_timeout_and_resume_prompt_id_through(self, tmp_path):
|
||||
tool = ComfyUIVideo()
|
||||
tool._client.is_available = lambda: True
|
||||
seen = {}
|
||||
|
||||
def fake_generate(workflow, output_node, dest, **kwargs):
|
||||
seen.update(kwargs)
|
||||
return [Path(dest)]
|
||||
|
||||
tool._client.generate = fake_generate
|
||||
|
||||
result = tool.execute({
|
||||
"prompt": "test",
|
||||
"workflow_json": json.dumps({"42": {"inputs": {}}}),
|
||||
"output_node": "42",
|
||||
"output_path": str(tmp_path / "video.mp4"),
|
||||
"timeout_seconds": 7200,
|
||||
"resume_prompt_id": "already-running-id",
|
||||
})
|
||||
|
||||
assert result.success is True
|
||||
assert seen["timeout"] == 7200
|
||||
assert seen["resume_prompt_id"] == "already-running-id"
|
||||
|
||||
def test_video_default_timeout_is_generous_not_900s(self, tmp_path):
|
||||
tool = ComfyUIVideo()
|
||||
tool._client.is_available = lambda: True
|
||||
seen = {}
|
||||
|
||||
def fake_generate(workflow, output_node, dest, **kwargs):
|
||||
seen.update(kwargs)
|
||||
return [Path(dest)]
|
||||
|
||||
tool._client.generate = fake_generate
|
||||
|
||||
tool.execute({
|
||||
"prompt": "test",
|
||||
"workflow_json": json.dumps({"42": {"inputs": {}}}),
|
||||
"output_node": "42",
|
||||
"output_path": str(tmp_path / "video.mp4"),
|
||||
})
|
||||
|
||||
# Regression guard: the old hardcoded 900s timeout false-failed real
|
||||
# renders on modest local GPUs (observed ~1360-1630s for non-accelerated
|
||||
# custom Wan 1.3B workflows at 832x480/81-97 frames).
|
||||
assert seen["timeout"] > 900
|
||||
|
||||
def test_image_missing_models_are_structured(self):
|
||||
tool = ComfyUIImage()
|
||||
tool._client.is_available = lambda: True
|
||||
@@ -466,6 +895,216 @@ class TestCustomWorkflowContract:
|
||||
assert any(item["role"] == "vae" for item in provenance["model_stack"])
|
||||
|
||||
|
||||
class TestComfyUIMusic:
|
||||
|
||||
def test_capability_and_provider(self):
|
||||
tool = ComfyUIMusic()
|
||||
assert tool.capability == "music_generation"
|
||||
assert tool.provider == "comfyui"
|
||||
|
||||
def test_bundled_path_requires_no_workflow_json_or_output_node(self, tmp_path):
|
||||
"""Without workflow_json/workflow_path it should attempt the bundled
|
||||
ACE-Step workflow, not demand a custom one."""
|
||||
tool = ComfyUIMusic()
|
||||
tool._client.is_available = lambda: True
|
||||
tool._client.check_models = lambda required: (list(required), [])
|
||||
tool._client.generate = lambda workflow, output_node, dest, **kwargs: [Path(dest)]
|
||||
|
||||
result = tool.execute({
|
||||
"prompt": "ambient pad",
|
||||
"output_path": str(tmp_path / "music.mp3"),
|
||||
})
|
||||
|
||||
assert result.success is True
|
||||
assert result.data["workflow_provenance"]["source"] == "bundled"
|
||||
|
||||
def test_custom_workflow_without_output_node_errors(self):
|
||||
tool = ComfyUIMusic()
|
||||
tool._client.is_available = lambda: True
|
||||
|
||||
result = tool.execute({
|
||||
"prompt": "ambient pad",
|
||||
"workflow_json": json.dumps({"9": {"inputs": {}}}),
|
||||
})
|
||||
|
||||
assert result.success is False
|
||||
assert "output_node" in result.error
|
||||
|
||||
def test_bundled_missing_models_returns_structured_payload(self):
|
||||
tool = ComfyUIMusic()
|
||||
tool._client.is_available = lambda: True
|
||||
tool._client.check_models = lambda required: ([], list(required))
|
||||
|
||||
result = tool.execute({"prompt": "ambient pad"})
|
||||
|
||||
assert result.success is False
|
||||
assert result.data["missing_models"][0]["name"] == "ace_step_v1_3.5b.safetensors"
|
||||
assert result.data["missing_models"][0]["download_url"]
|
||||
|
||||
def test_bundled_generation_patches_tags_lyrics_and_seed(self, tmp_path):
|
||||
tool = ComfyUIMusic()
|
||||
tool._client.is_available = lambda: True
|
||||
tool._client.check_models = lambda required: (list(required), [])
|
||||
seen = {}
|
||||
|
||||
def fake_generate(workflow, output_node, dest, **kwargs):
|
||||
seen["workflow"] = workflow
|
||||
seen["output_node"] = output_node
|
||||
return [Path(dest)]
|
||||
|
||||
tool._client.generate = fake_generate
|
||||
|
||||
result = tool.execute({
|
||||
"prompt": "lofi hip hop, chill, rain sounds",
|
||||
"lyrics": "[verse]\nquiet streets",
|
||||
"duration_seconds": 45,
|
||||
"seed": 777,
|
||||
"output_path": str(tmp_path / "music.mp3"),
|
||||
})
|
||||
|
||||
assert result.success is True
|
||||
assert seen["output_node"] == "10"
|
||||
assert seen["workflow"]["2"]["inputs"]["tags"] == "lofi hip hop, chill, rain sounds"
|
||||
assert seen["workflow"]["2"]["inputs"]["lyrics"] == "[verse]\nquiet streets"
|
||||
assert seen["workflow"]["4"]["inputs"]["seconds"] == 45
|
||||
assert seen["workflow"]["8"]["inputs"]["seed"] == 777
|
||||
assert result.data["model"] == "ace-step-v1-3.5b"
|
||||
|
||||
def test_bundled_generation_preserves_seed_zero(self, tmp_path):
|
||||
tool = ComfyUIMusic()
|
||||
tool._client.is_available = lambda: True
|
||||
tool._client.check_models = lambda required: (list(required), [])
|
||||
seen = {}
|
||||
|
||||
def fake_generate(workflow, output_node, dest, **kwargs):
|
||||
seen["seed"] = workflow["8"]["inputs"]["seed"]
|
||||
return [Path(dest)]
|
||||
|
||||
tool._client.generate = fake_generate
|
||||
result = tool.execute({
|
||||
"prompt": "deterministic test",
|
||||
"seed": 0,
|
||||
"output_path": str(tmp_path / "music.mp3"),
|
||||
})
|
||||
|
||||
assert result.success is True
|
||||
assert result.seed == 0
|
||||
assert seen["seed"] == 0
|
||||
|
||||
def test_get_status_degraded_when_model_missing(self):
|
||||
tool = ComfyUIMusic()
|
||||
tool._client.is_available = lambda: True
|
||||
tool._client.check_models = lambda required: ([], list(required))
|
||||
assert tool.get_status() == ToolStatus.DEGRADED
|
||||
|
||||
def test_unavailable_server_reports_unavailable_reason(self):
|
||||
tool = ComfyUIMusic()
|
||||
tool._client.is_available = lambda: False
|
||||
tool._client.unavailable_reason = lambda: "no server here"
|
||||
|
||||
result = tool.execute({
|
||||
"prompt": "ambient pad",
|
||||
"workflow_json": json.dumps({"9": {"inputs": {}}}),
|
||||
"output_node": "9",
|
||||
})
|
||||
|
||||
assert result.success is False
|
||||
assert result.error == "no server here"
|
||||
|
||||
def test_successful_generation_returns_provenance_and_duration(self, tmp_path, monkeypatch):
|
||||
tool = ComfyUIMusic()
|
||||
tool._client.is_available = lambda: True
|
||||
|
||||
dest_file = tmp_path / "music.mp3"
|
||||
|
||||
def fake_generate(workflow, output_node, dest, **kwargs):
|
||||
Path(dest).write_bytes(b"fake-audio-bytes")
|
||||
return [Path(dest)]
|
||||
|
||||
tool._client.generate = fake_generate
|
||||
monkeypatch.setattr("shutil.which", lambda name: None) # no ffprobe in test env
|
||||
|
||||
result = tool.execute({
|
||||
"prompt": "upbeat synthwave",
|
||||
"workflow_json": json.dumps({"9": {"inputs": {}}}),
|
||||
"output_node": "9",
|
||||
"output_path": str(dest_file),
|
||||
"workflow_name": "my-ace-step-graph",
|
||||
"workflow_model": "ace-step-v1-3.5b",
|
||||
})
|
||||
|
||||
assert result.success is True
|
||||
assert result.data["provider"] == "comfyui"
|
||||
assert result.data["model"] == "ace-step-v1-3.5b"
|
||||
assert result.data["output"] == str(dest_file)
|
||||
assert result.data["format"] == "mp3"
|
||||
assert result.data["duration_seconds"] is None # ffprobe unavailable
|
||||
provenance = result.data["workflow_provenance"]
|
||||
assert provenance["source"] == "user_supplied"
|
||||
assert provenance["output_node"] == "9"
|
||||
assert provenance["workflow_hash_sha256"]
|
||||
|
||||
def test_timeout_surfaces_resumable_prompt_id(self, tmp_path):
|
||||
from tools._comfyui.client import ComfyUIError
|
||||
|
||||
tool = ComfyUIMusic()
|
||||
tool._client.is_available = lambda: True
|
||||
|
||||
def fake_generate(workflow, output_node, dest, **kwargs):
|
||||
raise ComfyUIError("timed out", prompt_id="music-prompt-id")
|
||||
|
||||
tool._client.generate = fake_generate
|
||||
|
||||
result = tool.execute({
|
||||
"prompt": "ambient pad",
|
||||
"workflow_json": json.dumps({"9": {"inputs": {}}}),
|
||||
"output_node": "9",
|
||||
"output_path": str(tmp_path / "music.mp3"),
|
||||
})
|
||||
|
||||
assert result.success is False
|
||||
assert result.data["prompt_id"] == "music-prompt-id"
|
||||
assert "resume_prompt_id" in result.error
|
||||
|
||||
def test_passes_timeout_and_resume_prompt_id_through(self, tmp_path):
|
||||
tool = ComfyUIMusic()
|
||||
tool._client.is_available = lambda: True
|
||||
seen = {}
|
||||
|
||||
def fake_generate(workflow, output_node, dest, **kwargs):
|
||||
seen.update(kwargs)
|
||||
return [Path(dest)]
|
||||
|
||||
tool._client.generate = fake_generate
|
||||
|
||||
tool.execute({
|
||||
"prompt": "ambient pad",
|
||||
"workflow_json": json.dumps({"9": {"inputs": {}}}),
|
||||
"output_node": "9",
|
||||
"output_path": str(tmp_path / "music.mp3"),
|
||||
"timeout_seconds": 3600,
|
||||
"resume_prompt_id": "already-running-id",
|
||||
})
|
||||
|
||||
assert seen["timeout"] == 3600
|
||||
assert seen["resume_prompt_id"] == "already-running-id"
|
||||
|
||||
def test_registry_discovers_comfyui_music_under_music_generation(self):
|
||||
registry = ToolRegistry()
|
||||
tool = ComfyUIMusic()
|
||||
registry.register(tool)
|
||||
registry._discovered_packages.add("tools")
|
||||
|
||||
by_capability = registry.get_by_capability("music_generation")
|
||||
assert any(t.name == "comfyui_music" for t in by_capability)
|
||||
|
||||
def test_uses_music_capability_env_var_for_multi_server(self, monkeypatch):
|
||||
monkeypatch.delenv("COMFYUI_SERVER_URL", raising=False)
|
||||
monkeypatch.setenv("COMFYUI_MUSIC_SERVER_URL", "http://music-gpu:8188")
|
||||
tool = ComfyUIMusic()
|
||||
assert tool._client.server_url == "http://music-gpu:8188"
|
||||
|
||||
|
||||
class TestComfyUISetupOffer:
|
||||
|
||||
def test_provider_menu_summary_includes_structured_setup_offer(self):
|
||||
|
||||
Reference in New Issue
Block a user