From a1082cb48b3598424c295a77df9e0f7666d66144 Mon Sep 17 00:00:00 2001 From: QL-boy <46901341+QL-boy@users.noreply.github.com> Date: Sat, 17 Jan 2026 19:22:32 +0800 Subject: [PATCH] Refactor tile merging logic with new blending methods --- nodes/tile_merger.py | 397 +++++++++++++++++++++++-------------------- 1 file changed, 215 insertions(+), 182 deletions(-) diff --git a/nodes/tile_merger.py b/nodes/tile_merger.py index 8de5737..2f1faee 100644 --- a/nodes/tile_merger.py +++ b/nodes/tile_merger.py @@ -1,6 +1,13 @@ import torch +import torch.nn.functional as F +import numpy as np +import scipy.ndimage +from ..utils.blending import ( + compute_optimal_seamline, + match_histograms, + multi_band_blend +) import comfy.utils -from ..utils.blending import generate_weight_mask, laplacian_pyramid_blending class CustomTileMerger: @classmethod @@ -8,222 +15,248 @@ class CustomTileMerger: return { "required": { "tile_config": ("TILE_CONFIG",), + "blend_mode": (["linear", "gaussian", "multi_band", "none"], {"default": "gaussian"}), + "blend_strength": ("FLOAT", {"default": 0.25, "min": 0.0, "max": 1.0, "step": 0.01}), + "seam_mode": (["optimal", "middle"], {"default": "optimal"}), + "histogram_matching": ("BOOLEAN", {"default": False}), + "seam_search_radius": ("INT", {"default": 0, "min": 0, "max": 100}), }, "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}), } } RETURN_TYPES = ("IMAGE", "LATENT") RETURN_NAMES = ("merged_image", "merged_latent") FUNCTION = "merge_tiles" - CATEGORY = "image/processing/tile" - + CATEGORY = "image/postprocessing" 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] if blend_mode else "gaussian" - feather = feather_percent[0] / 100.0 if feather_percent else 0.25 + def merge_tiles(self, tile_config, blend_mode, blend_strength, seam_mode, + histogram_matching, seam_search_radius, + processed_tiles_image=None, processed_tiles_latent=None): - merged_image = None - merged_latent = None + config = tile_config[0] + mode = blend_mode[0] + strength = blend_strength[0] + seam_type = seam_mode[0] + do_hist_match = histogram_matching[0] + search_radius = seam_search_radius[0] - splits = tile_config["splits"] - batch_size = tile_config["batch_size"] + splits = config["splits"] + batch_size = config["batch_size"] + H, W = config["padded_height"], config["padded_width"] + orig_H, orig_W = config["original_height"], config["original_width"] + overlap = config["overlap"] - # 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"] + # Optimization: Build a Grid Map for fast neighbor lookup + # key: (row, col) -> value: index in splits list + grid_map = {} + for idx, s in enumerate(splits): + grid_map[(s["row"], s["col"])] = idx - # ---------------------------------------------------------------------- - # 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 - - # ---------------------------------------------------------------------- - # Logic: Image Merging - # ---------------------------------------------------------------------- - tiles_img = flatten_input(processed_tiles_image, is_latent=False) + # ================= IMAGE MERGING ================= + merged_img = None - if len(tiles_img) > 0: + # processed_tiles_image is a LIST of [B, H, W, C] tensors (one per tile position) + if processed_tiles_image is not None and len(processed_tiles_image) > 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 len(processed_tiles_image) != len(splits): + print(f"Warning: Tile count mismatch. Config: {len(splits)}, Input: {len(processed_tiles_image)}") + + # Use first tile to determine device and channels + ref_tile = processed_tiles_image[0] + device = ref_tile.device + C = ref_tile.shape[-1] + + # Check actual batch size from input (in case it changed during processing) + actual_batch_size = ref_tile.shape[0] - 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 + # Initialize Canvas + canvas = torch.zeros((actual_batch_size, H, W, C), device=device) + pbar = comfy.utils.ProgressBar(actual_batch_size * len(splits)) + + for b in range(actual_batch_size): + # 1. Prepare Tile Clones for this Batch + # We need clones because Histogram Matching modifies them in-place + current_batch_tiles = [] + for i in range(len(splits)): + if i < len(processed_tiles_image): + # Extract the specific image for batch b + # Input List [TileIdx] -> Tensor [Batch, H, W, C] + # We want [1, H, W, C] for easier processing logic + t = processed_tiles_image[i][b:b+1].clone() + current_batch_tiles.append(t) + else: + current_batch_tiles.append(torch.zeros_like(ref_tile[0:1])) - # 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] + # Initialize Masks (All 1.0) + batch_masks = [torch.ones((s['h'], s['w'], 1), device=device) for s in splits] + + # 2. Process Tiles + for i, split in enumerate(splits): + y, x, h, w = split['y'], split['x'], split['h'], split['w'] + r, c = split['row'], split['col'] - # 1. Permute tiles to [1, C, H, W] - permuted_tiles = [t.permute(0, 3, 1, 2) for t in batch_tiles_list] + tile_i = current_batch_tiles[i] # [1, h, w, c] - merged_tensor = laplacian_pyramid_blending( - permuted_tiles, - batch_splits_info, - (H, W), - feather=feather, - device=device - ) + # --- A. Histogram Matching (Look at Left Neighbor) --- + if do_hist_match and c > 0: + left_idx = grid_map.get((r, c - 1)) + if left_idx is not None: + prev_tile = current_batch_tiles[left_idx] + + # Overlap: prev_tile Right vs tile_i Left + ref = prev_tile[0, :, -overlap:, :] + src = tile_i[0, :, :overlap, :] + + s_mean = torch.mean(src, dim=(0, 1), keepdim=True) + s_std = torch.std(src, dim=(0, 1), keepdim=True) + 1e-6 + r_mean = torch.mean(ref, dim=(0, 1), keepdim=True) + r_std = torch.std(ref, dim=(0, 1), keepdim=True) + 1e-6 + + # Apply to full tile + tile_i[0] = (tile_i[0] - s_mean) / s_std * r_std + r_mean + tile_i.clamp_(0.0, 1.0) # In-place clamp + current_batch_tiles[i] = tile_i # Update reference + + # --- B. Seam Calculation (Cut Future Neighbors) --- + + # 1. Right Neighbor (Forward Cut) + right_idx = grid_map.get((r, c + 1)) + if right_idx is not None: + tile_j = current_batch_tiles[right_idx] + + ov_left = tile_i[0, :, -overlap:, :] # Left Tile (i) Right part + ov_right = tile_j[0, :, :overlap, :] # Right Tile (j) Left part + + seam = compute_optimal_seamline(ov_left, ov_right, 'horizontal', search_radius, seam_type) + seam = seam.to(device).unsqueeze(-1) # [h, overlap, 1] + + # Tile J (Right) should be hidden where Seam is 1 (Left wins) + # So Mask J = 0 where Seam = 1 + batch_masks[right_idx][:, :overlap, :] *= (1.0 - seam) + + # 2. Bottom Neighbor (Forward Cut) + bottom_idx = grid_map.get((r + 1, c)) + if bottom_idx is not None: + tile_k = current_batch_tiles[bottom_idx] + + ov_top = tile_i[0, -overlap:, :, :] + ov_bottom = tile_k[0, :overlap, :, :] + + seam = compute_optimal_seamline(ov_top, ov_bottom, 'vertical', search_radius, seam_type) + seam = seam.to(device).unsqueeze(-1) + + batch_masks[bottom_idx][:overlap, :, :] *= (1.0 - seam) + + # 3. Feathering + if strength > 0 and mode != "none": + sigma = strength * (overlap / 4.0) + sigma = max(0.5, sigma) + for m_idx in range(len(batch_masks)): + m_tens = batch_masks[m_idx] + if m_tens.min() < 0.99: # Only blur if there are edges + m_np = m_tens.squeeze(-1).cpu().numpy() + # Nearest mode keeps borders valid + m_blur = scipy.ndimage.gaussian_filter(m_np, sigma=sigma, mode='nearest') + batch_masks[m_idx] = torch.from_numpy(m_blur).unsqueeze(-1).to(device) + + # 4. Blend to Canvas + for i, split in enumerate(splits): + y, x, h, w = split['y'], split['x'], split['h'], split['w'] + tile = current_batch_tiles[i] + mask = batch_masks[i] + + canvas_crop = canvas[b:b+1, y:y+h, x:x+w, :] + + if mode == "multi_band" and strength > 0: + t_nchw = tile.permute(0, 3, 1, 2) + c_nchw = canvas_crop.permute(0, 3, 1, 2) + m_nchw = mask.permute(2, 0, 1).unsqueeze(0) + + blended = multi_band_blend(c_nchw, t_nchw, m_nchw, levels=5) + canvas[b:b+1, y:y+h, x:x+w, :] = blended.permute(0, 2, 3, 1) + else: + # Standard blending + canvas[b:b+1, y:y+h, x:x+w, :] = canvas_crop * (1.0 - mask) + tile * mask - # 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) + merged_img = canvas[:, :orig_H, :orig_W, :] - 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) - - # 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_img)) - - for b in range(batch_size): - for i, (y, x, h, w) in enumerate(splits): - idx = b * len(splits) + i - if idx >= len(tiles_img): break - - tile = tiles_img[idx] - curr_h, curr_w = tile.shape[1], tile.shape[2] - - # 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: - 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 - - pbar.update(1) - - merged_image = canvas / (weight_map + 1e-6) - - # 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, :] - - # ---------------------------------------------------------------------- - # 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 + # ================= LATENT MERGING ================= + merged_lat = None + if processed_tiles_latent is not None and len(processed_tiles_latent) > 0: + # Latent List: [{"samples": Tensor[B, 4, H, W]}, ...] + ref_lat = processed_tiles_latent[0]['samples'] + device = ref_lat.device + actual_batch_size = ref_lat.shape[0] + lc = ref_lat.shape[1] + 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((actual_batch_size, lc, L_H, L_W), device=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 + for b in range(actual_batch_size): + for i, split in enumerate(splits): + if i >= len(processed_tiles_latent): break - tile = tiles_lat[idx] + # Get Latent Tile for Batch b + # processed_tiles_latent[i] is {"samples": [B, C, H, W]} + tile = processed_tiles_latent[i]['samples'][b:b+1] + + y, x = split['y'], split['x'] ly, lx = y // scale, x // scale lh, lw = tile.shape[2], tile.shape[3] + lo = overlap // scale - # Safety bounds - ly = min(ly, L_H) - lx = min(lx, L_W) - lh = min(lh, L_H - ly) - lw = min(lw, L_W - lx) + # Latent Blending: Hard Cut at Center of Overlap + crop_t, crop_b = 0, lh + crop_l, crop_r = 0, lw + paste_y, paste_x = ly, 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 + # Logic: If neighbor exists, crop half overlap + # We rely on grid indices from split dict + r, c = split['row'], split['col'] + + if c > 0: # Left neighbor exists + crop_l += lo // 2 + paste_x += lo // 2 + if r > 0: # Top neighbor exists + crop_t += lo // 2 + paste_y += lo // 2 + + # Determine right/bottom overlap presence based on grid size + # n_cols is stored in config["cols"] if needed, or check neighbor + if grid_map.get((r, c + 1)) is not None: + crop_r -= lo // 2 + if grid_map.get((r + 1, c)) is not None: + crop_b -= lo // 2 + + if crop_r > crop_l and crop_b > crop_t: + snippet = tile[:, :, crop_t:crop_b, crop_l:crop_r] + ph, pw = snippet.shape[2], snippet.shape[3] + + # Boundary safe paste + y1 = max(0, paste_y) + x1 = max(0, paste_x) + y2 = min(L_H, paste_y + ph) + x2 = min(L_W, paste_x + pw) + + # Corresponding snippet coords + sy1 = y1 - paste_y + sx1 = x1 - paste_x + sy2 = sy1 + (y2 - y1) + sx2 = sx1 + (x2 - x1) + + if y2 > y1 and x2 > x1: + canvas_l[b:b+1, :, y1:y2, x1:x2] = snippet[:, :, sy1:sy2, sx1:sx2] - merged_samples = canvas_l / (weight_l + 1e-6) + merged_samples = canvas_l[:, :, :L_orig_H, :L_orig_W] + merged_lat = {"samples": merged_samples} - # 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) + return (merged_img, merged_lat)