Fix custom tile size type (#2870)

* Fix custom tile size type * None -> 0
chaiNNer-org · May 18, 2024 · e01f8bc · e01f8bc
1 parent a971d74
commit e01f8bc
Show file tree

Hide file tree

Showing 5 changed files with 8 additions and 14 deletions.
diff --git a/backend/src/packages/chaiNNer_ncnn/ncnn/processing/upscale_image.py b/backend/src/packages/chaiNNer_ncnn/ncnn/processing/upscale_image.py
@@ -202,7 +202,7 @@ def upscale_image_node(
     img: np.ndarray,
     model: NcnnModelWrapper,
     tile_size: TileSize,
-    custom_tile_size: int | None,
+    custom_tile_size: int,
     separate_alpha: bool,
 ) -> np.ndarray:
     settings = get_settings(context)
@@ -219,9 +219,7 @@ def upscale(i: np.ndarray) -> np.ndarray:
             model,
             model.model.layers[0].outputs[0],
             model.model.layers[-1].outputs[0],
-            TileSize(custom_tile_size)
-            if tile_size == CUSTOM and custom_tile_size is not None
-            else tile_size,
+            TileSize(custom_tile_size) if tile_size == CUSTOM else tile_size,
         )
         if ic == 3:
             i = cv2.cvtColor(i, cv2.COLOR_RGB2BGR)

diff --git a/backend/src/packages/chaiNNer_ncnn/ncnn/utility/interpolate_models.py b/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, None, False)
+    result = upscale_image_node(context, fake_img, interp, NO_TILING, 0, False)
 
     mean_color = np.mean(result)
     del result

diff --git a/backend/src/packages/chaiNNer_onnx/onnx/processing/upscale_image.py b/backend/src/packages/chaiNNer_onnx/onnx/processing/upscale_image.py
@@ -109,7 +109,7 @@ def upscale_image_node(
     img: np.ndarray,
     model: OnnxModel,
     tile_size: TileSize,
-    custom_tile_size: int | None,
+    custom_tile_size: int,
     separate_alpha: bool,
 ) -> np.ndarray:
     settings = get_settings(context)
@@ -141,9 +141,7 @@ def upscale_image_node(
         lambda i: upscale(
             i,
             session,
-            TileSize(custom_tile_size)
-            if tile_size == CUSTOM and custom_tile_size is not None
-            else tile_size,
+            TileSize(custom_tile_size) if tile_size == CUSTOM else tile_size,
             change_shape,
             exact_size,
         ),

diff --git a/backend/src/packages/chaiNNer_onnx/onnx/utility/interpolate_models.py b/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, None, False)
+    result = upscale_image_node(context, fake_img, model, NO_TILING, 0, False)
 
     mean_color = np.mean(result)
     del result

diff --git a/backend/src/packages/chaiNNer_pytorch/pytorch/processing/upscale_image.py b/backend/src/packages/chaiNNer_pytorch/pytorch/processing/upscale_image.py
@@ -261,7 +261,7 @@ def upscale_image_node(
     use_custom_scale: bool,
     custom_scale: int,
     tile_size: TileSize,
-    custom_tile_size: int | None,
+    custom_tile_size: int,
     separate_alpha: bool,
 ) -> np.ndarray:
     exec_options = get_settings(context)
@@ -286,9 +286,7 @@ def inner_upscale(img: np.ndarray) -> np.ndarray:
             lambda i: upscale(
                 i,
                 model,
-                TileSize(custom_tile_size)
-                if tile_size == CUSTOM and custom_tile_size is not None
-                else tile_size,
+                TileSize(custom_tile_size) if tile_size == CUSTOM else tile_size,
                 exec_options,
                 context,
             ),