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.
472 lines
16 KiB
Python
472 lines
16 KiB
Python
import asyncio
|
||
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:
|
||
from comfy.model_management import processing_interrupted, InterruptProcessingException
|
||
INTERRUPT_AVAILABLE = True
|
||
except Exception:
|
||
INTERRUPT_AVAILABLE = False
|
||
processing_interrupted = lambda: False
|
||
|
||
class InterruptProcessingException(Exception):
|
||
pass
|
||
|
||
|
||
SUCCESS_STATUSES = {"succeed", "succeeded", "success", "completed", "done", "finished"}
|
||
FAILURE_STATUSES = {
|
||
"fail",
|
||
"failed",
|
||
"failure",
|
||
"error",
|
||
"expired",
|
||
"timeout",
|
||
"timed_out",
|
||
"cancel",
|
||
"canceled",
|
||
"cancelled",
|
||
"rejected",
|
||
}
|
||
|
||
|
||
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()
|
||
|
||
|
||
async def interruptible_sleep(seconds: float, step: float = 0.2) -> None:
|
||
elapsed = 0.0
|
||
while elapsed < seconds:
|
||
check_interrupt()
|
||
delay = min(step, seconds - elapsed)
|
||
await asyncio.sleep(delay)
|
||
elapsed += delay
|
||
check_interrupt()
|
||
|
||
|
||
async def run_with_interrupt(coro, step: float = 0.2):
|
||
task = asyncio.ensure_future(coro)
|
||
try:
|
||
while not task.done():
|
||
check_interrupt()
|
||
await asyncio.wait({task}, timeout=step)
|
||
check_interrupt()
|
||
return await task
|
||
except InterruptProcessingException:
|
||
task.cancel()
|
||
try:
|
||
await task
|
||
except BaseException:
|
||
pass
|
||
raise
|
||
|
||
|
||
def _as_dict(value: Any) -> Dict[str, Any]:
|
||
return value if isinstance(value, dict) else {}
|
||
|
||
|
||
def _nested_payloads(payload: Dict[str, Any]):
|
||
root = _as_dict(payload)
|
||
data = _as_dict(root.get("data"))
|
||
inner = _as_dict(data.get("data"))
|
||
return root, data, inner
|
||
|
||
|
||
def extract_status(payload: Dict[str, Any]) -> str:
|
||
root, data, inner = _nested_payloads(payload)
|
||
keys = ("status", "task_status", "state", "task_state")
|
||
statuses = []
|
||
for source in (data, inner, root):
|
||
for key in keys:
|
||
value = source.get(key)
|
||
if value is not None and str(value).strip():
|
||
statuses.append(str(value).strip().lower())
|
||
for status in statuses:
|
||
if status in FAILURE_STATUSES or any(
|
||
token in status for token in ("fail", "error", "reject", "timeout", "cancel")
|
||
):
|
||
return status
|
||
for status in statuses:
|
||
if status in SUCCESS_STATUSES:
|
||
return status
|
||
return statuses[0] if statuses else ""
|
||
|
||
|
||
def extract_progress(payload: Dict[str, Any]) -> int:
|
||
root, data, inner = _nested_payloads(payload)
|
||
for source in (data, inner, root):
|
||
value = source.get("progress")
|
||
if value is None:
|
||
continue
|
||
try:
|
||
return max(0, min(100, int(float(str(value).strip().rstrip("%")))))
|
||
except (TypeError, ValueError):
|
||
return 0
|
||
return 0
|
||
|
||
|
||
def extract_error_message(payload: Dict[str, Any], default: str = "未知错误") -> str:
|
||
root, data, inner = _nested_payloads(payload)
|
||
keys = (
|
||
"fail_reason",
|
||
"failure_reason",
|
||
"task_status_msg",
|
||
"status_msg",
|
||
"error_message",
|
||
"message",
|
||
"msg",
|
||
"reason",
|
||
"detail",
|
||
"details",
|
||
)
|
||
for source in (data, inner, root):
|
||
error = source.get("error")
|
||
if isinstance(error, dict):
|
||
for key in ("message", "msg", "detail", "reason", "code"):
|
||
value = error.get(key)
|
||
if value:
|
||
return str(value)
|
||
elif error:
|
||
return str(error)
|
||
|
||
for key in keys:
|
||
value = source.get(key)
|
||
if value:
|
||
return str(value)
|
||
return default
|
||
|
||
|
||
def extract_video_url(payload: Dict[str, Any]) -> str | None:
|
||
root, data, inner = _nested_payloads(payload)
|
||
for source in (data, inner, root):
|
||
for key in ("video_url", "result_url", "url", "download_url"):
|
||
value = source.get(key)
|
||
if value:
|
||
return str(value)
|
||
|
||
result = _as_dict(source.get("result"))
|
||
for key in ("video_url", "result_url", "url", "download_url"):
|
||
value = result.get(key)
|
||
if value:
|
||
return str(value)
|
||
|
||
content = _as_dict(source.get("content"))
|
||
value = content.get("video_url") or content.get("url")
|
||
if value:
|
||
return str(value)
|
||
|
||
task_result = _as_dict(source.get("task_result"))
|
||
videos = task_result.get("videos")
|
||
if isinstance(videos, list) and videos:
|
||
first = _as_dict(videos[0])
|
||
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
|
||
|
||
|
||
def is_success_status(status: str) -> bool:
|
||
return status in SUCCESS_STATUSES
|
||
|
||
|
||
def is_failure_status(status: str, payload: Dict[str, Any] | None = None) -> bool:
|
||
if status in FAILURE_STATUSES:
|
||
return True
|
||
if any(token in status for token in ("fail", "error", "reject", "timeout", "cancel")):
|
||
return True
|
||
if payload is None:
|
||
return False
|
||
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
|