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]>
This commit is contained in:
o1key
+11
-2
@@ -22,6 +22,9 @@ logging.getLogger().addFilter(_seeder_filter)
|
||||
|
||||
from .nodes import NanoBananaPro, BatchNanoBananaPro, GoogleGemini, LoadFile, ImageStitchPro, BatchCleanMetadata, VideoPreview, GoogleVeo, FluxImageEdit, UniversalLLMChat, KlingVideo, KlingFirstLastFrame, KlingMotionControlTest, AspectRatioPreset, BatchImagesO1key, Seedance, SeedanceMultiModal, StreamPreview, DoubaoImage, O1keyGPTImage, O1keyGrokImage, KVideoFirstLast, KVideoImage2Video
|
||||
from .nodes import K3Video, K3VideoFirstLast, K3MotionControl, K3MotionVideoCheck, NanoBananaV2, NanoBananaV2Batch, SaveImageFormat
|
||||
from .nodes import O1keySavePSD
|
||||
from .nodes import O1keyRemoveBackground
|
||||
from .nodes import O1keyColorRemoveBG
|
||||
|
||||
# 报错弹框友好文案（不修改原节点代码，仅在外层统一处理）
|
||||
_MSG_TIMEOUT = "API 请求超时，请稍后重试或检查网络。"
|
||||
@@ -92,6 +95,9 @@ NODE_CLASS_MAPPINGS = {
|
||||
"NanoBananaV2": NanoBananaV2,
|
||||
"NanoBananaV2Batch": NanoBananaV2Batch,
|
||||
"SaveImageFormat": SaveImageFormat,
|
||||
"O1keySavePSD": O1keySavePSD,
|
||||
"O1keyRemoveBackground": O1keyRemoveBackground,
|
||||
"O1keyColorRemoveBG": O1keyColorRemoveBG,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
@@ -127,6 +133,9 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"NanoBananaV2": "Nano Banana V2",
|
||||
"NanoBananaV2Batch": "Nano Banana V2（批量）",
|
||||
"SaveImageFormat": "保存图像（格式转换）",
|
||||
"O1keySavePSD": "保存 PSD（分层）",
|
||||
"O1keyRemoveBackground": "去背景（rembg）",
|
||||
"O1keyColorRemoveBG": "颜色去背景",
|
||||
}
|
||||
|
||||
WEB_DIRECTORY = "./web"
|
||||
@@ -181,7 +190,7 @@ try:
|
||||
test_key = data.get("api_key", "").strip()
|
||||
if not test_key:
|
||||
return web.json_response({"valid": False, "error": "密钥不能为空"})
|
||||
base_url = NETWORK_ROUTES.get("全球加速", "https://api.o1key.cn")
|
||||
base_url = NETWORK_ROUTES.get("CF加速", "https://cf-api.o1key.com")
|
||||
url = f"{base_url}/v1/models"
|
||||
headers = {"Authorization": f"Bearer {test_key}"}
|
||||
try:
|
||||
@@ -500,7 +509,7 @@ try:
|
||||
if not api_key:
|
||||
return web.json_response({"error": "未配置 API Key"}, status=401)
|
||||
|
||||
route = data.get("route", "全球加速")
|
||||
route = data.get("route", "CF加速")
|
||||
base_url = NETWORK_ROUTES.get(route, "https://cf-api.o1key.com")
|
||||
model = data.get("model", "gpt-5.5")
|
||||
messages = data.get("messages", [])
|
||||
|
||||
@@ -173,7 +173,20 @@ class GptImageClient:
|
||||
arr = np.array(img.convert("RGBA")).astype(np.float32) / 255.0
|
||||
tensors.append(torch.from_numpy(arr))
|
||||
|
||||
return torch.stack(tensors, dim=0) # [B, H, W, 4]
|
||||
# 批量模式下 API 可能返回不同尺寸，统一 resize 到最大尺寸
|
||||
max_h = max(t.shape[0] for t in tensors)
|
||||
max_w = max(t.shape[1] for t in tensors)
|
||||
aligned = []
|
||||
for t in tensors:
|
||||
if t.shape[0] != max_h or t.shape[1] != max_w:
|
||||
t = t.permute(2, 0, 1).unsqueeze(0) # [1, C, H, W]
|
||||
t = torch.nn.functional.interpolate(
|
||||
t, size=(max_h, max_w), mode="bilinear", align_corners=False
|
||||
)
|
||||
t = t.squeeze(0).permute(1, 2, 0) # [H, W, C]
|
||||
aligned.append(t)
|
||||
|
||||
return torch.stack(aligned, dim=0) # [B, H, W, 4]
|
||||
|
||||
# ── 响应解析（通用） ─────────────────────────────────────────────────────
|
||||
|
||||
|
||||
+4
-1
@@ -28,5 +28,8 @@ from .K3_video import K3Video
|
||||
from .K3_video_firstlast import K3VideoFirstLast
|
||||
from .K3_motion_control import K3MotionControl, K3MotionVideoCheck
|
||||
from .save_image_format import SaveImageFormat
|
||||
from .save_psd import O1keySavePSD
|
||||
from .remove_bg import O1keyRemoveBackground
|
||||
from .color_remove_bg import O1keyColorRemoveBG
|
||||
|
||||
__all__ = ['NanoBananaV2', 'NanoBananaV2Batch', 'NanoBanana', 'BatchNanoBananaPro', 'GoogleGemini', 'LoadFile', 'ImageStitchPro', 'BatchCleanMetadata', 'VideoPreview', 'KlingVideo', 'KlingFirstLastFrame', 'KlingMotionControlTest', 'AspectRatioPreset', 'GoogleVeo', 'FluxImageEdit', 'UniversalLLMChat', 'BatchImagesO1key', 'Seedance', 'SeedanceMultiModal', 'StreamPreview', 'DoubaoImage', 'O1keyGPTImage', 'O1keyGrokImage', 'KVideoFirstLast', 'KVideoImage2Video', 'K3Video', 'K3VideoFirstLast', 'K3MotionControl', 'K3MotionVideoCheck', 'AsyncImageGenerator', 'BatchAsyncImageGenerator', 'SaveImageFormat']
|
||||
__all__ = ['NanoBananaV2', 'NanoBananaV2Batch', 'NanoBanana', 'BatchNanoBananaPro', 'GoogleGemini', 'LoadFile', 'ImageStitchPro', 'BatchCleanMetadata', 'VideoPreview', 'KlingVideo', 'KlingFirstLastFrame', 'KlingMotionControlTest', 'AspectRatioPreset', 'GoogleVeo', 'FluxImageEdit', 'UniversalLLMChat', 'BatchImagesO1key', 'Seedance', 'SeedanceMultiModal', 'StreamPreview', 'DoubaoImage', 'O1keyGPTImage', 'O1keyGrokImage', 'KVideoFirstLast', 'KVideoImage2Video', 'K3Video', 'K3VideoFirstLast', 'K3MotionControl', 'K3MotionVideoCheck', 'AsyncImageGenerator', 'BatchAsyncImageGenerator', 'SaveImageFormat', 'O1keySavePSD', 'O1keyRemoveBackground', 'O1keyColorRemoveBG']
|
||||
|
||||
@@ -0,0 +1,93 @@
|
||||
"""
|
||||
o1key 颜色去背景节点
|
||||
基于颜色距离计算，精确可控，不依赖 AI 模型
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
|
||||
class O1keyColorRemoveBG:
|
||||
"""
|
||||
颜色去背景 - 精确移除纯色背景
|
||||
|
||||
模式说明：
|
||||
- 白色(white): 移除白色背景，适合大多数场景
|
||||
- 白色保护(white-preserve): 移除白底但保护浅色前景物体
|
||||
- 自动检测(corner): 自动采样四角颜色作为背景色
|
||||
- 指定颜色(color): 手动指定要移除的背景颜色
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE",),
|
||||
"模式": (["白色", "白色保护", "自动检测", "指定颜色"], {
|
||||
"default": "白色",
|
||||
}),
|
||||
"容差": ("FLOAT", {
|
||||
"default": 8.0,
|
||||
"min": 0.0,
|
||||
"max": 100.0,
|
||||
"step": 1.0,
|
||||
"tooltip": "颜色距离阈值，越大去除范围越广",
|
||||
}),
|
||||
"羽化": ("FLOAT", {
|
||||
"default": 45.0,
|
||||
"min": 0.0,
|
||||
"max": 200.0,
|
||||
"step": 1.0,
|
||||
"tooltip": "边缘过渡范围，越大边缘越柔和",
|
||||
}),
|
||||
},
|
||||
"optional": {
|
||||
"背景色R": ("INT", {"default": 255, "min": 0, "max": 255}),
|
||||
"背景色G": ("INT", {"default": 255, "min": 0, "max": 255}),
|
||||
"背景色B": ("INT", {"default": 255, "min": 0, "max": 255}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("RGBA图像",)
|
||||
FUNCTION = "remove_bg"
|
||||
CATEGORY = "o1key/image"
|
||||
|
||||
_MODE_MAP = {
|
||||
"白色": "white",
|
||||
"白色保护": "white-preserve",
|
||||
"自动检测": "corner",
|
||||
"指定颜色": "color",
|
||||
}
|
||||
|
||||
def remove_bg(self, image, 模式, 容差, 羽化, 背景色R=255, 背景色G=255, 背景色B=255):
|
||||
from ..utils.color_key import remove_background
|
||||
|
||||
mode = self._MODE_MAP.get(模式, "white")
|
||||
bg_color = (背景色R, 背景色G, 背景色B)
|
||||
|
||||
batch_size = image.shape[0]
|
||||
results = []
|
||||
|
||||
for i in range(batch_size):
|
||||
frame = image[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_background(
|
||||
pil_img, mode=mode, bg_color=bg_color,
|
||||
tolerance=容差, feather=羽化,
|
||||
)
|
||||
|
||||
result_arr = np.array(result.convert("RGBA")).astype(np.float32) / 255.0
|
||||
results.append(torch.from_numpy(result_arr))
|
||||
|
||||
output = torch.stack(results, dim=0)
|
||||
print(f"[o1key 颜色去背景] 模式={模式}, 容差={容差}, 羽化={羽化}, "
|
||||
f"处理 {batch_size} 张")
|
||||
return (output,)
|
||||
@@ -0,0 +1,36 @@
|
||||
"""
|
||||
o1key 去背景节点
|
||||
基于 rembg 实现，支持 CPU 推理
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
|
||||
class O1keyRemoveBackground:
|
||||
"""
|
||||
移除图像背景，输出 RGBA 透明图层
|
||||
|
||||
基于 rembg (ISNet-General-Use) 模型，支持 CPU 推理。
|
||||
首次运行会自动下载模型（约 170MB）。
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("RGBA图像",)
|
||||
FUNCTION = "remove_bg"
|
||||
CATEGORY = "o1key/image"
|
||||
|
||||
def remove_bg(self, image):
|
||||
from ..utils.rembg_utils import remove_background_tensor
|
||||
print("[o1key 去背景] 正在处理...")
|
||||
result = remove_background_tensor(image)
|
||||
print(f"[o1key 去背景] 完成，输出 {result.shape[0]} 张 RGBA")
|
||||
return (result,)
|
||||
@@ -0,0 +1,237 @@
|
||||
"""
|
||||
o1key SavePSD 节点
|
||||
将多个 IMAGE 图层合成为分层 PSD 文件
|
||||
手写 PSD 二进制格式，零外部依赖（仅 numpy + Pillow）
|
||||
"""
|
||||
|
||||
import os
|
||||
import struct
|
||||
import time
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
import folder_paths
|
||||
|
||||
|
||||
def _pad_even(data: bytes) -> bytes:
|
||||
if len(data) % 2:
|
||||
return data + b"\x00"
|
||||
return data
|
||||
|
||||
|
||||
def _pad4(data: bytes) -> bytes:
|
||||
return data + (b"\x00" * ((4 - (len(data) % 4)) % 4))
|
||||
|
||||
|
||||
def _pascal_name(name: str) -> bytes:
|
||||
raw = name.encode("macroman", errors="replace")[:255]
|
||||
data = bytes([len(raw)]) + raw
|
||||
return _pad4(data)
|
||||
|
||||
|
||||
def _unicode_name_block(name: str) -> bytes:
|
||||
payload = struct.pack(">I", len(name)) + name.encode("utf-16be")
|
||||
block = b"8BIM" + b"luni" + struct.pack(">I", len(payload)) + _pad_even(payload)
|
||||
return block
|
||||
|
||||
|
||||
def _layer_extra_data(name: str) -> bytes:
|
||||
data = b""
|
||||
data += struct.pack(">I", 0) # layer mask data length
|
||||
data += struct.pack(">I", 0) # layer blending ranges length
|
||||
data += _pascal_name(name)
|
||||
data += _unicode_name_block(name)
|
||||
return data
|
||||
|
||||
|
||||
def _alpha_bbox(rgba_arr: np.ndarray):
|
||||
"""找到 RGBA 数组中非透明区域的 bounding box。"""
|
||||
alpha = rgba_arr[:, :, 3]
|
||||
rows = np.any(alpha > 0, axis=1)
|
||||
cols = np.any(alpha > 0, axis=0)
|
||||
if not rows.any():
|
||||
return None
|
||||
top = int(np.argmax(rows))
|
||||
bottom = int(len(rows) - np.argmax(rows[::-1]))
|
||||
left = int(np.argmax(cols))
|
||||
right = int(len(cols) - np.argmax(cols[::-1]))
|
||||
return top, left, bottom, right
|
||||
|
||||
|
||||
def write_psd(filepath: str, layers: list, canvas_w: int, canvas_h: int):
|
||||
"""
|
||||
写入 PSD 文件。
|
||||
|
||||
layers: [(name, rgba_array), ...] 从底到顶排列
|
||||
rgba_array: numpy uint8 [H, W, 4]
|
||||
"""
|
||||
records = []
|
||||
channel_data_blocks = []
|
||||
layers_top_to_bottom = list(reversed(layers))
|
||||
|
||||
for name, rgba in layers_top_to_bottom:
|
||||
bbox = _alpha_bbox(rgba)
|
||||
if not bbox:
|
||||
continue
|
||||
top, left, bottom, right = bbox
|
||||
cropped = rgba[top:bottom, left:right]
|
||||
|
||||
# PLACEHOLDER_CHANNELS
|
||||
|
||||
channels = [
|
||||
(0, cropped[:, :, 0].tobytes(order="C")),
|
||||
(1, cropped[:, :, 1].tobytes(order="C")),
|
||||
(2, cropped[:, :, 2].tobytes(order="C")),
|
||||
(-1, cropped[:, :, 3].tobytes(order="C")),
|
||||
]
|
||||
channel_info = b""
|
||||
data_block = b""
|
||||
for channel_id, data in channels:
|
||||
channel_info += struct.pack(">hI", channel_id, 2 + len(data))
|
||||
data_block += struct.pack(">H", 0) + data # raw compression
|
||||
|
||||
extra = _layer_extra_data(name)
|
||||
record = b""
|
||||
record += struct.pack(">iiii", top, left, bottom, right)
|
||||
record += struct.pack(">H", len(channels))
|
||||
record += channel_info
|
||||
record += b"8BIM" + b"norm"
|
||||
record += bytes([255, 0, 0, 0]) # opacity=255, clipping, flags, filler
|
||||
record += struct.pack(">I", len(extra)) + extra
|
||||
records.append(record)
|
||||
channel_data_blocks.append(data_block)
|
||||
|
||||
if not records:
|
||||
raise ValueError("所有图层均为空（完全透明），无法生成 PSD")
|
||||
|
||||
# Layer and Mask Information
|
||||
layer_info = struct.pack(">h", len(records))
|
||||
layer_info += b"".join(records) + b"".join(channel_data_blocks)
|
||||
layer_info = _pad_even(layer_info)
|
||||
layer_info_block = struct.pack(">I", len(layer_info)) + layer_info
|
||||
global_mask = struct.pack(">I", 0)
|
||||
layer_mask_payload = layer_info_block + global_mask
|
||||
layer_and_mask = struct.pack(">I", len(layer_mask_payload)) + layer_mask_payload
|
||||
|
||||
# PLACEHOLDER_COMPOSITE
|
||||
|
||||
# Composite preview (flattened image for compatibility)
|
||||
comp = Image.new("RGBA", (canvas_w, canvas_h), (255, 255, 255, 255))
|
||||
for name, rgba in layers:
|
||||
layer_img = Image.fromarray(rgba, "RGBA")
|
||||
comp.alpha_composite(layer_img)
|
||||
comp_rgb = np.asarray(comp.convert("RGB"), dtype=np.uint8)
|
||||
composite_data = (
|
||||
struct.pack(">H", 0)
|
||||
+ comp_rgb[:, :, 0].tobytes(order="C")
|
||||
+ comp_rgb[:, :, 1].tobytes(order="C")
|
||||
+ comp_rgb[:, :, 2].tobytes(order="C")
|
||||
)
|
||||
|
||||
# Write PSD file
|
||||
with open(filepath, "wb") as f:
|
||||
# Header
|
||||
f.write(b"8BPS")
|
||||
f.write(struct.pack(">H", 1)) # version
|
||||
f.write(b"\x00" * 6) # reserved
|
||||
f.write(struct.pack(">HIIHH", 3, canvas_h, canvas_w, 8, 3))
|
||||
# Color Mode Data
|
||||
f.write(struct.pack(">I", 0))
|
||||
# Image Resources
|
||||
f.write(struct.pack(">I", 0))
|
||||
# Layer and Mask
|
||||
f.write(layer_and_mask)
|
||||
# Composite Image Data
|
||||
f.write(composite_data)
|
||||
|
||||
|
||||
# PLACEHOLDER_NODE
|
||||
|
||||
class O1keySavePSD:
|
||||
"""
|
||||
将多个 IMAGE 输入合成为分层 PSD 文件
|
||||
|
||||
每个输入作为独立图层，支持 RGBA 透明通道。
|
||||
图层从下到上排列（图层1在最底部）。
|
||||
使用 bbox 裁剪优化文件大小，包含合成预览层。
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"批次图像": ("IMAGE", {
|
||||
"tooltip": "批次图像输入，每张图自动作为独立图层（支持RGBA透明）",
|
||||
}),
|
||||
},
|
||||
"optional": {
|
||||
"图层名称": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": True,
|
||||
"tooltip": "每行一个图层名称，与图层顺序对应。留空则自动命名。",
|
||||
}),
|
||||
"文件名前缀": ("STRING", {
|
||||
"default": "o1key_layers",
|
||||
"tooltip": "输出 PSD 文件名前缀",
|
||||
}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("文件路径",)
|
||||
FUNCTION = "save_psd"
|
||||
CATEGORY = "o1key/image"
|
||||
OUTPUT_NODE = True
|
||||
|
||||
def save_psd(self, 批次图像, 图层名称: str = "", 文件名前缀: str = "o1key_layers", **kwargs):
|
||||
# 将批次 tensor [B, H, W, C] 拆为单张列表
|
||||
if 批次图像.dim() == 3:
|
||||
layer_tensors = [批次图像]
|
||||
else:
|
||||
layer_tensors = [批次图像[i] for i in range(批次图像.shape[0])]
|
||||
|
||||
names = [n.strip() for n in 图层名称.split("\n") if n.strip()]
|
||||
|
||||
# 确定画布尺寸
|
||||
max_h, max_w = 0, 0
|
||||
for t in layer_tensors:
|
||||
h, w = t.shape[0], t.shape[1]
|
||||
max_h = max(max_h, h)
|
||||
max_w = max(max_w, w)
|
||||
|
||||
# 转换为 [(name, rgba_array), ...] 格式
|
||||
layers = []
|
||||
for idx, tensor in enumerate(layer_tensors):
|
||||
arr = (tensor.cpu().numpy() * 255).clip(0, 255).astype(np.uint8)
|
||||
h, w = arr.shape[0], arr.shape[1]
|
||||
channels = arr.shape[2] if arr.ndim == 3 else 1
|
||||
|
||||
if channels == 3:
|
||||
rgba = np.zeros((max_h, max_w, 4), dtype=np.uint8)
|
||||
rgba[:h, :w, :3] = arr
|
||||
rgba[:h, :w, 3] = 255
|
||||
elif channels == 4:
|
||||
rgba = np.zeros((max_h, max_w, 4), dtype=np.uint8)
|
||||
rgba[:h, :w] = arr
|
||||
else:
|
||||
rgba = np.zeros((max_h, max_w, 4), dtype=np.uint8)
|
||||
rgba[:h, :w, 0] = rgba[:h, :w, 1] = rgba[:h, :w, 2] = arr[:, :, 0] if arr.ndim == 3 else arr
|
||||
rgba[:h, :w, 3] = 255
|
||||
|
||||
name = names[idx] if idx < len(names) else f"图层 {idx + 1}"
|
||||
layers.append((name, rgba))
|
||||
print(f"[o1key SavePSD] 图层 '{name}': {w}×{h}")
|
||||
|
||||
# 写入 PSD
|
||||
output_dir = folder_paths.get_output_directory()
|
||||
timestamp = time.strftime("%Y%m%d_%H%M%S")
|
||||
filename = f"{文件名前缀}_{timestamp}.psd"
|
||||
filepath = os.path.join(output_dir, filename)
|
||||
|
||||
write_psd(filepath, layers, max_w, max_h)
|
||||
|
||||
size_kb = os.path.getsize(filepath) / 1024
|
||||
print(f"[o1key SavePSD] 完成: {filepath} ({size_kb:.0f}KB, "
|
||||
f"{len(layers)} 层, {max_w}×{max_h})")
|
||||
return (filepath,)
|
||||
@@ -1,3 +1,4 @@
|
||||
aiohttp>=3.9.0
|
||||
Pillow>=10.0.0
|
||||
requests>=2.31.0
|
||||
rembg[cpu]>=2.0.50
|
||||
|
||||
@@ -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)
|
||||
+57
-5
@@ -182,6 +182,22 @@ const CSS = `
|
||||
.o1k-conv-item:hover .conv-rename,.o1k-conv-item:hover .conv-del{opacity:1}
|
||||
.o1k-conv-item .conv-rename:hover{background:rgba(126,184,247,.2);color:#7eb8f7}
|
||||
.o1k-conv-item .conv-del:hover{background:rgba(220,60,60,.6);color:#fff}
|
||||
.o1k-conv-item .conv-rename-input{flex:1;min-width:0;padding:2px 6px;border:1px solid rgba(126,184,247,.5);border-radius:4px;background:rgba(0,0,0,.3);color:#ddd;font-size:12px;outline:none}
|
||||
.o1k-conv-item .conv-rename-input:focus{border-color:#7eb8f7}
|
||||
.o1k-conv-confirm{display:flex;align-items:center;gap:6px;flex:1;min-width:0}
|
||||
.o1k-conv-confirm span{font-size:11px;color:#ccc;white-space:nowrap}
|
||||
.o1k-conv-confirm button{padding:2px 8px;border:none;border-radius:4px;font-size:11px;cursor:pointer;transition:all .12s}
|
||||
.o1k-conv-confirm .confirm-yes{background:rgba(220,60,60,.7);color:#fff}
|
||||
.o1k-conv-confirm .confirm-yes:hover{background:rgba(220,60,60,.9)}
|
||||
.o1k-conv-confirm .confirm-no{background:rgba(255,255,255,.1);color:#aaa}
|
||||
.o1k-conv-confirm .confirm-no:hover{background:rgba(255,255,255,.15);color:#ddd}
|
||||
.o1k-msg-del-confirm{display:flex;align-items:center;gap:6px;padding:4px 8px;margin-top:4px;border-radius:6px;background:rgba(220,60,60,.08);border:1px solid rgba(220,60,60,.2)}
|
||||
.o1k-msg-del-confirm span{font-size:11px;color:#ccc}
|
||||
.o1k-msg-del-confirm button{padding:2px 8px;border:none;border-radius:4px;font-size:11px;cursor:pointer;transition:all .12s}
|
||||
.o1k-msg-del-confirm .confirm-yes{background:rgba(220,60,60,.7);color:#fff}
|
||||
.o1k-msg-del-confirm .confirm-yes:hover{background:rgba(220,60,60,.9)}
|
||||
.o1k-msg-del-confirm .confirm-no{background:rgba(255,255,255,.1);color:#aaa}
|
||||
.o1k-msg-del-confirm .confirm-no:hover{background:rgba(255,255,255,.15);color:#ddd}
|
||||
#o1key-chat-input-area{flex-shrink:0;padding:10px 14px 14px;background:transparent}
|
||||
#o1key-chat-previews{display:flex;gap:6px;padding:0 0 8px;flex-wrap:wrap}
|
||||
#o1key-chat-previews .preview-thumb{position:relative;width:40px;height:40px;border-radius:6px;overflow:hidden;border:1px solid rgba(255,255,255,.1)}
|
||||
@@ -394,35 +410,71 @@ function renderHistory() {
|
||||
const id = btn.dataset.id;
|
||||
const conv = conversations.find(c => c.id === id);
|
||||
if (!conv) return;
|
||||
const newTitle = prompt("重命名对话", conv.title || "");
|
||||
if (newTitle === null || !newTitle.trim()) return;
|
||||
conv.title = newTitle.trim();
|
||||
const item = btn.closest(".o1k-conv-item");
|
||||
const infoEl = item.querySelector(".conv-info");
|
||||
const oldHtml = infoEl.innerHTML;
|
||||
infoEl.innerHTML = `<input class="conv-rename-input" value="${escapeHtml(conv.title || "")}" />`;
|
||||
const input = infoEl.querySelector(".conv-rename-input");
|
||||
input.focus();
|
||||
input.select();
|
||||
const commit = () => {
|
||||
const val = input.value.trim();
|
||||
if (val && val !== conv.title) {
|
||||
conv.title = val;
|
||||
conv.updatedAt = Date.now();
|
||||
saveConversations();
|
||||
}
|
||||
renderHistory();
|
||||
};
|
||||
input.addEventListener("keydown", (ev) => {
|
||||
if (ev.key === "Enter") { ev.preventDefault(); commit(); }
|
||||
if (ev.key === "Escape") { ev.preventDefault(); renderHistory(); }
|
||||
});
|
||||
input.addEventListener("blur", commit);
|
||||
});
|
||||
});
|
||||
box.querySelectorAll(".conv-del").forEach(btn => {
|
||||
btn.addEventListener("click", (e) => {
|
||||
e.stopPropagation();
|
||||
if (!confirm("确定删除这个对话？")) return;
|
||||
const id = btn.dataset.id;
|
||||
const item = btn.closest(".o1k-conv-item");
|
||||
const infoEl = item.querySelector(".conv-info");
|
||||
infoEl.innerHTML = `<div class="o1k-conv-confirm"><span>确定删除？</span><button class="confirm-yes">删除</button><button class="confirm-no">取消</button></div>`;
|
||||
item.querySelector(".conv-rename").style.display = "none";
|
||||
btn.style.display = "none";
|
||||
infoEl.querySelector(".confirm-yes").addEventListener("click", (ev) => {
|
||||
ev.stopPropagation();
|
||||
conversations = conversations.filter(c => c.id !== id);
|
||||
if (activeConvId === id) activeConvId = conversations[0]?.id || null;
|
||||
saveConversations();
|
||||
renderHistory();
|
||||
});
|
||||
infoEl.querySelector(".confirm-no").addEventListener("click", (ev) => {
|
||||
ev.stopPropagation();
|
||||
renderHistory();
|
||||
});
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
function deleteMessage(idx) {
|
||||
if (!confirm("确定删除这条消息？")) return;
|
||||
const wrap = chatContainer.querySelector(`.o1k-msg-wrap[data-idx="${idx}"]`);
|
||||
if (!wrap || wrap.querySelector(".o1k-msg-del-confirm")) return;
|
||||
const confirmEl = document.createElement("div");
|
||||
confirmEl.className = "o1k-msg-del-confirm";
|
||||
confirmEl.innerHTML = `<span>确定删除？</span><button class="confirm-yes">删除</button><button class="confirm-no">取消</button>`;
|
||||
wrap.appendChild(confirmEl);
|
||||
confirmEl.querySelector(".confirm-yes").addEventListener("click", () => {
|
||||
const conv = getActiveConv();
|
||||
if (!conv) return;
|
||||
conv.messages.splice(idx, 1);
|
||||
conv.updatedAt = Date.now();
|
||||
saveConversations();
|
||||
renderMessages();
|
||||
});
|
||||
confirmEl.querySelector(".confirm-no").addEventListener("click", () => {
|
||||
confirmEl.remove();
|
||||
});
|
||||
}
|
||||
|
||||
function copyMessage(idx) {
|
||||
|
||||
Reference in New Issue
Block a user