InpaintRegionEditor Dev
19d25f5f61
- 使用 folder_paths.get_filename_list('input') 替代手动遍历
- 递归搜索包含子目录文件（如 clipspace/）
- 添加 image_upload 参数支持上传功能
- 删除前端不必要的 clipspace 处理代码
147 lines
4.4 KiB
Python
147 lines
4.4 KiB
Python
"""
|
||
InpaintRegionEditor Node
|
||
|
||
增强版遮罩编辑器，支持 Photopea 集成和可拖动选区
|
||
|
||
概念说明：
|
||
- 遮罩 (Mask): 定义重绘区域，AI 在这个区域内生成新内容
|
||
- 选区 (Region): 定义参考区域，AI 参考这个区域内的原图内容
|
||
|
||
约束：选区必须 >= 遮罩
|
||
"""
|
||
|
||
import torch
|
||
import numpy as np
|
||
import folder_paths
|
||
from PIL import Image
|
||
import os
|
||
import json
|
||
|
||
|
||
class InpaintRegionEditor:
|
||
"""
|
||
增强版遮罩编辑器，支持 Photopea 集成和可拖动选区
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(s):
|
||
files = folder_paths.get_filename_list("input")
|
||
|
||
return {
|
||
"required": {
|
||
"image": (sorted(files), {"image_upload": True}),
|
||
"padding": ("INT", {"default": 64, "min": 0, "max": 512, "step": 32}),
|
||
},
|
||
"hidden": {
|
||
"region_coords": "STRING",
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE", "MASK", "INT", "INT", "INT", "INT")
|
||
RETURN_NAMES = (
|
||
"image",
|
||
"mask",
|
||
"region_top",
|
||
"region_left",
|
||
"region_width",
|
||
"region_height",
|
||
)
|
||
FUNCTION = "process"
|
||
CATEGORY = "image/inpaint"
|
||
|
||
def process(self, image, padding, region_coords=None):
|
||
image_path = folder_paths.get_annotated_filepath(image)
|
||
img = Image.open(image_path)
|
||
|
||
img_width, img_height = img.size
|
||
|
||
# 提取 MASK（从 Alpha 通道）
|
||
if "A" in img.getbands():
|
||
mask = np.array(img.getchannel("A")).astype(np.float32) / 255.0
|
||
mask = 1.0 - mask # ComfyUI 约定：1=遮罩区域
|
||
else:
|
||
raise ValueError(
|
||
"请先创建蒙版（局部重绘需要蒙版）\n"
|
||
"可以通过以下方式创建：\n"
|
||
"1. 上传带 Alpha 通道的 PNG\n"
|
||
"2. 使用 Photopea 编辑后保存\n"
|
||
"3. 使用 ComfyUI 系统蒙版工具绘制"
|
||
)
|
||
|
||
# 计算蒙版边界框
|
||
mask_bounds = self._calculate_mask_bounds(mask)
|
||
|
||
# 解析前端坐标（如果用户拖动过）
|
||
region_rect = None
|
||
if region_coords:
|
||
try:
|
||
region_rect = json.loads(region_coords)
|
||
except:
|
||
pass
|
||
|
||
if not region_rect:
|
||
region_rect = self._calculate_region_rect(mask_bounds, padding)
|
||
|
||
# 验证
|
||
if (
|
||
region_rect["width"] < mask_bounds["width"]
|
||
or region_rect["height"] < mask_bounds["height"]
|
||
):
|
||
raise ValueError(
|
||
f"选区太小！必须 >= 遮罩\n"
|
||
f"选区：{region_rect['width']}×{region_rect['height']}\n"
|
||
f"遮罩：{mask_bounds['width']}×{mask_bounds['height']}"
|
||
)
|
||
|
||
region_rect = self._clamp_to_image(region_rect, img_width, img_height)
|
||
|
||
# 转换为 tensor
|
||
image_tensor = torch.from_numpy(
|
||
np.array(img.convert("RGB")).astype(np.float32) / 255.0
|
||
).unsqueeze(0)
|
||
|
||
mask_tensor = torch.from_numpy(mask).unsqueeze(0)
|
||
|
||
return (
|
||
image_tensor,
|
||
mask_tensor,
|
||
int(region_rect["y"]),
|
||
int(region_rect["x"]),
|
||
int(region_rect["width"]),
|
||
int(region_rect["height"]),
|
||
)
|
||
|
||
def _calculate_mask_bounds(self, mask):
|
||
rows = np.any(mask > 0.5, axis=1)
|
||
cols = np.any(mask > 0.5, axis=0)
|
||
|
||
if not np.any(rows) or not np.any(cols):
|
||
return {"x": 0, "y": 0, "width": 0, "height": 0}
|
||
|
||
rmin, rmax = np.where(rows)[0][[0, -1]]
|
||
cmin, cmax = np.where(cols)[0][[0, -1]]
|
||
|
||
return {
|
||
"x": int(cmin),
|
||
"y": int(rmin),
|
||
"width": int(cmax - cmin + 1),
|
||
"height": int(rmax - rmin + 1),
|
||
}
|
||
|
||
def _calculate_region_rect(self, mask_bounds, padding):
|
||
return {
|
||
"x": mask_bounds["x"] - padding,
|
||
"y": mask_bounds["y"] - padding,
|
||
"width": mask_bounds["width"] + padding * 2,
|
||
"height": mask_bounds["height"] + padding * 2,
|
||
}
|
||
|
||
def _clamp_to_image(self, rect, img_width, img_height):
|
||
x = max(0, min(rect["x"], img_width - rect["width"]))
|
||
y = max(0, min(rect["y"], img_height - rect["height"]))
|
||
return {"x": x, "y": y, "width": rect["width"], "height": rect["height"]}
|
||
|
||
@classmethod
|
||
def IS_CHANGED(s, image, padding, region_coords=None):
|
||
return float("NaN")
|