Update image and video workflow nodes

utils/http_error.py

@@ -28,12 +28,15 @@ HTTP_ERROR_MESSAGES = {
 
 # 错误内容关键词 → 用户友好文案（优先于状态码匹配）
 ERROR_CONTENT_MESSAGES = {
+    "Your request was rejected by the safety system": "请求被安全系统拦截：请调整提示词，避免敏感、违规、血腥、色情、仇恨、未成年人或真实人物等高风险内容。",
+    "safety system": "请求被安全系统拦截：请调整提示词，避免敏感、违规、血腥、色情、仇恨、未成年人或真实人物等高风险内容。",
+    "unexpected end of JSON input": "通常重试能解决；反复出现就降低分辨率、数量或换网络线路。",
     "The current model has a high load": "模型过载，请稍后重试！",
     "system error": "系统错误，请稍后重试。",
 }
 
 # 可退避重试的状态码
-RETRYABLE_STATUS_CODES = {429, 502, 503, 504}
+RETRYABLE_STATUS_CODES = {429, 502, 503, 504, 524}
 
 # 退避重试默认参数
 DEFAULT_MAX_RETRIES = 3
@@ -44,10 +47,15 @@ DEFAULT_BACKOFF_FACTOR = 2.0  # 指数退避因子
 
 def get_friendly_message(status_code: int, raw_message: str = "") -> str:
     """根据状态码/错误内容返回友好文案，未匹配则返回原始信息"""
+    if status_code == 524:
+        return "Gateway timed out while waiting for upstream image generation. Please retry, lower resolution/count, or switch network route."
     if raw_message:
+        raw_message_lower = raw_message.lower()
         for keyword, friendly_msg in ERROR_CONTENT_MESSAGES.items():
-            if keyword in raw_message:
+            if keyword.lower() in raw_message_lower:
                 return friendly_msg
+    if status_code == 500:
+        return "服务器返回 500：上游生成失败或服务端临时异常。请稍后重试；如果多次出现，请降低分辨率/数量，或调整提示词。"
     friendly = HTTP_ERROR_MESSAGES.get(status_code)
     if friendly:
         return friendly
@@ -88,7 +96,7 @@ async def async_request_with_retry(
     """
     带退避重试的 aiohttp 请求。
 
-    仅对 RETRYABLE_STATUS_CODES (429/503/504) 进行重试。
+    仅对 RETRYABLE_STATUS_CODES (429/502/503/504/524) 进行重试。
     超过最大重试次数后抛出友好 RuntimeError。
     成功时返回 response 对象（调用者需在 async with 外自行处理 body）。
 

utils/image_utils.py

@@ -5,7 +5,8 @@
 
 import base64
 from io import BytesIO
-from typing import List
+import json
+from typing import Callable, List, Tuple
 
 import numpy as np
 import torch
@@ -110,7 +111,130 @@ def encode_image_to_base64(image: Image.Image, format: str = "PNG") -> str:
     return base64.b64encode(img_bytes).decode('utf-8')
 
 
-_MAX_IMAGE_BYTES = 10 * 1024 * 1024  # 10MB base64 上限
+_MAX_REQUEST_BODY_BYTES = 50 * 1024 * 1024  # 50MB 请求体上限
+
+
+def _encode_image_to_base64_with_quality(image: Image.Image, quality: int) -> str:
+    buffered = BytesIO()
+    working = image
+    if working.mode != 'RGB':
+        working = working.convert('RGB')
+
+    working.save(
+        buffered,
+        format="JPEG",
+        quality=quality,
+        optimize=True,
+        subsampling=2,
+    )
+    return base64.b64encode(buffered.getvalue()).decode('utf-8')
+
+
+def encode_images_for_request_body_limit(
+    images: List[Image.Image],
+    build_body: Callable[[List[Tuple[str, str]]], dict],
+    max_body_bytes: int = _MAX_REQUEST_BODY_BYTES,
+) -> List[Tuple[str, str]]:
+    """
+    为请求体编码图片，并保证完整 JSON 请求体不超过 max_body_bytes。
+
+    策略：
+    - 先按原始 PNG 编码估算完整请求体；
+    - 若超过限制，改用 JPEG 质量压缩，逐步降低 quality；
+    - 全程不缩放图片尺寸。
+
+    Returns:
+        [(mime_type, base64), ...]
+    """
+    encoded = [("image/png", encode_image_to_base64(img, format="PNG")) for img in images]
+    body_size = len(json.dumps(build_body(encoded)).encode("utf-8"))
+    if body_size <= max_body_bytes:
+        return encoded
+
+    for quality in [95, 90, 85, 80, 75, 70, 65, 60, 55, 50, 45, 40, 35, 30, 25, 20, 15, 10, 5, 1]:
+        encoded = [
+            ("image/jpeg", _encode_image_to_base64_with_quality(img, quality))
+            for img in images
+        ]
+        body_size = len(json.dumps(build_body(encoded)).encode("utf-8"))
+        if body_size <= max_body_bytes:
+            print(
+                f"输入图片已通过 JPEG 质量压缩控制请求体积: "
+                f"quality={quality}, 请求体积={body_size / 1024 / 1024:.2f}MB "
+                f"(限制 {max_body_bytes / 1024 / 1024:.0f}MB)"
+            )
+            return encoded
+
+    raise ValueError(
+        f"请求体超过 {max_body_bytes / 1024 / 1024:.0f}MB，"
+        "即使压缩到最低图片质量仍无法满足限制；请减少参考图数量或输入图片内容复杂度"
+    )
+
+
+_MAX_IMAGE_BYTES = 10 * 1024 * 1024  # 10MB 单张图片上限
+
+
+def _encode_image_to_bytes(image: Image.Image, format: str = "PNG", quality: int = None) -> bytes:
+    buffered = BytesIO()
+    working = image
+
+    if format.upper() == "JPEG" and working.mode != 'RGB':
+        working = working.convert('RGB')
+    elif working.mode == 'RGBA':
+        working = working.convert('RGB')
+
+    save_kwargs = {"format": format}
+    if quality is not None:
+        save_kwargs.update({
+            "quality": quality,
+            "optimize": True,
+            "subsampling": 2,
+        })
+
+    working.save(buffered, **save_kwargs)
+    return buffered.getvalue()
+
+
+def encode_images_for_image_size_limit(
+    images: List[Image.Image],
+    max_image_bytes: int = _MAX_IMAGE_BYTES,
+) -> List[Tuple[str, str]]:
+    """
+    将图片编码为 base64，并保证每张编码前的图片文件体积不超过 max_image_bytes。
+
+    策略：
+    - 先尝试 PNG 原图尺寸编码；
+    - 单张超过限制时，改用 JPEG 质量压缩；
+    - 全程不缩放图片尺寸。
+
+    Returns:
+        [(mime_type, base64), ...]
+    """
+    encoded = []
+
+    for idx, img in enumerate(images, start=1):
+        png_bytes = _encode_image_to_bytes(img, format="PNG")
+        if len(png_bytes) <= max_image_bytes:
+            encoded.append(("image/png", base64.b64encode(png_bytes).decode('utf-8')))
+            continue
+
+        for quality in [95, 90, 85, 80, 75, 70, 65, 60, 55, 50, 45, 40, 35, 30, 25, 20, 15, 10, 5, 1]:
+            jpg_bytes = _encode_image_to_bytes(img, format="JPEG", quality=quality)
+            if len(jpg_bytes) <= max_image_bytes:
+                print(
+                    f"输入图片 {idx} 已通过 JPEG 质量压缩控制单图体积: "
+                    f"quality={quality}, 图片体积={len(jpg_bytes) / 1024 / 1024:.2f}MB "
+                    f"(限制 {max_image_bytes / 1024 / 1024:.0f}MB)，尺寸保持 {img.width}x{img.height}"
+                )
+                encoded.append(("image/jpeg", base64.b64encode(jpg_bytes).decode('utf-8')))
+                break
+        else:
+            raise ValueError(
+                f"输入图片 {idx} 超过 {max_image_bytes / 1024 / 1024:.0f}MB，"
+                "即使压缩到最低图片质量仍无法满足限制；请减少图片内容复杂度或手动处理图片"
+            )
+
+    return encoded
 
 
 def encode_image_to_base64_limited(
@@ -241,4 +365,4 @@ def parse_batch_prompts(prompt: str) -> List[str]:
     if not filtered_prompts:
         raise ValueError("批量提示词模式下，所有提示词都为空，请至少提供一个有效的提示词")
     
-    return filtered_prompts
+    return filtered_prompts

utils/rembg_utils.py

@@ -3,6 +3,7 @@
 基于 rembg 库实现，支持 CPU 推理
 """
 
+import os
 import numpy as np
 import torch
 from PIL import Image
@@ -14,6 +15,13 @@ def _get_session():
     """懒加载 rembg session，避免启动时加载模型"""
     global _session
     if _session is None:
+        try:
+            import folder_paths
+            models_dir = os.path.join(folder_paths.models_dir, "rembg")
+            os.makedirs(models_dir, exist_ok=True)
+            os.environ["U2NET_HOME"] = models_dir
+        except Exception:
+            pass
         try:
             from rembg import new_session
             _session = new_session("isnet-general-use")

utils/video_task.py

@@ -0,0 +1,181 @@
+import asyncio
+from typing import Any, Dict
+
+
+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",
+}
+
+
+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)
+    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))