Joseph IbrahimandClaude Opus 4.5
cd90ee01c6
RTX 4090 performance optimization nodes for ComfyUI. Nodes: - BatchImageProcessor: Efficient multi-image processing - BatchLatentProcessor: Batch latent operations - RTX4090Optimizer: TF32, cuDNN, AMP settings - MemoryManager: VRAM management and cleanup - GPUMonitor: Real-time GPU metrics - TensorRTAutoConverter: Model conversion Optimized for RTX 4090 + Threadripper PRO. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
239 lines
8.3 KiB
Python
239 lines
8.3 KiB
Python
"""
|
|
Batch Processing Nodes for ComfyUI
|
|
Optimized for RTX 4090 parallel processing
|
|
"""
|
|
|
|
import torch
|
|
import numpy as np
|
|
from typing import List, Tuple, Dict, Any
|
|
import folder_paths
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
import asyncio
|
|
|
|
class BatchImageProcessor:
|
|
"""Process multiple images in parallel using RTX 4090"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"batch_size": ("INT", {"default": 4, "min": 1, "max": 16}),
|
|
"operation": (["resize", "crop", "rotate", "flip", "enhance"],),
|
|
"parallel_workers": ("INT", {"default": 4, "min": 1, "max": 8}),
|
|
},
|
|
"optional": {
|
|
"width": ("INT", {"default": 1024, "min": 64, "max": 8192}),
|
|
"height": ("INT", {"default": 1024, "min": 64, "max": 8192}),
|
|
"angle": ("FLOAT", {"default": 0.0, "min": -360.0, "max": 360.0}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "INT")
|
|
RETURN_NAMES = ("images", "count")
|
|
FUNCTION = "process_batch"
|
|
CATEGORY = "RTX4090/batch"
|
|
|
|
def process_batch(self, images, batch_size, operation, parallel_workers,
|
|
width=1024, height=1024, angle=0.0):
|
|
"""Process images in optimized batches"""
|
|
|
|
# Ensure we're using CUDA
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
images = images.to(device)
|
|
|
|
# Split into batches
|
|
total_images = images.shape[0]
|
|
batches = []
|
|
for i in range(0, total_images, batch_size):
|
|
batch = images[i:i+batch_size]
|
|
batches.append(batch)
|
|
|
|
# Process batches in parallel
|
|
processed_batches = []
|
|
|
|
with ThreadPoolExecutor(max_workers=parallel_workers) as executor:
|
|
futures = []
|
|
|
|
for batch in batches:
|
|
if operation == "resize":
|
|
future = executor.submit(self._resize_batch, batch, width, height)
|
|
elif operation == "crop":
|
|
future = executor.submit(self._crop_batch, batch, width, height)
|
|
elif operation == "rotate":
|
|
future = executor.submit(self._rotate_batch, batch, angle)
|
|
elif operation == "flip":
|
|
future = executor.submit(self._flip_batch, batch)
|
|
elif operation == "enhance":
|
|
future = executor.submit(self._enhance_batch, batch)
|
|
|
|
futures.append(future)
|
|
|
|
# Collect results
|
|
for future in futures:
|
|
result = future.result()
|
|
processed_batches.append(result)
|
|
|
|
# Concatenate all processed batches
|
|
processed_images = torch.cat(processed_batches, dim=0)
|
|
|
|
return (processed_images, processed_images.shape[0])
|
|
|
|
def _resize_batch(self, batch, width, height):
|
|
"""Resize batch using RTX 4090 tensor cores"""
|
|
import torch.nn.functional as F
|
|
|
|
# Use bicubic interpolation for quality
|
|
resized = F.interpolate(
|
|
batch.permute(0, 3, 1, 2), # BHWC to BCHW
|
|
size=(height, width),
|
|
mode='bicubic',
|
|
align_corners=False
|
|
).permute(0, 2, 3, 1) # Back to BHWC
|
|
|
|
return resized.clamp(0, 1)
|
|
|
|
def _crop_batch(self, batch, width, height):
|
|
"""Center crop batch"""
|
|
b, h, w, c = batch.shape
|
|
start_h = (h - height) // 2
|
|
start_w = (w - width) // 2
|
|
|
|
return batch[:, start_h:start_h+height, start_w:start_w+width, :]
|
|
|
|
def _rotate_batch(self, batch, angle):
|
|
"""Rotate batch using GPU acceleration"""
|
|
import torchvision.transforms.functional as TF
|
|
|
|
# Convert to appropriate format for rotation
|
|
batch_rotated = []
|
|
for img in batch:
|
|
img_chw = img.permute(2, 0, 1) # HWC to CHW
|
|
rotated = TF.rotate(img_chw, angle)
|
|
batch_rotated.append(rotated.permute(1, 2, 0)) # Back to HWC
|
|
|
|
return torch.stack(batch_rotated)
|
|
|
|
def _flip_batch(self, batch):
|
|
"""Flip batch horizontally"""
|
|
return torch.flip(batch, dims=[2])
|
|
|
|
def _enhance_batch(self, batch):
|
|
"""AI-based enhancement using RTX 4090"""
|
|
# Simple enhancement - can be replaced with more sophisticated model
|
|
# Increase contrast and saturation
|
|
mean = batch.mean(dim=(1, 2), keepdim=True)
|
|
enhanced = (batch - mean) * 1.2 + mean
|
|
|
|
return enhanced.clamp(0, 1)
|
|
|
|
|
|
class BatchLatentProcessor:
|
|
"""Process multiple latents efficiently"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"latents": ("LATENT",),
|
|
"operation": (["denoise", "blend", "interpolate", "upscale"],),
|
|
"strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0}),
|
|
},
|
|
"optional": {
|
|
"reference_latent": ("LATENT",),
|
|
"blend_factor": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
FUNCTION = "process_latents"
|
|
CATEGORY = "RTX4090/batch"
|
|
|
|
def process_latents(self, latents, operation, strength,
|
|
reference_latent=None, blend_factor=0.5):
|
|
"""Process latents with GPU optimization"""
|
|
|
|
samples = latents["samples"]
|
|
device = samples.device
|
|
|
|
if operation == "denoise":
|
|
# Simple denoising in latent space
|
|
processed = self._denoise_latent(samples, strength)
|
|
|
|
elif operation == "blend" and reference_latent is not None:
|
|
# Blend with reference
|
|
ref_samples = reference_latent["samples"]
|
|
processed = samples * (1 - blend_factor) + ref_samples * blend_factor
|
|
|
|
elif operation == "interpolate" and reference_latent is not None:
|
|
# Spherical interpolation
|
|
processed = self._slerp(samples, reference_latent["samples"], blend_factor)
|
|
|
|
elif operation == "upscale":
|
|
# Latent upscaling
|
|
import torch.nn.functional as F
|
|
processed = F.interpolate(
|
|
samples,
|
|
scale_factor=strength,
|
|
mode='bilinear',
|
|
align_corners=False
|
|
)
|
|
else:
|
|
processed = samples
|
|
|
|
return ({"samples": processed},)
|
|
|
|
def _denoise_latent(self, latent, strength):
|
|
"""Apply denoising in latent space"""
|
|
# Simple Gaussian blur for denoising
|
|
import torch.nn.functional as F
|
|
|
|
# Create Gaussian kernel
|
|
kernel_size = int(3 + strength * 2) | 1 # Ensure odd
|
|
sigma = strength
|
|
|
|
# Apply Gaussian blur
|
|
channels = latent.shape[1]
|
|
kernel = self._gaussian_kernel(kernel_size, sigma).to(latent.device)
|
|
kernel = kernel.expand(channels, 1, kernel_size, kernel_size)
|
|
|
|
# Apply convolution
|
|
denoised = F.conv2d(
|
|
latent, kernel,
|
|
padding=kernel_size//2,
|
|
groups=channels
|
|
)
|
|
|
|
return denoised
|
|
|
|
def _gaussian_kernel(self, size, sigma):
|
|
"""Create Gaussian kernel"""
|
|
coords = torch.arange(size).float() - (size - 1) / 2
|
|
g = coords**2
|
|
g = (-(g.unsqueeze(0) + g.unsqueeze(1)) / (2 * sigma**2)).exp()
|
|
g /= g.sum()
|
|
return g.unsqueeze(0).unsqueeze(0)
|
|
|
|
def _slerp(self, a, b, t):
|
|
"""Spherical linear interpolation"""
|
|
# Flatten to vectors
|
|
a_flat = a.flatten(start_dim=1)
|
|
b_flat = b.flatten(start_dim=1)
|
|
|
|
# Compute angle
|
|
dots = (a_flat * b_flat).sum(dim=1, keepdim=True)
|
|
norms_a = a_flat.norm(dim=1, keepdim=True)
|
|
norms_b = b_flat.norm(dim=1, keepdim=True)
|
|
|
|
cos_angle = dots / (norms_a * norms_b + 1e-8)
|
|
angle = cos_angle.acos()
|
|
|
|
# SLERP formula
|
|
sin_angle = angle.sin()
|
|
wa = ((1 - t) * angle).sin() / (sin_angle + 1e-8)
|
|
wb = (t * angle).sin() / (sin_angle + 1e-8)
|
|
|
|
# Interpolate
|
|
result = wa.unsqueeze(-1).unsqueeze(-1) * a + wb.unsqueeze(-1).unsqueeze(-1) * b
|
|
|
|
return result |