Skip to content

jonstreeter/ComfyUI-Reference-Based-Video-Colorization

Folders and files

NameName
Last commit message
Last commit date

Latest commit

Β 

History

57 Commits
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 
Β 

Repository files navigation

ComfyUI Reference-Based Video Colorization

Reference-Based Video Colorization Workflow

A comprehensive ComfyUI implementation featuring two state-of-the-art reference-based video colorization methods:

  • 🎨 ColorMNet (2024) - Modern memory-based approach with DINOv2 features
  • 🎬 Deep Exemplar (2019) - Classic CVPR method with temporal propagation

Transform black & white videos and images into vibrant color using reference images!

🎬 Demo

See ColorMNet in action colorizing classic black & white footage:

Example: Colorized using ColorMNet with color matching post-processing. The model successfully transfers colors from a reference image while maintaining temporal consistency.


✨ Features

ColorMNet Nodes (New)

  • ColorMNet Video Colorization - Memory-based temporal coherent colorization
  • ColorMNet Image Colorization - Single image colorization
  • DINOv2-based feature extraction for superior quality
  • Multiple memory modes (balanced, low memory, high quality)
  • FP16 support for faster processing
  • torch.compile optimization for 15-25% speedup (PyTorch 2.0+)
  • Performance reports with timing and FPS metrics

Deep Exemplar Nodes (Original)

  • Deep Exemplar Video Colorization - Frame propagation for temporal consistency
  • Deep Exemplar Image Colorization - Classic exemplar-based method
  • WLS (Weighted Least Squares) filtering for smoother results
  • Configurable lambda and sigma parameters
  • torch.compile optimization for 15-25% speedup (PyTorch 2.0+)
  • SageAttention for 20-30% faster attention operations
  • Performance reports for benchmarking

Common Features

  • βœ… Automatic model download - No manual setup required!
  • βœ… Progress bars - Real-time processing feedback in ComfyUI
  • βœ… Performance metrics - Compare speed and quality between methods
  • βœ… Flexible resolution - Process at any resolution
  • βœ… ComfyUI native - Fully integrated workflow support

πŸ“¦ Installation

Method 1: ComfyUI Manager (Recommended)

  1. Install ComfyUI Manager
  2. Open ComfyUI β†’ Manager β†’ Install Custom Nodes
  3. Search for "Deep Exemplar Video Colorization"
  4. Click Install
  5. Restart ComfyUI

Models will download automatically on first use (~700MB total).

Method 2: Manual Installation

cd ComfyUI/custom_nodes/
git clone https://github.com/jonstreeter/ComfyUI-Reference-Based-Video-Colorization.git
cd ComfyUI-Reference-Based-Video-Colorization/
pip install -r requirements.txt

Restart ComfyUI. Models download automatically on first use.


πŸš€ Quick Start

Example Workflow

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
  3. Choose Method - Try both ColorMNet and Deep Exemplar
  4. Compare Results - Use the performance reports to benchmark
  5. Save Output - Export colorized video with VHS Video Combine

Nodes Overview

ColorMNet Nodes

ColorMNet Video (ColorMNet/Video)

  • video_frames - Input grayscale video frames
  • reference_image - Color reference image
  • target_width/height - Output resolution
  • memory_mode - balanced | low_memory | high_quality
  • use_fp16 - Enable FP16 for speed (default: True)
  • use_torch_compile - Enable PyTorch 2.0+ compilation (default: False)
  • Outputs: colorized_frames, performance_report

ColorMNet Image (ColorMNet/Image)

  • image - Input grayscale image
  • reference_image - Color reference image
  • target_width/height - Output resolution
  • use_fp16 - Enable FP16 (default: True)
  • use_torch_compile - Enable PyTorch 2.0+ compilation (default: False)
  • Outputs: colorized_image, performance_report

Deep Exemplar Nodes

Deep Exemplar Video (DeepExemplar/Video)

  • video_frames - Input grayscale video frames
  • reference_image - Color reference image
  • target_width/height - Output resolution
  • frame_propagate - Use temporal consistency (default: True)
  • use_half_resolution - Process at half resolution (default: True)
  • use_torch_compile - Enable PyTorch 2.0+ compilation (default: False)
  • use_sage_attention - Enable SageAttention optimization (default: False)
  • Outputs: colorized_frames, performance_report

Deep Exemplar Image (DeepExemplar/Image)

  • image_to_colorize - Input grayscale image
  • reference_image - Color reference image
  • target_width/height - Output resolution
  • use_torch_compile - Enable PyTorch 2.0+ compilation (default: False)
  • use_sage_attention - Enable SageAttention optimization (default: False)
  • Outputs: colorized_image, performance_report

πŸ“Š Performance Reports

All nodes output optional performance reports with timing metrics:

Example ColorMNet Video Report

ColorMNet Video Colorization Report
==================================================
Frames Processed: 240
Resolution: 768x432
Total Time: 45.23 seconds
Average FPS: 5.31
Time per Frame: 0.188 seconds
Memory Mode: balanced
FP16 Enabled: True
Torch Compile: False
==================================================

Example Deep Exemplar Video Report

Deep Exemplar Video Colorization Report
==================================================
Frames Processed: 240
Resolution: 768x432
Total Time: 52.34 seconds
Average FPS: 4.59
Time per Frame: 0.218 seconds
Frame Propagation: Enabled
Half Resolution: Enabled
WLS Filter: Enabled
Lambda: 500.0
Sigma Color: 4.0
Torch Compile: False
SageAttention: False
==================================================

Connect the performance_report output to a text display node or save to file for benchmarking!


πŸ’‘ Tips & Best Practices

Reference Image Selection

  • Choose references semantically similar to your content
  • Match the color palette you want to achieve
  • Higher quality references = better results
  • Try multiple references to find the best match

Resolution Settings

  • ColorMNet: Processes at full target resolution
  • Deep Exemplar: Internally uses half resolution by default
  • Start with 768x432 for good speed/quality balance
  • Increase for final high-quality renders

Memory Management

ColorMNet Memory Modes:

  • balanced - Good quality, moderate memory (recommended)
  • low_memory - Reduced memory usage, slight quality trade-off
  • high_quality - Best quality, higher memory requirements

Deep Exemplar:

  • Enable use_half_resolution to reduce memory
  • Disable frame_propagate for independent frame processing
  • Process shorter clips if encountering OOM errors

Quality vs Speed

For Best Quality:

  • ColorMNet: memory_mode=high_quality, use_fp16=False, use_torch_compile=False
  • Deep Exemplar: use_half_resolution=False, use_torch_compile=False, use_sage_attention=False

For Best Speed:

  • ColorMNet: memory_mode=low_memory, use_fp16=True, use_torch_compile=True
  • Deep Exemplar: use_half_resolution=True, use_torch_compile=True, use_sage_attention=True

Optimization Notes:

  • torch.compile requires PyTorch 2.0+ and provides 15-25% speedup after initial warmup
  • use_sage_attention (Deep Exemplar only) provides 20-30% faster attention with sageattention installed
  • Both optimizations maintain identical quality to non-optimized versions

πŸ”§ Advanced Configuration

Custom Model Paths

Models are automatically downloaded to:

custom_nodes/ComfyUI-Reference-Based-Video-Colorization/
β”œβ”€β”€ checkpoints/
β”‚   β”œβ”€β”€ DINOv2FeatureV6_LocalAtten_s2_154000.pth  # ColorMNet (~500MB)
β”‚   └── video_moredata_l1/
β”‚       β”œβ”€β”€ nonlocal_net_iter_76000.pth            # Deep Exemplar
β”‚       └── colornet_iter_76000.pth                # Deep Exemplar
└── data/
    β”œβ”€β”€ vgg19_conv.pth                              # Shared VGG19
    └── vgg19_gray.pth                              # Deep Exemplar VGG

Optimizations

torch.compile (PyTorch 2.0+): Provides 15-25% speedup through graph compilation and optimization.

  • Available for all 4 nodes (ColorMNet + Deep Exemplar)
  • Enable via use_torch_compile=True parameter
  • First run includes warmup compilation (slower), subsequent runs benefit from speedup
  • No additional installation required if using PyTorch 2.0+

SageAttention (Deep Exemplar only): INT8-quantized attention for 20-30% faster attention operations.

Installation:

pip install sageattention

Requirements:

  • CUDA-capable GPU
  • PyTorch with CUDA support
  • Enable via use_sage_attention=True parameter
  • Automatically falls back to standard attention if unavailable

CUDA Correlation Sampler (Optional for ColorMNet): ColorMNet can use optimized CUDA correlation operations if available.

Requirements:

  • CUDA Toolkit installed
  • Visual Studio Build Tools (Windows)
  • Will be attempted automatically on first run

OpenCV Contrib (Required for WLS):

pip install opencv-contrib-python

πŸ“š Documentation


πŸŽ“ Citation

If you use these methods in your research, please cite the original papers:

ColorMNet

@article{yang2024colormnet,
  title={ColorMNet: A Memory-based Deep Spatial-Temporal Feature Propagation Network for Video Colorization},
  author={Yang, Yixin and Zhou, Xiaoyu and Liu, Chao and Chen, Junchi and Wang, Zhiwen},
  journal={arXiv preprint arXiv:2404.06251},
  year={2024}
}

Deep Exemplar

@inproceedings{zhang2019deep,
  title={Deep exemplar-based video colorization},
  author={Zhang, Bo and He, Mingming and Liao, Jing and Sander, Pedro V and Yuan, Lu and Bermak, Amine and Chen, Dong},
  booktitle={Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition},
  pages={8052--8061},
  year={2019}
}

πŸ”— Related Projects


πŸ“ License

This ComfyUI implementation is licensed under the MIT License.

Note on Model Licenses:

  • ColorMNet model: CC BY-NC-SA 4.0 (Non-commercial use only)
  • Deep Exemplar model: Subject to original project license

πŸ™ Acknowledgments


πŸ› Issues & Support

Found a bug or have a feature request? Please open an issue on GitHub!

For general questions:

  1. Check the documentation
  2. Review existing issues
  3. Open a new issue with details

🎯 Roadmap

  • Batch processing optimization
  • Memory-efficient streaming for long videos
  • Additional colorization methods
  • Color transfer utilities
  • Integration with other ComfyUI nodes

Star ⭐ this repo if you find it useful!

About

Dual implementation of reference-based video colorization: ColorMNet (2024) + Deep Exemplar (2019) for ComfyUI

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Contributors 3

  •  
  •  
  •  

Languages