From 70ca2a63ee8db7476283be0d3f5572c30742c26b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Toni=20F=C3=B6rster?= Date: Tue, 15 Aug 2023 17:55:09 +0200 Subject: [PATCH] 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 --- .../nodes/impl/upscale/auto_split_tiles.py | 2 +- .../ncnn/processing/upscale_image.py | 73 ++++-- .../pytorch/processing/upscale_image.py | 74 ++++-- src/common/common-types.ts | 5 + src/common/env.ts | 2 + src/renderer/components/SettingsModal.tsx | 215 ++++++++++++++++-- .../components/settings/components.tsx | 49 ++-- src/renderer/contexts/SettingsContext.tsx | 9 + 8 files changed, 351 insertions(+), 78 deletions(-) diff --git a/backend/src/nodes/impl/upscale/auto_split_tiles.py b/backend/src/nodes/impl/upscale/auto_split_tiles.py index 7c055cb679..eb0a2b38e8 100644 --- a/backend/src/nodes/impl/upscale/auto_split_tiles.py +++ b/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}" ) diff --git a/backend/src/packages/chaiNNer_ncnn/ncnn/processing/upscale_image.py b/backend/src/packages/chaiNNer_ncnn/ncnn/processing/upscale_image.py index dc4e1a6196..2cd8fc34cd 100644 --- a/backend/src/packages/chaiNNer_ncnn/ncnn/processing/upscale_image.py +++ b/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,8 +157,12 @@ 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), @@ -139,10 +170,17 @@ def estimate_cpu(): 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.", ) ), ], diff --git a/backend/src/packages/chaiNNer_pytorch/pytorch/processing/upscale_image.py b/backend/src/packages/chaiNNer_pytorch/pytorch/processing/upscale_image.py index b8c066676b..5277ce52b7 100644 --- a/backend/src/packages/chaiNNer_pytorch/pytorch/processing/upscale_image.py +++ b/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.", ) ), ], diff --git a/src/common/common-types.ts b/src/common/common-types.ts index d2f7c6050b..d93846f461 100644 --- a/src/common/common-types.ts +++ b/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; @@ -330,6 +332,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; } diff --git a/src/common/env.ts b/src/common/env.ts index 18ee934481..4640414cfd 100644 --- a/src/common/env.ts +++ b/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(); diff --git a/src/renderer/components/SettingsModal.tsx b/src/renderer/components/SettingsModal.tsx index 9d2c25491a..e8fdc7f2c8 100644 --- a/src/renderer/components/SettingsModal.tsx +++ b/src/renderer/components/SettingsModal.tsx @@ -1,5 +1,11 @@ import { LinkIcon, SettingsIcon, SmallCloseIcon } from '@chakra-ui/icons'; import { + AlertDialog, + AlertDialogBody, + AlertDialogContent, + AlertDialogFooter, + AlertDialogHeader, + AlertDialogOverlay, Button, HStack, Icon, @@ -7,6 +13,7 @@ import { Input, InputGroup, InputLeftElement, + InputRightElement, Modal, ModalBody, ModalCloseButton, @@ -27,12 +34,12 @@ import { } from '@chakra-ui/react'; import { produce } from 'immer'; import path from 'path'; -import { memo, useCallback, useState } from 'react'; +import { memo, useCallback, useRef, useState } from 'react'; import { BsFillPencilFill, BsPaletteFill } from 'react-icons/bs'; import { FaPython, FaTools } from 'react-icons/fa'; import { useContext } from 'use-context-selector'; import { SettingKey, SettingValue } from '../../common/common-types'; -import { isMac } from '../../common/env'; +import { isArmMac, isMac, totalMemory } from '../../common/env'; import { ipcRenderer } from '../../common/safeIpc'; import { BackendContext } from '../contexts/BackendContext'; import { SettingsContext } from '../contexts/SettingsContext'; @@ -115,7 +122,12 @@ const AppearanceSettings = memo(() => { }); const EnvironmentSettings = memo(() => { - const { useStartupTemplate } = useContext(SettingsContext); + const { + useStartupTemplate, + useMemoryForUpscaling, + useIsSystemMemory, + useMemoryForUpscalingGPU, + } = useContext(SettingsContext); const [startupTemplate, setStartupTemplate] = useStartupTemplate; @@ -142,11 +154,145 @@ const EnvironmentSettings = memo(() => { } }, [startupTemplate, lastDirectory, setStartupTemplate]); + const [memoryForUpscaling, setMemoryForUpscaling] = useMemoryForUpscaling; + const [isSystemMemory, setIsSystemMemory] = useIsSystemMemory; + const [memoryForUpscalingGPU, setMemoryForUpscalingGPU] = useMemoryForUpscalingGPU; + + const systemMemory = () => totalMemory / 1024 ** 3; + + const [isConfirmationOpen, setIsConfirmationOpen] = useState(false); + const cancelRef = useRef(null); + + const handleConfirmationConfirm = () => { + setIsConfirmationOpen(false); + setIsSystemMemory(true); + }; + + const handleConfirmationCancel = () => { + setIsConfirmationOpen(false); + }; + + const handleToggle = () => { + if (!isSystemMemory) { + setIsConfirmationOpen(true); + } else { + setIsSystemMemory(false); + } + }; + + const confirmationDialog = ( + + + + + Are you sure? + + + This action may result in heavy swapping and could potentially{' '} + render your system unusable when the limit is + set too high! Do you really want to continue? + + + + + + + + ); + return ( } w="full" > + {!isArmMac && ( + + )} + + + + + Up to{' '} + + {(systemMemory() * (memoryForUpscaling / 100)).toFixed(1)} GiB + {' '} + out of a total of{' '} + + {systemMemory()} GiB + {' '} + RAM will be used. Use with caution, may result in heavy swapping! + + ), + label: `Apply ${memoryForUpscaling.toFixed( + 1 + )}% memory limit to total system memory`, + disabled: false, + }} + value={isSystemMemory} + /> + {confirmationDialog} + + { placeholder="Select a file..." textOverflow="ellipsis" value={startupTemplate ? path.parse(startupTemplate).base : ''} + width={183} // eslint-disable-next-line @typescript-eslint/no-misused-promises onClick={onButtonClick} /> + + + } + size="xs" + onClick={() => setStartupTemplate('')} + /> + - } - size="xs" - onClick={() => setStartupTemplate('')} - /> Looking for the CPU and FP16 settings? They moved to the Python tab. @@ -270,7 +423,29 @@ const PythonSettings = memo(() => { /> {isSystemPython && ( + If wanted, use a specific python binary rather than the + default one invoked by{' '} + + python3 + {' '} + or{' '} + + python + + . This is useful if you have multiple python versions + installed and want to pick a specific one. + + } title="System Python location (optional)" > @@ -293,6 +468,7 @@ const PythonSettings = memo(() => { className="nodrag" cursor="pointer" draggable={false} + marginLeft="1.5" placeholder="Select a file..." textOverflow="ellipsis" value={ @@ -300,17 +476,24 @@ const PythonSettings = memo(() => { ? path.parse(systemPythonLocation).base : '' } + width={183} // eslint-disable-next-line @typescript-eslint/no-misused-promises onClick={onButtonClick} /> + + + } + size="xs" + onClick={() => setSystemPythonLocation(null)} + /> + - } - size="xs" - onClick={() => setSystemPythonLocation(null)} - /> )} diff --git a/src/renderer/components/settings/components.tsx b/src/renderer/components/settings/components.tsx index 44224c8922..9f2519b1b4 100644 --- a/src/renderer/components/settings/components.tsx +++ b/src/renderer/components/settings/components.tsx @@ -1,6 +1,8 @@ import { Button, HStack, + InputGroup, + InputRightElement, NumberDecrementStepper, NumberIncrementStepper, NumberInput, @@ -82,31 +84,38 @@ export const DropdownSetting = memo(({ setting, value, setValue }: SettingsProps }); export const NumberSetting = memo(({ setting, value, setValue }: SettingsProps<'number'>) => { + const isStepFloatingPoint = setting.step && Number(setting.step) % 1 !== 0; + const decimalPlaces = isStepFloatingPoint ? 1 : 0; + const formattedValue = isStepFloatingPoint ? value.toFixed(decimalPlaces) : Math.round(value); + return ( - { - const newValue = parseFloat(v); - if (!Number.isNaN(newValue)) { - setValue(newValue); - } - }} - > - - - - - - + + { + const newValue = parseFloat(v); + if (!Number.isNaN(newValue)) { + setValue(newValue); + } + }} + > + + {setting.unit} + + + + + + ); }); diff --git a/src/renderer/contexts/SettingsContext.tsx b/src/renderer/contexts/SettingsContext.tsx index 57cef769db..ee453994ac 100644 --- a/src/renderer/contexts/SettingsContext.tsx +++ b/src/renderer/contexts/SettingsContext.tsx @@ -19,6 +19,9 @@ interface Settings { setSnapToGridAmount: SetState ]; useStartupTemplate: GetSetState; + useMemoryForUpscaling: GetSetState; + useIsSystemMemory: GetSetState; + useMemoryForUpscalingGPU: GetSetState; useSelectTheme: GetSetState; useAnimateChain: GetSetState; useExperimentalFeatures: GetSetState; @@ -45,6 +48,9 @@ export const SettingsProvider = memo(({ children }: React.PropsWithChildren