Optimize noise nodes with GPU-accelerated tensors

[djdarcy]

TL;DR version: This PR exclusively focuses on adding GPU-accelerated noise generation to help speed up the generation time. No other changes are added other than optimizations.

Basic problem I kept running into ...

Noise generation nodes (Plasma, Random, Grey, Pink, Brown) experienced severe performance degradation on modern hardware. On a Windows system with RTX 5090 GPU, generating a 1024×1024 noise image took 5 to 10+ seconds, creating workflow bottlenecks. The nodes remained entirely CPU-bound despite available GPU compute, with the bottleneck traced to PIL's putpixel() method called in nested Python loops (1M+ calls per 1024×1024 image).

Problem symptoms:

• Multi-second generation times at moderate resolutions
• GPU idle while CPU struggles with nested loops
• Exponential slowdown at higher resolutions (60-120s for 4096×4096)
• Slow on modern hardware despite a high-end GPU

System specs where issue observed:

• OS: Windows 11
• GPU: NVIDIA RTX 5090 (24GB VRAM)
• ComfyUI: Current version (2025-11-11)
• Python: ComfyUI bundled version

Attempted solution

This PR attempts to address the performance bottleneck by replacing PIL putpixel() nested loops with GPU-accelerated PyTorch tensor operations. The approach quantifiably should provide a substantial speedup on systems with GPU acceleration while maintaining the original API compatibility; and the new code also provides automatic CPU fallback.

Main improvements

GPU-Accelerated Tensor Generation:

• Replace outimage.putpixel((x,y), (nr,ng,nb)) loops with torch.randint() vectorized operations
• Leverage GPU compute via comfy.model_management.get_torch_device()
• Automatic CPU fallback if GPU unavailable or comfy device detection fails
• Direct tensor output eliminates PIL→tensor conversion overhead

Nodes Optimized:

• PlasmaNoise: Optimized output stage (converts Python lists to tensors, vectorizes remap). Recursive subdivision algorithm preserved but output much faster.
• RandNoise: Independent RGB random values per pixel
• GreyNoise: Single greyscale value remapped to RGB channels (correlated)
• PinkNoise: Cube root power transformation for brightness bias
• BrownNoise: Double cube root transformation for stronger brightness bias

Code Reduction:

• PlasmaNoise: ~11 lines of putpixel loops → ~28 lines vectorized (output stage)
• RandNoise: ~80 lines → ~35 lines (56% reduction)
• GreyNoise: ~80 lines → ~48 lines (40% reduction)
• PinkNoise: ~80 lines → ~45 lines (44% reduction)
• BrownNoise: ~80 lines → ~45 lines (44% reduction)
• Total: ~340 lines → ~201 lines (41% reduction)

Implementation details

Device Selection Pattern:

try:
    device = comfy.model_management.get_torch_device()
except:
    device = torch.device("cpu")

Falls back to CPU if ComfyUI device management unavailable.

Vectorized Generation (RandNoise example):

# Original: 262,144 putpixel() calls for 512×512
for y in range(height):
    for x in range(width):
        nr = random.randint(lr, mr)
        outimage.putpixel((x,y), (nr, ng, nb))

# Optimized: 3 GPU kernel launches for 512×512
r = torch.randint(lr, mr+1, (height,width), device=device) / 255.0
g = torch.randint(lg, mg+1, (height,width), device=device) / 255.0
b = torch.randint(lb, mb+1, (height,width), device=device) / 255.0

PlasmaNoise Output Stage (new):

# Original: nested putpixel() loops for final output
for y in range(ah):
    for x in range(aw):
        nr = int(remap(r[x][y], 0, 255, lr, mr))
        ng = int(remap(g[x][y], 0, 255, lg, mg))
        nb = int(remap(b[x][y], 0, 255, lb, mb))
        outimage.putpixel((x,y), (nr, ng, nb))

# Optimized: convert lists to tensors, vectorize remap
r_array = torch.tensor([[r[x][y] for y in range(ah)] for x in range(aw)], device=device).T
# ... same for g, b
r_channel = ((r_array - 0) / 255.0 * (mr - lr) + lr) / 255.0
# Vectorized operations on entire arrays at once

Seed Management:

torch.manual_seed(seed)
if device.type == "cuda":
    torch.cuda.manual_seed(seed)

Maintains reproducibility within new implementation.

Edge Case Handling (GreyNoise):

if mv > lv:  # Avoid division by zero
    r_channel = ((grey - lv) / (mv - lv) * (mr - lr) + lr) / 255.0
else:  # min == max → constant value
    r_channel = torch.full((height, width), lr/255.0, device=device)

Performance results

Testing (Win11 + RTX 5090):

• Before: 5-10+ seconds for 512x512
• After: so fast I barely even see it processing now even with 1024x1024 and larger resolutions
• Estimated: 100-200x speedup (not precisely benchmarked)

Improvement achieved by:

• Eliminating nested Python loops (1M+ → 3-4 GPU kernels per image)
• GPU parallel compute vs sequential CPU operations
• Direct tensor output (no PIL intermediate)
• Vectorized math operations (torch.pow for Pink/Brown transformations)

PlasmaNoise Note: Still slowest overall due to recursive subdivision algorithm (not optimized), but output stage now much faster. Full optimization would require rewriting subdivision algorithm in PyTorch tensors (future work).

Note: Exact speedup ratios not measured with formal benchmarks. Estimates based on observation and operation count reduction.

Limitations

Breaking change: seed incompatibility

Issue: Different random number generator means same seed produces different output.

Reason: Switched from Python random.randint() to PyTorch torch.randint() for GPU acceleration.

Impact:

• Users with saved workflows cannot reproduce exact previous images
• Must adjust seed values to find new desired outputs
• Reproducibility maintained within new implementation (same seed → same output) but not backward compatible

Rationale: Performance gain (100-200x) justifies breaking change. Users can pin to previous version if exact reproducibility required.

No Migration Path Provided: Did not implement seed conversion tool or legacy compatibility mode (would double code complexity).

Testing limitations

Tested environment:

• Win11 Build 26200.6899 (25H2 Insider) + RTX 5090 (driver 581.57) + GPU PyTorch (2.7.0+cu128, CUDA 12.8, cuDNN 9.0.701)
• Only did visual inspection of noise quality
• Parameter controls (per-channel clamping, seed, etc.)
• Subjective performance evaluation

Not tested:

• Linux envs
• macOS systems
• CPU-only systems (fallback should work but unverified)
• AMD GPUs
• Older NVIDIA GPUs
• Very large resolutions (8K+)
• Exhaustive parameter combinations
• Edge cases (min > max, all -1 sentinels, etc.)

No formal benchmarks: Performance improvement based on personal observation and operation count analysis, not precise timing measurements.

Scope limitations

Not included:

• PlasmaNoise recursive subdivision algorithm (only output stage optimized, would require major rewrite)
• PowerImage node (lines 311-357, similar pattern but out of scope)
• Seed migration tool
• Legacy compatibility mode
• Performance benchmark suite
• Automated regression tests

Unknown factors

Unclear:

• Actual speedup ratio on different hardware (only tested RTX 5090)
• CPU fallback performance characteristics
• Behavior on edge case parameter combinations
• Performance at extreme resolutions (8K+)
• Memory usage characteristics on lower-VRAM GPUs
• Compatibility across all PyTorch/ComfyUI versions
• PlasmaNoise speedup magnitude (output stage only, generation still slow)

Potential issues:

• Device detection may fail on some systems
• Float32 precision may introduce minor numeric differences (visually imperceptible for noise)
• Large resolution tensor allocations (~200MB for 8K)
• PlasmaNoise list→tensor conversion overhead may reduce benefit for small images

Changes

nodes.py

• Modified 5 classes: PlasmaNoise, RandNoise, GreyNoise, PinkNoise, BrownNoise
• Line changes: ~340 lines replaced → ~201 lines (41% reduction)
• Removed: PIL putpixel() nested loops, Python random module usage (except in PlasmaNoise subdivision)
• Added: GPU device detection, torch.randint() generation, vectorized operations, tensor conversions
• Preserved: All input parameters, per-channel clamping, -1 sentinel handling, API compatibility, turbulence parameter

Dependencies

• No new dependencies required (PyTorch already used by ComfyUI)
• Still imports PIL (used elsewhere in nodes.py)
• Still imports numpy (used elsewhere in nodes.py)

Technical notes

Important torch.randint() Difference:

# Python: random.randint(a, b) is [a, b] inclusive on both ends
# PyTorch: torch.randint(a, b) is [a, b) exclusive on upper end

# Must add 1 to upper bound:
torch.randint(lr, mr + 1, ...)  # Equivalent to random.randint(lr, mr)

ComfyUI IMAGE format:

• Expected: float32 tensor in [0, 1] range
• Shape: (batch, height, width, channels)
• This PR outputs tensors directly, eliminating PIL→tensor conversion

Memory efficiency:

• Original: PIL Image intermediate + conversion overhead
• Optimized: Direct tensor generation on GPU, minimal CPU↔GPU transfers

PlasmaNoise Hybrid Approach:

• Keeps recursive subdivision in Python lists (complex algorithm, hard to vectorize)
• Converts final pixmaps to tensors for output stage
• Eliminates putpixel() bottleneck while preserving algorithm

Edit to add after some more experimenting

After testing, confirmed:

• Plasma Noise: Output stage faster, but remains slowest overall (subdivision algorithm unchanged)
• Random Noise: so fast I barely even see it processing now
• Brown Noise: Working correctly, fast
• Pink Noise: Working correctly, fast
• Grey Noise: Working correctly, fast

No crashes, errors, or visual artifacts observed. Parameter controls all functional.

---

Questions / feedback welcome

1. Is seed incompatibility acceptable for this performance gain?
2. Would you like more formal benchmarks to be collected before merge?
3. Any concerns about CPU fallback implementation?
4. Interest in fully optimizing PlasmaNoise subdivision algorithm as follow-up PR (but would require major rewrite)?
5. Preference for seed migration tool or documentation?

While the optimizations here attempt to modernize the performance for users with GPU hardware, the goal was to also respect your original CPU-based design and API. Hopefully I managed that balance.

Thanks for creating comfy-plasma!