add FLUX_SCHNELL, add Default offloading mode

This commit is contained in:
unknown
2025-06-02 12:29:21 -07:00
parent 8345d194a7
commit 3169e6a6cb
2 changed files with 7 additions and 4 deletions
+1 -1
View File
@@ -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.")
+6 -3
View File
@@ -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