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.
This commit is contained in:
Jony
+31
-35
@@ -1,37 +1,33 @@
|
||||
"""
|
||||
工具模块
|
||||
包含图像处理、配置管理、文件处理等通用工具函数
|
||||
"""
|
||||
"""Shared helpers exposed through lazy imports."""
|
||||
|
||||
from .image_utils import (
|
||||
tensor_to_pil,
|
||||
pil_to_tensor,
|
||||
encode_image_to_base64,
|
||||
decode_base64_to_pil
|
||||
)
|
||||
from .config import load_config, get_api_key
|
||||
from .file_utils import (
|
||||
ImageInfo,
|
||||
load_images_from_folder,
|
||||
pair_images_indexed,
|
||||
pair_images_cartesian,
|
||||
generate_timestamp_filename,
|
||||
save_image,
|
||||
get_folder_image_count
|
||||
)
|
||||
from importlib import import_module
|
||||
|
||||
__all__ = [
|
||||
'tensor_to_pil',
|
||||
'pil_to_tensor',
|
||||
'encode_image_to_base64',
|
||||
'decode_base64_to_pil',
|
||||
'load_config',
|
||||
'get_api_key',
|
||||
'ImageInfo',
|
||||
'load_images_from_folder',
|
||||
'pair_images_indexed',
|
||||
'pair_images_cartesian',
|
||||
'generate_timestamp_filename',
|
||||
'save_image',
|
||||
'get_folder_image_count'
|
||||
]
|
||||
|
||||
_EXPORTS = {
|
||||
"tensor_to_pil": ("image_utils", "tensor_to_pil"),
|
||||
"pil_to_tensor": ("image_utils", "pil_to_tensor"),
|
||||
"encode_image_to_base64": ("image_utils", "encode_image_to_base64"),
|
||||
"decode_base64_to_pil": ("image_utils", "decode_base64_to_pil"),
|
||||
"load_config": ("config", "load_config"),
|
||||
"get_api_key": ("config", "get_api_key"),
|
||||
"ImageInfo": ("file_utils", "ImageInfo"),
|
||||
"load_images_from_folder": ("file_utils", "load_images_from_folder"),
|
||||
"pair_images_indexed": ("file_utils", "pair_images_indexed"),
|
||||
"pair_images_cartesian": ("file_utils", "pair_images_cartesian"),
|
||||
"generate_timestamp_filename": ("file_utils", "generate_timestamp_filename"),
|
||||
"save_image": ("file_utils", "save_image"),
|
||||
"get_folder_image_count": ("file_utils", "get_folder_image_count"),
|
||||
}
|
||||
|
||||
__all__ = list(_EXPORTS)
|
||||
|
||||
|
||||
def __getattr__(name):
|
||||
try:
|
||||
module_name, attribute_name = _EXPORTS[name]
|
||||
except KeyError as exc:
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from exc
|
||||
|
||||
value = getattr(import_module(f".{module_name}", __name__), attribute_name)
|
||||
globals()[name] = value
|
||||
return value
|
||||
|
||||
@@ -0,0 +1,691 @@
|
||||
"""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]）。"
|
||||
"若搜索结果与问题无关或不足以回答，请说明这一点，再依据自身知识谨慎补充。"
|
||||
)
|
||||
@@ -1,174 +0,0 @@
|
||||
"""
|
||||
颜色去背景工具模块
|
||||
基于颜色距离计算实现精确可控的背景移除，不依赖 AI 模型。
|
||||
|
||||
支持模式：
|
||||
- white: 白色背景去除
|
||||
- white-preserve: 白色背景但保护浅色前景物体
|
||||
- corner: 自动采样四角颜色作为背景色
|
||||
- color: 指定任意颜色去除
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
|
||||
def background_to_alpha(
|
||||
image: Image.Image,
|
||||
bg_color: tuple = (255, 255, 255),
|
||||
tolerance: float = 8.0,
|
||||
feather: float = 45.0,
|
||||
strength: float = 1.0,
|
||||
min_alpha: int = 2,
|
||||
) -> Image.Image:
|
||||
"""
|
||||
将纯色背景转为透明。
|
||||
|
||||
对白色背景使用 white-to-alpha 恢复算法，保持彩色文字和抗锯齿边缘清晰。
|
||||
对其他颜色使用欧氏距离计算。
|
||||
"""
|
||||
rgba = np.asarray(image.convert("RGBA")).astype(np.float32)
|
||||
rgb = rgba[:, :, :3] / 255.0
|
||||
existing_alpha = rgba[:, :, 3] / 255.0
|
||||
bg = np.array(bg_color, dtype=np.float32) / 255.0
|
||||
|
||||
if max(bg_color) >= 245 and min(bg_color) >= 245:
|
||||
alpha = (1.0 - np.min(rgb, axis=2)) * float(strength)
|
||||
if tolerance > 0:
|
||||
dist = np.linalg.norm((1.0 - rgb) * 255.0, axis=2)
|
||||
gate = np.clip(
|
||||
(dist - float(tolerance)) / max(1.0, float(feather) * 0.25),
|
||||
0.0, 1.0,
|
||||
)
|
||||
alpha *= gate
|
||||
else:
|
||||
dist = np.linalg.norm((rgb - bg) * 255.0, axis=2)
|
||||
denom = max(1.0, float(feather))
|
||||
alpha = np.clip((dist - float(tolerance)) / denom, 0.0, 1.0)
|
||||
alpha *= float(strength)
|
||||
|
||||
alpha = np.clip(alpha, 0.0, 1.0) * existing_alpha
|
||||
alpha[alpha < (float(min_alpha) / 255.0)] = 0.0
|
||||
|
||||
# 从 alpha 混合中恢复前景色，避免白边
|
||||
out_rgb = rgb.copy()
|
||||
mask = alpha > 1e-6
|
||||
out_rgb[mask] = (rgb[mask] - bg * (1.0 - alpha[mask, None])) / alpha[mask, None]
|
||||
out_rgb = np.clip(out_rgb, 0.0, 1.0)
|
||||
|
||||
out = np.dstack([
|
||||
(out_rgb * 255.0).astype(np.uint8),
|
||||
(alpha * 255.0).astype(np.uint8),
|
||||
])
|
||||
return Image.fromarray(out, "RGBA")
|
||||
|
||||
|
||||
def corner_color(image: Image.Image, sample: int = 12) -> tuple:
|
||||
"""采样图片四角像素的中位数颜色，用于自动检测背景色。"""
|
||||
rgb = np.asarray(image.convert("RGB"))
|
||||
h, w = rgb.shape[:2]
|
||||
sample = max(1, min(sample, h, w))
|
||||
patches = [
|
||||
rgb[:sample, :sample],
|
||||
rgb[:sample, w - sample:],
|
||||
rgb[h - sample:, :sample],
|
||||
rgb[h - sample:, w - sample:],
|
||||
]
|
||||
merged = np.concatenate([p.reshape(-1, 3) for p in patches], axis=0)
|
||||
return tuple(np.median(merged, axis=0).astype(int))
|
||||
|
||||
|
||||
# PLACEHOLDER_PRESERVE
|
||||
|
||||
def preserve_light_foreground_to_alpha(
|
||||
image: Image.Image,
|
||||
tolerance: float = 10.0,
|
||||
preserve_opacity: float = 0.72,
|
||||
min_area_ratio: float = 0.00025,
|
||||
) -> Image.Image:
|
||||
"""
|
||||
白底去除 + 浅色前景保护。
|
||||
|
||||
适用于前景包含白色/浅色物体（白盘子、白帆、白色包装）的场景。
|
||||
使用 OpenCV 连通区域分析保护大面积浅色前景结构。
|
||||
如果 OpenCV 不可用，回退到普通 white-to-alpha。
|
||||
"""
|
||||
base = background_to_alpha(image, (255, 255, 255), tolerance=tolerance)
|
||||
try:
|
||||
import cv2
|
||||
except ImportError:
|
||||
return base
|
||||
|
||||
rgb_u8 = np.asarray(image.convert("RGB"))
|
||||
h, w = rgb_u8.shape[:2]
|
||||
dist = np.sqrt(np.sum((255.0 - rgb_u8.astype(np.float32)) ** 2, axis=2))
|
||||
rough = (dist > float(tolerance)).astype(np.uint8) * 255
|
||||
|
||||
kernel_open = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
|
||||
kernel_close = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (17, 17))
|
||||
rough = cv2.morphologyEx(rough, cv2.MORPH_OPEN, kernel_open, iterations=1)
|
||||
rough = cv2.morphologyEx(rough, cv2.MORPH_CLOSE, kernel_close, iterations=2)
|
||||
|
||||
count, labels, stats, _ = cv2.connectedComponentsWithStats(rough, 8)
|
||||
keep = np.zeros_like(rough)
|
||||
min_area = max(24, int(w * h * float(min_area_ratio)))
|
||||
for idx in range(1, count):
|
||||
if stats[idx, cv2.CC_STAT_AREA] >= min_area:
|
||||
keep[labels == idx] = 255
|
||||
|
||||
# PLACEHOLDER_FLOOD
|
||||
|
||||
flood = keep.copy()
|
||||
ff_mask = np.zeros((h + 2, w + 2), dtype=np.uint8)
|
||||
cv2.floodFill(flood, ff_mask, (0, 0), 255)
|
||||
filled = cv2.bitwise_or(keep, cv2.bitwise_not(flood))
|
||||
soft = cv2.GaussianBlur(filled, (0, 0), 5).astype(np.float32) / 255.0
|
||||
|
||||
near_kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (29, 29))
|
||||
near = cv2.dilate(
|
||||
(dist > (float(tolerance) * 0.65)).astype(np.uint8) * 255,
|
||||
near_kernel, iterations=1,
|
||||
)
|
||||
near = cv2.GaussianBlur(near, (0, 0), 8).astype(np.float32) / 255.0
|
||||
lift = np.minimum(soft, near) * float(preserve_opacity)
|
||||
|
||||
arr = np.asarray(base.convert("RGBA")).copy()
|
||||
alpha = arr[:, :, 3].astype(np.float32) / 255.0
|
||||
alpha = np.maximum(alpha, lift)
|
||||
alpha[alpha < (2.0 / 255.0)] = 0.0
|
||||
|
||||
original = np.asarray(image.convert("RGB"))
|
||||
very_light = (np.mean(original, axis=2) > 224) & (lift > 0.12)
|
||||
arr[:, :, :3][very_light] = original[very_light]
|
||||
arr[:, :, 3] = np.clip(alpha * 255.0, 0, 255).astype(np.uint8)
|
||||
return Image.fromarray(arr, "RGBA")
|
||||
|
||||
|
||||
def remove_background(
|
||||
image: Image.Image,
|
||||
mode: str = "white",
|
||||
bg_color: tuple = (255, 255, 255),
|
||||
tolerance: float = 8.0,
|
||||
feather: float = 45.0,
|
||||
strength: float = 1.0,
|
||||
) -> Image.Image:
|
||||
"""
|
||||
统一入口：根据模式移除背景。
|
||||
|
||||
mode:
|
||||
- white: 白色背景去除
|
||||
- white-preserve: 白底 + 保护浅色前景
|
||||
- corner: 自动采样四角颜色
|
||||
- color: 使用指定 bg_color
|
||||
"""
|
||||
if mode == "white":
|
||||
return background_to_alpha(image, (255, 255, 255), tolerance, feather, strength)
|
||||
elif mode == "white-preserve":
|
||||
return preserve_light_foreground_to_alpha(image, tolerance)
|
||||
elif mode == "corner":
|
||||
bg = corner_color(image)
|
||||
return background_to_alpha(image, bg, tolerance, feather, strength)
|
||||
elif mode == "color":
|
||||
return background_to_alpha(image, bg_color, tolerance, feather, strength)
|
||||
else:
|
||||
return image.convert("RGBA")
|
||||
+63
-4
@@ -4,22 +4,24 @@
|
||||
"""
|
||||
|
||||
import os
|
||||
import threading
|
||||
from typing import Dict, Optional
|
||||
|
||||
|
||||
# 获取插件根目录
|
||||
PLUGIN_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
CONFIG_FILE = os.path.join(PLUGIN_ROOT, ".config")
|
||||
_CONFIG_LOCK = threading.RLock()
|
||||
|
||||
|
||||
# ============ API 基础配置 ============
|
||||
# 所有 API 客户端的统一基础 URL
|
||||
# 可通过环境变量 O1KEY_API_BASE_URL 覆盖
|
||||
DEFAULT_API_BASE_URL = "https://api.o1key.com"
|
||||
DEFAULT_API_BASE_URL = "https://api.o1key.cn"
|
||||
|
||||
# 异步 API 基础 URL（用于异步提交+轮询模式）
|
||||
# 可通过环境变量 O1KEY_ASYNC_API_BASE_URL 覆盖
|
||||
DEFAULT_ASYNC_API_BASE_URL = "https://cf-api.o1key.com"
|
||||
DEFAULT_ASYNC_API_BASE_URL = "https://api.o1key.cn"
|
||||
|
||||
# ============ 网络线路配置 ============
|
||||
NETWORK_ROUTES = {
|
||||
@@ -28,6 +30,8 @@ NETWORK_ROUTES = {
|
||||
"美国直连": "https://api.o1key.com",
|
||||
}
|
||||
NETWORK_ROUTE_OPTIONS = ["全球加速", "CF加速", "美国直连"]
|
||||
DEFAULT_NETWORK_ROUTE = "全球加速"
|
||||
NETWORK_ROUTE_CONFIG_KEY = "O1KEY_NETWORK_ROUTE"
|
||||
|
||||
|
||||
def load_config(config_path: Optional[str] = None) -> Dict[str, str]:
|
||||
@@ -76,6 +80,45 @@ def load_config(config_path: Optional[str] = None) -> Dict[str, str]:
|
||||
return config
|
||||
|
||||
|
||||
def save_config(config: Dict[str, str], config_path: Optional[str] = None) -> None:
|
||||
"""原子写入配置，避免界面快速切换时产生半写入文件。"""
|
||||
if config_path is None:
|
||||
config_path = CONFIG_FILE
|
||||
|
||||
temp_path = f"{config_path}.tmp"
|
||||
with _CONFIG_LOCK:
|
||||
with open(temp_path, "w", encoding="utf-8", newline="\n") as config_file:
|
||||
for key, value in config.items():
|
||||
config_file.write(f"{key}={value}\n")
|
||||
os.replace(temp_path, config_path)
|
||||
|
||||
|
||||
def update_config(
|
||||
updates: Optional[Dict[str, str]] = None,
|
||||
remove: Optional[list[str]] = None,
|
||||
) -> Dict[str, str]:
|
||||
"""以一次原子写入更新配置，并返回更新后的配置。"""
|
||||
with _CONFIG_LOCK:
|
||||
config = load_config()
|
||||
for key, value in (updates or {}).items():
|
||||
config[key] = value
|
||||
for key in remove or []:
|
||||
config.pop(key, None)
|
||||
save_config(config)
|
||||
return config
|
||||
|
||||
|
||||
def get_runtime_config_signature() -> tuple[str, str, str, str]:
|
||||
"""用于复用客户端；仅当运行时 API 配置变化时才重建客户端。"""
|
||||
config = load_config()
|
||||
return (
|
||||
config.get("O1KEY_API_KEY", ""),
|
||||
config.get(NETWORK_ROUTE_CONFIG_KEY, DEFAULT_NETWORK_ROUTE),
|
||||
config.get("O1KEY_API_BASE_URL", ""),
|
||||
config.get("O1KEY_ASYNC_API_BASE_URL", ""),
|
||||
)
|
||||
|
||||
|
||||
def get_api_key(key_name: str = "O1KEY_API_KEY") -> Optional[str]:
|
||||
"""
|
||||
获取 API 密钥
|
||||
@@ -112,6 +155,14 @@ def get_api_key_or_raise(key_name: str = "O1KEY_API_KEY") -> str:
|
||||
return api_key
|
||||
|
||||
|
||||
def get_network_route() -> str:
|
||||
"""Return the globally configured network route."""
|
||||
route = load_config().get(NETWORK_ROUTE_CONFIG_KEY, DEFAULT_NETWORK_ROUTE)
|
||||
if isinstance(route, (list, tuple)):
|
||||
route = route[0] if route else DEFAULT_NETWORK_ROUTE
|
||||
return route if route in NETWORK_ROUTES else DEFAULT_NETWORK_ROUTE
|
||||
|
||||
|
||||
def get_api_base_url() -> str:
|
||||
"""
|
||||
获取 API 基础 URL
|
||||
@@ -121,6 +172,9 @@ def get_api_base_url() -> str:
|
||||
API 基础 URL 字符串
|
||||
"""
|
||||
config = load_config()
|
||||
route = config.get(NETWORK_ROUTE_CONFIG_KEY)
|
||||
if route in NETWORK_ROUTES:
|
||||
return NETWORK_ROUTES[route].rstrip('/')
|
||||
base_url = config.get("O1KEY_API_BASE_URL")
|
||||
|
||||
if base_url:
|
||||
@@ -138,6 +192,9 @@ def get_async_api_base_url() -> str:
|
||||
异步 API 基础 URL 字符串
|
||||
"""
|
||||
config = load_config()
|
||||
route = config.get(NETWORK_ROUTE_CONFIG_KEY)
|
||||
if route in NETWORK_ROUTES:
|
||||
return NETWORK_ROUTES[route].rstrip('/')
|
||||
base_url = config.get("O1KEY_ASYNC_API_BASE_URL")
|
||||
|
||||
if base_url:
|
||||
@@ -146,8 +203,10 @@ def get_async_api_base_url() -> str:
|
||||
return DEFAULT_ASYNC_API_BASE_URL
|
||||
|
||||
|
||||
def get_base_url_by_route(route: str) -> str:
|
||||
"""根据网络线路选项返回对应域名，未匹配则走 config 垫底"""
|
||||
def get_base_url_by_route(route: Optional[str] = None) -> str:
|
||||
"""Resolve an explicit route or fall back to the global route setting."""
|
||||
if isinstance(route, (list, tuple)):
|
||||
route = route[0] if route else None
|
||||
if route is None:
|
||||
route = get_network_route()
|
||||
return NETWORK_ROUTES.get(route, get_api_base_url())
|
||||
|
||||
@@ -0,0 +1,371 @@
|
||||
"""HTTP/2-first client with an aiohttp fallback for dependency-light installs."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import importlib.util
|
||||
import logging
|
||||
import os
|
||||
from typing import Any, AsyncIterator, Optional
|
||||
|
||||
import aiohttp
|
||||
|
||||
try:
|
||||
import httpx
|
||||
except ImportError: # Manual plugin copies may not install new requirements.
|
||||
httpx = None
|
||||
|
||||
|
||||
def _verbose_http_logging_enabled() -> bool:
|
||||
value = os.environ.get("O1KEY_VERBOSE_LOG", "")
|
||||
return value.strip().lower() not in ("", "0", "false", "no", "off")
|
||||
|
||||
|
||||
# httpx emits one INFO line for every polling request. Those access logs drown
|
||||
# out the batch lifecycle and add no information on successful requests. Keep
|
||||
# warnings/errors, and allow the existing verbose switch to restore raw logs.
|
||||
if not _verbose_http_logging_enabled():
|
||||
logging.getLogger("httpx").setLevel(logging.WARNING)
|
||||
logging.getLogger("httpcore").setLevel(logging.WARNING)
|
||||
|
||||
|
||||
HTTPX_AVAILABLE = httpx is not None
|
||||
HTTP_CLIENT_ERRORS = (aiohttp.ClientError,)
|
||||
HTTP_STREAM_ERRORS = (aiohttp.ClientPayloadError,)
|
||||
if HTTPX_AVAILABLE:
|
||||
HTTP_CLIENT_ERRORS += (httpx.HTTPError,)
|
||||
HTTP_STREAM_ERRORS += (httpx.StreamError,)
|
||||
|
||||
|
||||
class ResponseBodyIntegrityError(OSError):
|
||||
"""The response body ended early or disagreed with Content-Length."""
|
||||
|
||||
|
||||
class ResponseTaskIdMismatchError(OSError):
|
||||
"""A task query returned a different task ID than the one requested."""
|
||||
|
||||
|
||||
def _response_header(response: Any, name: str) -> str:
|
||||
headers = getattr(response, "headers", None) or {}
|
||||
value = headers.get(name)
|
||||
if value is None:
|
||||
value = headers.get(name.lower())
|
||||
if value is None:
|
||||
for key, candidate in headers.items():
|
||||
if str(key).lower() == name.lower():
|
||||
value = candidate
|
||||
break
|
||||
return "" if value is None else str(value).strip()
|
||||
|
||||
|
||||
def response_body_diagnostics(response: Any, received_bytes: int) -> dict[str, Any]:
|
||||
"""Describe whether a fully-read response agrees with its declared length."""
|
||||
declared_text = _response_header(response, "Content-Length")
|
||||
content_encoding = _response_header(response, "Content-Encoding") or "identity"
|
||||
transfer_encoding = _response_header(response, "Transfer-Encoding") or "<none>"
|
||||
declared_bytes = None
|
||||
length_check = "not-declared"
|
||||
|
||||
if declared_text:
|
||||
try:
|
||||
declared_bytes = int(declared_text)
|
||||
except ValueError:
|
||||
length_check = "invalid-header"
|
||||
else:
|
||||
if content_encoding.lower() not in ("", "identity"):
|
||||
# aiohttp/httpx expose decoded bytes while Content-Length can
|
||||
# describe the compressed wire representation.
|
||||
length_check = "skipped-compressed"
|
||||
elif declared_bytes == received_bytes:
|
||||
length_check = "match"
|
||||
else:
|
||||
length_check = "mismatch"
|
||||
|
||||
return {
|
||||
"http_version": str(getattr(response, "http_version", "") or "HTTP"),
|
||||
"status": int(getattr(response, "status", 0) or 0),
|
||||
"content_length": declared_text or "<none>",
|
||||
"declared_bytes": declared_bytes,
|
||||
"received_bytes": int(received_bytes),
|
||||
"content_encoding": content_encoding,
|
||||
"transfer_encoding": transfer_encoding,
|
||||
"length_check": length_check,
|
||||
}
|
||||
|
||||
|
||||
def format_response_body_diagnostics(diagnostics: dict[str, Any]) -> str:
|
||||
return (
|
||||
f"{diagnostics['http_version']} {diagnostics['status']}"
|
||||
f" | Content-Length={diagnostics['content_length']}"
|
||||
f" | received={diagnostics['received_bytes']}B"
|
||||
f" | Content-Encoding={diagnostics['content_encoding']}"
|
||||
f" | Transfer-Encoding={diagnostics['transfer_encoding']}"
|
||||
f" | length_check={diagnostics['length_check']}"
|
||||
)
|
||||
|
||||
|
||||
async def read_response_body_with_diagnostics(
|
||||
response: Any,
|
||||
*,
|
||||
chunk_size: int = 64 * 1024,
|
||||
) -> tuple[bytes, dict[str, Any]]:
|
||||
"""Read a body while retaining the received byte count if the stream fails."""
|
||||
buffer = bytearray()
|
||||
try:
|
||||
content = getattr(response, "content", None)
|
||||
iter_chunked = getattr(content, "iter_chunked", None)
|
||||
if callable(iter_chunked):
|
||||
async for chunk in iter_chunked(chunk_size):
|
||||
buffer.extend(chunk)
|
||||
else:
|
||||
read = getattr(response, "read", None)
|
||||
if callable(read):
|
||||
buffer.extend(await read())
|
||||
else:
|
||||
text = getattr(response, "text", None)
|
||||
if not callable(text):
|
||||
raise TypeError("response does not expose a readable body")
|
||||
buffer.extend((await text()).encode("utf-8"))
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
diagnostics = response_body_diagnostics(response, len(buffer))
|
||||
raise ResponseBodyIntegrityError(
|
||||
"响应体读取提前中断 | "
|
||||
f"{format_response_body_diagnostics(diagnostics)}"
|
||||
f" | cause={type(exc).__name__}: {exc}"
|
||||
) from exc
|
||||
|
||||
body = bytes(buffer)
|
||||
diagnostics = response_body_diagnostics(response, len(body))
|
||||
if diagnostics["length_check"] == "mismatch":
|
||||
raise ResponseBodyIntegrityError(
|
||||
"响应体长度与服务端声明不一致，可能在传输中提前中断 | "
|
||||
f"{format_response_body_diagnostics(diagnostics)}"
|
||||
)
|
||||
return body, diagnostics
|
||||
|
||||
|
||||
def response_task_id(payload: Any) -> Optional[str]:
|
||||
"""Return an explicit task_id/taskId without mistaking result item IDs for it."""
|
||||
if not isinstance(payload, dict):
|
||||
return None
|
||||
sources = [payload]
|
||||
for key in ("data", "result", "task"):
|
||||
nested = payload.get(key)
|
||||
if isinstance(nested, dict):
|
||||
sources.append(nested)
|
||||
for source in sources:
|
||||
for key in ("task_id", "taskId"):
|
||||
value = source.get(key)
|
||||
if value is not None and str(value).strip():
|
||||
return str(value).strip()
|
||||
return None
|
||||
|
||||
|
||||
def validate_response_task_id(payload: Any, expected_task_id: str) -> Optional[str]:
|
||||
actual_task_id = response_task_id(payload)
|
||||
if actual_task_id is not None and actual_task_id != str(expected_task_id):
|
||||
raise ResponseTaskIdMismatchError(
|
||||
f"任务查询响应 ID 不匹配 | requested_task_id={expected_task_id}"
|
||||
f" | response_task_id={actual_task_id}"
|
||||
)
|
||||
return actual_task_id
|
||||
|
||||
|
||||
def http2_runtime_available() -> bool:
|
||||
return HTTPX_AVAILABLE and importlib.util.find_spec("h2") is not None
|
||||
|
||||
|
||||
def create_timeout(
|
||||
total: float,
|
||||
*,
|
||||
connect: float,
|
||||
read: float,
|
||||
write: float,
|
||||
pool: float,
|
||||
):
|
||||
"""Return a timeout object understood by the selected HTTP backend."""
|
||||
if HTTPX_AVAILABLE:
|
||||
return httpx.Timeout(total, connect=connect, read=read, write=write, pool=pool)
|
||||
return aiohttp.ClientTimeout(
|
||||
total=total,
|
||||
connect=connect,
|
||||
sock_connect=connect,
|
||||
sock_read=read,
|
||||
)
|
||||
|
||||
|
||||
class _ResponseContent:
|
||||
def __init__(self, response: httpx.Response):
|
||||
self._response = response
|
||||
|
||||
async def iter_chunked(self, chunk_size: int) -> AsyncIterator[bytes]:
|
||||
async for chunk in self._response.aiter_bytes(chunk_size):
|
||||
yield chunk
|
||||
|
||||
|
||||
class HttpResponse:
|
||||
def __init__(self, response: httpx.Response):
|
||||
self._response = response
|
||||
self.content = _ResponseContent(response)
|
||||
|
||||
@property
|
||||
def status(self) -> int:
|
||||
return self._response.status_code
|
||||
|
||||
@property
|
||||
def headers(self):
|
||||
return self._response.headers
|
||||
|
||||
@property
|
||||
def raw_headers(self):
|
||||
return tuple(self._response.headers.raw)
|
||||
|
||||
@property
|
||||
def http_version(self) -> str:
|
||||
return self._response.http_version
|
||||
|
||||
async def text(self) -> str:
|
||||
await self._response.aread()
|
||||
return self._response.text
|
||||
|
||||
|
||||
class _RequestContext:
|
||||
def __init__(self, context):
|
||||
self._context = context
|
||||
|
||||
async def __aenter__(self) -> HttpResponse:
|
||||
response = await self._context.__aenter__()
|
||||
return HttpResponse(response)
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return await self._context.__aexit__(exc_type, exc, tb)
|
||||
|
||||
|
||||
class O1keyAsyncHttpClient:
|
||||
"""HTTP/2-first client; falls back to httpx HTTP/1.1, then aiohttp."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
max_connections: int = 32,
|
||||
max_keepalive_connections: int = 16,
|
||||
keepalive_expiry: float = 30.0,
|
||||
http2: bool = True,
|
||||
):
|
||||
self.http2_enabled = bool(http2 and http2_runtime_available())
|
||||
self.backend = "httpx" if HTTPX_AVAILABLE else "aiohttp"
|
||||
self._max_connections = max_connections
|
||||
self._keepalive_expiry = keepalive_expiry
|
||||
self._client = None
|
||||
if HTTPX_AVAILABLE:
|
||||
self._client = httpx.AsyncClient(
|
||||
http2=self.http2_enabled,
|
||||
verify=True,
|
||||
limits=httpx.Limits(
|
||||
max_connections=max_connections,
|
||||
max_keepalive_connections=max_keepalive_connections,
|
||||
keepalive_expiry=keepalive_expiry,
|
||||
),
|
||||
follow_redirects=False,
|
||||
)
|
||||
|
||||
async def __aenter__(self):
|
||||
if self.backend == "aiohttp":
|
||||
connector = aiohttp.TCPConnector(
|
||||
limit=self._max_connections,
|
||||
limit_per_host=self._max_connections,
|
||||
keepalive_timeout=self._keepalive_expiry,
|
||||
)
|
||||
self._client = aiohttp.ClientSession(connector=connector, trust_env=True)
|
||||
await self._client.__aenter__()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return await self._client.__aexit__(exc_type, exc, tb)
|
||||
|
||||
def post(
|
||||
self,
|
||||
url: str,
|
||||
*,
|
||||
headers: Optional[dict[str, str]] = None,
|
||||
data: Any = None,
|
||||
files: Any = None,
|
||||
timeout: Any = None,
|
||||
) -> _RequestContext:
|
||||
if self._client is None:
|
||||
raise RuntimeError("HTTP client must be entered before use")
|
||||
if self.backend == "httpx":
|
||||
return _RequestContext(self._client.stream(
|
||||
"POST",
|
||||
url,
|
||||
headers=headers,
|
||||
content=data if files is None else None,
|
||||
files=files,
|
||||
timeout=timeout,
|
||||
))
|
||||
|
||||
payload = data
|
||||
if files:
|
||||
payload = aiohttp.FormData()
|
||||
for field_name, (filename, source, content_type) in files.items():
|
||||
payload.add_field(
|
||||
field_name,
|
||||
source,
|
||||
filename=filename,
|
||||
content_type=content_type,
|
||||
)
|
||||
return self._client.post(
|
||||
url,
|
||||
headers=headers,
|
||||
data=payload,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
def get(
|
||||
self,
|
||||
url: str,
|
||||
*,
|
||||
headers: Optional[dict[str, str]] = None,
|
||||
allow_redirects: bool = False,
|
||||
timeout: Any = None,
|
||||
) -> _RequestContext:
|
||||
if self._client is None:
|
||||
raise RuntimeError("HTTP client must be entered before use")
|
||||
if self.backend == "httpx":
|
||||
return _RequestContext(self._client.stream(
|
||||
"GET",
|
||||
url,
|
||||
headers=headers,
|
||||
follow_redirects=allow_redirects,
|
||||
timeout=timeout,
|
||||
))
|
||||
return self._client.get(
|
||||
url,
|
||||
headers=headers,
|
||||
allow_redirects=allow_redirects,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
|
||||
def create_http_client(**kwargs) -> O1keyAsyncHttpClient:
|
||||
return O1keyAsyncHttpClient(**kwargs)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"HTTPX_AVAILABLE",
|
||||
"HTTP_CLIENT_ERRORS",
|
||||
"HTTP_STREAM_ERRORS",
|
||||
"O1keyAsyncHttpClient",
|
||||
"ResponseBodyIntegrityError",
|
||||
"ResponseTaskIdMismatchError",
|
||||
"create_http_client",
|
||||
"create_timeout",
|
||||
"format_response_body_diagnostics",
|
||||
"http2_runtime_available",
|
||||
"read_response_body_with_diagnostics",
|
||||
"response_body_diagnostics",
|
||||
"response_task_id",
|
||||
"validate_response_task_id",
|
||||
]
|
||||
+201
-3
@@ -11,6 +11,7 @@
|
||||
import asyncio
|
||||
import json
|
||||
import random
|
||||
import re
|
||||
from typing import Any, Optional
|
||||
|
||||
import aiohttp
|
||||
@@ -21,23 +22,125 @@ import aiohttp
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
HTTP_ERROR_MESSAGES = {
|
||||
402: "账户余额或模型额度不足，请充值或检查令牌额度。",
|
||||
429: "模型速率超限或额度不足！",
|
||||
422: "输入内容未通过安全检查，请调整提示词或参考素材。",
|
||||
502: "网关超时。请重试或将网络切换为美国直连",
|
||||
503: "模型超载。请稍后重试！",
|
||||
504: "网关超时。请稍后重试。",
|
||||
529: "MiniMax 上游模型过载，请稍后重试。",
|
||||
}
|
||||
|
||||
# 【o1key 图片生成】节点的模型无关错误封装。GPT Image 与所有
|
||||
# Nano Banana 分支都必须使用同一份映射，不能按模型拆分。
|
||||
O1KEY_IMAGE_ERROR_CONTENT_MESSAGES = {
|
||||
"content rejected: the image was flagged as unsafe by the content safety system": "内容被拒绝：该图像被内容安全系统标记为不安全。",
|
||||
"Your request was rejected by the safety system": "您的请求已被安全系统拒绝",
|
||||
"insufficient balance": "上游额度不足！",
|
||||
"Image generation returned empty response": "图片生成过程中被内容审查机制拒绝！",
|
||||
"The provided prompt is considered unsafe and it cannot be used to generate content": "提供的提示被认为是不安全的，不能用于生成内容。",
|
||||
}
|
||||
|
||||
# 错误内容关键词 → 用户友好文案（优先于状态码匹配）
|
||||
ERROR_CONTENT_MESSAGES = {
|
||||
"Your request was rejected by the safety system": "请求被安全系统拦截：请调整提示词，避免敏感、违规、血腥、色情、仇恨、未成年人或真实人物等高风险内容。",
|
||||
**O1KEY_IMAGE_ERROR_CONTENT_MESSAGES,
|
||||
"safety system": "请求被安全系统拦截：请调整提示词，避免敏感、违规、血腥、色情、仇恨、未成年人或真实人物等高风险内容。",
|
||||
"unexpected end of JSON input": "通常重试能解决；反复出现就降低分辨率、数量或换网络线路。",
|
||||
"unexpected end of JSON input": "通常重试能解决；反复出现就降低分辨率、数量，或在 API密钥设置中切换全局线路。",
|
||||
"The current model has a high load": "模型过载，请稍后重试！",
|
||||
"system error": "系统错误，请稍后重试。",
|
||||
}
|
||||
|
||||
|
||||
def format_o1key_image_error(value: Any) -> str:
|
||||
"""Format provider errors for every model in the o1key image generator."""
|
||||
message = str(value or "生成失败")
|
||||
message_lower = message.lower()
|
||||
for keyword, friendly_message in O1KEY_IMAGE_ERROR_CONTENT_MESSAGES.items():
|
||||
if keyword.lower() in message_lower:
|
||||
return friendly_message
|
||||
return message
|
||||
|
||||
|
||||
O1KEY_VIDEO_COPYRIGHT_MESSAGES = {
|
||||
"audio": "请求失败，输出视频中音频触发版权限制！",
|
||||
"video": "请求失败，输出视频触发版权限制！",
|
||||
"content": "请求失败，提示词触发版权限制！",
|
||||
"real": "请求失败，真人内容触发版权限制！",
|
||||
"unknown": "请求失败，生成内容触发版权限制！",
|
||||
}
|
||||
O1KEY_VIDEO_REVIEW_MESSAGES = {
|
||||
"audio": "请求失败，输出视频中音频触发审查！",
|
||||
"video": "请求失败，输出视频触发审查！",
|
||||
"content": "请求失败，提示词触发审查！",
|
||||
"real": "请求失败，真人内容触发审查！",
|
||||
"unknown": "请求失败，生成内容触发审查！",
|
||||
}
|
||||
O1KEY_VIDEO_REVIEW_MARKERS = (
|
||||
"sensitive",
|
||||
"safety",
|
||||
"moderation",
|
||||
"policy violation",
|
||||
"policy_violation",
|
||||
"policyviolation",
|
||||
"unsafe",
|
||||
"censor",
|
||||
"review",
|
||||
)
|
||||
O1KEY_VIDEO_SUBJECT_REVIEW_MARKERS = ("rejected", "blocked")
|
||||
_O1KEY_VIDEO_FIELD_RE = re.compile(
|
||||
r"(?:[\"']?(?:field|type|category|source)[\"']?\s*[:=]\s*[\"']?)"
|
||||
r"(audio|video|content|real)\b",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
|
||||
def _o1key_video_error_subject(message: str) -> str:
|
||||
"""Identify which part of a video request triggered an upstream review."""
|
||||
explicit_field = _O1KEY_VIDEO_FIELD_RE.search(message)
|
||||
if explicit_field:
|
||||
return explicit_field.group(1).lower()
|
||||
|
||||
message_lower = message.lower()
|
||||
if re.search(r"\baudio\b", message_lower):
|
||||
return "audio"
|
||||
if re.search(r"\breal\b|\breal[-_ ]?person\b|真人", message_lower):
|
||||
return "real"
|
||||
if re.search(r"\bprompt\b|input[-_ ]?content|提示词", message_lower):
|
||||
return "content"
|
||||
if "outputvideo" in message_lower or "output_video" in message_lower:
|
||||
return "video"
|
||||
if re.search(r"output[-_ ]+video|\bvideo\b", message_lower):
|
||||
return "video"
|
||||
if re.search(r"\bcontent\b", message_lower):
|
||||
return "content"
|
||||
return "unknown"
|
||||
|
||||
|
||||
def format_o1key_video_error(value: Any) -> str:
|
||||
"""Format review errors for every model in the o1key video generator.
|
||||
|
||||
Upstream providers use several envelope shapes, but their error text or
|
||||
field values consistently identify the reviewed subject as audio, video,
|
||||
content (the prompt), or real-person content. Copyright takes precedence
|
||||
over the broader safety-review markers.
|
||||
"""
|
||||
message = str(value or "视频生成失败")
|
||||
message_lower = message.lower()
|
||||
subject = _o1key_video_error_subject(message)
|
||||
if "copyright" in message_lower:
|
||||
return O1KEY_VIDEO_COPYRIGHT_MESSAGES[subject]
|
||||
has_review_marker = any(
|
||||
marker in message_lower for marker in O1KEY_VIDEO_REVIEW_MARKERS
|
||||
)
|
||||
has_subject_rejection = subject != "unknown" and any(
|
||||
marker in message_lower for marker in O1KEY_VIDEO_SUBJECT_REVIEW_MARKERS
|
||||
)
|
||||
if has_review_marker or has_subject_rejection:
|
||||
return O1KEY_VIDEO_REVIEW_MESSAGES[subject]
|
||||
return message
|
||||
|
||||
# 可退避重试的状态码
|
||||
RETRYABLE_STATUS_CODES = {429, 502, 503, 504, 524}
|
||||
RETRYABLE_STATUS_CODES = {429, 502, 503, 504, 524, 529}
|
||||
|
||||
# 退避重试默认参数
|
||||
DEFAULT_MAX_RETRIES = 3
|
||||
@@ -125,6 +228,101 @@ def is_retryable(status_code: int) -> bool:
|
||||
return status_code in RETRYABLE_STATUS_CODES
|
||||
|
||||
|
||||
def extract_error_detail(payload: Any) -> dict:
|
||||
"""Extract error_detail from async task payloads."""
|
||||
if not isinstance(payload, dict):
|
||||
return {}
|
||||
|
||||
queue = [payload]
|
||||
seen = set()
|
||||
while queue:
|
||||
current = queue.pop(0)
|
||||
if not isinstance(current, dict):
|
||||
continue
|
||||
|
||||
obj_id = id(current)
|
||||
if obj_id in seen:
|
||||
continue
|
||||
seen.add(obj_id)
|
||||
|
||||
for key in ("error_detail", "errorDetail", "error_details", "errorDetails"):
|
||||
detail = current.get(key)
|
||||
if isinstance(detail, dict):
|
||||
return detail
|
||||
|
||||
for key in ("data", "result", "response", "output", "task_result", "content"):
|
||||
nested = current.get(key)
|
||||
if isinstance(nested, dict):
|
||||
queue.append(nested)
|
||||
|
||||
return {}
|
||||
|
||||
|
||||
def _coerce_int(value: Any) -> Optional[int]:
|
||||
if isinstance(value, bool) or value is None:
|
||||
return None
|
||||
if isinstance(value, int):
|
||||
return value
|
||||
if isinstance(value, float) and value.is_integer():
|
||||
return int(value)
|
||||
if isinstance(value, str):
|
||||
text = value.strip()
|
||||
if not text:
|
||||
return None
|
||||
try:
|
||||
return int(float(text))
|
||||
except ValueError:
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
def extract_error_status_code(error_detail: Any) -> Optional[int]:
|
||||
"""Extract the real upstream HTTP status from error_detail."""
|
||||
if not isinstance(error_detail, dict):
|
||||
return None
|
||||
for key in (
|
||||
"upstream_status",
|
||||
"upstreamStatus",
|
||||
"upstream_status_code",
|
||||
"status_code",
|
||||
"statusCode",
|
||||
"http_status",
|
||||
"httpStatus",
|
||||
"status",
|
||||
):
|
||||
status_code = _coerce_int(error_detail.get(key))
|
||||
if status_code is not None:
|
||||
return status_code
|
||||
return None
|
||||
|
||||
|
||||
def is_error_detail_retryable(error_detail: Any) -> bool:
|
||||
if not isinstance(error_detail, dict):
|
||||
return False
|
||||
retryable = error_detail.get("retryable")
|
||||
if retryable is True:
|
||||
return True
|
||||
if isinstance(retryable, str):
|
||||
return retryable.strip().lower() in ("true", "1", "yes")
|
||||
return False
|
||||
|
||||
|
||||
def extract_retry_after_seconds(error_detail: Any) -> Optional[float]:
|
||||
if not isinstance(error_detail, dict):
|
||||
return None
|
||||
for key in ("retry_after_seconds", "retryAfterSeconds", "retry_after", "retryAfter"):
|
||||
value = error_detail.get(key)
|
||||
if isinstance(value, bool) or value is None:
|
||||
continue
|
||||
try:
|
||||
seconds = float(value)
|
||||
except (TypeError, ValueError):
|
||||
continue
|
||||
if seconds >= 0:
|
||||
return seconds
|
||||
return None
|
||||
|
||||
|
||||
def _compute_delay(attempt: int, base_delay: float, max_delay: float, backoff_factor: float) -> float:
|
||||
"""计算第 attempt 次重试的等待时间（含 jitter）"""
|
||||
delay = base_delay * (backoff_factor ** attempt)
|
||||
|
||||
+132
-2
@@ -6,7 +6,7 @@
|
||||
import base64
|
||||
from io import BytesIO
|
||||
import json
|
||||
from typing import Callable, List, Tuple
|
||||
from typing import Any, Callable, List, Tuple
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
@@ -29,6 +29,7 @@ def tensor_to_pil(tensor: torch.Tensor) -> List[Image.Image]:
|
||||
... img.save(f"output_{i}.png")
|
||||
"""
|
||||
images = []
|
||||
source_metadata = getattr(tensor, "_o1key_source_metadata", None)
|
||||
|
||||
# 转换为 numpy 数组
|
||||
np_images = tensor.cpu().numpy()
|
||||
@@ -42,6 +43,22 @@ def tensor_to_pil(tensor: torch.Tensor) -> List[Image.Image]:
|
||||
|
||||
# 创建 PIL Image
|
||||
img = Image.fromarray(img_array)
|
||||
if isinstance(source_metadata, list) and i < len(source_metadata):
|
||||
metadata = source_metadata[i]
|
||||
if isinstance(metadata, dict):
|
||||
source_format = metadata.get("format")
|
||||
if source_format:
|
||||
img.format = source_format
|
||||
setattr(img, "_o1key_original_format", source_format)
|
||||
source_path = metadata.get("path")
|
||||
if source_path:
|
||||
setattr(img, "_o1key_original_path", source_path)
|
||||
source_filename = metadata.get("filename")
|
||||
if source_filename:
|
||||
setattr(img, "_o1key_original_filename", source_filename)
|
||||
source_bytes = metadata.get("bytes")
|
||||
if isinstance(source_bytes, bytes):
|
||||
setattr(img, "_o1key_original_bytes", source_bytes)
|
||||
images.append(img)
|
||||
|
||||
return images
|
||||
@@ -63,8 +80,16 @@ def pil_to_tensor(images: List[Image.Image]) -> torch.Tensor:
|
||||
>>> print(tensor.shape) # [1, H, W, 3]
|
||||
"""
|
||||
tensors = []
|
||||
source_metadata = []
|
||||
|
||||
for img in images:
|
||||
source_metadata.append({
|
||||
"format": getattr(img, "_o1key_original_format", None) or img.format,
|
||||
"path": getattr(img, "_o1key_original_path", None),
|
||||
"filename": getattr(img, "_o1key_original_filename", None),
|
||||
"bytes": getattr(img, "_o1key_original_bytes", None),
|
||||
"modified": bool(getattr(img, "_o1key_pixels_modified", False)),
|
||||
})
|
||||
# 确保是 RGB 模式
|
||||
if img.mode != 'RGB':
|
||||
img = img.convert('RGB')
|
||||
@@ -81,7 +106,9 @@ def pil_to_tensor(images: List[Image.Image]) -> torch.Tensor:
|
||||
batch_tensor = np.stack(tensors, axis=0)
|
||||
|
||||
# 转换为 torch tensor
|
||||
return torch.from_numpy(batch_tensor)
|
||||
tensor = torch.from_numpy(batch_tensor)
|
||||
tensor._o1key_source_metadata = source_metadata
|
||||
return tensor
|
||||
|
||||
|
||||
def encode_image_to_base64(image: Image.Image, format: str = "PNG") -> str:
|
||||
@@ -366,3 +393,106 @@ def parse_batch_prompts(prompt: str) -> List[str]:
|
||||
raise ValueError("批量提示词模式下，所有提示词都为空，请至少提供一个有效的提示词")
|
||||
|
||||
return filtered_prompts
|
||||
|
||||
|
||||
def expand_batch_prompt_tasks(prompt: str, images_per_prompt: int) -> List[str]:
|
||||
"""Expand one or more prompts into prompt-major generation tasks."""
|
||||
count = int(images_per_prompt)
|
||||
if count < 1:
|
||||
raise ValueError("每条提示词的生图数量必须大于 0")
|
||||
prompts = parse_batch_prompts(prompt) or [prompt.strip()]
|
||||
return [task_prompt for task_prompt in prompts for _ in range(count)]
|
||||
|
||||
|
||||
IMAGE_BATCH_MODE_GROUP_TO_MODELS = "一组搭配+多模特"
|
||||
IMAGE_BATCH_MODE_CARTESIAN = "全部搭配×全部模特"
|
||||
IMAGE_BATCH_MODE_SINGLE_REFERENCES = "单图素材批量"
|
||||
IMAGE_BATCH_MODES = (
|
||||
IMAGE_BATCH_MODE_GROUP_TO_MODELS,
|
||||
IMAGE_BATCH_MODE_CARTESIAN,
|
||||
IMAGE_BATCH_MODE_SINGLE_REFERENCES,
|
||||
)
|
||||
|
||||
|
||||
def expand_image_generation_tasks(
|
||||
prompt: str,
|
||||
images_per_pair: int,
|
||||
*,
|
||||
batch_enabled: bool = False,
|
||||
batch_mode: str = IMAGE_BATCH_MODE_GROUP_TO_MODELS,
|
||||
reference_count: int = 0,
|
||||
model_reference_count: int = 0,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Expand prompts and reference pairing into stable prompt-major tasks.
|
||||
|
||||
Normal mode keeps the historical ``prompt x image_count`` ordering. In
|
||||
batch mode, ``reference_indices`` point into the outfit/reference list and
|
||||
``model_reference_indices`` point into the separately uploaded model list.
|
||||
Single-reference batch mode creates one task per source image and ignores
|
||||
the separately uploaded target list. Keeping indexes instead of file data
|
||||
makes the plan safe to serialize and lets direct and background execution
|
||||
share the exact same ordering.
|
||||
"""
|
||||
count = int(images_per_pair)
|
||||
if count < 1:
|
||||
raise ValueError("每个组合的生图数量必须大于 0")
|
||||
prompts = parse_batch_prompts(prompt) or [prompt.strip()]
|
||||
if not batch_enabled:
|
||||
pairings = [{
|
||||
"reference_indices": tuple(range(max(0, int(reference_count)))),
|
||||
"model_reference_indices": (),
|
||||
"outfit_index": None,
|
||||
"model_index": None,
|
||||
}]
|
||||
else:
|
||||
if batch_mode not in IMAGE_BATCH_MODES:
|
||||
raise ValueError("批量模式无效")
|
||||
outfit_total = max(0, int(reference_count))
|
||||
model_total = max(0, int(model_reference_count))
|
||||
if batch_mode == IMAGE_BATCH_MODE_SINGLE_REFERENCES:
|
||||
if outfit_total < 1:
|
||||
raise ValueError("单图批量至少需要上传1张素材图")
|
||||
pairings = [
|
||||
{
|
||||
"reference_indices": (outfit_index,),
|
||||
"model_reference_indices": (),
|
||||
"outfit_index": outfit_index,
|
||||
"model_index": None,
|
||||
}
|
||||
for outfit_index in range(outfit_total)
|
||||
]
|
||||
elif model_total < 1:
|
||||
raise ValueError("批量出图至少需要上传1张目标图")
|
||||
elif outfit_total < 1:
|
||||
raise ValueError("批量出图至少需要上传1张素材图")
|
||||
elif batch_mode == IMAGE_BATCH_MODE_GROUP_TO_MODELS:
|
||||
pairings = [
|
||||
{
|
||||
"reference_indices": tuple(range(outfit_total)),
|
||||
"model_reference_indices": (model_index,),
|
||||
"outfit_index": None,
|
||||
"model_index": model_index,
|
||||
}
|
||||
for model_index in range(model_total)
|
||||
]
|
||||
else:
|
||||
pairings = [
|
||||
{
|
||||
"reference_indices": (outfit_index,),
|
||||
"model_reference_indices": (model_index,),
|
||||
"outfit_index": outfit_index,
|
||||
"model_index": model_index,
|
||||
}
|
||||
for outfit_index in range(outfit_total)
|
||||
for model_index in range(model_total)
|
||||
]
|
||||
|
||||
return [
|
||||
{
|
||||
"prompt": task_prompt,
|
||||
**pairing,
|
||||
}
|
||||
for task_prompt in prompts
|
||||
for pairing in pairings
|
||||
for _ in range(count)
|
||||
]
|
||||
|
||||
@@ -0,0 +1,231 @@
|
||||
"""Media validation helpers for MiniMax-H3 reference inputs."""
|
||||
|
||||
import io
|
||||
import os
|
||||
from typing import Any, Dict, Iterable
|
||||
|
||||
|
||||
MB = 1024 * 1024
|
||||
MIN_MEDIA_DIMENSION = 256
|
||||
MAX_MEDIA_DIMENSION = 5760
|
||||
MIN_MEDIA_RATIO = 0.4
|
||||
MAX_MEDIA_RATIO = 2.5
|
||||
|
||||
MAX_IMAGE_BYTES = 30 * MB
|
||||
MAX_REFERENCE_IMAGES = 9
|
||||
|
||||
MAX_VIDEO_BYTES = 50 * MB
|
||||
MAX_REFERENCE_VIDEOS = 3
|
||||
MIN_REFERENCE_DURATION = 2.0
|
||||
MAX_REFERENCE_DURATION = 15.0
|
||||
MAX_TOTAL_VIDEO_DURATION = 15.0
|
||||
MIN_VIDEO_FPS = 23.976
|
||||
MAX_VIDEO_FPS = 60.0
|
||||
|
||||
MAX_AUDIO_BYTES = 15 * MB
|
||||
MAX_REFERENCE_AUDIOS = 3
|
||||
MAX_TOTAL_AUDIO_DURATION = 15.0
|
||||
|
||||
|
||||
def _validate_dimensions(width: int, height: int, label: str) -> None:
|
||||
if not (
|
||||
MIN_MEDIA_DIMENSION <= width <= MAX_MEDIA_DIMENSION
|
||||
and MIN_MEDIA_DIMENSION <= height <= MAX_MEDIA_DIMENSION
|
||||
):
|
||||
raise ValueError(
|
||||
f"{label}宽高必须均在 {MIN_MEDIA_DIMENSION}～{MAX_MEDIA_DIMENSION}px，"
|
||||
f"当前为 {width}x{height}。"
|
||||
)
|
||||
ratio = width / height
|
||||
if not MIN_MEDIA_RATIO <= ratio <= MAX_MEDIA_RATIO:
|
||||
raise ValueError(
|
||||
f"{label}宽高比必须在 0.4～2.5，当前为 {ratio:.3f}（{width}:{height}）。"
|
||||
)
|
||||
|
||||
|
||||
def validate_image(image, label: str = "图片") -> Dict[str, Any]:
|
||||
"""Validate the exact PNG bytes sent by the shared uploader."""
|
||||
width, height = image.size
|
||||
_validate_dimensions(int(width), int(height), label)
|
||||
|
||||
buffer = io.BytesIO()
|
||||
image.save(buffer, format="PNG")
|
||||
size = buffer.tell()
|
||||
if size > MAX_IMAGE_BYTES:
|
||||
raise ValueError(
|
||||
f"{label}转为 PNG 后不能超过 30MB，当前为 {size / MB:.2f}MB。"
|
||||
)
|
||||
return {"width": width, "height": height, "size": size}
|
||||
|
||||
|
||||
def _get_video_source(video):
|
||||
if hasattr(video, "get_stream_source"):
|
||||
source = video.get_stream_source()
|
||||
elif isinstance(video, dict):
|
||||
source = (
|
||||
video.get("video")
|
||||
or video.get("path")
|
||||
or video.get("file")
|
||||
or video.get("filename")
|
||||
or video.get("source_path")
|
||||
)
|
||||
elif isinstance(video, (str, os.PathLike, io.BytesIO)):
|
||||
source = video
|
||||
else:
|
||||
source = None
|
||||
for attr in ("source_path", "path", "video", "file", "filename"):
|
||||
if hasattr(video, attr):
|
||||
source = getattr(video, attr)
|
||||
if source:
|
||||
break
|
||||
if source is None:
|
||||
raise ValueError("无法获取参考视频数据。")
|
||||
return source
|
||||
|
||||
|
||||
def _source_for_probe(source):
|
||||
if isinstance(source, io.BytesIO):
|
||||
data = source.getvalue()
|
||||
return io.BytesIO(data), len(data)
|
||||
try:
|
||||
path = os.fspath(source)
|
||||
except TypeError:
|
||||
raise ValueError(f"无法读取参考视频来源：{type(source).__name__}") from None
|
||||
if not os.path.isfile(path):
|
||||
raise ValueError(f"参考视频文件不存在：{path}")
|
||||
return path, os.path.getsize(path)
|
||||
|
||||
|
||||
def probe_video(video) -> Dict[str, Any]:
|
||||
"""Inspect a ComfyUI VIDEO using PyAV without transcoding it."""
|
||||
try:
|
||||
import av
|
||||
except ImportError:
|
||||
raise RuntimeError("当前 ComfyUI 环境缺少 PyAV，无法校验 MiniMax H3 参考视频。") from None
|
||||
|
||||
source = _get_video_source(video)
|
||||
probe_source, size = _source_for_probe(source)
|
||||
try:
|
||||
container = av.open(probe_source)
|
||||
except Exception as exc:
|
||||
raise ValueError(f"无法读取参考视频：{exc}") from None
|
||||
|
||||
try:
|
||||
video_streams = list(container.streams.video)
|
||||
if not video_streams:
|
||||
raise ValueError("参考视频不包含视频轨道。")
|
||||
stream = video_streams[0]
|
||||
width = int(stream.codec_context.width or 0)
|
||||
height = int(stream.codec_context.height or 0)
|
||||
video_codec = str(stream.codec_context.name or "").lower()
|
||||
rate = stream.average_rate or stream.base_rate or stream.guessed_rate
|
||||
fps = float(rate) if rate else 0.0
|
||||
|
||||
duration = None
|
||||
if stream.duration is not None and stream.time_base is not None:
|
||||
duration = float(stream.duration * stream.time_base)
|
||||
elif container.duration is not None:
|
||||
duration = float(container.duration / av.time_base)
|
||||
duration = float(duration or 0.0)
|
||||
|
||||
audio_codecs = {
|
||||
str(audio_stream.codec_context.name or "").lower()
|
||||
for audio_stream in container.streams.audio
|
||||
}
|
||||
format_names = {
|
||||
name.strip().lower()
|
||||
for name in str(container.format.name or "").split(",")
|
||||
if name.strip()
|
||||
}
|
||||
finally:
|
||||
container.close()
|
||||
|
||||
return {
|
||||
"size": size,
|
||||
"width": width,
|
||||
"height": height,
|
||||
"duration": duration,
|
||||
"fps": fps,
|
||||
"video_codec": video_codec,
|
||||
"audio_codecs": audio_codecs,
|
||||
"format_names": format_names,
|
||||
}
|
||||
|
||||
|
||||
def validate_video_info(info: Dict[str, Any], label: str = "参考视频") -> None:
|
||||
size = int(info.get("size") or 0)
|
||||
if size <= 0:
|
||||
raise ValueError(f"{label}文件为空。")
|
||||
if size > MAX_VIDEO_BYTES:
|
||||
raise ValueError(f"{label}不能超过 50MB，当前为 {size / MB:.2f}MB。")
|
||||
|
||||
format_names = set(info.get("format_names") or ())
|
||||
if not format_names.intersection({"mp4", "mov"}):
|
||||
raise ValueError(f"{label}容器仅支持 MP4/MOV，当前为 {sorted(format_names) or '未知'}。")
|
||||
|
||||
duration = float(info.get("duration") or 0.0)
|
||||
if not MIN_REFERENCE_DURATION <= duration <= MAX_REFERENCE_DURATION:
|
||||
raise ValueError(f"{label}时长必须为 2～15 秒，当前为 {duration:.3f} 秒。")
|
||||
|
||||
fps = float(info.get("fps") or 0.0)
|
||||
if not MIN_VIDEO_FPS <= fps <= MAX_VIDEO_FPS:
|
||||
raise ValueError(f"{label}帧率必须为 23.976～60 FPS，当前为 {fps:.3f} FPS。")
|
||||
|
||||
def validate_reference_videos(videos: Iterable[Any]) -> list[Dict[str, Any]]:
|
||||
videos = list(videos)
|
||||
if len(videos) > MAX_REFERENCE_VIDEOS:
|
||||
raise ValueError(f"参考视频最多 {MAX_REFERENCE_VIDEOS} 个，当前为 {len(videos)} 个。")
|
||||
infos = []
|
||||
for index, video in enumerate(videos, start=1):
|
||||
info = probe_video(video)
|
||||
validate_video_info(info, f"参考视频{index}")
|
||||
infos.append(info)
|
||||
total_duration = sum(float(info["duration"]) for info in infos)
|
||||
if total_duration > MAX_TOTAL_VIDEO_DURATION:
|
||||
raise ValueError(f"参考视频总时长不能超过 15 秒，当前为 {total_duration:.3f} 秒。")
|
||||
return infos
|
||||
|
||||
|
||||
def inspect_audio(audio, label: str = "参考音频") -> Dict[str, Any]:
|
||||
if not isinstance(audio, dict):
|
||||
raise ValueError(f"{label}数据格式无效。")
|
||||
waveform = audio.get("waveform")
|
||||
sample_rate = int(audio.get("sample_rate") or 0)
|
||||
if waveform is None or sample_rate <= 0:
|
||||
raise ValueError(f"{label}缺少 waveform 或 sample_rate。")
|
||||
|
||||
shape = tuple(int(v) for v in waveform.shape)
|
||||
if not shape:
|
||||
raise ValueError(f"{label}波形为空。")
|
||||
if len(shape) == 3 and shape[0] != 1:
|
||||
raise ValueError(f"{label}仅支持一个音频批次，当前批次为 {shape[0]}。")
|
||||
samples = shape[-1]
|
||||
if samples <= 0:
|
||||
raise ValueError(f"{label}波形为空。")
|
||||
duration = samples / sample_rate
|
||||
|
||||
# The shared uploader currently downmixes to mono 16-bit PCM WAV.
|
||||
encoded_size = 44 + samples * 2
|
||||
if encoded_size > MAX_AUDIO_BYTES:
|
||||
raise ValueError(
|
||||
f"{label}编码为 WAV 后不能超过 15MB，预计为 {encoded_size / MB:.2f}MB。"
|
||||
)
|
||||
if not MIN_REFERENCE_DURATION <= duration <= MAX_REFERENCE_DURATION:
|
||||
raise ValueError(f"{label}时长必须为 2～15 秒，当前为 {duration:.3f} 秒。")
|
||||
return {
|
||||
"sample_rate": sample_rate,
|
||||
"samples": samples,
|
||||
"duration": duration,
|
||||
"encoded_size": encoded_size,
|
||||
}
|
||||
|
||||
|
||||
def validate_reference_audios(audios: Iterable[Any]) -> list[Dict[str, Any]]:
|
||||
audios = list(audios)
|
||||
if len(audios) > MAX_REFERENCE_AUDIOS:
|
||||
raise ValueError(f"参考音频最多 {MAX_REFERENCE_AUDIOS} 个，当前为 {len(audios)} 个。")
|
||||
infos = [inspect_audio(audio, f"参考音频{index}") for index, audio in enumerate(audios, 1)]
|
||||
total_duration = sum(float(info["duration"]) for info in infos)
|
||||
if total_duration > MAX_TOTAL_AUDIO_DURATION:
|
||||
raise ValueError(f"参考音频总时长不能超过 15 秒，当前为 {total_duration:.3f} 秒。")
|
||||
return infos
|
||||
+1357
-213
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,49 @@
|
||||
"""Nano Banana 节点的展示模型、线路与 API 模型名映射。"""
|
||||
|
||||
NANO_BANANA_MODEL_OPTIONS = [
|
||||
"Nano Banana 2",
|
||||
"Nano Banana Pro",
|
||||
"Nano Banana 2 Lite",
|
||||
"Nano Banana",
|
||||
]
|
||||
|
||||
NANO_BANANA_ROUTE_OPTIONS = ["畅速", "直连", "专线"]
|
||||
|
||||
NANO_BANANA_MODEL_MATRIX = {
|
||||
("Nano Banana Pro", "畅速"): "gemini-3-pro-image-c-sp",
|
||||
("Nano Banana 2", "畅速"): "gemini-3.1-flash-image-c-sp",
|
||||
("Nano Banana 2 Lite", "畅速"): "gemini-3.1-flash-lite-image-c-sp",
|
||||
("Nano Banana", "畅速"): "nano-banana",
|
||||
("Nano Banana Pro", "直连"): "gemini-3-pro-image-c-sd",
|
||||
("Nano Banana 2", "直连"): "gemini-3.1-flash-image-c-sd",
|
||||
("Nano Banana 2 Lite", "直连"): "gemini-3.1-flash-lite-image-c-sd",
|
||||
("Nano Banana", "直连"): "nano-banana",
|
||||
("Nano Banana Pro", "专线"): "gemini-3-pro-image",
|
||||
("Nano Banana 2", "专线"): "gemini-3.1-flash-image",
|
||||
("Nano Banana 2 Lite", "专线"): "gemini-3.1-flash-lite-image",
|
||||
("Nano Banana", "专线"): "gemini-2.5-flash-image",
|
||||
}
|
||||
|
||||
_LEGACY_BILLING_ROUTE_MAP = {
|
||||
"特价": "畅速",
|
||||
"官方": "专线",
|
||||
}
|
||||
|
||||
|
||||
def normalize_nano_banana_route(route: str) -> str:
|
||||
"""兼容旧“计费”值，并校验新模型线路。"""
|
||||
normalized = _LEGACY_BILLING_ROUTE_MAP.get(route, route)
|
||||
if normalized not in NANO_BANANA_ROUTE_OPTIONS:
|
||||
raise ValueError(
|
||||
f"模型线路 '{route}' 无效，支持的线路：{', '.join(NANO_BANANA_ROUTE_OPTIONS)}"
|
||||
)
|
||||
return normalized
|
||||
|
||||
|
||||
def resolve_nano_banana_model(model_name: str, route: str) -> str:
|
||||
"""由界面模型和模型线路解析实际 API 模型名。"""
|
||||
normalized_route = normalize_nano_banana_route(route)
|
||||
model = NANO_BANANA_MODEL_MATRIX.get((model_name, normalized_route))
|
||||
if model is None:
|
||||
raise ValueError(f"模型 '{model_name}' 不支持模型线路 '{normalized_route}'")
|
||||
return model
|
||||
@@ -0,0 +1,229 @@
|
||||
"""Shared capability catalog for the unified o1key image generator."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from .nano_banana_models import (
|
||||
NANO_BANANA_MODEL_OPTIONS,
|
||||
NANO_BANANA_ROUTE_OPTIONS,
|
||||
)
|
||||
|
||||
|
||||
GPT_IMAGE_MODEL_OPTIONS = [
|
||||
"gpt-image-2",
|
||||
"gpt-image-2.5-sunburst",
|
||||
"gpt-image-2.5-flare",
|
||||
]
|
||||
SEEDREAM_MODEL_OPTIONS = ["Seedream 5.0 Pro"]
|
||||
UNIFIED_IMAGE_MODEL_OPTIONS = [
|
||||
*NANO_BANANA_MODEL_OPTIONS[:2],
|
||||
*GPT_IMAGE_MODEL_OPTIONS,
|
||||
*SEEDREAM_MODEL_OPTIONS,
|
||||
*NANO_BANANA_MODEL_OPTIONS[2:],
|
||||
]
|
||||
UNIFIED_IMAGE_ROUTE_OPTIONS = list(NANO_BANANA_ROUTE_OPTIONS)
|
||||
|
||||
UNIFIED_IMAGE_COUNTS = [1, 2, 4, 9]
|
||||
GPT_IMAGE_COUNTS = list(range(1, 9))
|
||||
MAX_UNIFIED_REFERENCE_IMAGES = 10
|
||||
MAX_UNIFIED_BATCH_IMAGES = 50
|
||||
MAX_UNIFIED_IMAGE_TASKS = 1000
|
||||
UNIFIED_IMAGE_SMART_RESOLUTION = "智能"
|
||||
|
||||
BANANA_RESOLUTION_OPTIONS = ["512", "1K", "2K", "4K"]
|
||||
BANANA_ASPECT_RATIO_OPTIONS = [
|
||||
"智能", "1:1", "1:4", "1:8", "2:3", "3:2", "3:4",
|
||||
"4:1", "4:3", "4:5", "5:4", "8:1", "9:16", "16:9", "21:9",
|
||||
]
|
||||
|
||||
GPT_IMAGE_RESOLUTION_OPTIONS = ["1K", "2K", "4K"]
|
||||
GPT_IMAGE_ASPECT_RATIO_OPTIONS = [
|
||||
"智能", "1:1", "3:2", "2:3", "4:3", "3:4", "16:9", "9:16",
|
||||
]
|
||||
GPT_IMAGE_EXACT_SIZE_OPTIONS = [
|
||||
"智能",
|
||||
"1024x1024（1K 正方形 1:1）",
|
||||
"1536x1024（1K 横版 3:2）",
|
||||
"1024x1536（1K 竖版 2:3）",
|
||||
"1360x1024（1K 横版 4:3）",
|
||||
"1024x1360（1K 竖版 3:4）",
|
||||
"1824x1024（1K 横版 16:9）",
|
||||
"1024x1824（1K 竖版 9:16）",
|
||||
"2048x2048（2K 正方形 1:1）",
|
||||
"3072x2048（2K 横版 3:2）",
|
||||
"2048x3072（2K 竖版 2:3）",
|
||||
"2736x2048（2K 横版 4:3）",
|
||||
"2048x2736（2K 竖版 3:4）",
|
||||
"3648x2048（2K 横版 16:9）",
|
||||
"2048x3648（2K 竖版 9:16）",
|
||||
"2880x2880（4K 正方形 1:1）",
|
||||
"3504x2336（4K 横版 3:2）",
|
||||
"2336x3504（4K 竖版 2:3）",
|
||||
"3264x2448（4K 横版 4:3）",
|
||||
"2448x3264（4K 竖版 3:4）",
|
||||
"3840x2160（4K 横版 16:9）",
|
||||
"2160x3840（4K 竖版 9:16）",
|
||||
]
|
||||
GPT_IMAGE_QUALITY_OPTIONS = ["高", "中", "低", "自动"]
|
||||
GPT_IMAGE_25_QUALITY_OPTIONS = [*GPT_IMAGE_QUALITY_OPTIONS, "超高", "最高"]
|
||||
GPT_IMAGE_OUTPUT_FORMAT_OPTIONS = ["jpeg", "png", "webp"]
|
||||
GPT_IMAGE_BACKGROUND_OPTIONS = ["auto", "transparent", "opaque"]
|
||||
GPT_IMAGE_SIZE_MATRIX = {
|
||||
("1K", "1:1"): "1024x1024",
|
||||
("1K", "3:2"): "1536x1024",
|
||||
("1K", "2:3"): "1024x1536",
|
||||
("1K", "4:3"): "1360x1024",
|
||||
("1K", "3:4"): "1024x1360",
|
||||
("1K", "16:9"): "1824x1024",
|
||||
("1K", "9:16"): "1024x1824",
|
||||
("2K", "1:1"): "2048x2048",
|
||||
("2K", "3:2"): "3072x2048",
|
||||
("2K", "2:3"): "2048x3072",
|
||||
("2K", "4:3"): "2736x2048",
|
||||
("2K", "3:4"): "2048x2736",
|
||||
("2K", "16:9"): "3648x2048",
|
||||
("2K", "9:16"): "2048x3648",
|
||||
("4K", "1:1"): "2880x2880",
|
||||
("4K", "3:2"): "3504x2336",
|
||||
("4K", "2:3"): "2336x3504",
|
||||
("4K", "4:3"): "3264x2448",
|
||||
("4K", "3:4"): "2448x3264",
|
||||
("4K", "16:9"): "3840x2160",
|
||||
("4K", "9:16"): "2160x3840",
|
||||
}
|
||||
|
||||
SEEDREAM_RESOLUTION_OPTIONS = ["1K", "2K"]
|
||||
SEEDREAM_LAYER_RESOLUTION_OPTIONS = ["auto", "1K", "1.5K", "2K"]
|
||||
SEEDREAM_ASPECT_RATIO_OPTIONS = [
|
||||
"智能", "1:1", "4:3", "3:4", "16:9", "9:16", "3:2", "2:3", "21:9",
|
||||
]
|
||||
SEEDREAM_OUTPUT_FORMAT_OPTIONS = ["png", "jpeg"]
|
||||
SEEDREAM_SIZE_MATRIX = {
|
||||
("1K", "1:1"): "1024x1024",
|
||||
("1K", "4:3"): "1152x864",
|
||||
("1K", "3:4"): "864x1152",
|
||||
("1K", "16:9"): "1424x800",
|
||||
("1K", "9:16"): "800x1424",
|
||||
("1K", "3:2"): "1248x832",
|
||||
("1K", "2:3"): "832x1248",
|
||||
("1K", "21:9"): "1568x672",
|
||||
("2K", "1:1"): "2048x2048",
|
||||
("2K", "4:3"): "2368x1776",
|
||||
("2K", "3:4"): "1776x2368",
|
||||
("2K", "16:9"): "2816x1584",
|
||||
("2K", "9:16"): "1584x2816",
|
||||
("2K", "3:2"): "2496x1664",
|
||||
("2K", "2:3"): "1664x2496",
|
||||
("2K", "21:9"): "3136x1344",
|
||||
}
|
||||
|
||||
|
||||
def is_gpt_image_model(model: str) -> bool:
|
||||
return model in GPT_IMAGE_MODEL_OPTIONS
|
||||
|
||||
|
||||
def resolve_gpt_image_quality(model_name: str, quality: str) -> str:
|
||||
"""Resolve a displayed GPT quality for the selected model family."""
|
||||
if quality in {"超高", "最高"}:
|
||||
if model_name not in GPT_IMAGE_MODEL_OPTIONS or not model_name.startswith("gpt-image-2.5-"):
|
||||
raise ValueError("超高和最高质量仅支持 GPT Image 2.5 系列模型")
|
||||
return {"超高": "xhigh", "最高": "max"}[quality]
|
||||
try:
|
||||
return {"高": "high", "中": "medium", "低": "low", "自动": "auto"}[quality]
|
||||
except KeyError:
|
||||
raise ValueError("GPT Image 质量参数无效") from None
|
||||
|
||||
|
||||
def is_seedream_model(model: str) -> bool:
|
||||
return model in SEEDREAM_MODEL_OPTIONS
|
||||
|
||||
|
||||
def resolve_gpt_image_size(value: str, aspect_ratio: str = "智能") -> str:
|
||||
"""Convert a GPT Image display-size value into the API value."""
|
||||
value = (value or "").strip()
|
||||
aspect_ratio = (aspect_ratio or "智能").strip()
|
||||
normalized_resolution = value.upper()
|
||||
if normalized_resolution in GPT_IMAGE_RESOLUTION_OPTIONS:
|
||||
effective_ratio = "1:1" if aspect_ratio == "智能" else aspect_ratio
|
||||
mapped = GPT_IMAGE_SIZE_MATRIX.get((normalized_resolution, effective_ratio))
|
||||
if mapped is None:
|
||||
raise ValueError(
|
||||
f"GPT Image 不支持分辨率 {value} 与宽高比 {aspect_ratio} 的组合"
|
||||
)
|
||||
return mapped
|
||||
if not value or value == "智能" or value.lower() == "auto":
|
||||
return "auto"
|
||||
|
||||
first_part = value.split("（")[0].strip()
|
||||
normalized_size = first_part.lower().replace("*", "x").replace("×", "x")
|
||||
size_parts = [part.strip() for part in normalized_size.split("x")]
|
||||
if len(size_parts) == 2 and all(part.isdigit() for part in size_parts):
|
||||
return f"{int(size_parts[0])}x{int(size_parts[1])}"
|
||||
|
||||
if first_part in {"auto", "1024x1024", "1K", "2K", "4K"}:
|
||||
return first_part
|
||||
if "4K" in value:
|
||||
return "4K"
|
||||
if "2K" in value:
|
||||
return "2K"
|
||||
if "1K" in value:
|
||||
return "1K"
|
||||
return "auto"
|
||||
|
||||
|
||||
def resolve_seedream_size(resolution: str, aspect_ratio: str = "智能") -> str:
|
||||
"""Resolve Seedream's UI resolution and ratio into its exact size value."""
|
||||
normalized_resolution = str(resolution or "").strip().upper()
|
||||
normalized_ratio = str(aspect_ratio or "智能").strip()
|
||||
if normalized_resolution not in SEEDREAM_RESOLUTION_OPTIONS:
|
||||
raise ValueError(f"Seedream 分辨率无效：{resolution}")
|
||||
effective_ratio = "1:1" if normalized_ratio == "智能" else normalized_ratio
|
||||
mapped = SEEDREAM_SIZE_MATRIX.get((normalized_resolution, effective_ratio))
|
||||
if mapped is None:
|
||||
raise ValueError(
|
||||
f"Seedream 不支持分辨率 {resolution} 与宽高比 {aspect_ratio} 的组合"
|
||||
)
|
||||
return mapped
|
||||
|
||||
|
||||
def resolve_seedream_layer_size(resolution: str) -> str:
|
||||
"""Normalize Seedream's layer-decomposition size without applying a ratio."""
|
||||
value = str(resolution or "auto").strip()
|
||||
normalized = "auto" if value.lower() == "auto" else value.upper()
|
||||
if normalized not in SEEDREAM_LAYER_RESOLUTION_OPTIONS:
|
||||
raise ValueError(f"Seedream 图层拆分分辨率无效：{resolution}")
|
||||
return normalized
|
||||
|
||||
|
||||
__all__ = [
|
||||
"BANANA_ASPECT_RATIO_OPTIONS",
|
||||
"BANANA_RESOLUTION_OPTIONS",
|
||||
"GPT_IMAGE_ASPECT_RATIO_OPTIONS",
|
||||
"GPT_IMAGE_BACKGROUND_OPTIONS",
|
||||
"GPT_IMAGE_COUNTS",
|
||||
"GPT_IMAGE_EXACT_SIZE_OPTIONS",
|
||||
"GPT_IMAGE_MODEL_OPTIONS",
|
||||
"GPT_IMAGE_25_QUALITY_OPTIONS",
|
||||
"GPT_IMAGE_OUTPUT_FORMAT_OPTIONS",
|
||||
"GPT_IMAGE_QUALITY_OPTIONS",
|
||||
"GPT_IMAGE_RESOLUTION_OPTIONS",
|
||||
"GPT_IMAGE_SIZE_MATRIX",
|
||||
"SEEDREAM_ASPECT_RATIO_OPTIONS",
|
||||
"SEEDREAM_LAYER_RESOLUTION_OPTIONS",
|
||||
"SEEDREAM_MODEL_OPTIONS",
|
||||
"SEEDREAM_OUTPUT_FORMAT_OPTIONS",
|
||||
"SEEDREAM_RESOLUTION_OPTIONS",
|
||||
"SEEDREAM_SIZE_MATRIX",
|
||||
"MAX_UNIFIED_BATCH_IMAGES",
|
||||
"MAX_UNIFIED_IMAGE_TASKS",
|
||||
"MAX_UNIFIED_REFERENCE_IMAGES",
|
||||
"UNIFIED_IMAGE_COUNTS",
|
||||
"UNIFIED_IMAGE_MODEL_OPTIONS",
|
||||
"UNIFIED_IMAGE_ROUTE_OPTIONS",
|
||||
"UNIFIED_IMAGE_SMART_RESOLUTION",
|
||||
"is_gpt_image_model",
|
||||
"is_seedream_model",
|
||||
"resolve_gpt_image_size",
|
||||
"resolve_gpt_image_quality",
|
||||
"resolve_seedream_size",
|
||||
"resolve_seedream_layer_size",
|
||||
]
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,598 @@
|
||||
"""Format-aware saving for the unified o1key image workflow.
|
||||
|
||||
The generator keeps provider bytes in memory or in ComfyUI's temp directory.
|
||||
Only the save node promotes those results into a permanent save directory.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import struct
|
||||
import threading
|
||||
import uuid
|
||||
import zlib
|
||||
from io import BytesIO
|
||||
from typing import Any, Iterable
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from PIL.PngImagePlugin import PngInfo
|
||||
|
||||
|
||||
SAVE_FORMAT_OPTIONS = ("原始", "png", "jpg", "webp")
|
||||
SAVE_NAMING_RULE_OPTIONS = ("和主图一致", "自然数字", "自定义前缀")
|
||||
DEFAULT_SAVE_NAMING_RULE = "自定义前缀"
|
||||
_FORMAT_EXTENSIONS = {"PNG": "png", "JPEG": "jpg", "WEBP": "webp"}
|
||||
_PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"
|
||||
_SAVE_LOCK = threading.Lock()
|
||||
_UNSAFE_FILENAME_CHARS = re.compile(r'[<>:"/\\|?*\x00-\x1f]')
|
||||
_WINDOWS_RESERVED_STEMS = {
|
||||
"CON", "PRN", "AUX", "NUL",
|
||||
*(f"COM{index}" for index in range(1, 10)),
|
||||
*(f"LPT{index}" for index in range(1, 10)),
|
||||
}
|
||||
|
||||
|
||||
def normalize_save_format(value: Any) -> str:
|
||||
normalized = str(value or "原始").strip().lower()
|
||||
if normalized == "jpeg":
|
||||
normalized = "jpg"
|
||||
if normalized == "原始":
|
||||
return "原始"
|
||||
if normalized not in {"png", "jpg", "webp"}:
|
||||
raise ValueError("保存格式仅支持：原始、png、jpg、webp")
|
||||
return normalized
|
||||
|
||||
|
||||
def normalize_save_location(value: Any) -> str:
|
||||
"""Return an output-relative directory or an explicit absolute directory."""
|
||||
|
||||
raw = str(value or "").strip().replace("\\", "/")
|
||||
if not raw or raw.lower() == "output" or raw == ".":
|
||||
return ""
|
||||
drive, _ = os.path.splitdrive(raw)
|
||||
if os.path.isabs(raw):
|
||||
return os.path.normpath(raw)
|
||||
if drive:
|
||||
raise ValueError("盘符路径必须填写完整绝对路径，例如 D:/图片")
|
||||
parts = [part for part in raw.split("/") if part not in {"", "."}]
|
||||
if not parts or any(part == ".." for part in parts):
|
||||
raise ValueError("相对保存位置必须是 output 下的子目录，不能包含 ..")
|
||||
return os.path.join(*parts)
|
||||
|
||||
|
||||
def _external_preview_descriptor(
|
||||
encoded: bytes,
|
||||
filename: str,
|
||||
folder_paths_module,
|
||||
) -> dict[str, Any]:
|
||||
"""Keep external saves previewable without exposing local paths to workflows."""
|
||||
|
||||
temp_root = os.path.abspath(folder_paths_module.get_temp_directory())
|
||||
preview_subfolder = os.path.join("o1key_external_preview", uuid.uuid4().hex)
|
||||
preview_directory = os.path.join(temp_root, preview_subfolder)
|
||||
os.makedirs(preview_directory, exist_ok=False)
|
||||
_atomic_write(os.path.join(preview_directory, filename), encoded)
|
||||
return {
|
||||
"filename": filename,
|
||||
"subfolder": preview_subfolder,
|
||||
"type": "temp",
|
||||
"external_saved": True,
|
||||
}
|
||||
|
||||
|
||||
def normalize_naming_rule(value: Any) -> str:
|
||||
normalized = str(value or DEFAULT_SAVE_NAMING_RULE).strip()
|
||||
if normalized not in SAVE_NAMING_RULE_OPTIONS:
|
||||
raise ValueError("命名规则仅支持：和主图一致、自然数字、自定义前缀")
|
||||
return normalized
|
||||
|
||||
|
||||
def _safe_filename_stem(value: Any) -> str:
|
||||
filename = os.path.basename(str(value or "").strip().replace("\\", "/"))
|
||||
stem, _ = os.path.splitext(filename)
|
||||
stem = _UNSAFE_FILENAME_CHARS.sub("_", stem).strip(" .")
|
||||
stem = stem[:240] or "o1key"
|
||||
if stem.upper() in _WINDOWS_RESERVED_STEMS:
|
||||
stem = f"_{stem}"
|
||||
return stem
|
||||
|
||||
|
||||
def detect_image_format(data: bytes) -> str | None:
|
||||
if data.startswith(_PNG_SIGNATURE):
|
||||
return "PNG"
|
||||
if data.startswith(b"\xff\xd8\xff"):
|
||||
return "JPEG"
|
||||
if len(data) >= 12 and data[:4] == b"RIFF" and data[8:12] == b"WEBP":
|
||||
return "WEBP"
|
||||
return None
|
||||
|
||||
|
||||
def _metadata_enabled() -> bool:
|
||||
try:
|
||||
from comfy.cli_args import args
|
||||
|
||||
return not bool(args.disable_metadata)
|
||||
except Exception:
|
||||
return True
|
||||
|
||||
|
||||
def _metadata_items(
|
||||
prompt: Any = None,
|
||||
extra_pnginfo: dict[str, Any] | None = None,
|
||||
) -> list[tuple[str, str]]:
|
||||
if not _metadata_enabled():
|
||||
return []
|
||||
items: list[tuple[str, str]] = []
|
||||
if prompt is not None:
|
||||
items.append(("prompt", json.dumps(prompt)))
|
||||
if isinstance(extra_pnginfo, dict):
|
||||
for key, value in extra_pnginfo.items():
|
||||
items.append((str(key), json.dumps(value)))
|
||||
return items
|
||||
|
||||
|
||||
def _has_workflow_metadata(extra_pnginfo: dict[str, Any] | None) -> bool:
|
||||
"""Match ComfyUI's recoverable image contract.
|
||||
|
||||
The frontend metadata parser restores workflows from PNG and WebP, but not
|
||||
from JPEG. Native SaveImage always writes PNG, so metadata-bearing JPEG
|
||||
results must use that same container instead of an EXIF-only JPEG.
|
||||
"""
|
||||
|
||||
return (
|
||||
_metadata_enabled()
|
||||
and isinstance(extra_pnginfo, dict)
|
||||
and extra_pnginfo.get("workflow") is not None
|
||||
)
|
||||
|
||||
|
||||
def _png_text_chunk(key: str, value: str) -> bytes:
|
||||
payload = key.encode("latin-1", "replace") + b"\0" + value.encode(
|
||||
"latin-1", "replace"
|
||||
)
|
||||
chunk_type = b"tEXt"
|
||||
return (
|
||||
struct.pack(">I", len(payload))
|
||||
+ chunk_type
|
||||
+ payload
|
||||
+ struct.pack(">I", zlib.crc32(chunk_type + payload) & 0xFFFFFFFF)
|
||||
)
|
||||
|
||||
|
||||
def inject_png_metadata(data: bytes, metadata: Iterable[tuple[str, str]]) -> bytes:
|
||||
"""Append tEXt chunks without decoding or recompressing PNG pixels."""
|
||||
chunks = [_png_text_chunk(key, value) for key, value in metadata]
|
||||
if not chunks or not data.startswith(_PNG_SIGNATURE):
|
||||
return data
|
||||
|
||||
offset = len(_PNG_SIGNATURE)
|
||||
while offset + 12 <= len(data):
|
||||
length = struct.unpack(">I", data[offset : offset + 4])[0]
|
||||
end = offset + 12 + length
|
||||
if end > len(data):
|
||||
return data
|
||||
if data[offset + 4 : offset + 8] == b"IEND":
|
||||
return data[:offset] + b"".join(chunks) + data[offset:]
|
||||
offset = end
|
||||
return data
|
||||
|
||||
|
||||
def _create_exif_bytes(metadata: Iterable[tuple[str, str]]) -> bytes:
|
||||
items = list(metadata)
|
||||
if not items:
|
||||
return b""
|
||||
exif = Image.Exif()
|
||||
extra_index = 0
|
||||
for key, value in items:
|
||||
if key == "prompt":
|
||||
exif[0x0110] = f"prompt:{value}"
|
||||
else:
|
||||
exif[0x010F - extra_index] = f"{key}:{value}"
|
||||
extra_index += 1
|
||||
try:
|
||||
return exif.tobytes()
|
||||
except (OSError, OverflowError, ValueError):
|
||||
# JPEG APP1 and some Pillow EXIF writers have practical size limits.
|
||||
# The image must still save; PNG remains the fully lossless metadata path.
|
||||
return b""
|
||||
|
||||
|
||||
def inject_jpeg_metadata(data: bytes, metadata: Iterable[tuple[str, str]]) -> bytes:
|
||||
"""Insert an EXIF APP1 segment without recompressing JPEG pixels."""
|
||||
if not data.startswith(b"\xff\xd8"):
|
||||
return data
|
||||
exif = _create_exif_bytes(metadata)
|
||||
if not exif or len(exif) + 2 > 0xFFFF:
|
||||
return data
|
||||
segment = b"\xff\xe1" + struct.pack(">H", len(exif) + 2) + exif
|
||||
return data[:2] + segment + data[2:]
|
||||
|
||||
|
||||
def _webp_chunks(data: bytes) -> list[tuple[bytes, bytes]] | None:
|
||||
if len(data) < 12 or data[:4] != b"RIFF" or data[8:12] != b"WEBP":
|
||||
return None
|
||||
chunks: list[tuple[bytes, bytes]] = []
|
||||
offset = 12
|
||||
while offset + 8 <= len(data):
|
||||
kind = data[offset : offset + 4]
|
||||
length = struct.unpack("<I", data[offset + 4 : offset + 8])[0]
|
||||
start = offset + 8
|
||||
end = start + length
|
||||
if end > len(data):
|
||||
return None
|
||||
chunks.append((kind, data[start:end]))
|
||||
offset = end + (length & 1)
|
||||
return chunks
|
||||
|
||||
|
||||
def _pack_webp_chunk(kind: bytes, payload: bytes) -> bytes:
|
||||
padding = b"\0" if len(payload) & 1 else b""
|
||||
return kind + struct.pack("<I", len(payload)) + payload + padding
|
||||
|
||||
|
||||
def inject_webp_metadata(data: bytes, metadata: Iterable[tuple[str, str]]) -> bytes:
|
||||
"""Add/replace a WebP EXIF chunk without recompressing image data."""
|
||||
exif = _create_exif_bytes(metadata)
|
||||
chunks = _webp_chunks(data)
|
||||
if not exif or chunks is None:
|
||||
return data
|
||||
|
||||
chunks = [(kind, payload) for kind, payload in chunks if kind != b"EXIF"]
|
||||
vp8x_index = next((i for i, item in enumerate(chunks) if item[0] == b"VP8X"), None)
|
||||
if vp8x_index is None:
|
||||
try:
|
||||
with Image.open(BytesIO(data)) as image:
|
||||
width, height = image.size
|
||||
has_alpha = "A" in image.getbands()
|
||||
except Exception:
|
||||
return data
|
||||
flags = 0x08 | (0x10 if has_alpha else 0)
|
||||
vp8x = bytes((flags, 0, 0, 0)) + (width - 1).to_bytes(3, "little") + (
|
||||
height - 1
|
||||
).to_bytes(3, "little")
|
||||
chunks.insert(0, (b"VP8X", vp8x))
|
||||
else:
|
||||
kind, payload = chunks[vp8x_index]
|
||||
if len(payload) != 10:
|
||||
return data
|
||||
chunks[vp8x_index] = (kind, bytes((payload[0] | 0x08,)) + payload[1:])
|
||||
|
||||
chunks.append((b"EXIF", exif))
|
||||
body = b"WEBP" + b"".join(_pack_webp_chunk(kind, payload) for kind, payload in chunks)
|
||||
return b"RIFF" + struct.pack("<I", len(body)) + body
|
||||
|
||||
|
||||
def inject_workflow_metadata(
|
||||
data: bytes,
|
||||
image_format: str,
|
||||
prompt: Any = None,
|
||||
extra_pnginfo: dict[str, Any] | None = None,
|
||||
) -> bytes:
|
||||
metadata = _metadata_items(prompt, extra_pnginfo)
|
||||
if not metadata:
|
||||
return data
|
||||
if image_format == "PNG":
|
||||
return inject_png_metadata(data, metadata)
|
||||
if image_format == "JPEG":
|
||||
return inject_jpeg_metadata(data, metadata)
|
||||
if image_format == "WEBP":
|
||||
return inject_webp_metadata(data, metadata)
|
||||
return data
|
||||
|
||||
|
||||
def _atomic_write(path: str, data: bytes) -> None:
|
||||
temporary = f"{path}.{uuid.uuid4().hex}.tmp"
|
||||
reserved = False
|
||||
try:
|
||||
descriptor = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL)
|
||||
os.close(descriptor)
|
||||
reserved = True
|
||||
with open(temporary, "wb") as handle:
|
||||
handle.write(data)
|
||||
os.replace(temporary, path)
|
||||
reserved = False
|
||||
finally:
|
||||
if os.path.exists(temporary):
|
||||
try:
|
||||
os.remove(temporary)
|
||||
except OSError:
|
||||
pass
|
||||
if reserved and os.path.exists(path):
|
||||
try:
|
||||
os.remove(path)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def _tensor_uint8(image_tensor) -> np.ndarray:
|
||||
return np.clip(255.0 * image_tensor.detach().cpu().numpy(), 0, 255).astype(np.uint8)
|
||||
|
||||
|
||||
def _raw_bytes_match_tensor(data: bytes, image_tensor) -> bool:
|
||||
try:
|
||||
array = _tensor_uint8(image_tensor)
|
||||
mode = "RGBA" if array.shape[-1] == 4 else "RGB"
|
||||
with Image.open(BytesIO(data)) as source:
|
||||
source.load()
|
||||
decoded = np.asarray(source.convert(mode), dtype=np.uint8)
|
||||
return decoded.shape == array.shape and np.array_equal(decoded, array)
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def _source_for_tensor(images, index: int) -> tuple[bytes | None, str | None, str | None]:
|
||||
metadata = getattr(images, "_o1key_source_metadata", None)
|
||||
if not isinstance(metadata, list) or index >= len(metadata):
|
||||
return None, None, None
|
||||
item = metadata[index]
|
||||
if not isinstance(item, dict):
|
||||
return None, None, None
|
||||
source_filename = item.get("filename")
|
||||
if not isinstance(source_filename, str) or not source_filename.strip():
|
||||
source_filename = None
|
||||
if item.get("modified"):
|
||||
return None, None, source_filename
|
||||
|
||||
data = item.get("bytes")
|
||||
if not isinstance(data, bytes):
|
||||
source_path = item.get("path")
|
||||
if isinstance(source_path, str) and os.path.isfile(source_path):
|
||||
try:
|
||||
with open(source_path, "rb") as handle:
|
||||
data = handle.read()
|
||||
except OSError:
|
||||
data = None
|
||||
if not isinstance(data, bytes) or not _raw_bytes_match_tensor(data, images[index]):
|
||||
data = None
|
||||
|
||||
image_format = detect_image_format(data) if data else None
|
||||
if image_format is None:
|
||||
hinted = str(item.get("format") or "").upper()
|
||||
if hinted == "JPG":
|
||||
hinted = "JPEG"
|
||||
image_format = hinted if hinted in _FORMAT_EXTENSIONS else None
|
||||
return data, image_format, source_filename
|
||||
|
||||
|
||||
def _pil_from_tensor(image_tensor) -> Image.Image:
|
||||
return Image.fromarray(_tensor_uint8(image_tensor))
|
||||
|
||||
|
||||
def _encode_image(
|
||||
image: Image.Image,
|
||||
image_format: str,
|
||||
prompt: Any = None,
|
||||
extra_pnginfo: dict[str, Any] | None = None,
|
||||
) -> bytes:
|
||||
output = BytesIO()
|
||||
metadata = _metadata_items(prompt, extra_pnginfo)
|
||||
working = image
|
||||
converted = None
|
||||
try:
|
||||
if image_format == "PNG":
|
||||
pnginfo = None
|
||||
if metadata:
|
||||
pnginfo = PngInfo()
|
||||
for key, value in metadata:
|
||||
pnginfo.add_text(key, value)
|
||||
working.save(output, format="PNG", pnginfo=pnginfo, compress_level=4)
|
||||
elif image_format == "JPEG":
|
||||
if working.mode != "RGB":
|
||||
converted = working.convert("RGB")
|
||||
working = converted
|
||||
kwargs: dict[str, Any] = {
|
||||
"format": "JPEG",
|
||||
"quality": 100,
|
||||
"subsampling": 0,
|
||||
"optimize": True,
|
||||
}
|
||||
exif = _create_exif_bytes(metadata)
|
||||
if exif:
|
||||
kwargs["exif"] = exif
|
||||
working.save(output, **kwargs)
|
||||
elif image_format == "WEBP":
|
||||
kwargs = {"format": "WEBP", "lossless": True, "quality": 100, "method": 4}
|
||||
exif = _create_exif_bytes(metadata)
|
||||
if exif:
|
||||
kwargs["exif"] = exif
|
||||
working.save(output, **kwargs)
|
||||
else:
|
||||
raise ValueError(f"不支持的图像格式：{image_format}")
|
||||
return output.getvalue()
|
||||
finally:
|
||||
if converted is not None:
|
||||
converted.close()
|
||||
|
||||
|
||||
def _save_sources(
|
||||
sources: list[tuple[bytes | None, str | None, Image.Image, str | None]],
|
||||
filename_prefix: str,
|
||||
save_format: str,
|
||||
output_directory: str,
|
||||
folder_paths_module,
|
||||
prompt: Any = None,
|
||||
extra_pnginfo: dict[str, Any] | None = None,
|
||||
save_location: str = "",
|
||||
naming_rule: str = DEFAULT_SAVE_NAMING_RULE,
|
||||
main_filename: str | None = None,
|
||||
) -> list[dict[str, Any]]:
|
||||
if not sources:
|
||||
return []
|
||||
width, height = sources[0][2].size
|
||||
normalized_location = normalize_save_location(save_location)
|
||||
external_location = bool(normalized_location and os.path.isabs(normalized_location))
|
||||
relative_location = "" if external_location else normalized_location
|
||||
save_root = normalized_location if external_location else output_directory
|
||||
selected_rule = normalize_naming_rule(naming_rule)
|
||||
filename = ""
|
||||
counter = 1
|
||||
if selected_rule == "自定义前缀":
|
||||
effective_prefix = filename_prefix or "o1key"
|
||||
if relative_location:
|
||||
effective_prefix = os.path.join(relative_location, effective_prefix)
|
||||
full_output_folder, filename, counter, subfolder, _ = (
|
||||
folder_paths_module.get_save_image_path(
|
||||
effective_prefix,
|
||||
save_root,
|
||||
width,
|
||||
height,
|
||||
)
|
||||
)
|
||||
else:
|
||||
full_output_folder = os.path.join(save_root, relative_location)
|
||||
subfolder = relative_location
|
||||
os.makedirs(full_output_folder, exist_ok=True)
|
||||
selected = normalize_save_format(save_format)
|
||||
main_stem = _safe_filename_stem(
|
||||
main_filename or next((item[3] for item in sources if item[3]), None)
|
||||
)
|
||||
results: list[dict[str, Any]] = []
|
||||
natural_number = 1
|
||||
main_number = 0
|
||||
for batch_number, (raw, source_format, image, _source_filename) in enumerate(sources):
|
||||
target_format = source_format if selected == "原始" else {
|
||||
"png": "PNG",
|
||||
"jpg": "JPEG",
|
||||
"webp": "WEBP",
|
||||
}[selected]
|
||||
if target_format not in _FORMAT_EXTENSIONS:
|
||||
target_format = "PNG"
|
||||
if target_format == "JPEG" and _has_workflow_metadata(extra_pnginfo):
|
||||
target_format = "PNG"
|
||||
extension = _FORMAT_EXTENSIONS[target_format]
|
||||
if selected == "原始" and raw is not None and detect_image_format(raw) == target_format:
|
||||
encoded = inject_workflow_metadata(
|
||||
raw,
|
||||
target_format,
|
||||
prompt=prompt,
|
||||
extra_pnginfo=extra_pnginfo,
|
||||
)
|
||||
else:
|
||||
encoded = _encode_image(
|
||||
image,
|
||||
target_format,
|
||||
prompt=prompt,
|
||||
extra_pnginfo=extra_pnginfo,
|
||||
)
|
||||
|
||||
with _SAVE_LOCK:
|
||||
while True:
|
||||
if selected_rule == "自然数字":
|
||||
saved_name = f"{natural_number}.{extension}"
|
||||
natural_number += 1
|
||||
elif selected_rule == "和主图一致":
|
||||
suffix = "" if main_number == 0 else str(main_number)
|
||||
saved_name = f"{main_stem}{suffix}.{extension}"
|
||||
main_number += 1
|
||||
else:
|
||||
filename_with_batch_num = filename.replace(
|
||||
"%batch_num%", str(batch_number)
|
||||
)
|
||||
saved_name = f"{filename_with_batch_num}_{counter:05}_.{extension}"
|
||||
counter += 1
|
||||
output_path = os.path.join(full_output_folder, saved_name)
|
||||
if os.path.exists(output_path):
|
||||
continue
|
||||
try:
|
||||
_atomic_write(output_path, encoded)
|
||||
break
|
||||
except FileExistsError:
|
||||
continue
|
||||
if external_location:
|
||||
results.append(
|
||||
_external_preview_descriptor(encoded, saved_name, folder_paths_module)
|
||||
)
|
||||
else:
|
||||
results.append({"filename": saved_name, "subfolder": subfolder, "type": "output"})
|
||||
return results
|
||||
|
||||
|
||||
def save_tensor_images(
|
||||
images,
|
||||
filename_prefix: str,
|
||||
save_format: str,
|
||||
output_directory: str,
|
||||
folder_paths_module,
|
||||
prompt: Any = None,
|
||||
extra_pnginfo: dict[str, Any] | None = None,
|
||||
save_location: str = "",
|
||||
naming_rule: str = DEFAULT_SAVE_NAMING_RULE,
|
||||
main_filename: str | None = None,
|
||||
) -> list[dict[str, Any]]:
|
||||
sources: list[tuple[bytes | None, str | None, Image.Image, str | None]] = []
|
||||
try:
|
||||
for index, image_tensor in enumerate(images):
|
||||
raw, source_format, source_filename = _source_for_tensor(images, index)
|
||||
sources.append((raw, source_format, _pil_from_tensor(image_tensor), source_filename))
|
||||
return _save_sources(
|
||||
sources,
|
||||
filename_prefix,
|
||||
save_format,
|
||||
output_directory,
|
||||
folder_paths_module,
|
||||
prompt,
|
||||
extra_pnginfo,
|
||||
save_location,
|
||||
naming_rule,
|
||||
main_filename or getattr(images, "_o1key_main_filename", None),
|
||||
)
|
||||
finally:
|
||||
for _, _, image, _ in sources:
|
||||
image.close()
|
||||
|
||||
|
||||
def save_temp_images(
|
||||
source_paths: list[str],
|
||||
filename_prefix: str,
|
||||
save_format: str,
|
||||
output_directory: str,
|
||||
folder_paths_module,
|
||||
prompt: Any = None,
|
||||
extra_pnginfo: dict[str, Any] | None = None,
|
||||
save_location: str = "",
|
||||
naming_rule: str = DEFAULT_SAVE_NAMING_RULE,
|
||||
main_filename: str | None = None,
|
||||
) -> list[dict[str, Any]]:
|
||||
sources: list[tuple[bytes | None, str | None, Image.Image, str | None]] = []
|
||||
try:
|
||||
for source_path in source_paths:
|
||||
with open(source_path, "rb") as handle:
|
||||
raw = handle.read()
|
||||
image_format = detect_image_format(raw)
|
||||
if image_format not in _FORMAT_EXTENSIONS:
|
||||
raise ValueError(f"临时结果不是受支持的图像：{os.path.basename(source_path)}")
|
||||
with Image.open(BytesIO(raw)) as opened:
|
||||
opened.load()
|
||||
image = opened.copy()
|
||||
sources.append((raw, image_format, image, None))
|
||||
return _save_sources(
|
||||
sources,
|
||||
filename_prefix,
|
||||
save_format,
|
||||
output_directory,
|
||||
folder_paths_module,
|
||||
prompt,
|
||||
extra_pnginfo,
|
||||
save_location,
|
||||
naming_rule,
|
||||
main_filename,
|
||||
)
|
||||
finally:
|
||||
for _, _, image, _ in sources:
|
||||
image.close()
|
||||
|
||||
|
||||
__all__ = [
|
||||
"SAVE_FORMAT_OPTIONS",
|
||||
"SAVE_NAMING_RULE_OPTIONS",
|
||||
"detect_image_format",
|
||||
"inject_workflow_metadata",
|
||||
"normalize_save_format",
|
||||
"normalize_save_location",
|
||||
"normalize_naming_rule",
|
||||
"save_temp_images",
|
||||
"save_tensor_images",
|
||||
]
|
||||
@@ -0,0 +1,140 @@
|
||||
"""Bounded reference thumbnails for the unified image-generator panel."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from io import BytesIO
|
||||
import os
|
||||
from typing import Any
|
||||
|
||||
from PIL import Image, ImageOps
|
||||
|
||||
|
||||
THUMBNAIL_MAX_SIZE = 256
|
||||
THUMBNAIL_QUALITY = 82
|
||||
THUMBNAIL_CACHE_SECONDS = 3600
|
||||
_ALLOWED_FOLDER_TYPES = frozenset({"input", "output", "temp"})
|
||||
|
||||
|
||||
def _is_within(root: str, candidate: str) -> bool:
|
||||
try:
|
||||
return os.path.commonpath(
|
||||
[os.path.normcase(root), os.path.normcase(candidate)]
|
||||
) == os.path.normcase(root)
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
|
||||
def resolve_thumbnail_source(
|
||||
folder_paths_module: Any,
|
||||
filename: str,
|
||||
subfolder: str = "",
|
||||
folder_type: str = "input",
|
||||
) -> str:
|
||||
"""Resolve one public ComfyUI image descriptor without allowing traversal."""
|
||||
|
||||
folder_type = str(folder_type or "input").strip()
|
||||
if folder_type not in _ALLOWED_FOLDER_TYPES:
|
||||
raise ValueError("图片目录类型无效")
|
||||
|
||||
filename = str(filename or "").strip()
|
||||
subfolder = str(subfolder or "").strip()
|
||||
if not filename or os.path.basename(filename) != filename:
|
||||
raise ValueError("图片文件名无效")
|
||||
if os.path.isabs(subfolder) or ".." in subfolder.replace("\\", "/").split("/"):
|
||||
raise ValueError("图片子目录无效")
|
||||
|
||||
root_value = folder_paths_module.get_directory_by_type(folder_type)
|
||||
if not root_value:
|
||||
raise ValueError("图片目录不存在")
|
||||
root = os.path.realpath(os.path.abspath(root_value))
|
||||
source = os.path.realpath(os.path.abspath(os.path.join(root, subfolder, filename)))
|
||||
if not _is_within(root, source):
|
||||
raise ValueError("图片路径不安全")
|
||||
if not os.path.isfile(source):
|
||||
raise FileNotFoundError(filename)
|
||||
return source
|
||||
|
||||
|
||||
def render_reference_thumbnail(
|
||||
source: str,
|
||||
max_size: int = THUMBNAIL_MAX_SIZE,
|
||||
quality: int = THUMBNAIL_QUALITY,
|
||||
) -> bytes:
|
||||
"""Render a small WebP without modifying or replacing the source image."""
|
||||
|
||||
max_size = max(32, min(int(max_size), THUMBNAIL_MAX_SIZE))
|
||||
quality = max(1, min(int(quality), 100))
|
||||
with Image.open(source) as opened:
|
||||
if getattr(opened, "is_animated", False):
|
||||
opened.seek(0)
|
||||
# JPEG decoders can use draft mode to avoid materialising every source
|
||||
# pixel before the final thumbnail resize. Other formats still remain
|
||||
# bounded by the route-level semaphore.
|
||||
opened.draft("RGB", (max_size, max_size))
|
||||
image = None
|
||||
try:
|
||||
image = ImageOps.exif_transpose(opened)
|
||||
image.thumbnail(
|
||||
(max_size, max_size),
|
||||
Image.Resampling.LANCZOS,
|
||||
reducing_gap=3.0,
|
||||
)
|
||||
has_alpha = "A" in image.getbands() or "transparency" in image.info
|
||||
converted = image.convert("RGBA" if has_alpha else "RGB")
|
||||
try:
|
||||
output = BytesIO()
|
||||
converted.save(output, format="WEBP", quality=quality, method=4)
|
||||
return output.getvalue()
|
||||
finally:
|
||||
converted.close()
|
||||
finally:
|
||||
if image is not None and image is not opened:
|
||||
image.close()
|
||||
|
||||
|
||||
def _thumbnail_etag(source: str) -> str:
|
||||
stat = os.stat(source)
|
||||
return f'"o1key-thumb-{stat.st_mtime_ns:x}-{stat.st_size:x}"'
|
||||
|
||||
|
||||
def register_o1key_image_thumbnail_route(PromptServer, web, folder_paths_module):
|
||||
"""Register the thumbnail endpoint used only by lightweight panel previews."""
|
||||
|
||||
render_semaphore = asyncio.Semaphore(2)
|
||||
|
||||
@PromptServer.instance.routes.get("/o1key/image/thumbnail")
|
||||
async def get_o1key_image_thumbnail(request):
|
||||
try:
|
||||
source = resolve_thumbnail_source(
|
||||
folder_paths_module,
|
||||
request.query.get("filename", ""),
|
||||
request.query.get("subfolder", ""),
|
||||
request.query.get("type", "input"),
|
||||
)
|
||||
etag = await asyncio.to_thread(_thumbnail_etag, source)
|
||||
headers = {
|
||||
"Cache-Control": f"private, max-age={THUMBNAIL_CACHE_SECONDS}",
|
||||
"ETag": etag,
|
||||
"X-Content-Type-Options": "nosniff",
|
||||
}
|
||||
if request.headers.get("If-None-Match") == etag:
|
||||
return web.Response(status=304, headers=headers)
|
||||
async with render_semaphore:
|
||||
body = await asyncio.to_thread(render_reference_thumbnail, source)
|
||||
return web.Response(body=body, content_type="image/webp", headers=headers)
|
||||
except FileNotFoundError:
|
||||
return web.Response(status=404)
|
||||
except ValueError:
|
||||
return web.Response(status=400)
|
||||
except Exception:
|
||||
# Do not expose local paths or decoder details to the browser.
|
||||
return web.Response(status=415)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"THUMBNAIL_MAX_SIZE",
|
||||
"register_o1key_image_thumbnail_route",
|
||||
"render_reference_thumbnail",
|
||||
"resolve_thumbnail_source",
|
||||
]
|
||||
@@ -0,0 +1,308 @@
|
||||
"""Canonical capability catalog for the panel-style o1key video generator."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
|
||||
VIDEO_PROVIDER_OPTIONS = ["seedance"]
|
||||
SEEDANCE_MODEL_OPTIONS = [
|
||||
"seedance-2.0",
|
||||
"seedance-2.0-fast",
|
||||
"seedance-2.0-mini",
|
||||
"seedance-2.5",
|
||||
]
|
||||
SEEDANCE_ROUTE_OPTIONS = ["domestic", "overseas_hc"]
|
||||
SEEDANCE_ASSET_CREATION_MODE_OPTIONS = ["auto", "manual"]
|
||||
VIDEO_GENERATION_MODE_OPTIONS = [
|
||||
"text",
|
||||
"first_frame",
|
||||
"first_last_frame",
|
||||
"multimodal",
|
||||
]
|
||||
VIDEO_RESOLUTION_OPTIONS = ["480p", "720p", "1080p", "4k"]
|
||||
VIDEO_ASPECT_RATIO_OPTIONS = ["auto", "16:9", "9:16", "4:3", "3:4", "1:1", "21:9"]
|
||||
VIDEO_DURATION_OPTIONS = ["auto", *[str(value) for value in range(4, 31)]]
|
||||
|
||||
MAX_VIDEO_PROMPT_CHARS = 20_000
|
||||
SEEDANCE_REFERENCE_IMAGE_MAX_BYTES = 30 * 1024 * 1024
|
||||
SEEDANCE_REFERENCE_VIDEO_MAX_BYTES = 512 * 1024 * 1024
|
||||
SEEDANCE_REFERENCE_AUDIO_MAX_BYTES = 100 * 1024 * 1024
|
||||
SEEDANCE_REFERENCE_MIN_DIMENSION = 300
|
||||
SEEDANCE_REFERENCE_MAX_DIMENSION = 6000
|
||||
SEEDANCE_REFERENCE_MIN_ASPECT_RATIO = 0.4
|
||||
SEEDANCE_REFERENCE_MAX_ASPECT_RATIO = 2.5
|
||||
SEEDANCE_REFERENCE_VIDEO_MIN_PIXELS = 407_696
|
||||
SEEDANCE_REFERENCE_VIDEO_MAX_PIXELS = 8_295_044
|
||||
|
||||
SEEDANCE_MODEL_MATRIX = {
|
||||
("seedance-2.0", "domestic"): "doubao-seedance-2-0-260128-max",
|
||||
("seedance-2.0-fast", "domestic"): "doubao-seedance-2-0-fast-260128-max",
|
||||
("seedance-2.0-mini", "domestic"): "doubao-seedance-2-0-mini-260615-max",
|
||||
("seedance-2.5", "domestic"): "doubao-seedance-2-5-260628-max",
|
||||
("seedance-2.0", "overseas_hc"): "dreamina-seedance-2-0-hc",
|
||||
("seedance-2.0-fast", "overseas_hc"): "dreamina-seedance-2-0-fast-hc",
|
||||
("seedance-2.0-mini", "overseas_hc"): "dreamina-seedance-2-0-mini-hc",
|
||||
("seedance-2.5", "overseas_hc"): "dreamina-seedance-2-5-hc",
|
||||
}
|
||||
|
||||
SEEDANCE_CAPABILITIES = {
|
||||
"seedance-2.0": {
|
||||
"duration_min": 4,
|
||||
"duration_max": 15,
|
||||
"images": 9,
|
||||
"videos": 3,
|
||||
"audios": 3,
|
||||
"resolutions": ["480p", "720p", "1080p", "4k"],
|
||||
},
|
||||
"seedance-2.0-fast": {
|
||||
"duration_min": 4,
|
||||
"duration_max": 15,
|
||||
"images": 9,
|
||||
"videos": 3,
|
||||
"audios": 3,
|
||||
"resolutions": ["480p", "720p"],
|
||||
},
|
||||
"seedance-2.0-mini": {
|
||||
"duration_min": 4,
|
||||
"duration_max": 15,
|
||||
"images": 9,
|
||||
"videos": 3,
|
||||
"audios": 3,
|
||||
"resolutions": ["480p", "720p"],
|
||||
},
|
||||
"seedance-2.5": {
|
||||
"duration_min": 4,
|
||||
"duration_max": 30,
|
||||
"images": 30,
|
||||
"videos": 10,
|
||||
"audios": 10,
|
||||
"resolutions": ["480p", "720p", "1080p", "4k"],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _as_bool(value: Any) -> bool:
|
||||
if isinstance(value, bool):
|
||||
return value
|
||||
return str(value or "").strip().lower() in {"1", "true", "yes", "on", "打开", "开启"}
|
||||
|
||||
|
||||
def resolve_seedance_model(model: str, route: str) -> str:
|
||||
try:
|
||||
return SEEDANCE_MODEL_MATRIX[(model, route)]
|
||||
except KeyError:
|
||||
raise ValueError(f"Seedance 模型与线路组合无效：{model} / {route}") from None
|
||||
|
||||
|
||||
def normalize_seedance_parameters(payload: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Validate scalar Seedance controls before media upload or paid requests."""
|
||||
|
||||
provider = str(payload.get("provider") or "seedance").strip()
|
||||
if provider != "seedance":
|
||||
raise ValueError(f"暂不支持的视频提供方：{provider}")
|
||||
|
||||
model = str(payload.get("model") or "seedance-2.0").strip()
|
||||
if model not in SEEDANCE_MODEL_OPTIONS:
|
||||
raise ValueError(f"不支持的 Seedance 模型：{model}")
|
||||
route = str(payload.get("route") or "domestic").strip()
|
||||
if route not in SEEDANCE_ROUTE_OPTIONS:
|
||||
raise ValueError(f"不支持的 Seedance 线路：{route}")
|
||||
mode = str(payload.get("generation_mode") or "text").strip()
|
||||
if mode not in VIDEO_GENERATION_MODE_OPTIONS:
|
||||
raise ValueError(f"不支持的视频生成模式：{mode}")
|
||||
asset_creation_mode = str(payload.get("asset_creation_mode") or "auto").strip()
|
||||
if asset_creation_mode not in SEEDANCE_ASSET_CREATION_MODE_OPTIONS:
|
||||
raise ValueError(f"不支持的素材创建模式：{asset_creation_mode}")
|
||||
|
||||
prompt = str(payload.get("prompt") or "").strip()
|
||||
if len(prompt) > MAX_VIDEO_PROMPT_CHARS:
|
||||
raise ValueError(f"提示词最多支持 {MAX_VIDEO_PROMPT_CHARS} 个字符")
|
||||
if mode == "text" and not prompt:
|
||||
raise ValueError("文生视频模式下提示词不能为空")
|
||||
|
||||
resolution = str(payload.get("resolution") or "720p").strip().lower()
|
||||
capabilities = SEEDANCE_CAPABILITIES[model]
|
||||
if resolution not in capabilities["resolutions"]:
|
||||
supported = " / ".join(capabilities["resolutions"])
|
||||
raise ValueError(f"{model} 分辨率仅支持 {supported}")
|
||||
|
||||
aspect_ratio = str(payload.get("aspect_ratio") or "auto").strip()
|
||||
if aspect_ratio not in VIDEO_ASPECT_RATIO_OPTIONS:
|
||||
raise ValueError(f"不支持的宽高比：{aspect_ratio}")
|
||||
|
||||
raw_duration = str(payload.get("duration") or "5").strip().lower()
|
||||
if raw_duration == "auto":
|
||||
duration = 5
|
||||
else:
|
||||
try:
|
||||
duration = int(raw_duration.removesuffix("秒"))
|
||||
except ValueError:
|
||||
raise ValueError(f"无效的视频时长：{raw_duration}") from None
|
||||
if not capabilities["duration_min"] <= duration <= capabilities["duration_max"]:
|
||||
raise ValueError(
|
||||
f"{model} 的时长仅支持 {capabilities['duration_min']}-{capabilities['duration_max']} 秒"
|
||||
)
|
||||
|
||||
try:
|
||||
seed = int(payload.get("seed") or 0)
|
||||
except (TypeError, ValueError):
|
||||
raise ValueError("seed 必须是非负整数") from None
|
||||
if not 0 <= seed <= 0xFFFFFFFFFFFFFFFF:
|
||||
raise ValueError("seed 必须是 0 到 18446744073709551615 之间的整数")
|
||||
|
||||
filename_prefix = str(payload.get("filename_prefix") or "o1key_video").strip()
|
||||
if not filename_prefix:
|
||||
filename_prefix = "o1key_video"
|
||||
if len(filename_prefix) > 240:
|
||||
raise ValueError("文件名前缀最多支持 240 个字符")
|
||||
|
||||
return {
|
||||
"provider": provider,
|
||||
"model": model,
|
||||
"actual_model": resolve_seedance_model(model, route),
|
||||
"route": route,
|
||||
"generation_mode": mode,
|
||||
"asset_creation_mode": asset_creation_mode,
|
||||
"prompt": prompt,
|
||||
"resolution": resolution,
|
||||
"aspect_ratio": aspect_ratio,
|
||||
"duration": duration,
|
||||
"generate_audio": _as_bool(payload.get("generate_audio")),
|
||||
"return_last_frame": _as_bool(payload.get("return_last_frame")),
|
||||
"seed": seed,
|
||||
"filename_prefix": filename_prefix,
|
||||
"save_location": str(payload.get("save_location") or "").strip(),
|
||||
"capabilities": capabilities,
|
||||
}
|
||||
|
||||
|
||||
def validate_seedance_media_counts(
|
||||
model: str,
|
||||
image_count: int,
|
||||
video_count: int,
|
||||
audio_count: int,
|
||||
*,
|
||||
model_label: str | None = None,
|
||||
) -> None:
|
||||
"""Apply the shared per-type Seedance reference-item limits."""
|
||||
|
||||
if model not in SEEDANCE_CAPABILITIES:
|
||||
raise ValueError(f"不支持的 Seedance 模型：{model}")
|
||||
capabilities = SEEDANCE_CAPABILITIES[model]
|
||||
label = model_label or model
|
||||
for media_label, count, limit in (
|
||||
("参考图片", int(image_count), capabilities["images"]),
|
||||
("参考视频", int(video_count), capabilities["videos"]),
|
||||
("参考音频", int(audio_count), capabilities["audios"]),
|
||||
):
|
||||
if count > limit:
|
||||
raise ValueError(f"{label} 最多支持 {limit} 个{media_label}，当前输入 {count} 个")
|
||||
|
||||
|
||||
def validate_seedance_reference_dimensions(
|
||||
width: int,
|
||||
height: int,
|
||||
label: str,
|
||||
*,
|
||||
require_video_pixel_range: bool = False,
|
||||
) -> None:
|
||||
"""Apply the shared Seedance reference image/video dimension limits."""
|
||||
|
||||
width = int(width or 0)
|
||||
height = int(height or 0)
|
||||
if not (
|
||||
SEEDANCE_REFERENCE_MIN_DIMENSION <= width <= SEEDANCE_REFERENCE_MAX_DIMENSION
|
||||
and SEEDANCE_REFERENCE_MIN_DIMENSION <= height <= SEEDANCE_REFERENCE_MAX_DIMENSION
|
||||
):
|
||||
raise ValueError(
|
||||
f"{label}宽高必须分别在 {SEEDANCE_REFERENCE_MIN_DIMENSION}～"
|
||||
f"{SEEDANCE_REFERENCE_MAX_DIMENSION}px，当前为 {width}×{height}"
|
||||
)
|
||||
|
||||
aspect_ratio = width / height
|
||||
if not (
|
||||
SEEDANCE_REFERENCE_MIN_ASPECT_RATIO
|
||||
<= aspect_ratio
|
||||
<= SEEDANCE_REFERENCE_MAX_ASPECT_RATIO
|
||||
):
|
||||
raise ValueError(
|
||||
f"{label}宽高比必须在 {SEEDANCE_REFERENCE_MIN_ASPECT_RATIO}～"
|
||||
f"{SEEDANCE_REFERENCE_MAX_ASPECT_RATIO}，当前为 {aspect_ratio:.3f}（{width}:{height}）"
|
||||
)
|
||||
|
||||
if require_video_pixel_range:
|
||||
pixel_count = width * height
|
||||
if not (
|
||||
SEEDANCE_REFERENCE_VIDEO_MIN_PIXELS
|
||||
<= pixel_count
|
||||
<= SEEDANCE_REFERENCE_VIDEO_MAX_PIXELS
|
||||
):
|
||||
raise ValueError(
|
||||
f"{label}总像素必须在 {SEEDANCE_REFERENCE_VIDEO_MIN_PIXELS:,}～"
|
||||
f"{SEEDANCE_REFERENCE_VIDEO_MAX_PIXELS:,} 之间，当前为 "
|
||||
f"{width}×{height}（{pixel_count:,} 像素）"
|
||||
)
|
||||
|
||||
|
||||
def public_video_capabilities() -> dict[str, Any]:
|
||||
return {
|
||||
"providers": [
|
||||
{
|
||||
"value": "seedance",
|
||||
"label": "Seedance",
|
||||
"models": [
|
||||
{
|
||||
"value": model,
|
||||
"label": model.replace("seedance", "Seedance", 1),
|
||||
**SEEDANCE_CAPABILITIES[model],
|
||||
}
|
||||
for model in SEEDANCE_MODEL_OPTIONS
|
||||
],
|
||||
"routes": [
|
||||
{"value": "domestic", "label": "国内"},
|
||||
{"value": "overseas_hc", "label": "海外"},
|
||||
],
|
||||
"modes": [
|
||||
{"value": "text", "label": "文生视频"},
|
||||
{"value": "first_frame", "label": "首帧图生视频"},
|
||||
{"value": "first_last_frame", "label": "首尾帧生视频"},
|
||||
{"value": "multimodal", "label": "多模态参考生视频"},
|
||||
],
|
||||
"asset_creation_modes": [
|
||||
{"value": "auto", "label": "自动创建"},
|
||||
{"value": "manual", "label": "手动"},
|
||||
],
|
||||
"aspect_ratios": VIDEO_ASPECT_RATIO_OPTIONS,
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
__all__ = [
|
||||
"MAX_VIDEO_PROMPT_CHARS",
|
||||
"SEEDANCE_CAPABILITIES",
|
||||
"SEEDANCE_ASSET_CREATION_MODE_OPTIONS",
|
||||
"SEEDANCE_MODEL_MATRIX",
|
||||
"SEEDANCE_MODEL_OPTIONS",
|
||||
"SEEDANCE_REFERENCE_AUDIO_MAX_BYTES",
|
||||
"SEEDANCE_REFERENCE_IMAGE_MAX_BYTES",
|
||||
"SEEDANCE_REFERENCE_MAX_ASPECT_RATIO",
|
||||
"SEEDANCE_REFERENCE_MAX_DIMENSION",
|
||||
"SEEDANCE_REFERENCE_MIN_ASPECT_RATIO",
|
||||
"SEEDANCE_REFERENCE_MIN_DIMENSION",
|
||||
"SEEDANCE_REFERENCE_VIDEO_MAX_BYTES",
|
||||
"SEEDANCE_REFERENCE_VIDEO_MAX_PIXELS",
|
||||
"SEEDANCE_REFERENCE_VIDEO_MIN_PIXELS",
|
||||
"SEEDANCE_ROUTE_OPTIONS",
|
||||
"VIDEO_ASPECT_RATIO_OPTIONS",
|
||||
"VIDEO_DURATION_OPTIONS",
|
||||
"VIDEO_GENERATION_MODE_OPTIONS",
|
||||
"VIDEO_PROVIDER_OPTIONS",
|
||||
"VIDEO_RESOLUTION_OPTIONS",
|
||||
"normalize_seedance_parameters",
|
||||
"public_video_capabilities",
|
||||
"resolve_seedance_model",
|
||||
"validate_seedance_media_counts",
|
||||
"validate_seedance_reference_dimensions",
|
||||
]
|
||||
File diff suppressed because it is too large
Load Diff
+172
-58
@@ -1,77 +1,175 @@
|
||||
"""
|
||||
R2 文件上传工具（通过 o1key 后端预签名接口）
|
||||
- 插件内零 R2 凭证，仅使用用户的 O1KEY_API_KEY
|
||||
- 流程：请求预签名 URL → PUT 直传 R2 → 返回公网 URL
|
||||
"""o1key 临时素材上传工具。
|
||||
|
||||
只供后续生成请求本来就使用公网 URL 的节点调用。图片、视频和音频会通过
|
||||
``POST /v1/o1key/uploads`` 以 multipart/form-data 上传，并返回临时 HTTPS URL。
|
||||
各节点已有的数量、尺寸、时长和体积校验仍由节点自身负责；本模块不新增统一
|
||||
文件大小限制，也不替调用方缩放或转码素材。
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import uuid
|
||||
from typing import Optional
|
||||
|
||||
import aiohttp
|
||||
|
||||
from .config import get_api_key_or_raise, get_api_base_url
|
||||
from .video_task import check_interrupt, run_with_interrupt
|
||||
|
||||
|
||||
async def _presign(filename: str, content_type: str) -> tuple:
|
||||
"""向 o1key 后端请求预签名 URL，返回 (upload_url, public_url)"""
|
||||
_UPLOAD_ENDPOINT = "/v1/o1key/uploads"
|
||||
_UPLOAD_RETRY_DELAYS = (2.0, 4.0, 8.0)
|
||||
_UPLOAD_TIMEOUT = aiohttp.ClientTimeout(total=180)
|
||||
|
||||
|
||||
def _redact_upload_text(value: object) -> str:
|
||||
text = str(value or "")
|
||||
text = re.sub(r"data:[^\s,;]+;base64,[A-Za-z0-9+/=_-]+", "<base64 omitted>", text)
|
||||
return re.sub(r"https?://[^\s\"'<>]+", "<temporary URL omitted>", text)[:500]
|
||||
|
||||
|
||||
def _upload_error_detail(text: str) -> str:
|
||||
try:
|
||||
payload = json.loads(text)
|
||||
except Exception:
|
||||
return _redact_upload_text(text)
|
||||
if isinstance(payload, dict):
|
||||
error = payload.get("error") or payload.get("message") or payload
|
||||
if isinstance(error, dict):
|
||||
error = error.get("message") or error.get("detail") or error
|
||||
return _redact_upload_text(error)
|
||||
return _redact_upload_text(payload)
|
||||
|
||||
|
||||
def _retry_after_seconds(value: Optional[str], fallback: float) -> float:
|
||||
try:
|
||||
return max(0.0, min(float(value), 120.0))
|
||||
except (TypeError, ValueError):
|
||||
return fallback
|
||||
|
||||
|
||||
async def _upload_file(
|
||||
data: bytes,
|
||||
filename: str,
|
||||
content_type: str,
|
||||
base_url: Optional[str] = None,
|
||||
) -> str:
|
||||
"""上传一个附件并返回新版接口提供的临时 HTTPS URL。"""
|
||||
check_interrupt()
|
||||
api_key = get_api_key_or_raise()
|
||||
base_url = get_api_base_url()
|
||||
api_base_url = (base_url or get_api_base_url()).rstrip("/")
|
||||
upload_url = f"{api_base_url}{_UPLOAD_ENDPOINT}"
|
||||
last_error: Optional[BaseException] = None
|
||||
|
||||
connector = aiohttp.TCPConnector(ssl=False)
|
||||
async with aiohttp.ClientSession(connector=connector) as session:
|
||||
async with session.post(
|
||||
f"{base_url}/v1/storage/presign",
|
||||
headers={"Authorization": f"Bearer {api_key}"},
|
||||
json={"filename": filename, "content_type": content_type},
|
||||
timeout=aiohttp.ClientTimeout(total=10),
|
||||
) as resp:
|
||||
if resp.status != 200:
|
||||
text = await resp.text()
|
||||
raise RuntimeError(f"预签名请求失败 ({resp.status}): {text}")
|
||||
data = await resp.json()
|
||||
for attempt in range(len(_UPLOAD_RETRY_DELAYS) + 1):
|
||||
check_interrupt()
|
||||
form = aiohttp.FormData()
|
||||
form.add_field(
|
||||
"file",
|
||||
data,
|
||||
filename=filename,
|
||||
content_type=content_type,
|
||||
)
|
||||
|
||||
return data["upload_url"], data["public_url"]
|
||||
connector = aiohttp.TCPConnector(ssl=False)
|
||||
try:
|
||||
async with aiohttp.ClientSession(connector=connector) as session:
|
||||
async def _request():
|
||||
async with session.post(
|
||||
upload_url,
|
||||
headers={"Authorization": f"Bearer {api_key}"},
|
||||
data=form,
|
||||
timeout=_UPLOAD_TIMEOUT,
|
||||
) as response:
|
||||
return response.status, await response.text(), dict(response.headers)
|
||||
|
||||
status, text, response_headers = await run_with_interrupt(_request())
|
||||
|
||||
if status in (200, 201):
|
||||
try:
|
||||
payload = json.loads(text)
|
||||
except Exception:
|
||||
raise RuntimeError("临时素材上传响应不是有效 JSON。") from None
|
||||
public_url = str(payload.get("url") or "").strip()
|
||||
if not public_url.startswith("https://"):
|
||||
raise RuntimeError("临时素材上传响应缺少有效 HTTPS URL。")
|
||||
check_interrupt()
|
||||
print(
|
||||
f"[素材上传] {filename} 上传完成：{len(data) / 1024:.1f} KB"
|
||||
)
|
||||
return public_url
|
||||
|
||||
detail = _upload_error_detail(text)
|
||||
retryable = status == 429 or 500 <= status < 600
|
||||
if not retryable or attempt >= len(_UPLOAD_RETRY_DELAYS):
|
||||
retry_after = response_headers.get("Retry-After")
|
||||
suffix = f"，Retry-After={retry_after}s" if retry_after else ""
|
||||
raise RuntimeError(
|
||||
f"临时素材上传失败 HTTP {status}{suffix}: {detail}"
|
||||
)
|
||||
|
||||
delay = _retry_after_seconds(
|
||||
response_headers.get("Retry-After"),
|
||||
_UPLOAD_RETRY_DELAYS[attempt],
|
||||
)
|
||||
except (aiohttp.ClientError, asyncio.TimeoutError) as error:
|
||||
last_error = error
|
||||
if attempt >= len(_UPLOAD_RETRY_DELAYS):
|
||||
break
|
||||
delay = _UPLOAD_RETRY_DELAYS[attempt]
|
||||
|
||||
print(
|
||||
f"[素材上传] 上传暂时失败，{delay:.0f}s 后重试 "
|
||||
f"({attempt + 1}/{len(_UPLOAD_RETRY_DELAYS)})..."
|
||||
)
|
||||
await asyncio.sleep(delay)
|
||||
|
||||
raise RuntimeError(
|
||||
f"临时素材上传失败（已重试 {len(_UPLOAD_RETRY_DELAYS)} 次）: {last_error}"
|
||||
) from None
|
||||
|
||||
|
||||
async def _put_upload(upload_url: str, data: bytes, content_type: str):
|
||||
"""用预签名 URL 直传文件到 R2（不带 Authorization）"""
|
||||
connector = aiohttp.TCPConnector(ssl=False)
|
||||
async with aiohttp.ClientSession(connector=connector) as session:
|
||||
async with session.put(
|
||||
upload_url,
|
||||
data=data,
|
||||
headers={"Content-Type": content_type},
|
||||
timeout=aiohttp.ClientTimeout(total=120),
|
||||
) as resp:
|
||||
if resp.status not in (200, 204):
|
||||
text = await resp.text()
|
||||
raise RuntimeError(f"文件上传失败 ({resp.status}): {text}")
|
||||
|
||||
|
||||
async def upload_image(pil_image) -> str:
|
||||
async def upload_image(pil_image, base_url: Optional[str] = None) -> str:
|
||||
"""
|
||||
接受 PIL Image 对象，编码为 PNG 上传到 R2，返回公网 URL。
|
||||
接受 PIL Image 对象，编码为 PNG 上传并返回临时公网 URL。
|
||||
"""
|
||||
check_interrupt()
|
||||
import io as _io
|
||||
buf = _io.BytesIO()
|
||||
pil_image.save(buf, format="PNG")
|
||||
data = buf.getvalue()
|
||||
filename = f"{uuid.uuid4()}.png"
|
||||
|
||||
upload_url, public_url = await _presign(filename, "image/png")
|
||||
await _put_upload(upload_url, data, "image/png")
|
||||
|
||||
print(f"[R2] 图片已上传: {public_url}")
|
||||
return public_url
|
||||
return await _upload_file(data, filename, "image/png", base_url=base_url)
|
||||
|
||||
|
||||
async def upload_video(video) -> str:
|
||||
async def upload_video(video, base_url: Optional[str] = None) -> str:
|
||||
"""
|
||||
接受 ComfyUI VIDEO 对象，上传到 R2，返回公网 URL。
|
||||
接受 ComfyUI VIDEO 对象，上传并返回临时公网 URL。
|
||||
支持 mp4 / mov 格式。
|
||||
"""
|
||||
source = video.get_stream_source()
|
||||
check_interrupt()
|
||||
if hasattr(video, "get_stream_source"):
|
||||
source = video.get_stream_source()
|
||||
elif isinstance(video, dict):
|
||||
source = (
|
||||
video.get("video")
|
||||
or video.get("path")
|
||||
or video.get("file")
|
||||
or video.get("filename")
|
||||
or video.get("source_path")
|
||||
)
|
||||
elif isinstance(video, str):
|
||||
source = video
|
||||
else:
|
||||
source = None
|
||||
for attr in ("source_path", "path", "video", "file", "filename"):
|
||||
if hasattr(video, attr):
|
||||
source = getattr(video, attr)
|
||||
break
|
||||
|
||||
if isinstance(source, io.BytesIO):
|
||||
source.seek(0)
|
||||
@@ -87,21 +185,41 @@ async def upload_video(video) -> str:
|
||||
with open(video_path, "rb") as f:
|
||||
data = f.read()
|
||||
|
||||
check_interrupt()
|
||||
content_type = "video/mp4" if ext == "mp4" else "video/quicktime"
|
||||
filename = f"{uuid.uuid4()}.{ext}"
|
||||
|
||||
upload_url, public_url = await _presign(filename, content_type)
|
||||
await _put_upload(upload_url, data, content_type)
|
||||
|
||||
print(f"[R2] 视频已上传: {public_url}")
|
||||
return public_url
|
||||
return await _upload_file(data, filename, content_type, base_url=base_url)
|
||||
|
||||
|
||||
async def upload_audio(audio) -> str:
|
||||
async def upload_audio(audio, base_url: Optional[str] = None) -> str:
|
||||
"""
|
||||
接受 ComfyUI AUDIO dict（waveform tensor + sample_rate），
|
||||
编码为 WAV 后上传到 R2，返回公网 URL。
|
||||
接受 ComfyUI AUDIO dict（waveform tensor + sample_rate）或音频文件路径。
|
||||
AUDIO 对象编码为 WAV；文件路径保留原格式上传。
|
||||
"""
|
||||
check_interrupt()
|
||||
if isinstance(audio, (str, os.PathLike)):
|
||||
audio_path = os.fspath(audio)
|
||||
if not os.path.isfile(audio_path):
|
||||
raise ValueError(f"无法获取参考音频文件路径（当前路径：{audio_path}）")
|
||||
ext = os.path.splitext(audio_path)[1].lower()
|
||||
content_types = {
|
||||
".wav": "audio/wav",
|
||||
".mp3": "audio/mpeg",
|
||||
".m4a": "audio/mp4",
|
||||
".aac": "audio/aac",
|
||||
".flac": "audio/flac",
|
||||
".ogg": "audio/ogg",
|
||||
}
|
||||
content_type = content_types.get(ext)
|
||||
if content_type is None:
|
||||
supported = "/".join(item.removeprefix(".") for item in content_types)
|
||||
raise ValueError(f"参考音频格式须为 {supported}，当前为 {ext or '无扩展名'}")
|
||||
with open(audio_path, "rb") as file:
|
||||
data = file.read()
|
||||
filename = f"{uuid.uuid4()}{ext}"
|
||||
return await _upload_file(data, filename, content_type, base_url=base_url)
|
||||
|
||||
import struct
|
||||
import numpy as np
|
||||
|
||||
@@ -138,8 +256,4 @@ async def upload_audio(audio) -> str:
|
||||
data = buf.getvalue()
|
||||
filename = f"{uuid.uuid4()}.wav"
|
||||
|
||||
upload_url, public_url = await _presign(filename, "audio/wav")
|
||||
await _put_upload(upload_url, data, "audio/wav")
|
||||
|
||||
print(f"[R2] 音频已上传: {public_url}")
|
||||
return public_url
|
||||
return await _upload_file(data, filename, "audio/wav", base_url=base_url)
|
||||
|
||||
@@ -0,0 +1,332 @@
|
||||
"""Reference-guided, geometry-preserving colour correction for generated images."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
import math
|
||||
from typing import Sequence
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
COLOR_CORRECTION_MODES = ("不纠正", "智能纠正")
|
||||
_D65_WHITE = (0.95047, 1.0, 1.08883)
|
||||
_LAB_EPSILON = 216.0 / 24389.0
|
||||
_LAB_KAPPA = 24389.0 / 27.0
|
||||
_ANALYSIS_LONG_EDGE = 384
|
||||
_MAX_CHROMA_SHIFT = 12.0
|
||||
_MIN_APPLY_CONFIDENCE = 0.18
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ColorCorrectionReport:
|
||||
applied: bool
|
||||
confidence: float
|
||||
shift_a: float
|
||||
shift_b: float
|
||||
scale_a: float
|
||||
scale_b: float
|
||||
|
||||
|
||||
def _srgb_to_lab(image: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
rgb = image[..., :3].float().clamp(0.0, 1.0)
|
||||
linear = torch.where(
|
||||
rgb <= 0.04045,
|
||||
rgb / 12.92,
|
||||
((rgb + 0.055) / 1.055).pow(2.4),
|
||||
)
|
||||
red, green, blue = linear.unbind(dim=-1)
|
||||
x = (0.4124564 * red + 0.3575761 * green + 0.1804375 * blue) / _D65_WHITE[0]
|
||||
y = 0.2126729 * red + 0.7151522 * green + 0.0721750 * blue
|
||||
z = (0.0193339 * red + 0.1191920 * green + 0.9503041 * blue) / _D65_WHITE[2]
|
||||
|
||||
def pivot(value: torch.Tensor) -> torch.Tensor:
|
||||
return torch.where(
|
||||
value > _LAB_EPSILON,
|
||||
value.clamp_min(0.0).pow(1.0 / 3.0),
|
||||
(_LAB_KAPPA * value + 16.0) / 116.0,
|
||||
)
|
||||
|
||||
fx, fy, fz = pivot(x), pivot(y), pivot(z)
|
||||
return 116.0 * fy - 16.0, 500.0 * (fx - fy), 200.0 * (fy - fz)
|
||||
|
||||
|
||||
def _lab_to_srgb_unclamped(
|
||||
lightness: torch.Tensor,
|
||||
a: torch.Tensor,
|
||||
b: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
fy = (lightness + 16.0) / 116.0
|
||||
fx = fy + a / 500.0
|
||||
fz = fy - b / 200.0
|
||||
|
||||
def inverse_pivot(value: torch.Tensor) -> torch.Tensor:
|
||||
cubed = value.pow(3.0)
|
||||
return torch.where(
|
||||
cubed > _LAB_EPSILON,
|
||||
cubed,
|
||||
(116.0 * value - 16.0) / _LAB_KAPPA,
|
||||
)
|
||||
|
||||
x = _D65_WHITE[0] * inverse_pivot(fx)
|
||||
y = inverse_pivot(fy)
|
||||
z = _D65_WHITE[2] * inverse_pivot(fz)
|
||||
red = 3.2404542 * x - 1.5371385 * y - 0.4985314 * z
|
||||
green = -0.9692660 * x + 1.8760108 * y + 0.0415560 * z
|
||||
blue = 0.0556434 * x - 0.2040259 * y + 1.0572252 * z
|
||||
linear = torch.stack((red, green, blue), dim=-1)
|
||||
positive = linear.clamp_min(0.0)
|
||||
return torch.where(
|
||||
linear <= 0.0031308,
|
||||
12.92 * linear,
|
||||
1.055 * positive.pow(1.0 / 2.4) - 0.055,
|
||||
)
|
||||
|
||||
|
||||
def _smoothstep(value: torch.Tensor) -> torch.Tensor:
|
||||
value = value.clamp(0.0, 1.0)
|
||||
return value * value * (3.0 - 2.0 * value)
|
||||
|
||||
|
||||
def _analysis_size(height: int, width: int) -> tuple[int, int]:
|
||||
scale = min(1.0, _ANALYSIS_LONG_EDGE / max(height, width))
|
||||
return max(8, round(height * scale)), max(8, round(width * scale))
|
||||
|
||||
|
||||
def _resize(image: torch.Tensor, size: tuple[int, int]) -> torch.Tensor:
|
||||
return F.interpolate(
|
||||
image[..., :3].permute(0, 3, 1, 2).float(),
|
||||
size=size,
|
||||
mode="bilinear",
|
||||
align_corners=False,
|
||||
antialias=True,
|
||||
).permute(0, 2, 3, 1)
|
||||
|
||||
|
||||
def _resize_cover(image: torch.Tensor, size: tuple[int, int]) -> torch.Tensor:
|
||||
target_h, target_w = size
|
||||
source_h, source_w = image.shape[1:3]
|
||||
scale = max(target_h / source_h, target_w / source_w)
|
||||
resized_h = max(target_h, math.ceil(source_h * scale))
|
||||
resized_w = max(target_w, math.ceil(source_w * scale))
|
||||
resized = _resize(image, (resized_h, resized_w))
|
||||
top = max(0, (resized_h - target_h) // 2)
|
||||
left = max(0, (resized_w - target_w) // 2)
|
||||
return resized[:, top:top + target_h, left:left + target_w]
|
||||
|
||||
|
||||
def _masked_channel_stats(
|
||||
a: torch.Tensor,
|
||||
b: torch.Tensor,
|
||||
mask: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
values = torch.stack((a[mask], b[mask]), dim=1)
|
||||
center = values.median(dim=0).values
|
||||
q25 = torch.quantile(values, 0.25, dim=0)
|
||||
q75 = torch.quantile(values, 0.75, dim=0)
|
||||
return center, (q75 - q25).clamp_min(1.0)
|
||||
|
||||
|
||||
def _estimate_mapping(
|
||||
target: torch.Tensor,
|
||||
reference: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, float]:
|
||||
size = _analysis_size(target.shape[1], target.shape[2])
|
||||
target_small = _resize(target, size)[0]
|
||||
reference_small = _resize_cover(reference, size)[0]
|
||||
target_l, target_a, target_b = _srgb_to_lab(target_small)
|
||||
reference_l, reference_a, reference_b = _srgb_to_lab(reference_small)
|
||||
target_chroma = torch.hypot(target_a, target_b)
|
||||
reference_chroma = torch.hypot(reference_a, reference_b)
|
||||
|
||||
valid = (
|
||||
(target_l >= 5.0)
|
||||
& (target_l <= 97.0)
|
||||
& (reference_l >= 5.0)
|
||||
& (reference_l <= 97.0)
|
||||
& (target_chroma <= 110.0)
|
||||
& (reference_chroma <= 110.0)
|
||||
)
|
||||
if int(valid.sum().item()) < 64:
|
||||
valid = torch.ones_like(valid, dtype=torch.bool)
|
||||
paired = valid & ((target_l - reference_l).abs() <= 20.0)
|
||||
minimum = max(64, int(paired.numel() * 0.01))
|
||||
|
||||
target_center, target_spread = _masked_channel_stats(
|
||||
target_a,
|
||||
target_b,
|
||||
valid,
|
||||
)
|
||||
reference_center, reference_spread = _masked_channel_stats(
|
||||
reference_a,
|
||||
reference_b,
|
||||
valid,
|
||||
)
|
||||
unpaired_shift = reference_center - target_center
|
||||
|
||||
paired_count = int(paired.sum().item())
|
||||
if paired_count >= minimum:
|
||||
paired_diff = torch.stack(
|
||||
((reference_a - target_a)[paired], (reference_b - target_b)[paired]),
|
||||
dim=1,
|
||||
)
|
||||
preliminary = paired_diff.median(dim=0).values
|
||||
distance = torch.linalg.vector_norm(paired_diff - preliminary, dim=1)
|
||||
cutoff = torch.quantile(distance, 0.8)
|
||||
robust = paired_diff[distance <= cutoff]
|
||||
paired_shift = robust.median(dim=0).values if robust.numel() else preliminary
|
||||
match_ratio = paired_count / paired.numel()
|
||||
luma_error = float((target_l[paired] - reference_l[paired]).abs().median().item())
|
||||
target_ratio = target.shape[2] / target.shape[1]
|
||||
reference_ratio = reference.shape[2] / reference.shape[1]
|
||||
aspect_confidence = min(target_ratio, reference_ratio) / max(target_ratio, reference_ratio)
|
||||
confidence = min(1.0, match_ratio / 0.35) * max(0.0, 1.0 - luma_error / 24.0)
|
||||
confidence *= aspect_confidence
|
||||
else:
|
||||
paired_shift = unpaired_shift
|
||||
confidence = 0.0
|
||||
|
||||
blend = max(0.0, min(1.0, confidence))
|
||||
shift = paired_shift * blend + unpaired_shift * (1.0 - blend)
|
||||
shift_norm = float(torch.linalg.vector_norm(shift).item())
|
||||
if shift_norm > _MAX_CHROMA_SHIFT:
|
||||
shift = shift * (_MAX_CHROMA_SHIFT / shift_norm)
|
||||
|
||||
scale = (reference_spread / target_spread).clamp(0.88, 1.12)
|
||||
return target_center, shift, scale, blend
|
||||
|
||||
|
||||
def _apply_mapping(
|
||||
target: torch.Tensor,
|
||||
center: torch.Tensor,
|
||||
shift: torch.Tensor,
|
||||
scale: torch.Tensor,
|
||||
confidence: float,
|
||||
) -> torch.Tensor:
|
||||
output_chunks: list[torch.Tensor] = []
|
||||
effective_strength = 0.30 + 0.55 * confidence
|
||||
for top in range(0, target.shape[1], 256):
|
||||
chunk = target[:, top:top + 256].float()
|
||||
lightness, a, b = _srgb_to_lab(chunk)
|
||||
chroma = torch.hypot(a, b)
|
||||
tonal_weight = _smoothstep((lightness - 3.0) / 12.0) * _smoothstep(
|
||||
(97.0 - lightness) / 12.0
|
||||
)
|
||||
saturation_protection = 1.0 - 0.55 * _smoothstep((chroma - 45.0) / 35.0)
|
||||
weight = effective_strength * tonal_weight * saturation_protection
|
||||
|
||||
mapped_a = center[0] + (a - center[0]) * scale[0] + shift[0]
|
||||
mapped_b = center[1] + (b - center[1]) * scale[1] + shift[1]
|
||||
candidate_a = a + (mapped_a - a) * weight
|
||||
candidate_b = b + (mapped_b - b) * weight
|
||||
|
||||
# Compress only the proposed chroma movement when it exits sRGB gamut.
|
||||
# This avoids the hue discontinuities caused by hard per-channel clipping.
|
||||
for _ in range(5):
|
||||
raw = _lab_to_srgb_unclamped(lightness, candidate_a, candidate_b)
|
||||
outside = ((raw < 0.0) | (raw > 1.0)).any(dim=-1)
|
||||
if not bool(outside.any()):
|
||||
break
|
||||
candidate_a = torch.where(outside, a + (candidate_a - a) * 0.72, candidate_a)
|
||||
candidate_b = torch.where(outside, b + (candidate_b - b) * 0.72, candidate_b)
|
||||
|
||||
corrected_rgb = _lab_to_srgb_unclamped(
|
||||
lightness,
|
||||
candidate_a,
|
||||
candidate_b,
|
||||
).clamp(0.0, 1.0)
|
||||
if chunk.shape[-1] > 3:
|
||||
corrected_rgb = torch.cat((corrected_rgb, chunk[..., 3:]), dim=-1)
|
||||
output_chunks.append(corrected_rgb)
|
||||
return torch.cat(output_chunks, dim=1).to(dtype=target.dtype)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def correct_tensor_with_reference(
|
||||
images: torch.Tensor,
|
||||
reference: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, list[ColorCorrectionReport]]:
|
||||
"""Correct BHWC images against the first reference without spatial resampling."""
|
||||
if images.ndim != 4 or images.shape[-1] < 3:
|
||||
raise ValueError("色彩纠正需要 BHWC 格式的图像张量")
|
||||
if reference.ndim != 4 or reference.shape[0] < 1 or reference.shape[-1] < 3:
|
||||
raise ValueError("色彩纠正参考图格式无效")
|
||||
|
||||
reference = reference[:1].to(device=images.device, dtype=torch.float32)
|
||||
corrected: list[torch.Tensor] = []
|
||||
reports: list[ColorCorrectionReport] = []
|
||||
for index in range(images.shape[0]):
|
||||
target = images[index:index + 1]
|
||||
center, shift, scale, confidence = _estimate_mapping(target.float(), reference)
|
||||
meaningful = bool(
|
||||
confidence >= _MIN_APPLY_CONFIDENCE
|
||||
and (
|
||||
torch.linalg.vector_norm(shift).item() >= 0.05
|
||||
or torch.max(torch.abs(scale - 1.0)).item() >= 0.005
|
||||
)
|
||||
)
|
||||
output = (
|
||||
_apply_mapping(target, center, shift, scale, confidence)
|
||||
if meaningful
|
||||
else target.clone()
|
||||
)
|
||||
corrected.append(output)
|
||||
reports.append(ColorCorrectionReport(
|
||||
applied=meaningful,
|
||||
confidence=float(confidence),
|
||||
shift_a=float(shift[0].item()),
|
||||
shift_b=float(shift[1].item()),
|
||||
scale_a=float(scale[0].item()),
|
||||
scale_b=float(scale[1].item()),
|
||||
))
|
||||
return torch.cat(corrected, dim=0), reports
|
||||
|
||||
|
||||
def _pil_to_tensor(image: Image.Image) -> torch.Tensor:
|
||||
mode = "RGBA" if "A" in image.getbands() else "RGB"
|
||||
array = np.asarray(image.convert(mode), dtype=np.float32).copy() / 255.0
|
||||
return torch.from_numpy(array).unsqueeze(0)
|
||||
|
||||
|
||||
def _tensor_to_pil(image: torch.Tensor, mode: str) -> Image.Image:
|
||||
array = (
|
||||
image[0].detach().cpu().clamp(0.0, 1.0).mul(255.0).round().byte().numpy()
|
||||
)
|
||||
return Image.fromarray(array, mode=mode)
|
||||
|
||||
|
||||
def correct_pil_with_reference(
|
||||
image: Image.Image,
|
||||
reference: Image.Image,
|
||||
) -> tuple[Image.Image, ColorCorrectionReport]:
|
||||
target_mode = "RGBA" if "A" in image.getbands() else "RGB"
|
||||
corrected, reports = correct_tensor_with_reference(
|
||||
_pil_to_tensor(image),
|
||||
_pil_to_tensor(reference.convert("RGB")),
|
||||
)
|
||||
return _tensor_to_pil(corrected, target_mode), reports[0]
|
||||
|
||||
|
||||
def correct_pil_sequence_with_reference(
|
||||
images: Sequence[Image.Image],
|
||||
reference: Image.Image,
|
||||
) -> tuple[list[Image.Image], list[ColorCorrectionReport]]:
|
||||
corrected: list[Image.Image] = []
|
||||
reports: list[ColorCorrectionReport] = []
|
||||
for image in images:
|
||||
result, report = correct_pil_with_reference(image, reference)
|
||||
corrected.append(result)
|
||||
reports.append(report)
|
||||
return corrected, reports
|
||||
|
||||
|
||||
__all__ = [
|
||||
"COLOR_CORRECTION_MODES",
|
||||
"ColorCorrectionReport",
|
||||
"correct_pil_sequence_with_reference",
|
||||
"correct_pil_with_reference",
|
||||
"correct_tensor_with_reference",
|
||||
]
|
||||
@@ -0,0 +1,257 @@
|
||||
"""Shared Seedance asset creation and safe ID reuse.
|
||||
|
||||
The cache deliberately stores only content fingerprints and provider asset IDs.
|
||||
Upload URLs, local paths, credentials, and response envelopes never cross this
|
||||
boundary.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Any
|
||||
|
||||
from ..clients.seedance_element_client import SeedanceElementClient
|
||||
|
||||
|
||||
_ROUTE_REQUEST_TYPES = {
|
||||
"overseas_hc": "hc",
|
||||
"domestic": "doubao",
|
||||
}
|
||||
_CACHE_VERSION = 1
|
||||
_MAX_CACHE_ENTRIES = 512
|
||||
_SAFE_ASSET_ID = re.compile(r"^[A-Za-z0-9._:-]{1,256}$")
|
||||
_CACHE_LOCKS: dict[tuple[int, str, str], asyncio.Lock] = {}
|
||||
_FILE_LOCKS: dict[str, threading.Lock] = {}
|
||||
_FILE_LOCKS_GUARD = threading.Lock()
|
||||
|
||||
|
||||
def seedance_asset_request_type(route: str) -> str:
|
||||
"""Map the unified video route to the material API request type."""
|
||||
|
||||
try:
|
||||
return _ROUTE_REQUEST_TYPES[str(route).strip()]
|
||||
except KeyError:
|
||||
raise ValueError(f"不支持的 Seedance 素材线路：{route}") from None
|
||||
|
||||
|
||||
def seedance_asset_fingerprint(path: str, request_type: str, asset_type: str) -> str:
|
||||
"""Hash content plus the provider namespace used to create the asset."""
|
||||
|
||||
digest = hashlib.sha256()
|
||||
digest.update(str(request_type).strip().lower().encode("utf-8"))
|
||||
digest.update(b"\0")
|
||||
digest.update(str(asset_type).strip().lower().encode("utf-8"))
|
||||
digest.update(b"\0")
|
||||
with open(path, "rb") as handle:
|
||||
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
|
||||
digest.update(chunk)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
class SeedanceAssetCache:
|
||||
"""Small atomic cache containing no provider URLs or local media paths."""
|
||||
|
||||
def __init__(self, path: str):
|
||||
self.path = os.path.abspath(path)
|
||||
with _FILE_LOCKS_GUARD:
|
||||
self._lock = _FILE_LOCKS.setdefault(self.path, threading.Lock())
|
||||
|
||||
def _ensure_parent(self) -> None:
|
||||
parent = os.path.dirname(self.path)
|
||||
if parent:
|
||||
os.makedirs(parent, exist_ok=True)
|
||||
|
||||
def _read_unlocked(self) -> dict[str, dict[str, Any]]:
|
||||
if not os.path.isfile(self.path):
|
||||
return {}
|
||||
try:
|
||||
with open(self.path, "r", encoding="utf-8") as handle:
|
||||
payload = json.load(handle)
|
||||
except (OSError, json.JSONDecodeError):
|
||||
return {}
|
||||
if not isinstance(payload, dict) or payload.get("version") != _CACHE_VERSION:
|
||||
return {}
|
||||
entries = payload.get("entries")
|
||||
return entries if isinstance(entries, dict) else {}
|
||||
|
||||
def get(self, fingerprint: str, request_type: str, asset_type: str) -> str | None:
|
||||
with self._lock:
|
||||
item = self._read_unlocked().get(fingerprint)
|
||||
if not isinstance(item, dict):
|
||||
return None
|
||||
if item.get("request_type") != request_type or item.get("asset_type") != asset_type:
|
||||
return None
|
||||
asset_id = str(item.get("asset_id") or "").strip()
|
||||
return asset_id if _SAFE_ASSET_ID.fullmatch(asset_id) else None
|
||||
|
||||
def put(self, fingerprint: str, asset_id: str, request_type: str, asset_type: str) -> None:
|
||||
with self._lock:
|
||||
entries = self._read_unlocked()
|
||||
entries[fingerprint] = {
|
||||
"asset_id": str(asset_id),
|
||||
"request_type": request_type,
|
||||
"asset_type": asset_type,
|
||||
"updated_at": int(time.time()),
|
||||
}
|
||||
if len(entries) > _MAX_CACHE_ENTRIES:
|
||||
entries = dict(
|
||||
sorted(
|
||||
entries.items(),
|
||||
key=lambda pair: int(pair[1].get("updated_at") or 0),
|
||||
reverse=True,
|
||||
)[:_MAX_CACHE_ENTRIES]
|
||||
)
|
||||
self._ensure_parent()
|
||||
temporary = f"{self.path}.{uuid.uuid4().hex}.tmp"
|
||||
try:
|
||||
with open(temporary, "w", encoding="utf-8") as handle:
|
||||
json.dump(
|
||||
{"version": _CACHE_VERSION, "entries": entries},
|
||||
handle,
|
||||
ensure_ascii=False,
|
||||
indent=2,
|
||||
)
|
||||
os.replace(temporary, self.path)
|
||||
finally:
|
||||
if os.path.exists(temporary):
|
||||
try:
|
||||
os.remove(temporary)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
def discard(self, fingerprint: str) -> None:
|
||||
with self._lock:
|
||||
entries = self._read_unlocked()
|
||||
if fingerprint not in entries:
|
||||
return
|
||||
entries.pop(fingerprint, None)
|
||||
self._ensure_parent()
|
||||
temporary = f"{self.path}.{uuid.uuid4().hex}.tmp"
|
||||
try:
|
||||
with open(temporary, "w", encoding="utf-8") as handle:
|
||||
json.dump(
|
||||
{"version": _CACHE_VERSION, "entries": entries},
|
||||
handle,
|
||||
ensure_ascii=False,
|
||||
indent=2,
|
||||
)
|
||||
os.replace(temporary, self.path)
|
||||
finally:
|
||||
if os.path.exists(temporary):
|
||||
try:
|
||||
os.remove(temporary)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
class SeedanceAssetService:
|
||||
"""Create Active assets through the same service used by the node and panel."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
base_url: str | None = None,
|
||||
client: SeedanceElementClient | None = None,
|
||||
cache_path: str | None = None,
|
||||
):
|
||||
self.client = client or SeedanceElementClient(base_url=base_url)
|
||||
self.cache = SeedanceAssetCache(cache_path) if cache_path else None
|
||||
|
||||
async def create_from_url(
|
||||
self,
|
||||
*,
|
||||
name: str,
|
||||
asset_url: str | None = None,
|
||||
asset_url_factory: Callable[[], Awaitable[str]] | None = None,
|
||||
asset_type: str,
|
||||
request_type: str,
|
||||
fingerprint: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
normalized_type = str(asset_type).strip().lower()
|
||||
normalized_request = str(request_type).strip().lower()
|
||||
|
||||
async def resolve_asset_url() -> str:
|
||||
value = asset_url
|
||||
if value is None and asset_url_factory is not None:
|
||||
value = await asset_url_factory()
|
||||
if not value or not str(value).startswith("https://"):
|
||||
raise ValueError("创建 Seedance 素材必须使用 HTTPS 上传地址")
|
||||
return str(value)
|
||||
|
||||
def finalized(result: dict[str, Any], *, reused: bool) -> dict[str, Any]:
|
||||
asset_id = str(result.get("Id") or "").strip()
|
||||
if not _SAFE_ASSET_ID.fullmatch(asset_id):
|
||||
raise RuntimeError("Seedance 素材已激活但未返回安全有效的 ID")
|
||||
value = dict(result)
|
||||
value["Id"] = asset_id
|
||||
value["_reused"] = reused
|
||||
return value
|
||||
|
||||
if not fingerprint or self.cache is None:
|
||||
result = await self.client.create_hc_asset_and_wait(
|
||||
name=name,
|
||||
asset_url=await resolve_asset_url(),
|
||||
asset_type=asset_type,
|
||||
request_type=normalized_request,
|
||||
)
|
||||
return finalized(result, reused=False)
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
lock_key = (id(loop), self.cache.path, fingerprint)
|
||||
lock = _CACHE_LOCKS.setdefault(lock_key, asyncio.Lock())
|
||||
async with lock:
|
||||
cached_id = await asyncio.to_thread(
|
||||
self.cache.get,
|
||||
fingerprint,
|
||||
normalized_request,
|
||||
normalized_type,
|
||||
)
|
||||
if cached_id:
|
||||
try:
|
||||
cached = await self.client.get_hc_asset(
|
||||
cached_id,
|
||||
request_type=normalized_request,
|
||||
)
|
||||
except RuntimeError as exc:
|
||||
if "(404)" not in str(exc) and "(410)" not in str(exc):
|
||||
raise
|
||||
cached = None
|
||||
if cached is not None:
|
||||
if str(cached.get("Status") or "").strip().lower() == "active":
|
||||
result = dict(cached)
|
||||
result["Id"] = cached_id
|
||||
return finalized(result, reused=True)
|
||||
try:
|
||||
await asyncio.to_thread(self.cache.discard, fingerprint)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
result = await self.client.create_hc_asset_and_wait(
|
||||
name=name,
|
||||
asset_url=await resolve_asset_url(),
|
||||
asset_type=asset_type,
|
||||
request_type=normalized_request,
|
||||
)
|
||||
result = finalized(result, reused=False)
|
||||
asset_id = result["Id"]
|
||||
try:
|
||||
await asyncio.to_thread(
|
||||
self.cache.put,
|
||||
fingerprint,
|
||||
asset_id,
|
||||
normalized_request,
|
||||
normalized_type,
|
||||
)
|
||||
except OSError:
|
||||
# A local reuse optimization must not turn an already-created
|
||||
# provider asset into a failed video job.
|
||||
pass
|
||||
return result
|
||||
@@ -1,148 +0,0 @@
|
||||
"""
|
||||
更新检查工具
|
||||
在插件加载时检查是否有新版本
|
||||
"""
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
from typing import Optional
|
||||
|
||||
|
||||
def get_current_version() -> Optional[str]:
|
||||
"""
|
||||
获取当前版本号
|
||||
|
||||
Returns:
|
||||
版本号字符串，如果读取失败返回 None
|
||||
"""
|
||||
version_file = os.path.join(os.path.dirname(os.path.dirname(__file__)), "version.txt")
|
||||
try:
|
||||
with open(version_file, 'r', encoding='utf-8') as f:
|
||||
return f.read().strip()
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def check_for_updates() -> bool:
|
||||
"""
|
||||
检查是否有更新
|
||||
|
||||
Returns:
|
||||
True 如果有更新，False 如果已是最新或检查失败
|
||||
"""
|
||||
try:
|
||||
# 获取当前目录
|
||||
plugin_dir = os.path.dirname(os.path.dirname(__file__))
|
||||
|
||||
# 检查是否是 Git 仓库
|
||||
git_dir = os.path.join(plugin_dir, '.git')
|
||||
if not os.path.exists(git_dir):
|
||||
return False
|
||||
|
||||
# 执行 git fetch（禁止弹出认证弹框，失败时静默处理）
|
||||
env = os.environ.copy()
|
||||
env['GIT_TERMINAL_PROMPT'] = '0'
|
||||
subprocess.run(
|
||||
['git', 'fetch', 'origin'],
|
||||
cwd=plugin_dir,
|
||||
capture_output=True,
|
||||
timeout=10,
|
||||
env=env
|
||||
)
|
||||
|
||||
# 检查本地和远程版本
|
||||
local = subprocess.run(
|
||||
['git', 'rev-parse', '@'],
|
||||
cwd=plugin_dir,
|
||||
capture_output=True,
|
||||
text=True
|
||||
).stdout.strip()
|
||||
|
||||
remote = subprocess.run(
|
||||
['git', 'rev-parse', '@{u}'],
|
||||
cwd=plugin_dir,
|
||||
capture_output=True,
|
||||
text=True
|
||||
).stdout.strip()
|
||||
|
||||
return local != remote
|
||||
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def get_update_changelog() -> list:
|
||||
"""从远程 CHANGELOG.md 最新版本块中提取更新内容（最多5条）"""
|
||||
try:
|
||||
plugin_dir = os.path.dirname(os.path.dirname(__file__))
|
||||
result = subprocess.run(
|
||||
['git', 'show', 'origin/main:CHANGELOG.md'],
|
||||
cwd=plugin_dir,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
encoding='utf-8'
|
||||
)
|
||||
lines = result.stdout.splitlines()
|
||||
|
||||
in_block = False
|
||||
items = []
|
||||
for line in lines:
|
||||
if line.startswith('## [') and not line.startswith('## [Unreleased]'):
|
||||
if in_block:
|
||||
break
|
||||
in_block = True
|
||||
continue
|
||||
if in_block:
|
||||
stripped = line.strip()
|
||||
if stripped and not stripped.startswith('#') and not stripped.startswith('---'):
|
||||
text = stripped.lstrip('- ').replace('**', '').strip()
|
||||
if text and len(text) > 3:
|
||||
items.append(text)
|
||||
if len(items) >= 5:
|
||||
break
|
||||
|
||||
return items
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
|
||||
def notify_new_version():
|
||||
"""检测到新版本时，推送蓝色更新通知弹框"""
|
||||
changelog = get_update_changelog()
|
||||
|
||||
try:
|
||||
import threading
|
||||
from server import PromptServer
|
||||
|
||||
def _send():
|
||||
try:
|
||||
PromptServer.instance.send_sync(
|
||||
"o1key.new_version",
|
||||
{"changelog": changelog}
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
threading.Timer(3.0, _send).start()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def notify_update_available():
|
||||
"""通知用户有更新可用（前端弹窗）"""
|
||||
try:
|
||||
import threading
|
||||
from server import PromptServer
|
||||
|
||||
def _send():
|
||||
try:
|
||||
PromptServer.instance.send_sync(
|
||||
"o1key.update_available",
|
||||
{"message": "欢迎使用o1key工作流，祝您马年，马上有福，马上有钱，马到成功！！！"}
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
threading.Timer(3.0, _send).start()
|
||||
except Exception:
|
||||
pass
|
||||
+57
-22
@@ -1,4 +1,4 @@
|
||||
"""Safely fast-forward a Git installation of this node package."""
|
||||
"""Fast-forward a clean Git installation from the public O1Key release repo."""
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
@@ -6,10 +6,22 @@ from pathlib import Path
|
||||
|
||||
|
||||
PLUGIN_DIR = Path(__file__).resolve().parent.parent
|
||||
RELEASE_REPOSITORY_URL = "https://git.o1key.com/publisher/comfyui_o1key.git"
|
||||
|
||||
|
||||
class UpdateError(Exception):
|
||||
pass
|
||||
def __init__(self, code, message, suggestion, status=409):
|
||||
super().__init__(message)
|
||||
self.code = code
|
||||
self.suggestion = suggestion
|
||||
self.status = status
|
||||
|
||||
def as_dict(self):
|
||||
return {
|
||||
"code": self.code,
|
||||
"error": str(self),
|
||||
"suggestion": self.suggestion,
|
||||
}
|
||||
|
||||
|
||||
def _git(*args, timeout=60, check=True):
|
||||
@@ -18,49 +30,72 @@ def _git(*args, timeout=60, check=True):
|
||||
env["GCM_INTERACTIVE"] = "Never"
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["git", *args],
|
||||
cwd=PLUGIN_DIR,
|
||||
env=env,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
encoding="utf-8",
|
||||
errors="replace",
|
||||
timeout=timeout,
|
||||
["git", *args], cwd=PLUGIN_DIR, env=env,
|
||||
capture_output=True, text=True, encoding="utf-8",
|
||||
errors="replace", timeout=timeout,
|
||||
)
|
||||
except FileNotFoundError as exc:
|
||||
raise UpdateError("未找到 Git，请先安装 Git。") from exc
|
||||
raise UpdateError(
|
||||
"git_missing", "未找到 Git。", "安装 Git 后重启 ComfyUI，再重新检查更新。", 503,
|
||||
) from exc
|
||||
except subprocess.TimeoutExpired as exc:
|
||||
raise UpdateError("Git 操作超时，请检查网络后重试。") from exc
|
||||
raise UpdateError(
|
||||
"timeout", "连接发布仓库超时。", "检查网络连接，稍后再试。", 504,
|
||||
) from exc
|
||||
if check and result.returncode:
|
||||
detail = (result.stderr or result.stdout).strip().splitlines()
|
||||
raise UpdateError(detail[-1] if detail else "Git 操作失败。")
|
||||
raise UpdateError(
|
||||
"git_failed", "Git 操作未完成。", "检查插件目录的 Git 状态后重试。",
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
def update_package():
|
||||
"""Update origin/main without discarding local changes or switching branches."""
|
||||
"""Fetch the release main branch and fast-forward only a clean local main."""
|
||||
if not (PLUGIN_DIR / ".git").exists():
|
||||
raise UpdateError("当前节点包不是 Git 安装。请通过 Git 安装后再使用界面更新。")
|
||||
raise UpdateError(
|
||||
"not_git", "当前插件不是 Git 安装。",
|
||||
"请从发布仓库重新以 Git 安装；现有配置文件先单独备份。",
|
||||
)
|
||||
|
||||
branch = _git("symbolic-ref", "--quiet", "--short", "HEAD", check=False)
|
||||
if branch.returncode or branch.stdout.strip() != "main":
|
||||
raise UpdateError("当前不在 main 分支，请手动检查分支后更新。")
|
||||
raise UpdateError(
|
||||
"wrong_branch", "当前不在 main 分支。",
|
||||
"请先检查并切换分支；本地分支上的修改不会自动合并。",
|
||||
)
|
||||
|
||||
if _git("status", "--porcelain", "--untracked-files=no").stdout.strip():
|
||||
raise UpdateError("节点包有本地修改，请先保存或处理修改后再更新。")
|
||||
raise UpdateError(
|
||||
"local_changes", "插件目录有未提交的代码修改。",
|
||||
"请先保存、提交或暂存修改；更新不会覆盖这些文件。",
|
||||
)
|
||||
|
||||
old_commit = _git("rev-parse", "HEAD").stdout.strip()
|
||||
old_requirements = _git("show", "HEAD:requirements.txt", check=False).stdout
|
||||
_git("fetch", "origin", "main")
|
||||
fetched = _git("fetch", "--no-tags", RELEASE_REPOSITORY_URL, "main", timeout=90, check=False)
|
||||
if fetched.returncode:
|
||||
raise UpdateError(
|
||||
"fetch_failed", "无法获取发布仓库的 main 分支。",
|
||||
"检查 git.o1key.com 的网络连接与仓库读取权限，稍后重试。", 503,
|
||||
)
|
||||
new_commit = _git("rev-parse", "FETCH_HEAD").stdout.strip()
|
||||
if old_commit == new_commit:
|
||||
return {"updated": False, "version": old_commit[:7], "requirements_changed": False}
|
||||
|
||||
if _git("merge-base", "--is-ancestor", "HEAD", "FETCH_HEAD", check=False).returncode:
|
||||
raise UpdateError("本地与 origin/main 已分叉，无法安全快进。请手动处理。")
|
||||
raise UpdateError(
|
||||
"diverged", "本地提交与发布仓库已分叉，无法自动快进。",
|
||||
"请手动比较两个分支并合并，不要强制重置本地文件。",
|
||||
)
|
||||
|
||||
_git("merge", "--ff-only", "FETCH_HEAD")
|
||||
requirements_changed = old_requirements != (PLUGIN_DIR / "requirements.txt").read_text(encoding="utf-8")
|
||||
new_requirements = _git("show", "FETCH_HEAD:requirements.txt", check=False).stdout
|
||||
merged = _git("merge", "--ff-only", "FETCH_HEAD", check=False)
|
||||
if merged.returncode:
|
||||
raise UpdateError(
|
||||
"merge_blocked", "更新被本地文件阻止。",
|
||||
"检查是否有与新版本重名的未跟踪文件，保留文件后手动处理。",
|
||||
)
|
||||
requirements_changed = old_requirements != new_requirements
|
||||
return {
|
||||
"updated": True,
|
||||
"version": new_commit[:7],
|
||||
|
||||
+291
-1
@@ -1,5 +1,40 @@
|
||||
import asyncio
|
||||
from typing import Any, Dict
|
||||
import os
|
||||
import time
|
||||
from typing import Any, Callable, Dict, Optional
|
||||
|
||||
import aiohttp
|
||||
|
||||
|
||||
# 视频任务轮询总时长上限（秒）。超过后停止等待并抛出 TimeoutError，
|
||||
# 避免任务在服务端长时间无进展时无限轮询、迫使用户手动中断。
|
||||
POLL_DEADLINE_SECONDS = 2000
|
||||
|
||||
|
||||
class PollDeadline:
|
||||
"""轮询看门狗：累计等待超过上限即抛出 TimeoutError。
|
||||
|
||||
用法：
|
||||
deadline = PollDeadline(label="K3 图生视频")
|
||||
while True:
|
||||
deadline.check()
|
||||
...
|
||||
"""
|
||||
|
||||
def __init__(self, seconds: float = POLL_DEADLINE_SECONDS, label: str = "视频任务"):
|
||||
self.seconds = seconds
|
||||
self.label = label
|
||||
self.start = time.time()
|
||||
|
||||
def elapsed(self) -> float:
|
||||
return time.time() - self.start
|
||||
|
||||
def check(self) -> None:
|
||||
if self.elapsed() >= self.seconds:
|
||||
raise TimeoutError(
|
||||
f"{self.label} 轮询已超过 {self.seconds:.0f}s 仍未完成，已停止等待。"
|
||||
f"任务可能仍在服务端生成，请稍后重试。"
|
||||
)
|
||||
|
||||
|
||||
try:
|
||||
@@ -29,6 +64,25 @@ FAILURE_STATUSES = {
|
||||
}
|
||||
|
||||
|
||||
SEEDANCE_COPYRIGHT_RESTRICTION_ERRORS = (
|
||||
"The request failed because the output video may be related to copyright restriction",
|
||||
"The request failed because the output video may be related to copyright restrictions",
|
||||
)
|
||||
SEEDANCE_COPYRIGHT_RESTRICTION_MESSAGE = "输出视频触发版权审查被拒绝生成！"
|
||||
|
||||
|
||||
def format_seedance_generation_error(value: Any) -> str:
|
||||
"""Normalize errors shared by the two interactive Seedance video nodes."""
|
||||
message = str(value or "视频生成失败")
|
||||
normalized_message = message.lower()
|
||||
if any(
|
||||
marker.lower() in normalized_message
|
||||
for marker in SEEDANCE_COPYRIGHT_RESTRICTION_ERRORS
|
||||
):
|
||||
return SEEDANCE_COPYRIGHT_RESTRICTION_MESSAGE
|
||||
return message
|
||||
|
||||
|
||||
def check_interrupt() -> None:
|
||||
if INTERRUPT_AVAILABLE and processing_interrupted():
|
||||
raise InterruptProcessingException()
|
||||
@@ -162,6 +216,12 @@ def extract_video_url(payload: Dict[str, Any]) -> str | None:
|
||||
value = first.get("url") or first.get("video_url")
|
||||
if value:
|
||||
return str(value)
|
||||
|
||||
# 腾讯 Kling（v3-t）渠道：完成时视频地址在 metadata.url
|
||||
metadata = _as_dict(source.get("metadata"))
|
||||
value = metadata.get("url") or metadata.get("video_url")
|
||||
if value:
|
||||
return str(value)
|
||||
return None
|
||||
|
||||
|
||||
@@ -179,3 +239,233 @@ def is_failure_status(status: str, payload: Dict[str, Any] | None = None) -> boo
|
||||
root, data, inner = _nested_payloads(payload)
|
||||
failure_keys = ("error", "fail_reason", "failure_reason", "task_status_msg", "error_message")
|
||||
return any(any(source.get(key) for key in failure_keys) for source in (data, inner, root))
|
||||
|
||||
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
# 视频下载：抗超时 / 可断点续传 / 无限重试 / 可随时取消
|
||||
# ═══════════════════════════════════════════════════════════════════════════════
|
||||
#
|
||||
# 设计目标（针对“视频已在服务端生成成功、后台已扣费，必须把成品拿到手”的场景）：
|
||||
# 1. 不用固定 total 超时一刀切大文件——只要持续有数据就一直下载；
|
||||
# 用 sock_read 检测“卡死”（连续 N 秒收不到任何字节）才判定异常。
|
||||
# 2. 网络抖动 / 超时 / 5xx / 连接中断 → 退避后无限重试，直到成功。
|
||||
# 3. 已下载的字节用 HTTP Range 断点续传，不从 0 重来。
|
||||
# 4. 永久性错误（403/404/410 等）快速失败，不做无意义的死循环。
|
||||
# 5. 全程可被 ComfyUI 的“取消”随时打断（每个 attempt 包在 run_with_interrupt 中，
|
||||
# 分块写入与退避等待都会检查中断）。
|
||||
|
||||
# 连续多少秒收不到任何数据就判定当前连接卡死（触发重试，而非整体失败）
|
||||
DOWNLOAD_SOCK_READ_TIMEOUT = 120
|
||||
# 建立连接的超时
|
||||
DOWNLOAD_CONNECT_TIMEOUT = 30
|
||||
# 重试退避：起始 / 上限（秒）
|
||||
DOWNLOAD_RETRY_BASE_DELAY = 2.0
|
||||
DOWNLOAD_RETRY_MAX_DELAY = 30.0
|
||||
# 视为“永久失败、无需重试”的 HTTP 状态码
|
||||
DOWNLOAD_PERMANENT_STATUS = {400, 401, 403, 404, 405, 410, 451}
|
||||
|
||||
|
||||
class _PermanentDownloadError(RuntimeError):
|
||||
"""不可重试的下载错误（如 403/404）。"""
|
||||
|
||||
|
||||
class _IncompleteDownloadError(RuntimeError):
|
||||
"""连接被提前关闭、文件未下完，需要续传重试。"""
|
||||
|
||||
|
||||
def _download_timeout(sock_read: float) -> aiohttp.ClientTimeout:
|
||||
# total=None：不限制总时长，让缓慢但持续的大文件下载得以完成；
|
||||
# sock_connect/sock_read：分别约束“连接建立”和“两次收包之间”的最大间隔。
|
||||
return aiohttp.ClientTimeout(
|
||||
total=None,
|
||||
connect=None,
|
||||
sock_connect=DOWNLOAD_CONNECT_TIMEOUT,
|
||||
sock_read=sock_read,
|
||||
)
|
||||
|
||||
|
||||
async def _stream_once(
|
||||
session: aiohttp.ClientSession,
|
||||
url: str,
|
||||
fobj,
|
||||
*,
|
||||
headers: Optional[dict],
|
||||
resume_from: int,
|
||||
sock_read: float,
|
||||
chunk_size: int,
|
||||
on_bytes: Optional[Callable[[int], None]],
|
||||
) -> int:
|
||||
"""发起一次 GET 并把响应体写入已打开的文件对象 fobj。
|
||||
|
||||
返回“本次结束后文件应当达到的总字节数”（已知时），未知时返回 -1。
|
||||
若服务端支持 Range 且 resume_from>0，则带上 Range 头从断点继续；
|
||||
否则从头下载（必要时先 truncate）。出错时抛异常，由上层决定是否重试。
|
||||
"""
|
||||
req_headers = dict(headers) if headers else {}
|
||||
if resume_from > 0:
|
||||
req_headers["Range"] = f"bytes={resume_from}-"
|
||||
|
||||
timeout = _download_timeout(sock_read)
|
||||
async with session.get(
|
||||
url, headers=req_headers or None, timeout=timeout, allow_redirects=True
|
||||
) as resp:
|
||||
status = resp.status
|
||||
|
||||
# 206：服务端接受断点续传，从 resume_from 续写。
|
||||
# 200：服务端忽略 Range（或本就从头下），需从文件起点重写。
|
||||
base = resume_from
|
||||
if resume_from > 0 and status == 200:
|
||||
fobj.seek(0)
|
||||
fobj.truncate(0)
|
||||
base = 0
|
||||
elif status not in (200, 206):
|
||||
text = ""
|
||||
try:
|
||||
text = (await resp.text())[:500]
|
||||
except Exception:
|
||||
pass
|
||||
if status in DOWNLOAD_PERMANENT_STATUS:
|
||||
raise _PermanentDownloadError(f"视频下载失败 ({status})：{text}")
|
||||
raise RuntimeError(f"视频下载失败 ({status})：{text}")
|
||||
|
||||
# 解析“完整文件总大小”，用于检测连接被提前关闭导致的截断。
|
||||
expected_total = _expected_total_size(resp, base)
|
||||
|
||||
async for chunk in resp.content.iter_chunked(chunk_size):
|
||||
check_interrupt()
|
||||
if chunk:
|
||||
fobj.write(chunk)
|
||||
if on_bytes:
|
||||
on_bytes(len(chunk))
|
||||
fobj.flush()
|
||||
return expected_total
|
||||
|
||||
|
||||
def _expected_total_size(resp: aiohttp.ClientResponse, base: int) -> int:
|
||||
"""根据响应头推断完整文件总字节数；无法判断时返回 -1。"""
|
||||
# Content-Range: bytes start-end/total → total 即完整大小
|
||||
cr = resp.headers.get("Content-Range", "")
|
||||
if "/" in cr:
|
||||
tail = cr.rsplit("/", 1)[-1].strip()
|
||||
if tail.isdigit():
|
||||
return int(tail)
|
||||
# Content-Length 是“本次响应体长度”，加上已续传的 base 即完整大小
|
||||
cl = resp.headers.get("Content-Length")
|
||||
if cl is not None and cl.isdigit():
|
||||
return base + int(cl)
|
||||
return -1
|
||||
|
||||
|
||||
async def download_video_to_file(
|
||||
session: aiohttp.ClientSession,
|
||||
url: str,
|
||||
save_path: str,
|
||||
*,
|
||||
headers: Optional[dict] = None,
|
||||
sock_read: float = DOWNLOAD_SOCK_READ_TIMEOUT,
|
||||
chunk_size: int = 1024 * 1024,
|
||||
label: str = "视频",
|
||||
on_bytes: Optional[Callable[[int], None]] = None,
|
||||
max_retries: Optional[int] = None,
|
||||
) -> str:
|
||||
"""把远程视频下载到 save_path，抗超时 + 断点续传 + 无限重试 + 可取消。
|
||||
|
||||
- 已生成成功的视频务必拿到手：默认 max_retries=None 表示对“可重试错误”
|
||||
（网络中断 / 超时 / 5xx / 连接失败）一直重试，直到成功。
|
||||
- 永久性错误（403/404/410 等）立即抛出，不做无意义重试。
|
||||
- 通过 HTTP Range 从已落盘的字节处续传，不重复下载。
|
||||
- 通过 check_interrupt() 全程响应 ComfyUI 取消。
|
||||
|
||||
返回 save_path。
|
||||
"""
|
||||
parent = os.path.dirname(save_path)
|
||||
if parent:
|
||||
os.makedirs(parent, exist_ok=True)
|
||||
|
||||
# 起始先清空目标文件：避免对调用方残留的旧文件做错误续传（续传只针对本次调用
|
||||
# 内部已写入的字节）。
|
||||
try:
|
||||
with open(save_path, "wb"):
|
||||
pass
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
attempt = 0
|
||||
while True:
|
||||
check_interrupt()
|
||||
# 仅对“本次调用已落盘”的字节做断点续传。
|
||||
resume_from = 0
|
||||
if os.path.isfile(save_path):
|
||||
try:
|
||||
resume_from = os.path.getsize(save_path)
|
||||
except OSError:
|
||||
resume_from = 0
|
||||
|
||||
# 已有部分文件 → 追加续写；否则新建。
|
||||
mode = "r+b" if resume_from > 0 else "wb"
|
||||
try:
|
||||
with open(save_path, mode) as f:
|
||||
if resume_from > 0:
|
||||
f.seek(0, os.SEEK_END)
|
||||
expected_total = await run_with_interrupt(
|
||||
_stream_once(
|
||||
session, url, f,
|
||||
headers=headers,
|
||||
resume_from=resume_from,
|
||||
sock_read=sock_read,
|
||||
chunk_size=chunk_size,
|
||||
on_bytes=on_bytes,
|
||||
)
|
||||
)
|
||||
# 走到这里说明本次 GET 的响应体已读尽、连接已关闭。
|
||||
size = os.path.getsize(save_path) if os.path.isfile(save_path) else 0
|
||||
if size <= 0:
|
||||
raise RuntimeError(f"{label}下载失败：保存后的文件为空。")
|
||||
# 连接被提前关闭（截断）：实际大小 < 服务端声明的完整大小 → 续传重试。
|
||||
if expected_total > 0 and size < expected_total:
|
||||
raise _IncompleteDownloadError(
|
||||
f"{label}下载不完整：{size}/{expected_total} 字节，将续传。"
|
||||
)
|
||||
return save_path
|
||||
|
||||
except InterruptProcessingException:
|
||||
raise
|
||||
except _PermanentDownloadError:
|
||||
raise
|
||||
except (_IncompleteDownloadError, aiohttp.ClientError, asyncio.TimeoutError, OSError) as e:
|
||||
attempt += 1
|
||||
if max_retries is not None and attempt > max_retries:
|
||||
raise RuntimeError(f"{label}下载失败（已重试 {max_retries} 次）：{e}") from None
|
||||
delay = min(DOWNLOAD_RETRY_BASE_DELAY * (2 ** (attempt - 1)), DOWNLOAD_RETRY_MAX_DELAY)
|
||||
done = "续传" if (os.path.isfile(save_path) and os.path.getsize(save_path) > 0) else "重连"
|
||||
print(f"[{label}] 下载中断（{type(e).__name__}），{delay:.0f}s 后{done}重试（第 {attempt} 次）...")
|
||||
await interruptible_sleep(delay)
|
||||
|
||||
|
||||
async def download_video_bytes(
|
||||
session: aiohttp.ClientSession,
|
||||
url: str,
|
||||
*,
|
||||
headers: Optional[dict] = None,
|
||||
label: str = "视频",
|
||||
**kwargs,
|
||||
) -> bytes:
|
||||
"""与 download_video_to_file 相同的健壮性，但返回内存中的 bytes。
|
||||
|
||||
内部仍落盘到临时文件以支持断点续传，读出后删除。
|
||||
"""
|
||||
import tempfile
|
||||
|
||||
fd, tmp_path = tempfile.mkstemp(suffix=".bin", prefix="o1key_dl_")
|
||||
os.close(fd)
|
||||
try:
|
||||
await download_video_to_file(
|
||||
session, url, tmp_path, headers=headers, label=label, **kwargs
|
||||
)
|
||||
with open(tmp_path, "rb") as f:
|
||||
return f.read()
|
||||
finally:
|
||||
try:
|
||||
os.remove(tmp_path)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
Reference in New Issue
Block a user