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
@@ -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,)
|
||||
Reference in New Issue
Block a user