Files
Jony ba920f2b66 Publish current ComfyUI O1Key code baseline
Replace the prior release tree with the current plugin, frontend, tests, and documentation. Document retired node IDs and the public Gitea update source.
2026-09-24 19:56:48 +08:00

692 lines
28 KiB
Python

"""Helpers for Agent chat attachments and web search."""
import base64
import html
import io
import json
import math
import os
import posixpath
import re
import xml.etree.ElementTree as ET
import zipfile
from urllib.parse import quote_plus
from PIL import Image, ImageOps
MAX_EXTRACTED_CHARS = 200_000
SEARCH_TIMEOUT_SECONDS = 12
REWRITE_TIMEOUT_SECONDS = 15
PROMPT_OPTIMIZER_MODEL = "gpt-5.6-sol"
PROMPT_OPTIMIZER_REASONING_EFFORT = "high"
PROMPT_OPTIMIZER_TIMEOUT_SECONDS = 300
PROMPT_OPTIMIZER_BODY_LIMIT_BYTES = 18 * 1024 * 1024
PROMPT_OPTIMIZER_MAX_REFERENCES = 10
PROMPT_OPTIMIZER_MAX_PROMPT_CHARS = 50_000
PROMPT_OPTIMIZER_MAX_IMAGE_BYTES = 1_200_000
PROMPT_OPTIMIZER_MAX_LONG_EDGE = 1536
PROMPT_OPTIMIZER_MIN_LONG_EDGE = 384
_USER_AGENT = (
"Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
"AppleWebKit/537.36 (KHTML, like Gecko) Chrome/126.0 Safari/537.36"
)
_PROMPT_OPTIMIZER_SYSTEM_PROMPT = """你是专业的 AI 图像生成提示词优化器。请根据用户当前提示词和按上传顺序提供的参考图,推断用户真实的文生图或图生图需求,并只输出一份可直接交给图像生成模型的最终提示词。
优化优先级:
1. 视觉元素绑定优先。把主体身份、外貌、服装、材质、动作、道具、相对位置、环境、光线、色彩、镜头和构图明确绑定到对应对象,避免形成互不关联的形容词堆叠。
2. 其次才使用指令型语言。对于图生图或编辑需求,明确区分“必须保持不变”的视觉要素和“需要改变”的目标;只强调真正重要的约束,不要反复使用强硬措辞。
3. 参考图按编号理解。识别每张参考图承担的身份、造型、风格、构图、背景或局部细节角色,不要把不同参考图中的元素错误混合。
4. 保留用户原始意图、专有名词、数量关系和明确限制;补足有助于生成的视觉信息,但不要擅自增加与需求冲突的主体或情节。
5. 文生图时,补足主体、场景、构图、镜头、光影、色彩、材质和画面层次;图生图时,优先说明参考图用途、保留项、修改项及修改后的视觉关系。
6. 使用与用户当前提示词相同的主要语言。不要解释优化过程,不要输出标题、Markdown、引号、分析、备选版本或其他附加内容,只输出最终提示词。"""
_VIDEO_PROMPT_WRITER_SYSTEM_PROMPT = """你是专业的 AI 视频生成提示词编写助手。请结合用户当前描述、生成模式、视频参数及按顺序提供的参考素材,理解用户真正想生成的视频,并只输出一份可直接提交给视频生成模型的最终提示词。
编写原则:
1. 保留用户明确指定的主体、身份、数量、外观、服装、道具、环境、动作、风格和限制,不要擅自增加会改变故事含义的新人物、新物体或新情节。
2. 每个动作都要写清楚由谁完成、作用于什么对象、动作如何开始和结束,避免无法确定主体的模糊描述。
3. 建立连续的时间过程,明确初始状态、主要动作和变化、结束状态,避免角色瞬移、物体突然出现、动作跳跃、肢体异常、身份变化或背景无原因切换。
4. 根据用户意图补充适量的景别、摄像机角度、镜头运动、对焦关系、主体运动速度和画面节奏;镜头语言必须服务于主体动作,不要堆砌互相冲突的运镜术语。
5. 强调人物身份、面部、服装、物体外观、材质、颜色、空间位置、光线方向和场景结构在整个视频中保持一致。
6. 文生视频时补全主体、场景、动作过程、镜头、光线、风格和结束状态。首帧图生视频时把首帧视为必须保持的初始状态,描述画面如何从首帧自然发展。首尾帧生视频时把两帧视为严格的开始与结束状态,设计合理连续的中间动作。多模态参考时严格按素材编号理解用途,不要混淆不同素材中的人物、外观、动作、镜头、节奏或声音。
7. 启用生成音频时,描述需要的环境声、动作声、对白或音乐氛围,并让声音与画面事件同步;未启用时不要主动添加声音要求。
8. 结合视频时长安排动作数量,短视频只保留一个清晰核心动作,较长视频可包含多个连续阶段。结合宽高比安排主体位置和镜头运动,但不要重复输出分辨率、时长等接口参数。
9. 只依据用户文字和实际提供的图片判断视觉内容。参考视频或参考音频如果没有可分析内容,只保留用户明确给出的绑定关系,不猜测素材内容。
10. 使用与用户输入相同的主要语言。只输出最终视频提示词,不要输出分析过程、标题、Markdown、参数表、备选方案、解释或空泛的画质宣传词。"""
def _is_within_directory(root, candidate):
try:
return os.path.commonpath([root, candidate]) == root
except ValueError:
return False
def _normalize_prompt_references(references):
if references in (None, ""):
return []
if not isinstance(references, list):
raise ValueError("参考图清单必须是数组")
if len(references) > PROMPT_OPTIMIZER_MAX_REFERENCES:
raise ValueError(f"参考图最多支持 {PROMPT_OPTIMIZER_MAX_REFERENCES} 张")
normalized = []
for item in references:
if not isinstance(item, dict):
raise ValueError("参考图清单包含无效项目")
name = str(item.get("name") or "").strip()
subfolder = str(item.get("subfolder") or "").strip()
folder_type = str(item.get("type") or "input").strip()
if not name:
raise ValueError("参考图文件名不能为空")
if folder_type != "input":
raise ValueError("参考图必须来自 ComfyUI input 目录")
normalized.append({"name": name, "subfolder": subfolder, "type": "input"})
return normalized
def _resolve_prompt_reference(item, input_directory):
input_root = os.path.realpath(os.path.abspath(input_directory))
candidate = os.path.realpath(os.path.abspath(
os.path.join(input_root, item.get("subfolder", ""), item["name"])
))
if not _is_within_directory(input_root, candidate) or not os.path.isfile(candidate):
raise ValueError(f"参考图不存在或路径不安全:{item['name']}")
return candidate
def _flatten_prompt_reference(image):
if image.mode in ("RGBA", "LA") or "transparency" in image.info:
rgba = image.convert("RGBA")
background = Image.new("RGB", rgba.size, "white")
background.paste(rgba, mask=rgba.getchannel("A"))
rgba.close()
return background
return image.convert("RGB")
def _encode_prompt_reference(path):
"""Create an aspect-preserving analysis JPEG without mutating source pixels."""
try:
with Image.open(path) as opened:
opened.seek(0)
transposed = ImageOps.exif_transpose(opened)
try:
source = _flatten_prompt_reference(transposed)
finally:
if transposed is not opened:
transposed.close()
except Exception as exc:
raise ValueError(f"无法读取参考图:{os.path.basename(path)}") from exc
try:
source_long_edge = max(source.size)
target_long_edge = min(source_long_edge, PROMPT_OPTIMIZER_MAX_LONG_EDGE)
encoded = b""
while True:
scale = min(1.0, target_long_edge / source_long_edge)
size = (
max(1, round(source.width * scale)),
max(1, round(source.height * scale)),
)
candidate = (
source.resize(size, Image.Resampling.LANCZOS)
if size != source.size
else source
)
try:
buffer = io.BytesIO()
candidate.save(buffer, format="JPEG", quality=92, optimize=True)
encoded = buffer.getvalue()
finally:
if candidate is not source:
candidate.close()
if (
len(encoded) <= PROMPT_OPTIMIZER_MAX_IMAGE_BYTES
or target_long_edge <= PROMPT_OPTIMIZER_MIN_LONG_EDGE
):
return encoded
estimated_scale = math.sqrt(
PROMPT_OPTIMIZER_MAX_IMAGE_BYTES / len(encoded)
) * 0.96
target_long_edge = max(
PROMPT_OPTIMIZER_MIN_LONG_EDGE,
min(target_long_edge - 1, round(target_long_edge * estimated_scale)),
)
finally:
source.close()
def build_prompt_optimization_payload(prompt, references, input_directory):
prompt = str(prompt or "").strip()
if not prompt:
raise ValueError("请先输入需要优化的提示词")
if len(prompt) > PROMPT_OPTIMIZER_MAX_PROMPT_CHARS:
raise ValueError(
f"提示词过长,最多支持 {PROMPT_OPTIMIZER_MAX_PROMPT_CHARS} 个字符"
)
normalized_references = _normalize_prompt_references(references)
content = [{
"type": "text",
"text": (
f"用户当前提示词:\n{prompt}\n\n"
+ (
f"下面共有 {len(normalized_references)} 张参考图,请严格按编号和上传顺序分析。"
if normalized_references
else "当前没有参考图,请按文生图需求优化。"
)
),
}]
for index, item in enumerate(normalized_references, 1):
path = _resolve_prompt_reference(item, input_directory)
encoded = _encode_prompt_reference(path)
content.extend([
{"type": "text", "text": f"参考图 {index}(上传顺序第 {index} 张)"},
{
"type": "image_url",
"image_url": {
"url": "data:image/jpeg;base64," + base64.b64encode(encoded).decode("ascii"),
"detail": "high",
},
},
])
payload = {
"model": PROMPT_OPTIMIZER_MODEL,
"reasoning_effort": PROMPT_OPTIMIZER_REASONING_EFFORT,
"stream": False,
"max_tokens": 8192,
"messages": [
{"role": "system", "content": _PROMPT_OPTIMIZER_SYSTEM_PROMPT},
{"role": "user", "content": content},
],
}
body_size = len(
json.dumps(payload, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
)
if body_size > PROMPT_OPTIMIZER_BODY_LIMIT_BYTES:
raise ValueError(
f"提示词优化请求体 {body_size / (1024 * 1024):.2f} MiB 超过 "
f"{PROMPT_OPTIMIZER_BODY_LIMIT_BYTES / (1024 * 1024):.0f} MiB 上限,"
"请减少参考图数量或在上游缩小图片"
)
return payload
def build_video_prompt_writing_payload(
prompt,
references,
input_directory,
context=None,
):
"""Build the dedicated multimodal request used by the video AI-writing action."""
prompt = str(prompt or "").strip()
if not prompt:
raise ValueError("请先输入视频创意或基础提示词")
if len(prompt) > PROMPT_OPTIMIZER_MAX_PROMPT_CHARS:
raise ValueError(
f"提示词过长,最多支持 {PROMPT_OPTIMIZER_MAX_PROMPT_CHARS} 个字符"
)
context = context if isinstance(context, dict) else {}
generation_mode = str(context.get("generation_mode") or "text").strip()
mode_labels = {
"text": "文生视频",
"first_frame": "首帧图生视频",
"first_last_frame": "首尾帧生视频",
"multimodal": "多模态参考",
}
if generation_mode not in mode_labels:
raise ValueError("不支持的视频生成模式")
duration = str(context.get("duration") or "auto").strip()[:32]
aspect_ratio = str(context.get("aspect_ratio") or "auto").strip()[:32]
generate_audio = context.get("generate_audio") is True
def _reference_count(name):
try:
return max(0, min(1000, int(context.get(name) or 0)))
except (TypeError, ValueError):
return 0
video_count = _reference_count("reference_video_count")
audio_count = _reference_count("reference_audio_count")
normalized_references = _normalize_prompt_references(references)
if generation_mode == "text":
normalized_references = []
elif generation_mode == "first_frame":
normalized_references = normalized_references[:1]
elif generation_mode == "first_last_frame":
normalized_references = normalized_references[:2]
material_note = (
f"可分析图片 {len(normalized_references)} 张;"
f"另有参考视频 {video_count} 个、参考音频 {audio_count} 个。"
)
if video_count or audio_count:
material_note += "参考视频和参考音频仅提供数量与顺序语义,不得猜测未提供的内容。"
content = [{
"type": "text",
"text": (
f"用户当前视频描述:\n{prompt}\n\n"
f"生成模式:{mode_labels[generation_mode]}\n"
f"目标时长:{duration}\n"
f"画面宽高比:{aspect_ratio}\n"
f"生成音频:{'开启' if generate_audio else '关闭'}\n"
f"参考素材:{material_note}"
),
}]
for index, item in enumerate(normalized_references, 1):
if generation_mode == "first_frame":
role = "首帧"
elif generation_mode == "first_last_frame":
role = "首帧" if index == 1 else "尾帧"
else:
role = f"参考图 {index}"
path = _resolve_prompt_reference(item, input_directory)
encoded = _encode_prompt_reference(path)
content.extend([
{"type": "text", "text": f"{role}(第 {index} 张分析图)"},
{
"type": "image_url",
"image_url": {
"url": "data:image/jpeg;base64," + base64.b64encode(encoded).decode("ascii"),
"detail": "high",
},
},
])
payload = {
"model": PROMPT_OPTIMIZER_MODEL,
"reasoning_effort": PROMPT_OPTIMIZER_REASONING_EFFORT,
"stream": False,
"max_tokens": 8192,
"messages": [
{"role": "system", "content": _VIDEO_PROMPT_WRITER_SYSTEM_PROMPT},
{"role": "user", "content": content},
],
}
body_size = len(
json.dumps(payload, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
)
if body_size > PROMPT_OPTIMIZER_BODY_LIMIT_BYTES:
raise ValueError(
f"视频 AI帮写请求体 {body_size / (1024 * 1024):.2f} MiB 超过 "
f"{PROMPT_OPTIMIZER_BODY_LIMIT_BYTES / (1024 * 1024):.0f} MiB 上限,"
"请减少参考图数量或在上游缩小图片"
)
return payload
def _extract_optimized_prompt(payload, label="提示词优化"):
try:
content = payload["choices"][0]["message"]["content"]
except (KeyError, IndexError, TypeError) as exc:
raise RuntimeError(f"{label}响应缺少文本内容") from exc
if isinstance(content, list):
content = "".join(
str(item.get("text") or "")
for item in content
if isinstance(item, dict) and item.get("type") in {None, "text", "output_text"}
)
optimized = str(content or "").strip()
fenced = re.fullmatch(r"```(?:\w+)?\s*(.*?)\s*```", optimized, re.S)
if fenced:
optimized = fenced.group(1).strip()
if not optimized:
raise RuntimeError(f"{label}结果为空")
return optimized
async def optimize_image_prompt(session, base_url, api_key, prompt, references, input_directory):
payload = build_prompt_optimization_payload(prompt, references, input_directory)
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {api_key}",
}
try:
async with session.post(
f"{base_url}/v1/chat/completions",
headers=headers,
json=payload,
timeout=PROMPT_OPTIMIZER_TIMEOUT_SECONDS,
) as response:
if response.status != 200:
raise RuntimeError(f"提示词优化请求失败 (HTTP {response.status})")
result = await response.json(content_type=None)
except RuntimeError:
raise
except Exception as exc:
raise RuntimeError("提示词优化请求失败,请检查网络后重试") from exc
return _extract_optimized_prompt(result)
async def write_video_prompt(
session,
base_url,
api_key,
prompt,
references,
input_directory,
context=None,
):
payload = build_video_prompt_writing_payload(
prompt,
references,
input_directory,
context,
)
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {api_key}",
}
try:
async with session.post(
f"{base_url}/v1/chat/completions",
headers=headers,
json=payload,
timeout=PROMPT_OPTIMIZER_TIMEOUT_SECONDS,
) as response:
if response.status != 200:
raise RuntimeError(f"视频 AI帮写请求失败 (HTTP {response.status})")
result = await response.json(content_type=None)
except RuntimeError:
raise
except Exception as exc:
raise RuntimeError("视频 AI帮写请求失败,请检查网络后重试") from exc
return _extract_optimized_prompt(result, "视频 AI帮写")
def _xml_local_name(tag):
return tag.rsplit("}", 1)[-1]
def _column_index(cell_ref):
letters = re.match(r"[A-Za-z]+", cell_ref or "")
if not letters:
return None
value = 0
for char in letters.group(0).upper():
value = value * 26 + ord(char) - ord("A") + 1
return value - 1
def _shared_string_text(node):
return "".join(
child.text or ""
for child in node.iter()
if _xml_local_name(child.tag) == "t"
)
def _sheet_entries(archive):
fallback = sorted(
name
for name in archive.namelist()
if re.fullmatch(r"xl/worksheets/sheet\d+\.xml", name)
)
try:
workbook = ET.fromstring(archive.read("xl/workbook.xml"))
relationships = ET.fromstring(archive.read("xl/_rels/workbook.xml.rels"))
except (KeyError, ET.ParseError):
return [(posixpath.basename(name).removesuffix(".xml"), name) for name in fallback]
targets = {
rel.attrib.get("Id"): rel.attrib.get("Target")
for rel in relationships
if rel.attrib.get("Id") and rel.attrib.get("Target")
}
sheets = []
for sheet in workbook.iter():
if _xml_local_name(sheet.tag) != "sheet":
continue
rel_id = next(
(value for key, value in sheet.attrib.items() if _xml_local_name(key) == "id"),
None,
)
target = targets.get(rel_id)
if not target:
continue
normalized = posixpath.normpath(posixpath.join("xl", target))
if normalized in archive.namelist():
sheets.append((sheet.attrib.get("name") or "工作表", normalized))
return sheets or [(posixpath.basename(name).removesuffix(".xml"), name) for name in fallback]
def extract_xlsx_text(data, max_chars=MAX_EXTRACTED_CHARS):
"""Convert an xlsx/xlsm OOXML workbook to readable tab-separated text."""
try:
archive = zipfile.ZipFile(io.BytesIO(data))
except (OSError, zipfile.BadZipFile) as exc:
raise ValueError("不是有效的 XLSX 文件") from exc
with archive:
shared_strings = []
try:
shared_root = ET.fromstring(archive.read("xl/sharedStrings.xml"))
shared_strings = [
_shared_string_text(node)
for node in shared_root
if _xml_local_name(node.tag) == "si"
]
except (KeyError, ET.ParseError):
pass
blocks = []
for sheet_name, sheet_path in _sheet_entries(archive):
try:
sheet_root = ET.fromstring(archive.read(sheet_path))
except (KeyError, ET.ParseError):
continue
lines = []
for row in sheet_root.iter():
if _xml_local_name(row.tag) != "row":
continue
values = []
next_column = 0
for cell in row:
if _xml_local_name(cell.tag) != "c":
continue
column = _column_index(cell.attrib.get("r"))
if column is None:
column = next_column
while len(values) < column:
values.append("")
cell_type = cell.attrib.get("t")
value = ""
if cell_type == "inlineStr":
value = _shared_string_text(cell)
else:
raw = next(
(
child.text or ""
for child in cell
if _xml_local_name(child.tag) == "v"
),
"",
)
if cell_type == "s":
try:
value = shared_strings[int(raw)]
except (ValueError, IndexError):
value = raw
elif cell_type == "b":
value = "TRUE" if raw == "1" else "FALSE"
else:
value = raw
while len(values) <= column:
values.append("")
values[column] = value
next_column = column + 1
while values and values[-1] == "":
values.pop()
if values:
lines.append("\t".join(values))
blocks.append(f"[{sheet_name}]\n" + "\n".join(lines))
if not blocks:
raise ValueError("XLSX 中没有可读取的工作表")
text = "\n\n".join(blocks).strip()
if len(text) > max_chars:
text = text[:max_chars] + "\n\n[内容过长,已截断]"
return text
def expand_xlsx_attachments(messages):
"""Replace inline xlsx/xlsm file parts with extracted text parts."""
expanded = []
for message in messages:
if not isinstance(message, dict):
expanded.append(message)
continue
content = message.get("content")
if not isinstance(content, list):
expanded.append(message)
continue
new_content = []
for part in content:
file_info = part.get("file") if isinstance(part, dict) else None
filename = file_info.get("filename", "") if isinstance(file_info, dict) else ""
file_data = file_info.get("file_data") if isinstance(file_info, dict) else None
extension = filename.rsplit(".", 1)[-1].lower() if "." in filename else ""
if not isinstance(part, dict) or part.get("type") != "file" or extension not in {"xlsx", "xlsm"} or not file_data:
new_content.append(part)
continue
try:
encoded = file_data.split(",", 1)[1] if "," in file_data else file_data
workbook = base64.b64decode(encoded, validate=True)
extracted = extract_xlsx_text(workbook)
except Exception as exc:
raise ValueError(f'无法读取 Excel 文件 "{filename}": {exc}') from exc
new_content.append({
"type": "text",
"text": f"[Excel 文件: {filename}]\n{extracted}",
})
expanded.append({**message, "content": new_content})
return expanded
def extract_search_query(messages, max_length=100):
for message in reversed(messages):
if not isinstance(message, dict) or message.get("role") != "user":
continue
content = message.get("content")
if isinstance(content, str):
return content.strip()[:max_length]
if isinstance(content, list):
text = " ".join(
part.get("text", "")
for part in content
if isinstance(part, dict) and part.get("type") == "text"
)
return text.strip()[:max_length]
return ""
return ""
def _strip_html(value):
value = re.sub(r"<[^>]+>", "", value or "")
return re.sub(r"\s+", " ", html.unescape(value)).strip()
async def web_search(session, query, count=6):
url = f"https://www.bing.com/search?q={quote_plus(query)}&mkt=zh-CN"
headers = {"User-Agent": _USER_AGENT, "Accept-Language": "zh-CN,zh;q=0.9"}
async with session.get(
url,
headers=headers,
allow_redirects=True,
timeout=SEARCH_TIMEOUT_SECONDS,
) as response:
if response.status != 200:
raise RuntimeError(f"搜索请求失败 (HTTP {response.status})")
page = await response.text(errors="ignore")
results = []
for block in re.split(r'<li class="b_algo[^\"]*"', page)[1:]:
link = re.search(r'<h2[^>]*><a[^>]*href="([^"]+)"[^>]*>(.*?)</a>', block, re.S)
if not link:
continue
target = _strip_html(link.group(1))
if not target.startswith(("http://", "https://")):
continue
snippet_match = re.search(r"<p[^>]*>(.*?)</p>", block, re.S)
results.append({
"title": _strip_html(link.group(2))[:120],
"url": target,
"snippet": _strip_html(snippet_match.group(1))[:320] if snippet_match else "",
})
if len(results) >= count:
break
if not results:
raise RuntimeError("搜索结果解析失败")
return results
async def rewrite_search_query(session, base_url, api_key, question):
payload = {
"model": "gpt-5.6-sol",
"reasoning_effort": "low",
"stream": False,
"max_tokens": 1024,
"messages": [
{
"role": "system",
"content": (
"你是搜索查询生成器。把用户的问题改写成一条适合搜索引擎的简洁查询词:"
"保留关键实体和意图,去掉口语、疑问词和时间副词。只输出查询词本身,不要引号,不要解释。"
),
},
{"role": "user", "content": question},
],
}
headers = {"Content-Type": "application/json", "Authorization": f"Bearer {api_key}"}
try:
async with session.post(
f"{base_url}/v1/chat/completions",
headers=headers,
json=payload,
timeout=REWRITE_TIMEOUT_SECONDS,
) as response:
if response.status != 200:
return ""
result = await response.json(content_type=None)
except Exception:
return ""
try:
query = result["choices"][0]["message"]["content"].strip().strip('"\'「」『』')
except (KeyError, IndexError, TypeError, AttributeError):
return ""
return query if query and len(query) <= 80 and "\n" not in query else ""
def build_search_context(query, results):
items = []
for index, result in enumerate(results, 1):
item = f"[{index}] {result['title']}\n来源: {result['url']}"
if result.get("snippet"):
item += f"\n摘要: {result['snippet']}"
items.append(item)
return (
f"以下是针对用户最新问题的联网搜索结果(查询词:{query}):\n\n"
+ "\n\n".join(items)
+ "\n\n请优先基于以上搜索结果回答用户的最新问题,引用某条结果时标注其编号(如 [1])。"
"若搜索结果与问题无关或不足以回答,请说明这一点,再依据自身知识谨慎补充。"
)