diff --git a/nodes/image_nodes.py b/nodes/image_nodes.py index cfec880..2542844 100644 --- a/nodes/image_nodes.py +++ b/nodes/image_nodes.py @@ -257,6 +257,57 @@ class ImageBlend_GPU: return (output_bhwc.cpu(),) +class ImageBlendGrid: + """ + Grid configuration node for Image Blend GPU Advanced. + + Plug its output into the 'enable_grid' input of Image Blend GPU Advanced. + When 'enable' is not "false", that node builds an image grid out of the + background frames: + + - enable "3x2" -> 3 columns x 2 rows (6 cells), "3x3" -> 9 cells. + - The background video length is normalized to 'frames': the last frame + is duplicated if the video is shorter, extra frames are dropped if it + is longer. + - One grid is built per output frame (sliding storyboard): frame t's + cells are evenly sampled over the frames from t to the end, so the + first cell is always the current frame and the window slides forward + until the last frame fills every cell. Each cell keeps the background + aspect ratio. + - The output canvas fits the grid exactly (e.g. 3x2 is only 2 cell-rows + tall), so there are no empty rows/columns around the grid. + - horizontal / vertical select which grid cell the mask output marks + (e.g. left/top -> the top-left cell is white in the mask, everything + else is black). The mask has one frame per output frame. + """ + NODE_NAME = "Image Blend Grid" + + GRID_LAYOUTS = ["false", "3x2", "3x3"] + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "enable": (cls.GRID_LAYOUTS, {"default": "false"}), + "horizontal": (["left", "middle", "right"], {"default": "left"}), + "vertical": (["top", "middle", "bottom"], {"default": "top"}), + "frames": ("INT", {"default": 22, "min": 5, "max": 1000000, "step": 17}), + } + } + + RETURN_TYPES = ("GRID_CONFIG", "INT") + RETURN_NAMES = ("enable_grid", "frames") + FUNCTION = "build_grid_config" + CATEGORY = 'WanVideoWrapper_QQ/image' + + def build_grid_config(self, enable, horizontal, vertical, frames): + return ({"enable": enable, + "horizontal": horizontal, + "vertical": vertical, + "frames": int(frames)}, + int(frames)) + + class ImageBlend_GPU_advanced(ImageBlend_GPU): """ Advanced image blend node. @@ -278,6 +329,12 @@ class ImageBlend_GPU_advanced(ImageBlend_GPU): to the given megapixels and also output it (plus its mask) through the image_low / mask_low outputs. - mask output: shows where the layer_image was placed on the background. + - layer_image is optional: when not connected, no layer is composited and + the background alone goes through high_MP / resize_to_32 / low_MP. + - enable_grid (optional): plug the "Image Blend Grid" node here to build an + image grid out of the frames (background only, or the blended result when + a layer is connected) instead of the plain output. high_MP / resize_to_32 / + low_MP are applied after the grid. """ NODE_NAME = "Image Blend GPU Advanced" @@ -294,7 +351,6 @@ class ImageBlend_GPU_advanced(ImageBlend_GPU): return { "required": { "background_image": ("IMAGE",), # BHWC tensor - "layer_image": ("IMAGE",), # BHWC tensor "blend_mode": (cls.BLEND_MODES,), "blend_corner": (cls.BLEND_CORNERS, {"default": "none"}), "opacity": ("INT", {"default": 100, "min": 0, "max": 100, "step": 1}), @@ -303,7 +359,11 @@ class ImageBlend_GPU_advanced(ImageBlend_GPU): "low_MP": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 100.0, "step": 0.01}), }, "optional": { + "layer_image": ("IMAGE",), # BHWC tensor (if omitted, only the + # background is processed: high_MP / + # resize_to_32 / low_MP still apply) "layer_mask": ("MASK",), # BHW tensor + "enable_grid": ("GRID_CONFIG",), # output of the Image Blend Grid node } } @@ -343,17 +403,37 @@ class ImageBlend_GPU_advanced(ImageBlend_GPU): x0 = (w - target_w) // 2 return tensor_bhwc[:, y0:y0 + target_h, x0:x0 + target_w, :], True - def image_blend_gpu_advanced(self, background_image, layer_image, - blend_mode, blend_corner, opacity, - layer_mask=None, resize_to_32=False, - high_MP=0.0, low_MP=0.0): + def image_blend_gpu_advanced(self, background_image, layer_image=None, + blend_mode="normal", blend_corner="none", + opacity=100, layer_mask=None, + resize_to_32=False, high_MP=0.0, low_MP=0.0, + enable_grid=None): device = background_image.device bg_bhwc = background_image.to(device, dtype=torch.float32) - layer_bhwc = layer_image.to(device, dtype=torch.float32) b_frames, bg_h, bg_w = bg_bhwc.shape[0], bg_bhwc.shape[1], bg_bhwc.shape[2] + + # --- Ensure RGB channel counts --- + if bg_bhwc.shape[3] == 4: + print("[WARNING] Background has alpha, removing it.") + bg_bhwc = bg_bhwc[..., :3] + elif bg_bhwc.shape[3] == 1: + bg_bhwc = bg_bhwc.repeat(1, 1, 1, 3) + + # --- No layer connected: pass the background through post-processing only --- + if layer_image is None: + if layer_mask is not None: + print("[WARNING] layer_mask is ignored because layer_image is not connected.") + print("[INFO] No layer_image connected, processing background only.") + output_bhwc = torch.clamp(bg_bhwc, 0.0, 1.0) + output_mask = torch.zeros((b_frames, bg_h, bg_w), dtype=torch.float32, device=device) + return self._grid_and_finalize(output_bhwc, output_mask, enable_grid, + resize_to_32, high_MP, low_MP) + + layer_bhwc = layer_image.to(device, dtype=torch.float32) + l_frames, l_h, l_w = layer_bhwc.shape[0], layer_bhwc.shape[1], layer_bhwc.shape[2] max_frames = max(b_frames, l_frames) @@ -367,13 +447,7 @@ class ImageBlend_GPU_advanced(ImageBlend_GPU): f"Repeating last layer frame to {max_frames}.") layer_bhwc = self._pad_batch_to(layer_bhwc, max_frames) - # --- Ensure RGB channel counts --- - if bg_bhwc.shape[3] == 4: - print("[WARNING] Background has alpha, removing it.") - bg_bhwc = bg_bhwc[..., :3] - elif bg_bhwc.shape[3] == 1: - bg_bhwc = bg_bhwc.repeat(1, 1, 1, 3) - + # --- Ensure RGB channel counts (layer) --- if layer_bhwc.shape[3] == 4: layer_bhwc = layer_bhwc[..., :3] elif layer_bhwc.shape[3] == 1: @@ -478,6 +552,101 @@ class ImageBlend_GPU_advanced(ImageBlend_GPU): output_mask[:, y0:y0 + crop_h, x0:x0 + crop_w] = mask_crop.squeeze(-1).clamp(0.0, 1.0) + # --- Optional grid, then shared post-processing: high_MP / resize_to_32 / low_MP --- + return self._grid_and_finalize(output_bhwc, output_mask, enable_grid, + resize_to_32, high_MP, low_MP) + + def _grid_and_finalize(self, output_bhwc, output_mask, enable_grid, + resize_to_32, high_MP, low_MP): + """Apply the optional grid (from the Image Blend Grid node), then + run the shared high_MP / resize_to_32 / low_MP post-processing.""" + if enable_grid is not None and str(enable_grid.get("enable", "false")).lower() != "false": + output_bhwc, output_mask = self._build_grid(output_bhwc, output_mask, enable_grid) + return self._finalize_output(output_bhwc, output_mask, + resize_to_32, high_MP, low_MP) + + def _build_grid(self, output_bhwc, output_mask, grid_cfg): + """ + Build a single-frame image grid (e.g. 3x2 / 3x3) from the output frames. + + - The source is normalized to 'frames' frames (last frame duplicated + if shorter, truncated if longer). + - One grid is built per output frame: frame t's cells are evenly + sampled over [t .. frames-1] (first cell = current frame, window + slides forward, last frame fills every cell at the end). Each cell + keeps the source aspect ratio. + - The canvas fits the grid block exactly, so no empty rows/columns + are left around the grid (e.g. 3x2 -> canvas is 2 cell-rows tall). + - The returned mask is white only inside the cell selected by + horizontal + vertical (e.g. left/top -> top-left cell), black + elsewhere, with one frame per output frame. + """ + layout = str(grid_cfg.get("enable", "false")).lower() + horizontal = str(grid_cfg.get("horizontal", "left")).lower() + vertical = str(grid_cfg.get("vertical", "top")).lower() + frames = int(grid_cfg.get("frames", 22)) + + try: + cols, rows = (int(v) for v in layout.split("x")) + except ValueError: + print(f"[WARNING] Invalid grid layout '{layout}', skipping grid.") + return output_bhwc, output_mask + if cols < 1 or rows < 1: + print(f"[WARNING] Invalid grid layout '{layout}', skipping grid.") + return output_bhwc, output_mask + + src_h, src_w = output_bhwc.shape[1], output_bhwc.shape[2] + cell_w = max(1, src_w // cols) + cell_h = max(1, int(round(cell_w * (src_h / src_w)))) + n_cells = cols * rows + + # Canvas fits the grid block exactly (e.g. 3x2 -> 2 cell-rows tall), + # so no empty rows/columns are left around the grid. + canvas_h = rows * cell_h + canvas_w = cols * cell_w + + # --- Normalize the source to exactly `frames` frames: + # - video shorter than `frames`: last frame is duplicated + # - video longer than `frames`: truncated to `frames` --- + src = self._pad_batch_to(output_bhwc, frames) + n_src = src.shape[0] + + # --- Build one grid per output frame (sliding storyboard): + # frame t's cells are evenly sampled over [t .. frames-1], + # so the first cell is always the current frame and the window + # slides forward until the last frame fills every cell --- + device = src.device + t_range = torch.arange(frames, device=device).unsqueeze(1) # (F, 1) + base = torch.linspace(0, frames - 1, n_cells, device=device).unsqueeze(0) # (1, C) + idx = (base + t_range).clamp(max=frames - 1).long() # (F, C) + sampled = src[idx] # (F, C, H, W, Ch) + sampled = sampled.reshape(frames * n_cells, src_h, src_w, sampled.shape[-1]) + cells = self._resize_bhwc(sampled, cell_h, cell_w) + cells = cells.reshape(frames, rows, cols, cell_h, cell_w, -1) + canvas = cells.permute(0, 1, 3, 2, 4, 5).reshape(frames, canvas_h, canvas_w, -1) + + # --- Mask only the cell selected by horizontal + vertical --- + # (e.g. left/top -> top-left cell; middle maps to the centre column/row, + # or the last of two when there is no true middle) + col_idx = {"left": 0, "middle": cols // 2, "right": cols - 1}.get(horizontal, 0) + row_idx = {"top": 0, "middle": rows // 2, "bottom": rows - 1}.get(vertical, 0) + my, mx = row_idx * cell_h, col_idx * cell_w + grid_mask = torch.zeros((1, canvas_h, canvas_w), dtype=torch.float32, device=device) + grid_mask[0, my:my + cell_h, mx:mx + cell_w] = 1.0 + grid_mask = grid_mask.repeat(frames, 1, 1) + + print(f"[INFO] Grid '{layout}': sliding per-frame grid, {n_cells} cells of " + f"{cell_w}x{cell_h} on a {canvas_w}x{canvas_h} canvas " + f"(source had {output_bhwc.shape[0]} frame(s), normalized to {n_src}); " + f"mask marks the {vertical}/{horizontal} cell " + f"(row {row_idx + 1}/{rows}, col {col_idx + 1}/{cols}).") + + return canvas, grid_mask + + def _finalize_output(self, output_bhwc, output_mask, + resize_to_32, high_MP, low_MP): + """Apply high_MP scaling, resize_to_32 cropping and the low_MP copy + to a composited (image, mask) pair and build the return tuple.""" # --- Optional high-megapixel scale of the main output (before resize_to_32) --- if high_MP > 0.0: src_h, src_w = output_bhwc.shape[1], output_bhwc.shape[2]