Add files via upload

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)

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)