diff --git a/nodes/tile_merger.py b/nodes/tile_merger.py new file mode 100644 index 0000000..b1e28fc --- /dev/null +++ b/nodes/tile_merger.py @@ -0,0 +1,163 @@ +import torch +import comfy.utils +from ..utils.blending import generate_weight_mask + +class CustomTileMerger: + @classmethod + def INPUT_TYPES(cls): + 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",), + } + } + + RETURN_TYPES = ("IMAGE", "LATENT") + RETURN_NAMES = ("merged_image", "merged_latent") + 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，所有输入参数都会变成列表。 + # 我们需要解包非列表类型的参数（即那些我们只期望接收一个值的参数）。 + + # 解包配置和参数 + tile_config = tile_config[0] + blend_mode = blend_mode[0] + feather = feather_percent[0] / 100.0 + + 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]) + + # 验证数量 + 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.") + + # 初始化画布 [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 + + 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) + base_mask = base_mask.unsqueeze(0).unsqueeze(-1) # [1, H, W, 1] + + pbar = comfy.utils.ProgressBar(len(tiles)) + + 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 + + tile = tiles[tile_idx] # [1, Th, Tw, C] + + # 动态 Resize 支持 (处理 Tiled Upscale 的情况) + 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) + else: + tgt_y, tgt_x, tgt_h, tgt_w = y, x, h, w + curr_mask = base_mask + + # 累加 + 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 + + pbar.update(1) + + 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} + + return (merged_image, merged_latent) \ No newline at end of file diff --git a/nodes/tile_splitter.py b/nodes/tile_splitter.py new file mode 100644 index 0000000..50dbdcf --- /dev/null +++ b/nodes/tile_splitter.py @@ -0,0 +1,113 @@ +import torch +from ..utils.blending import get_tiled_splits + +class CustomTileSplitter: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "tile_size": ("INT", {"default": 512, "min": 64, "max": 8192, "step": 64}), + "overlap": ("INT", {"default": 64, "min": 0, "max": 512, "step": 8}), + }, + "optional": { + "image": ("IMAGE",), # [B, H, W, C] + "latent": ("LATENT",), # {'samples': [B, C, H, W]} + } + } + + # 修复：将 LIST 类型改为实际的 IMAGE/LATENT 类型，以便 UI 正确连接 + RETURN_TYPES = ("IMAGE", "IMAGE", "LATENT", "LATENT", "TILE_CONFIG") + RETURN_NAMES = ("tiles_image_batch", "tiles_image_list", "tiles_latent_batch", "tiles_latent_list", "tile_config") + + # 新增：明确定义哪些输出是列表。 + # 对应位置：(False, True, False, True, False) + # tiles_image_batch (False), tiles_image_list (True), tiles_latent_batch (False), tiles_latent_list (True), tile_config (False) + OUTPUT_IS_LIST = (False, True, False, True, False) + + FUNCTION = "split_tiles" + CATEGORY = "image/processing/tile" + + def split_tiles(self, tile_size=512, overlap=64, image=None, latent=None): + tiles_image_batch = None + tiles_image_list = [] + tiles_latent_batch = None + tiles_latent_list = [] + + full_h, full_w = 0, 0 + processing_image = False + processing_latent = False + batch_size = 1 + + if image is not None: + processing_image = True + batch_size, full_h, full_w, channels = image.shape + + if latent is not None: + processing_latent = True + # latent samples: [B, C, H, W] + l_batch, l_c, l_h, l_w = latent['samples'].shape + if not processing_image: + # 如果没有 Image 参照，根据 Latent 估算像素尺寸 (x8) + scale_factor = 8 + l_tile_size = tile_size // scale_factor + l_overlap = overlap // scale_factor + + full_h, full_w = l_h, l_w + tile_size_act = l_tile_size + overlap_act = l_overlap + else: + tile_size_act = tile_size + overlap_act = overlap + else: + if processing_image: + tile_size_act = tile_size + overlap_act = overlap + + if not processing_image and not processing_latent: + raise ValueError("No input provided (Image or Latent required)") + + splits = get_tiled_splits(full_h, full_w, tile_size_act, overlap_act) + + tile_config = { + "original_height": full_h, + "original_width": full_w, + "tile_height": tile_size_act, + "tile_width": tile_size_act, + "overlap": overlap_act, + "splits": splits, + "is_latent_config": (not processing_image and processing_latent), + "batch_size": batch_size + } + + # Image Processing + if processing_image: + img_tiles = [] + for (y, x, h, w) in splits: + tile = image[:, y:y+h, x:x+w, :] + img_tiles.append(tile) + + tiles_image_list = img_tiles + try: + tiles_image_batch = torch.cat(img_tiles, dim=0) + except: + pass + + # Latent Processing + if processing_latent: + lat_tiles = [] + samples = latent['samples'] + is_pixel_config = processing_image + scale = 8 if is_pixel_config else 1 + + for (y, x, h, w) in splits: + ly, lx = y // scale, x // scale + lh, lw = h // scale, w // scale + tile = samples[:, :, ly:ly+lh, lx:lx+lw] + lat_tiles.append(tile) + + # Latent 列表中的每个元素都需要是标准的 Latent 字典格式 + tiles_latent_list = [{"samples": t} for t in lat_tiles] + if len(lat_tiles) > 0: + tiles_latent_batch = {"samples": torch.cat(lat_tiles, dim=0)} + + return (tiles_image_batch, tiles_image_list, tiles_latent_batch, tiles_latent_list, tile_config) \ No newline at end of file