From ae0ab02effcdb5faf8a4293d81d48d7ac7a1523f Mon Sep 17 00:00:00 2001
From: Sean Kilbride <110414370+SKilbride@users.noreply.github.com>
Date: Tue, 1 Apr 2025 14:33:05 -0700
Subject: [PATCH] Update nim.py

---
 nim.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/nim.py b/nim.py
index b24bc69..e0fc064 100644
--- a/nim.py
+++ b/nim.py
@@ -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)
\ No newline at end of file
+    manager.stop_nim(model_name)