Update nim.py

This commit is contained in:
Sean Kilbride
2025-04-01 14:33:05 -07:00
committed by GitHub
parent 45286bb4c4
commit ae0ab02eff
+4 -4
View File
@@ -38,9 +38,9 @@ class NIMManager:
# Registry paths for different model types
MODEL_REGISTRY: dict[ModelType, str] = {
ModelType.FLUX_DEV: "nvcr.io/nvstaging/nim/flux.1-dev:1.0.0-rc.5.25953947",
ModelType.FLUX_CANNY: "nvcr.io/nvstaging/nim/flux.1-dev:1.0.0-rc.5.25953947",
ModelType.FLUX_DEPTH: "nvcr.io/nvstaging/nim/flux.1-dev:1.0.0-rc.5.25953947",
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",
}
PORT = 5000
@@ -361,4 +361,4 @@ if __name__ == "__main__":
manager = NIMManager()
manager.deploy_nim(model_name)
manager.stop_nim(model_name)
manager.stop_nim(model_name)