Set the amount of memory to use for upscaling

CPU upscaling & upscaling on Apple Silicon (CPU & GPU)

A value between 20% and 80% of the freely available memory can be chosen for upscaling. If desired, instead of the freely available memory, these values can be applied to the entire available RAM. If a user chooses to do so, a warning will be presented and when upscaling the  settings and amount of used RAM will be logged.

For GPU upscaling the amount of freely available can be set. This setting is only available on Windows and Linux.

MaxTileSize for NCNN on Apple Silicon has been added.

add description for memory limit to upscaling node

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

@@ -25,7 +25,7 @@ def estimate_tile_size(
     required_mem = f"{mem_required_estimation/GB_AMT:.2f}"
     budget_mem = f"{budget/GB_AMT:.2f}"
     logger.info(
-        f"Estimating memory required: {required_mem} GB, {budget_mem} GB free."
+        f"Estimating memory required: {required_mem} GB, {budget_mem} GB available."
         f" Estimated tile size: {tile_size}"
     )
 

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

@@ -4,6 +4,7 @@
 
 import cv2
 import numpy as np
+import psutil
 
 try:
     from ncnn_vulkan import ncnn
@@ -34,7 +35,7 @@
 )
 from nodes.properties.outputs import ImageOutput
 from nodes.utils.utils import get_h_w_c
-from system import is_mac
+from system import is_arm_mac, is_mac
 
 from ...settings import get_settings
 from .. import processing_group
@@ -80,13 +81,39 @@ def upscale_impl(
             vkdev = ncnn.get_gpu_device(settings.gpu_index)
 
             def estimate_gpu():
-                if is_mac:
+                if is_arm_mac:
+                    memory_for_upscaling = get_settings.memory_for_upscaling / 100
+
+                    if get_settings.is_system_memory:
+                        available_memory = psutil.virtual_memory().total
+
+                        logger.info(
+                            f"Memory limit set to {memory_for_upscaling * 100}% of"
+                            " total system memory."
+                            f" ({available_memory / (1024 ** 3)} GB)"
+                        )
+                    else:
+                        available_memory = psutil.virtual_memory().available
+
+                elif is_mac:
                     # the actual estimate frequently crashes on mac, so we just use 256
                     return MaxTileSize(256)
+                else:
+                    available_memory = vkdev.get_heap_budget() * 1024 * 1024
+                    memory_for_upscaling = get_settings.memory_for_upscaling_gpu / 100
+
+                budget = int(
+                    max(
+                        available_memory * 0.2,
+                        min(
+                            available_memory * memory_for_upscaling,
+                            available_memory * 0.8,
+                        ),
+                    )
+                )
 
-                heap_budget = vkdev.get_heap_budget() * 1024 * 1024 * 0.8
                 return MaxTileSize(
-                    estimate_tile_size(heap_budget, model.model.bin_length, img, 4)
+                    estimate_tile_size(budget, model.model.bin_length, img, 4)
                 )
 
             with ncnn_allocators(vkdev) as (
@@ -130,19 +157,30 @@ def estimate_cpu():
 @processing_group.register(
     schema_id="chainner:ncnn:upscale_image",
     name="Upscale Image",
-    description="Upscale an image with NCNN. Unlike PyTorch, NCNN has GPU support on all devices, assuming your drivers support Vulkan. \
-            Select a manual number of tiles if you are having issues with the automatic mode.",
+    description=(
+        "Upscale an image with NCNN. Unlike PyTorch, NCNN has GPU support on all"
+        " devices, assuming your drivers support Vulkan. Select a lower memory limit or"
+        " manually adjust the number of tiles if you are experiencing issues with the"
+        " automatic mode."
+    ),
     icon="NCNN",
     inputs=[
         ImageInput().with_id(1),
         NcnnModelInput().with_id(0),
         TileSizeDropdown()
         .with_id(2)
         .with_docs(
-            "Tiled upscaling is used to allow large images to be upscaled without hitting memory limits.",
-            "This works by splitting the image into tiles (with overlap), upscaling each tile individually, and seamlessly recombining them.",
-            "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. On certain machines, a very small tile size such as 256 or 128 might be required for it to work at all.",
+            "Tiled upscaling is used to allow large images to be upscaled without"
+            " hitting memory limits.",
+            "This works by splitting the image into tiles (with overlap), upscaling"
+            " each tile individually, and seamlessly recombining them.",
+            "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 either set a"
+            " lower memory limit or manually select a tile size. On certain machines, a"
+            " very small tile size such as 256 or 128 might be required for it to work"
+            " at all.",
         ),
         if_group(
             Condition.type(1, "Image { channels: 4 } ")
@@ -154,13 +192,14 @@ def estimate_cpu():
             )
         )(
             BoolInput("Separate Alpha", default=False).with_docs(
-                "Upscale alpha separately from color. Enabling this option will cause the alpha of"
-                " the upscaled image to be less noisy and more accurate to the alpha of the original"
-                " image, but the image may suffer from dark borders near transparency edges"
-                " (transition from fully transparent to fully opaque).",
-                "Whether enabling this option will improve the upscaled image depends on the original"
-                " image. We generally recommend this option for images with smooth transitions between"
-                " transparent and opaque regions.",
+                "Upscale alpha separately from color. Enabling this option will cause"
+                " the alpha of the upscaled image to be less noisy and more accurate to"
+                " the alpha of the original image, but the image may suffer from dark"
+                " borders near transparency edges (transition from fully transparent to"
+                " fully opaque).",
+                "Whether enabling this option will improve the upscaled image depends"
+                " on the original image. We generally recommend this option for images"
+                " with smooth transitions between transparent and opaque regions.",
             )
         ),
     ],

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

@@ -3,6 +3,7 @@
 from typing import Tuple
 
 import numpy as np
+import psutil
 import torch
 from sanic.log import logger
 
@@ -45,22 +46,44 @@ def upscale(
         device = options.device
 
         def estimate():
+            element_size = 2 if use_fp16 else 4
+            model_bytes = sum(p.numel() * element_size for p in model.parameters())
+
             if "cuda" in device.type:
                 mem_info: Tuple[int, int] = torch.cuda.mem_get_info(device)  # type: ignore
-                free, _total = mem_info
-                element_size = 2 if use_fp16 else 4
-                model_bytes = sum(p.numel() * element_size for p in model.parameters())
-                budget = int(free * 0.8)
-
-                return MaxTileSize(
-                    estimate_tile_size(
-                        budget,
-                        model_bytes,
-                        img,
-                        element_size,
+                available_memory, _total = mem_info
+                memory_for_upscaling = get_settings.memory_for_upscaling_gpu / 100
+            else:
+                memory_for_upscaling = get_settings.memory_for_upscaling / 100
+
+                if get_settings.is_system_memory:
+                    available_memory = psutil.virtual_memory().total
+
+                    logger.info(
+                        f"Memory limit set to {memory_for_upscaling * 100}% of total"
+                        f" system memory. ({available_memory / (1024 ** 3)} GB)"
                     )
+                else:
+                    available_memory = psutil.virtual_memory().available
+
+            budget = int(
+                max(
+                    available_memory * 0.2,
+                    min(
+                        available_memory * memory_for_upscaling,
+                        available_memory * 0.8,
+                    ),
                 )
-            return MaxTileSize()
+            )
+
+            return MaxTileSize(
+                estimate_tile_size(
+                    budget,
+                    model_bytes,
+                    img,
+                    element_size,
+                )
+            )
 
         # Disable tiling for SCUNet
         upscale_tile_size = tile_size
@@ -83,8 +106,9 @@ def estimate():
     schema_id="chainner:pytorch:upscale_image",
     name="Upscale Image",
     description=(
-        "Upscales an image using a PyTorch Super-Resolution model. Select a"
-        " manual number of tiles if you are having issues with the automatic mode. "
+        "Upscales an image using a PyTorch Super-Resolution model. Select a lower"
+        " memory limit or manually adjust the number of tiles if you are experiencing"
+        " issues with the automatic mode."
     ),
     icon="PyTorch",
     inputs=[
@@ -105,9 +129,10 @@ 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. Sometimes, a manually selected tile size may be"
-                " faster than what the automatic mode picks.",
+                "If you are having issues with the automatic mode, you can either"
+                " select a lower memory limit or manually select a tile size."
+                " Sometimes, a manually selected tile size may be faster than what the"
+                " automatic mode picks.",
                 hint=True,
             )
         ),
@@ -123,13 +148,14 @@ def estimate():
             )
         )(
             BoolInput("Separate Alpha", default=False).with_docs(
-                "Upscale alpha separately from color. Enabling this option will cause the alpha of"
-                " the upscaled image to be less noisy and more accurate to the alpha of the original"
-                " image, but the image may suffer from dark borders near transparency edges"
-                " (transition from fully transparent to fully opaque).",
-                "Whether enabling this option will improve the upscaled image depends on the original"
-                " image. We generally recommend this option for images with smooth transitions between"
-                " transparent and opaque regions.",
+                "Upscale alpha separately from color. Enabling this option will cause"
+                " the alpha of the upscaled image to be less noisy and more accurate to"
+                " the alpha of the original image, but the image may suffer from dark"
+                " borders near transparency edges (transition from fully transparent to"
+                " fully opaque).",
+                "Whether enabling this option will improve the upscaled image depends"
+                " on the original image. We generally recommend this option for images"
+                " with smooth transitions between transparent and opaque regions.",
             )
         ),
     ],

src/common/common-types.ts

@@ -101,6 +101,8 @@ export interface NumberInput extends InputBase {
     readonly def: number;
     readonly min?: number | null;
     readonly max?: number | null;
+    readonly step: number | null;
+    readonly width: number;
     readonly precision: number;
     readonly controlsStep: number;
     readonly unit?: string | null;
@@ -332,6 +334,9 @@ export interface NumberSetting extends SettingBase {
     readonly type: 'number';
     readonly min: number;
     readonly max: number;
+    readonly step?: number;
+    readonly width?: number;
+    readonly unit?: string;
     readonly default: number;
 }
 

src/common/env.ts

@@ -31,3 +31,5 @@ export const sanitizedEnv = env;
 export const getCacheLocation = (userDataPath: string, cacheKey: string) => {
     return path.join(userDataPath, '/cache/', cacheKey);
 };
+
+export const totalMemory = os.totalmem();