diff --git a/README.md b/README.md index 6a325d8..38b18d7 100644 --- a/README.md +++ b/README.md @@ -1,129 +1,151 @@ -# Deep Exemplar-based Video Colorization (Pytorch Implementation) +# ComfyUI Deep Exemplar-based Colorization Nodes - +ComfyUI Node Teaser -### [Paper](https://arxiv.org/abs/1906.09909) | [Pretrained Model](https://facevc.blob.core.windows.net/zhanbo/old_photo/colorization_checkpoint.zip) | [Youtube video](https://youtu.be/HXWR5h5vVYI) :fire: | [Colab demo](https://colab.research.google.com/drive/1Y1XTlTdUG-2LzrH1Vnr_osg9BQavfYsz?usp=sharing) +This repository provides custom nodes for [ComfyUI](https://github.com/comfyanonymous/ComfyUI) that implement the **Deep Exemplar-based Video Colorization** method (CVPR 2019). It allows you to colorize grayscale images or video frames using a color reference image directly within the ComfyUI workflow. + +Based on the original paper and Pytorch implementation: +**[Paper](https://arxiv.org/abs/1906.09909) | [Original Project](https://github.com/zhangmozhe/Deep-Exemplar-based-Video-Colorization) | [Original Pretrained Model](https://github.com/zhangmozhe/Deep-Exemplar-based-Video-Colorization/releases/download/v1.0/colorization_checkpoint.zip)** **Deep Exemplar-based Video Colorization, CVPR2019** +[Bo Zhang](https://www.microsoft.com/en-us/research/people/zhanbo/)1,3, [Mingming He](http://mingminghe.com/)1,5, [Jing Liao](https://liaojing.github.io/html/)2, [Pedro V. Sander](https://www.cse.ust.hk/~psander/)1, [Lu Yuan](https://www.microsoft.com/en-us/research/people/luyuan/)4, [Amine Bermak](https://eebermak.home.ece.ust.hk/)1, [Dong Chen](https://www.microsoft.com/en-us/research/people/doch/)3
+1HKUST, 2CityU HK, 3Microsoft Research Asia, 4Microsoft Cloud&AI, 5USC ICT -[Bo Zhang](https://www.microsoft.com/en-us/research/people/zhanbo/)1,3, -[Mingming He](http://mingminghe.com/)1,5, -[Jing Liao](https://liaojing.github.io/html/)2, -[Pedro V. Sander](https://www.cse.ust.hk/~psander/)1, -[Lu Yuan](https://www.microsoft.com/en-us/research/people/luyuan/)4, -[Amine Bermak](https://eebermak.home.ece.ust.hk/)1, -[Dong Chen](https://www.microsoft.com/en-us/research/people/doch/)3
-1Hong Kong University of Science and Technology,2City University of Hong Kong, -3Microsoft Research Asia, 4Microsoft Cloud&AI, 5USC Institute for Creative Technologies +## Features - - -## Prerequisites - -- Python 3.6+ -- Nvidia GPU + CUDA, CuDNN +* **DeepExColorImageNode:** Colorizes a single grayscale image using a color reference image. +* **DeepExColorVideoNode:** Colorizes a batch of grayscale video frames using a color reference image, with optional frame-to-frame propagation. +* Integration with ComfyUI workflows. +* Automatic download of pretrained models via ComfyUI Manager or included `install.py`. +* Adjustable target resolution (internally processed at half-resolution). +* Optional WLS (Fast Global Smoother) filter post-processing (requires `opencv-contrib-python`). ## Installation -First use the following commands to prepare the environment: +### Using ComfyUI Manager (Recommended) -```bash -conda create -n ColorVid python=3.6 -source activate ColorVid -pip install -r requirements.txt -``` +1. Install [ComfyUI Manager](https://github.com/ltdrdata/ComfyUI-Manager) if you haven't already. +2. Open ComfyUI and navigate to the "Manager" menu. +3. Click "Install Custom Nodes". +4. Search for "DeepExemplarColorization". +5. Click "Install". +6. **Wait for installation to complete.** The manager will run the `install.py` script to download the required model checkpoints (~500MB) into the `ComfyUI/custom_nodes/ComfyUI-Deep-Exemplar-based-Video-Colorization/` directory. +7. **Install Dependencies:** Ensure `opencv-contrib-python` is installed in your ComfyUI environment if you want to use the WLS filter. The Manager will attempt to install packages listed in `requirements.txt`. You can verify/install manually if needed: + ```bash + # In ComfyUI's Python environment: + # pip install opencv-contrib-python + ``` +8. **Restart ComfyUI.** -Then, download the pretrained models from [this link](https://github.com/zhangmozhe/Deep-Exemplar-based-Video-Colorization/releases/download/v1.0/colorization_checkpoint.zip), unzip the file and place the files into the corresponding folders: +### Manual Installation -- `video_moredata_l1` under the `checkpoints` folder -- `vgg19_conv.pth` and `vgg19_gray.pth` under the `data` folder +1. Navigate to your ComfyUI `custom_nodes` directory: + ```bash + cd ComfyUI/custom_nodes/ + ``` +2. Clone this repository: + ```bash + git clone https://github.com/jonstreeter/ComfyUI-Deep-Exemplar-based-Video-Colorization # <-- Replace with your repo URL! + ``` +3. Navigate into the cloned directory: + ```bash + cd ComfyUI-Deep-Exemplar-based-Video-Colorization/ + ``` +4. **Download Model Checkpoints:** Run the installation script using the Python environment ComfyUI uses: + ```bash + # Example using ComfyUI's portable python: + # ../../python_embeded/python.exe install.py + # Or if using a venv: + # path/to/your/venv/Scripts/python.exe install.py + ``` + This will download `colorization_checkpoint.zip` and extract the necessary files into the `checkpoints/` and `data/` subdirectories within this custom node folder. +5. **Install Dependencies:** Install the required Python packages using the same Python environment: + ```bash + # Example using ComfyUI's portable python: + # ../../python_embeded/python.exe -m pip install -r requirements.txt + # Or if using a venv: + # path/to/your/venv/Scripts/python.exe -m pip install -r requirements.txt + ``` + *Note:* This includes `opencv-contrib-python`, which is required for the WLS filter. +6. **Restart ComfyUI.** -## Data Preparation +## Usage -In order to colorize your own video, it requires to extract the video frames, and provide a reference image as an example. +After installation and restarting ComfyUI, you can find the nodes under the "DeepExemplar" category when adding nodes (Right Click -> Add Node -> DeepExemplar). -- Place your video frames into one folder, _e.g._, `./sample_videos/v32_180` -- Place your reference images into another folder, _e.g._, `./sample_videos/v32` +### Nodes -If you want to _automatically_ retrieve color images, you can try the retrieval algorithm from [this link](https://github.com/hmmlillian/Gray-Image-Retrieval) which will retrieve similar images from the ImageNet dataset. Or you can try [this link](https://github.com/pochih/CBIR) on your own image database. +* **`DeepExColorImageNode`**: + * Inputs: + * `image_to_colorize`: The single grayscale image (standard ComfyUI IMAGE format). + * `reference_image`: The single color reference image (standard ComfyUI IMAGE format). + * `target_width`/`target_height`: Desired output resolution (internal processing is at half this size; will be adjusted to be multiples of 32, min 64x64). + * `wls_filter_on`: Enable WLS post-processing filter (requires `opencv-contrib-python`). + * `lambda_value`: WLS filter parameter (controls smoothing amount). + * `sigma_color`: WLS filter parameter (controls color sensitivity). + * Output: + * `IMAGE`: The colorized image. -## Test +* **`DeepExColorVideoNode`**: + * Inputs: + * `video_frames`: A batch of grayscale images (standard ComfyUI IMAGE format, BxHxWxC tensor). Use ComfyUI's `Load Video` or `Load Image Sequence` nodes upstream. + * `reference_image`: The single color reference image. + * `frame_propagate`: If true, uses the previous frame's colorization to guide the current frame (helps temporal consistency). + * `use_half_resolution`: If true, performs core computations at half the target resolution (faster, matches original paper). *Note: This was the default and only behavior in the reverted code, but the parameter exists.* + * `target_width`/`target_height`: Desired output resolution. + * `wls_filter_on`: Enable WLS post-processing filter (requires `opencv-contrib-python`). + * `lambda_value`: WLS filter parameter. + * `sigma_color`: WLS filter parameter. + * Output: + * `IMAGE`: A batch containing the colorized frames. Use ComfyUI's `Save Animated PNG/WEBP/GIF` or `Video Combine` nodes downstream. -```bash -python test.py --image-size [image-size] \ - --clip_path [path-to-video-frames] \ - --ref_path [path-to-reference] \ - --output_path [path-to-output] -``` +### Example Workflows -We provide several sample video clips with corresponding references. For example, one can colorize one sample legacy video using: + +*Screenshot showing: Load Image (grayscale) -> Load Image (color ref) -> DeepExColorImageNode -> Preview Image* -```bash -python test.py --clip_path ./sample_videos/clips/v32 \ - --ref_path ./sample_videos/ref/v32 \ - --output_path ./sample_videos/output -``` + +*Screenshot showing: Load Video -> DeepExColorVideoNode (with color ref loaded) -> Save Animated WEBP* -Note that we use 216\*384 images for training, which has aspect ratio of 1:2. During inference, we scale the input to this size and then rescale the output back to the original size. +### Tips -## Train +* **Reference Image:** The quality and color palette of the reference image significantly impact the result. Choose a reference that is semantically similar and has the desired color style. See the original project's README for links to image retrieval algorithms if needed. +* **Resolution:** The nodes process images internally based on half the `target_width` and `target_height`. Final output is scaled up to the target resolution. +* **Frame Propagation:** Enabling `frame_propagate` in the video node generally improves temporal stability but might slightly change the color compared to processing each frame independently. +* **WLS Filter:** The WLS (Fast Global Smoother) filter helps smooth colors while respecting edges. It requires `opencv-contrib-python`. If unavailable or disabled, colors might appear less smooth or blocky. +* **Memory:** Processing high-resolution videos can be memory-intensive. Adjust `target_width`/`target_height` or process shorter clips if you encounter out-of-memory errors. -We also provide training code for reference. The training can be started by running: +## Results Comparison -```bash -python train.py --data_root [root of video samples] \ - --data_root_imagenet [root of image samples] \ - --gpu_ids [gpu ids] \ -``` + + +Comparison Example 1 -We do not provide the full video dataset due to the copyright issue. For image samples, we retrieve semantically similar images from ImageNet using [this repository](https://github.com/hmmlillian/Gray-Image-Retrieval). Still, one can refer to our code to understand the detailed procedure of augmenting the image dataset to mimic the video frames. +## More Examples -## Comparison with State-of-the-Arts + +Result Example 2 +Result Example 3 -
- -
+ +Check out this [Sample Video Clip]() showing the node in action! -## More results +## Citation (Original Paper) -
- -
- -
- -
- -:star: Please check our [Youtube demo](https://youtu.be/HXWR5h5vVYI) for results of video colorization. - -## Citation - -If you use this code for your research, please cite our paper. - -``` -@inproceedings{zhang2019deep, +If you use the underlying method in your research, please cite the original paper: @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} } -``` -## Old Photo Restoration :fire: -If you are also interested in restoring the artifacts in the legacy photo, please check our recent work, [bringing old photo back to life](https://github.com/microsoft/Bringing-Old-Photos-Back-to-Life). +## Related Projects -``` -@inproceedings{wan2020bringing, -title={Bringing Old Photos Back to Life}, -author={Wan, Ziyu and Zhang, Bo and Chen, Dongdong and Zhang, Pan and Chen, Dong and Liao, Jing and Wen, Fang}, -booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition}, -pages={2747--2757}, -year={2020} -} -``` +* Original Implementation: [https://github.com/zhangmozhe/Deep-Exemplar-based-Video-Colorization](https://github.com/zhangmozhe/Deep-Exemplar-based-Video-Colorization) +* Bringing Old Photos Back to Life: [https://github.com/microsoft/Bringing-Old-Photos-Back-to-Life](https://github.com/microsoft/Bringing-Old-Photos-Back-to-Life) ## License -This project is licensed under the MIT license. +This ComfyUI node implementation is licensed under the MIT license. The underlying model and original code are subject to their respective licenses.