Merge pull request #1 from Comfy-Org/implement

Implement initial version
2025-01-17 21:07:06 +00:00
parent 309d3672c1 2edbcecadb
commit 5ad1c3e9ca
9 changed files with 262 additions and 124 deletions
@@ -0,0 +1 @@
+__pycache__
@@ -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)  
@@ -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"
+}
@@ -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
+}
@@ -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"
-}
@@ -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"]