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
+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)
|
||||
|
||||
Reference in New Issue
Block a user