#2654 add custom tiling option

backend/src/nodes/impl/upscale/auto_split_tiles.py

@@ -37,6 +37,7 @@ def estimate_tile_size(
 ESTIMATE = TileSize(0)
 NO_TILING = TileSize(-1)
 MAX_TILE_SIZE = TileSize(-2)
+CUSTOM = TileSize(-3)
 TILE_SIZE_256 = TileSize(256)
 
 

backend/src/nodes/properties/inputs/generic_inputs.py

@@ -17,7 +17,13 @@
 from ...impl.blend import BlendMode
 from ...impl.color.color import Color
 from ...impl.image_utils import FillColor
-from ...impl.upscale.auto_split_tiles import TileSize
+from ...impl.upscale.auto_split_tiles import (
+    CUSTOM,
+    ESTIMATE,
+    MAX_TILE_SIZE,
+    NO_TILING,
+    TileSize,
+)
 from ...utils.format import format_color_with_channels
 from ...utils.seed import Seed
 from ...utils.utils import (
@@ -527,14 +533,16 @@ def TileSizeDropdown(
 ) -> DropDownInput:
     options = []
     if estimate:
-        options.append({"option": "Auto (estimate)", "value": 0})
+        options.append({"option": "Auto (estimate)", "value": ESTIMATE})
 
-    options.append({"option": "Maximum", "value": -2})
-    options.append({"option": "No Tiling", "value": -1})
+    options.append({"option": "Maximum", "value": MAX_TILE_SIZE})
+    options.append({"option": "No Tiling", "value": NO_TILING})
 
-    for size in [128, 192, 256, 384, 512, 768, 1024, 2048, 4096]:
+    for size in [128, 192, 256, 384, 512, 768, 1024, 1536, 2048, 3072, 4096]:
         options.append({"option": str(size), "value": size})
 
+    options.append({"option": "Custom", "value": CUSTOM})
+
     return DropDownInput(
         input_type="TileSize",
         label=label,

backend/src/packages/chaiNNer_ncnn/ncnn/processing/upscale_image.py

@@ -16,11 +16,13 @@
 from sanic.log import logger
 
 from api import NodeContext
-from nodes.groups import Condition, if_group
+from nodes.groups import Condition, if_enum_group, if_group
 from nodes.impl.ncnn.auto_split import ncnn_auto_split
 from nodes.impl.ncnn.model import NcnnModelWrapper
 from nodes.impl.ncnn.session import get_ncnn_net
 from nodes.impl.upscale.auto_split_tiles import (
+    CUSTOM,
+    TILE_SIZE_256,
     TileSize,
     estimate_tile_size,
     parse_tile_size_input,
@@ -31,6 +33,7 @@
     BoolInput,
     ImageInput,
     NcnnModelInput,
+    NumberInput,
     TileSizeDropdown,
 )
 from nodes.properties.outputs import ImageOutput
@@ -158,6 +161,18 @@ def estimate():
             "Generally it's recommended to use the largest tile size possible for best performance (with the ideal scenario being no tiling at all), but depending on the model and image size, this may not be possible.",
             "If you are having issues with the automatic mode, you can manually select a tile size, or set a memory budget limit. On certain machines, a very small tile size such as 256 or 128 might be required for it to work at all.",
         ),
+        if_enum_group(2, CUSTOM)(
+            NumberInput(
+                "Custom Tile Size",
+                minimum=1,
+                maximum=None,
+                default=TILE_SIZE_256,
+                precision=0,
+                controls_step=1,
+                unit="px",
+                has_handle=False,
+            )
+        ),
         if_group(
             Condition.type(1, "Image { channels: 4 } ")
             & (
@@ -189,6 +204,7 @@ def upscale_image_node(
     img: np.ndarray,
     model: NcnnModelWrapper,
     tile_size: TileSize,
+    custom_tile_size: int | None,
     separate_alpha: bool,
 ) -> np.ndarray:
     settings = get_settings(context)
@@ -205,7 +221,9 @@ def upscale(i: np.ndarray) -> np.ndarray:
             model,
             model.model.layers[0].outputs[0],
             model.model.layers[-1].outputs[0],
-            tile_size,
+            TileSize(custom_tile_size)
+            if tile_size == CUSTOM and custom_tile_size is not None
+            else tile_size,
         )
         if ic == 3:
             i = cv2.cvtColor(i, cv2.COLOR_RGB2BGR)

backend/src/packages/chaiNNer_ncnn/ncnn/utility/interpolate_models.py

@@ -14,7 +14,7 @@
 
 def check_will_upscale(context: NodeContext, interp: NcnnModelWrapper):
     fake_img = np.ones((3, 3, 3), dtype=np.float32, order="F")
-    result = upscale_image_node(context, fake_img, interp, NO_TILING, False)
+    result = upscale_image_node(context, fake_img, interp, NO_TILING, None, False)
 
     mean_color = np.mean(result)
     del result

backend/src/packages/chaiNNer_onnx/onnx/processing/upscale_image.py

@@ -5,12 +5,13 @@
 from sanic.log import logger
 
 from api import NodeContext
-from nodes.groups import Condition, if_group
+from nodes.groups import Condition, if_enum_group, if_group
 from nodes.impl.onnx.auto_split import onnx_auto_split
 from nodes.impl.onnx.model import OnnxModel
 from nodes.impl.onnx.session import get_onnx_session
 from nodes.impl.onnx.utils import get_input_shape, get_output_shape
 from nodes.impl.upscale.auto_split_tiles import (
+    CUSTOM,
     TILE_SIZE_256,
     TileSize,
     parse_tile_size_input,
@@ -20,6 +21,7 @@
 from nodes.properties.inputs import (
     BoolInput,
     ImageInput,
+    NumberInput,
     OnnxGenericModelInput,
     TileSizeDropdown,
 )
@@ -75,6 +77,18 @@ def estimate():
             "ONNX upscaling does not support an automatic mode, meaning you may need to"
             " manually select a tile size for it to work.",
         ),
+        if_enum_group(2, CUSTOM)(
+            NumberInput(
+                "Custom Tile Size",
+                minimum=1,
+                maximum=None,
+                default=TILE_SIZE_256,
+                precision=0,
+                controls_step=1,
+                unit="px",
+                has_handle=False,
+            )
+        ),
         if_group(Condition.type(1, "Image { channels: 4 } "))(
             BoolInput("Separate Alpha", default=False).with_docs(
                 "Upscale alpha separately from color. Enabling this option will cause the alpha of"
@@ -97,6 +111,7 @@ def upscale_image_node(
     img: np.ndarray,
     model: OnnxModel,
     tile_size: TileSize,
+    custom_tile_size: int | None,
     separate_alpha: bool,
 ) -> np.ndarray:
     settings = get_settings(context)
@@ -125,6 +140,14 @@ def upscale_image_node(
         img,
         in_nc,
         out_nc,
-        lambda i: upscale(i, session, tile_size, change_shape, exact_size),
+        lambda i: upscale(
+            i,
+            session,
+            TileSize(custom_tile_size)
+            if tile_size == CUSTOM and custom_tile_size is not None
+            else tile_size,
+            change_shape,
+            exact_size,
+        ),
         separate_alpha,
     )

backend/src/packages/chaiNNer_onnx/onnx/utility/interpolate_models.py

@@ -49,7 +49,7 @@ def perform_interp(
 
 def check_will_upscale(context: NodeContext, model: OnnxModel):
     fake_img = np.ones((3, 3, 3), dtype=np.float32, order="F")
-    result = upscale_image_node(context, fake_img, model, NO_TILING, False)
+    result = upscale_image_node(context, fake_img, model, NO_TILING, None, False)
 
     mean_color = np.mean(result)
     del result

backend/src/packages/chaiNNer_pytorch/pytorch/processing/upscale_image.py

@@ -7,10 +7,12 @@
 from spandrel import ImageModelDescriptor, ModelTiling
 
 from api import KeyInfo, NodeContext, Progress
-from nodes.groups import Condition, if_group
+from nodes.groups import Condition, if_enum_group, if_group
 from nodes.impl.pytorch.auto_split import pytorch_auto_split
 from nodes.impl.upscale.auto_split_tiles import (
+    CUSTOM,
     NO_TILING,
+    TILE_SIZE_256,
     TileSize,
     estimate_tile_size,
     parse_tile_size_input,
@@ -169,6 +171,18 @@ def estimate():
                 hint=True,
             ),
         ),
+        if_enum_group(2, CUSTOM)(
+            NumberInput(
+                "Custom Tile Size",
+                minimum=1,
+                maximum=None,
+                default=TILE_SIZE_256,
+                precision=0,
+                controls_step=1,
+                unit="px",
+                has_handle=False,
+            ).with_id(6),
+        ),
         if_group(
             Condition.type(1, "Image { channels: 4 } ")
             & (
@@ -243,6 +257,7 @@ def upscale_image_node(
     use_custom_scale: bool,
     custom_scale: int,
     tile_size: TileSize,
+    custom_tile_size: int | None,
     separate_alpha: bool,
 ) -> np.ndarray:
     exec_options = get_settings(context)
@@ -262,7 +277,15 @@ def inner_upscale(img: np.ndarray) -> np.ndarray:
             img,
             in_nc,
             out_nc,
-            lambda i: upscale(i, model, tile_size, exec_options, context),
+            lambda i: upscale(
+                i,
+                model,
+                TileSize(custom_tile_size)
+                if tile_size == CUSTOM and custom_tile_size is not None
+                else tile_size,
+                exec_options,
+                context,
+            ),
             separate_alpha,
             clip=False,  # pytorch_auto_split already does clipping internally
         )