Add PyTorch memory budget limit setting (#2472)

* add pytorch memory budget setting * Make memory budget setting actually do something * specify VRAM * also use limit for CPU * change condition * change desc
chaiNNer-org · Jan 21, 2024 · 213f734 · 213f734
1 parent 3f6ef53
commit 213f734
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 4 deletions.
diff --git a/backend/src/packages/chaiNNer_pytorch/pytorch/processing/upscale_image.py b/backend/src/packages/chaiNNer_pytorch/pytorch/processing/upscale_image.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import numpy as np
+import psutil
 import torch
 from sanic.log import logger
 from spandrel import ImageModelDescriptor, ModelTiling
@@ -48,15 +49,32 @@ def upscale(
             tile_size = NO_TILING
 
         def estimate():
+            element_size = 2 if use_fp16 else 4
+            model_bytes = sum(
+                p.numel() * element_size for p in model.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.model.parameters()
-                )
+                if options.budget_limit > 0:
+                    free = min(options.budget_limit * 1024**3, free)
                 budget = int(free * 0.8)
 
+                return MaxTileSize(
+                    estimate_tile_size(
+                        budget,
+                        model_bytes,
+                        img,
+                        element_size,
+                    )
+                )
+            elif device.type == "cpu":
+                free = psutil.virtual_memory().available
+                if options.budget_limit > 0:
+                    free = min(options.budget_limit * 1024**3, free)
+                budget = int(free * 0.8)
                 return MaxTileSize(
                     estimate_tile_size(
                         budget,

diff --git a/backend/src/packages/chaiNNer_pytorch/settings.py b/backend/src/packages/chaiNNer_pytorch/settings.py
@@ -3,7 +3,7 @@
 import torch
 from sanic.log import logger
 
-from api import DropdownSetting, NodeContext, ToggleSetting
+from api import DropdownSetting, NodeContext, NumberSetting, ToggleSetting
 from gpu import get_nvidia_helper
 from system import is_arm_mac
 
@@ -66,12 +66,24 @@
     ),
 )
 
+package.add_setting(
+    NumberSetting(
+        label="Memory Budget Limit (GiB)",
+        key="budget_limit",
+        description="Maximum memory (VRAM if GPU, RAM if CPU) to use for PyTorch inference. 0 means no limit. Memory usage measurement is not completely accurate yet; you may need to significantly adjust this budget limit via trial-and-error if it's not having the effect you want.",
+        default=0,
+        min=0,
+        max=1024**2,
+    )
+)
+
 
 @dataclass(frozen=True)
 class PyTorchSettings:
     use_cpu: bool
     use_fp16: bool
     gpu_index: int
+    budget_limit: int
 
     # PyTorch 2.0 does not support FP16 when using CPU
     def __post_init__(self):
@@ -111,4 +123,5 @@ def get_settings(context: NodeContext) -> PyTorchSettings:
         use_cpu=settings.get_bool("use_cpu", False),
         use_fp16=settings.get_bool("use_fp16", False),
         gpu_index=settings.get_int("gpu_index", 0, parse_str=True),
+        budget_limit=settings.get_int("budget_limit", 0, parse_str=True),
     )