diff --git a/__init__.py b/__init__.py index e41ed44..25d922d 100644 --- a/__init__.py +++ b/__init__.py @@ -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", []) diff --git a/clients/gpt_image_client.py b/clients/gpt_image_client.py index 585c71c..ce9484f 100644 --- a/clients/gpt_image_client.py +++ b/clients/gpt_image_client.py @@ -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] # ── 响应解析(通用) ───────────────────────────────────────────────────── diff --git a/nodes/__init__.py b/nodes/__init__.py index 453e265..2c5608c 100644 --- a/nodes/__init__.py +++ b/nodes/__init__.py @@ -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'] diff --git a/nodes/color_remove_bg.py b/nodes/color_remove_bg.py new file mode 100644 index 0000000..decda48 --- /dev/null +++ b/nodes/color_remove_bg.py @@ -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,) diff --git a/nodes/remove_bg.py b/nodes/remove_bg.py new file mode 100644 index 0000000..3ffdbf4 --- /dev/null +++ b/nodes/remove_bg.py @@ -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,) diff --git a/nodes/save_psd.py b/nodes/save_psd.py new file mode 100644 index 0000000..139b134 --- /dev/null +++ b/nodes/save_psd.py @@ -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,) diff --git a/requirements.txt b/requirements.txt index 407fdd4..8f0c03a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,4 @@ aiohttp>=3.9.0 Pillow>=10.0.0 requests>=2.31.0 +rembg[cpu]>=2.0.50 diff --git a/utils/color_key.py b/utils/color_key.py new file mode 100644 index 0000000..834f305 --- /dev/null +++ b/utils/color_key.py @@ -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") diff --git a/utils/http_error.py b/utils/http_error.py index 488f05c..548f379 100644 --- a/utils/http_error.py +++ b/utils/http_error.py @@ -29,6 +29,7 @@ HTTP_ERROR_MESSAGES = { # 错误内容关键词 → 用户友好文案(优先于状态码匹配) ERROR_CONTENT_MESSAGES = { "The current model has a high load": "模型过载,请稍后重试!", + "system error": "系统错误,请稍后重试。", } # 可退避重试的状态码 diff --git a/utils/rembg_utils.py b/utils/rembg_utils.py new file mode 100644 index 0000000..783634d --- /dev/null +++ b/utils/rembg_utils.py @@ -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) diff --git a/web/js/chatPanel.js b/web/js/chatPanel.js index 342436a..fe786b2 100644 --- a/web/js/chatPanel.js +++ b/web/js/chatPanel.js @@ -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(); - conv.updatedAt = Date.now(); - saveConversations(); - renderHistory(); + const item = btn.closest(".o1k-conv-item"); + const infoEl = item.querySelector(".conv-info"); + const oldHtml = infoEl.innerHTML; + infoEl.innerHTML = ``; + 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; - conversations = conversations.filter(c => c.id !== id); - if (activeConvId === id) activeConvId = conversations[0]?.id || null; - saveConversations(); - renderHistory(); + const item = btn.closest(".o1k-conv-item"); + const infoEl = item.querySelector(".conv-info"); + infoEl.innerHTML = `