refactor execution settings

backend/src/nodes/impl/pytorch/utils.py

@@ -6,18 +6,17 @@
 import torch
 from torch import Tensor
 
-from ...utils.exec_options import PackageExecutionOptions
 from ..image_utils import as_3d
 from ..onnx.np_tensor_utils import MAX_VALUES_BY_DTYPE, np_denorm
 
 
-def get_pytorch_device(pytorch_options: PackageExecutionOptions) -> torch.device:
+def get_pytorch_device(use_cpu: bool = False, gpu_index: int = 0) -> torch.device:
     # CPU override
-    if pytorch_options.get("cpu_mode", False):
+    if use_cpu:
         device = "cpu"
     # Check for Nvidia CUDA
     elif torch.cuda.is_available() and torch.cuda.device_count() > 0:
-        device = "cuda:" + pytorch_options.get("gpu", "0")
+        device = f"cuda:{gpu_index}"
     # Check for Apple MPS
     elif hasattr(torch, "backends") and hasattr(torch.backends, "mps") and torch.backends.mps.is_built() and torch.backends.mps.is_available():  # type: ignore -- older pytorch versions dont support this technically
         device = "mps"

backend/src/packages/chaiNNer_pytorch/__init__.py

@@ -1,5 +1,5 @@
 import sys
-from typing import TypedDict
+from typing import TypedDict, cast
 
 from sanic.log import logger
 
@@ -169,10 +169,15 @@ def get_pytorch():
     {
         "cpu_mode": bool,
         "fp16_mode": bool,
-        "gpu": str,
+        "gpu": int,
     },
 )
 
+
+def get_pytorch_settings() -> PyTorchSettings:
+    return cast(PyTorchSettings, package.get_execution_settings())
+
+
 pytorch_category = package.add_category(
     name="PyTorch",
     description="Nodes for using the PyTorch Neural Network Framework with images.",

backend/src/packages/chaiNNer_pytorch/pytorch/io/load_model.py

@@ -11,11 +11,10 @@
 from nodes.impl.pytorch.utils import get_pytorch_device
 from nodes.properties.inputs import PthFileInput
 from nodes.properties.outputs import DirectoryOutput, FileNameOutput, ModelOutput
-from nodes.utils.exec_options import get_execution_options
 from nodes.utils.unpickler import RestrictedUnpickle
 from nodes.utils.utils import split_file_path
 
-from ... import package as pytorch_package
+from ... import get_pytorch_settings
 from .. import io_group
 
 
@@ -48,8 +47,10 @@ def load_model_node(path: str) -> Tuple[PyTorchModel, str, str]:
 
     assert os.path.isfile(path), f"Path {path} is not a file"
 
-    exec_options = pytorch_package.get_execution_settings()
-    pytorch_device = get_pytorch_device(exec_options)
+    exec_options = get_pytorch_settings()
+    pytorch_device = get_pytorch_device(
+        exec_options.get("cpu_mode", False), exec_options.get("gpu", 0)
+    )
     logger.info(f"Execution options: {exec_options}")
 
     try:

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

@@ -1,17 +1,15 @@
 from __future__ import annotations
 
 import numpy as np
-import torch
 
 from nodes.impl.pytorch.pix_transform.auto_split import pix_transform_auto_split
 from nodes.impl.pytorch.pix_transform.pix_transform import Params
 from nodes.impl.pytorch.utils import get_pytorch_device
 from nodes.impl.upscale.grayscale import SplitMode
 from nodes.properties.inputs import EnumInput, ImageInput, SliderInput
 from nodes.properties.outputs import ImageOutput
-from nodes.utils.exec_options import get_execution_options
 
-from ... import package as pytorch_package
+from ... import get_pytorch_settings
 from .. import processing_group
 
 
@@ -76,12 +74,15 @@ def guided_upscale_node(
     iterations: float,
     split_mode: SplitMode,
 ) -> np.ndarray:
-    exec_options = pytorch_package.get_execution_settings()
+    exec_options = get_pytorch_settings()
+    pytorch_device = get_pytorch_device(
+        exec_options.get("cpu_mode", False), exec_options.get("gpu", 0)
+    )
 
     return pix_transform_auto_split(
         source=source,
         guide=guide,
-        device=get_pytorch_device(exec_options),
+        device=pytorch_device,
         params=Params(iteration=int(iterations * 1000)),
         split_mode=split_mode,
     )

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

@@ -18,10 +18,9 @@
 from nodes.properties.inputs import ImageInput
 from nodes.properties.inputs.pytorch_inputs import InpaintModelInput
 from nodes.properties.outputs import ImageOutput
-from nodes.utils.exec_options import PackageExecutionOptions
 from nodes.utils.utils import get_h_w_c
 
-from ... import package as pytorch_package
+from ... import PyTorchSettings, get_pytorch_settings
 from .. import processing_group
 
 
@@ -59,12 +58,15 @@ def inpaint(
     img: np.ndarray,
     mask: np.ndarray,
     model: PyTorchInpaintModel,
-    options: PackageExecutionOptions,
+    options: PyTorchSettings,
 ):
     with torch.no_grad():
         # TODO: use bfloat16 if RTX
         use_fp16 = options.get("fp16_mode", False) and model.supports_fp16
-        device = get_pytorch_device(options)
+        device = get_pytorch_device(
+            options.get("cpu_mode", False), options.get("gpu", 0)
+        )
+
         model = model.to(device)
         model = model.half() if use_fp16 else model.float()
 
@@ -158,6 +160,6 @@ def inpaint_node(
         img.shape[:2] == mask.shape[:2]
     ), "Input image and mask must have the same resolution"
 
-    exec_options = pytorch_package.get_execution_settings()
+    exec_options = get_pytorch_settings()
 
     return inpaint(img, mask, model, exec_options)

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

@@ -8,7 +8,7 @@
 
 from nodes.impl.pytorch.auto_split import pytorch_auto_split
 from nodes.impl.pytorch.types import PyTorchSRModel
-from nodes.impl.pytorch.utils import to_pytorch_execution_options
+from nodes.impl.pytorch.utils import get_pytorch_device
 from nodes.impl.upscale.auto_split_tiles import (
     TileSize,
     estimate_tile_size,
@@ -18,28 +18,30 @@
 from nodes.impl.upscale.tiler import MaxTileSize
 from nodes.properties.inputs import ImageInput, SrModelInput, TileSizeDropdown
 from nodes.properties.outputs import ImageOutput
-from nodes.utils.exec_options import ExecutionOptions, get_execution_options
 from nodes.utils.utils import get_h_w_c
 
+from ... import PyTorchSettings, get_pytorch_settings
 from .. import processing_group
 
 
 def upscale(
     img: np.ndarray,
     model: PyTorchSRModel,
     tile_size: TileSize,
-    options: ExecutionOptions,
+    options: PyTorchSettings,
 ):
     with torch.no_grad():
         # Borrowed from iNNfer
         logger.debug("Upscaling image")
 
         # TODO: use bfloat16 if RTX
-        use_fp16 = options.fp16 and model.supports_fp16
-        device = torch.device(options.full_device)
+        use_fp16 = options.get("fp16_mode", False) and model.supports_fp16
+        device = get_pytorch_device(
+            options.get("cpu_mode", False), options.get("gpu", 0)
+        )
 
         def estimate():
-            if "cuda" in options.full_device:
+            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
@@ -101,7 +103,7 @@ def upscale_image_node(
 ) -> np.ndarray:
     """Upscales an image with a pretrained model"""
 
-    exec_options = to_pytorch_execution_options(get_execution_options())
+    exec_options = get_pytorch_settings()
 
     logger.debug(f"Upscaling image...")
 

backend/src/packages/chaiNNer_pytorch/pytorch/restoration/upscale_face.py

@@ -15,16 +15,17 @@
 from nodes.impl.image_utils import to_uint8
 from nodes.impl.pytorch.types import PyTorchFaceModel
 from nodes.impl.pytorch.utils import (
+    get_pytorch_device,
     np2tensor,
     safe_cuda_cache_empty,
     tensor2np,
-    to_pytorch_execution_options,
 )
 from nodes.properties.inputs import FaceModelInput, ImageInput, NumberInput, SliderInput
 from nodes.properties.outputs import ImageOutput
 from nodes.utils.exec_options import get_execution_options
 from nodes.utils.utils import get_h_w_c
 
+from ... import PyTorchSettings, get_pytorch_settings
 from .. import restoration_group
 
 
@@ -45,10 +46,9 @@ def upscale(
     face_helper: FaceRestoreHelper,
     face_model: PyTorchFaceModel,
     weight: float,
+    exec_options: PyTorchSettings,
+    device: torch.device,
 ):
-    exec_options = to_pytorch_execution_options(get_execution_options())
-
-    device = torch.device(exec_options.full_device)
     face_helper.clean_all()
 
     face_helper.read_image(img)
@@ -59,7 +59,7 @@ def upscale(
     # align and warp each face
     face_helper.align_warp_face()
 
-    should_use_fp16 = exec_options.fp16 and face_model.supports_fp16
+    should_use_fp16 = exec_options.get("fp16_mode", False) and face_model.supports_fp16
     if should_use_fp16:
         face_model = face_model.half()
     else:
@@ -161,8 +161,10 @@ def upscale_face_node(
     try:
         img = denormalize(img)
 
-        exec_options = to_pytorch_execution_options(get_execution_options())
-        device = torch.device(exec_options.full_device)
+        exec_options = get_pytorch_settings()
+        device = get_pytorch_device(
+            exec_options.get("cpu_mode", False), exec_options.get("gpu", 0)
+        )
 
         with torch.no_grad():
             appdata_path = user_data_dir(roaming=True)
@@ -187,6 +189,8 @@ def upscale_face_node(
                 face_helper,
                 face_model,
                 weight,
+                exec_options,
+                device,
             )
 
             return result

backend/src/packages/chaiNNer_pytorch/pytorch/utility/convert_to_onnx.py

@@ -7,11 +7,11 @@
 
 from nodes.impl.onnx.model import OnnxGeneric
 from nodes.impl.pytorch.types import PyTorchSRModel
-from nodes.impl.pytorch.utils import to_pytorch_execution_options
+from nodes.impl.pytorch.utils import get_pytorch_device
 from nodes.properties.inputs import OnnxFpDropdown, SrModelInput
 from nodes.properties.outputs import OnnxModelOutput, TextOutput
-from nodes.utils.exec_options import get_execution_options
 
+from ... import get_pytorch_settings
 from .. import utility_group
 
 
@@ -37,23 +37,28 @@ def convert_to_onnx_node(
     model: PyTorchSRModel, is_fp16: int
 ) -> Tuple[OnnxGeneric, str]:
     fp16 = bool(is_fp16)
-    exec_options = to_pytorch_execution_options(get_execution_options())
+    exec_options = get_pytorch_settings()
+    device = get_pytorch_device(
+        exec_options.get("cpu_mode", False), exec_options.get("gpu", 0)
+    )
     if fp16:
-        assert (
-            exec_options.fp16
+        assert exec_options.get(
+            "fp16_mode", False
         ), "PyTorch fp16 mode must be supported and turned on in settings to convert model as fp16."
 
     model = model.eval()
-    model = model.to(torch.device(exec_options.full_device))
+    model = model.to(device)
     # https://github.com/onnx/onnx/issues/654
     dynamic_axes = {
         "input": {0: "batch_size", 2: "height", 3: "width"},
         "output": {0: "batch_size", 2: "height", 3: "width"},
     }
     dummy_input = torch.rand(1, model.in_nc, 64, 64)
-    dummy_input = dummy_input.to(torch.device(exec_options.full_device))
+    dummy_input = dummy_input.to(device)
 
-    should_use_fp16 = exec_options.fp16 and model.supports_fp16 and fp16
+    should_use_fp16 = (
+        exec_options.get("fp16_mode", False) and model.supports_fp16 and fp16
+    )
     if should_use_fp16:
         model = model.half()
         dummy_input = dummy_input.half()

backend/src/server.py

@@ -237,7 +237,6 @@ async def run_individual(request: Request):
         logger.debug(full_data)
         exec_opts = parse_execution_options(full_data["options"])
         set_execution_options(exec_opts)
-        logger.debug(f"Using device: {exec_opts.full_device}")
         # Create node based on given category/name information
         node_instance = api.registry.get_node(full_data["schemaId"])