Jony
ba920f2b66
Replace the prior release tree with the current plugin, frontend, tests, and documentation. Document retired node IDs and the public Gitea update source.
260 lines
9.0 KiB
Python
260 lines
9.0 KiB
Python
"""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
|
||
|
||
|
||
_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()
|
||
api_base_url = (base_url or get_api_base_url()).rstrip("/")
|
||
upload_url = f"{api_base_url}{_UPLOAD_ENDPOINT}"
|
||
last_error: Optional[BaseException] = None
|
||
|
||
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,
|
||
)
|
||
|
||
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 upload_image(pil_image, base_url: Optional[str] = None) -> str:
|
||
"""
|
||
接受 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"
|
||
|
||
return await _upload_file(data, filename, "image/png", base_url=base_url)
|
||
|
||
|
||
async def upload_video(video, base_url: Optional[str] = None) -> str:
|
||
"""
|
||
接受 ComfyUI VIDEO 对象，上传并返回临时公网 URL。
|
||
支持 mp4 / mov 格式。
|
||
"""
|
||
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)
|
||
data = source.read()
|
||
ext = "mp4"
|
||
else:
|
||
video_path = source
|
||
if not video_path or not os.path.isfile(video_path):
|
||
raise ValueError(f"无法获取参考视频文件路径（当前路径：{video_path}）")
|
||
ext = os.path.splitext(video_path)[1].lower().lstrip(".")
|
||
if ext not in ("mp4", "mov"):
|
||
raise ValueError(f"参考视频格式须为 mp4 或 mov，当前为 .{ext}")
|
||
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}"
|
||
|
||
return await _upload_file(data, filename, content_type, base_url=base_url)
|
||
|
||
|
||
async def upload_audio(audio, base_url: Optional[str] = None) -> str:
|
||
"""
|
||
接受 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
|
||
|
||
waveform = audio["waveform"] # shape: [B, C, N] or [C, N]
|
||
sample_rate = int(audio["sample_rate"])
|
||
|
||
if waveform.dim() == 3:
|
||
waveform = waveform[0]
|
||
|
||
wav_np = waveform.cpu().numpy()
|
||
if wav_np.ndim == 2:
|
||
wav_np = wav_np.mean(axis=0)
|
||
wav_np = np.clip(wav_np, -1.0, 1.0)
|
||
pcm = (wav_np * 32767).astype(np.int16)
|
||
|
||
num_samples = len(pcm)
|
||
num_channels = 1
|
||
bits_per_sample = 16
|
||
byte_rate = sample_rate * num_channels * bits_per_sample // 8
|
||
block_align = num_channels * bits_per_sample // 8
|
||
data_size = num_samples * block_align
|
||
|
||
buf = io.BytesIO()
|
||
buf.write(b"RIFF")
|
||
buf.write(struct.pack("<I", 36 + data_size))
|
||
buf.write(b"WAVE")
|
||
buf.write(b"fmt ")
|
||
buf.write(struct.pack("<IHHIIHH", 16, 1, num_channels, sample_rate,
|
||
byte_rate, block_align, bits_per_sample))
|
||
buf.write(b"data")
|
||
buf.write(struct.pack("<I", data_size))
|
||
buf.write(pcm.tobytes())
|
||
|
||
data = buf.getvalue()
|
||
filename = f"{uuid.uuid4()}.wav"
|
||
|
||
return await _upload_file(data, filename, "audio/wav", base_url=base_url)
|