nodes/QL-boy-ComfyUI-Advanced-Tile-Processing
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add new merge scheme, other optimization processing添加新合并方案，其它优化处理

Browse Source 
Added Laplacian pyramid blend mode (laplacian)
Node port optimization
Other optimization
新增拉普拉斯金字塔混合模式（laplacian）
节点端口优化
其它优化
This commit is contained in:
QL-boy
2026-01-13 05:13:04 +08:00
committed by GitHub
parent 113d066e77
commit 9a7f797f0d
1 changed files with 175 additions and 109 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes/tile_merger.py
+175 -109
View File
@@ -1,6 +1,6 @@
import torch
import comfy.utils
from ..utils.blending import generate_weight_mask
from ..utils.blending import generate_weight_mask, laplacian_pyramid_blending
class CustomTileMerger:
@classmethod
@@ -8,12 +8,12 @@ class CustomTileMerger:
return {
"required": {
"tile_config": ("TILE_CONFIG",),
"blend_mode": (["gaussian", "linear", "cosine", "none"], {"default": "gaussian"}),
"feather_percent": ("FLOAT", {"default": 25.0, "min": 0.0, "max": 50.0, "step": 1.0}),
},
"optional": {
"processed_tiles_image": ("IMAGE",),
"processed_tiles_latent": ("LATENT",),
"blend_mode": (["gaussian", "linear", "cosine", "laplacian", "none"], {"default": "gaussian"}),
"feather_percent": ("FLOAT", {"default": 25.0, "min": 0.0, "max": 50.0, "step": 1.0}),
}
}
@@ -22,95 +22,148 @@ class CustomTileMerger:
FUNCTION = "merge_tiles"
CATEGORY = "image/processing/tile"
# 关键修复：添加 INPUT_IS_LIST = True
# 这告诉 ComfyUI 不要自动遍历列表，而是将整个列表一次性传给 merge_tiles 函数
INPUT_IS_LIST = True
def merge_tiles(self, tile_config, blend_mode, feather_percent, processed_tiles_image=None, processed_tiles_latent=None):
# 由于开启了 INPUT_IS_LIST = True，所有输入参数都会变成列表。
# 我们需要解包非列表类型的参数（即那些我们只期望接收一个值的参数）。
# 解包配置和参数
def merge_tiles(self, tile_config, processed_tiles_image=None, processed_tiles_latent=None, blend_mode=None, feather_percent=None):
# Unpack params (since INPUT_IS_LIST=True, even single inputs are lists)
tile_config = tile_config[0]
blend_mode = blend_mode[0]
feather = feather_percent[0] / 100.0
blend_mode = blend_mode[0] if blend_mode else "gaussian"
feather = feather_percent[0] / 100.0 if feather_percent else 0.25
merged_image = None
merged_latent = None
splits = tile_config["splits"]
batch_size = tile_config["batch_size"]
# --- 处理图像合并 ---
# processed_tiles_image 是一个列表，可能包含：
# 情况 A (List输入): [Tensor(1,H,W,C), Tensor(1,H,W,C), ...]
# 情况 B (Batch输入): [Tensor(N,H,W,C)]
# 我们将其统一展平为 [Tensor(1,H,W,C), ...] 的列表
if processed_tiles_image is not None:
tiles = []
for t in processed_tiles_image:
# t 是 Tensor [B, H, W, C]
# 即使输入是 List，列表中的每个元素也是 Tensor
# 如果输入是 Batch，列表只有一个元素，但该 Tensor 的 B > 1
for b in range(t.shape[0]):
tiles.append(t[b:b+1])
# Restore padded dimensions for canvas
H, W = tile_config["padded_height"], tile_config["padded_width"]
orig_H, orig_W = tile_config["original_height"], tile_config["original_width"]
# 验证数量
expected_tiles = len(splits) * batch_size
if len(tiles) != expected_tiles:
print(f"Warning: Expected {expected_tiles} tiles, got {len(tiles)}. Merge might be incomplete or misaligned.")
# ----------------------------------------------------------------------
# Helper: Flatten Input Lists
# ----------------------------------------------------------------------
def flatten_input(input_data, is_latent=False):
flat = []
if input_data is None: return flat
for item in input_data:
if item is None: continue
if is_latent:
# item is {'samples': Tensor}
samples = item['samples']
for b in range(samples.shape[0]):
flat.append(samples[b:b+1]) # Keep [1, C, H, W]
else:
# item is Tensor [B, H, W, C]
for b in range(item.shape[0]):
flat.append(item[b:b+1]) # Keep [1, H, W, C]
return flat
# 初始化画布 [B, H, W, C]
H, W = tile_config["original_height"], tile_config["original_width"]
# 确保 tiles 非空
if len(tiles) > 0:
C = tiles[0].shape[-1]
device = tiles[0].device
# ----------------------------------------------------------------------
# Logic: Image Merging
# ----------------------------------------------------------------------
tiles_img = flatten_input(processed_tiles_image, is_latent=False)
if len(tiles_img) > 0:
# Check consistency
expected = len(splits) * batch_size
if len(tiles_img) != expected:
print(f"[Warning] TileMerger: Expected {expected} tiles, got {len(tiles_img)}.")
if blend_mode == "laplacian":
# --- Laplacian Pyramid Blending ---
# This mode accumulates whole pyramids, so it is VRAM heavy.
# We process one Batch item at a time to save memory.
final_canvases = []
device = tiles_img[0].device
pbar = comfy.utils.ProgressBar(batch_size)
for b in range(batch_size):
# Gather tiles for this batch
batch_tiles_list = []
batch_splits_info = []
start_idx = b * len(splits)
for i, (y, x, h, w) in enumerate(splits):
t_idx = start_idx + i
if t_idx < len(tiles_img):
batch_tiles_list.append(tiles_img[t_idx])
batch_splits_info.append((y, x, h, w))
if not batch_tiles_list: continue
# Execute Pyramid Blending
# shape: (H, W, C) -> (C, H, W) for calc, then back
# The util function expects tensors in [C, H, W] or [1, C, H, W]
# 1. Permute tiles to [1, C, H, W]
permuted_tiles = [t.permute(0, 3, 1, 2) for t in batch_tiles_list]
merged_tensor = laplacian_pyramid_blending(
permuted_tiles,
batch_splits_info,
(H, W),
feather=feather,
device=device
)
# 2. Permute back to [1, H, W, C]
final_canvases.append(merged_tensor.permute(0, 2, 3, 1))
pbar.update(1)
if final_canvases:
merged_image = torch.cat(final_canvases, dim=0)
else:
# --- Standard Weighted Average (Gaussian/Linear/Cosine) ---
C = tiles_img[0].shape[-1]
device = tiles_img[0].device
canvas = torch.zeros((batch_size, H, W, C), device=device)
weight_map = torch.zeros((batch_size, H, W, C), device=device)
# 预计算 Mask
tile_h, tile_w = tiles[0].shape[1], tiles[0].shape[2]
base_mask = generate_weight_mask((tile_h, tile_w), blend_mode, feather, device)
# Precompute base mask (assuming uniform tile size)
th, tw = tiles_img[0].shape[1], tiles_img[0].shape[2]
base_mask = generate_weight_mask((th, tw), blend_mode, feather, device)
base_mask = base_mask.unsqueeze(0).unsqueeze(-1) # [1, H, W, 1]
pbar = comfy.utils.ProgressBar(len(tiles))
pbar = comfy.utils.ProgressBar(len(tiles_img))
for b in range(batch_size):
for i, (y, x, h, w) in enumerate(splits):
tile_idx = b * len(splits) + i
if tile_idx >= len(tiles): break
idx = b * len(splits) + i
if idx >= len(tiles_img): break
tile = tiles[tile_idx] # [1, Th, Tw, C]
# 动态 Resize 支持 (处理 Tiled Upscale 的情况)
tile = tiles_img[idx]
curr_h, curr_w = tile.shape[1], tile.shape[2]
if curr_h != h or curr_w != w:
scale_h = curr_h / tile_config["tile_height"]
scale_w = curr_w / tile_config["tile_width"]
target_H = int(H * scale_h)
target_W = int(W * scale_w)
# 如果画布尺寸需要调整（仅第一次检测到时）
if canvas.shape[1] != target_H:
canvas = torch.zeros((batch_size, target_H, target_W, C), device=device)
weight_map = torch.zeros((batch_size, target_H, target_W, C), device=device)
tgt_y = int(y * scale_h)
tgt_x = int(x * scale_w)
tgt_h, tgt_w = curr_h, curr_w
curr_mask = generate_weight_mask((tgt_h, tgt_w), blend_mode, feather, device)
curr_mask = curr_mask.unsqueeze(0).unsqueeze(-1)
# Handle Dynamic Resize (Upscale)
scale_h = curr_h / tile_config["tile_height"]
scale_w = curr_w / tile_config["tile_width"]
# Adjust target coords
tgt_y = int(y * scale_h)
tgt_x = int(x * scale_w)
tgt_h, tgt_w = curr_h, curr_w
# If canvas size changed due to upscale
if canvas.shape[1] < int(H * scale_h):
new_H = int(H * scale_h)
new_W = int(W * scale_w)
# Re-init canvas
canvas = torch.zeros((batch_size, new_H, new_W, C), device=device)
weight_map = torch.zeros((batch_size, new_H, new_W, C), device=device)
if curr_h != th or curr_w != tw:
curr_mask = generate_weight_mask((curr_h, curr_w), blend_mode, feather, device)
curr_mask = curr_mask.unsqueeze(0).unsqueeze(-1)
else:
tgt_y, tgt_x, tgt_h, tgt_w = y, x, h, w
curr_mask = base_mask
# 累加
curr_mask = base_mask
# Accumulate
canvas[b:b+1, tgt_y:tgt_y+tgt_h, tgt_x:tgt_x+tgt_w, :] += tile * curr_mask
weight_map[b:b+1, tgt_y:tgt_y+tgt_h, tgt_x:tgt_x+tgt_w, :] += curr_mask
@@ -118,46 +171,59 @@ class CustomTileMerger:
merged_image = canvas / (weight_map + 1e-6)
# --- 处理 Latent 合并 ---
if processed_tiles_latent is not None:
lat_tiles = []
# 同样展平 Latent 列表
for t_dict in processed_tiles_latent:
# ComfyUI 的 LATENT 类型是字典 {'samples': Tensor}
if 'samples' in t_dict:
full_samples = t_dict['samples'] # [B, C, H, W]
for b in range(full_samples.shape[0]):
lat_tiles.append(full_samples[b:b+1])
if len(lat_tiles) > 0:
is_image_config = not tile_config.get("is_latent_config", False)
scale = 8 if is_image_config else 1
orig_H = tile_config["original_height"] // scale
orig_W = tile_config["original_width"] // scale
C = lat_tiles[0].shape[1]
device = lat_tiles[0].device
canvas_l = torch.zeros((batch_size, C, orig_H, orig_W), device=device)
weight_map_l = torch.zeros((batch_size, 1, orig_H, orig_W), device=device)
th, tw = lat_tiles[0].shape[2], lat_tiles[0].shape[3]
mask_l = generate_weight_mask((C, th, tw), blend_mode, feather, device)
for b in range(batch_size):
for i, (y, x, h, w) in enumerate(splits):
tile_idx = b * len(splits) + i
if tile_idx >= len(lat_tiles): break
tile = lat_tiles[tile_idx]
ly, lx = y // scale, x // scale
lh, lw = tile.shape[2], tile.shape[3]
canvas_l[b:b+1, :, ly:ly+lh, lx:lx+lw] += tile * mask_l
weight_map_l[b:b+1, :, ly:ly+lh, lx:lx+lw] += mask_l
merged_samples = canvas_l / (weight_map_l + 1e-6)
merged_latent = {"samples": merged_samples}
# Crop Padding if exists
if merged_image is not None:
if orig_H != H or orig_W != W:
merged_image = merged_image[:, :orig_H, :orig_W, :]
return (merged_image, merged_latent)
# ----------------------------------------------------------------------
# Logic: Latent Merging
# ----------------------------------------------------------------------
# (Latent merging rarely benefits from Laplacian blending due to non-spatial nature of VAE channels,
# so we default to Standard Weighted for stability, or mirror simple blending)
tiles_lat = flatten_input(processed_tiles_latent, is_latent=True)
if len(tiles_lat) > 0:
# Latent typically 1/8 scale
scale = 8
L_H, L_W = H // scale, W // scale
L_orig_H, L_orig_W = orig_H // scale, orig_W // scale
C = tiles_lat[0].shape[1]
device = tiles_lat[0].device
canvas_l = torch.zeros((batch_size, C, L_H, L_W), device=device)
weight_l = torch.zeros((batch_size, 1, L_H, L_W), device=device)
th, tw = tiles_lat[0].shape[2], tiles_lat[0].shape[3]
# Use same blend mode but simpler mask gen
mask_l = generate_weight_mask((C, th, tw), blend_mode, feather, device)
for b in range(batch_size):
for i, (y, x, h, w) in enumerate(splits):
idx = b * len(splits) + i
if idx >= len(tiles_lat): break
tile = tiles_lat[idx]
ly, lx = y // scale, x // scale
lh, lw = tile.shape[2], tile.shape[3]
# Safety bounds
ly = min(ly, L_H)
lx = min(lx, L_W)
lh = min(lh, L_H - ly)
lw = min(lw, L_W - lx)
# Accumulate
canvas_l[b:b+1, :, ly:ly+lh, lx:lx+lw] += tile * mask_l
weight_l[b:b+1, :, ly:ly+lh, lx:lx+lw] += mask_l
merged_samples = canvas_l / (weight_l + 1e-6)
# Crop Latent Padding
if L_orig_H != L_H or L_orig_W != L_W:
merged_samples = merged_samples[:, :, :L_orig_H, :L_orig_W]
merged_latent = {"samples": merged_samples}
return (merged_image, merged_latent)