From 3169e6a6cbf3a03b37e6cbc861007c6cef747778 Mon Sep 17 00:00:00 2001 From: unknown Date: Mon, 2 Jun 2025 12:29:21 -0700 Subject: [PATCH] add FLUX_SCHNELL, add Default offloading mode --- __init__.py | 2 +- nim.py | 9 ++++++--- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/__init__.py b/__init__.py index 5e9420a..1263024 100644 --- a/__init__.py +++ b/__init__.py @@ -216,7 +216,7 @@ class InstallNIMNode: "NIM Installer", 1 | 48) if res == 1: import webbrowser - webbrowser.open("https://assets.ngc.nvidia.com/products/api-catalog/rtx/NIM_Prerequisites_Installer_03052025.zip") + webbrowser.open("https://assets.ngc.nvidia.com/products/api-catalog/rtx/NIMSetup.exe") raise Exception("Please install NVIDIA NIM first and try again.") diff --git a/nim.py b/nim.py index 1227d59..f15d7cf 100644 --- a/nim.py +++ b/nim.py @@ -20,12 +20,14 @@ class ModelType(Enum): FLUX_DEV = "FLUX_DEV" FLUX_CANNY = "FLUX_CANNY" FLUX_DEPTH = "FLUX_DEPTH" + FLUX_SCHNELL = "FLUX_SCHNELL" class OffloadingPolicy(Enum): NONE = "None" SYS = "System RAM" DISK = "Disk" + DEFAULT = "Default" class NIMManager: ''' @@ -38,9 +40,10 @@ class NIMManager: # Registry paths for different model types MODEL_REGISTRY: dict[ModelType, str] = { - ModelType.FLUX_DEV: "nvcr.io/nim/black-forest-labs/flux.1-dev:1.0.0", - ModelType.FLUX_CANNY: "nvcr.io/nim/black-forest-labs/flux.1-dev:1.0.0", - ModelType.FLUX_DEPTH: "nvcr.io/nim/black-forest-labs/flux.1-dev:1.0.0", + ModelType.FLUX_DEV: "nvcr.io/nim/black-forest-labs/flux.1-dev:1.1.0", + ModelType.FLUX_CANNY: "nvcr.io/nim/black-forest-labs/flux.1-dev:1.1.0", + ModelType.FLUX_DEPTH: "nvcr.io/nim/black-forest-labs/flux.1-dev:1.1.0", + ModelType.FLUX_SCHNELL: "nvcr.io/nim/black-forest-labs/flux.1-schnell:1.0.0", } PORT = 5000