o1keyandClaude Opus 4.7
3f0f4099fb
- 新增 O1keyRemoveBackground 节点（基于 rembg CPU 推理） - 新增 O1keyColorRemoveBG 节点（颜色距离去背景，支持多模式） - 新增 O1keySavePSD 节点（手写 PSD 二进制，零外部依赖） - GPT Image 批量输出不同尺寸时自动 resize 对齐 - 聊天面板大幅增强（多模态/交互优化） - 重启按钮绕过 beforeunload 弹窗强制刷新 - 默认路由切换为 CF加速 - http_error 新增 system error 友好文案 Co-Authored-By: Claude Opus 4.7 <[email protected]>
238 lines
7.8 KiB
Python
238 lines
7.8 KiB
Python
"""
|
||
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,)
|