jonstreeter-ComfyUI-Reference-Based-Video-Colorization/PERFORMANCE.md

Performance Optimization Guide

This guide covers all available performance optimizations for ColorMNet and Deep Exemplar nodes.

Quick Summary

Optimization	Speedup	Nodes	Installation
torch.compile	15-25%	All 4 nodes	Built-in (PyTorch 2.0+)
SageAttention	20-30%	Deep Exemplar only	pip install sageattention
CUDA Correlation	2-5x attention	ColorMNet only	Auto-install (requires CUDA toolkit)
FP16	~2x	ColorMNet only	Built-in parameter

---

torch.compile (PyTorch 2.0+)

Available for: All 4 nodes (ColorMNet Video, ColorMNet Image, Deep Exemplar Video, Deep Exemplar Image)

What is it?

PyTorch 2.0's JIT compiler that optimizes model execution through graph fusion, kernel optimization, and memory layout improvements.

Performance Impact

• Expected speedup: 15-25% after warmup
• First run: Slower due to compilation overhead (30-60s additional time)
• Subsequent runs: Full speedup benefit
• Quality: Identical to non-compiled version

Usage

Enable via node parameter:

use_torch_compile = True

Requirements

• PyTorch 2.0 or later
• No additional installation needed
• Works on both CPU and GPU

Troubleshooting

If compilation fails, the nodes automatically fall back to standard mode with a warning message. Check console output for details.

---

SageAttention

Available for: Deep Exemplar Video and Image nodes only

What is it?

INT8-quantized attention implementation that provides significant speedup for self-attention operations with minimal quality impact.

Performance Impact

• Expected speedup: 20-30% on attention operations
• Overall speedup: ~15-20% for full pipeline
• Quality: Negligible difference due to INT8 quantization
• VRAM: Slightly reduced usage

Installation

pip install sageattention

Usage

Enable via node parameter:

use_sage_attention = True

Requirements

• CUDA-capable GPU
• PyTorch with CUDA support
• sageattention package installed

How it Works

SageAttention integrates into the Self_Attn layer in NonlocalNet:

• Automatically converts standard attention to INT8-quantized version
• Falls back to standard attention if unavailable or if errors occur
• No model retraining required

Troubleshooting

If SageAttention fails to load:

1. Verify CUDA is available: python -c "import torch; print(torch.cuda.is_available())"
2. Check package installation: pip show sageattention
3. The node will automatically fall back to standard attention with equivalent quality

---

Automatic Performance Checks

ColorMNet automatically checks for and attempts to install performance optimizations on first run.

CUDA Correlation Sampler

The spatial_correlation_sampler CUDA extension provides 2-5x faster short-term attention processing.

Requirements Check

On first run, ColorMNet will automatically check for:

• ✅ CUDA Toolkit installed and CUDA_HOME set
• ✅ NVIDIA CUDA compiler (nvcc)
• ✅ Visual Studio C++ compiler (cl.exe)
• ✅ PyTorch CUDA availability

Automatic Installation

If all requirements are met, ColorMNet will:

1. Automatically compile and install spatial_correlation_sampler
2. Takes 2-5 minutes on first run
3. Enable optimized mode automatically

Performance Modes

Mode	Speed	Requirements
Optimized	100% (fastest)	CUDA extension installed
Fallback	60-70% speed	No extra requirements

Note: Both modes produce identical quality - the difference is only speed!

Installation Instructions

If automatic installation fails, you can install manually:

1. Install CUDA Toolkit

Download from: https://developer.nvidia.com/cuda-toolkit

Important: Match your PyTorch CUDA version

• Check with: python -c "import torch; print(torch.version.cuda)"
• Install matching CUDA Toolkit (e.g., 12.8 for PyTorch CUDA 12.8)

2. Set CUDA_HOME Environment Variable

Windows:

# System Properties → Advanced → Environment Variables → New
CUDA_HOME = C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.8

Linux/Mac:

export CUDA_HOME=/usr/local/cuda

3. Install Visual Studio Build Tools (Windows Only)

Download: https://visualstudio.microsoft.com/downloads/

Select: Desktop development with C++

4. Manual Installation

After setting up requirements:

pip install git+https://github.com/ClementPinard/Pytorch-Correlation-extension.git

5. Verify Installation

Restart ComfyUI and check the console output:

[ColorMNet] ✓ spatial_correlation_sampler already installed
[ColorMNet] 🚀 Optimized CUDA mode enabled (2-5x faster)
Correlation sampler: enabled (CUDA optimized)

Troubleshooting

"CUDA_HOME not set"

Set the environment variable to your CUDA installation directory.

"Visual Studio C++ compiler not found"

Install Visual Studio Build Tools with C++ workload.

"Compilation error"

Ensure CUDA Toolkit version matches your PyTorch CUDA version:

python -c "import torch; print(torch.version.cuda)"

Still not working?

The fallback mode works perfectly! You'll get:

• Same quality results
• 20-40% slower processing
• No additional setup required

Performance Comparison

Example: 196 frames @ 480x640 on RTX 5090

Mode	Time	FPS
Optimized	~45s	4.4
Fallback	~65s	3.0

The speed difference is noticeable but not critical for most use cases.

---

Combining Optimizations

You can enable multiple optimizations simultaneously for maximum performance:

ColorMNet Maximum Performance

use_fp16 = True
use_torch_compile = True
memory_mode = "low_memory"  # or "balanced"

Expected combined speedup: 2.5-3x vs baseline (FP32, no compile)

Deep Exemplar Maximum Performance

use_torch_compile = True
use_sage_attention = True
use_half_resolution = True
frame_propagate = True

Expected combined speedup: 1.4-1.6x vs baseline (no optimizations)

Performance Benchmark Example

Test: 240 frames @ 768x432 on RTX 5090

Configuration	ColorMNet	Deep Exemplar
Baseline (no opts)	~75s (3.2 FPS)	~68s (3.5 FPS)
FP16 only	~45s (5.3 FPS)	N/A
torch.compile only	~58s (4.1 FPS)	~55s (4.4 FPS)
All optimizations	~36s (6.7 FPS)	~47s (5.1 FPS)

Note: Actual performance varies by GPU, resolution, and content.

---

Optimization Recommendations by GPU

Low-end (6-8GB VRAM)

ColorMNet:

• use_fp16 = True
• use_torch_compile = True
• memory_mode = "low_memory"

Deep Exemplar:

• use_half_resolution = True
• use_torch_compile = True
• use_sage_attention = True

Mid-range (8-16GB VRAM)

ColorMNet:

• use_fp16 = True
• use_torch_compile = True
• memory_mode = "balanced"

Deep Exemplar:

• use_half_resolution = True (optional)
• use_torch_compile = True
• use_sage_attention = True

High-end (16GB+ VRAM)

ColorMNet:

• use_fp16 = False (best quality)
• use_torch_compile = True
• memory_mode = "high_quality"

Deep Exemplar:

• use_half_resolution = False
• use_torch_compile = True
• use_sage_attention = True

---

Quality Impact Summary

Optimization	Quality Impact
torch.compile	None - Identical output
SageAttention	Negligible - INT8 quantization barely visible
FP16	Minimal - Slight numerical differences
CUDA Correlation	None - Identical output

Recommendation: Enable all applicable optimizations for your hardware. The quality impact is minimal to none while providing significant speed improvements.