diff --git a/kie_api/grid.py b/kie_api/grid.py index 80afeaa..c70da37 100644 --- a/kie_api/grid.py +++ b/kie_api/grid.py @@ -8,6 +8,13 @@ import torch from .nanobanana import _log +VALID_GRIDS = { + "2x2": (2, 2), + "3x3": (3, 3), + "2x3": (2, 3), +} + + def slice_grid_tensor( image: torch.Tensor, grid: str, @@ -21,7 +28,7 @@ def slice_grid_tensor( Args: image: Input batch tensor with shape [B, H, W, 3]. - grid: "2x2" or "3x3". + grid: "2x2", "3x3", or "2x3". outer_crop_px: Pixels trimmed from each side before slicing. gutter_px: Pixels removed between tiles inside the grid. order: "row-major" or "column-major". @@ -31,7 +38,10 @@ def slice_grid_tensor( Returns: Tensor of tiles stacked on batch dimension. """ - rows, cols = (2, 2) if grid == "2x2" else (3, 3) + try: + rows, cols = VALID_GRIDS[grid] + except KeyError as exc: + raise RuntimeError(f"Unsupported grid '{grid}'. Choose from {', '.join(VALID_GRIDS)}.") from exc if image.dim() != 4 or image.shape[-1] != 3: raise RuntimeError("IMAGE input must have shape [B, H, W, 3].") @@ -103,4 +113,3 @@ def slice_grid_tensor( _log(log, f"Processed image {idx + 1}/{len(targets)}.") return torch.stack(all_tiles, dim=0) - diff --git a/nodes.py b/nodes.py index 96d3f17..ee0c3d0 100644 --- a/nodes.py +++ b/nodes.py @@ -330,7 +330,7 @@ Slice a single grid image into equal tiles, optionally cropping borders and remo Inputs: - IMAGE: Grid image batch (BHWC float32 0–1) -- grid: 2x2 or 3x3 layout +- grid: 2x2, 2x3, or 3x3 layout - outer_crop_px: Pixels to trim from each side before slicing - gutter_px: Pixels to remove between tiles inside the grid - order: Tile ordering (row-major or column-major) @@ -338,7 +338,7 @@ Inputs: - log: Console logging on/off Outputs: -- IMAGE: Tile batch (4 or 9 per processed image) +- IMAGE: Tile batch (4, 6, or 9 per processed image) """ @classmethod @@ -348,7 +348,7 @@ Outputs: "image": ("IMAGE",), }, "optional": { - "grid": ("COMBO", {"options": ["2x2", "3x3"], "default": "2x2"}), + "grid": ("COMBO", {"options": ["2x2", "2x3", "3x3"], "default": "2x2"}), "outer_crop_px": ("INT", {"default": 0, "min": 0, "max": 8192, "step": 1}), "gutter_px": ("INT", {"default": 0, "min": 0, "max": 8192, "step": 1}), "order": ("COMBO", {"options": ["row-major", "column-major"], "default": "row-major"}),