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.
README.md
+19 -1
View File
@@ -77,13 +77,29 @@ pip install -r requirements.txt
Restart ComfyUI. Models download automatically on first use.
### Recent changes
- Fixed ColorMNet pipeline normalization to match training (L-only input to encoder, normalized ab masks).
- Added auto-install for git-based CUDA deps (py-thin-plate-spline, correlation sampler) on node load.
- New sample: `Workflows/ColorMNet_Image_Workflow.json` for single-image colorization.
### ColorMNet git-based dependencies
ColorMNet uses two CUDA extensions that are shipped as git repos, not PyPI packages. The node now installs them automatically when it loads, so no manual steps are required. If your environment blocks installs and you need to do it yourself, the commands are:
```bash
pip install git+https://github.com/cheind/py-thin-plate-spline.git
pip install git+https://github.com/ClementPinard/Pytorch-Correlation-extension.git
```
If compilation fails on Windows, install the **Desktop development with C++** workload in Visual Studio Build Tools and ensure `CUDA_HOME` points to your CUDA Toolkit path. The nodes will still run without these extensions, just a bit slower.
---
## 🚀 Quick Start
### Example Workflow
Load the example workflow from `workflows/Colorize_Video_Workflow.json`:
Load the example workflow from `Workflows/Colorize_Video_Workflow.json`:
1. **Load Video** - Use VHS Video Loader to load your grayscale video
2. **Load Reference** - Load a color reference image
@@ -91,6 +107,8 @@ Load the example workflow from `workflows/Colorize_Video_Workflow.json`:
4. **Compare Results** - Use the performance reports to benchmark
5. **Save Output** - Export colorized video with VHS Video Combine
For single-frame work, use the lightweight image workflow at `Workflows/ColorMNet_Image_Workflow.json`. It loads a target image, loads a reference image, runs the `ColorMNetImage` node, and saves the result.
### Nodes Overview
#### ColorMNet Nodes
Workflows/ColorMNet_Image_Workflow.json
+190
View File
@@ -0,0 +1,190 @@
{
"last_node_id": 4,
"last_link_id": 3,
"nodes": [
{
"id": 1,
"type": "LoadImage",
"pos": [
-640,
-120
],
"size": [
210,
58
],
"flags": {},
"order": 0,
"mode": 0,
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
1
],
"slot_index": 0
},
{
"name": "MASK",
"type": "MASK",
"links": null,
"slot_index": 1
}
],
"properties": {
"Node name for S&R": "Load Target"
},
"widgets_values": [
"grayscale.png",
false
]
},
{
"id": 2,
"type": "LoadImage",
"pos": [
-640,
80
],
"size": [
210,
58
],
"flags": {},
"order": 1,
"mode": 0,
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
2
],
"slot_index": 0
},
{
"name": "MASK",
"type": "MASK",
"links": null,
"slot_index": 1
}
],
"properties": {
"Node name for S&R": "Load Reference"
},
"widgets_values": [
"reference.png",
false
]
},
{
"id": 3,
"type": "ColorMNetImage",
"pos": [
-180,
-40
],
"size": [
320,
150
],
"flags": {},
"order": 2,
"mode": 0,
"inputs": [
{
"name": "image",
"type": "IMAGE",
"link": 1
},
{
"name": "reference_image",
"type": "IMAGE",
"link": 2
}
],
"outputs": [
{
"name": "colorized_image",
"type": "IMAGE",
"links": [
3
],
"slot_index": 0
},
{
"name": "performance_report",
"type": "STRING",
"links": null,
"slot_index": 1
}
],
"properties": {
"Node name for S&R": "ColorMNet Image"
},
"widgets_values": [
768,
432,
true,
false
]
},
{
"id": 4,
"type": "SaveImage",
"pos": [
260,
-40
],
"size": [
210,
58
],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 3
}
],
"properties": {},
"widgets_values": [
"ColorMNet",
"png"
]
}
],
"links": [
[
1,
1,
0,
3,
0
],
[
2,
2,
0,
3,
1
],
[
3,
3,
0,
4,
0
]
],
"groups": [],
"config": {},
"extra": {
"frontendVersion": "1.28.8"
},
"version": 0.4
}
__init__.py
+38 -19
View File
@@ -14,17 +14,30 @@ __version__ = "2.0.0"
__author__ = "ComfyUI Deep Exemplar Contributors"
__license__ = "MIT"
import sys
import subprocess
import importlib
from pathlib import Path
# Add current directory to Python path so submodules can be imported
current_dir = Path(__file__).parent
if str(current_dir) not in sys.path:
sys.path.insert(0, str(current_dir))
print("[ColorMNet] Initializing...")
import sys
import subprocess
import importlib
from pathlib import Path
# Add current directory to Python path so submodules can be imported
current_dir = Path(__file__).parent
if str(current_dir) not in sys.path:
sys.path.insert(0, str(current_dir))
# Import after path is set so relative modules always resolve
from .colormnet.downloader import ensure_dependencies as ensure_git_dependencies
# Backward compatibility: older revisions imported the upstream code as `model.network`
# We vendor that code under `colormnet_model`, so alias it early to avoid ModuleNotFoundError
try:
import colormnet_model
if "model" not in sys.modules:
sys.modules["model"] = colormnet_model
except Exception as alias_error:
print(f"[ColorMNet] WARNING: Couldn't register compatibility alias for 'model': {alias_error}")
print("[ColorMNet] Initializing...")
# Check and install dependencies automatically
def check_and_install_dependencies():
@@ -78,14 +91,20 @@ def check_git_dependencies():
print(f"[ColorMNet] These will be auto-installed on first model run if needed")
return missing
# Auto-install dependencies
try:
check_and_install_dependencies()
missing_git_deps = check_git_dependencies()
except Exception as e:
print(f"[ColorMNet] WARNING: Error checking dependencies: {e}")
missing_git_deps = []
# Auto-install dependencies
try:
check_and_install_dependencies()
try:
# Install git-based deps (py-thin-plate-spline, Pytorch-Correlation-extension) on startup
ensure_git_dependencies()
missing_git_deps = []
except Exception as git_err:
print(f"[ColorMNet] WARNING: Git dependency check failed: {git_err}")
missing_git_deps = ["py-thin-plate-spline", "Pytorch-Correlation-extension"]
except Exception as e:
print(f"[ColorMNet] WARNING: Error checking dependencies: {e}")
missing_git_deps = []
# Import nodes
try:
colormnet/inference.py
+11 -32
View File
@@ -73,7 +73,7 @@ class ColorMNetInference:
target_size: Optional target (H, W)
Returns:
LAB frame [3, H, W] on device
Normalized LAB frame [3, H, W] on device (L ~ [-1,1], ab ~ [-1,1])
"""
# Convert to LAB
lab = rgb_to_lab(frame, normalize=True)
@@ -171,21 +171,11 @@ class ColorMNetInference:
# Create inference processor
self.processor = InferenceCore(model, config=self.config.to_dict())
# Prepare reference LAB
# Prepare reference LAB (normalized to training stats)
ref_lab = self._prepare_frame(reference, (proc_height, proc_width))
ref_L = ref_lab[0:1].unsqueeze(0) # [1, 1, H, W]
ref_ab_raw = ref_lab[1:3].unsqueeze(0) # [1, 2, H, W] - raw ab channels
# Normalize ab channels for the model (model expects ab in [-1, 1] range)
# ab channels are typically in [-128, 128] range, normalize by 110
ref_ab = ref_ab_raw / 110.0 # [1, 2, H, W] - normalized ab channels
# Create RGB version of reference for feature extraction
ref_rgb = lab_to_rgb(
ref_lab.permute(1, 2, 0), # CHW -> HWC
denormalize=True
).permute(2, 0, 1) # HWC -> CHW -> [C, H, W]
ref_rgb = self.device_manager.to_device(ref_rgb)
ref_ab = ref_lab[1:3].unsqueeze(0) # [1, 2, H, W] normalized
ref_lll = ref_L[0].repeat(3, 1, 1) # [3, H, W] normalized L replicated
# Process frames
colorized_frames = []
@@ -198,26 +188,19 @@ class ColorMNetInference:
# Prepare frame
frame_lab = self._prepare_frame(frame, (proc_height, proc_width))
frame_L = frame_lab[0:1].unsqueeze(0) # [1, 1, H, W]
# Create RGB for feature extraction
frame_rgb = lab_to_rgb(
frame_lab.permute(1, 2, 0), # CHW -> HWC
denormalize=True
).permute(2, 0, 1) # HWC -> CHW -> [C, H, W] (no batch dimension)
frame_rgb = self.device_manager.to_device(frame_rgb)
frame_lll = frame_L[0].repeat(3, 1, 1) # [3, H, W] normalized L replicated
# First frame: initialize with reference
if frame_idx == 0:
# Pass the full RGB reference image for feature extraction
msk_lll = ref_rgb # [3, H, W] - full color reference
# Pass only normalized L replicated (matches training pipeline)
msk_lll = ref_lll # [3, H, W] normalized L
msk_ab = ref_ab.squeeze(0) # [1, 2, H, W] -> [2, H, W]
labels = [1, 2] # Two color channels
self.processor.set_all_labels(labels)
# Process first frame with reference exemplar
prob = self.processor.step_AnyExemplar(
frame_rgb,
frame_lll,
msk_lll=msk_lll,
msk_ab=msk_ab,
valid_labels=labels,
@@ -227,7 +210,7 @@ class ColorMNetInference:
else:
# Subsequent frames: use memory
prob = self.processor.step_AnyExemplar(
frame_rgb,
frame_lll,
msk_lll=None,
msk_ab=None,
valid_labels=None,
@@ -236,24 +219,20 @@ class ColorMNetInference:
)
# Prob is [2, H, W] - predicted ab channels in normalized range [-1, 1]
# Scale from [-1, 1] to LAB ab range (using factor of 110, common in colorization)
prob_scaled = prob * 110.0
# Debug: Check if prob contains color information
if frame_idx == 0:
self.logger.info(f"First frame - prob (normalized) shape: {prob.shape}")
self.logger.info(f"First frame - prob (normalized) min: {prob.min():.3f}, max: {prob.max():.3f}, mean: {prob.mean():.3f}")
self.logger.info(f"First frame - prob_scaled (denormalized) min: {prob_scaled.min():.3f}, max: {prob_scaled.max():.3f}, mean: {prob_scaled.mean():.3f}")
self.logger.info(f"Reference ab (normalized) - min: {ref_ab.min():.3f}, max: {ref_ab.max():.3f}, mean: {ref_ab.mean():.3f}")
self.logger.info(f"Reference ab (raw) - min: {ref_ab_raw.min():.3f}, max: {ref_ab_raw.max():.3f}, mean: {ref_ab_raw.mean():.3f}")
# Combine with L channel
colorized_lab = torch.cat([frame_L[0], prob_scaled], dim=0) # [3, H, W]
colorized_lab = torch.cat([frame_L[0], prob], dim=0) # normalized LAB [3, H, W]
# Convert back to RGB
colorized_rgb = lab_to_rgb(
colorized_lab.permute(1, 2, 0), # CHW -> HWC
denormalize=True
normalized=True
)
# Resize to original resolution if needed
colormnet/model.py
+6 -1
View File
@@ -10,7 +10,12 @@ from core.device import DeviceManager
from .config import ColorMNetConfig
# Import ColorMNet network from embedded source
from colormnet_model.network import ColorMNet
# Prefer the vendored package name but keep a fallback for older paths.
try:
from colormnet_model.network import ColorMNet
except ImportError:
# Older revisions referenced `model.network`; alias is registered in __init__, but double-protect here.
from model.network import ColorMNet
class ColorMNetModel:
core/transforms.py
+14 -7
View File
@@ -15,7 +15,8 @@ def rgb_to_lab(
Args:
rgb: RGB tensor in range [0, 1], shape [..., H, W, 3]
normalize: Whether to normalize L to [-50, 50] and ab to [-128, 128]
normalize: Whether to normalize using training-time stats:
L -> (L-50)/50, ab -> ab/110 (range ~[-1, 1])
Returns:
LAB tensor, shape [..., H, W, 3]
@@ -60,9 +61,9 @@ def rgb_to_lab(
lab_tensor = torch.from_numpy(lab_np).to(rgb.device)
if normalize:
# L: [0, 100] -> [-50, 50]
# ab: [-128, 128] (already in this range from skimage)
lab_tensor[..., 0] = lab_tensor[..., 0] - 50.0
# Training-time normalization: L/50 and ab/110 to ~[-1, 1]
lab_tensor[..., 0] = (lab_tensor[..., 0] - 50.0) / 50.0
lab_tensor[..., 1:] = lab_tensor[..., 1:] / 110.0
return lab_tensor
@@ -70,20 +71,26 @@ def rgb_to_lab(
def lab_to_rgb(
lab: torch.Tensor,
denormalize: bool = True,
normalized: bool = False,
) -> torch.Tensor:
"""Convert LAB to RGB color space.
Args:
lab: LAB tensor, shape [..., H, W, 3]
denormalize: Whether L is in [-50, 50] and needs to be denormalized to [0, 100]
denormalize: Legacy flag; kept for compatibility (adds 50 to L).
normalized: If True, assumes training-time normalization (L: (x-50)/50, ab: /110).
Returns:
RGB tensor in range [0, 1], shape [..., H, W, 3]
"""
lab = lab.clone()
if denormalize:
# L: [-50, 50] -> [0, 100]
if normalized:
# Inverse of training-time normalization
lab[..., 0] = lab[..., 0] * 50.0 + 50.0
lab[..., 1:] = lab[..., 1:] * 110.0
elif denormalize:
# Legacy path: L was in [-50, 50]
lab[..., 0] = lab[..., 0] + 50.0
# Convert to numpy