27 KiB
27 KiB
Example Custom Node Implementations
Complete Working Examples
These examples demonstrate various custom node patterns and best practices.
1. Image Processing Node
A complete image filter node with multiple options:
import torch
import torch.nn.functional as F
import numpy as np
from PIL import Image, ImageFilter
class AdvancedImageFilter:
"""
Advanced image filtering with multiple algorithms and options.
Demonstrates image processing, batch handling, and mask support.
"""
def __init__(self):
self.type = "AdvancedImageFilter"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE", {"tooltip": "Input image to filter"}),
"filter_type": ([
"gaussian_blur",
"motion_blur",
"unsharp_mask",
"edge_enhance",
"emboss",
"find_edges",
"contour"
], {
"default": "gaussian_blur"
}),
"strength": ("FLOAT", {
"default": 1.0,
"min": 0.0,
"max": 10.0,
"step": 0.1,
"display": "slider"
}),
},
"optional": {
"mask": ("MASK", {
"tooltip": "Optional mask for selective filtering"
}),
"blend_mode": (["normal", "multiply", "screen", "overlay"], {
"default": "normal"
}),
}
}
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("filtered_image",)
OUTPUT_TOOLTIPS = ("The filtered image",)
FUNCTION = "apply_filter"
CATEGORY = "image/filters"
DESCRIPTION = "Apply various image filters with adjustable strength"
def apply_filter(self, image, filter_type, strength, mask=None, blend_mode="normal"):
# Handle batch processing
batch_size = image.shape[0]
results = []
for i in range(batch_size):
# Get single image
single_image = image[i]
# Apply filter
filtered = self._apply_single_filter(single_image, filter_type, strength)
# Apply mask if provided
if mask is not None:
single_mask = mask[i] if len(mask.shape) == 3 else mask
filtered = self._apply_mask(single_image, filtered, single_mask)
# Blend with original
if blend_mode != "normal":
filtered = self._blend_images(single_image, filtered, blend_mode, strength)
results.append(filtered)
# Stack results back into batch
result_batch = torch.stack(results)
return (result_batch,)
def _apply_single_filter(self, image, filter_type, strength):
"""Apply filter to single image"""
if filter_type == "gaussian_blur":
sigma = strength * 2.0
kernel_size = int(sigma * 4) | 1 # Ensure odd
return self._gaussian_blur(image, kernel_size, sigma)
elif filter_type == "motion_blur":
return self._motion_blur(image, int(strength * 10) + 1)
elif filter_type == "unsharp_mask":
return self._unsharp_mask(image, strength)
elif filter_type in ["edge_enhance", "emboss", "find_edges", "contour"]:
# Convert to PIL for these filters
pil_image = self._tensor_to_pil(image)
if filter_type == "edge_enhance":
filtered = pil_image.filter(ImageFilter.EDGE_ENHANCE_MORE)
elif filter_type == "emboss":
filtered = pil_image.filter(ImageFilter.EMBOSS)
elif filter_type == "find_edges":
filtered = pil_image.filter(ImageFilter.FIND_EDGES)
else: # contour
filtered = pil_image.filter(ImageFilter.CONTOUR)
# Convert back to tensor
result = self._pil_to_tensor(filtered)
# Blend with original based on strength
return image * (1 - strength) + result * strength
return image
def _gaussian_blur(self, image, kernel_size, sigma):
"""Apply Gaussian blur using PyTorch"""
# Add batch and channel dimensions for conv2d
image = image.permute(2, 0, 1).unsqueeze(0) # [1, C, H, W]
# Create Gaussian kernel
kernel = self._create_gaussian_kernel(kernel_size, sigma)
kernel = kernel.expand(image.shape[1], 1, kernel_size, kernel_size)
# Apply convolution
blurred = F.conv2d(image, kernel, padding=kernel_size//2, groups=image.shape[1])
# Convert back to [H, W, C]
return blurred.squeeze(0).permute(1, 2, 0)
def _create_gaussian_kernel(self, size, sigma):
"""Create 2D Gaussian kernel"""
coords = torch.arange(size, dtype=torch.float32)
coords -= size // 2
g = torch.exp(-(coords ** 2) / (2 * sigma ** 2))
g /= g.sum()
return g.unsqueeze(0) * g.unsqueeze(1)
def _motion_blur(self, image, size):
"""Apply motion blur"""
# Create motion kernel (horizontal)
kernel = torch.zeros(size, size)
kernel[size // 2, :] = 1.0 / size
# Apply to each channel
image = image.permute(2, 0, 1).unsqueeze(0)
kernel = kernel.unsqueeze(0).unsqueeze(0).expand(image.shape[1], 1, size, size)
blurred = F.conv2d(image, kernel, padding=size//2, groups=image.shape[1])
return blurred.squeeze(0).permute(1, 2, 0)
def _unsharp_mask(self, image, strength):
"""Apply unsharp mask sharpening"""
# Blur the image
blurred = self._gaussian_blur(image, 5, 1.0)
# Create mask
mask = image - blurred
# Apply sharpening
sharpened = image + mask * strength
# Clamp values
return torch.clamp(sharpened, 0, 1)
def _apply_mask(self, original, filtered, mask):
"""Apply mask to blend original and filtered images"""
# Ensure mask has 3 channels
if len(mask.shape) == 2:
mask = mask.unsqueeze(2).expand(-1, -1, 3)
return original * (1 - mask) + filtered * mask
def _blend_images(self, img1, img2, mode, opacity):
"""Blend two images using various modes"""
if mode == "multiply":
blended = img1 * img2
elif mode == "screen":
blended = 1 - (1 - img1) * (1 - img2)
elif mode == "overlay":
mask = img1 > 0.5
blended = torch.where(
mask,
1 - 2 * (1 - img1) * (1 - img2),
2 * img1 * img2
)
else: # normal
blended = img2
# Apply opacity
return img1 * (1 - opacity) + blended * opacity
def _tensor_to_pil(self, tensor):
"""Convert tensor to PIL Image"""
array = (tensor.cpu().numpy() * 255).astype(np.uint8)
return Image.fromarray(array)
def _pil_to_tensor(self, pil_image):
"""Convert PIL Image to tensor"""
array = np.array(pil_image).astype(np.float32) / 255.0
return torch.from_numpy(array).to(tensor.device)
@classmethod
def IS_CHANGED(cls, **kwargs):
# Always re-execute to see changes
return float("nan")
# Register the node
NODE_CLASS_MAPPINGS = {
"AdvancedImageFilter": AdvancedImageFilter
}
NODE_DISPLAY_NAME_MAPPINGS = {
"AdvancedImageFilter": "Advanced Image Filter"
}
2. Batch Processing Node
Handle multiple images with different operations:
class BatchImageProcessor:
"""
Process batches of images with different operations per image.
Demonstrates dynamic inputs and batch handling.
"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"images": ("IMAGE", {"tooltip": "Batch of images to process"}),
"operation": ([
"rotate_sequence",
"brightness_gradient",
"color_shift",
"zoom_sequence"
], {"default": "rotate_sequence"}),
"start_value": ("FLOAT", {
"default": 0.0,
"min": -1.0,
"max": 1.0,
"step": 0.1
}),
"end_value": ("FLOAT", {
"default": 1.0,
"min": -1.0,
"max": 1.0,
"step": 0.1
}),
}
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "process_batch"
CATEGORY = "image/batch"
def process_batch(self, images, operation, start_value, end_value):
batch_size = images.shape[0]
results = []
# Calculate value for each image in sequence
for i in range(batch_size):
t = i / max(1, batch_size - 1) # Normalize to 0-1
value = start_value + (end_value - start_value) * t
image = images[i]
if operation == "rotate_sequence":
angle = value * 360
result = self._rotate_image(image, angle)
elif operation == "brightness_gradient":
result = torch.clamp(image * (1 + value), 0, 1)
elif operation == "color_shift":
result = self._color_shift(image, value)
elif operation == "zoom_sequence":
scale = 1 + value
result = self._zoom_image(image, scale)
results.append(result)
return (torch.stack(results),)
def _rotate_image(self, image, angle):
"""Rotate image by angle degrees"""
# Convert to PIL for rotation
pil_img = Image.fromarray((image.cpu().numpy() * 255).astype(np.uint8))
rotated = pil_img.rotate(angle, expand=False, fillcolor=(0, 0, 0))
# Convert back
array = np.array(rotated).astype(np.float32) / 255.0
return torch.from_numpy(array).to(image.device)
def _color_shift(self, image, shift):
"""Shift colors in HSV space"""
# Simple RGB channel rotation
shifted = torch.zeros_like(image)
if shift > 0:
shifted[:, :, 0] = image[:, :, 1] * shift + image[:, :, 0] * (1 - shift)
shifted[:, :, 1] = image[:, :, 2] * shift + image[:, :, 1] * (1 - shift)
shifted[:, :, 2] = image[:, :, 0] * shift + image[:, :, 2] * (1 - shift)
else:
shift = -shift
shifted[:, :, 0] = image[:, :, 2] * shift + image[:, :, 0] * (1 - shift)
shifted[:, :, 1] = image[:, :, 0] * shift + image[:, :, 1] * (1 - shift)
shifted[:, :, 2] = image[:, :, 1] * shift + image[:, :, 2] * (1 - shift)
return shifted
def _zoom_image(self, image, scale):
"""Zoom image by scale factor"""
h, w = image.shape[:2]
new_h, new_w = int(h * scale), int(w * scale)
# Add batch dimension for interpolation
img_batch = image.permute(2, 0, 1).unsqueeze(0)
# Resize
resized = F.interpolate(img_batch, size=(new_h, new_w), mode='bilinear', align_corners=False)
resized = resized.squeeze(0).permute(1, 2, 0)
# Crop or pad to original size
if scale > 1:
# Crop center
start_h = (new_h - h) // 2
start_w = (new_w - w) // 2
result = resized[start_h:start_h+h, start_w:start_w+w]
else:
# Pad
pad_h = (h - new_h) // 2
pad_w = (w - new_w) // 2
result = torch.zeros_like(image)
result[pad_h:pad_h+new_h, pad_w:pad_w+new_w] = resized
return result
3. Model Integration Node
Integrate with external models or APIs:
import json
import requests
import base64
from io import BytesIO
class ExternalAPINode:
"""
Integrate with external APIs for processing.
Demonstrates async operations and error handling.
"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE", {}),
"api_endpoint": ("STRING", {
"default": "http://localhost:8000/process",
"multiline": False
}),
"api_key": ("STRING", {
"default": "",
"multiline": False
}),
"parameters": ("STRING", {
"default": "{}",
"multiline": True,
"tooltip": "JSON parameters for the API"
}),
},
"optional": {
"timeout": ("INT", {
"default": 30,
"min": 1,
"max": 300
}),
}
}
RETURN_TYPES = ("IMAGE", "STRING")
RETURN_NAMES = ("processed_image", "api_response")
FUNCTION = "call_api"
CATEGORY = "api/external"
def call_api(self, image, api_endpoint, api_key, parameters, timeout=30):
try:
# Parse parameters
params = json.loads(parameters)
except json.JSONDecodeError as e:
raise ValueError(f"Invalid JSON parameters: {e}")
# Convert image to base64
pil_image = self._tensor_to_pil(image[0]) # First image in batch
# Prepare request
buffered = BytesIO()
pil_image.save(buffered, format="PNG")
img_base64 = base64.b64encode(buffered.getvalue()).decode()
# Build request payload
payload = {
"image": img_base64,
"parameters": params
}
headers = {
"Content-Type": "application/json"
}
if api_key:
headers["Authorization"] = f"Bearer {api_key}"
try:
# Make API request
response = requests.post(
api_endpoint,
json=payload,
headers=headers,
timeout=timeout
)
response.raise_for_status()
# Parse response
result = response.json()
# Extract processed image
if "image" in result:
processed_img = self._base64_to_tensor(result["image"])
else:
raise ValueError("No image in API response")
# Return processed image and response
return (processed_img, json.dumps(result, indent=2))
except requests.exceptions.RequestException as e:
raise RuntimeError(f"API request failed: {e}")
def _tensor_to_pil(self, tensor):
"""Convert tensor to PIL Image"""
array = (tensor.cpu().numpy() * 255).astype(np.uint8)
return Image.fromarray(array)
def _base64_to_tensor(self, base64_str):
"""Convert base64 string to tensor"""
img_data = base64.b64decode(base64_str)
pil_image = Image.open(BytesIO(img_data))
array = np.array(pil_image).astype(np.float32) / 255.0
tensor = torch.from_numpy(array)
# Add batch dimension
return tensor.unsqueeze(0)
4. Custom Widget Node
Node with custom UI components:
class ColorPaletteNode:
"""
Generate color palettes with custom UI.
Demonstrates custom widgets and client-server communication.
"""
def __init__(self):
self.colors = []
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"palette_size": ("INT", {
"default": 5,
"min": 2,
"max": 10
}),
"base_color": ("STRING", {
"default": "#ff0000",
"multiline": False,
"tooltip": "Base color in hex format"
}),
"harmony_type": ([
"complementary",
"analogous",
"triadic",
"split_complementary",
"monochromatic"
], {"default": "complementary"}),
},
"hidden": {
"unique_id": "UNIQUE_ID"
}
}
RETURN_TYPES = ("IMAGE", "STRING")
RETURN_NAMES = ("palette_image", "color_list")
FUNCTION = "generate_palette"
CATEGORY = "color/palette"
def generate_palette(self, palette_size, base_color, harmony_type, unique_id):
# Convert hex to RGB
base_rgb = self._hex_to_rgb(base_color)
# Generate palette based on harmony type
colors = self._generate_colors(base_rgb, harmony_type, palette_size)
# Create palette image
palette_image = self._create_palette_image(colors)
# Create color list string
color_list = json.dumps([self._rgb_to_hex(c) for c in colors])
# Send to frontend for custom display
from server import PromptServer
PromptServer.instance.send_sync(
"color_palette.update",
{
"node_id": unique_id,
"colors": [self._rgb_to_hex(c) for c in colors]
}
)
return (palette_image, color_list)
def _hex_to_rgb(self, hex_color):
"""Convert hex color to RGB tuple"""
hex_color = hex_color.lstrip('#')
return tuple(int(hex_color[i:i+2], 16) / 255.0 for i in (0, 2, 4))
def _rgb_to_hex(self, rgb):
"""Convert RGB tuple to hex color"""
return '#{:02x}{:02x}{:02x}'.format(
int(rgb[0] * 255),
int(rgb[1] * 255),
int(rgb[2] * 255)
)
def _generate_colors(self, base_rgb, harmony_type, count):
"""Generate color harmony"""
import colorsys
# Convert to HSV for easier manipulation
h, s, v = colorsys.rgb_to_hsv(*base_rgb)
colors = [base_rgb]
if harmony_type == "complementary":
# Add complementary color
comp_h = (h + 0.5) % 1.0
colors.append(colorsys.hsv_to_rgb(comp_h, s, v))
# Fill rest with variations
for i in range(2, count):
var_s = s * (0.5 + (i / count) * 0.5)
var_v = v * (0.5 + (i / count) * 0.5)
colors.append(colorsys.hsv_to_rgb(h if i % 2 == 0 else comp_h, var_s, var_v))
elif harmony_type == "analogous":
# Colors adjacent on color wheel
step = 0.083 # 30 degrees
for i in range(1, count):
new_h = (h + step * (i - count//2)) % 1.0
colors.append(colorsys.hsv_to_rgb(new_h, s, v))
elif harmony_type == "triadic":
# Three colors equally spaced
colors.append(colorsys.hsv_to_rgb((h + 0.333) % 1.0, s, v))
colors.append(colorsys.hsv_to_rgb((h + 0.667) % 1.0, s, v))
# Fill rest with variations
for i in range(3, count):
idx = i % 3
base_h = [h, (h + 0.333) % 1.0, (h + 0.667) % 1.0][idx]
var_s = s * (0.7 + (i / count) * 0.3)
colors.append(colorsys.hsv_to_rgb(base_h, var_s, v))
elif harmony_type == "split_complementary":
# Base + two colors adjacent to complement
comp_h = (h + 0.5) % 1.0
colors.append(colorsys.hsv_to_rgb((comp_h - 0.083) % 1.0, s, v))
colors.append(colorsys.hsv_to_rgb((comp_h + 0.083) % 1.0, s, v))
# Fill rest
for i in range(3, count):
var_v = v * (0.5 + (i / count) * 0.5)
colors.append(colorsys.hsv_to_rgb(h, s, var_v))
else: # monochromatic
# Variations in saturation and value
for i in range(1, count):
var_s = s * (i / count)
var_v = v * (0.3 + (i / count) * 0.7)
colors.append(colorsys.hsv_to_rgb(h, var_s, var_v))
return colors[:count]
def _create_palette_image(self, colors):
"""Create visual representation of palette"""
swatch_width = 128
swatch_height = 128
width = swatch_width * len(colors)
height = swatch_height
# Create image tensor
image = torch.zeros(1, height, width, 3)
for i, color in enumerate(colors):
x_start = i * swatch_width
x_end = (i + 1) * swatch_width
for c in range(3):
image[0, :, x_start:x_end, c] = color[c]
return image
5. Data Processing Node
Process non-image data:
class DataAnalysisNode:
"""
Analyze and visualize data from various sources.
Demonstrates handling of non-image data types.
"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"data": ("STRING", {
"multiline": True,
"default": "1,2,3,4,5\n2,4,6,8,10",
"tooltip": "CSV data or JSON array"
}),
"data_type": (["csv", "json", "time_series"], {
"default": "csv"
}),
"visualization": (["line_chart", "bar_chart", "scatter", "heatmap"], {
"default": "line_chart"
}),
"width": ("INT", {
"default": 512,
"min": 256,
"max": 2048
}),
"height": ("INT", {
"default": 512,
"min": 256,
"max": 2048
}),
}
}
RETURN_TYPES = ("IMAGE", "STRING")
RETURN_NAMES = ("visualization", "statistics")
FUNCTION = "analyze_data"
CATEGORY = "data/analysis"
def analyze_data(self, data, data_type, visualization, width, height):
# Parse data
parsed_data = self._parse_data(data, data_type)
# Calculate statistics
stats = self._calculate_statistics(parsed_data)
# Create visualization
vis_image = self._create_visualization(
parsed_data, visualization, width, height
)
# Format statistics as JSON
stats_json = json.dumps(stats, indent=2)
return (vis_image, stats_json)
def _parse_data(self, data, data_type):
"""Parse input data based on type"""
if data_type == "csv":
lines = data.strip().split('\n')
parsed = []
for line in lines:
values = [float(x.strip()) for x in line.split(',')]
parsed.append(values)
return parsed
elif data_type == "json":
return json.loads(data)
else: # time_series
values = [float(x.strip()) for x in data.split(',')]
return [list(range(len(values))), values]
def _calculate_statistics(self, data):
"""Calculate basic statistics"""
import statistics
flat_data = []
for row in data:
if isinstance(row, list):
flat_data.extend(row)
else:
flat_data.append(row)
if not flat_data:
return {}
return {
"count": len(flat_data),
"mean": statistics.mean(flat_data),
"median": statistics.median(flat_data),
"stdev": statistics.stdev(flat_data) if len(flat_data) > 1 else 0,
"min": min(flat_data),
"max": max(flat_data)
}
def _create_visualization(self, data, vis_type, width, height):
"""Create visualization using matplotlib"""
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
fig, ax = plt.subplots(figsize=(width/100, height/100))
if vis_type == "line_chart":
for i, row in enumerate(data):
ax.plot(row, label=f"Series {i+1}")
ax.legend()
elif vis_type == "bar_chart":
if len(data) == 1:
ax.bar(range(len(data[0])), data[0])
else:
x = np.arange(len(data[0]))
width_bar = 0.8 / len(data)
for i, row in enumerate(data):
ax.bar(x + i * width_bar, row, width_bar, label=f"Series {i+1}")
ax.legend()
elif vis_type == "scatter":
if len(data) >= 2:
ax.scatter(data[0], data[1])
else:
ax.scatter(range(len(data[0])), data[0])
else: # heatmap
im = ax.imshow(data, cmap='viridis', aspect='auto')
plt.colorbar(im, ax=ax)
ax.set_title(f"{vis_type.replace('_', ' ').title()}")
plt.tight_layout()
# Convert to image tensor
fig.canvas.draw()
buf = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
buf = buf.reshape(fig.canvas.get_width_height()[::-1] + (3,))
plt.close(fig)
# Convert to tensor and normalize
tensor = torch.from_numpy(buf.copy()).float() / 255.0
# Add batch dimension
return tensor.unsqueeze(0)
Testing Your Nodes
Create a test file to verify your nodes work correctly:
# test_nodes.py
import sys
import torch
sys.path.append('../..') # Path to ComfyUI
from custom_nodes.my_nodes import *
def test_image_filter():
# Create test image
test_image = torch.rand(1, 256, 256, 3)
# Create node instance
node = AdvancedImageFilter()
# Test each filter type
for filter_type in ["gaussian_blur", "edge_enhance", "emboss"]:
result = node.apply_filter(
image=test_image,
filter_type=filter_type,
strength=0.5
)
assert result[0].shape == test_image.shape
print(f"✓ {filter_type} test passed")
if __name__ == "__main__":
test_image_filter()
print("All tests passed!")
Next Steps
- Learn about Frontend Integration
- Explore Publishing and Distribution
- Read Best Practices and Tips