feat: 新增去背景/PSD分层导出节点，优化聊天面板与重启逻辑

- 新增 O1keyRemoveBackground 节点（基于 rembg CPU 推理） - 新增 O1keyColorRemoveBG 节点（颜色距离去背景，支持多模式） - 新增 O1keySavePSD 节点（手写 PSD 二进制，零外部依赖） - GPT Image 批量输出不同尺寸时自动 resize 对齐 - 聊天面板大幅增强（多模态/交互优化） - 重启按钮绕过 beforeunload 弹窗强制刷新 - 默认路由切换为 CF加速 - http_error 新增 system error 友好文案 Co-Authored-By: Claude Opus 4.7 <[email protected]>
2026-05-26 18:05:07 +08:00
parent c974df1b5e
commit 3f0f4099fb
11 changed files with 706 additions and 22 deletions
@@ -0,0 +1,174 @@
+"""
+颜色去背景工具模块
+基于颜色距离计算实现精确可控的背景移除，不依赖 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")
@@ -29,6 +29,7 @@ HTTP_ERROR_MESSAGES = {
 # 错误内容关键词 → 用户友好文案（优先于状态码匹配）
 ERROR_CONTENT_MESSAGES = {
    "The current model has a high load": "模型过载，请稍后重试！",
+    "system error": "系统错误，请稍后重试。",
 }

 # 可退避重试的状态码
@@ -0,0 +1,65 @@
+"""
+背景移除工具模块
+基于 rembg 库实现，支持 CPU 推理
+"""
+
+import numpy as np
+import torch
+from PIL import Image
+
+_session = None
+
+
+def _get_session():
+    """懒加载 rembg session，避免启动时加载模型"""
+    global _session
+    if _session is None:
+        try:
+            from rembg import new_session
+            _session = new_session("isnet-general-use")
+            print("[o1key] rembg 模型加载完成 (isnet-general-use)")
+        except ImportError:
+            raise RuntimeError(
+                "未安装 rembg，请执行: pip install rembg[cpu]>=2.0.50"
+            )
+    return _session
+
+
+def remove_background_pil(image: Image.Image) -> Image.Image:
+    """
+    移除 PIL Image 背景，返回 RGBA 图像（背景透明）
+    """
+    from rembg import remove
+    session = _get_session()
+    result = remove(image, session=session)
+    return result.convert("RGBA")
+
+
+def remove_background_tensor(tensor: torch.Tensor) -> torch.Tensor:
+    """
+    移除 ComfyUI IMAGE tensor 的背景
+    输入: [B, H, W, C] (3或4通道)
+    输出: [B, H, W, 4] RGBA tensor
+    """
+    from rembg import remove
+    session = _get_session()
+
+    results = []
+    batch_size = tensor.shape[0]
+
+    for i in range(batch_size):
+        frame = tensor[i]  # [H, W, C]
+        arr = (frame.cpu().numpy() * 255).clip(0, 255).astype(np.uint8)
+
+        if arr.shape[2] == 4:
+            pil_img = Image.fromarray(arr, mode="RGBA")
+        else:
+            pil_img = Image.fromarray(arr, mode="RGB")
+
+        result = remove(pil_img, session=session)
+        result_rgba = result.convert("RGBA")
+
+        result_arr = np.array(result_rgba).astype(np.float32) / 255.0
+        results.append(torch.from_numpy(result_arr))
+
+    return torch.stack(results, dim=0)