mirror of
https://github.com/calesthio/OpenMontage.git
synced 2026-08-25 01:20:18 +08:00
Merge remote-tracking branch 'origin/main' into codex/repair-pr-457
# Conflicts: # tools/graphics/image_selector.py
This commit is contained in:
@@ -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