diff --git a/__init__.py b/__init__.py index f33d1ab..912f06b 100644 --- a/__init__.py +++ b/__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" } diff --git a/nim.py b/nim.py new file mode 100644 index 0000000..792f4d4 --- /dev/null +++ b/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")