diff --git a/PERFORMANCE.md b/PERFORMANCE.md index 951adb5..2f4aa1a 100644 --- a/PERFORMANCE.md +++ b/PERFORMANCE.md @@ -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. diff --git a/README.md b/README.md index a3e9c48..26870be 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/Workflows/ColorMNet_Image_Workflow.json b/Workflows/ColorMNet_Image_Workflow.json new file mode 100644 index 0000000..aee2341 --- /dev/null +++ b/Workflows/ColorMNet_Image_Workflow.json @@ -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 +} diff --git a/__init__.py b/__init__.py index b6ab899..f0de4ee 100644 --- a/__init__.py +++ b/__init__.py @@ -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: diff --git a/colormnet/inference.py b/colormnet/inference.py index 0044d71..90292e3 100644 --- a/colormnet/inference.py +++ b/colormnet/inference.py @@ -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 diff --git a/colormnet/model.py b/colormnet/model.py index 8484228..e2bb796 100644 --- a/colormnet/model.py +++ b/colormnet/model.py @@ -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: diff --git a/core/transforms.py b/core/transforms.py index 841993f..6afd591 100644 --- a/core/transforms.py +++ b/core/transforms.py @@ -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