""" 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