added an option to only use available memory

instead of assigning a fixed amount of memory to the system. 80% of the available memory are used for upscaling.
chaiNNer-org · Aug 14, 2023 · 5126eb8 · 5126eb8
1 parent 73d682c
commit 5126eb8
Show file tree

Hide file tree

Showing 8 changed files with 56 additions and 5 deletions.
diff --git a/backend/src/nodes/impl/pytorch/utils.py b/backend/src/nodes/impl/pytorch/utils.py
@@ -38,6 +38,7 @@ def to_pytorch_execution_options(options: ExecutionOptions):
         onnx_tensorrt_cache_path=options.onnx_tensorrt_cache_path,
         onnx_should_tensorrt_fp16=options.onnx_should_tensorrt_fp16,
         reserved_system_memory=options.reserved_system_memory,
+        available_memory_only=options.available_memory_only,
     )
 
 

diff --git a/backend/src/nodes/utils/exec_options.py b/backend/src/nodes/utils/exec_options.py
@@ -17,6 +17,7 @@ def __init__(
         onnx_tensorrt_cache_path: str,
         onnx_should_tensorrt_fp16: bool,
         reserved_system_memory: int,
+        available_memory_only: bool,
     ) -> None:
         self.__device = device
         self.__fp16 = fp16
@@ -28,6 +29,7 @@ def __init__(
         self.__onnx_tensorrt_cache_path = onnx_tensorrt_cache_path
         self.__onnx_should_tensorrt_fp16 = onnx_should_tensorrt_fp16
         self.__reserved_system_memory = reserved_system_memory
+        self.__available_memory_only = available_memory_only
 
         if (
             not os.path.exists(onnx_tensorrt_cache_path)
@@ -43,7 +45,7 @@ def __init__(
             f" {onnx_should_tensorrt_cache}, tensorrt_cache_path:"
             f" {onnx_tensorrt_cache_path}, should_tensorrt_fp16:"
             f" {onnx_should_tensorrt_fp16}, reserved_system_memory:"
-            f" {reserved_system_memory}"
+            f" {reserved_system_memory}, available_memory_only {available_memory_only}"
         )
 
     @property
@@ -88,9 +90,13 @@ def onnx_should_tensorrt_fp16(self):
     def reserved_system_memory(self):
         return self.__reserved_system_memory
 
+    @property
+    def available_memory_only(self):
+        return self.__available_memory_only
+
 
 __global_exec_options = ExecutionOptions(
-    "cpu", False, 0, 0, 0, "CPUExecutionProvider", False, "", False, 1024
+    "cpu", False, 0, 0, 0, "CPUExecutionProvider", False, "", False, 1024, False
 )
 
 
@@ -116,6 +122,7 @@ class JsonExecutionOptions(TypedDict):
     onnxTensorRtCachePath: str
     onnxShouldTensorRtFp16: bool
     reservedSystemMemory: int
+    availableMemoryOnly: bool
 
 
 def parse_execution_options(json: JsonExecutionOptions) -> ExecutionOptions:
@@ -130,4 +137,5 @@ def parse_execution_options(json: JsonExecutionOptions) -> ExecutionOptions:
         onnx_tensorrt_cache_path=json["onnxTensorRtCachePath"],
         onnx_should_tensorrt_fp16=json["onnxShouldTensorRtFp16"],
         reserved_system_memory=json["reservedSystemMemory"],
+        available_memory_only=json["availableMemoryOnly"],
     )
diff --git a/backend/src/packages/chaiNNer_pytorch/pytorch/processing/upscale_image.py b/backend/src/packages/chaiNNer_pytorch/pytorch/processing/upscale_image.py
@@ -59,10 +59,16 @@ def estimate():
 
             if is_arm_mac:
                 total_memory = psutil.virtual_memory().total
+                available_memory = psutil.virtual_memory().available
                 reserved_system_memory = options.reserved_system_memory * (1024**2)
                 pre_budget = int(total_memory - reserved_system_memory)
 
-                budget = max(total_memory * 0.2, min(pre_budget, total_memory * 0.8))
+                if options.available_memory_only:
+                    budget = available_memory * 0.8
+                else:
+                    budget = max(
+                        total_memory * 0.2, min(pre_budget, total_memory * 0.8)
+                    )
 
                 return MaxTileSize(
                     estimate_tile_size(

diff --git a/src/common/Backend.ts b/src/common/Backend.ts
@@ -78,6 +78,7 @@ export interface BackendExecutionOptions {
     onnxTensorRtCachePath: string;
     onnxShouldTensorRtFp16: boolean;
     reservedSystemMemory: number;
+    availableMemoryOnly: boolean;
 }
 export interface BackendRunRequest {
     data: BackendJsonNode[];

diff --git a/src/main/cli/run.ts b/src/main/cli/run.ts
@@ -139,6 +139,7 @@ const getExecutionOptions = (): BackendExecutionOptions => {
         onnxTensorRtCachePath: getOnnxTensorRtCacheLocation(app.getPath('userData')),
         onnxShouldTensorRtFp16: getSetting('onnx-should-tensorrt-fp16', false),
         reservedSystemMemory: getSetting('reserved-system-memory', 0),
+        availableMemoryOnly: getSetting('use-available-memory-only', false),
     };
 };
 

diff --git a/src/renderer/components/SettingsModal.tsx b/src/renderer/components/SettingsModal.tsx
@@ -271,7 +271,8 @@ const AppearanceSettings = memo(() => {
 });
 
 const EnvironmentSettings = memo(() => {
-    const { useStartupTemplate, useReservedSystemMemory } = useContext(SettingsContext);
+    const { useStartupTemplate, useReservedSystemMemory, useAvailableMemoryOnly } =
+        useContext(SettingsContext);
 
     const [startupTemplate, setStartupTemplate] = useStartupTemplate;
 
@@ -299,6 +300,7 @@ const EnvironmentSettings = memo(() => {
     }, [startupTemplate, lastDirectory, setStartupTemplate]);
 
     const [reservedSystemMemory, setReservedSystemMemory] = useReservedSystemMemory;
+    const [availableMemoryOnly, setAvailableMemoryOnly] = useAvailableMemoryOnly;
 
     // Maximum amount reserved for the system is 80 % of the total memory
     const calculateMaxValue = () => (totalMemory / 1024 ** 2) * 0.8;
@@ -319,6 +321,7 @@ const EnvironmentSettings = memo(() => {
                 >
                     <InputGroup>
                         <NumberInput
+                            isDisabled={availableMemoryOnly}
                             max={calculateMaxValue()}
                             min={calculateMinValue()}
                             step={512}
@@ -336,7 +339,16 @@ const EnvironmentSettings = memo(() => {
                                 paddingRight="3.7rem"
                                 textAlign="right"
                             />
-                            <InputRightElement marginRight="1.5rem">MB</InputRightElement>
+                            {availableMemoryOnly ? (
+                                <InputRightElement
+                                    marginRight="1.5rem"
+                                    opacity="0.4"
+                                >
+                                    MB
+                                </InputRightElement>
+                            ) : (
+                                <InputRightElement marginRight="1.5rem">MB</InputRightElement>
+                            )}
                             <NumberInputStepper>
                                 <NumberIncrementStepper />
                                 <NumberDecrementStepper />
@@ -347,6 +359,20 @@ const EnvironmentSettings = memo(() => {
             ) : (
                 []
             )}
+
+            {isArmMac ? (
+                <Toggle
+                    description="Instead of reserving a fixed amount of total system memory, chaiNNer will use at maximum 80% of the memory available for upscaling. (Overwrites above settings.)"
+                    title="Only use available memory"
+                    value={availableMemoryOnly}
+                    onToggle={() => {
+                        setAvailableMemoryOnly((prev) => !prev);
+                    }}
+                />
+            ) : (
+                []
+            )}
+
             <SettingsItem
                 description="Set a chain template to use by default when chaiNNer starts up."
                 title="Startup Template"

diff --git a/src/renderer/contexts/SettingsContext.tsx b/src/renderer/contexts/SettingsContext.tsx
@@ -28,6 +28,7 @@ interface Settings {
     ];
     useStartupTemplate: GetSetState<string>;
     useReservedSystemMemory: GetSetState<number>;
+    useAvailableMemoryOnly: GetSetState<boolean>;
     useSelectTheme: GetSetState<string>;
     useAnimateChain: GetSetState<boolean>;
     useExperimentalFeatures: GetSetState<boolean>;
@@ -72,6 +73,9 @@ export const SettingsProvider = memo(({ children }: React.PropsWithChildren<unkn
     const useReservedSystemMemory = useMemoArray(
         useLocalStorage('reserved-system-memory', calculatedReservedMemory())
     );
+    const useAvailableMemoryOnly = useMemoArray(
+        useLocalStorage('use-available-memory-only', false)
+    );
 
     const { setColorMode } = useColorMode();
     const [selectThemeColor] = useSelectTheme;
@@ -124,6 +128,7 @@ export const SettingsProvider = memo(({ children }: React.PropsWithChildren<unkn
         useCheckUpdOnStrtUp,
         useStartupTemplate,
         useReservedSystemMemory,
+        useAvailableMemoryOnly,
         useSelectTheme,
         useAnimateChain,
         useExperimentalFeatures,

diff --git a/src/renderer/hooks/useBackendExecutionOptions.ts b/src/renderer/hooks/useBackendExecutionOptions.ts
@@ -17,6 +17,7 @@ export const useBackendExecutionOptions = (): BackendExecutionOptions => {
         useOnnxShouldTensorRtCache,
         useOnnxShouldTensorRtFp16,
         useReservedSystemMemory,
+        useAvailableMemoryOnly,
     } = useContext(SettingsContext);
 
     const [isCpu] = useIsCpu;
@@ -39,6 +40,7 @@ export const useBackendExecutionOptions = (): BackendExecutionOptions => {
 
     const [onnxShouldTensorRtFp16] = useOnnxShouldTensorRtFp16;
     const [reservedSystemMemory] = useReservedSystemMemory;
+    const [availableMemoryOnly] = useAvailableMemoryOnly;
 
     return {
         isCpu,
@@ -51,5 +53,6 @@ export const useBackendExecutionOptions = (): BackendExecutionOptions => {
         onnxTensorRtCachePath,
         onnxShouldTensorRtFp16,
         reservedSystemMemory,
+        availableMemoryOnly,
     };
 };