diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..bee8a64 --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +__pycache__ diff --git a/README.md b/README.md index e69de29..6b01b4b 100644 --- a/README.md +++ b/README.md @@ -0,0 +1,48 @@ +# NVIDIA NIM SDXL + +### What is Stable Diffusion XL NIM? + +The Stable Diffusion XL NIM is a container that allows you to run Stability.AI’s Stable Diffusion XL model - one of the most popular visual generative AI models in the world - in the most optimal manner. + +## Getting Started with Stable Diffusion XL NIM with ComfyUI + +Before installing, ensure your system meets the following requirements: +Operating System: Windows 11 (22H1 or later) +GPU: AD100 or above +GPU Driver: Version 565.xx or later +Virtualization Settings: Enabled in SBIOS - [Instructions to enable virtualization if it is not enabled](https://support.microsoft.com/en-gb/windows/enable-virtualization-on-windows-c5578302-6e43-4b4b-a449-8ced115f58e1) + +1. Download the installer from {todo: URL} +2. Run the installer +3. Open the Terminal app +![Search for Terminal in your Start Menu](assets/terminal-startmenu.png) +4. Open the NVIDIA-Workbench profile by clicking on the arrow along the top +![Load the NVIDIA-Workbench profile in Terminal](assets/terminal-workbench.png) +5. Setup NIM directories by running the following commands +``` +export LOCAL_NIM_CACHE=~/.cache/nim +mkdir -p "$LOCAL_NIM_CACHE" +chmod -R a+w "$LOCAL_NIM_CACHE" +``` + +### Install node in ComfyUI +The recommended way to install these nodes is to use the [ComfyUI Manager](https://github.com/ltdrdata/ComfyUI-Manager) to easily install them to your ComfyUI instance. +You can also manually install them by git cloning the repo to your ComfyUI/custom_nodes folder + +## Usage +You will need to start the container each time you want to use the node. +To do this, navigate to the NVIDIA-Workbench profile in Terminal (see getting started section) and run the following command: +``` +podman run -it --rm \ + --device=nvidia.com/gpu=all \ + --shm-size=16GB \ + -e NGC_API_KEY=$NGC_API_KEY \ + -v "$LOCAL_NIM_CACHE:/opt/nim/.cache" \ + -u $(id -u) \ + -p 8000:8000 \ + {todo: URL} + ``` + The first time you run this it will download and configure the container, this may take a while. + +After the container is running you can then simply add the NIM SDXL node to generate images, download the [example workflow here](example_workflows/workflow.json) or you can drag the image below into ComfyUI to view the embedded workflow: +![Example workflow](assets/workflow.png) diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..14c345a --- /dev/null +++ b/__init__.py @@ -0,0 +1,117 @@ +import base64 +from PIL import Image +from io import BytesIO +import requests +import numpy as np +import torch + +invoke_url = "http://localhost:8003/v1/infer" + +class NIMSDXLNode: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "width": (["768", "832", "896", "960", "1024", "1088", "1152", "1216", "1280", "1344"], { + "default": "1024", + "tooltip": "Width of the image to generate, in pixels." + }), + "height": (["768", "832", "896", "960", "1024", "1088", "1152", "1216", "1280", "1344"], { + "default": "1024", + "tooltip": "Height of the image to generate, in pixels." + }), + "positive": ("STRING", { + "multiline": True, + "default": "beautiful scenery nature glass bottle landscape, purple galaxy bottle", + "tooltip": "The attributes you want to include in the image." + }), + "negative": ("STRING", { + "multiline": True, + "default": "text, watermark", + "tooltip": "The attributes you want to exclude from the image." + }), + "cfg_scale": ("FLOAT", { + "default": 5.0, + "min": 1.0, + "max": 20.0, + "step": 0.5, + "display": "slider", + "tooltip": "How strictly the diffusion process adheres to the prompt text (higher values keep your image closer to your prompt)" + }), + "sampler": (["DDIM", "K_EULER_ANCESTRAL", "K_DPM_2_ANCESTRAL"], { + "default": "K_DPM_2_ANCESTRAL", + "tooltip": "The sampler to use for generation. Varying diffusion samplers will vary outputs significantly." + }), + "seed": ("INT", { + "default": 0, + "min": 0, + "max": 4294967295, + "display": "number", + "tooltip": "The seed which governs generation. Use 0 for a random seed" + }), + "steps": ("INT", { + "default": 25, + "min": 1, + "max": 100, + "step": 1, + "display": "slider", + "tooltip": "Number of diffusion steps to run" + }), + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "generate" + CATEGORY = "image generation" + + def generate(self, width, height, positive, negative, cfg_scale, sampler, seed, steps): + payload = { + "width": width, + "height": height, + "text_prompts": [ + { + "text": positive, + "weight": 1 + }, + { + "text": negative, + "weight": -1 + } + ], + "cfg_scale": cfg_scale, + "sampler": sampler, + "seed": seed, + "steps": steps + } + + try: + response = requests.post(invoke_url, json=payload) + except requests.exceptions.ConnectionError: + raise ConnectionError("Unable to connect to NIM API. Please ensure that you have started the container via podman.") + + data = response.json() + response.raise_for_status() + img_base64 = data["artifacts"][0]["base64"] + img_bytes = base64.b64decode(img_base64) + + print("Result: " + data["artifacts"][0]["finishReason"]) + + image = Image.open(BytesIO(img_bytes)) + image = image.convert("RGB") + image = np.array(image).astype(np.float32) / 255.0 + image = torch.from_numpy(image)[None,] + + return (image,) + + +# Update the mappings +NODE_CLASS_MAPPINGS = { + "NIMSDXLNode": NIMSDXLNode +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "NIMSDXLNode": "NIM SDXL" +} diff --git a/assets/terminal-startmenu.png b/assets/terminal-startmenu.png new file mode 100644 index 0000000..cf0f415 Binary files /dev/null and b/assets/terminal-startmenu.png differ diff --git a/assets/terminal-workbench.png b/assets/terminal-workbench.png new file mode 100644 index 0000000..eb0b8f8 Binary files /dev/null and b/assets/terminal-workbench.png differ diff --git a/assets/workflow.png b/assets/workflow.png new file mode 100644 index 0000000..5baa214 Binary files /dev/null and b/assets/workflow.png differ diff --git a/example_workflows/sdxl_nim_workflow.json b/example_workflows/sdxl_nim_workflow.json new file mode 100644 index 0000000..f743885 --- /dev/null +++ b/example_workflows/sdxl_nim_workflow.json @@ -0,0 +1,95 @@ +{ + "last_node_id": 8, + "last_link_id": 3, + "nodes": [ + { + "id": 6, + "type": "NIMSDXLNode", + "pos": [ + 2175.919921875, + 684.1629638671875 + ], + "size": [ + 400, + 294 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 3 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "NIMSDXLNode" + }, + "widgets_values": [ + "1024", + "1024", + "beautiful scenery nature glass bottle landscape, purple galaxy bottle", + "text, watermark", + 5, + "K_DPM_2_ANCESTRAL", + 2124418934, + "randomize", + 25 + ] + }, + { + "id": 8, + "type": "SaveImage", + "pos": [ + 2596.94873046875, + 683.8137817382812 + ], + "size": [ + 315, + 270 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 3 + } + ], + "outputs": [], + "properties": {}, + "widgets_values": [ + "ComfyUI" + ] + } + ], + "links": [ + [ + 3, + 6, + 0, + 8, + 0, + "IMAGE" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 1, + "offset": [ + -804.919921875, + -169.81378173828125 + ] + } + }, + "version": 0.4 +} \ No newline at end of file diff --git a/node.py b/node.py deleted file mode 100644 index 66d8c0d..0000000 --- a/node.py +++ /dev/null @@ -1,123 +0,0 @@ -import base64 - -import requests - -invoke_url = "http://localhost:8003/v1/infer" - - - -# payload = { -# "text_prompts": [ -# { -# "text": "realistic futuristic city-downtown with short buildings, sunset", -# "weight": 1 -# }, -# { -# "text": "" , -# "weight": -1 -# } -# ], -# "cfg_scale": 5, -# "sampler": "K_DPM_2_ANCESTRAL", -# "seed": 0, -# "steps": 25 -# } - -# # Size is fixed for now -# payload_two = { "cfg_scale": 5, "clip_guidance_preset": "NONE", -# "disable_safety_checker": 'false', # Remain false for now. NIM_ALLOW_UNCHECKED_GENERATION=true -# "height": 1024, -# "sampler": "K_DPM_2_ANCESTRAL", -# "samples": 1, "seed": 0, "steps": 25, -# "style_preset": "none", -# "text_prompts": [ { -# "text": "A photo of a Shiba Inu dog with a backpack riding a bike", "weight": 1 -# } ], -# "use_refiner": 'false', -# "width": 1024 -# } - -# response = requests.post(invoke_url, json=payload) - -# response.raise_for_status() - -# data = response.json() - -# img_base64 = data['artifacts'][0]["base64"] - -# img_bytes = base64.b64decode(img_base64) - -class SDXLNIMNode: - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "text": ("STRING", { - "multiline": True, - "default": "realistic futuristic city-downtown with short buildings, sunset" - }), - "cfg_scale": ("FLOAT", { - "default": 5.0, - "min": 1.0, - "max": 20.0, - "step": 0.5, - "display": "slider" - }), - "sampler": (["K_DPM_2_ANCESTRAL"],), - "seed": ("INT", { - "default": 0, - "min": 0, - "max": 0xffffffffffffffff, # Max 64-bit unsigned int - "display": "number" - }), - "steps": ("INT", { - "default": 25, - "min": 1, - "max": 100, - "step": 1, - "display": "slider" - }), - }, - } - - RETURN_TYPES = ("IMAGE",) - FUNCTION = "generate" - CATEGORY = "image generation" - - def generate(self, text, cfg_scale, sampler, seed, steps): - payload = { - "text_prompts": [ - { - "text": text, - "weight": 1 - } - ], - "cfg_scale": cfg_scale, - "sampler": sampler, - "seed": seed, - "steps": steps - } - - # Your existing API call code here - response = requests.post(invoke_url, json=payload) - response.raise_for_status() - data = response.json() - img_base64 = data['artifacts'][0]["base64"] - img_bytes = base64.b64decode(img_base64) - - # Convert to the expected format for ComfyUI nodes - # Note: You'll need to convert img_bytes to the proper tensor format - # This is a placeholder - you'll need to implement the actual image conversion - return (img_bytes,) - -# Update the mappings -NODE_CLASS_MAPPINGS = { - "SDXLNIMNode": SDXLNIMNode -} - -NODE_DISPLAY_NAME_MAPPINGS = { - "SDXLNIMNode": "SDXL NIM Node" -} \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 42c8dbf..d961863 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "nimnodes" version = "0.1.0" -description = "Add your description here" +description = "Provides nodes to utilise NVIDIA NIM, a set of accelerated inference microservices that allow you to run AI models on NVIDIA GPUs anywhere." readme = "README.md" requires-python = ">=3.12" dependencies = ["requests>=2.32.3", "ruff>=0.8.2"]