280 lines
11 KiB
Python
280 lines
11 KiB
Python
import os
|
|
import base64
|
|
import json
|
|
import re
|
|
from typing import Tuple
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
|
|
import requests
|
|
|
|
|
|
class Batch_API_caption:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"input_dir": ("STRING", {"default": ""}),
|
|
"output_dir": ("STRING", {"default": ""}),
|
|
"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"}),
|
|
"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}),
|
|
"concurrency": ("INT", {"default": 6, "min": 1, "max": 32}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING")
|
|
RETURN_NAMES = ("log", "save_path")
|
|
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
|
|
|
|
@staticmethod
|
|
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)
|
|
|
|
@staticmethod
|
|
def _build_payload(model_name: str, prompt_full: str, img_data_uri: str, temperature: float, max_tokens: int):
|
|
messages = [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "text", "text": prompt_full},
|
|
{"type": "image_url", "image_url": {"url": img_data_uri}},
|
|
],
|
|
}
|
|
]
|
|
return {
|
|
"model": model_name,
|
|
"stream": False,
|
|
"messages": messages,
|
|
"max_tokens": int(max_tokens),
|
|
"temperature": float(temperature),
|
|
}
|
|
|
|
@staticmethod
|
|
def _post(api_url: str, api_key: str, payload: dict, timeout_sec: float = 90.0, retries: int = 0):
|
|
headers = {
|
|
"Authorization": f"Bearer {api_key}",
|
|
"Content-Type": "application/json",
|
|
"Accept": "application/json",
|
|
"User-Agent": "ComfyUI-DaNodes/1.0",
|
|
}
|
|
attempt = 0
|
|
last_err = None
|
|
result = None
|
|
thinking_fallback_used = False
|
|
while attempt <= max(0, int(retries)):
|
|
try:
|
|
r = requests.post(api_url, json=payload, headers=headers, timeout=float(timeout_sec))
|
|
except Exception as e:
|
|
last_err = f"请求失败: {e}"
|
|
attempt += 1
|
|
continue
|
|
|
|
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
|
|
try:
|
|
body = r.json()
|
|
body = json.dumps(body, ensure_ascii=False)
|
|
except Exception:
|
|
body = r.text
|
|
last_err = f"HTTP {r.status_code}: {body}"
|
|
attempt += 1
|
|
continue
|
|
|
|
try:
|
|
result = r.json()
|
|
except Exception as e:
|
|
last_err = f"解析 JSON 失败: {e}; body: {r.text}"
|
|
attempt += 1
|
|
continue
|
|
# 成功
|
|
break
|
|
|
|
if result is None:
|
|
return False, (last_err or "请求失败")
|
|
|
|
# 解析文本
|
|
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 = Batch_API_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 = Batch_API_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 = Batch_API_caption._strip_thinking_blocks(caption)
|
|
if caption == "":
|
|
try:
|
|
pretty = json.dumps(result, ensure_ascii=False)
|
|
except Exception:
|
|
pretty = str(result)
|
|
return False, f"无法解析文本，返回 JSON: {pretty}"
|
|
|
|
return True, caption
|
|
|
|
def generate(self, input_dir: str, output_dir: str, api_type: str, api_url: str, API_Key: str, model_name: str,
|
|
prompt: str, output_language: str, thinking_mode: bool = False,
|
|
temperature: float = 0.5, max_tokens: int = 512,
|
|
concurrency: int = 4) -> Tuple[str, str]:
|
|
# 目录校验
|
|
if not input_dir:
|
|
return ("输入地址未填写", output_dir)
|
|
if not os.path.isdir(input_dir):
|
|
return (f"输入地址不存在: {input_dir}", output_dir)
|
|
# API Key 校验：为空或仍为默认占位时直接报错
|
|
if not API_Key or API_Key.strip() in ("<your_key>", "your_key", "请输入你的 Key"):
|
|
return ("API 密钥未填写，请在节点中填写有效的 API_Key 后再运行", output_dir)
|
|
if not output_dir:
|
|
output_dir = input_dir
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
# 选择 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 类型", output_dir)
|
|
|
|
# 列出图片
|
|
exts = (".jpg", ".jpeg", ".png", ".bmp", ".webp")
|
|
files = [f for f in os.listdir(input_dir) if f.lower().endswith(exts)]
|
|
if not files:
|
|
return ("输入地址中没有图片文件", output_dir)
|
|
|
|
log = {"total": len(files), "succeeded": 0, "failed": 0}
|
|
first_error = None # type: ignore[assignment]
|
|
|
|
def process_one(name: str):
|
|
in_path = os.path.join(input_dir, name)
|
|
out_txt = os.path.join(output_dir, os.path.splitext(name)[0] + ".txt")
|
|
try:
|
|
# 直接嵌入原始文件字节，不做缩放与重新编码
|
|
with open(in_path, 'rb') as fbin:
|
|
raw = fbin.read()
|
|
img_b64 = base64.b64encode(raw).decode()
|
|
ext = os.path.splitext(in_path)[1].lower()
|
|
mime = "image/jpeg" if ext in (".jpg", ".jpeg") else (
|
|
"image/png" if ext == ".png" else (
|
|
"image/webp" if ext == ".webp" else (
|
|
"image/bmp" if ext == ".bmp" else "application/octet-stream"))
|
|
)
|
|
img_data_uri = f"data:{mime};base64,{img_b64}"
|
|
|
|
if output_language == "Chinese":
|
|
prompt_full = f"{prompt} 请用Chinese返回描述"
|
|
else:
|
|
prompt_full = f"{prompt} Please return the description in English."
|
|
|
|
payload = self._build_payload(model_name, prompt_full, img_data_uri, temperature, max_tokens)
|
|
payload = self._apply_thinking_config(payload, api_type, bool(thinking_mode))
|
|
ok, resp = self._post(api_url, API_Key, payload, timeout_sec=90.0, retries=1)
|
|
|
|
if ok:
|
|
with open(out_txt, "w", encoding="utf-8") as f:
|
|
f.write(resp)
|
|
return True, name, None
|
|
else:
|
|
# 失败时不创建/不写入任何 txt 文件
|
|
return False, name, resp
|
|
except Exception as e:
|
|
return False, name, str(e)
|
|
|
|
workers = max(1, int(concurrency))
|
|
files_sorted = list(files)
|
|
try:
|
|
files_sorted.sort()
|
|
except Exception:
|
|
pass
|
|
with ThreadPoolExecutor(max_workers=workers) as ex:
|
|
futures = {ex.submit(process_one, n): n for n in files_sorted}
|
|
for fut in as_completed(futures):
|
|
ok, name, err = fut.result()
|
|
if ok:
|
|
log["succeeded"] += 1
|
|
else:
|
|
log["failed"] += 1
|
|
if first_error is None:
|
|
first_error = (err or "unknown")
|
|
|
|
if first_error is not None:
|
|
log["error"] = first_error
|
|
return (json.dumps(log, ensure_ascii=False), output_dir)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"Batch_API_caption": Batch_API_caption,
|
|
}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"Batch_API_caption": "Batch_API_caption☀",
|
|
}
|