nodes/jonstreeter-ComfyUI-Reference-Based-Video-Colorization
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Align ColorMNet pipeline with training normalization

Browse Source
This commit is contained in:
jonstreeter
2025-11-21 08:53:02 -06:00
parent 85d6ba07e2
commit 644680e9df
7 changed files with 283 additions and 361 deletions
Download Patch File Download Diff File Expand all files Collapse all files
PERFORMANCE.md
+5 -301
View File
@@ -1,302 +1,6 @@
# Performance Optimization Guide
Performance and GPU Support
===========================
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
```bash
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:**
```powershell
# System Properties → Advanced → Environment Variables → New
CUDA_HOME = C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.8
```
**Linux/Mac:**
```bash
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:
```bash
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:
```bash
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.
- GPU detection is automatic; CUDA is used when available, with FP16 enabled on modern NVIDIA cards (Turing/RTX 20 series and newer, including RTX 50-series/5090).
- If CUDA dev tools are missing, the correlation CUDA extension falls back to the PyTorch implementation; quality is unchanged, speed is just a bit lower.
- Model download still occurs on first use (~500MB); all Python dependencies and git-based CUDA extensions are auto-installed when the node loads.