268 lines
10 KiB
Python
268 lines
10 KiB
Python
import base64
|
|
import io
|
|
import json
|
|
import re
|
|
import time
|
|
from typing import Tuple
|
|
|
|
import numpy as np
|
|
import requests
|
|
from PIL import Image
|
|
from requests.adapters import HTTPAdapter
|
|
import torch
|
|
|
|
|
|
def tensor2pil(t_image: torch.Tensor) -> Image.Image:
|
|
if not isinstance(t_image, torch.Tensor):
|
|
raise TypeError("输入必须是 torch.Tensor")
|
|
arr = t_image.cpu().numpy()
|
|
if arr.ndim == 4:
|
|
arr = arr[0]
|
|
arr = np.clip(255.0 * arr.squeeze(), 0, 255).astype(np.uint8)
|
|
if arr.ndim == 3 and arr.shape[0] in (1, 3):
|
|
arr = np.transpose(arr, (1, 2, 0))
|
|
return Image.fromarray(arr)
|
|
|
|
|
|
class API_caption:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"api_type": (["Siliconflow", "T8zhenzhen", "OpenRouter", "Other"], {"default": "Siliconflow"}),
|
|
"api_url": ("STRING", {"default": "<url>"}),
|
|
"API_Key": ("STRING", {"default": "<your_key>"}),
|
|
"model_name": ("STRING", {"default": "Qwen/Qwen3-VL-32B-Instruct"}),
|
|
"image": ("IMAGE",),
|
|
"prompt": (
|
|
"STRING",
|
|
{
|
|
"default": "You are a professional AI image generation prompt engineer. Please describe in detail the main body, foreground, mid-ground, background, composition, visual guidance, color tone, and light and shadow atmosphere of this image, and create an image prompt with depth, atmosphere, and artistic appeal. Requirements: Chinese prompt, no description of image watermark, no irrelevant words or symbols, no summary, limited to 800 words.",
|
|
"multiline": True,
|
|
"rows": 4,
|
|
},
|
|
),
|
|
"output_language": (["Chinese", "English"], {"default": "Chinese"}),
|
|
"thinking_mode": ("BOOLEAN", {"default": False}),
|
|
"temperature": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 2.0, "step": 0.01}),
|
|
"max_tokens": ("INT", {"default": 258, "min": 125, "max": 4096}),
|
|
# 使用 control_after_generate 支持 fixed / increment / decrement / randomize 等控制选项
|
|
"noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "control_after_generate": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("text",)
|
|
FUNCTION = "generate"
|
|
CATEGORY = "DaNodes/API"
|
|
|
|
@staticmethod
|
|
def _strip_thinking_blocks(text: str) -> str:
|
|
if not isinstance(text, str):
|
|
return text
|
|
return re.sub(r"(?is)<think>.*?</think>\s*", "", text).strip()
|
|
|
|
@staticmethod
|
|
def _apply_thinking_config(payload: dict, api_type: str, thinking_mode: bool) -> dict:
|
|
if api_type == "OpenRouter":
|
|
payload["reasoning"] = (
|
|
{"enabled": True, "exclude": False}
|
|
if thinking_mode
|
|
else {"effort": "none", "exclude": True}
|
|
)
|
|
else:
|
|
payload["enable_thinking"] = bool(thinking_mode)
|
|
return payload
|
|
|
|
def generate(
|
|
self,
|
|
api_type: str,
|
|
API_Key: str,
|
|
model_name: str,
|
|
image: torch.Tensor,
|
|
prompt: str,
|
|
output_language: str,
|
|
thinking_mode: bool = False,
|
|
temperature: float = 0.5,
|
|
max_tokens: int = 512,
|
|
api_url: str = "",
|
|
noise_seed: int = 0,
|
|
) -> Tuple[str]:
|
|
try:
|
|
# 复用 HTTP 连接以减少握手开销
|
|
session = getattr(self, "_session", None)
|
|
if session is None:
|
|
session = requests.Session()
|
|
adapter = HTTPAdapter(pool_connections=4, pool_maxsize=4)
|
|
session.mount("https://", adapter)
|
|
session.mount("http://", adapter)
|
|
self._session = session
|
|
|
|
# 根据 provider 选择实际的 URL
|
|
provider_map = {
|
|
"OpenRouter": "https://openrouter.ai/api/v1/chat/completions",
|
|
"Siliconflow": "https://api.siliconflow.cn/v1/chat/completions",
|
|
"T8zhenzhen": "https://api.bltcy.ai/v1/chat/completions",
|
|
}
|
|
if api_type == "Other":
|
|
api_url = api_url or ""
|
|
else:
|
|
api_url = provider_map.get(api_type, api_url)
|
|
if not api_url:
|
|
return ("请求地址为空，请填写或选择有效的 API 类型",)
|
|
|
|
# image -> base64（保持原始尺寸，使用无损 PNG 编码）
|
|
pil = tensor2pil(image)
|
|
buf = io.BytesIO()
|
|
pil.save(buf, format="PNG")
|
|
img_b64 = base64.b64encode(buf.getvalue()).decode()
|
|
|
|
# 提示词
|
|
if output_language == "Chinese":
|
|
prompt_full = f"{prompt} 请用Chinese返回描述"
|
|
else:
|
|
prompt_full = f"{prompt} Please return the description in English."
|
|
|
|
# 按照示例将图片放在 messages[0].content 中
|
|
image_data_uri = f"data:image/png;base64,{img_b64}"
|
|
messages = [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "text", "text": prompt_full},
|
|
{"type": "image_url", "image_url": {"url": image_data_uri}},
|
|
],
|
|
}
|
|
]
|
|
|
|
payload = {
|
|
"model": model_name,
|
|
"stream": False,
|
|
"messages": messages,
|
|
"max_tokens": int(max_tokens),
|
|
"temperature": float(temperature),
|
|
}
|
|
payload = self._apply_thinking_config(payload, api_type, bool(thinking_mode))
|
|
# 如果后端支持 seed 字段，则可用此随机种子控制多次运行的一致性/随机性
|
|
try:
|
|
payload["seed"] = int(noise_seed)
|
|
except Exception:
|
|
pass
|
|
|
|
headers = {
|
|
"Authorization": f"Bearer {API_Key}",
|
|
"Content-Type": "application/json",
|
|
"Accept": "application/json",
|
|
"User-Agent": "ComfyUI-DaNodes/1.0",
|
|
}
|
|
|
|
# 固定使用统一 JSON 风格请求，遇到 429/5xx/读超时简单重试一次
|
|
r = None
|
|
last_err = None
|
|
thinking_fallback_used = False
|
|
for attempt in range(3):
|
|
try:
|
|
r = session.post(api_url, json=payload, headers=headers, timeout=90)
|
|
except Exception as e:
|
|
last_err = e
|
|
if attempt < 2:
|
|
time.sleep(1.0)
|
|
continue
|
|
return (f"API 请求失败（网络连接错误）: {e}",)
|
|
if not r.ok:
|
|
if (
|
|
not thinking_fallback_used
|
|
and r.status_code in (400, 422)
|
|
and ("enable_thinking" in payload or "reasoning" in payload)
|
|
):
|
|
payload = dict(payload)
|
|
payload.pop("enable_thinking", None)
|
|
payload.pop("reasoning", None)
|
|
thinking_fallback_used = True
|
|
continue
|
|
if r.status_code in (429, 500, 502, 503, 504) and attempt < 2:
|
|
time.sleep(1.0)
|
|
continue
|
|
break
|
|
|
|
if r is None:
|
|
return (f"API 请求失败: {last_err}",)
|
|
|
|
if not r.ok:
|
|
body = None
|
|
try:
|
|
body = r.json()
|
|
body = json.dumps(body, ensure_ascii=False)
|
|
except Exception:
|
|
body = r.text
|
|
return (f"API 返回错误状态 {r.status_code}: {body}",)
|
|
|
|
try:
|
|
result = r.json()
|
|
except Exception as e:
|
|
return (f"无法解析 API 返回的 JSON: {e}; body: {r.text}",)
|
|
|
|
# 提取文本回复
|
|
def extract_text_from_content(content):
|
|
if isinstance(content, str):
|
|
return content
|
|
if isinstance(content, dict):
|
|
return content.get("text") or content.get("content") or str(content)
|
|
if isinstance(content, list):
|
|
parts = []
|
|
for block in content:
|
|
if not isinstance(block, dict):
|
|
parts.append(str(block))
|
|
continue
|
|
btype = block.get("type", "")
|
|
if btype in ("text", "paragraph"):
|
|
parts.append(block.get("text", ""))
|
|
elif btype == "image_url":
|
|
continue
|
|
else:
|
|
parts.append(block.get("text") or block.get("content") or "")
|
|
return " ".join(p for p in parts if p).strip()
|
|
return str(content)
|
|
|
|
caption = ""
|
|
if isinstance(result, dict):
|
|
choices = result.get("choices") or []
|
|
if choices:
|
|
first = choices[0]
|
|
if isinstance(first, dict):
|
|
msg = first.get("message") or first.get("delta") or None
|
|
if isinstance(msg, dict) and "content" in msg:
|
|
caption = extract_text_from_content(msg.get("content"))
|
|
else:
|
|
if "text" in first and isinstance(first.get("text"), str):
|
|
caption = first.get("text")
|
|
else:
|
|
caption = extract_text_from_content(first)
|
|
else:
|
|
caption = str(first)
|
|
else:
|
|
caption = result.get("caption") or result.get("text") or ""
|
|
else:
|
|
caption = str(result)
|
|
|
|
caption = self._strip_thinking_blocks(caption)
|
|
if caption == "":
|
|
try:
|
|
pretty = json.dumps(result, ensure_ascii=False)
|
|
except Exception:
|
|
pretty = str(result)
|
|
return (f"无法从 API 响应解析出文本。返回 JSON: {pretty}",)
|
|
|
|
return (caption,)
|
|
except Exception as e:
|
|
# 捕获整个函数体的未处理异常，避免 ComfyUI 无输出
|
|
return (f"节点内部异常: {e}",)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"API_caption": API_caption
|
|
}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"API_caption": "API_caption☀"
|
|
}
|