Add nim nodes.

__init__.py

@@ -1,5 +1,6 @@
 import base64
 from io import BytesIO
+from nim import ModelType
 
 import numpy as np
 import requests
@@ -126,13 +127,51 @@ class FetchNGCApiKey:
         api_key =ngc.get_ngc_key()
         return (api_key,)
 
+class LoadNimNode:
+    def __init__(self):
+        pass
+    
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "model_type": ([e.value for e in ModelType], {
+                    "default": ModelType.SDXL.value,
+                    "tooltip": "The type of NIM model to load"
+                }),
+                "api_key": ("STRING", {
+                    "default": "",
+                    "tooltip": "NGC API Key for authentication"
+                }),
+                "port": ("INT", {
+                    "default": 8003,
+                    "min": 1,
+                    "max": 65535,
+                    "tooltip": "Port number where NIM is running"
+                })
+            }
+        }
+    
+    RETURN_TYPES = ("STRING",)  # Returns success/failure message
+    FUNCTION = "load_nim"
+    
+    def load_nim(self, api_key, port):
+        global invoke_url
+        invoke_url = f"http://localhost:{port}/v1/infer"
+        
+        # Here you could add logic to validate the API key and connection
+        # For now, we'll just return a success message
+        return (f"NIM configured on port {port}",)
+
 # Update the mappings
 NODE_CLASS_MAPPINGS = {
     "NIMSDXLNode": NIMSDXLNode,
-    "FetchNGCApiKey": FetchNGCApiKey
+    "FetchNGCApiKey": FetchNGCApiKey,
+    "LoadNimNode": LoadNimNode
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
     "NIMSDXLNode": "NIM SDXL",
-    "FetchNGCApiKey": "NGC API Key"
+    "FetchNGCApiKey": "NGC API Key",
+    "LoadNimNode": "Load NIM"
 }

nim.py

@@ -0,0 +1,159 @@
+import os
+import subprocess
+from enum import Enum
+from pathlib import Path
+
+
+class ModelType(Enum):
+    SDXL = "SDXL"
+    FLUX = "FLUX"
+    StudioVoice = "StudioVoice"
+
+class NIMManager:
+    '''
+    This class is responsible for managing the NIM containers.
+
+    Supports: 
+    - SDXL
+    - FLUX
+    - StudioVoice (https://build.nvidia.com/nvidia/studiovoice/docker)
+    '''
+    # Registry paths for different model types
+    MODEL_REGISTRY: dict[ModelType, str] = {
+        ModelType.SDXL: "nvcr.io/nvstaging/nim/sdxl:1.0.0",
+        ModelType.FLUX: "nvcr.io/nvstaging/nim/flux:1.0.0",
+        ModelType.StudioVoice: "nvcr.io/nim/nvidia/maxine-studio-voice:latest"
+    }
+
+    def __init__(self):
+        self._nim_server_proc_dict: dict[str, subprocess.Popen] = {}
+        
+    def setup_directories(self, model_name: str, version: str) -> bool:
+        """Create necessary directories for NIM cache"""
+        try:
+            home = os.path.expanduser("~")
+            cache_path = Path(f"{home}/nimcache/{model_name}/{version}/.cache")
+            cache_path.parent.mkdir(parents=True, exist_ok=True)
+            
+            # Set permissions
+            chmod_command = f"chmod 777 -R {cache_path}"
+            result = subprocess.run(chmod_command, shell=True, check=True)
+            
+            print(f"Directory setup completed for {model_name}")
+            return True
+            
+        except Exception as e:
+            print(f"Failed to setup directories: {str(e)}")
+            return False
+    
+    def pull_nim_image(self, nim_id: str, registry_path: str) -> bool:
+        """Pull NIM image from the registry"""
+        try:
+            command = f"wsl -d NVIDIA-Workbench -- podman pull {registry_path}"
+            result = subprocess.run(
+                command,
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.DEVNULL,
+                shell=True
+            )
+            
+            if result.returncode == 0:
+                print(f"{nim_id} image downloaded successfully")
+                return True
+            else:
+                print(f"Failed to download image {nim_id}")
+                return False
+                
+        except Exception as e:
+            print(f"Error pulling image: {str(e)}")
+            return False
+    
+    def start_nim_container(self, 
+                          nim_id: str,
+                          model_name: str,
+                          version: str,
+                          port: int,
+                          api_key: str,
+                          registry_path: str) -> bool:
+        """Start a NIM container with the specified configuration"""
+        try:
+            home = os.path.expanduser("~")
+            cache_path = f"{home}/nimonwsl2/{model_name}/{version}/.cache"
+            
+            command = (
+                f"wsl -d NVIDIA-Workbench -- podman run -it --rm "
+                f"--device=nvidia.com/gpu=all "
+                f"--name={model_name} "
+                f"--shm-size=8GB "
+                f"-e NGC_API_KEY={api_key} "
+                f"-v {cache_path}:/opt/nim/.cache "
+                f"-e NIM_RELAX_MEM_CONSTRAINTS=1 "
+                f"-p {port}:8000 "
+                f"{registry_path}"
+            )
+            
+            run_process = subprocess.Popen(
+                command,
+                stdin=subprocess.PIPE,
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.DEVNULL,
+                shell=True
+            )
+            
+            if run_process:
+                self._nim_server_proc_dict[nim_id] = run_process
+                print(f"Started NIM container for {nim_id} on port {port}")
+                return True
+            else:
+                print(f"Failed to start container for {nim_id}")
+                return False
+                
+        except Exception as e:
+            print(f"Error starting container: {str(e)}")
+            return False
+    
+    def deploy_nim(self, 
+                   model_name: ModelType,
+                   version: str,
+                   port: int,
+                   api_key: str) -> bool:
+        """Deploy a NIM model with all necessary setup"""
+        nim_id = f"{model_name.value}-{version}"
+        
+        # Get registry path from MODEL_REGISTRY
+        registry_path = self.MODEL_REGISTRY[model_name]
+        
+        # Setup directories
+        if not self.setup_directories(model_name.value, version):
+            return False
+            
+        # Pull image
+        if not self.pull_nim_image(nim_id, registry_path):
+            return False
+            
+        # Start container
+        return self.start_nim_container(
+            nim_id=nim_id,
+            model_name=model_name.value,
+            version=version,
+            port=port,
+            api_key=api_key,
+            registry_path=registry_path
+        )
+
+# Example usage
+if __name__ == "__main__":
+    manager = NIMManager()
+    
+    # Example deployment
+    success = manager.deploy_nim(
+        model_name=ModelType.SDXL,
+        version="1.5.1-rtx.rc3",
+        port=18001,
+        api_key="your-api-key-here"
+    )
+    
+    if success:
+        print("NIM deployment successful")
+    else:
+        print("NIM deployment failed")