mirror of
https://github.com/calesthio/OpenMontage.git
synced 2026-08-15 13:14:38 +08:00
Merge remote-tracking branch 'origin/main' into codex/repair-pr-457
# Conflicts: # tools/graphics/image_selector.py
This commit is contained in:
calesthio
@@ -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.
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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 <KLING_API_KEY>` API, provider name `kling_official`, and direct Classic/Turbo/Omni task protocols.
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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):
|
||||
|
||||
291
tests/tools/test_minimax_image.py
Normal file
291
tests/tools/test_minimax_image.py
Normal file
@@ -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 "")
|
||||
299
tests/tools/test_seedream_image.py
Normal file
299
tests/tools/test_seedream_image.py
Normal file
@@ -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",
|
||||
}
|
||||
@@ -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}. "
|
||||
|
||||
@@ -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"
|
||||
),
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
|
||||
80
tools/_comfyui/workflows/ace-step-1-t2a.json
Normal file
80
tools/_comfyui/workflows/ace-step-1-t2a.json
Normal file
@@ -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"
|
||||
}
|
||||
}
|
||||
}
|
||||
367
tools/audio/comfyui_music.py
Normal file
367
tools/audio/comfyui_music.py
Normal file
@@ -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
|
||||
@@ -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():
|
||||
|
||||
@@ -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)
|
||||
|
||||
298
tools/graphics/minimax_image.py
Normal file
298
tools/graphics/minimax_image.py
Normal file
@@ -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"],
|
||||
)
|
||||
275
tools/graphics/seedream_image.py
Normal file
275
tools/graphics/seedream_image.py
Normal file
@@ -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",
|
||||
)
|
||||
@@ -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}")
|
||||
|
||||
|
||||
Reference in New Issue
Block a user