diff --git a/.agents/skills/comfyui/SKILL.md b/.agents/skills/comfyui/SKILL.md index 64950870..05bec8c2 100644 --- a/.agents/skills/comfyui/SKILL.md +++ b/.agents/skills/comfyui/SKILL.md @@ -1,17 +1,19 @@ --- name: comfyui -description: Use when working with ComfyUI workflows in OpenMontage, including comfyui_image/comfyui_video, custom workflow_json/workflow_path inputs, output_node selection, missing model setup, LoRAs, low-VRAM workflow choices, and community workflow imports. +description: Use when working with ComfyUI workflows in OpenMontage, including comfyui_image/comfyui_video/comfyui_music, custom workflow_json/workflow_path inputs, output_node selection, missing model setup, LoRAs, low-VRAM workflow choices, and community workflow imports. --- # ComfyUI Workflows in OpenMontage -Use this skill before calling `comfyui_image` or `comfyui_video`, and when converting a community ComfyUI workflow into an OpenMontage tool call. +Use this skill before calling `comfyui_image`, `comfyui_video`, or `comfyui_music`, and when converting a community ComfyUI workflow into an OpenMontage tool call. ## Server Contract - ComfyUI must be running before the tool can generate. The default server is `http://localhost:8188`; override it with `COMFYUI_SERVER_URL`. +- Running separate ComfyUI instances per capability (different GPU, different model set)? `COMFYUI_IMAGE_SERVER_URL` / `COMFYUI_VIDEO_SERVER_URL` / `COMFYUI_MUSIC_SERVER_URL` each override `COMFYUI_SERVER_URL` for that one tool only. Optional -- a single-server setup needs none of these. - Health and hardware status come from `GET /system_stats`. - Jobs are submitted to `POST /prompt`, completed outputs are read from `GET /history/{prompt_id}`, and artifact bytes are downloaded with `GET /view`. +- Long waits (video, music) prefer ComfyUI's websocket feed for immediate completion/error detection and transparently fall back to REST polling if `websocket-client` isn't installed. Either way, a timeout is recoverable: pass the error's `prompt_id` back in as `resume_prompt_id` to resume waiting on the same job instead of resubmitting it. - Export workflows with ComfyUI's API-format JSON, not the UI layout format. If a downloaded workflow will not submit, re-export it from ComfyUI with API format enabled. ## Choosing a Workflow @@ -52,4 +54,13 @@ Use this skill before calling `comfyui_image` or `comfyui_video`, and when conve - If the server is unavailable, surface the structured setup offer. Starting ComfyUI or setting `COMFYUI_SERVER_URL` is the first fix. - If models are missing, read `data.missing_models[]`; each item should include the file name, role, destination hint, and download URL when OpenMontage knows it. - If custom nodes are missing, ask the user to install them through ComfyUI Manager or the workflow author's documented install path, then restart ComfyUI. -- If a long render times out locally, check ComfyUI history before retrying from scratch; the server may still have completed the prompt. +- If a long render times out locally, check ComfyUI history before retrying from scratch; the server may still have completed the prompt -- or just call again with `resume_prompt_id` set to the `prompt_id` from the timeout error. + +## Music (`comfyui_music`) + +- Bundled default is ACE-Step v1 (3.5B) text-to-audio, built from ComfyUI's *native* `TextEncodeAceStepAudio`/`EmptyAceStepLatentAudio` nodes (core, not a third-party pack) -- unlike ACE-Step 1.5 or other custom node packs, v1's interface is standardized enough to bundle safely. +- `prompt` maps to the bundled workflow's `tags` field (style/genre/mood, e.g. `"upbeat electronic pop, female vocals"`), matching the same "prompt = music description" convention `suno_music` uses. `lyrics` is a separate optional field -- leave empty for instrumental, or use `[verse]`/`[chorus]`/`[bridge]` structure tags and `[zh]`/`[ja]`/`[ko]`-style language-code prefixes for non-English lines. +- `duration_seconds`, `steps`, `cfg`, `lyrics_strength`, and `seed` are patchable on the bundled workflow. Missing `ace_step_v1_3.5b.safetensors` surfaces through the same `data.missing_models[]` contract as image/video. +- Need ACE-Step 1.5, a different node pack, or a non-ACE-Step audio model? Fall back to `workflow_json`/`workflow_path` + `output_node`, exactly like a custom image/video workflow -- in that mode `prompt` becomes provenance/logging only again and must already be baked into the graph. +- `output_node` (bundled or custom) should be the node that writes the final audio -- the bundled workflow's is `SaveAudioMP3`. The client reads artifacts from that node's `"audio"` output key (parallel to `"images"` for image/video savers). +- For custom workflows, provide `workflow_name`/`workflow_model`/`workflow_model_stack` for provenance exactly as you would for a custom image/video workflow. diff --git a/.env.example b/.env.example index 51d973e4..f4788da5 100644 --- a/.env.example +++ b/.env.example @@ -8,6 +8,13 @@ FAL_KEY= # Alias for FAL_KEY (some SDKs/docs use this name); either one is read. FAL_AI_API_KEY= +# --- MiniMax official direct API --- +# First-party image generation (image-01 / image-01-live). +# Get one at https://platform.minimax.io/user-center/basic-information/interface-key +MINIMAX_API_KEY= +# Optional: global (default) or cn. +MINIMAX_REGION=global + # --- Replicate --- # Replicate-hosted video gen (seedance_replicate). Needed to make the # Replicate-backed Seedance path selectable alongside the fal.ai one. diff --git a/docs/PROVIDERS.md b/docs/PROVIDERS.md index 0afd5af1..e81a4348 100644 --- a/docs/PROVIDERS.md +++ b/docs/PROVIDERS.md @@ -50,6 +50,7 @@ AZURE_SPEECH_REGION= # Speech resource region, e.g. eastus # MULTI-MODEL GATEWAY (one key, 6+ tools) FAL_KEY= # FLUX, Recraft, Kling, Veo, MiniMax video +MINIMAX_API_KEY= # MiniMax first-party image generation # KLING OFFICIAL DIRECT API KLING_API_KEY= # Official Kling video, image, TTS, avatar, lip sync @@ -192,7 +193,7 @@ The ASR tool (`qwen3-asr-flash-filetrans`) uses an async submit-poll pattern. Au > **Broad single-key coverage.** One API key unlocks image and video providers across multiple models. -**Tools unlocked:** `flux_image`, `recraft_image`, `kling_video`, `veo_video`, `minimax_video` +**Tools unlocked:** `flux_image`, `recraft_image`, `seedream_image`, `kling_video`, `veo_video`, `minimax_video` **Env var:** `FAL_KEY` #### Setup @@ -213,6 +214,8 @@ No subscription — pure pay-as-you-go, no minimum spend. | FLUX Pro v1.1 | $0.05/image | 20 images | | FLUX Dev | $0.03/image | 33 images | | Recraft v3 | ~$0.04/image | 25 images | +| Seedream 5 Pro (up to 1536x1536) | $0.0675/image | ~14 images | +| Seedream 5 Pro (up to 2048x2048) | $0.135/image | ~7 images | **Video generation:** @@ -227,6 +230,40 @@ No subscription — pure pay-as-you-go, no minimum spend. --- +### MiniMax — Official Direct Image API + +> **Low-cost first-party image generation.** The direct MiniMax API supports +> seeded text-to-image, character subject references, custom dimensions, and +> global or mainland-China routing without a gateway. + +**Tool unlocked:** `minimax_image` + +**Env var:** `MINIMAX_API_KEY` + +**Optional region:** `MINIMAX_REGION=global` (default) or `cn` + +#### Setup + +1. Create a MiniMax Open Platform account. +2. Generate an API key in the account's API-key page. +3. Add `MINIMAX_API_KEY=...` to `.env`. +4. For a mainland-China account, also set `MINIMAX_REGION=cn`. + +#### Pricing + +| Models | Global pay-as-you-go price | +|--------|----------------------------| +| `image-01`, `image-01-live` | $0.0035 per generated image | + +MiniMax also offers subscription token plans with included daily image quota. +OpenMontage conservatively reports the standard pay-as-you-go amount in cost +estimates and generation results. + +The tool is automatically discoverable through `image_selector`; choose it +with `preferred_provider: "minimax"`. + +--- + ### Kling Official — Direct API > **Official Kling path.** This is separate from `kling_video` via fal.ai: it uses Kling's official `Authorization: Bearer ` API, provider name `kling_official`, and direct Classic/Turbo/Omni task protocols. diff --git a/docs/comfyui-adapter-plan.md b/docs/comfyui-adapter-plan.md index d60fb89d..0e29d601 100644 --- a/docs/comfyui-adapter-plan.md +++ b/docs/comfyui-adapter-plan.md @@ -257,21 +257,36 @@ output_node: string # required for custom workflows workflow_name: string # optional custom workflow provenance label workflow_model: string # optional custom model/provenance label workflow_model_stack: [] # optional custom dependency provenance +timeout_seconds: integer # optional, default 3600 (see below) +resume_prompt_id: string # optional, resume a timed-out job without resubmitting ``` **execute() flow (i2v):** 1. Upload reference image via `client.upload_image()` 2. Deep-copy i2v workflow template 3. Inject prompt, uploaded image name, seed, dimensions -4. `client.generate(workflow, output_node="108", dest=output_path, timeout=900)` +4. `client.generate(workflow, output_node="108", dest=output_path, timeout=inputs.get("timeout_seconds", 3600), resume_prompt_id=inputs.get("resume_prompt_id"))` 5. Return `ToolResult` **execute() flow (t2v):** 1. Deep-copy t2v workflow template 2. Inject prompt, seed, dimensions -3. `client.generate(workflow, output_node="16", dest=output_path, timeout=900)` +3. `client.generate(workflow, output_node="16", dest=output_path, timeout=inputs.get("timeout_seconds", 3600), resume_prompt_id=inputs.get("resume_prompt_id"))` 4. Return `ToolResult` +**Timeout and resume (added after real-world local-GPU testing):** the +default client wait was raised from 900s to 3600s — non-accelerated custom +Wan 1.3B workflows on modest local GPUs were observed taking ~1360-1630s at +832x480/81-97 frames, and the old 900s default false-failed those jobs even +though ComfyUI kept rendering server-side. `ComfyUIError` now carries a +`prompt_id` on both execution errors and timeouts (`ComfyUIError.prompt_id`), +and `ComfyUIVideo`'s `ToolResult.error`/`.data` surface it on timeout so the +caller isn't left guessing whether the job is dead. Callers recover a +timed-out-but-still-running job by calling `execute()` again with +`resume_prompt_id` set to that `prompt_id` (and a longer `timeout_seconds` if +needed) — `client.generate()` then skips `submit()` entirely and just resumes +polling/downloading the existing job instead of queuing a duplicate. + `comfyui_video` publishes `operation_statuses` in `get_info()` and implements `is_operation_available(operation)` for selector routing. This keeps partial ComfyUI installs useful for the installed mode without advertising unavailable @@ -281,19 +296,52 @@ not promote ComfyUI for an operation whose bundled models are missing. --- -### `comfyui_music` -- Music Generation (not shipped) +### `comfyui_music` -- Music Generation (shipped, with a native-node bundled workflow) -We explored adding a `comfyui_music` tool using the ACE-Step 3.5B model. -The model runs well in ComfyUI, but the ComfyUI node interface for -ACE-Step is not standardized -- there are multiple custom node packs with -different class names (`AceStepModelLoader` vs native `TextEncodeAceStepAudio`, -etc.). Shipping a workflow that only works with one specific custom node -pack would break for most users. +`tools/audio/comfyui_music.py`. `capability="music_generation"`, `provider="comfyui"`. -**Future path:** ACE-Step support should be revisited once OpenMontage decides -the music-generation routing shape and a portable ComfyUI audio workflow -contract. Current image/video workflow overrides are intentionally scoped to -image and video artifacts, not arbitrary audio workflows. +**Bundled default:** ACE-Step v1 (3.5B) text-to-audio, via `tools/_comfyui/workflows/ace-step-1-t2a.json`. +The node-pack fragmentation that originally blocked this tool (`AceStepModelLoader` +vs native `TextEncodeAceStepAudio`, etc.) turned out to be moot for ACE-Step v1: +ComfyUI ships `TextEncodeAceStepAudio`/`EmptyAceStepLatentAudio` as **native core +nodes** (`comfy_extras/nodes_ace.py`), not a third-party pack, and Comfy-Org's own +[`workflow_templates`](https://github.com/Comfy-Org/workflow_templates) repo bundles +an official ACE-Step-v1 template built entirely from those native nodes plus +long-stable core nodes (`CheckpointLoaderSimple`, `KSampler`, `ModelSamplingSD3`, +`VAEDecodeAudio`, `SaveAudioMP3`). Every node's `class_type` and input names in +`ace-step-1-t2a.json` were cross-checked against ComfyUI's own source +(`comfy_extras/nodes_ace.py`, `nodes_audio.py`, `nodes_latent.py`, `nodes.py`) -- +not guessed from the UI export -- since the UI-format template Comfy-Org ships +isn't directly usable as the API-format JSON this client submits. + +`prompt` maps to ACE-Step's `tags` field (style/genre/mood description, matching +the "prompt = description of desired music" convention `suno_music` already uses). +`lyrics` is a separate optional field (empty for instrumental). `duration_seconds`, +`steps`, `cfg`, `lyrics_strength`, and `seed` are all patchable; `shift` and the +tonemap `multiplier` stay at the official template's defaults. + +Newer/different setups aren't locked out: `workflow_json`/`workflow_path` + +`output_node` still works exactly like the image/video tools' override path -- +for ACE-Step 1.5, a different node pack, or a non-ACE-Step audio model entirely. + +**Selector integration:** no dedicated `music_selector` exists in OpenMontage +(unlike `tts_selector`/`image_selector`/`video_selector`) -- music tools are +already routed directly via `registry.get_by_capability("music_generation")`, +and `comfyui_music` participates in that the same way `suno_music`/`music_gen` +do. `fallback_tools = ["suno_music", "music_gen"]`. + +**Audio artifact schema:** `ToolResult.data` follows the same shape as the +image/video tools (`provider`, `model`, `output`, `format`, `workflow_provenance`), +plus `lyrics` and `duration_seconds` -- the latter a best-effort `ffprobe` probe +of the downloaded file (`None` if `ffprobe` isn't on PATH), since even the bundled +workflow doesn't report actual rendered duration back through `/history`. + +**Workflow/output-node contract:** identical to image/video -- `output_node` +must be the ID of the node that writes the final artifact (the bundled workflow's +is `SaveAudioMP3`, ComfyUI's native audio saver). `ComfyUIClient.generate()`'s +artifact extraction now also checks the `"audio"` output key (previously only +`"images"`/`"gifs"`), which is what `SaveAudioMP3`/`SaveAudio` write to in +ComfyUI's `/history` response. --- @@ -358,6 +406,19 @@ COMFYUI_POLL_TIMEOUT=600 # max wait for image gen COMFYUI_VIDEO_TIMEOUT=900 # max wait for video gen ``` +**Multi-server (optional):** point `comfyui_image`, `comfyui_video`, and +`comfyui_music` at separate ComfyUI instances -- e.g. one GPU running FLUX 2, +another running WAN 2.2, another running ACE-Step -- by setting a +per-capability override. Each takes priority over `COMFYUI_SERVER_URL` for +its own tool only; leave all three unset and everything talks to the single +shared server. + +```bash +COMFYUI_IMAGE_SERVER_URL=http://gpu-a:8188 +COMFYUI_VIDEO_SERVER_URL=http://gpu-b:8188 +COMFYUI_MUSIC_SERVER_URL=http://gpu-c:8188 +``` + **For Docker Compose setups** (ComfyUI in a container): ```bash @@ -459,15 +520,35 @@ pipeline definition, or any schema. user-provided via a config directory? Bundling gives reproducibility; external gives flexibility. -2. **Async generation:** ComfyUI supports websocket connections for real-time - progress. Worth implementing for long video generations, or is polling - sufficient? +2. ~~**Async generation:**~~ **Resolved.** `ComfyUIClient.generate()` now + waits via ComfyUI's websocket feed (`wait_ws()`) by default, reacting to + `executing`/`execution_error` events immediately instead of sleeping + between REST polls — completion and errors are caught without the + `interval`-seconds lag, and an optional `on_progress` callback gets live + `progress` events (`comfyui_video` uses this to print step progress on + long renders). No new hard dependency: `websocket-client` is an optional + import, and `_wait()` transparently falls back to the original + `poll()` REST loop when it isn't installed or the connection fails — + `resume_prompt_id` recovery behaves identically either way. -3. **Multi-server:** Should the adapter support multiple ComfyUI instances - (e.g., one for images, one for video) via per-capability URLs? +3. ~~**Multi-server:**~~ **Resolved.** `ComfyUIClient(capability="image"|"video"|"music")` + resolves its server URL from a per-capability env var first + (`COMFYUI_IMAGE_SERVER_URL` / `COMFYUI_VIDEO_SERVER_URL` / `COMFYUI_MUSIC_SERVER_URL`), + then the shared `COMFYUI_SERVER_URL`, then the `http://localhost:8188` default. + All three tools pass their capability at construction, so image, video, and + music generation can each point at different ComfyUI instances (different GPUs, + different model sets) with zero code changes -- single-server setups need no extra + configuration since all three env vars are optional. `client.capability`/ + `client.is_default_url`/`client.unavailable_reason()` all account for the + override, and `COMFYUI_SETUP_OFFER.per_capability_env_var_overrides` documents + it for the setup-offer surfacing in `provider_menu()`. -4. **Music generation:** ACE-Step works in ComfyUI but OpenMontage needs a - dedicated music-generation routing contract before adding `comfyui_music`. - The follow-up should decide selector integration, audio artifact schemas, and - a portable workflow/output-node contract rather than treating music as a - hidden image/video workflow override. +4. ~~**Music generation:**~~ **Resolved -- shipped with a bundled ACE-Step v1 workflow.** + `comfyui_music` is a real tool now (not a hidden image/video override), routed + through the existing `registry.get_by_capability("music_generation")` path + like `suno_music`/`music_gen`. The node-pack fragmentation that originally + blocked this turned out not to apply to ACE-Step v1: its ComfyUI nodes are + native core nodes, not a third-party pack, so `ace-step-1-t2a.json` ships as + the default, verified node-by-node against ComfyUI's own source. Custom + `workflow_json`/`workflow_path` + `output_node` remains available for other + versions/packs. See the `comfyui_music` section above for the full contract. diff --git a/tests/contracts/test_comfyui_tools.py b/tests/contracts/test_comfyui_tools.py index a8ee8304..cf678a85 100644 --- a/tests/contracts/test_comfyui_tools.py +++ b/tests/contracts/test_comfyui_tools.py @@ -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): diff --git a/tests/tools/test_minimax_image.py b/tests/tools/test_minimax_image.py new file mode 100644 index 00000000..7fcf4777 --- /dev/null +++ b/tests/tools/test_minimax_image.py @@ -0,0 +1,291 @@ +"""Contract tests for the MiniMax image generation tool.""" + +from __future__ import annotations + +import base64 + +import pytest +import requests + +from tools.base_tool import ToolStatus +from tools.graphics import minimax_image +from tools.graphics.minimax_image import MiniMaxImage +from tools.tool_registry import ToolRegistry + + +class FakeResponse: + def __init__(self, *, json_data=None, content: bytes = b"") -> None: + self._json_data = json_data + self.content = content + + def raise_for_status(self) -> None: + return None + + def json(self): + return self._json_data + + +@pytest.fixture(autouse=True) +def clear_minimax_env(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("MINIMAX_API_KEY", raising=False) + monkeypatch.delenv("MINIMAX_REGION", raising=False) + monkeypatch.delenv("MINIMAX_BASE_URL", raising=False) + + +def test_registry_registers_minimax_image_tool() -> None: + registry = ToolRegistry() + assert registry.register_module(minimax_image) == ["minimax_image"] + + tool = registry.get("minimax_image") + assert tool is not None + assert tool.provider == "minimax" + assert tool.capability == "image_generation" + assert tool.input_schema["properties"]["model"]["enum"] == [ + "image-01", + "image-01-live", + ] + + +def test_status_requires_api_key(monkeypatch: pytest.MonkeyPatch) -> None: + tool = MiniMaxImage() + assert tool.get_status() == ToolStatus.UNAVAILABLE + + monkeypatch.setenv("MINIMAX_API_KEY", "test-key") + assert tool.get_status() == ToolStatus.AVAILABLE + + +def test_cost_estimate_and_result_report_paid_images( + monkeypatch: pytest.MonkeyPatch, tmp_path +) -> None: + monkeypatch.setenv("MINIMAX_API_KEY", "test-key") + monkeypatch.setattr( + requests, + "post", + lambda *args, **kwargs: FakeResponse( + json_data={ + "data": { + "image_base64": [ + base64.b64encode(b"one").decode("ascii"), + base64.b64encode(b"two").decode("ascii"), + ] + }, + "base_resp": {"status_code": 0}, + } + ), + ) + + tool = MiniMaxImage() + inputs = { + "prompt": "A lighthouse at dusk", + "response_format": "base64", + "n": 2, + "output_path": str(tmp_path / "image.png"), + } + assert tool.estimate_cost(inputs) == pytest.approx(0.007) + result = tool.execute(inputs) + assert result.success, result.error + assert result.cost_usd == pytest.approx(0.007) + + +def test_image_selector_can_route_to_minimax( + monkeypatch: pytest.MonkeyPatch, tmp_path +) -> None: + from tools.graphics.image_selector import ImageSelector + + monkeypatch.setenv("MINIMAX_API_KEY", "test-key") + monkeypatch.setattr( + requests, + "post", + lambda *args, **kwargs: FakeResponse( + json_data={ + "data": { + "image_base64": [base64.b64encode(b"image").decode("ascii")] + }, + "base_resp": {"status_code": 0}, + } + ), + ) + tool = MiniMaxImage() + selector = ImageSelector() + monkeypatch.setattr(selector, "_providers", lambda: [tool]) + + result = selector.execute( + { + "prompt": "A lighthouse at dusk", + "preferred_provider": "minimax", + "response_format": "base64", + "output_path": str(tmp_path / "selected.png"), + } + ) + + assert result.success, result.error + assert result.data["selected_provider"] == "minimax" + assert result.data["selected_tool"] == "minimax_image" + + +@pytest.mark.parametrize( + ("region", "expected_base_url"), + [ + ("global", "https://api.minimax.io"), + ("global_en", "https://api.minimax.io"), + ("cn", "https://api.minimaxi.com"), + ("cn_zh", "https://api.minimaxi.com"), + ], +) +def test_region_routes_to_official_endpoint( + monkeypatch: pytest.MonkeyPatch, region: str, expected_base_url: str +) -> None: + monkeypatch.setenv("MINIMAX_REGION", region) + assert MiniMaxImage()._base_url() == expected_base_url + + +def test_url_response_downloads_all_images( + monkeypatch: pytest.MonkeyPatch, tmp_path +) -> None: + monkeypatch.setenv("MINIMAX_API_KEY", "test-key") + monkeypatch.setenv("MINIMAX_REGION", "cn") + captured = {} + + def fake_post(url, *, headers, json, timeout): + captured.update(url=url, headers=headers, payload=json, timeout=timeout) + return FakeResponse( + json_data={ + "id": "request-1", + "data": { + "image_urls": [ + "https://example.test/one.png", + "https://example.test/two.png", + ] + }, + "metadata": {"success_count": 2, "failed_count": 0}, + "base_resp": {"status_code": 0, "status_msg": "success"}, + } + ) + + def fake_get(url, *, timeout): + assert timeout == 120 + return FakeResponse(content=url.rsplit("/", 1)[-1].encode()) + + monkeypatch.setattr(requests, "post", fake_post) + monkeypatch.setattr(requests, "get", fake_get) + + output_path = tmp_path / "image.png" + result = MiniMaxImage().execute( + { + "prompt": "A lighthouse at dusk", + "model": "image-01-live", + "subject_reference": [ + {"type": "character", "image_file": "https://example.test/ref.png"} + ], + "aspect_ratio": "16:9", + "response_format": "url", + "seed": 42, + "n": 2, + "prompt_optimizer": True, + "output_path": str(output_path), + } + ) + + assert result.success + assert captured["url"] == "https://api.minimaxi.com/v1/image_generation" + assert captured["headers"]["Authorization"] == "Bearer test-key" + assert captured["payload"] == { + "model": "image-01-live", + "prompt": "A lighthouse at dusk", + "response_format": "url", + "n": 2, + "prompt_optimizer": True, + "subject_reference": [ + {"type": "character", "image_file": "https://example.test/ref.png"} + ], + "aspect_ratio": "16:9", + "seed": 42, + } + assert result.artifacts == [ + str(tmp_path / "image_1.png"), + str(tmp_path / "image_2.png"), + ] + assert (tmp_path / "image_1.png").read_bytes() == b"one.png" + assert (tmp_path / "image_2.png").read_bytes() == b"two.png" + assert result.data["metadata"] == {"success_count": 2, "failed_count": 0} + + +def test_base64_response_writes_inline_images( + monkeypatch: pytest.MonkeyPatch, tmp_path +) -> None: + monkeypatch.setenv("MINIMAX_API_KEY", "test-key") + image_bytes = b"inline image" + + monkeypatch.setattr( + requests, + "post", + lambda *args, **kwargs: FakeResponse( + json_data={ + "data": { + "image_base64": [ + "data:image/png;base64," + + base64.b64encode(image_bytes).decode("ascii") + ] + }, + "metadata": {"success_count": "1", "failed_count": "0"}, + "base_resp": {"status_code": 0}, + } + ), + ) + monkeypatch.setattr( + requests, + "get", + lambda *args, **kwargs: pytest.fail("base64 output must not be downloaded"), + ) + + output_path = tmp_path / "inline.png" + result = MiniMaxImage().execute( + { + "prompt": "A paper-cut forest", + "response_format": "base64", + "output_path": str(output_path), + } + ) + + assert result.success + assert output_path.read_bytes() == image_bytes + assert result.data["response_format"] == "base64" + + +def test_base_response_error_is_returned(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("MINIMAX_API_KEY", "test-key") + monkeypatch.setattr( + requests, + "post", + lambda *args, **kwargs: FakeResponse( + json_data={ + "base_resp": { + "status_code": 1008, + "status_msg": "insufficient balance", + } + } + ), + ) + + result = MiniMaxImage().execute({"prompt": "A mountain cabin"}) + + assert not result.success + assert result.error == "MiniMax API error 1008: insufficient balance" + + +def test_width_and_height_must_be_provided_together( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("MINIMAX_API_KEY", "test-key") + monkeypatch.setattr( + requests, + "post", + lambda *args, **kwargs: pytest.fail("invalid inputs must not call the API"), + ) + + result = MiniMaxImage().execute( + {"prompt": "A mountain cabin", "width": 1024} + ) + + assert not result.success + assert "width and height must be set together" in (result.error or "") diff --git a/tests/tools/test_seedream_image.py b/tests/tools/test_seedream_image.py new file mode 100644 index 00000000..4932a41d --- /dev/null +++ b/tests/tools/test_seedream_image.py @@ -0,0 +1,299 @@ +"""Regression tests: seedream_image must return every image it requests and bills for. + +Covers: +- Multi-image output: all requested images must be written and returned +- Cost estimation: billed count matches delivered artifacts +- Single-image output: exact output path preserved +- Async polling: COMPLETED / FAILED / CANCELLED / timeout paths +- API key validation: graceful failure when FAL_KEY is unset +""" + +import sys +import types +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + +PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent +sys.path.insert(0, str(PROJECT_ROOT)) + + +class _FakeResponse: + def __init__(self, json_data: dict | None = None, status_code: int = 200, content: bytes = b""): + self._json_data = json_data or {} + self.status_code = status_code + self.content = content + + def raise_for_status(self): + if self.status_code >= 400: + import requests + raise requests.HTTPError(response=self) + + def json(self): + return self._json_data + + +def _build_submit_response(request_id: str = "req_123") -> _FakeResponse: + return _FakeResponse({"request_id": request_id}) + + +def _build_status_response(status: str, error: str | None = None) -> _FakeResponse: + data = {"status": status} + if error: + data["error"] = error + return _FakeResponse(data) + + +def _build_result_response(image_urls: list[str]) -> _FakeResponse: + images = [{"url": url} for url in image_urls] + return _FakeResponse({"images": images}) + + +def _build_image_content(index: int) -> bytes: + return f"SEEDREAM_IMAGE_{index}".encode() + + +@pytest.fixture +def seedream_tool(monkeypatch): + monkeypatch.setenv("FAL_KEY", "test-fal-key") + from tools.graphics.seedream_image import SeedreamImage + return SeedreamImage() + + +@pytest.fixture +def mock_requests(monkeypatch): + mock_post = MagicMock() + mock_get = MagicMock() + fake_requests = types.ModuleType("requests") + fake_requests.post = mock_post + fake_requests.get = mock_get + fake_requests.HTTPError = type("HTTPError", (Exception,), {}) + monkeypatch.setitem(sys.modules, "requests", fake_requests) + return mock_post, mock_get + + +def _setup_mock_execution(mock_post, mock_get, num_images: int = 1, status: str = "COMPLETED", + error: str | None = None, extra_gets: list = None): + """Helper to setup common mock execution flow.""" + mock_post.return_value = _build_submit_response() + + side_effects = [_build_status_response(status, error)] + if status == "COMPLETED": + urls = [f"http://img.url/{i}" for i in range(num_images)] + side_effects.append(_build_result_response(urls)) + side_effects.extend([_FakeResponse(content=_build_image_content(i)) for i in range(num_images)]) + elif extra_gets: + side_effects.extend(extra_gets) + + mock_get.side_effect = side_effects + + +# ========== Core Regression Tests ========== + +class TestMultiOutputRegression: + def test_all_requested_images_are_written(self, seedream_tool, tmp_path, mock_requests): + mock_post, mock_get = mock_requests + _setup_mock_execution(mock_post, mock_get, num_images=3) + + result = seedream_tool.execute({ + "prompt": "test", "num_images": 3, + "output_format": "jpeg", "output_path": str(tmp_path / "gen.jpeg"), + }) + + assert result.success + assert result.data["image_count"] == 3 + assert len(result.artifacts) == 3 + + files = sorted(tmp_path.glob("*.jpeg")) + assert len(files) == 3 + contents = {f.read_bytes() for f in files} + assert contents == {b"SEEDREAM_IMAGE_0", b"SEEDREAM_IMAGE_1", b"SEEDREAM_IMAGE_2"} + + def test_artifacts_match_billed_count(self, seedream_tool, tmp_path, mock_requests): + mock_post, mock_get = mock_requests + _setup_mock_execution(mock_post, mock_get, num_images=4) + + inputs = {"prompt": "t", "num_images": 4, "output_path": str(tmp_path / "out.png")} + result = seedream_tool.execute(inputs) + billed = seedream_tool.estimate_cost(inputs) + + assert len(result.artifacts) == 4 + assert billed == pytest.approx(0.135 * 4) + + +class TestSingleOutput: + def test_single_image_keeps_exact_path(self, seedream_tool, tmp_path, mock_requests): + mock_post, mock_get = mock_requests + _setup_mock_execution(mock_post, mock_get, num_images=1) + + out = tmp_path / "single.png" + result = seedream_tool.execute({"prompt": "s", "num_images": 1, "output_path": str(out)}) + + assert result.success + assert result.artifacts == [str(out)] + assert out.read_bytes() == b"SEEDREAM_IMAGE_0" + + +# ========== Cost Estimation (Parameterized) ========== + +class TestCostEstimation: + @pytest.mark.parametrize("size,expected", [ + ("square", 0.0675), ("landscape_4_3", 0.0675), + ("portrait_4_3", 0.0675), ("auto_1K", 0.0675), + ]) + def test_small_size_pricing(self, seedream_tool, size, expected): + cost = seedream_tool.estimate_cost({"image_size": size, "num_images": 1}) + assert cost == pytest.approx(expected) + + @pytest.mark.parametrize("size,expected", [ + ("square_hd", 0.135), ("landscape_16_9", 0.135), + ("portrait_16_9", 0.135), ("auto_2K", 0.135), + ]) + def test_large_size_pricing(self, seedream_tool, size, expected): + cost = seedream_tool.estimate_cost({"image_size": size, "num_images": 1}) + assert cost == pytest.approx(expected) + + @pytest.mark.parametrize("n", [1, 2, 3, 4]) + def test_cost_scales_with_num_images(self, seedream_tool, n): + cost = seedream_tool.estimate_cost({"image_size": "auto_2K", "num_images": n}) + assert cost == pytest.approx(round(0.135 * n, 4)) + + def test_unknown_size_falls_back_to_high_price(self, seedream_tool): + cost = seedream_tool.estimate_cost({"image_size": "unknown", "num_images": 1}) + assert cost == pytest.approx(0.135) + + def test_default_values(self, seedream_tool): + cost = seedream_tool.estimate_cost({}) + assert cost == pytest.approx(0.135) + + +# ========== Async Polling States ========== + +class TestAsyncPolling: + def test_completed_on_first_poll(self, seedream_tool, tmp_path, mock_requests): + mock_post, mock_get = mock_requests + _setup_mock_execution(mock_post, mock_get, num_images=1) + + result = seedream_tool.execute({"prompt": "q", "output_path": str(tmp_path / "q.png")}) + assert result.success + assert result.data["request_id"] + + @pytest.mark.parametrize("status,error_msg", [ + ("FAILED", "Content policy violation"), + ("CANCELLED", None), + ]) + def test_failed_states_return_error(self, seedream_tool, mock_requests, status, error_msg): + mock_post, mock_get = mock_requests + _setup_mock_execution(mock_post, mock_get, status=status, error=error_msg) + + result = seedream_tool.execute({"prompt": "bad"}) + assert not result.success + assert status in result.error + + def test_timeout_returns_error(self, seedream_tool, mock_requests): + mock_post, mock_get = mock_requests + mock_post.return_value = _build_submit_response() + mock_get.side_effect = [_build_status_response("IN_PROGRESS")] * 100 + + with patch("tools.graphics.seedream_image.time.sleep"): + result = seedream_tool.execute({"prompt": "timeout"}) + assert not result.success + assert "timed out" in result.error.lower() + + +# ========== Validation & Error Handling ========== + +class TestValidation: + @pytest.mark.parametrize("value", [0, 5, 1.5, True]) + def test_num_images_rejects_invalid_values( + self, seedream_tool, mock_requests, value + ): + mock_post, _ = mock_requests + result = seedream_tool.execute({"prompt": "t", "num_images": value}) + assert not result.success + assert "num_images" in (result.error or "") + mock_post.assert_not_called() + + def test_missing_api_key_returns_error(self, monkeypatch): + monkeypatch.delenv("FAL_KEY", raising=False) + monkeypatch.delenv("FAL_AI_API_KEY", raising=False) + from tools.graphics.seedream_image import SeedreamImage + result = SeedreamImage().execute({"prompt": "t"}) + assert not result.success + assert "FAL_KEY" in result.error + + def test_status_available_with_key(self, seedream_tool): + assert seedream_tool.get_status().name == "AVAILABLE" + + def test_status_unavailable_without_key(self, monkeypatch): + monkeypatch.delenv("FAL_KEY", raising=False) + monkeypatch.delenv("FAL_AI_API_KEY", raising=False) + from tools.graphics.seedream_image import SeedreamImage + assert SeedreamImage().get_status().name == "UNAVAILABLE" + + def test_missing_request_id_raises_error(self, seedream_tool, mock_requests): + mock_post, mock_get = mock_requests + mock_post.return_value = _FakeResponse({}) + result = seedream_tool.execute({"prompt": "no id"}) + assert not result.success + assert "request_id" in result.error.lower() + + def test_completed_without_images_raises_error(self, seedream_tool, mock_requests): + mock_post, mock_get = mock_requests + mock_post.return_value = _build_submit_response() + mock_get.side_effect = [ + _build_status_response("COMPLETED"), + _FakeResponse({"images": []}), + ] + result = seedream_tool.execute({"prompt": "empty"}) + assert not result.success + assert "no images" in result.error.lower() + + +# ========== Metadata & Integration ========== + +class TestMetadata: + def test_provider_and_model_info(self, seedream_tool, tmp_path, mock_requests): + mock_post, mock_get = mock_requests + _setup_mock_execution(mock_post, mock_get, num_images=1) + + result = seedream_tool.execute({"prompt": "m", "output_path": str(tmp_path / "m.png")}) + assert result.data["provider"] == "seedream" + assert result.data["model"] == "seedream_v5" + assert result.model == "fal-ai/bytedance/seedream/v5" + + def test_cost_matches_estimate(self, seedream_tool, tmp_path, mock_requests): + mock_post, mock_get = mock_requests + _setup_mock_execution(mock_post, mock_get, num_images=2) + + inputs = {"prompt": "c", "image_size": "square", "num_images": 2, "output_path": str(tmp_path / "c.jpeg")} + result = seedream_tool.execute(inputs) + assert result.cost_usd == pytest.approx(seedream_tool.estimate_cost(inputs)) + + def test_image_selector_routes_count_and_returns_distinct_artifacts( + self, seedream_tool, tmp_path, mock_requests, monkeypatch + ): + from tools.graphics.image_selector import ImageSelector + + mock_post, mock_get = mock_requests + _setup_mock_execution(mock_post, mock_get, num_images=2) + selector = ImageSelector() + monkeypatch.setattr(selector, "_providers", lambda: [seedream_tool]) + + result = selector.execute( + { + "prompt": "campaign artwork", + "preferred_provider": "bytedance", + "n": 2, + "output_path": str(tmp_path / "selected.png"), + } + ) + + assert result.success, result.error + assert result.data["selected_tool"] == "seedream_image" + assert len(set(result.artifacts)) == 2 + assert {Path(path).name for path in result.artifacts} == { + "selected_1.png", + "selected_2.png", + } diff --git a/tools/_comfyui/client.py b/tools/_comfyui/client.py index fa80e843..e37bc2b4 100644 --- a/tools/_comfyui/client.py +++ b/tools/_comfyui/client.py @@ -11,14 +11,25 @@ import json import os import random import time +import uuid from pathlib import Path -from typing import Any +from typing import Any, Callable import requests class ComfyUIError(Exception): - """Raised when ComfyUI returns an error or times out.""" + """Raised when ComfyUI returns an error or times out. + + ``prompt_id`` is set when the error follows a successful ``submit()``, + so callers can recover a timed-out-but-still-running job instead of + losing track of it: poll ``GET /history/{prompt_id}`` directly, or + pass ``resume_prompt_id`` back into ``ComfyUIVideo.execute()``. + """ + + def __init__(self, message: str, prompt_id: str | None = None) -> None: + super().__init__(message) + self.prompt_id = prompt_id class ComfyUIClient: @@ -31,11 +42,32 @@ class ComfyUIClient: 4. POST /upload/image → stage a local image for I2V workflows """ - def __init__(self, server_url: str | None = None) -> None: - self.server_url = ( - server_url - or os.environ.get("COMFYUI_SERVER_URL", "http://localhost:8188") - ).rstrip("/") + def __init__( + self, server_url: str | None = None, capability: str | None = None + ) -> None: + """*capability*, if given (e.g. ``"image"``, ``"video"``), lets a + per-capability env var (``COMFYUI_{CAPABILITY}_SERVER_URL``) point + this client at its own ComfyUI instance -- useful when image and + video generation are split across separate servers/GPUs. Falls back + to the shared ``COMFYUI_SERVER_URL`` when the capability-specific + var isn't set, so single-server setups need no extra configuration. + """ + self.capability = capability + self._capability_env_var = ( + f"COMFYUI_{capability.upper()}_SERVER_URL" if capability else None + ) + resolved = server_url or self._capability_url() or os.environ.get( + "COMFYUI_SERVER_URL" + ) + self.server_url = (resolved or "http://localhost:8188").rstrip("/") + # Scopes websocket execution events to this client (see wait_ws) and + # is echoed back on /prompt so the server targets messages to us. + self.client_id = str(uuid.uuid4()) + + def _capability_url(self) -> str | None: + if self._capability_env_var: + return os.environ.get(self._capability_env_var) + return None # ------------------------------------------------------------------ # Health @@ -43,8 +75,25 @@ class ComfyUIClient: @property def is_default_url(self) -> bool: - """True if using the fallback URL (user didn't set COMFYUI_SERVER_URL).""" - return not os.environ.get("COMFYUI_SERVER_URL") + """True if neither the capability-specific nor shared env var is set.""" + return not (self._capability_url() or os.environ.get("COMFYUI_SERVER_URL")) + + def unavailable_reason(self) -> str: + """Human-readable explanation of why the server can't be reached.""" + env_var_hint = self._capability_env_var or "COMFYUI_SERVER_URL" + if self._capability_env_var: + env_var_hint += " (or the shared COMFYUI_SERVER_URL)" + if self.is_default_url: + return ( + f"No ComfyUI server found at {self.server_url} " + f"(default — no server URL configured).\n" + f"Set {env_var_hint} in your .env file to the address of " + f"your ComfyUI server (e.g. http://localhost:8188)." + ) + return ( + f"ComfyUI server not reachable at {self.server_url}.\n" + f"Check that ComfyUI is running and the URL is correct." + ) def is_available(self) -> bool: """Return True if the ComfyUI server is reachable.""" @@ -56,20 +105,6 @@ class ComfyUIClient: except Exception: return False - def unavailable_reason(self) -> str: - """Human-readable explanation of why the server can't be reached.""" - if self.is_default_url: - return ( - f"No ComfyUI server found at {self.server_url} " - f"(default — no COMFYUI_SERVER_URL configured).\n" - f"Set COMFYUI_SERVER_URL in your .env file to the address of " - f"your ComfyUI server (e.g. http://localhost:8188)." - ) - return ( - f"ComfyUI server not reachable at {self.server_url}.\n" - f"Check that ComfyUI is running and the URL is correct." - ) - # ------------------------------------------------------------------ # Model discovery # ------------------------------------------------------------------ @@ -137,7 +172,7 @@ class ComfyUIClient: """Queue a workflow for execution. Returns the ``prompt_id``.""" resp = requests.post( f"{self.server_url}/prompt", - json={"prompt": workflow}, + json={"prompt": workflow, "client_id": self.client_id}, timeout=30, ) try: @@ -164,23 +199,170 @@ class ComfyUIClient: """Block until *prompt_id* finishes. Returns the history entry.""" deadline = time.time() + timeout while time.time() < deadline: - resp = requests.get( - f"{self.server_url}/history/{prompt_id}", timeout=10 - ) - resp.raise_for_status() - history = resp.json() - if prompt_id in history: - entry = history[prompt_id] - status = entry.get("status", {}) - if status.get("status_str") == "error": - msgs = status.get("messages", []) - raise ComfyUIError(f"Execution error: {msgs}") + entry = self._history_entry(prompt_id) + if entry is not None: return entry time.sleep(interval) raise ComfyUIError( - f"Prompt {prompt_id} did not complete within {timeout}s" + f"Prompt {prompt_id} did not complete within {timeout}s. " + f"The job is very likely still running on the ComfyUI server " + f"(local/custom workflows on modest GPUs routinely exceed the " + f"client wait) — it was not cancelled. Poll " + f"GET {{server_url}}/history/{prompt_id} directly, or call " + f"generate()/execute() again with a longer timeout and this " + f"prompt_id to resume waiting without resubmitting.", + prompt_id=prompt_id, ) + def _history_entry(self, prompt_id: str) -> dict | None: + """Return a completed history entry, or ``None`` while it is absent.""" + resp = requests.get(f"{self.server_url}/history/{prompt_id}", timeout=10) + resp.raise_for_status() + entry = resp.json().get(prompt_id) + if entry is None: + return None + status = entry.get("status", {}) + if status.get("status_str") == "error": + msgs = status.get("messages", []) + raise ComfyUIError(f"Execution error: {msgs}", prompt_id=prompt_id) + return entry + + def _history_entry_if_reachable(self, prompt_id: str) -> dict | None: + """Best-effort history probe while the websocket remains usable.""" + try: + return self._history_entry(prompt_id) + except ComfyUIError: + raise + except Exception: + return None + + def wait_ws( + self, + prompt_id: str, + *, + timeout: int = 600, + interval: int = 5, + on_progress: Callable[[dict], None] | None = None, + ) -> dict: + """Block until *prompt_id* finishes, watching ComfyUI's websocket feed. + + Reacts to server-pushed ``executing``/``progress``/``execution_error`` + events instead of sleeping between REST polls, so completion and + errors are detected immediately rather than up to *interval* seconds + late. *on_progress*, if given, is called with each ``progress`` + message's ``data`` dict (``value``, ``max``, ``node``, ``prompt_id``). + + Requires the optional ``websocket-client`` package. Any transport + failure (missing dependency, connection refused, dropped socket, + malformed frame) propagates as a plain exception — callers should + catch it and fall back to :meth:`poll`, which is what :meth:`generate` + does. A genuine ComfyUI-side execution error or an unmet deadline is + raised as :class:`ComfyUIError` with ``prompt_id`` set, exactly like + :meth:`poll`, so ``resume_prompt_id`` recovery works the same way + regardless of which wait strategy was used. + """ + import websocket # websocket-client; optional, see docstring + + # History is authoritative and websocket events are not replayed. The + # job may already have finished between submit() and this wait call. + entry = self._history_entry_if_reachable(prompt_id) + if entry is not None: + return entry + + ws_url = self.server_url.replace("http://", "ws://", 1).replace( + "https://", "wss://", 1 + ) + conn = websocket.create_connection( + f"{ws_url}/ws?clientId={self.client_id}", timeout=10 + ) + try: + conn.settimeout(interval) + deadline = time.time() + timeout + finished = False + # Close the remaining race between the first history probe and + # websocket connection establishment. Events after this point are + # queued on the open socket; earlier completion is in history. + entry = self._history_entry_if_reachable(prompt_id) + if entry is not None: + return entry + while time.time() < deadline: + try: + raw = conn.recv() + except websocket.WebSocketTimeoutException: + entry = self._history_entry_if_reachable(prompt_id) + if entry is not None: + return entry + continue + if not isinstance(raw, str): + continue # binary preview-image frame, not a status message + try: + message = json.loads(raw) + except json.JSONDecodeError: + continue + data = message.get("data", {}) + if data.get("prompt_id") not in (None, prompt_id): + continue # another job sharing this connection + msg_type = message.get("type") + if msg_type == "progress": + if on_progress: + on_progress(data) + elif msg_type == "execution_error": + raise ComfyUIError( + f"Execution error: {data}", prompt_id=prompt_id + ) + elif msg_type == "executing" and data.get("node") is None: + finished = True + break + finally: + conn.close() + + if not finished: + entry = self._history_entry_if_reachable(prompt_id) + if entry is not None: + return entry + raise ComfyUIError( + f"Prompt {prompt_id} did not complete within {timeout}s " + f"(websocket wait). The job was not cancelled — resume with " + f"resume_prompt_id={prompt_id!r} and a longer timeout.", + prompt_id=prompt_id, + ) + + entry = self._history_entry(prompt_id) + if entry is None: + raise ComfyUIError( + f"No history entry for {prompt_id} after completion", + prompt_id=prompt_id, + ) + return entry + + def _wait( + self, + prompt_id: str, + *, + timeout: int, + interval: int, + on_progress: Callable[[dict], None] | None = None, + ) -> dict: + """Wait for *prompt_id*, preferring the websocket feed over polling. + + Falls back to :meth:`poll` when ``websocket-client`` isn't installed + or the websocket can't be established/maintained. A genuine + :class:`ComfyUIError` (execution error or deadline reached) is never + swallowed by the fallback — only transport-level failures are. The + fallback gets whatever's left of *timeout*, not a fresh budget, so a + mid-wait websocket drop can't double the caller's worst-case wait. + """ + started = time.time() + try: + return self.wait_ws( + prompt_id, timeout=timeout, interval=interval, on_progress=on_progress + ) + except ComfyUIError: + raise + except Exception: + remaining = max(timeout - (time.time() - started), 0) + return self.poll(prompt_id, timeout=remaining, interval=interval) + def download( self, filename: str, @@ -229,16 +411,41 @@ class ComfyUIClient: *, timeout: int = 600, interval: int = 5, + resume_prompt_id: str | None = None, + on_progress: Callable[[dict], None] | None = None, ) -> list[Path]: - """Submit → poll → download. Returns list of artifact paths.""" - prompt_id = self.submit(workflow) - entry = self.poll(prompt_id, timeout=timeout, interval=interval) + """Submit → wait → download. Returns list of artifact paths. + + Pass ``resume_prompt_id`` (from a previous ``ComfyUIError.prompt_id``) + to skip re-submitting an already-queued/running job and just resume + waiting on it — the common recovery path after a timeout. + + Waiting prefers ComfyUI's websocket feed (immediate completion/error + detection, optional live ``on_progress`` callback) and transparently + falls back to REST polling if ``websocket-client`` isn't installed or + the connection can't be used. See :meth:`_wait`. + """ + prompt_id = resume_prompt_id or self.submit(workflow) + if resume_prompt_id: + # A prompt resumed by a new client instance was submitted with the + # original instance's client_id, so its websocket events are not + # guaranteed to reach this socket. Poll authoritative history. + entry = self.poll(prompt_id, timeout=timeout, interval=interval) + else: + entry = self._wait( + prompt_id, timeout=timeout, interval=interval, on_progress=on_progress + ) outputs = entry.get("outputs", {}) node_output = outputs.get(output_node, {}) - # ComfyUI stores images and videos under the "images" key - items = node_output.get("images", []) or node_output.get("gifs", []) + # ComfyUI stores images/video frames under "images", legacy GIFs + # under "gifs", and the native SaveAudio node's output under "audio". + items = ( + node_output.get("images", []) + or node_output.get("gifs", []) + or node_output.get("audio", []) + ) if not items: raise ComfyUIError( f"No output artifacts on node {output_node}. " diff --git a/tools/_comfyui/metadata.py b/tools/_comfyui/metadata.py index dcca514b..084fbeb3 100644 --- a/tools/_comfyui/metadata.py +++ b/tools/_comfyui/metadata.py @@ -18,6 +18,14 @@ COMFYUI_SETUP_OFFER: dict[str, Any] = { "free local video generation through ComfyUI workflows", "community workflow_json/workflow_path execution", ], + # Optional: point image/video generation at separate ComfyUI instances + # (e.g. different GPUs). Each overrides COMFYUI_SERVER_URL for its own + # tool only; single-server setups can ignore this entirely. + "per_capability_env_var_overrides": { + "comfyui_image": "COMFYUI_IMAGE_SERVER_URL", + "comfyui_video": "COMFYUI_VIDEO_SERVER_URL", + "comfyui_music": "COMFYUI_MUSIC_SERVER_URL", + }, } @@ -176,6 +184,17 @@ BUNDLED_MODEL_STACKS: dict[str, list[dict[str, Any]]] = { ), }, ], + "ace-step-1-t2a": [ + { + "role": "checkpoint", + "name": "ace_step_v1_3.5b.safetensors", + "destination_hint": "ComfyUI/models/checkpoints/", + "download_url": ( + "https://huggingface.co/Comfy-Org/ACE-Step_ComfyUI_repackaged/" + "blob/main/all_in_one/ace_step_v1_3.5b.safetensors" + ), + }, + ], } diff --git a/tools/_comfyui/workflows/ace-step-1-t2a.json b/tools/_comfyui/workflows/ace-step-1-t2a.json new file mode 100644 index 00000000..fdeba405 --- /dev/null +++ b/tools/_comfyui/workflows/ace-step-1-t2a.json @@ -0,0 +1,80 @@ +{ + "1": { + "class_type": "CheckpointLoaderSimple", + "inputs": { + "ckpt_name": "ace_step_v1_3.5b.safetensors" + } + }, + "2": { + "class_type": "TextEncodeAceStepAudio", + "inputs": { + "clip": ["1", 1], + "tags": "", + "lyrics": "", + "lyrics_strength": 0.99 + } + }, + "3": { + "class_type": "ConditioningZeroOut", + "inputs": { + "conditioning": ["2", 0] + } + }, + "4": { + "class_type": "EmptyAceStepLatentAudio", + "inputs": { + "seconds": 120, + "batch_size": 1 + } + }, + "5": { + "class_type": "ModelSamplingSD3", + "inputs": { + "model": ["1", 0], + "shift": 5.0 + } + }, + "6": { + "class_type": "LatentOperationTonemapReinhard", + "inputs": { + "multiplier": 1.0 + } + }, + "7": { + "class_type": "LatentApplyOperationCFG", + "inputs": { + "model": ["5", 0], + "operation": ["6", 0] + } + }, + "8": { + "class_type": "KSampler", + "inputs": { + "model": ["7", 0], + "positive": ["2", 0], + "negative": ["3", 0], + "latent_image": ["4", 0], + "seed": 0, + "steps": 50, + "cfg": 5.0, + "sampler_name": "euler", + "scheduler": "simple", + "denoise": 1.0 + } + }, + "9": { + "class_type": "VAEDecodeAudio", + "inputs": { + "samples": ["8", 0], + "vae": ["1", 2] + } + }, + "10": { + "class_type": "SaveAudioMP3", + "inputs": { + "audio": ["9", 0], + "filename_prefix": "openmontage", + "quality": "V0" + } + } +} diff --git a/tools/audio/comfyui_music.py b/tools/audio/comfyui_music.py new file mode 100644 index 00000000..05dfddf0 --- /dev/null +++ b/tools/audio/comfyui_music.py @@ -0,0 +1,367 @@ +"""ComfyUI music generation via a local or remote ComfyUI server. + +Default workflow: ACE-Step v1 (3.5B) text-to-audio using ComfyUI's native +``TextEncodeAceStepAudio``/``EmptyAceStepLatentAudio`` nodes (built into +ComfyUI core, not a third-party pack). Custom workflows are still accepted +via ``workflow_json``/``workflow_path`` for other ACE-Step node packs, other +versions (e.g. ACE-Step 1.5), or entirely different audio models -- the same +override contract ``comfyui_image``/``comfyui_video`` offer. +""" + +from __future__ import annotations + +import json +import shutil +import subprocess +import time +from pathlib import Path +from typing import Any + +from tools.base_tool import ( + BaseTool, + Determinism, + ExecutionMode, + ResourceProfile, + RetryPolicy, + ToolResult, + ToolRuntime, + ToolStability, + ToolStatus, + ToolTier, +) +from tools._comfyui.client import ComfyUIClient, ComfyUIError +from tools._comfyui.metadata import ( + BUNDLED_MODEL_STACKS, + COMFYUI_SETUP_OFFER, + missing_models_payload, + model_stack, + workflow_hash, +) + +_WORKFLOWS = Path(__file__).resolve().parent.parent / "_comfyui" / "workflows" + +# Model required by the bundled ACE-Step v1 workflow +_REQUIRED_MODELS = ["ace_step_v1_3.5b.safetensors"] + + +class ComfyUIMusic(BaseTool): + name = "comfyui_music" + version = "0.2.0" + tier = ToolTier.GENERATE + capability = "music_generation" + provider = "comfyui" + stability = ToolStability.EXPERIMENTAL + execution_mode = ExecutionMode.SYNC + determinism = Determinism.SEEDED + runtime = ToolRuntime.LOCAL_GPU + + dependencies = [] # checked at runtime via server health + setup_offer = COMFYUI_SETUP_OFFER + install_instructions = ( + "Start a ComfyUI server and set COMFYUI_SERVER_URL " + "(default http://localhost:8188).\n" + "Requires ace_step_v1_3.5b.safetensors in ComfyUI's checkpoints " + "directory for the bundled workflow.\n" + "Running a separate ComfyUI instance for music? Set " + "COMFYUI_MUSIC_SERVER_URL instead -- it takes priority over " + "COMFYUI_SERVER_URL for this tool only." + ) + agent_skills = ["comfyui"] + + capabilities = ["generate_background_music", "generate_song", "generate_instrumental"] + supports = { + "seed": True, + "lyrics": True, + "custom_workflow": True, + "custom_output_node": True, + "offline": True, + } + best_for = [ + "local GPU music generation without API costs", + "instrumentals and songs with lyrics via the bundled ACE-Step v1 workflow", + "full control over sampling or other ACE-Step versions/node packs via custom ComfyUI workflows", + ] + not_good_for = [ + "setups without a running ComfyUI server", + "CPU-only machines", + ] + fallback_tools = ["suno_music", "music_gen"] + + input_schema = { + "type": "object", + "required": ["prompt"], + "properties": { + "prompt": { + "type": "string", + "description": ( + "Style/mood/genre description (ACE-Step 'tags'), e.g. " + "'upbeat electronic pop, female vocals, driving bassline'. " + "Comma-separated tags work best. Not injected for custom workflows." + ), + }, + "lyrics": { + "type": "string", + "default": "", + "description": ( + "Optional lyrics. Leave empty for instrumental. Supports structure " + "tags like [verse]/[chorus]/[bridge] and language-code prefixes " + "(e.g. [zh], [ja]) for non-English lines." + ), + }, + "duration_seconds": {"type": "number", "default": 120.0}, + "steps": {"type": "integer", "default": 50}, + "cfg": {"type": "number", "default": 5.0}, + "lyrics_strength": {"type": "number", "default": 0.99}, + "seed": {"type": "integer", "description": "Random if omitted"}, + "output_path": {"type": "string", "description": "Where to save the audio"}, + "workflow_json": { + "type": "string", + "description": "Optional full ComfyUI workflow JSON. Requires output_node.", + }, + "workflow_path": { + "type": "string", + "description": "Optional path to a ComfyUI workflow JSON file. Requires output_node.", + }, + "output_node": { + "type": "string", + "description": "ComfyUI output node ID for custom workflow_json/workflow_path.", + }, + "workflow_name": { + "type": "string", + "description": "Optional human-readable provenance label for a custom workflow.", + }, + "workflow_model": { + "type": "string", + "description": "Optional model/provenance label for a custom workflow.", + }, + "workflow_model_stack": { + "type": "array", + "description": ( + "Optional provenance metadata for custom workflow dependencies. " + "Items should include name, role, and node-pack origin when known." + ), + "items": {"type": "object"}, + }, + "timeout_seconds": { + "type": "integer", + "description": "How long to wait for the ComfyUI job before giving up. Default 1800s (30min).", + }, + "resume_prompt_id": { + "type": "string", + "description": "A prompt_id from a previous timed-out call. Skips resubmission and resumes waiting/downloading.", + }, + }, + } + + resource_profile = ResourceProfile( + cpu_cores=2, ram_mb=8000, vram_mb=8000, disk_mb=500, network_required=False, + ) + retry_policy = RetryPolicy(max_retries=1, retryable_errors=["timeout"]) + idempotency_key_fields = ["prompt", "lyrics", "duration_seconds", "seed"] + side_effects = ["writes audio file to output_path"] + user_visible_verification = ["Listen to generated audio for mood, genre accuracy, and quality"] + + def __init__(self) -> None: + self._client = ComfyUIClient(capability="music") + self._last_progress_log = 0.0 + + def get_status(self) -> ToolStatus: + if not self._client.is_available(): + return ToolStatus.UNAVAILABLE + _, missing = self._client.check_models(_REQUIRED_MODELS) + if missing: + return ToolStatus.DEGRADED + return ToolStatus.AVAILABLE + + def estimate_cost(self, inputs: dict[str, Any]) -> float: + return 0.0 + + def estimate_runtime(self, inputs: dict[str, Any]) -> float: + return float(inputs.get("steps", 50)) * 2.0 + + def get_info(self) -> dict[str, Any]: + info = super().get_info() + info["setup_offer"] = self.setup_offer + info["bundled_model_stack"] = BUNDLED_MODEL_STACKS["ace-step-1-t2a"] + return info + + def _log_progress(self, data: dict) -> None: + """Throttled progress line (see comfyui_video for rationale).""" + now = time.monotonic() + if now - self._last_progress_log < 10: + return + self._last_progress_log = now + value, max_value = data.get("value"), data.get("max") + if value is not None and max_value: + print(f"[comfyui_music] step {value}/{max_value}") + + def execute(self, inputs: dict[str, Any]) -> ToolResult: + custom_workflow = bool(inputs.get("workflow_json") or inputs.get("workflow_path")) + if custom_workflow and not inputs.get("output_node"): + return ToolResult( + success=False, + error=( + "Custom ComfyUI workflows require output_node so OpenMontage " + "knows which ComfyUI node to download artifacts from." + ), + ) + + if not self._client.is_available(): + return ToolResult(success=False, error=self._client.unavailable_reason()) + + if not custom_workflow: + _, missing = self._client.check_models(_REQUIRED_MODELS) + if missing: + return ToolResult( + success=False, + data=missing_models_payload( + missing, + workflow_key="ace-step-1-t2a", + workflow_name="ace-step-1-t2a.json", + ), + error=( + f"ComfyUI server is running but missing required models: " + f"{', '.join(missing)}.\n" + f"See data.missing_models for destination hints and download URLs." + ), + ) + + start = time.time() + seed = inputs.get("seed") + if seed is None: + seed = ComfyUIClient.random_seed() + output_path = Path(inputs.get("output_path", f"comfyui_music_{seed}.mp3")) + + try: + if custom_workflow: + workflow = self._load_custom_workflow(inputs) + output_node = str(inputs["output_node"]) + else: + workflow = ComfyUIClient.load_workflow(_WORKFLOWS / "ace-step-1-t2a.json") + workflow = ComfyUIClient.patch_workflow(workflow, { + "2": { + "tags": inputs["prompt"], + "lyrics": inputs.get("lyrics", ""), + "lyrics_strength": inputs.get("lyrics_strength", 0.99), + }, + "4": {"seconds": inputs.get("duration_seconds", 120.0)}, + "8": { + "seed": seed, + "steps": inputs.get("steps", 50), + "cfg": inputs.get("cfg", 5.0), + }, + "10": {"filename_prefix": output_path.stem}, + }) + output_node = "10" + + provenance = self._workflow_provenance(inputs, custom_workflow, output_node, workflow) + paths = self._client.generate( + workflow, + output_node=output_node, + dest=output_path, + timeout=inputs.get("timeout_seconds", 1800), + interval=10, + resume_prompt_id=inputs.get("resume_prompt_id"), + on_progress=self._log_progress, + ) + + except ComfyUIError as exc: + data = {"prompt_id": exc.prompt_id} if exc.prompt_id else {} + if exc.prompt_id: + error_msg = ( + f"{exc}\n\nThis job was NOT cancelled and is very likely still " + f"running server-side. To recover it without resubmitting, call " + f"execute() again with resume_prompt_id={exc.prompt_id!r} " + f"(and a longer timeout_seconds if it needs more time), or poll " + f"GET {{COMFYUI_SERVER_URL}}/history/{exc.prompt_id} directly." + ) + else: + error_msg = str(exc) + return ToolResult(success=False, error=error_msg, data=data) + except Exception as exc: + return ToolResult(success=False, error=f"ComfyUI music generation failed: {exc}") + + duration = self._probe_duration(paths[0]) + model_name = self._model_name(inputs, custom_workflow) + return ToolResult( + success=True, + data={ + "provider": "comfyui", + "model": model_name, + "prompt": inputs["prompt"], + "lyrics": inputs.get("lyrics", ""), + "duration_seconds": duration, + "output": str(paths[0]), + "format": paths[0].suffix.lstrip("."), + "workflow_provenance": provenance, + }, + artifacts=[str(p) for p in paths], + cost_usd=0.0, + duration_seconds=round(time.time() - start, 2), + seed=seed, + model=model_name, + ) + + @staticmethod + def _load_custom_workflow(inputs: dict[str, Any]) -> dict: + if inputs.get("workflow_json"): + return json.loads(inputs["workflow_json"]) + return ComfyUIClient.load_workflow(Path(inputs["workflow_path"])) + + @staticmethod + def _model_name(inputs: dict[str, Any], custom_workflow: bool) -> str: + if not custom_workflow: + return "ace-step-v1-3.5b" + return ( + inputs.get("workflow_model") + or inputs.get("model") + or inputs.get("workflow_name") + or "custom-comfyui-workflow" + ) + + @staticmethod + def _workflow_provenance( + inputs: dict[str, Any], + custom_workflow: bool, + output_node: str, + workflow: dict[str, Any], + ) -> dict[str, Any]: + if not custom_workflow: + return { + "source": "bundled", + "workflow": "ace-step-1-t2a.json", + "workflow_hash_sha256": workflow_hash(workflow), + "model_stack": model_stack("ace-step-1-t2a", inputs), + "output_node": output_node, + } + stack = inputs.get("workflow_model_stack") + return { + "source": "user_supplied", + "workflow_name": inputs.get("workflow_name"), + "workflow_path": inputs.get("workflow_path"), + "model": inputs.get("workflow_model") or inputs.get("model"), + "workflow_hash_sha256": workflow_hash(workflow), + "model_stack": stack if isinstance(stack, list) else [], + "model_stack_source": "caller_supplied" if stack else "unknown_custom_workflow", + "output_node": output_node, + } + + @staticmethod + def _probe_duration(path: Path) -> float | None: + """Best-effort track duration via ffprobe; None if unavailable.""" + if shutil.which("ffprobe") is None: + return None + try: + out = subprocess.run( + [ + "ffprobe", "-v", "error", + "-show_entries", "format=duration", + "-of", "default=noprint_wrappers=1:nokey=1", + str(path), + ], + capture_output=True, text=True, timeout=15, check=True, + ) + value = out.stdout.strip() + return round(float(value), 2) if value else None + except (subprocess.SubprocessError, ValueError): + return None diff --git a/tools/graphics/comfyui_image.py b/tools/graphics/comfyui_image.py index 4c55d94b..e91e71c3 100644 --- a/tools/graphics/comfyui_image.py +++ b/tools/graphics/comfyui_image.py @@ -58,7 +58,9 @@ class ComfyUIImage(BaseTool): install_instructions = ( "Start a ComfyUI server and set COMFYUI_SERVER_URL " "(default http://localhost:8188).\n" - "See https://github.com/comfyanonymous/ComfyUI for setup." + "See https://github.com/comfyanonymous/ComfyUI for setup.\n" + "Running a separate ComfyUI instance for images? Set COMFYUI_IMAGE_SERVER_URL " + "instead -- it takes priority over COMFYUI_SERVER_URL for this tool only." ) agent_skills = ["comfyui", "flux-best-practices"] @@ -133,7 +135,7 @@ class ComfyUIImage(BaseTool): user_visible_verification = ["Inspect generated image for quality and prompt adherence"] def __init__(self) -> None: - self._client = ComfyUIClient() + self._client = ComfyUIClient(capability="image") def get_status(self) -> ToolStatus: if not self._client.is_available(): diff --git a/tools/graphics/image_selector.py b/tools/graphics/image_selector.py index 965679cc..3bab0bd8 100644 --- a/tools/graphics/image_selector.py +++ b/tools/graphics/image_selector.py @@ -250,6 +250,8 @@ class ImageSelector(BaseTool): and "model" not in adapted ): adapted["model"] = adapted["model_name"] + if "n" in adapted and "num_images" in props and "num_images" not in adapted: + adapted["num_images"] = adapted["n"] # Strip selector-only keys that downstream tools don't understand adapted.pop("preferred_provider", None) diff --git a/tools/graphics/minimax_image.py b/tools/graphics/minimax_image.py new file mode 100644 index 00000000..2c879223 --- /dev/null +++ b/tools/graphics/minimax_image.py @@ -0,0 +1,298 @@ +"""MiniMax image generation through the first-party API.""" + +from __future__ import annotations + +import base64 +import os +import time +from pathlib import Path +from typing import Any + +from tools.base_tool import ( + BaseTool, + Determinism, + ExecutionMode, + ResourceProfile, + RetryPolicy, + ToolResult, + ToolRuntime, + ToolStability, + ToolTier, +) + + +MODELS = ["image-01", "image-01-live"] +DEFAULT_MODEL = "image-01" +DEFAULT_REGION = "global" +# Official global pay-as-you-go rate for image-01/image-01-live. +PRICE_PER_IMAGE_USD = 0.0035 +REGION_BASE_URLS = { + "global": "https://api.minimax.io", + "global_en": "https://api.minimax.io", + "cn": "https://api.minimaxi.com", + "cn_zh": "https://api.minimaxi.com", +} + + +class MiniMaxImage(BaseTool): + name = "minimax_image" + version = "0.1.0" + tier = ToolTier.GENERATE + capability = "image_generation" + provider = "minimax" + stability = ToolStability.BETA + execution_mode = ExecutionMode.SYNC + determinism = Determinism.SEEDED + runtime = ToolRuntime.API + + dependencies = ["env:MINIMAX_API_KEY"] + install_instructions = ( + "Set MINIMAX_API_KEY to your MiniMax API key. " + "Optionally set MINIMAX_REGION to global or cn." + ) + # MiniMax is not a FLUX model. Use the provider-neutral visual direction + # skill until a dedicated MiniMax prompting skill is available. + agent_skills = ["visual-style"] + + capabilities = ["generate_image", "text_to_image"] + supports = { + "multiple_outputs": True, + "aspect_ratio": True, + "custom_dimensions": True, + "seed": True, + "subject_reference": True, + "url_response": True, + "base64_response": True, + } + best_for = [ + "first-party MiniMax image generation", + "seeded multi-image generation", + "global and mainland China API routing", + ] + not_good_for = ["offline generation"] + + input_schema = { + "type": "object", + "required": ["prompt"], + "properties": { + "prompt": {"type": "string", "maxLength": 1500}, + "model": { + "type": "string", + "enum": MODELS, + "default": DEFAULT_MODEL, + }, + "subject_reference": { + "type": "array", + "items": { + "type": "object", + "required": ["type", "image_file"], + "properties": { + "type": {"type": "string", "enum": ["character"]}, + "image_file": {"type": "string"}, + }, + }, + }, + "aspect_ratio": { + "type": "string", + "enum": ["1:1", "16:9", "4:3", "3:2", "2:3", "3:4", "9:16", "21:9"], + "default": "1:1", + }, + "width": {"type": "integer", "minimum": 512, "maximum": 2048, "multipleOf": 8}, + "height": {"type": "integer", "minimum": 512, "maximum": 2048, "multipleOf": 8}, + "response_format": { + "type": "string", + "enum": ["url", "base64"], + "default": "url", + }, + "seed": {"type": "integer"}, + "n": {"type": "integer", "minimum": 1, "maximum": 9, "default": 1}, + "prompt_optimizer": {"type": "boolean", "default": False}, + "output_path": {"type": "string"}, + }, + } + + resource_profile = ResourceProfile( + cpu_cores=1, ram_mb=512, vram_mb=0, disk_mb=100, network_required=True + ) + retry_policy = RetryPolicy( + max_retries=2, retryable_errors=["rate_limit", "timeout"] + ) + idempotency_key_fields = [ + "prompt", + "model", + "subject_reference", + "aspect_ratio", + "width", + "height", + "response_format", + "seed", + "n", + "prompt_optimizer", + ] + side_effects = [ + "writes image files to output_path", + "calls the MiniMax image generation API", + ] + user_visible_verification = [ + "Inspect generated images for prompt adherence and visual quality" + ] + + @staticmethod + def _region() -> str: + region = os.environ.get("MINIMAX_REGION", DEFAULT_REGION).strip().lower() + return region if region in REGION_BASE_URLS else DEFAULT_REGION + + def _base_url(self) -> str: + override = os.environ.get("MINIMAX_BASE_URL") + if override: + return override.rstrip("/") + return REGION_BASE_URLS[self._region()] + + @staticmethod + def _base_resp_error(data: dict[str, Any]) -> str | None: + base_resp = data.get("base_resp") or {} + status_code = base_resp.get("status_code") + if status_code in (None, 0): + return None + status_msg = base_resp.get("status_msg") or "unknown error" + return f"MiniMax API error {status_code}: {status_msg}" + + @staticmethod + def _output_paths(output_path: str | None, count: int) -> list[Path]: + path = Path(output_path or "minimax_image.png") + if not path.suffix: + path = path.with_suffix(".png") + if count == 1: + return [path] + return [ + path.with_name(f"{path.stem}_{index}{path.suffix}") + for index in range(1, count + 1) + ] + + @staticmethod + def _build_payload(inputs: dict[str, Any]) -> dict[str, Any]: + model = inputs.get("model", DEFAULT_MODEL) + if model not in MODELS: + raise ValueError(f"Unsupported MiniMax image model '{model}'.") + + prompt = inputs.get("prompt") + if not isinstance(prompt, str) or not prompt: + raise ValueError("MiniMax image generation requires 'prompt'.") + if len(prompt) > 1500: + raise ValueError("MiniMax image prompt must not exceed 1500 characters.") + + width = inputs.get("width") + height = inputs.get("height") + if (width is None) != (height is None): + raise ValueError("MiniMax image width and height must be set together.") + + payload: dict[str, Any] = { + "model": model, + "prompt": prompt, + "response_format": inputs.get("response_format", "url"), + "n": inputs.get("n", 1), + "prompt_optimizer": inputs.get("prompt_optimizer", False), + } + for field in ( + "subject_reference", + "aspect_ratio", + "width", + "height", + "seed", + ): + if inputs.get(field) is not None: + payload[field] = inputs[field] + return payload + + @staticmethod + def _decode_base64_image(value: str) -> bytes: + encoded = value.split(",", 1)[1] if value.startswith("data:") else value + return base64.b64decode(encoded) + + @staticmethod + def _safe_error(exc: Exception, api_key: str) -> str: + return str(exc).replace(api_key, "[redacted]") if api_key else str(exc) + + def estimate_cost(self, inputs: dict[str, Any]) -> float: + return PRICE_PER_IMAGE_USD * int(inputs.get("n", 1)) + + def execute(self, inputs: dict[str, Any]) -> ToolResult: + api_key = os.environ.get("MINIMAX_API_KEY", "") + if not api_key: + return ToolResult( + success=False, + error="MINIMAX_API_KEY not set. " + self.install_instructions, + ) + + import requests + + start = time.time() + try: + payload = self._build_payload(inputs) + response = requests.post( + f"{self._base_url()}/v1/image_generation", + headers={ + "Authorization": f"Bearer {api_key}", + "Content-Type": "application/json", + }, + json=payload, + timeout=180, + ) + response.raise_for_status() + data = response.json() + + base_error = self._base_resp_error(data) + if base_error: + return ToolResult(success=False, error=base_error) + + response_format = payload["response_format"] + data_object = data.get("data") or {} + image_values = data_object.get( + "image_base64" if response_format == "base64" else "image_urls" + ) or [] + if not image_values: + return ToolResult( + success=False, + error=f"MiniMax returned no {response_format} image outputs.", + ) + + output_paths = self._output_paths( + inputs.get("output_path"), len(image_values) + ) + for path, value in zip(output_paths, image_values): + path.parent.mkdir(parents=True, exist_ok=True) + if response_format == "base64": + path.write_bytes(self._decode_base64_image(value)) + else: + download = requests.get(value, timeout=120) + download.raise_for_status() + path.write_bytes(download.content) + except Exception as exc: + return ToolResult( + success=False, + error=( + "MiniMax image generation failed: " + f"{self._safe_error(exc, api_key)}" + ), + ) + + outputs = [str(path) for path in output_paths] + return ToolResult( + success=True, + data={ + "provider": "minimax", + "model": payload["model"], + "prompt": payload["prompt"], + "region": self._region(), + "response_format": payload["response_format"], + "output": outputs[0], + "outputs": outputs, + "images_generated": len(outputs), + "metadata": data.get("metadata") or {}, + "request_id": data.get("id"), + }, + artifacts=outputs, + cost_usd=self.estimate_cost(inputs), + duration_seconds=round(time.time() - start, 2), + model=payload["model"], + ) diff --git a/tools/graphics/seedream_image.py b/tools/graphics/seedream_image.py new file mode 100644 index 00000000..64acc8d8 --- /dev/null +++ b/tools/graphics/seedream_image.py @@ -0,0 +1,275 @@ +"""Seedream V5 image generation via fal.ai API. +deep-thinking prompt understanding, native text in 14 languages, and precise control over dense layouts and structured designs. +""" +from __future__ import annotations + +import os +import time +from pathlib import Path +from typing import Any + +from tools.base_tool import ( + BaseTool, + Determinism, + ExecutionMode, + ResourceProfile, + RetryPolicy, + ToolResult, + ToolRuntime, + ToolStability, + ToolStatus, + ToolTier, +) +class SeedreamImage(BaseTool): + name = "seedream_image" + version = "0.1.0" + tier = ToolTier.GENERATE + capability = "image_generation" + provider = "bytedance" + stability = ToolStability.EXPERIMENTAL + execution_mode = ExecutionMode.ASYNC + determinism = Determinism.STOCHASTIC + runtime = ToolRuntime.API + + dependencies = ["env:FAL_KEY"] + install_instructions = ( + "Set FAL_KEY to your fal.ai API key.\n" + " Get one at https://fal.ai/dashboard/keys" + ) + agent_skills = ["visual-style"] + + capabilities = [ + "generate_image", + "text_to_image", + "structured_designs", + "dense_layouts", + "multi_language_text", + ] + supports = { + "text_rendering": True, + "color_palette": True, + "custom_size": True, + "structured_designs": True, + "dense_layouts": True, + "multi_language_text": True, + } + best_for = [ + "raster brand and campaign assets", + "images with accurate text rendering", + "structured designs and dense layouts", + "multi-language text rendering (14 languages)", + ] + input_schema = { + "type": "object", + "required": ["prompt"], + "properties": { + "prompt": {"type": "string"}, + "image_size": { + "type": "string", + "enum": [ + "square", "square_hd", + "landscape_4_3", "landscape_16_9", + "portrait_4_3", "portrait_16_9", + "auto_1K","auto_2K" + ], + "default": "auto_2K", + }, + "num_images": { + "type": "integer", + "minimum": 1, + "maximum": 4, + "default": 1, + }, + "output_format": { + "type": "string", + "enum": ["jpeg", "png"], + "description": "Output image format. Use 'jpeg' for smaller file size with lossy compression (suitable for web/preview), or 'png' for lossless quality with transparency support (suitable for design assets and further editing).", + }, + "enable_safety_checker": { + "type": "boolean", + "default": True, + "description": "If set to true, the safety checker will be enabled.", + }, + "output_path": {"type": "string"} + }, + } + resource_profile = ResourceProfile( + cpu_cores=1, ram_mb=512, vram_mb=0, disk_mb=100, network_required=True + ) + retry_policy = RetryPolicy(max_retries=2, retryable_errors=["rate_limit", "timeout"]) + idempotency_key_fields = [ + "prompt", + "image_size", + "output_format", + "num_images", + "enable_safety_checker", + ] + side_effects = ["writes image file to output_path", "calls fal.ai queue API"] + user_visible_verification = ["Inspect generated image for brand accuracy and text readability"] + + def _get_api_key(self) -> str | None: + return os.environ.get("FAL_KEY") or os.environ.get("FAL_AI_API_KEY") + + def get_status(self) -> ToolStatus: + if self._get_api_key(): + return ToolStatus.AVAILABLE + return ToolStatus.UNAVAILABLE + + def estimate_cost(self, inputs: dict[str, Any]) -> float: + image_size = inputs.get("image_size", "auto_2K") + num_images = inputs.get("num_images", 1) + size_price_map = { + "square": 0.0675, + "square_hd": 0.135, + "landscape_4_3": 0.0675, + "landscape_16_9": 0.135, + "portrait_4_3": 0.0675, + "portrait_16_9": 0.135, + "auto_1K": 0.0675, + "auto_2K": 0.135, + } + unit_price = size_price_map.get(image_size, 0.135) + return round(unit_price * num_images, 4) + + @staticmethod + def _output_paths( + output_path: str | None, count: int, output_format: str + ) -> list[Path]: + path = Path(output_path or f"seedream_image.{output_format}") + if not path.suffix: + path = path.with_suffix(f".{output_format}") + if count == 1: + return [path] + return [ + path.with_name(f"{path.stem}_{index}{path.suffix}") + for index in range(1, count + 1) + ] + + def execute(self, inputs: dict[str, Any]) -> ToolResult: + import requests + + api_key = self._get_api_key() + if not api_key: + return ToolResult( + success=False, + error="FAL_KEY not set. " + self.install_instructions, + ) + + start = time.time() + prompt = inputs["prompt"] + num_images = inputs.get("num_images", 1) + if isinstance(num_images, bool) or not isinstance(num_images, int): + return ToolResult( + success=False, error="num_images must be an integer from 1 to 4." + ) + if not 1 <= num_images <= 4: + return ToolResult( + success=False, error="num_images must be between 1 and 4." + ) + submit_url = "https://queue.fal.run/bytedance/seedream/v5/pro/text-to-image" + payload: dict[str, Any] = { + "prompt": prompt, + "image_size": inputs.get("image_size", "auto_2K"), + "output_format": inputs.get("output_format", "jpeg"), + "num_images": num_images, + "enable_safety_checker": inputs.get("enable_safety_checker", True), + } + + try: + headers = { + "Authorization": f"Key {api_key}", + "Content-Type": "application/json", + } + + submit_resp = requests.post( + submit_url, + headers=headers, + json=payload, + timeout=(10, 60), + ) + submit_resp.raise_for_status() + submit_data = submit_resp.json() + request_id = submit_data.get("request_id") + if not request_id: + raise RuntimeError( + "Seedream submit succeeded but did not return request_id" + ) + status_url = ( + f"https://queue.fal.run/bytedance/seedream/requests/" + f"{request_id}/status" + ) + elapsed = 0.0 + while elapsed < 300: + status_resp = requests.get( + status_url, + headers=headers, + timeout=30, + ) + status_resp.raise_for_status() + status_data = status_resp.json() + status = status_data.get("status") + + if status == "COMPLETED": + break + elif status in ("FAILED", "CANCELLED"): + error_msg = status_data.get("error", "Unknown error") + raise RuntimeError(f"Seedream task {status}: {error_msg}") + + time.sleep(10) + elapsed += 10 + + if elapsed >= 300: + raise RuntimeError( + f"Seedream task timed out after {300}s" + ) + + result_resp = requests.get( + f"https://queue.fal.run/bytedance/seedream/requests/" + f"{request_id}", + headers=headers, + timeout=30, + ) + result_resp.raise_for_status() + result_data = result_resp.json() + + images = result_data.get("images", []) + if not images: + raise RuntimeError("Seedream completed but no images returned") + + ext = inputs.get("output_format", "jpeg") + expected_paths = self._output_paths( + inputs.get("output_path"), len(images), ext + ) + output_paths = [] + for img, output_path in zip(images, expected_paths): + image_url = img.get("url") + if not image_url: + continue + image_resp = requests.get(image_url, timeout=60) + image_resp.raise_for_status() + + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_bytes(image_resp.content) + output_paths.append(str(output_path)) + + except Exception as e: + return ToolResult( + success=False, + error=f"Seedream generation failed: {e}", + ) + + return ToolResult( + success=True, + data={ + "provider": "seedream", + "model": "seedream_v5", + "prompt": prompt, + "request_id": request_id, + "image_count": len(output_paths), + "outputs": output_paths, + }, + artifacts=output_paths, + cost_usd=self.estimate_cost(inputs), + duration_seconds=round(time.time() - start, 2), + model="fal-ai/bytedance/seedream/v5", + ) diff --git a/tools/video/comfyui_video.py b/tools/video/comfyui_video.py index 409264f5..85ed3717 100644 --- a/tools/video/comfyui_video.py +++ b/tools/video/comfyui_video.py @@ -107,7 +107,9 @@ class ComfyUIVideo(BaseTool): install_instructions = ( "Start a ComfyUI server and set COMFYUI_SERVER_URL " "(default http://localhost:8188).\n" - "Requires WAN 2.2 models and LightX2V LoRAs in ComfyUI's model directory." + "Requires WAN 2.2 models and LightX2V LoRAs in ComfyUI's model directory.\n" + "Running a separate ComfyUI instance for video? Set COMFYUI_VIDEO_SERVER_URL " + "instead -- it takes priority over COMFYUI_SERVER_URL for this tool only." ) agent_skills = ["comfyui", "ai-video-gen", "ltx2"] @@ -186,6 +188,24 @@ class ComfyUIVideo(BaseTool): ), "items": {"type": "object"}, }, + "timeout_seconds": { + "type": "integer", + "description": ( + "How long to wait for the ComfyUI job to finish before giving up. " + "Default 3600s (1hr) covers slow/local GPUs and non-accelerated " + "custom workflows; raise it further for large frame counts or " + "high resolutions. On timeout the job is NOT cancelled server-side " + "and the error's data.prompt_id can be passed back via " + "resume_prompt_id to keep waiting without resubmitting." + ), + }, + "resume_prompt_id": { + "type": "string", + "description": ( + "A prompt_id from a previous timed-out call (see error data on " + "timeout). Skips resubmission and just resumes waiting/downloading." + ), + }, }, } @@ -198,7 +218,24 @@ class ComfyUIVideo(BaseTool): user_visible_verification = ["Watch generated clip for motion coherence and artifacts"] def __init__(self) -> None: - self._client = ComfyUIClient() + self._client = ComfyUIClient(capability="video") + self._last_progress_log = 0.0 + + def _log_progress(self, data: dict) -> None: + """Print a throttled progress line for long video renders. + + Video jobs can run for tens of minutes; without this the process + looks hung. Throttled to once per 10s since ComfyUI pushes a + ``progress`` event per sampling step, which would otherwise flood + stdout on fast GPUs. + """ + now = time.monotonic() + if now - self._last_progress_log < 10: + return + self._last_progress_log = now + value, max_value = data.get("value"), data.get("max") + if value is not None and max_value: + print(f"[comfyui_video] step {value}/{max_value}") def get_status(self) -> ToolStatus: if not self._client.is_available(): @@ -320,12 +357,25 @@ class ComfyUIVideo(BaseTool): workflow, output_node=output_node, dest=output_path, - timeout=900, + timeout=inputs.get("timeout_seconds", 3600), interval=10, + resume_prompt_id=inputs.get("resume_prompt_id"), + on_progress=self._log_progress, ) except ComfyUIError as exc: - return ToolResult(success=False, error=str(exc)) + data = {"prompt_id": exc.prompt_id} if exc.prompt_id else {} + if exc.prompt_id: + error_msg = ( + f"{exc}\n\nThis job was NOT cancelled and is very likely still " + f"running server-side. To recover it without resubmitting, call " + f"execute() again with resume_prompt_id={exc.prompt_id!r} " + f"(and a longer timeout_seconds if it needs more time), or poll " + f"GET {{COMFYUI_SERVER_URL}}/history/{exc.prompt_id} directly." + ) + else: + error_msg = str(exc) + return ToolResult(success=False, error=error_msg, data=data) except Exception as exc: return ToolResult(success=False, error=f"ComfyUI video generation failed: {exc}")