Backend-defined Dependencies

backend/src/api.py

@@ -2,6 +2,8 @@
 
 import importlib
 import os
+import platform
+import sys
 from dataclasses import dataclass, field
 from typing import Callable, Dict, Iterable, List, Tuple, TypedDict, TypeVar, Union
 
@@ -19,6 +21,10 @@
     typeValidateSchema,
 )
 
+KB = 1024**1
+MB = 1024**2
+GB = 1024**3
+
 
 def _process_inputs(base_inputs: Iterable[Union[BaseInput, NestedGroup]]):
     inputs: List[BaseInput] = []
@@ -174,11 +180,34 @@ def toDict(self):
         }
 
 
+@dataclass
+class Dependency:
+    display_name: str
+    package_name: str
+    version: str
+    size_estimate: int | float
+    auto_update: bool = False
+    extra_index_url: str | None = None
+
+    import_name: str | None = None
+
+    def toDict(self):
+        return {
+            "displayName": self.display_name,
+            "packageName": self.package_name,
+            "version": self.version,
+            "sizeEstimate": int(self.size_estimate),
+            "autoUpdate": self.auto_update,
+            "findLink": self.extra_index_url,
+        }
+
+
 @dataclass
 class Package:
     where: str
     name: str
-    dependencies: List[str] = field(default_factory=list)
+    description: str
+    dependencies: List[Dependency] = field(default_factory=list)
     categories: List[Category] = field(default_factory=list)
 
     def add_category(
@@ -200,6 +229,12 @@ def add_category(
         self.categories.append(result)
         return result
 
+    def add_dependency(
+        self,
+        dependency: Dependency,
+    ):
+        self.dependencies.append(dependency)
+
 
 def _iter_py_files(directory: str):
     for root, _, files in os.walk(directory):
@@ -271,5 +306,13 @@ def _refresh_nodes(self):
 registry = PackageRegistry()
 
 
-def add_package(where: str, name: str, dependencies: List[str]) -> Package:
-    return registry.add(Package(where, name, dependencies))
+def add_package(
+    where: str, name: str, description: str, dependencies: List[Dependency]
+) -> Package:
+    return registry.add(Package(where, name, description, dependencies))
+
+
+is_mac = sys.platform == "darwin"
+is_arm_mac = is_mac and platform.machine() == "arm64"
+is_windows = sys.platform == "win32"
+is_linux = sys.platform == "linux"

backend/src/gpu.py

@@ -0,0 +1,59 @@
+import pynvml as nv
+from sanic.log import logger
+
+nvidia_is_available = False
+
+try:
+    nv.nvmlInit()
+    nvidia_is_available = True
+    nv.nvmlShutdown()
+except nv.NVMLError as e:
+    logger.info("No Nvidia GPU found, or invalid driver installed.")
+except Exception as e:
+    logger.info(f"Unknown error occurred when trying to initialize Nvidia GPU: {e}")
+
+
+class NvidiaHelper:
+    def __init__(self):
+        self.nvidia_is_available = nvidia_is_available
+        if not nvidia_is_available:
+            raise RuntimeError("Nvidia GPU not found, or invalid driver installed.")
+
+        nv.nvmlInit()
+
+        self.__num_gpus = nv.nvmlDeviceGetCount()
+
+        self.__gpus = []
+        for i in range(self.__num_gpus):
+            handle = nv.nvmlDeviceGetHandleByIndex(i)
+            self.__gpus.append(
+                {
+                    "name": nv.nvmlDeviceGetName(handle),
+                    "uuid": nv.nvmlDeviceGetUUID(handle),
+                    "index": i,
+                    "handle": handle,
+                }
+            )
+
+    def __del__(self):
+        if nvidia_is_available:
+            nv.nvmlShutdown()
+
+    def list_gpus(self):
+        if not nvidia_is_available:
+            return None
+        return self.__gpus
+
+    def get_current_vram_usage(self, gpu_index=0):
+        if not nvidia_is_available:
+            return None
+
+        info = nv.nvmlDeviceGetMemoryInfo(self.__gpus[gpu_index]["handle"])
+
+        return info.total, info.used, info.free
+
+
+__all__ = [
+    "nvidia_is_available",
+    "NvidiaHelper",
+]

backend/src/installed_deps.py

@@ -0,0 +1,29 @@
+import subprocess
+import sys
+
+python_path = sys.executable
+
+# pylint: disable=global-at-module-level
+global installed_packages
+installed_packages = {}
+
+
+def install_dependency(package_name, version):
+    if package_name not in installed_packages:
+        subprocess.check_call(
+            [
+                python_path,
+                "-m",
+                "pip",
+                "install",
+                "--upgrade",
+                f"{package_name}=={version}",
+            ]
+        )
+        installed_packages[package_name] = version
+
+
+def set_installed_packages(packages):
+    # pylint: disable=global-statement
+    global installed_packages
+    installed_packages = packages

backend/src/packages/chaiNNer_external/__init__.py

@@ -2,7 +2,12 @@
 
 from api import add_package
 
-package = add_package(__file__, name="chaiNNer_external", dependencies=[])
+package = add_package(
+    __file__,
+    name="External",
+    description="Interact with an external Stable Diffusion API",
+    dependencies=[],
+)
 
 external_stable_diffusion_category = package.add_category(
     name="Stable Diffusion (External)",

backend/src/packages/chaiNNer_ncnn/__init__.py

@@ -1,8 +1,22 @@
 from sanic.log import logger
 
-from api import add_package
+from api import MB, Dependency, add_package, is_mac
 
-package = add_package(__file__, name="chaiNNer_ncnn", dependencies=[])
+package = add_package(
+    __file__,
+    name="NCNN",
+    description="NCNN uses .bin/.param models to upscale images. NCNN uses Vulkan for GPU acceleration, meaning it supports any modern GPU. Models can be converted from PyTorch to NCNN.",
+    dependencies=[
+        Dependency(
+            "NCNN",
+            "ncnn-vulkan",
+            "2022.9.12",
+            7 * MB if is_mac else 4 * MB,
+            auto_update=True,
+            import_name="ncnn_vulkan",
+        ),
+    ],
+)
 
 ncnn_category = package.add_category(
     name="NCNN",

backend/src/packages/chaiNNer_onnx/__init__.py

@@ -1,8 +1,85 @@
 from sanic.log import logger
 
-from api import add_package
+from api import KB, MB, Dependency, add_package, is_arm_mac
+from gpu import nvidia_is_available
 
-package = add_package(__file__, name="chaiNNer_onnx", dependencies=[])
+
+def get_onnx_runtime():
+    if is_arm_mac:
+        return Dependency(
+            display_name="ONNX Runtime (Silicon)",
+            package_name="onnxruntime-silicon",
+            version="1.13.1",
+            size_estimate=6 * MB,
+            import_name="onnxruntime",
+        )
+    elif nvidia_is_available:
+        return Dependency(
+            display_name="ONNX Runtime (GPU)",
+            package_name="onnxruntime-gpu",
+            version="1.13.1",
+            size_estimate=110 * MB,
+            import_name="onnxruntime",
+        )
+    else:
+        return Dependency(
+            display_name="ONNX Runtime",
+            package_name="onnxruntime",
+            version="1.13.1",
+            size_estimate=5 * MB,
+        )
+
+
+dependencies = [
+    Dependency(
+        display_name="ONNX",
+        package_name="onnx",
+        version="1.13.0",
+        size_estimate=12 * MB,
+    ),
+]
+
+if not is_arm_mac:
+    dependencies.append(
+        Dependency(
+            display_name="ONNX Optimizer",
+            package_name="onnxoptimizer",
+            version="0.3.6",
+            size_estimate=300 * KB,
+        )
+    )
+
+dependencies.extend(
+    [
+        get_onnx_runtime(),
+        Dependency(
+            display_name="Protobuf",
+            package_name="protobuf",
+            version="3.20.2",
+            size_estimate=500 * KB,
+        ),
+        Dependency(
+            display_name="SciPy",
+            package_name="scipy",
+            version="1.9.3",
+            size_estimate=42 * MB,
+        ),
+        Dependency(
+            display_name="Numba",
+            package_name="numba",
+            version="0.56.3",
+            size_estimate=2.5 * MB,
+        ),
+    ]
+)
+
+
+package = add_package(
+    __file__,
+    name="ONNX",
+    description="ONNX uses .onnx models to upscale images. It also helps to convert between PyTorch and NCNN. It is fastest when CUDA is supported. If TensorRT is installed on the system, it can also be configured to use that.",
+    dependencies=dependencies,
+)
 
 onnx_category = package.add_category(
     name="ONNX",

backend/src/packages/chaiNNer_pytorch/__init__.py

@@ -1,8 +1,83 @@
+import sys
+
 from sanic.log import logger
 
-from api import add_package
+from api import GB, KB, MB, Dependency, add_package
+from gpu import nvidia_is_available
+
+python_version = sys.version_info
 
-package = add_package(__file__, name="chaiNNer_pytorch", dependencies=[])
+dependencies = []
+if python_version.minor < 10:
+    dependencies.extend(
+        [
+            Dependency(
+                display_name="PyTorch",
+                package_name="torch",
+                version="1.10.2+cu113" if nvidia_is_available else "1.10.2",
+                size_estimate=2 * GB if nvidia_is_available else 140 * MB,
+                extra_index_url="https://download.pytorch.org/whl/cu113"
+                if nvidia_is_available
+                else None,
+            ),
+            Dependency(
+                display_name="TorchVision",
+                package_name="torchvision",
+                version="0.11.3+cu113" if nvidia_is_available else "0.11.3",
+                size_estimate=2 * MB if nvidia_is_available else 800 * KB,
+                extra_index_url="https://download.pytorch.org/whl/cu113"
+                if nvidia_is_available
+                else None,
+            ),
+        ]
+    )
+elif python_version.minor >= 10:
+    dependencies.extend(
+        [
+            Dependency(
+                display_name="PyTorch",
+                package_name="torch",
+                version="1.12.1+cu116" if nvidia_is_available else "1.12.1",
+                size_estimate=2 * GB if nvidia_is_available else 140 * MB,
+                extra_index_url="https://download.pytorch.org/whl/cu116"
+                if nvidia_is_available
+                else None,
+            ),
+            Dependency(
+                display_name="TorchVision",
+                package_name="torchvision",
+                version="0.13.1+cu116" if nvidia_is_available else "0.13.1",
+                size_estimate=2 * MB if nvidia_is_available else 800 * KB,
+                extra_index_url="https://download.pytorch.org/whl/cu116"
+                if nvidia_is_available
+                else None,
+            ),
+        ]
+    )
+
+dependencies.extend(
+    [
+        Dependency(
+            display_name="FaceXLib",
+            package_name="facexlib",
+            version="0.2.5",
+            size_estimate=1.1 * MB,
+        ),
+        Dependency(
+            display_name="Einops",
+            package_name="einops",
+            version="0.5.0",
+            size_estimate=36.5 * KB,
+        ),
+    ]
+)
+
+package = add_package(
+    __file__,
+    name="PyTorch",
+    description="PyTorch uses .pth models to upscale images, and is fastest when CUDA is supported (Nvidia GPU). If CUDA is unsupported, it will install with CPU support (which is very slow).",
+    dependencies=dependencies,
+)
 
 pytorch_category = package.add_category(
     name="PyTorch",