update grok
This commit is contained in:
@@ -1,4 +1,11 @@
|
||||
1. 你是一个经验丰富的python开发者,擅长制作commfyui的节点,现在需要创建一个新的节点
|
||||
2. 确保输入输出图像或者遮罩的张量形状分别为B H W C ,B H W
|
||||
3. 需要包含导入,节点注册等构建节点所需的完整代码
|
||||
4. 所有远程 API 节点必须接收隐藏的 UNIQUE_ID,并在等待请求、上传、轮询和下载期间,每秒在节点上方显示状态与已运行秒数。
|
||||
5. 所有阻塞等待必须至少每 0.2 秒调用 comfy.model_management.throw_exception_if_processing_interrupted()。检测到终止后必须立即关闭会话,以 wait=False 停止等待线程,并原样重新抛出 InterruptProcessingException;禁止把终止转换成普通错误输出。
|
||||
6. 用户终止后禁止继续重试、轮询、下载、解码或提交后续请求,并在节点上方显示包含耗时的 Cancelled 状态。
|
||||
7. 每次执行最多提交一次计费生成请求,除非用户明确要求多个任务。计费 POST 禁止自动重试,除非服务支持幂等键且每次重试复用同一个键,防止网络状态不明时重复扣费。
|
||||
8. 每次计费请求发送前必须再次检查 ComfyUI 终止状态;所有请求和轮询必须设置有限超时。
|
||||
9. 优先显示接口返回的实际 usage、credits 或价格;估算费用必须明确标注为 estimated。不得承诺远端已经受理请求后,本地终止能够退款或撤销计费。
|
||||
10. 付费 API 节点测试必须模拟慢请求,确认点击终止后 0.5 秒内返回,且不会提交第二次计费请求;同时验证前端能够看到运行秒数。
|
||||
|
||||
|
||||
@@ -0,0 +1,27 @@
|
||||
# Paid API Node Safety Rules
|
||||
|
||||
These rules apply to every remote API node in this plugin, including new nodes and changes to existing nodes.
|
||||
|
||||
## Cancellation and progress
|
||||
|
||||
- Every node that waits for a remote request, polling job, upload, or download must accept ComfyUI's hidden `UNIQUE_ID` and display a frontend status with elapsed seconds at least once per second while it is waiting.
|
||||
- Every blocking wait must call `comfy.model_management.throw_exception_if_processing_interrupted()` at least every 0.2 seconds. This includes generation requests, uploads, polling sleeps, retries, and result downloads.
|
||||
- A blocking HTTP call must not run directly on the ComfyUI execution thread. Run it in an interruptible wrapper and, on cancellation, close the HTTP session, cancel work that has not started, call executor shutdown with `wait=False`, and immediately re-raise `InterruptProcessingException`.
|
||||
- Never catch or convert `InterruptProcessingException` into a normal node result. Do not use a broad `except BaseException` that swallows cancellation.
|
||||
- After cancellation, do not submit another request, retry, poll, download, decode, or return a partial success result. Show a short `Cancelled` frontend status including elapsed seconds.
|
||||
- Clear or replace the running status with `Completed`, `Failed`, `Timed out`, or `Cancelled` when execution ends.
|
||||
|
||||
## Billing safety
|
||||
|
||||
- Call `throw_exception_if_processing_interrupted()` immediately before every request that may create a billable task.
|
||||
- Submit at most one billable generation request per node execution unless the user explicitly requested multiple jobs.
|
||||
- Do not automatically retry a billable `POST` unless the provider supports an idempotency key and the same key is reused for every attempt. Connection ambiguity is not proof that a paid request failed.
|
||||
- Use finite connection and response timeouts. Polling must also have a finite attempt or elapsed-time limit.
|
||||
- If an API returns a task ID, preserve it in status or node output when practical so charges can be audited.
|
||||
- Show actual usage, credits, or price returned by the service. If only an estimate is available, label it clearly as estimated.
|
||||
- Never claim that local cancellation refunds or cancels a provider charge after the remote service has accepted the request.
|
||||
|
||||
## Verification
|
||||
|
||||
- Tests for paid API nodes must simulate a slow request and verify that ComfyUI cancellation returns control within 0.5 seconds without submitting a second billable request.
|
||||
- Verify that elapsed time is visible through the frontend progress-text mechanism during the simulated wait.
|
||||
@@ -0,0 +1,107 @@
|
||||
import io
|
||||
from urllib.request import urlopen
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
|
||||
def _image_bytes(image):
|
||||
if image.ndim != 4 or image.shape[0] != 1:
|
||||
raise ValueError("阿里云换脸节点每个输入只支持一张图片")
|
||||
|
||||
array = image[0].detach().cpu().numpy()
|
||||
array = np.clip(array * 255.0, 0, 255).astype(np.uint8)
|
||||
output = io.BytesIO()
|
||||
Image.fromarray(array, "RGB").save(output, format="JPEG", quality=95)
|
||||
output.seek(0)
|
||||
return output
|
||||
|
||||
|
||||
def _image_tensor(data):
|
||||
with Image.open(io.BytesIO(data)) as image:
|
||||
array = np.asarray(image.convert("RGB"), dtype=np.float32) / 255.0
|
||||
return torch.from_numpy(array.copy()).unsqueeze(0)
|
||||
|
||||
|
||||
def _create_client():
|
||||
try:
|
||||
from alibabacloud_credentials.client import Client as CredentialClient
|
||||
from alibabacloud_facebody20191230.client import Client
|
||||
from alibabacloud_tea_openapi import models as open_api_models
|
||||
except ImportError as error:
|
||||
raise RuntimeError(
|
||||
"缺少阿里云 SDK,请运行: pip install alibabacloud_facebody20191230 alibabacloud_credentials"
|
||||
) from error
|
||||
|
||||
config = open_api_models.Config(credential=CredentialClient())
|
||||
config.endpoint = "facebody.cn-shanghai.aliyuncs.com"
|
||||
config.region_id = "cn-shanghai"
|
||||
return Client(config)
|
||||
|
||||
|
||||
class PDAliyunFaceSwap:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"template_image": ("IMAGE", {"tooltip": "保留姿势、表情和背景的底图"}),
|
||||
"face_image": ("IMAGE", {"tooltip": "提供新人物面部特征的人脸图"}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("换脸结果",)
|
||||
FUNCTION = "swap_face"
|
||||
CATEGORY = "PD/API/阿里云"
|
||||
DESCRIPTION = "将人脸图的面部特征融合到底图。请确保已获得两张图片中人物的授权。"
|
||||
|
||||
def swap_face(self, template_image, face_image):
|
||||
from alibabacloud_facebody20191230 import models
|
||||
from alibabacloud_tea_util import models as util_models
|
||||
|
||||
client = _create_client()
|
||||
runtime = util_models.RuntimeOptions()
|
||||
template_id = None
|
||||
|
||||
try:
|
||||
template_request = models.AddFaceImageTemplateAdvanceRequest(
|
||||
image_urlobject=_image_bytes(template_image)
|
||||
)
|
||||
template_response = client.add_face_image_template_advance(template_request, runtime)
|
||||
template_id = template_response.body.data.template_id
|
||||
if not template_id:
|
||||
raise RuntimeError("阿里云未返回人脸模板 ID")
|
||||
|
||||
merge_request = models.MergeImageFaceAdvanceRequest(
|
||||
template_id=template_id,
|
||||
image_urlobject=_image_bytes(face_image),
|
||||
)
|
||||
merge_response = client.merge_image_face_advance(merge_request, runtime)
|
||||
image_url = merge_response.body.data.image_url
|
||||
if not image_url:
|
||||
raise RuntimeError("阿里云未返回换脸结果地址")
|
||||
|
||||
with urlopen(image_url, timeout=60) as response:
|
||||
return (_image_tensor(response.read()),)
|
||||
except Exception as error:
|
||||
code = getattr(error, "code", None)
|
||||
message = getattr(error, "message", None) or str(error)
|
||||
detail = f" [{code}]" if code else ""
|
||||
raise RuntimeError(f"阿里云换脸失败{detail}: {message}") from error
|
||||
finally:
|
||||
if template_id:
|
||||
try:
|
||||
delete_request = models.DeleteFaceImageTemplateRequest(template_id=template_id)
|
||||
client.delete_face_image_template_with_options(delete_request, runtime)
|
||||
except Exception as error:
|
||||
print(f"[PD_NODE] 阿里云临时人脸模板清理失败: {error}")
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"PD_AliyunFaceSwap": PDAliyunFaceSwap,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"PD_AliyunFaceSwap": "PD 阿里云换脸(双图)",
|
||||
}
|
||||
@@ -0,0 +1,311 @@
|
||||
import base64
|
||||
import concurrent.futures
|
||||
import time
|
||||
|
||||
import comfy.model_management as mm
|
||||
import requests
|
||||
import torch
|
||||
from server import PromptServer
|
||||
|
||||
from ._image_gpt_common import (
|
||||
empty_image,
|
||||
format_credits_value,
|
||||
format_rmb_range,
|
||||
format_rmb_value,
|
||||
load_image_to_tensor,
|
||||
tensor_to_bytes,
|
||||
usd_to_comfy_credits,
|
||||
)
|
||||
|
||||
|
||||
BASE_URL = "https://api.comfy.org"
|
||||
EDIT_ENDPOINT = "/proxy/xai/v1/images/edits"
|
||||
MODELS = [
|
||||
"grok-imagine-image-2.0",
|
||||
"grok-imagine-image-quality",
|
||||
"grok-imagine-image-pro",
|
||||
"grok-imagine-image",
|
||||
]
|
||||
ASPECT_RATIOS = [
|
||||
"auto",
|
||||
"1:1",
|
||||
"2:3",
|
||||
"3:2",
|
||||
"3:4",
|
||||
"4:3",
|
||||
"9:16",
|
||||
"16:9",
|
||||
"9:19.5",
|
||||
"19.5:9",
|
||||
"9:20",
|
||||
"20:9",
|
||||
"1:2",
|
||||
"2:1",
|
||||
]
|
||||
QUALITY_MODELS = {"grok-imagine-image-2.0"}
|
||||
ASPECT_RATIO_NEEDS_MULTIPLE = {"grok-imagine-image-quality", "grok-imagine-image"}
|
||||
MAX_INPUT_IMAGES = {
|
||||
"grok-imagine-image-2.0": 3,
|
||||
"grok-imagine-image-quality": 3,
|
||||
"grok-imagine-image-pro": 1,
|
||||
"grok-imagine-image": 3,
|
||||
}
|
||||
|
||||
|
||||
def tensor_to_data_url(image):
|
||||
encoded = base64.b64encode(tensor_to_bytes(image).getvalue()).decode("utf-8")
|
||||
return f"data:image/png;base64,{encoded}"
|
||||
|
||||
|
||||
def flatten_images(*inputs):
|
||||
images = []
|
||||
for tensor in inputs:
|
||||
if tensor is None:
|
||||
continue
|
||||
if tensor.ndim == 4:
|
||||
images.extend(tensor[index] for index in range(tensor.shape[0]))
|
||||
else:
|
||||
images.append(tensor)
|
||||
return images
|
||||
|
||||
|
||||
def send_status(node_id, text):
|
||||
if node_id is None:
|
||||
return
|
||||
try:
|
||||
PromptServer.instance.send_progress_text(text, node_id)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def request_with_interrupt(session, method, url, node_id, status, start_time, **kwargs):
|
||||
executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
|
||||
future = executor.submit(method, url, **kwargs)
|
||||
last_elapsed = -1
|
||||
try:
|
||||
while True:
|
||||
try:
|
||||
return future.result(timeout=0.2)
|
||||
except concurrent.futures.TimeoutError:
|
||||
mm.throw_exception_if_processing_interrupted()
|
||||
elapsed = int(time.monotonic() - start_time)
|
||||
if elapsed != last_elapsed:
|
||||
send_status(node_id, f"{status}\nTime elapsed: {elapsed}s")
|
||||
last_elapsed = elapsed
|
||||
except mm.InterruptProcessingException:
|
||||
future.cancel()
|
||||
session.close()
|
||||
send_status(node_id, f"Cancelled\nTime elapsed: {int(time.monotonic() - start_time)}s")
|
||||
raise
|
||||
except BaseException:
|
||||
future.cancel()
|
||||
session.close()
|
||||
raise
|
||||
finally:
|
||||
executor.shutdown(wait=False, cancel_futures=True)
|
||||
|
||||
|
||||
def estimate_price(model, resolution, quality, number_of_images):
|
||||
is_1k = resolution.lower() == "1k"
|
||||
if model == "grok-imagine-image-2.0":
|
||||
rate = (0.04 if is_1k else 0.06) if quality == "low" else (0.06 if is_1k else 0.08)
|
||||
base = 0.01
|
||||
elif model == "grok-imagine-image-quality":
|
||||
rate = 0.05 if is_1k else 0.07
|
||||
base = 0.01
|
||||
elif model == "grok-imagine-image-pro":
|
||||
rate = 0.07
|
||||
base = 0.002
|
||||
else:
|
||||
rate = 0.02
|
||||
base = 0.002
|
||||
|
||||
output_price = rate * number_of_images
|
||||
if model == "grok-imagine-image-pro":
|
||||
price = base + output_price
|
||||
return price, price
|
||||
return base + output_price, 3 * base + output_price
|
||||
|
||||
|
||||
def format_price(result, response, model, resolution, quality, number_of_images):
|
||||
usage = result.get("usage") or {}
|
||||
ticks = usage.get("cost_in_usd_ticks")
|
||||
credits_header = response.headers.get("X-Comfy-Credits-Used")
|
||||
|
||||
credits = None
|
||||
try:
|
||||
credits = float(credits_header) if credits_header is not None else None
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
usd = ticks / 10_000_000_000 if isinstance(ticks, (int, float)) else None
|
||||
if usd is None and credits is not None:
|
||||
usd = credits / usd_to_comfy_credits(1.0)
|
||||
if credits is None and usd is not None:
|
||||
credits = usd_to_comfy_credits(usd)
|
||||
|
||||
if usd is not None:
|
||||
lines = [
|
||||
f"Actual price (USD): ${usd:.6f}",
|
||||
f"Actual price (RMB): {format_rmb_value(usd)}",
|
||||
]
|
||||
if credits is not None:
|
||||
lines.append(f"Credits used: {credits:.2f} credits")
|
||||
return "\n".join(lines)
|
||||
|
||||
min_usd, max_usd = estimate_price(model, resolution, quality, number_of_images)
|
||||
if min_usd == max_usd:
|
||||
return "\n".join([
|
||||
f"Estimated price (USD): ${min_usd:.6f}",
|
||||
f"Estimated price (RMB): {format_rmb_value(min_usd)}",
|
||||
f"Estimated price (Credits): {format_credits_value(min_usd)}",
|
||||
])
|
||||
return "\n".join([
|
||||
f"Estimated price (USD): ${min_usd:.6f}-${max_usd:.6f}",
|
||||
f"Estimated price (RMB): {format_rmb_range(min_usd, max_usd)}",
|
||||
f"Estimated price (Credits): 约 {usd_to_comfy_credits(min_usd):.2f}-{usd_to_comfy_credits(max_usd):.2f} credits",
|
||||
])
|
||||
|
||||
|
||||
class PDGrokImageEditAPIKey:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"api_key": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False,
|
||||
"placeholder": "Your ComfyUI API Key",
|
||||
}),
|
||||
"prompt": ("STRING", {"default": "", "multiline": True}),
|
||||
"model": (MODELS, {"default": "grok-imagine-image-2.0"}),
|
||||
"image_1": ("IMAGE",),
|
||||
"resolution": (["1K", "2K"], {"default": "1K"}),
|
||||
"quality": (["medium", "low"], {"default": "medium"}),
|
||||
"number_of_images": ("INT", {"default": 1, "min": 1, "max": 10, "step": 1}),
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": 2147483647, "step": 1}),
|
||||
"aspect_ratio": (ASPECT_RATIOS, {"default": "auto"}),
|
||||
},
|
||||
"optional": {
|
||||
"image_2": ("IMAGE",),
|
||||
"image_3": ("IMAGE",),
|
||||
},
|
||||
"hidden": {
|
||||
"unique_id": "UNIQUE_ID",
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "STRING")
|
||||
RETURN_NAMES = ("image", "show_price")
|
||||
FUNCTION = "edit_image"
|
||||
CATEGORY = "PD_Tools/Image_Generation"
|
||||
|
||||
def edit_image(
|
||||
self,
|
||||
api_key,
|
||||
prompt,
|
||||
model,
|
||||
image_1,
|
||||
resolution,
|
||||
quality,
|
||||
number_of_images,
|
||||
seed,
|
||||
aspect_ratio,
|
||||
unique_id,
|
||||
image_2=None,
|
||||
image_3=None,
|
||||
):
|
||||
if not api_key or not api_key.strip():
|
||||
return (empty_image(), "Error: ComfyUI API Key is required")
|
||||
if not prompt or not prompt.strip():
|
||||
return (empty_image(), "Error: prompt is required")
|
||||
|
||||
images = flatten_images(image_1, image_2, image_3)
|
||||
max_images = MAX_INPUT_IMAGES[model]
|
||||
if len(images) > max_images:
|
||||
return (empty_image(), f"Error: {model} supports at most {max_images} input image(s)")
|
||||
if aspect_ratio != "auto" and model in ASPECT_RATIO_NEEDS_MULTIPLE and len(images) == 1:
|
||||
return (empty_image(), "Error: custom aspect ratio requires multiple input images for this model")
|
||||
|
||||
payload = {
|
||||
"model": model,
|
||||
"images": [{"url": tensor_to_data_url(image)} for image in images],
|
||||
"prompt": prompt.strip(),
|
||||
"resolution": resolution.lower(),
|
||||
"n": number_of_images,
|
||||
"seed": seed,
|
||||
"aspect_ratio": None if aspect_ratio == "auto" or model == "grok-imagine-image-pro" else aspect_ratio,
|
||||
"quality": quality if model in QUALITY_MODELS else None,
|
||||
}
|
||||
headers = {
|
||||
"X-API-KEY": api_key.strip(),
|
||||
"Accept": "application/json",
|
||||
"Content-Type": "application/json",
|
||||
"Connection": "close",
|
||||
"User-Agent": "ComfyUI-PD-Node/2.1",
|
||||
}
|
||||
|
||||
start_time = time.monotonic()
|
||||
session = requests.Session()
|
||||
try:
|
||||
response = request_with_interrupt(
|
||||
session,
|
||||
session.post,
|
||||
f"{BASE_URL}{EDIT_ENDPOINT}",
|
||||
unique_id,
|
||||
"Grok image editing",
|
||||
start_time,
|
||||
headers=headers,
|
||||
json=payload,
|
||||
timeout=600,
|
||||
)
|
||||
if response.status_code != 200:
|
||||
return (empty_image(), f"API Error {response.status_code}: {response.text[:500]}")
|
||||
|
||||
result = response.json()
|
||||
output_images = []
|
||||
for item in result.get("data") or []:
|
||||
if item.get("url"):
|
||||
download = request_with_interrupt(
|
||||
session,
|
||||
session.get,
|
||||
item["url"],
|
||||
unique_id,
|
||||
"Downloading result",
|
||||
start_time,
|
||||
timeout=120,
|
||||
)
|
||||
download.raise_for_status()
|
||||
output_images.append(load_image_to_tensor(download.content))
|
||||
elif item.get("b64_json"):
|
||||
output_images.append(load_image_to_tensor(base64.b64decode(item["b64_json"])))
|
||||
|
||||
if not output_images:
|
||||
return (empty_image(), f"Error: response contained no images: {str(result)[:500]}")
|
||||
|
||||
image = torch.cat(output_images, dim=0)
|
||||
price = format_price(result, response, model, resolution, quality, number_of_images)
|
||||
elapsed = time.monotonic() - start_time
|
||||
send_status(unique_id, f"Completed\nTime elapsed: {int(elapsed)}s")
|
||||
print(f"[PD_GROK_EDIT] Completed {len(output_images)} image(s) in {elapsed:.2f}s")
|
||||
return (image, price)
|
||||
except requests.exceptions.Timeout:
|
||||
send_status(unique_id, f"Timed out\nTime elapsed: {int(time.monotonic() - start_time)}s")
|
||||
return (empty_image(), "Error: request timeout (600s)")
|
||||
except requests.exceptions.ConnectionError as error:
|
||||
send_status(unique_id, f"Connection failed\nTime elapsed: {int(time.monotonic() - start_time)}s")
|
||||
return (empty_image(), f"Error: connection failed - {error}")
|
||||
except Exception as error:
|
||||
send_status(unique_id, f"Failed\nTime elapsed: {int(time.monotonic() - start_time)}s")
|
||||
return (empty_image(), f"Error: {error}")
|
||||
finally:
|
||||
session.close()
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"PDGrokImageEditAPIKey": PDGrokImageEditAPIKey,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"PDGrokImageEditAPIKey": "grok edit (comfyapikey)",
|
||||
}
|
||||
@@ -0,0 +1,326 @@
|
||||
"""Seedream 5.0 Pro layer separation using a ComfyUI API key."""
|
||||
|
||||
import concurrent.futures
|
||||
import io
|
||||
import logging
|
||||
|
||||
import numpy as np
|
||||
import requests
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from ._image_gpt_common import tensor_to_bytes
|
||||
|
||||
|
||||
BASE_URL = "https://api.comfy.org"
|
||||
IMAGE_ENDPOINT = "/proxy/byteplus/api/v3/images/generations"
|
||||
MODEL = "seedream-5-0-pro-260628"
|
||||
|
||||
|
||||
def _headers(api_key):
|
||||
return {
|
||||
"X-API-KEY": api_key,
|
||||
"Accept": "application/json",
|
||||
"Content-Type": "application/json",
|
||||
"Connection": "close",
|
||||
"User-Agent": "ComfyUI-PD-Node/2.1",
|
||||
}
|
||||
|
||||
|
||||
def _response_json(response, operation):
|
||||
if response.status_code not in (200, 201):
|
||||
raise RuntimeError(f"{operation} error {response.status_code}: {response.text[:500]}")
|
||||
return response.json()
|
||||
|
||||
|
||||
def _upload_image(api_key, image):
|
||||
image_bytes = tensor_to_bytes(image, max_pixels=2048 * 2048).getvalue()
|
||||
create_response = requests.post(
|
||||
f"{BASE_URL}/customers/storage",
|
||||
headers=_headers(api_key),
|
||||
json={"file_name": "seedream-layer-input.png", "content_type": "image/png"},
|
||||
timeout=120,
|
||||
)
|
||||
upload = _response_json(create_response, "ComfyUI storage")
|
||||
upload_url = upload.get("upload_url")
|
||||
download_url = upload.get("download_url")
|
||||
if not upload_url or not download_url:
|
||||
raise RuntimeError(f"ComfyUI storage response is missing URLs: {str(upload)[:500]}")
|
||||
|
||||
put_response = requests.put(
|
||||
upload_url,
|
||||
headers={"Content-Type": "image/png"},
|
||||
data=image_bytes,
|
||||
timeout=300,
|
||||
)
|
||||
if put_response.status_code not in (200, 201, 204):
|
||||
raise RuntimeError(f"Image upload error {put_response.status_code}: {put_response.text[:500]}")
|
||||
return download_url
|
||||
|
||||
|
||||
def _download_image(url, api_key, mode):
|
||||
request_url = f"{BASE_URL}{url}" if url.startswith("/") else url
|
||||
headers = _headers(api_key) if url.startswith("/") else {}
|
||||
response = requests.get(request_url, headers=headers, timeout=180)
|
||||
response.raise_for_status()
|
||||
image = Image.open(io.BytesIO(response.content)).convert(mode)
|
||||
array = np.asarray(image).astype(np.float32) / 255.0
|
||||
return torch.from_numpy(array.copy())
|
||||
|
||||
|
||||
def _z_index(item):
|
||||
value = item.get("z_index")
|
||||
if isinstance(value, bool):
|
||||
return 1_000_000
|
||||
if isinstance(value, (int, float)):
|
||||
return int(value)
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
return int(value.strip())
|
||||
except ValueError:
|
||||
pass
|
||||
return 1_000_000
|
||||
|
||||
|
||||
class PDSeedream5LayerSeparationAPIKey:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"comfy_api_key": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False,
|
||||
"placeholder": "Your ComfyUI API Key",
|
||||
}),
|
||||
"image": ("IMAGE", {
|
||||
"tooltip": "Exactly one image, at least 512x512 pixels.",
|
||||
}),
|
||||
"prompt": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": True,
|
||||
"tooltip": "Leave empty to separate all major elements automatically.",
|
||||
}),
|
||||
"size": (["auto", "1K", "1.5K", "2K"], {"default": "auto"}),
|
||||
"seed": ("INT", {
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 2147483647,
|
||||
"step": 1,
|
||||
"control_after_generate": True,
|
||||
}),
|
||||
"prompt_optimization": (["standard", "fast"], {"default": "standard"}),
|
||||
"watermark": ("BOOLEAN", {"default": False}),
|
||||
"crop_layers": ("BOOLEAN", {
|
||||
"default": False,
|
||||
"label_on": "minimal size",
|
||||
"label_off": "full canvas",
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "MASK", "IMAGE", "MASK", "BOUNDING_BOX", "LAYERS")
|
||||
RETURN_NAMES = ("base_image", "base_mask", "layers", "masks", "bboxes", "layer_stack")
|
||||
FUNCTION = "separate_layers"
|
||||
CATEGORY = "PD_Tools/Image_Generation"
|
||||
|
||||
def separate_layers(
|
||||
self,
|
||||
comfy_api_key,
|
||||
image,
|
||||
prompt,
|
||||
size,
|
||||
seed,
|
||||
prompt_optimization,
|
||||
watermark,
|
||||
crop_layers,
|
||||
):
|
||||
api_key = (comfy_api_key or "").strip()
|
||||
if not api_key:
|
||||
raise ValueError("ComfyUI API Key is required")
|
||||
if image.ndim != 4 or image.shape[0] != 1:
|
||||
raise ValueError("Only a single input image is supported")
|
||||
|
||||
height, width = image.shape[1:3]
|
||||
if width < 512 or height < 512:
|
||||
raise ValueError(f"Input image must be at least 512x512 pixels, got {width}x{height}")
|
||||
ratio = width / height
|
||||
if ratio < 1 / 16 or ratio > 16:
|
||||
raise ValueError(f"Input image aspect ratio must be between 1:16 and 16:1, got {width}:{height}")
|
||||
|
||||
image_url = _upload_image(api_key, image)
|
||||
payload = {
|
||||
"model": MODEL,
|
||||
"prompt": prompt.strip() or None,
|
||||
"image": image_url,
|
||||
"size": size,
|
||||
"seed": int(seed),
|
||||
"response_format": "url",
|
||||
"output_format": "png",
|
||||
"layer_decomposition": True,
|
||||
"watermark": bool(watermark),
|
||||
"optimize_prompt_options": {"mode": prompt_optimization},
|
||||
}
|
||||
response = requests.post(
|
||||
f"{BASE_URL}{IMAGE_ENDPOINT}",
|
||||
headers=_headers(api_key),
|
||||
json=payload,
|
||||
timeout=1800,
|
||||
)
|
||||
result = _response_json(response, "Seedream layer separation")
|
||||
if result.get("error"):
|
||||
error = result["error"]
|
||||
raise RuntimeError(
|
||||
f"ByteDance request failed. Code: {error.get('code')}, message: {error.get('message')}"
|
||||
)
|
||||
|
||||
data = [item for item in (result.get("data") or []) if isinstance(item, dict)]
|
||||
if not data or "url" not in data[0]:
|
||||
raise RuntimeError("Unexpected response: no base image returned")
|
||||
base_item = data[0]
|
||||
if _z_index(base_item) not in (0, 1_000_000):
|
||||
raise RuntimeError("Unexpected response: the first item is not the base image")
|
||||
|
||||
layer_items = [item for item in data[1:] if "url" in item]
|
||||
if not layer_items:
|
||||
raise RuntimeError("The model returned no layers. Try a different prompt or input image")
|
||||
layer_items.sort(key=_z_index)
|
||||
|
||||
base_image = _download_image(str(base_item["url"]), api_key, "RGB").unsqueeze(0).contiguous()
|
||||
height, width = base_image.shape[1:3]
|
||||
specs = []
|
||||
for item in layer_items:
|
||||
flags = []
|
||||
bbox = item.get("bounding_box")
|
||||
absolute = bbox.get("absolute") if isinstance(bbox, dict) else None
|
||||
if (
|
||||
isinstance(absolute, (list, tuple))
|
||||
and len(absolute) == 4
|
||||
and all(isinstance(value, (int, float)) and not isinstance(value, bool) for value in absolute)
|
||||
):
|
||||
left, top, right, bottom = (int(round(value)) for value in absolute)
|
||||
rect_w, rect_h = right - left, bottom - top
|
||||
if rect_w > width or rect_h > height:
|
||||
rect_w, rect_h = min(rect_w, width), min(rect_h, height)
|
||||
flags.append("bbox_clamped")
|
||||
if rect_w <= 0 or rect_h <= 0:
|
||||
flags.append("bbox_degenerate")
|
||||
else:
|
||||
left, top, rect_w, rect_h = 0, 0, width, height
|
||||
flags.append("bbox_missing")
|
||||
specs.append({
|
||||
"item": item,
|
||||
"flags": flags,
|
||||
"left": left,
|
||||
"top": top,
|
||||
"rect_w": rect_w,
|
||||
"rect_h": rect_h,
|
||||
})
|
||||
|
||||
if crop_layers:
|
||||
canvas_w = max((spec["rect_w"] for spec in specs if "bbox_degenerate" not in spec["flags"]), default=1)
|
||||
canvas_h = max((spec["rect_h"] for spec in specs if "bbox_degenerate" not in spec["flags"]), default=1)
|
||||
else:
|
||||
canvas_w, canvas_h = width, height
|
||||
|
||||
base_mask = torch.zeros((1, height, width))
|
||||
layers = torch.zeros((len(specs), canvas_h, canvas_w, 3))
|
||||
masks = torch.ones((len(specs), canvas_h, canvas_w))
|
||||
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
|
||||
downloads = list(executor.map(
|
||||
lambda spec: _download_image(str(spec["item"]["url"]), api_key, "RGBA"),
|
||||
specs,
|
||||
))
|
||||
|
||||
stack_items = [{
|
||||
"image": base_image,
|
||||
"type": "raster",
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"z_index": 0,
|
||||
"name": "background",
|
||||
}]
|
||||
boxes = []
|
||||
for index, (spec, rgba) in enumerate(zip(specs, downloads)):
|
||||
item = spec["item"]
|
||||
flags = spec["flags"]
|
||||
left, top = spec["left"], spec["top"]
|
||||
rect_w, rect_h = spec["rect_w"], spec["rect_h"]
|
||||
native_size = f"{rgba.shape[1]}x{rgba.shape[0]}"
|
||||
|
||||
if "bbox_degenerate" not in flags:
|
||||
if (rgba.shape[1], rgba.shape[0]) != (rect_w, rect_h):
|
||||
rgba = rgba.clone()
|
||||
rgba[..., :3] *= rgba[..., 3:4]
|
||||
rgba = torch.nn.functional.interpolate(
|
||||
rgba.permute(2, 0, 1).unsqueeze(0),
|
||||
size=(rect_h, rect_w),
|
||||
mode="bilinear",
|
||||
antialias=True,
|
||||
).squeeze(0).permute(1, 2, 0)
|
||||
alpha = rgba[..., 3:4]
|
||||
rgba = torch.cat([rgba[..., :3] / alpha.clamp(min=1e-6), alpha], dim=-1).clamp(0, 1)
|
||||
flags.append("resized_to_bbox")
|
||||
|
||||
if crop_layers:
|
||||
layers[index, :rect_h, :rect_w] = rgba[..., :3]
|
||||
masks[index, :rect_h, :rect_w] = 1.0 - rgba[..., 3]
|
||||
else:
|
||||
x0, y0 = max(left, 0), max(top, 0)
|
||||
x1, y1 = min(left + rect_w, width), min(top + rect_h, height)
|
||||
if x0 < x1 and y0 < y1:
|
||||
patch = rgba[y0 - top:y1 - top, x0 - left:x1 - left]
|
||||
layers[index, y0:y1, x0:x1] = patch[..., :3]
|
||||
masks[index, y0:y1, x0:x1] = 1.0 - patch[..., 3]
|
||||
else:
|
||||
flags.append("bbox_out_of_canvas")
|
||||
|
||||
stack_item = {
|
||||
"image": rgba[..., :3].unsqueeze(0).contiguous(),
|
||||
"type": "raster",
|
||||
"x": left,
|
||||
"y": top,
|
||||
"z_index": _z_index(item) if _z_index(item) != 1_000_000 else index + 1,
|
||||
"mask": (1.0 - rgba[..., 3]).unsqueeze(0),
|
||||
}
|
||||
if isinstance(item.get("name"), str):
|
||||
stack_item["name"] = item["name"]
|
||||
stack_items.append(stack_item)
|
||||
|
||||
abnormal = [flag for flag in flags if flag != "resized_to_bbox"]
|
||||
if abnormal:
|
||||
logging.warning(
|
||||
"Seedream layer separation: layer %d (%r) flagged %s",
|
||||
index + 1,
|
||||
item.get("name"),
|
||||
", ".join(abnormal),
|
||||
)
|
||||
z_index = _z_index(item)
|
||||
rect_x, rect_y = (0, 0) if crop_layers else (left, top)
|
||||
boxes.append({
|
||||
"x": left if crop_layers else 0,
|
||||
"y": top if crop_layers else 0,
|
||||
"width": canvas_w,
|
||||
"height": canvas_h,
|
||||
"metadata": {
|
||||
"name": item.get("name"),
|
||||
"desc": item.get("description"),
|
||||
"z_index": z_index if z_index != 1_000_000 else None,
|
||||
"native_size": native_size,
|
||||
"content_rect": [rect_x, rect_y, max(rect_w, 0), max(rect_h, 0)],
|
||||
"flags": flags,
|
||||
},
|
||||
})
|
||||
|
||||
bboxes = [boxes]
|
||||
layer_stack = {"version": 1, "canvas": (width, height), "layers": stack_items}
|
||||
return (base_image, base_mask, layers, masks, bboxes, layer_stack)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"PDSeedream5LayerSeparationAPIKey": PDSeedream5LayerSeparationAPIKey,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"PDSeedream5LayerSeparationAPIKey": "PD: Seedream 5.0 Pro Layer Separation (comfyui_apikey)",
|
||||
}
|
||||
@@ -0,0 +1,276 @@
|
||||
"""Google Gemini multimodal text node using a ComfyUI API key."""
|
||||
|
||||
import queue
|
||||
import threading
|
||||
import time
|
||||
|
||||
import requests
|
||||
|
||||
from comfy_api.latest import Types
|
||||
from comfy_api_nodes.util import audio_to_base64_string, tensor_to_base64_string, video_to_base64_string
|
||||
from comfy_api_nodes.util._helpers import is_processing_interrupted
|
||||
from comfy_api_nodes.util.common_exceptions import ProcessingInterrupted
|
||||
|
||||
from ._image_gpt_common import USD_TO_CNY_RATE, format_credits_value
|
||||
|
||||
|
||||
GEMINI_MODELS = {
|
||||
"Gemini 3.1 Pro": "gemini-3.1-pro-preview",
|
||||
"Gemini 3.1 Flash-Lite": "gemini-3.1-flash-lite-preview",
|
||||
}
|
||||
|
||||
# USD per million tokens: (input, text/thinking output).
|
||||
GEMINI_TOKEN_PRICES = {
|
||||
"gemini-3.1-pro-preview": (2.0, 12.0),
|
||||
"gemini-3.1-flash-lite-preview": (0.25, 1.5),
|
||||
}
|
||||
|
||||
|
||||
def _part_to_dict(part):
|
||||
if hasattr(part, "model_dump"):
|
||||
return part.model_dump(by_alias=True, exclude_none=True)
|
||||
return part
|
||||
|
||||
|
||||
def _response_text(result):
|
||||
text = []
|
||||
for candidate in result.get("candidates") or []:
|
||||
for part in (candidate.get("content") or {}).get("parts") or []:
|
||||
if part.get("text"):
|
||||
text.append(str(part["text"]))
|
||||
return "\n".join(text)
|
||||
|
||||
|
||||
def _estimated_price(model_id, usage):
|
||||
rates = GEMINI_TOKEN_PRICES.get(model_id)
|
||||
if rates is None or not isinstance(usage, dict):
|
||||
return None
|
||||
|
||||
input_tokens = int(usage.get("promptTokenCount") or 0)
|
||||
output_tokens = int(usage.get("candidatesTokenCount") or 0)
|
||||
thinking_tokens = int(usage.get("thoughtsTokenCount") or 0)
|
||||
input_rate, output_rate = rates
|
||||
return (input_tokens * input_rate + (output_tokens + thinking_tokens) * output_rate) / 1_000_000.0
|
||||
|
||||
|
||||
def _post_with_interrupt(url, headers, payload):
|
||||
result_queue = queue.Queue(maxsize=1)
|
||||
|
||||
def send_request():
|
||||
try:
|
||||
result_queue.put(requests.post(url, headers=headers, json=payload, timeout=300))
|
||||
except Exception as error:
|
||||
result_queue.put(error)
|
||||
|
||||
threading.Thread(target=send_request, daemon=True).start()
|
||||
while True:
|
||||
if is_processing_interrupted():
|
||||
raise ProcessingInterrupted("Task cancelled")
|
||||
try:
|
||||
result = result_queue.get(timeout=0.25)
|
||||
except queue.Empty:
|
||||
continue
|
||||
if isinstance(result, Exception):
|
||||
raise result
|
||||
return result
|
||||
|
||||
|
||||
class PDGeminiAPIKey:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"comfy_api_key": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False,
|
||||
"placeholder": "Your ComfyUI API Key",
|
||||
"tooltip": "The key is stored in the workflow when the workflow is saved.",
|
||||
}),
|
||||
"prompt": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": True,
|
||||
}),
|
||||
"model": (list(GEMINI_MODELS), {
|
||||
"default": "Gemini 3.1 Pro",
|
||||
}),
|
||||
"thinking_level": (["LOW", "HIGH"], {
|
||||
"default": "HIGH",
|
||||
}),
|
||||
"temperature": ("FLOAT", {
|
||||
"default": 1.0,
|
||||
"min": 0.0,
|
||||
"max": 2.0,
|
||||
"step": 0.01,
|
||||
}),
|
||||
"top_p": ("FLOAT", {
|
||||
"default": 0.95,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.01,
|
||||
}),
|
||||
"max_output_tokens": ("INT", {
|
||||
"default": 32768,
|
||||
"min": 16,
|
||||
"max": 65536,
|
||||
"step": 1,
|
||||
}),
|
||||
"seed": ("INT", {
|
||||
"default": 42,
|
||||
"min": 0,
|
||||
"max": 2147483647,
|
||||
"control_after_generate": True,
|
||||
}),
|
||||
"system_prompt": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": True,
|
||||
}),
|
||||
},
|
||||
"optional": {
|
||||
"image_1": ("IMAGE",),
|
||||
"audio_1": ("AUDIO",),
|
||||
"video_1": ("VIDEO",),
|
||||
"files": ("GEMINI_INPUT_FILES",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING", "STRING")
|
||||
RETURN_NAMES = ("text", "price")
|
||||
FUNCTION = "generate_text"
|
||||
CATEGORY = "PD_Tools/Text_Generation"
|
||||
|
||||
def generate_text(
|
||||
self,
|
||||
comfy_api_key,
|
||||
prompt,
|
||||
model,
|
||||
thinking_level,
|
||||
temperature,
|
||||
top_p,
|
||||
max_output_tokens,
|
||||
seed,
|
||||
system_prompt,
|
||||
image_1=None,
|
||||
audio_1=None,
|
||||
video_1=None,
|
||||
files=None,
|
||||
):
|
||||
comfy_api_key = comfy_api_key.strip()
|
||||
if not comfy_api_key:
|
||||
return ("Error: ComfyUI API Key is required", "")
|
||||
if not prompt.strip():
|
||||
return ("Error: prompt is required", "")
|
||||
|
||||
model_id = GEMINI_MODELS[model]
|
||||
parts = [{"text": prompt}]
|
||||
|
||||
if image_1 is not None:
|
||||
if is_processing_interrupted():
|
||||
raise ProcessingInterrupted("Task cancelled")
|
||||
images = image_1 if image_1.ndim == 4 else image_1.unsqueeze(0)
|
||||
for image in images:
|
||||
parts.append({
|
||||
"inlineData": {
|
||||
"mimeType": "image/png",
|
||||
"data": tensor_to_base64_string(image),
|
||||
}
|
||||
})
|
||||
|
||||
if audio_1 is not None:
|
||||
if is_processing_interrupted():
|
||||
raise ProcessingInterrupted("Task cancelled")
|
||||
waveform = audio_1["waveform"]
|
||||
for index in range(waveform.shape[0]):
|
||||
audio = {
|
||||
"waveform": waveform[index].unsqueeze(0),
|
||||
"sample_rate": audio_1["sample_rate"],
|
||||
}
|
||||
parts.append({
|
||||
"inlineData": {
|
||||
"mimeType": "audio/mpeg",
|
||||
"data": audio_to_base64_string(audio, container_format="mp3", codec_name="libmp3lame"),
|
||||
}
|
||||
})
|
||||
|
||||
if video_1 is not None:
|
||||
if is_processing_interrupted():
|
||||
raise ProcessingInterrupted("Task cancelled")
|
||||
parts.append({
|
||||
"inlineData": {
|
||||
"mimeType": "video/mp4",
|
||||
"data": video_to_base64_string(
|
||||
video_1,
|
||||
container_format=Types.VideoContainer.MP4,
|
||||
codec=Types.VideoCodec.H264,
|
||||
),
|
||||
}
|
||||
})
|
||||
|
||||
if files:
|
||||
parts.extend(_part_to_dict(part) for part in files)
|
||||
|
||||
payload = {
|
||||
"contents": [{"role": "user", "parts": parts}],
|
||||
"generationConfig": {
|
||||
"temperature": float(temperature),
|
||||
"topP": float(top_p),
|
||||
"maxOutputTokens": int(max_output_tokens),
|
||||
"seed": int(seed) if seed > 0 else None,
|
||||
"thinkingConfig": {"thinkingLevel": thinking_level},
|
||||
},
|
||||
}
|
||||
if system_prompt.strip():
|
||||
payload["systemInstruction"] = {"parts": [{"text": system_prompt.strip()}]}
|
||||
|
||||
url = f"https://api.comfy.org/proxy/vertexai/gemini/{model_id}"
|
||||
headers = {
|
||||
"X-API-KEY": comfy_api_key,
|
||||
"Content-Type": "application/json",
|
||||
"User-Agent": "ComfyUI-PD-Node/2.1",
|
||||
}
|
||||
|
||||
started = time.time()
|
||||
try:
|
||||
response = _post_with_interrupt(url, headers, payload)
|
||||
if response.status_code != 200:
|
||||
return (f"API Error {response.status_code}: {response.text[:1000]}", "")
|
||||
|
||||
result = response.json()
|
||||
text = _response_text(result)
|
||||
if not text:
|
||||
return (f"Error: no text output in response: {str(result)[:1000]}", "")
|
||||
|
||||
usage = result.get("usageMetadata") or {}
|
||||
usd = _estimated_price(model_id, usage)
|
||||
price_lines = [
|
||||
f"Model: {result.get('modelVersion') or model_id}",
|
||||
f"Time: {time.time() - started:.2f}s",
|
||||
f"Input tokens: {int(usage.get('promptTokenCount') or 0)}",
|
||||
f"Output tokens: {int(usage.get('candidatesTokenCount') or 0)}",
|
||||
f"Thinking tokens: {int(usage.get('thoughtsTokenCount') or 0)}",
|
||||
]
|
||||
if usd is not None:
|
||||
price_lines.extend([
|
||||
f"Estimated price (USD): ${usd:.6f}",
|
||||
f"Estimated price (RMB): 约 ¥{usd * USD_TO_CNY_RATE:.6f}",
|
||||
f"Estimated price (Credits): {format_credits_value(usd)}",
|
||||
])
|
||||
else:
|
||||
price_lines.append("Estimated price: unavailable")
|
||||
return (text, "\n".join(price_lines))
|
||||
except requests.exceptions.Timeout:
|
||||
return ("Error: request timeout (300s)", "")
|
||||
except requests.exceptions.ConnectionError as error:
|
||||
return (f"Error: connection failed - {error}", "")
|
||||
except ProcessingInterrupted:
|
||||
raise
|
||||
except Exception as error:
|
||||
return (f"Error: {error}", "")
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"PDGeminiAPIKey": PDGeminiAPIKey,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"PDGeminiAPIKey": "PD: Google Gemini (ComfyUI API Key)",
|
||||
}
|
||||
@@ -0,0 +1,2 @@
|
||||
alibabacloud_credentials>=1.0.11
|
||||
alibabacloud_facebody20191230>=6.0.0
|
||||
Reference in New Issue
Block a user