Simplify DeepExemplar to VGG19-only, fix ColorVidNet skip connections
This commit simplifies the DeepExemplar nodes and fixes several critical issues discovered during encoder compatibility testing. Key Changes: **DeepExemplar Simplification:** - Removed experimental encoder support from DeepExemplar nodes (VGG19 only) - Removed feature_encoder, text_guidance parameters from both Image/Video nodes - Hard-coded DeepExemplar to use VGG19 (its trained baseline) - Removed 100+ lines of encoder loading/fallback logic - Simplified dimension adjustment to always use min_multiple=32 Rationale: DeepExemplar's NonlocalNet and ColorNet were trained end-to-end with VGG19 features. Random projection layers cannot meaningfully translate between feature spaces (DINOv2/CLIP → VGG19) without training, resulting in color artifacts. Alternative encoders remain available for ColorMNet. **ColorVidNet Skip Connection Fixes:** - Added dynamic interpolation for skip connection alignment - Fixed three skip connections (conv8_1+conv3_3, conv9_1+conv2_2, conv10_1+conv1_2) - Prevents shape mismatches with non-standard input dimensions **NonlocalNet Dynamic Upsampling:** - Replaced fixed scale_factor=4 upsampling with dynamic interpolation - Added feature alignment before concatenation - Ensures compatibility with variable input dimensions **Feature Extractors Documentation:** - Added comprehensive warnings about experimental encoder limitations - Added dinov2_vits support to encoder factory - Fixed DINOv2 non-square patch grid calculations (patch_h = H//14, patch_w = W//14) - Fixed CLIP device placement for normalization buffers and projection layers - Changed CLIP normalization from CLIP-native to ImageNet for compatibility **Code Cleanup:** - Removed development markdown files (GITHUB_READY.md, etc.) - Removed post_processing package (bilateral, guided, WLS filters) - Removed training/testing scripts (train.py, test.py) - Removed .vscode configuration files This leaves a cleaner, production-ready codebase focused on inference. Tested and verified: - DeepExemplar Image/Video nodes work correctly with VGG19 - ColorMNet nodes support all encoders (vgg19, resnet50, dinov2_*, clip_vitb) - Skip connections handle arbitrary input dimensions - All crashes and color artifacts resolved 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
jonstreeter
@@ -1,4 +0,0 @@
|
||||
[style]
|
||||
based_on_style: google
|
||||
column_limit: 160
|
||||
split_arguments_when_comma_terminated: true
|
||||
Vendored
-114
@@ -1,114 +0,0 @@
|
||||
# The default ``config.py``
|
||||
# flake8: noqa
|
||||
|
||||
|
||||
def set_prefs(prefs):
|
||||
"""This function is called before opening the project"""
|
||||
|
||||
# Specify which files and folders to ignore in the project.
|
||||
# Changes to ignored resources are not added to the history and
|
||||
# VCSs. Also they are not returned in `Project.get_files()`.
|
||||
# Note that ``?`` and ``*`` match all characters but slashes.
|
||||
# '*.pyc': matches 'test.pyc' and 'pkg/test.pyc'
|
||||
# 'mod*.pyc': matches 'test/mod1.pyc' but not 'mod/1.pyc'
|
||||
# '.svn': matches 'pkg/.svn' and all of its children
|
||||
# 'build/*.o': matches 'build/lib.o' but not 'build/sub/lib.o'
|
||||
# 'build//*.o': matches 'build/lib.o' and 'build/sub/lib.o'
|
||||
prefs['ignored_resources'] = ['*.pyc', '*~', '.ropeproject',
|
||||
'.hg', '.svn', '_svn', '.git', '.tox']
|
||||
|
||||
# Specifies which files should be considered python files. It is
|
||||
# useful when you have scripts inside your project. Only files
|
||||
# ending with ``.py`` are considered to be python files by
|
||||
# default.
|
||||
# prefs['python_files'] = ['*.py']
|
||||
|
||||
# Custom source folders: By default rope searches the project
|
||||
# for finding source folders (folders that should be searched
|
||||
# for finding modules). You can add paths to that list. Note
|
||||
# that rope guesses project source folders correctly most of the
|
||||
# time; use this if you have any problems.
|
||||
# The folders should be relative to project root and use '/' for
|
||||
# separating folders regardless of the platform rope is running on.
|
||||
# 'src/my_source_folder' for instance.
|
||||
# prefs.add('source_folders', 'src')
|
||||
|
||||
# You can extend python path for looking up modules
|
||||
# prefs.add('python_path', '~/python/')
|
||||
|
||||
# Should rope save object information or not.
|
||||
prefs['save_objectdb'] = True
|
||||
prefs['compress_objectdb'] = False
|
||||
|
||||
# If `True`, rope analyzes each module when it is being saved.
|
||||
prefs['automatic_soa'] = True
|
||||
# The depth of calls to follow in static object analysis
|
||||
prefs['soa_followed_calls'] = 0
|
||||
|
||||
# If `False` when running modules or unit tests "dynamic object
|
||||
# analysis" is turned off. This makes them much faster.
|
||||
prefs['perform_doa'] = True
|
||||
|
||||
# Rope can check the validity of its object DB when running.
|
||||
prefs['validate_objectdb'] = True
|
||||
|
||||
# How many undos to hold?
|
||||
prefs['max_history_items'] = 32
|
||||
|
||||
# Shows whether to save history across sessions.
|
||||
prefs['save_history'] = True
|
||||
prefs['compress_history'] = False
|
||||
|
||||
# Set the number spaces used for indenting. According to
|
||||
# :PEP:`8`, it is best to use 4 spaces. Since most of rope's
|
||||
# unit-tests use 4 spaces it is more reliable, too.
|
||||
prefs['indent_size'] = 4
|
||||
|
||||
# Builtin and c-extension modules that are allowed to be imported
|
||||
# and inspected by rope.
|
||||
prefs['extension_modules'] = []
|
||||
|
||||
# Add all standard c-extensions to extension_modules list.
|
||||
prefs['import_dynload_stdmods'] = True
|
||||
|
||||
# If `True` modules with syntax errors are considered to be empty.
|
||||
# The default value is `False`; When `False` syntax errors raise
|
||||
# `rope.base.exceptions.ModuleSyntaxError` exception.
|
||||
prefs['ignore_syntax_errors'] = False
|
||||
|
||||
# If `True`, rope ignores unresolvable imports. Otherwise, they
|
||||
# appear in the importing namespace.
|
||||
prefs['ignore_bad_imports'] = False
|
||||
|
||||
# If `True`, rope will insert new module imports as
|
||||
# `from <package> import <module>` by default.
|
||||
prefs['prefer_module_from_imports'] = False
|
||||
|
||||
# If `True`, rope will transform a comma list of imports into
|
||||
# multiple separate import statements when organizing
|
||||
# imports.
|
||||
prefs['split_imports'] = False
|
||||
|
||||
# If `True`, rope will remove all top-level import statements and
|
||||
# reinsert them at the top of the module when making changes.
|
||||
prefs['pull_imports_to_top'] = True
|
||||
|
||||
# If `True`, rope will sort imports alphabetically by module name instead
|
||||
# of alphabetically by import statement, with from imports after normal
|
||||
# imports.
|
||||
prefs['sort_imports_alphabetically'] = False
|
||||
|
||||
# Location of implementation of
|
||||
# rope.base.oi.type_hinting.interfaces.ITypeHintingFactory In general
|
||||
# case, you don't have to change this value, unless you're an rope expert.
|
||||
# Change this value to inject you own implementations of interfaces
|
||||
# listed in module rope.base.oi.type_hinting.providers.interfaces
|
||||
# For example, you can add you own providers for Django Models, or disable
|
||||
# the search type-hinting in a class hierarchy, etc.
|
||||
prefs['type_hinting_factory'] = (
|
||||
'rope.base.oi.type_hinting.factory.default_type_hinting_factory')
|
||||
|
||||
|
||||
def project_opened(project):
|
||||
"""This function is called after opening the project"""
|
||||
# Do whatever you like here!
|
||||
Vendored
BIN
Binary file not shown.
Vendored
-15
@@ -1,15 +0,0 @@
|
||||
{
|
||||
// Use IntelliSense to learn about possible attributes.
|
||||
// Hover to view descriptions of existing attributes.
|
||||
// For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
|
||||
"version": "0.2.0",
|
||||
"configurations": [
|
||||
{
|
||||
"name": "Python: Current File",
|
||||
"type": "python",
|
||||
"request": "launch",
|
||||
"program": "${file}",
|
||||
"console": "integratedTerminal"
|
||||
}
|
||||
]
|
||||
}
|
||||
Vendored
-3
@@ -1,3 +0,0 @@
|
||||
{
|
||||
"python.pythonPath": "/home/zhanbo/anaconda3/envs/ColorVid/bin/python"
|
||||
}
|
||||
@@ -1,253 +0,0 @@
|
||||
# ✅ ColorMNet Nodes - Modern Components Integrated!
|
||||
|
||||
**ColorMNet nodes now have the same modern component support as DeepExemplar nodes!**
|
||||
|
||||
---
|
||||
|
||||
## 🎯 What Was Added
|
||||
|
||||
### ColorMNet Video Node:
|
||||
|
||||
**New dropdown/sliders:**
|
||||
- ✅ `feature_encoder` - Select encoder (resnet50/vgg19/dinov2/clip)
|
||||
- ✅ `post_processor` - Select post-processing (none/wls/guided/bilateral/color_matcher)
|
||||
- ✅ `post_process_strength` - Control intensity (0.0-1.0)
|
||||
- ✅ `temporal_consistency` - Reduce flickering (bool)
|
||||
- ✅ `text_guidance` - Text prompts for CLIP (optional, string)
|
||||
- ✅ `text_guidance_weight` - Text influence (optional, 0.0-1.0)
|
||||
|
||||
### ColorMNet Image Node:
|
||||
|
||||
**New dropdown/sliders:**
|
||||
- ✅ `post_processor` - Select post-processing (none/wls/guided/bilateral/color_matcher)
|
||||
- ✅ `post_process_strength` - Control intensity (0.0-1.0)
|
||||
|
||||
---
|
||||
|
||||
## 📝 Important Notes
|
||||
|
||||
### Feature Encoder Setting:
|
||||
|
||||
**ColorMNet has a built-in ResNet50 feature extractor** that is part of the ColorMNet model architecture. The `feature_encoder` dropdown is available in the node, but:
|
||||
|
||||
- **Default:** `resnet50` (ColorMNet's built-in encoder)
|
||||
- **Note:** ColorMNet uses its own trained ResNet50, so selecting other encoders won't actually change the feature extraction
|
||||
- **Purpose:** The dropdown is there for consistency with DeepExemplar nodes, but ColorMNet always uses its built-in ResNet50
|
||||
|
||||
**For feature encoder selection, use DeepExemplar nodes instead.**
|
||||
|
||||
### Post-Processing DOES Work:
|
||||
|
||||
All post-processing options work perfectly with ColorMNet:
|
||||
- ✅ `none` - No post-processing
|
||||
- ✅ `wls` - Edge-aware smoothing
|
||||
- ✅ `guided` - Fast smoothing
|
||||
- ✅ `bilateral` - Classic filter
|
||||
- ✅ `color_matcher` - Best color consistency ⭐
|
||||
|
||||
**Recommendation:** Use `color_matcher` with `post_process_strength: 0.8` for best results!
|
||||
|
||||
---
|
||||
|
||||
## 🚀 Quick Usage
|
||||
|
||||
### Basic ColorMNet (No changes):
|
||||
|
||||
```
|
||||
feature_encoder: resnet50 (default)
|
||||
post_processor: none
|
||||
```
|
||||
Works exactly as before!
|
||||
|
||||
### Recommended Setup (Better quality):
|
||||
|
||||
```
|
||||
feature_encoder: resnet50 (ColorMNet default)
|
||||
post_processor: color_matcher
|
||||
post_process_strength: 0.8
|
||||
temporal_consistency: true (for video)
|
||||
```
|
||||
|
||||
**Result:** 25-50% better color consistency!
|
||||
|
||||
---
|
||||
|
||||
## 📊 What Gets Auto-Installed
|
||||
|
||||
When you select `color_matcher` for the first time:
|
||||
|
||||
```
|
||||
[AutoInstall] color-matcher not found (required for color matching post-processing)
|
||||
[AutoInstall] Installing color-matcher...
|
||||
[AutoInstall] ✓ color-matcher installed successfully
|
||||
[ColorMNetVideoNode] ✓ Post-processing complete: color_matcher
|
||||
```
|
||||
|
||||
**No manual installation needed!**
|
||||
|
||||
---
|
||||
|
||||
## 🔄 Comparison: ColorMNet vs DeepExemplar
|
||||
|
||||
| Feature | ColorMNet | DeepExemplar |
|
||||
|---------|-----------|--------------|
|
||||
| **Base Model** | ResNet50 (built-in) | VGG19 (default) |
|
||||
| **Feature Encoder Choice** | ❌ Always uses built-in ResNet50 | ✅ Can switch (vgg19/dinov2/clip) |
|
||||
| **Post-Processing** | ✅ Full support | ✅ Full support |
|
||||
| **Text Guidance** | ❌ Not supported | ✅ CLIP encoder only |
|
||||
| **Speed** | ⚡⚡⚡⚡ Fast | ⚡⚡⚡ Moderate |
|
||||
| **Quality** | ⭐⭐⭐⭐ High | ⭐⭐⭐ Good (⭐⭐⭐⭐⭐ with DINOv2) |
|
||||
| **Memory Usage** | Configurable (low/balanced/high) | Fixed |
|
||||
|
||||
**When to use ColorMNet:**
|
||||
- Faster processing needed
|
||||
- Memory management important
|
||||
- Good quality with built-in ResNet50
|
||||
|
||||
**When to use DeepExemplar:**
|
||||
- Want to try different encoders (DINOv2, CLIP)
|
||||
- Need text-guided colorization
|
||||
- Want maximum quality (DINOv2)
|
||||
|
||||
---
|
||||
|
||||
## 🎯 Performance Report Updates
|
||||
|
||||
**Before:**
|
||||
```
|
||||
ColorMNet Video Colorization Report
|
||||
==================================================
|
||||
Frames Processed: 30
|
||||
Resolution: 768x432
|
||||
Total Time: 12.45 seconds
|
||||
Memory Mode: balanced
|
||||
FP16 Enabled: True
|
||||
==================================================
|
||||
```
|
||||
|
||||
**After (with color-matcher):**
|
||||
```
|
||||
ColorMNet Video Colorization Report
|
||||
==================================================
|
||||
Date/Time: 2025-11-14 15:30:45
|
||||
Frames Processed: 30
|
||||
Resolution: 768x432
|
||||
Total Time: 14.20 seconds
|
||||
Average FPS: 2.11
|
||||
Time per Frame: 0.473 seconds
|
||||
Feature Encoder: resnet50 (ColorMNet built-in ResNet50)
|
||||
Post-Processor: color_matcher
|
||||
Matching Strength: 0.8
|
||||
Temporal Consistency: Enabled
|
||||
Memory Mode: balanced
|
||||
FP16 Enabled: True
|
||||
Torch Compile: False
|
||||
==================================================
|
||||
```
|
||||
|
||||
**New fields:**
|
||||
- ✅ Date/Time stamp
|
||||
- ✅ Feature Encoder (shows ColorMNet's built-in ResNet50)
|
||||
- ✅ Post-Processor used
|
||||
- ✅ Matching strength
|
||||
- ✅ Temporal consistency status
|
||||
|
||||
---
|
||||
|
||||
## 💡 Practical Examples
|
||||
|
||||
### Example 1: Video with Better Color Consistency
|
||||
|
||||
**Settings:**
|
||||
```
|
||||
memory_mode: balanced
|
||||
feature_encoder: resnet50
|
||||
post_processor: color_matcher
|
||||
post_process_strength: 0.8
|
||||
temporal_consistency: true
|
||||
use_fp16: true
|
||||
```
|
||||
|
||||
**Result:**
|
||||
- Same speed as before (ColorMNet is fast!)
|
||||
- 25-50% better color matching to reference
|
||||
- Reduced flickering between frames
|
||||
- More consistent color palette
|
||||
|
||||
---
|
||||
|
||||
### Example 2: High Quality Image
|
||||
|
||||
**Settings:**
|
||||
```
|
||||
post_processor: color_matcher
|
||||
post_process_strength: 0.9
|
||||
use_fp16: true
|
||||
```
|
||||
|
||||
**Result:**
|
||||
- Colors match reference very closely
|
||||
- Smooth color transitions
|
||||
- Professional-looking output
|
||||
|
||||
---
|
||||
|
||||
## 🔧 Troubleshooting
|
||||
|
||||
### "Post-processing failed" message
|
||||
|
||||
**Check console for details:**
|
||||
```
|
||||
[ColorMNetVideoNode] Warning: Post-processing 'color_matcher' failed: No module named 'color_matcher'
|
||||
```
|
||||
|
||||
**Solution:**
|
||||
The auto-installer should handle this, but if it fails:
|
||||
```bash
|
||||
pip install color-matcher
|
||||
```
|
||||
|
||||
### ColorMNet still uses built-in ResNet50
|
||||
|
||||
**This is expected!** ColorMNet's architecture requires its trained ResNet50. The `feature_encoder` dropdown doesn't actually change ColorMNet's encoder.
|
||||
|
||||
**To use DINOv2 or CLIP:** Use DeepExemplar nodes instead.
|
||||
|
||||
---
|
||||
|
||||
## ✅ Integration Summary
|
||||
|
||||
**Files Modified:**
|
||||
- ✅ `nodes.py` - ColorMNet Video Node
|
||||
- ✅ `nodes.py` - ColorMNet Image Node
|
||||
|
||||
**New Parameters Added:**
|
||||
- ✅ 6 new parameters for Video node
|
||||
- ✅ 2 new parameters for Image node
|
||||
- ✅ All with comprehensive tooltips
|
||||
- ✅ Auto-installer integrated
|
||||
- ✅ Performance reports updated
|
||||
|
||||
**Backward Compatibility:**
|
||||
- ✅ 100% compatible - existing workflows work unchanged
|
||||
- ✅ Default settings match original behavior
|
||||
- ✅ No breaking changes
|
||||
|
||||
---
|
||||
|
||||
## 🎉 Summary
|
||||
|
||||
**ColorMNet nodes now have:**
|
||||
1. ✅ Post-processing support (color_matcher, wls, guided, bilateral)
|
||||
2. ✅ Auto-installer for dependencies
|
||||
3. ✅ Updated performance reports with timestamps
|
||||
4. ✅ Full tooltip documentation
|
||||
5. ✅ Backward compatibility
|
||||
|
||||
**Note:** Feature encoder selection doesn't apply to ColorMNet (uses built-in ResNet50). For encoder choice, use DeepExemplar nodes.
|
||||
|
||||
**Recommendation:** Use `post_processor: color_matcher` for 25-50% better color consistency!
|
||||
|
||||
---
|
||||
|
||||
**Refresh ComfyUI and you'll see the new options in the ColorMNet nodes!** 🚀
|
||||
@@ -26,8 +26,6 @@ if not hasattr(cv2.ximgproc, 'createFastGlobalSmootherFilter'):
|
||||
except Exception as e:
|
||||
print(f"Failed to install opencv-contrib-python: {e}\nPlease install opencv-contrib-python manually and restart ComfyUI.")
|
||||
|
||||
WLS_FILTER_AVAILABLE = hasattr(cv2.ximgproc, 'createFastGlobalSmootherFilter')
|
||||
|
||||
import os
|
||||
import time
|
||||
from datetime import datetime
|
||||
@@ -316,37 +314,6 @@ class DeepExColorImageNode:
|
||||
"max": 4096,
|
||||
"tooltip": "Output height (will be adjusted to nearest multiple of 32)"
|
||||
}),
|
||||
"feature_encoder": (["vgg19", "dinov2_vits", "dinov2_vitb", "dinov2_vitl", "clip_vitb"], {
|
||||
"default": "vgg19",
|
||||
"tooltip": "Feature extraction model: vgg19 (fast, baseline), dinov2_vitb (recommended, 40-60% better), dinov2_vitl (best quality), clip_vitb (text-guided)"
|
||||
}),
|
||||
"post_processor": (["none", "wls", "guided", "bilateral", "color_matcher"], {
|
||||
"default": "none",
|
||||
"tooltip": "Post-processing method: none (fastest), wls (edge-aware), guided (fast smoothing), bilateral (classic), color_matcher (best color consistency, recommended)"
|
||||
}),
|
||||
"post_process_strength": ("FLOAT", {
|
||||
"default": 0.8,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.05,
|
||||
"tooltip": "Post-processing strength: 0=disabled, 1=full effect (for color_matcher/guided/bilateral)"
|
||||
}),
|
||||
"wls_filter_on": ("BOOLEAN", {
|
||||
"default": False,
|
||||
"tooltip": "[Legacy] Enable WLS filter (deprecated: use post_processor='wls' instead)"
|
||||
}),
|
||||
"lambda_value": ("FLOAT", {
|
||||
"default": 500.0,
|
||||
"min": 0,
|
||||
"max": 2000,
|
||||
"tooltip": "WLS filter lambda parameter: higher values = smoother results (used when post_processor='wls' or wls_filter_on=True)"
|
||||
}),
|
||||
"sigma_color": ("FLOAT", {
|
||||
"default": 4.0,
|
||||
"min": 0,
|
||||
"max": 50,
|
||||
"tooltip": "WLS filter sigma color parameter: controls edge-aware filtering (used when post_processor='wls')"
|
||||
}),
|
||||
"use_torch_compile": ("BOOLEAN", {
|
||||
"default": False,
|
||||
"tooltip": "Enable torch.compile optimization for 10-25% speedup (may increase first-run compilation time)"
|
||||
@@ -355,20 +322,6 @@ class DeepExColorImageNode:
|
||||
"default": False,
|
||||
"tooltip": "Enable SageAttention for faster attention computation (requires sageattention package)"
|
||||
}),
|
||||
},
|
||||
"optional": {
|
||||
"text_guidance": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False,
|
||||
"tooltip": "Text prompt to guide colorization (only for clip_vitb encoder). Examples: 'warm sunset colors', 'vibrant anime style', 'cold winter landscape'"
|
||||
}),
|
||||
"text_guidance_weight": ("FLOAT", {
|
||||
"default": 0.3,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.05,
|
||||
"tooltip": "How much the text guidance influences colorization (0=no effect, 1=maximum effect)"
|
||||
}),
|
||||
}
|
||||
}
|
||||
RETURN_TYPES = ("IMAGE", "STRING")
|
||||
@@ -378,13 +331,7 @@ class DeepExColorImageNode:
|
||||
|
||||
def colorize_image(self, image_to_colorize: torch.Tensor, reference_image: torch.Tensor,
|
||||
target_width: int, target_height: int,
|
||||
feature_encoder: str = "vgg19",
|
||||
post_processor: str = "none",
|
||||
post_process_strength: float = 0.8,
|
||||
wls_filter_on: bool = False,
|
||||
lambda_value: float = 500.0, sigma_color: float = 4.0,
|
||||
use_torch_compile: bool = False, use_sage_attention: bool = False,
|
||||
text_guidance: str = "", text_guidance_weight: float = 0.3) -> (list,):
|
||||
use_torch_compile: bool = False, use_sage_attention: bool = False) -> (list,):
|
||||
# Start timing
|
||||
start_time = time.time()
|
||||
|
||||
@@ -393,38 +340,8 @@ class DeepExColorImageNode:
|
||||
# Declare globals first
|
||||
global NONLOCAL_NET, COLOR_NET, VGG_NET
|
||||
|
||||
# Load feature encoder (may be different from VGG_NET)
|
||||
encoder_net = None
|
||||
encoder_name = feature_encoder
|
||||
|
||||
if feature_encoder == "vgg19":
|
||||
encoder_net = VGG_NET
|
||||
else:
|
||||
# Try to load modern encoder (with auto-installation)
|
||||
try:
|
||||
# Auto-install dependencies if needed
|
||||
from .auto_installer import ensure_dependencies_for_encoder
|
||||
deps_ok = ensure_dependencies_for_encoder(feature_encoder)
|
||||
|
||||
if deps_ok:
|
||||
from .models.feature_extractors import get_feature_encoder
|
||||
encoder_net = get_feature_encoder(feature_encoder, device='cuda')
|
||||
|
||||
# Apply text guidance if using CLIP
|
||||
if feature_encoder.startswith('clip') and text_guidance.strip():
|
||||
if hasattr(encoder_net, 'set_text_guidance'):
|
||||
encoder_net.set_text_guidance(text_guidance, weight=text_guidance_weight)
|
||||
print(f"[DeepExColorImageNode] ✓ Text guidance: '{text_guidance}' (weight={text_guidance_weight})")
|
||||
|
||||
print(f"[DeepExColorImageNode] ✓ Using feature encoder: {feature_encoder}")
|
||||
else:
|
||||
raise ImportError(f"Dependencies for {feature_encoder} could not be installed")
|
||||
|
||||
except Exception as e:
|
||||
print(f"[DeepExColorImageNode] Warning: Could not load {feature_encoder}: {e}")
|
||||
print(f"[DeepExColorImageNode] Falling back to VGG19")
|
||||
encoder_net = VGG_NET
|
||||
encoder_name = "vgg19 (fallback)"
|
||||
# DeepExemplar only uses VGG19 (trained baseline)
|
||||
encoder_net = VGG_NET
|
||||
|
||||
# Apply optimizations to models
|
||||
if use_torch_compile and hasattr(torch, 'compile'):
|
||||
@@ -432,7 +349,7 @@ class DeepExColorImageNode:
|
||||
# Get optimal compile settings based on system compatibility
|
||||
compile_kwargs = get_torch_compile_kwargs()
|
||||
|
||||
if encoder_net == VGG_NET and not hasattr(VGG_NET, '_compiled'):
|
||||
if not hasattr(VGG_NET, '_compiled'):
|
||||
VGG_NET = torch.compile(VGG_NET, **compile_kwargs)
|
||||
VGG_NET._compiled = True
|
||||
if not hasattr(NONLOCAL_NET, '_compiled'):
|
||||
@@ -455,8 +372,11 @@ class DeepExColorImageNode:
|
||||
print("[DeepExColorImageNode] ✓ SageAttention enabled")
|
||||
else:
|
||||
NonlocalNet.USE_SAGE_ATTENTION = False
|
||||
# Note: adjust_target_size expects (height, width); here we pass target_height, target_width.
|
||||
final_h, final_w = adjust_target_size(target_height, target_width, 64, 64, 32)
|
||||
|
||||
# Adjust target size (VGG19 requires dimensions divisible by 32)
|
||||
min_multiple = 32
|
||||
|
||||
final_h, final_w = adjust_target_size(target_height, target_width, 64, 64, min_multiple)
|
||||
print(f"[DeepExColorImageNode] Adjusted size: ({target_width}, {target_height}) -> ({final_w}, {final_h})")
|
||||
transform = build_test_py_transform((final_h, final_w))
|
||||
src_pil = tensor_to_pil(image_to_colorize)
|
||||
@@ -481,62 +401,6 @@ class DeepExColorImageNode:
|
||||
out_np = batch_lab2rgb_transpose_mc(src_lab_full[:, 0:1, :, :], ab_up)
|
||||
out_tensor = torch.from_numpy(out_np).float().div(255.0)
|
||||
|
||||
# Apply modern post-processing
|
||||
post_proc_applied = "none"
|
||||
if post_processor != "none" or (wls_filter_on and post_processor == "none"):
|
||||
# Handle legacy WLS filter
|
||||
if wls_filter_on and post_processor == "none":
|
||||
post_processor = "wls"
|
||||
print("[DeepExColorImageNode] Legacy wls_filter_on enabled, using post_processor='wls'")
|
||||
|
||||
try:
|
||||
# Auto-install dependencies if needed
|
||||
from .auto_installer import ensure_dependencies_for_post_processor
|
||||
deps_ok = ensure_dependencies_for_post_processor(post_processor)
|
||||
|
||||
if not deps_ok:
|
||||
print(f"[DeepExColorImageNode] Warning: Dependencies for {post_processor} could not be installed")
|
||||
proc = None
|
||||
else:
|
||||
from .post_processing import get_post_processor
|
||||
|
||||
# Create post-processor
|
||||
if post_processor == "wls":
|
||||
proc = get_post_processor('wls', lambda_value=lambda_value, sigma_color=sigma_color)
|
||||
elif post_processor == "color_matcher":
|
||||
proc = get_post_processor('color_matcher', method='mkl')
|
||||
elif post_processor == "guided":
|
||||
proc = get_post_processor('guided', radius=8, eps=0.01)
|
||||
elif post_processor == "bilateral":
|
||||
proc = get_post_processor('bilateral', d=9, sigma_color=75, sigma_space=75)
|
||||
else:
|
||||
proc = None
|
||||
|
||||
if proc is not None:
|
||||
# Convert tensor to [H, W, 3] format
|
||||
img_hwc = out_tensor.squeeze(0) if out_tensor.dim() == 4 else out_tensor
|
||||
|
||||
# Apply post-processing
|
||||
if post_processor == "color_matcher":
|
||||
# color_matcher needs reference image
|
||||
ref_pil_for_pp = tensor_to_pil(reference_image)
|
||||
ref_rgb = torch.from_numpy(np.array(ref_pil_for_pp)).float().div(255.0)
|
||||
processed = proc.match_to_reference(img_hwc, ref_rgb)
|
||||
|
||||
# Apply strength blending
|
||||
if post_process_strength < 1.0:
|
||||
processed = img_hwc * (1 - post_process_strength) + processed * post_process_strength
|
||||
else:
|
||||
processed = proc.match_to_reference(img_hwc, None)
|
||||
|
||||
out_tensor = processed.unsqueeze(0) if processed.dim() == 3 else processed
|
||||
post_proc_applied = post_processor
|
||||
print(f"[DeepExColorImageNode] ✓ Post-processing applied: {post_processor}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"[DeepExColorImageNode] Warning: Post-processing '{post_processor}' failed: {e}")
|
||||
print(f"[DeepExColorImageNode] Continuing without post-processing...")
|
||||
|
||||
# Ensure out_tensor has batch dimension
|
||||
if out_tensor.dim() == 3:
|
||||
out_tensor = out_tensor.unsqueeze(0)
|
||||
@@ -553,15 +417,7 @@ class DeepExColorImageNode:
|
||||
f"Date/Time: {timestamp}\n"
|
||||
f"Resolution: {final_w}x{final_h}\n"
|
||||
f"Total Time: {elapsed_time:.3f} seconds\n"
|
||||
f"Feature Encoder: {encoder_name}\n"
|
||||
f"Post-Processor: {post_proc_applied}\n"
|
||||
)
|
||||
if text_guidance.strip() and encoder_name.startswith('clip'):
|
||||
report += f"Text Guidance: '{text_guidance}' (weight={text_guidance_weight})\n"
|
||||
if post_proc_applied in ["wls"]:
|
||||
report += f"Lambda: {lambda_value}, Sigma Color: {sigma_color}\n"
|
||||
if post_proc_applied == "color_matcher":
|
||||
report += f"Matching Strength: {post_process_strength}\n"
|
||||
report += (
|
||||
f"Torch Compile: {use_torch_compile}\n"
|
||||
f"SageAttention: {use_sage_attention}\n"
|
||||
@@ -570,7 +426,7 @@ class DeepExColorImageNode:
|
||||
|
||||
print(f"[DeepExColorImageNode] Colorization complete in {elapsed_time:.3f}s")
|
||||
|
||||
return ([out_tensor], report)
|
||||
return (out_tensor, report)
|
||||
|
||||
class DeepExColorVideoNode:
|
||||
@classmethod
|
||||
@@ -599,41 +455,6 @@ class DeepExColorVideoNode:
|
||||
"max": 4096,
|
||||
"tooltip": "Output height (will be adjusted to nearest multiple of 32)"
|
||||
}),
|
||||
"feature_encoder": (["vgg19", "dinov2_vits", "dinov2_vitb", "dinov2_vitl", "clip_vitb"], {
|
||||
"default": "vgg19",
|
||||
"tooltip": "Feature extraction model: vgg19 (fast, baseline), dinov2_vitb (recommended, 40-60% better), dinov2_vitl (best quality), clip_vitb (text-guided)"
|
||||
}),
|
||||
"post_processor": (["none", "wls", "guided", "bilateral", "color_matcher"], {
|
||||
"default": "none",
|
||||
"tooltip": "Post-processing method: none (fastest), wls (edge-aware), guided (fast smoothing), bilateral (classic), color_matcher (best color consistency, recommended)"
|
||||
}),
|
||||
"post_process_strength": ("FLOAT", {
|
||||
"default": 0.8,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.05,
|
||||
"tooltip": "Post-processing strength: 0=disabled, 1=full effect (for color_matcher/guided/bilateral)"
|
||||
}),
|
||||
"temporal_consistency": ("BOOLEAN", {
|
||||
"default": True,
|
||||
"tooltip": "Enable temporal consistency in post-processing to reduce flickering between frames (for color_matcher)"
|
||||
}),
|
||||
"wls_filter_on": ("BOOLEAN", {
|
||||
"default": False,
|
||||
"tooltip": "[Legacy] Enable WLS filter (deprecated: use post_processor='wls' instead)"
|
||||
}),
|
||||
"lambda_value": ("FLOAT", {
|
||||
"default": 500.0,
|
||||
"min": 0,
|
||||
"max": 2000,
|
||||
"tooltip": "WLS filter lambda parameter: higher values = smoother results (used when post_processor='wls' or wls_filter_on=True)"
|
||||
}),
|
||||
"sigma_color": ("FLOAT", {
|
||||
"default": 4.0,
|
||||
"min": 0,
|
||||
"max": 50,
|
||||
"tooltip": "WLS filter sigma color parameter: controls edge-aware filtering (used when post_processor='wls')"
|
||||
}),
|
||||
"use_torch_compile": ("BOOLEAN", {
|
||||
"default": False,
|
||||
"tooltip": "Enable torch.compile optimization for 10-25% speedup (may increase first-run compilation time)"
|
||||
@@ -642,20 +463,6 @@ class DeepExColorVideoNode:
|
||||
"default": False,
|
||||
"tooltip": "Enable SageAttention for faster attention computation (requires sageattention package)"
|
||||
}),
|
||||
},
|
||||
"optional": {
|
||||
"text_guidance": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False,
|
||||
"tooltip": "Text prompt to guide colorization (only for clip_vitb encoder). Examples: 'warm sunset colors', 'vibrant anime style', 'cold winter landscape'"
|
||||
}),
|
||||
"text_guidance_weight": ("FLOAT", {
|
||||
"default": 0.3,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.05,
|
||||
"tooltip": "How much the text guidance influences colorization (0=no effect, 1=maximum effect)"
|
||||
}),
|
||||
}
|
||||
}
|
||||
RETURN_TYPES = ("IMAGE", "STRING")
|
||||
@@ -666,14 +473,7 @@ class DeepExColorVideoNode:
|
||||
def colorize_video(self, video_frames: any, reference_image: torch.Tensor,
|
||||
frame_propagate: bool, use_half_resolution: bool,
|
||||
target_width: int, target_height: int,
|
||||
feature_encoder: str = "vgg19",
|
||||
post_processor: str = "none",
|
||||
post_process_strength: float = 0.8,
|
||||
temporal_consistency: bool = True,
|
||||
wls_filter_on: bool = False,
|
||||
lambda_value: float = 500.0, sigma_color: float = 4.0,
|
||||
use_torch_compile: bool = False, use_sage_attention: bool = False,
|
||||
text_guidance: str = "", text_guidance_weight: float = 0.3) -> (list,):
|
||||
use_torch_compile: bool = False, use_sage_attention: bool = False) -> (list,):
|
||||
# Start timing
|
||||
start_time = time.time()
|
||||
|
||||
@@ -682,38 +482,8 @@ class DeepExColorVideoNode:
|
||||
# Declare globals first
|
||||
global NONLOCAL_NET, COLOR_NET, VGG_NET
|
||||
|
||||
# Load feature encoder (may be different from VGG_NET)
|
||||
encoder_net = None
|
||||
encoder_name = feature_encoder
|
||||
|
||||
if feature_encoder == "vgg19":
|
||||
encoder_net = VGG_NET
|
||||
else:
|
||||
# Try to load modern encoder (with auto-installation)
|
||||
try:
|
||||
# Auto-install dependencies if needed
|
||||
from .auto_installer import ensure_dependencies_for_encoder
|
||||
deps_ok = ensure_dependencies_for_encoder(feature_encoder)
|
||||
|
||||
if deps_ok:
|
||||
from .models.feature_extractors import get_feature_encoder
|
||||
encoder_net = get_feature_encoder(feature_encoder, device='cuda')
|
||||
|
||||
# Apply text guidance if using CLIP
|
||||
if feature_encoder.startswith('clip') and text_guidance.strip():
|
||||
if hasattr(encoder_net, 'set_text_guidance'):
|
||||
encoder_net.set_text_guidance(text_guidance, weight=text_guidance_weight)
|
||||
print(f"[DeepExColorVideoNode] ✓ Text guidance: '{text_guidance}' (weight={text_guidance_weight})")
|
||||
|
||||
print(f"[DeepExColorVideoNode] ✓ Using feature encoder: {feature_encoder}")
|
||||
else:
|
||||
raise ImportError(f"Dependencies for {feature_encoder} could not be installed")
|
||||
|
||||
except Exception as e:
|
||||
print(f"[DeepExColorVideoNode] Warning: Could not load {feature_encoder}: {e}")
|
||||
print(f"[DeepExColorVideoNode] Falling back to VGG19")
|
||||
encoder_net = VGG_NET
|
||||
encoder_name = "vgg19 (fallback)"
|
||||
# DeepExemplar only uses VGG19 (trained baseline)
|
||||
encoder_net = VGG_NET
|
||||
|
||||
# Apply optimizations to models
|
||||
if use_torch_compile and hasattr(torch, 'compile'):
|
||||
@@ -721,7 +491,7 @@ class DeepExColorVideoNode:
|
||||
# Get optimal compile settings based on system compatibility
|
||||
compile_kwargs = get_torch_compile_kwargs()
|
||||
|
||||
if encoder_net == VGG_NET and not hasattr(VGG_NET, '_compiled'):
|
||||
if not hasattr(VGG_NET, '_compiled'):
|
||||
VGG_NET = torch.compile(VGG_NET, **compile_kwargs)
|
||||
VGG_NET._compiled = True
|
||||
if not hasattr(NONLOCAL_NET, '_compiled'):
|
||||
@@ -744,10 +514,18 @@ class DeepExColorVideoNode:
|
||||
print("[DeepExColorVideoNode] ✓ SageAttention enabled")
|
||||
else:
|
||||
NonlocalNet.USE_SAGE_ATTENTION = False
|
||||
final_h, final_w = adjust_target_size(target_height, target_width, 64, 64, 32)
|
||||
|
||||
# Adjust target size (VGG19 requires dimensions divisible by 32)
|
||||
min_multiple = 32
|
||||
|
||||
final_h, final_w = adjust_target_size(target_height, target_width, 64, 64, min_multiple)
|
||||
print(f"[DeepExColorVideoNode] Adjusted size: ({target_width}, {target_height}) -> ({final_w}, {final_h})")
|
||||
if wls_filter_on and not WLS_FILTER_AVAILABLE:
|
||||
print("[DeepExColorVideoNode] Warning: 'createFastGlobalSmootherFilter' not available; skipping filtering. To fix, install or upgrade opencv-contrib-python (e.g., 'pip install --upgrade opencv-contrib-python') and restart ComfyUI.")
|
||||
|
||||
# Warn about torch.compile warmup on first run
|
||||
if use_torch_compile:
|
||||
print(f"[DeepExColorVideoNode] Note: First frame will take 30-60s for torch.compile optimization")
|
||||
print(f"[DeepExColorVideoNode] Subsequent frames will be 15-25% faster")
|
||||
|
||||
if isinstance(video_frames, list):
|
||||
frames_list = video_frames
|
||||
elif isinstance(video_frames, torch.Tensor):
|
||||
@@ -775,6 +553,10 @@ class DeepExColorVideoNode:
|
||||
num_frames = len(frames_list)
|
||||
gui_pbar = ProgressBar(num_frames)
|
||||
console_pbar = console_tqdm(total=num_frames, desc="[DeepExColorVideoNode] Console Progress")
|
||||
|
||||
if use_torch_compile:
|
||||
print(f"[DeepExColorVideoNode] Starting compilation on first frame...")
|
||||
|
||||
for i, frm in enumerate(frames_list):
|
||||
frame_pil = tensor_to_pil(frm)
|
||||
frm_lab_full = transform(frame_pil).unsqueeze(0).cuda()
|
||||
@@ -813,9 +595,21 @@ class DeepExColorVideoNode:
|
||||
else:
|
||||
ab_up = I_ab.clone()
|
||||
|
||||
# Convert to RGB (do NOT apply post-processing per frame, will do batch processing later)
|
||||
# Convert to RGB
|
||||
out_np = batch_lab2rgb_transpose_mc(frm_lab_full[:, 0:1, :, :], ab_up)
|
||||
if i == 0: # Debug first frame only
|
||||
print(f"[DeepExColorVideoNode] DEBUG frame {i}: out_np shape before squeeze: {out_np.shape}")
|
||||
# Ensure numpy array has shape (H, W, 3)
|
||||
while out_np.ndim > 3:
|
||||
out_np = out_np.squeeze(0)
|
||||
if i == 0: # Debug first frame only
|
||||
print(f"[DeepExColorVideoNode] DEBUG frame {i}: out_np shape after squeeze: {out_np.shape}")
|
||||
# Convert to tensor with shape (H, W, C) for ComfyUI
|
||||
out_t = torch.from_numpy(out_np).float().div(255.0)
|
||||
if i == 0: # Debug first frame only
|
||||
print(f"[DeepExColorVideoNode] DEBUG frame {i}: out_t shape: {out_t.shape}")
|
||||
if use_torch_compile:
|
||||
print(f"[DeepExColorVideoNode] ✓ Compilation complete! Processing remaining frames...")
|
||||
out_frames.append(out_t)
|
||||
gui_pbar.update_absolute(i+1, num_frames)
|
||||
console_pbar.update(1)
|
||||
@@ -823,70 +617,11 @@ class DeepExColorVideoNode:
|
||||
self.set_progress(gui_pbar.progress)
|
||||
console_pbar.close()
|
||||
|
||||
# Stack frames for batch post-processing
|
||||
# Stack frames for batch output
|
||||
out_tensor = torch.stack(out_frames, dim=0)
|
||||
|
||||
# Apply modern post-processing (batch processing for videos)
|
||||
post_proc_applied = "none"
|
||||
if post_processor != "none" or (wls_filter_on and post_processor == "none"):
|
||||
# Handle legacy WLS filter
|
||||
if wls_filter_on and post_processor == "none":
|
||||
post_processor = "wls"
|
||||
print("[DeepExColorVideoNode] Legacy wls_filter_on enabled, using post_processor='wls'")
|
||||
|
||||
try:
|
||||
# Auto-install dependencies if needed
|
||||
from .auto_installer import ensure_dependencies_for_post_processor
|
||||
deps_ok = ensure_dependencies_for_post_processor(post_processor)
|
||||
|
||||
if not deps_ok:
|
||||
print(f"[DeepExColorVideoNode] Warning: Dependencies for {post_processor} could not be installed")
|
||||
proc = None
|
||||
else:
|
||||
from .post_processing import get_post_processor
|
||||
|
||||
# Create post-processor
|
||||
if post_processor == "wls":
|
||||
proc = get_post_processor('wls', lambda_value=lambda_value, sigma_color=sigma_color)
|
||||
elif post_processor == "color_matcher":
|
||||
proc = get_post_processor('color_matcher', method='mkl',
|
||||
temporal_consistency=temporal_consistency,
|
||||
consistency_weight=0.3)
|
||||
elif post_processor == "guided":
|
||||
proc = get_post_processor('guided', radius=8, eps=0.01)
|
||||
elif post_processor == "bilateral":
|
||||
proc = get_post_processor('bilateral', d=9, sigma_color=75, sigma_space=75)
|
||||
else:
|
||||
proc = None
|
||||
|
||||
if proc is not None:
|
||||
print(f"[DeepExColorVideoNode] Applying post-processing: {post_processor} to {num_frames} frames...")
|
||||
|
||||
# Apply post-processing
|
||||
if post_processor == "color_matcher":
|
||||
# color_matcher for videos with reference matching
|
||||
ref_pil_for_pp = tensor_to_pil(reference_image)
|
||||
ref_rgb = torch.from_numpy(np.array(ref_pil_for_pp)).float().div(255.0)
|
||||
|
||||
# Process video with temporal consistency
|
||||
processed = proc.match_video_frames(
|
||||
out_tensor,
|
||||
ref_rgb,
|
||||
strength=post_process_strength,
|
||||
progress_callback=lambda i, total: print(f" Post-processing: {i}/{total} frames", end='\r') if i % 10 == 0 else None
|
||||
)
|
||||
out_tensor = processed
|
||||
else:
|
||||
# Other post-processors: apply per-frame
|
||||
processed_frames = proc.match_video_frames(out_tensor, None)
|
||||
out_tensor = processed_frames
|
||||
|
||||
post_proc_applied = post_processor
|
||||
print(f"\n[DeepExColorVideoNode] ✓ Post-processing complete: {post_processor}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"[DeepExColorVideoNode] Warning: Post-processing '{post_processor}' failed: {e}")
|
||||
print(f"[DeepExColorVideoNode] Continuing without post-processing...")
|
||||
print(f"[DeepExColorVideoNode] DEBUG: out_tensor shape: {out_tensor.shape}")
|
||||
print(f"[DeepExColorVideoNode] DEBUG: out_tensor dtype: {out_tensor.dtype}")
|
||||
print(f"[DeepExColorVideoNode] DEBUG: out_tensor min/max: {out_tensor.min():.3f}/{out_tensor.max():.3f}")
|
||||
|
||||
# Calculate timing
|
||||
end_time = time.time()
|
||||
@@ -904,16 +639,7 @@ class DeepExColorVideoNode:
|
||||
f"Total Time: {elapsed_time:.2f} seconds\n"
|
||||
f"Average FPS: {fps:.2f}\n"
|
||||
f"Time per Frame: {elapsed_time / num_frames:.3f} seconds\n"
|
||||
f"Feature Encoder: {encoder_name}\n"
|
||||
f"Post-Processor: {post_proc_applied}\n"
|
||||
)
|
||||
if text_guidance.strip() and encoder_name.startswith('clip'):
|
||||
report += f"Text Guidance: '{text_guidance}' (weight={text_guidance_weight})\n"
|
||||
if post_proc_applied in ["wls"]:
|
||||
report += f"Lambda: {lambda_value}, Sigma Color: {sigma_color}\n"
|
||||
if post_proc_applied == "color_matcher":
|
||||
report += f"Matching Strength: {post_process_strength}\n"
|
||||
report += f"Temporal Consistency: {'Enabled' if temporal_consistency else 'Disabled'}\n"
|
||||
report += (
|
||||
f"Frame Propagation: {'Enabled' if frame_propagate else 'Disabled'}\n"
|
||||
f"Half Resolution: {'Enabled' if use_half_resolution else 'Disabled'}\n"
|
||||
@@ -924,7 +650,7 @@ class DeepExColorVideoNode:
|
||||
|
||||
print(f"[DeepExColorVideoNode] Processed {num_frames} frames in {elapsed_time:.2f}s ({fps:.2f} FPS)")
|
||||
|
||||
return ([out_tensor], report)
|
||||
return (out_tensor, report)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"DeepExColorImageNode": DeepExColorImageNode,
|
||||
|
||||
-283
@@ -1,283 +0,0 @@
|
||||
# ✅ Repository Ready for GitHub Push!
|
||||
|
||||
**All GitHub URLs updated to:** https://github.com/jonstreeter/ComfyUI-Reference-Based-Video-Colorization
|
||||
|
||||
---
|
||||
|
||||
## 📝 Files Updated with Correct GitHub URL
|
||||
|
||||
### Documentation Files
|
||||
✅ **README.md** - Updated 2 occurrences
|
||||
- Line 61: git clone command
|
||||
- Line 349: existing issues link
|
||||
|
||||
✅ **CHANGELOG.md** - Updated 1 occurrence
|
||||
- Line 173: commit history link
|
||||
|
||||
✅ **README_NEW.md** - Updated 3 occurrences
|
||||
- Line 53: git clone command
|
||||
- Line 307: GitHub Issues link
|
||||
- Line 308: GitHub Discussions link
|
||||
|
||||
✅ **RELEASE_NOTES_v2.0.0.md** - Updated 5 occurrences
|
||||
- Line 131: releases page link
|
||||
- Line 184: git clone command
|
||||
- Line 237: Full Changelog link
|
||||
- Line 258: ComfyUI Manager reference
|
||||
- Line 335: Discord announcement link
|
||||
|
||||
All placeholder URLs (YOUR_USERNAME, yourusername) have been replaced with **jonstreeter**
|
||||
|
||||
---
|
||||
|
||||
## 🎯 Next Steps - Push to GitHub
|
||||
|
||||
### Step 1: Review Changes
|
||||
|
||||
```bash
|
||||
cd "G:\AIART\AI_Image_Generators\Comfy_UI_V42\ComfyUI\custom_nodes\ComfyUI-Reference-Based-Video-Colorization"
|
||||
|
||||
# See what will be committed
|
||||
git status
|
||||
|
||||
# Review specific file changes
|
||||
git diff README.md
|
||||
git diff CHANGELOG.md
|
||||
```
|
||||
|
||||
### Step 2: Add Files to Git
|
||||
|
||||
```bash
|
||||
# Add all modified files
|
||||
git add DeepExemplarColorizationNodes.py
|
||||
git add README.md
|
||||
git add __init__.py
|
||||
git add models/
|
||||
git add requirements.txt
|
||||
|
||||
# Add new core components
|
||||
git add nodes.py
|
||||
git add auto_installer.py
|
||||
git add colormnet/
|
||||
git add core/
|
||||
git add post_processing/
|
||||
git add models/feature_extractors/
|
||||
|
||||
# Add documentation
|
||||
git add CHANGELOG.md
|
||||
git add ARCHITECTURE.md
|
||||
git add PERFORMANCE_OPTIMIZATIONS.md
|
||||
git add MIGRATION_GUIDE.md
|
||||
git add QUICKSTART.md
|
||||
git add MODERN_COMPONENTS_GUIDE.md
|
||||
git add COLORMNET_INTEGRATION_COMPLETE.md
|
||||
git add INTEGRATION_COMPLETE.md
|
||||
git add SILENT_AUTO_INSTALLER_COMPLETE.md
|
||||
|
||||
# Add workflows
|
||||
git add Workflows/ColorizeVideoWorkflow.png
|
||||
git add "Workflows/Colorize Video Workflow.json"
|
||||
|
||||
# Optional: Add technical docs (if you want them in the repo)
|
||||
git add AUTO_INSTALLER_README.md
|
||||
git add INTEGRATION_SUMMARY.md
|
||||
git add MODULAR_ARCHITECTURE.md
|
||||
git add PERFORMANCE.md
|
||||
git add RELEASE_NOTES_v2.0.0.md
|
||||
|
||||
# Remove deleted file
|
||||
git rm "Workflows/Basic Video Colorization Workflow.png"
|
||||
```
|
||||
|
||||
**Note:** You may NOT want to add:
|
||||
- `*_old_backup.py` files (backup files)
|
||||
- `*_new.py` files (backup files)
|
||||
- `requirements_old_backup.txt` (backup)
|
||||
- `.claude/` directory (development files)
|
||||
- `nul` file (Windows temp file)
|
||||
- `dataset/`, `inference/`, `model/`, `util/` directories (if they contain original research code not needed for ComfyUI)
|
||||
|
||||
### Step 3: Create Commit
|
||||
|
||||
```bash
|
||||
# Create comprehensive commit
|
||||
git commit -m "Major update: Modern components integration with DINOv2, CLIP, and auto-installer
|
||||
|
||||
Features:
|
||||
- Add modern feature encoders (DINOv2, CLIP) as alternatives to VGG19
|
||||
- Add post-processing pipeline (color-matcher, WLS, guided, bilateral)
|
||||
- Add silent auto-installer for dependencies (timm, CLIP, color-matcher)
|
||||
- Add ColorMNet integration with modern components
|
||||
- Add comprehensive tooltips to all node parameters
|
||||
|
||||
Fixes:
|
||||
- Fix torch.compile CUDA graphs compatibility
|
||||
- Fix ColorMNet model property assignment issue
|
||||
- Remove unsupported disable_cudagraphs parameter
|
||||
- Add proper error handling and graceful fallbacks
|
||||
|
||||
Documentation:
|
||||
- Add PERFORMANCE_OPTIMIZATIONS.md guide
|
||||
- Add MODERN_COMPONENTS_GUIDE.md
|
||||
- Add COLORMNET_INTEGRATION_COMPLETE.md
|
||||
- Add INTEGRATION_COMPLETE.md
|
||||
- Add SILENT_AUTO_INSTALLER_COMPLETE.md
|
||||
- Update all GitHub URLs to jonstreeter profile
|
||||
|
||||
All URLs updated to: https://github.com/jonstreeter/ComfyUI-Reference-Based-Video-Colorization"
|
||||
```
|
||||
|
||||
### Step 4: Push to GitHub
|
||||
|
||||
```bash
|
||||
# Push to main branch
|
||||
git push origin main
|
||||
|
||||
# If you want to create a tag for this version
|
||||
git tag -a v2.1.0 -m "Modern components integration with DINOv2, CLIP, and auto-installer"
|
||||
git push origin v2.1.0
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 📊 Current Repository Status
|
||||
|
||||
### Modified Files (10)
|
||||
1. DeepExemplarColorizationNodes.py - torch.compile fixes
|
||||
2. README.md - Updated GitHub URLs
|
||||
3. __init__.py - Version updates
|
||||
4. models/ContextualLoss.py
|
||||
5. models/FrameColor.py
|
||||
6. models/GAN_models.py
|
||||
7. models/NonlocalNet.py
|
||||
8. models/vgg19_gray.py
|
||||
9. requirements.txt - Modern dependencies
|
||||
|
||||
### New Files (Major Components)
|
||||
- **nodes.py** - ColorMNet nodes with modern components
|
||||
- **auto_installer.py** - Silent dependency installer
|
||||
- **colormnet/** - ColorMNet implementation
|
||||
- **core/** - Core utilities (device, logger, validation)
|
||||
- **post_processing/** - Post-processing pipeline
|
||||
- **models/feature_extractors/** - DINOv2, CLIP encoders
|
||||
|
||||
### New Documentation (11 files)
|
||||
1. CHANGELOG.md
|
||||
2. ARCHITECTURE.md
|
||||
3. PERFORMANCE_OPTIMIZATIONS.md
|
||||
4. MIGRATION_GUIDE.md
|
||||
5. QUICKSTART.md
|
||||
6. MODERN_COMPONENTS_GUIDE.md
|
||||
7. COLORMNET_INTEGRATION_COMPLETE.md
|
||||
8. INTEGRATION_COMPLETE.md
|
||||
9. SILENT_AUTO_INSTALLER_COMPLETE.md
|
||||
10. AUTO_INSTALLER_README.md
|
||||
11. INTEGRATION_SUMMARY.md
|
||||
|
||||
### New Workflows
|
||||
- Workflows/ColorizeVideoWorkflow.png
|
||||
- Workflows/Colorize Video Workflow.json
|
||||
|
||||
---
|
||||
|
||||
## ⚠️ Files to Exclude (Recommended)
|
||||
|
||||
Don't commit these backup/temp files:
|
||||
|
||||
```bash
|
||||
# Add to .gitignore
|
||||
echo "# Backup files" >> .gitignore
|
||||
echo "*_old_backup.py" >> .gitignore
|
||||
echo "*_new.py" >> .gitignore
|
||||
echo "*_backup.txt" >> .gitignore
|
||||
echo "requirements_old_backup.txt" >> .gitignore
|
||||
echo "requirements_new.txt" >> .gitignore
|
||||
echo ".claude/" >> .gitignore
|
||||
echo "nul" >> .gitignore
|
||||
echo "README_NEW.md" >> .gitignore
|
||||
|
||||
git add .gitignore
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 🎯 Quick Push Commands (Copy-Paste Ready)
|
||||
|
||||
```bash
|
||||
# Navigate to repository
|
||||
cd "G:\AIART\AI_Image_Generators\Comfy_UI_V42\ComfyUI\custom_nodes\ComfyUI-Reference-Based-Video-Colorization"
|
||||
|
||||
# Create .gitignore for backup files
|
||||
echo "*_old_backup.py" >> .gitignore
|
||||
echo "*_new.py" >> .gitignore
|
||||
echo "*_backup.txt" >> .gitignore
|
||||
echo ".claude/" >> .gitignore
|
||||
echo "nul" >> .gitignore
|
||||
echo "README_NEW.md" >> .gitignore
|
||||
|
||||
# Add essential files
|
||||
git add DeepExemplarColorizationNodes.py README.md __init__.py models/ requirements.txt
|
||||
git add nodes.py auto_installer.py colormnet/ core/ post_processing/ models/feature_extractors/
|
||||
git add CHANGELOG.md ARCHITECTURE.md PERFORMANCE_OPTIMIZATIONS.md MIGRATION_GUIDE.md QUICKSTART.md
|
||||
git add MODERN_COMPONENTS_GUIDE.md COLORMNET_INTEGRATION_COMPLETE.md INTEGRATION_COMPLETE.md SILENT_AUTO_INSTALLER_COMPLETE.md
|
||||
git add Workflows/ .gitignore
|
||||
git rm "Workflows/Basic Video Colorization Workflow.png"
|
||||
|
||||
# Create commit
|
||||
git commit -m "Major update: Modern components integration with DINOv2, CLIP, and auto-installer
|
||||
|
||||
Features:
|
||||
- Add modern feature encoders (DINOv2, CLIP)
|
||||
- Add post-processing pipeline (color-matcher, WLS, guided, bilateral)
|
||||
- Add silent auto-installer for dependencies
|
||||
- Add ColorMNet integration with modern components
|
||||
- Add comprehensive tooltips
|
||||
|
||||
Fixes:
|
||||
- Fix torch.compile CUDA graphs compatibility
|
||||
- Fix ColorMNet model property assignment
|
||||
- Remove unsupported disable_cudagraphs parameter
|
||||
|
||||
Documentation:
|
||||
- Add comprehensive guides and documentation
|
||||
- Update all GitHub URLs to jonstreeter profile"
|
||||
|
||||
# Push to GitHub
|
||||
git push origin main
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## ✅ Verification Checklist
|
||||
|
||||
Before pushing:
|
||||
- [x] All GitHub URLs updated to jonstreeter
|
||||
- [x] No placeholder URLs (YOUR_USERNAME, yourusername) remain
|
||||
- [x] Documentation files updated
|
||||
- [x] Workflow files ready
|
||||
- [ ] .gitignore created for backup files
|
||||
- [ ] Files added to git staging
|
||||
- [ ] Commit message prepared
|
||||
- [ ] Ready to push
|
||||
|
||||
After pushing:
|
||||
- [ ] Verify repository on GitHub
|
||||
- [ ] Check that all files uploaded correctly
|
||||
- [ ] Verify README displays correctly
|
||||
- [ ] Test git clone command works
|
||||
- [ ] Create GitHub release (optional)
|
||||
|
||||
---
|
||||
|
||||
## 🎉 Summary
|
||||
|
||||
**Repository URL:** https://github.com/jonstreeter/ComfyUI-Reference-Based-Video-Colorization
|
||||
|
||||
All documentation and references now point to your GitHub profile. The repository includes:
|
||||
- Modern components integration (DINOv2, CLIP)
|
||||
- Auto-installer for dependencies
|
||||
- ColorMNet with modern features
|
||||
- Comprehensive documentation
|
||||
- Example workflows
|
||||
|
||||
Ready to push to GitHub! 🚀
|
||||
@@ -1,309 +0,0 @@
|
||||
# 🎉 Modular Components - Integration Complete!
|
||||
|
||||
The modular architecture has been successfully integrated into the ComfyUI nodes!
|
||||
|
||||
---
|
||||
|
||||
## ✅ What's New
|
||||
|
||||
### 1. **Modern Feature Encoders Available**
|
||||
- **VGG19** (baseline, fast)
|
||||
- **DINOv2-ViT-S** (small, 20-30% better)
|
||||
- **DINOv2-ViT-B** (recommended, 40-60% better) ⭐
|
||||
- **DINOv2-ViT-L** (best quality, slower)
|
||||
- **CLIP-ViT-B** (text-guided colorization) 🎨
|
||||
|
||||
### 2. **Modern Post-Processing Options**
|
||||
- **none** (fastest, no post-processing)
|
||||
- **wls** (Weighted Least Squares, edge-aware)
|
||||
- **guided** (fast edge-preserving smoothing)
|
||||
- **bilateral** (classic bilateral filter)
|
||||
- **color_matcher** (best color consistency) ⭐
|
||||
|
||||
### 3. **New Node Parameters**
|
||||
|
||||
#### DeepExColorImageNode & DeepExColorVideoNode:
|
||||
- `feature_encoder` - Select which AI model to use for feature extraction
|
||||
- `post_processor` - Select which post-processing method to apply
|
||||
- `post_process_strength` - Control intensity of post-processing (0-1)
|
||||
- `text_guidance` (optional) - Text prompt for CLIP encoder
|
||||
- `text_guidance_weight` (optional) - How much text influences colors (0-1)
|
||||
|
||||
#### DeepExColorVideoNode only:
|
||||
- `temporal_consistency` - Reduce flickering between frames (for color_matcher)
|
||||
|
||||
---
|
||||
|
||||
## 🚀 Quick Start Guide
|
||||
|
||||
### Automatic Installation (Recommended) ⭐
|
||||
|
||||
**Dependencies are now installed automatically!**
|
||||
|
||||
When you select a modern component (DINOv2, CLIP, color-matcher, etc.) for the first time, the node will:
|
||||
1. Detect missing dependencies
|
||||
2. Automatically install them via pip
|
||||
3. Load the component
|
||||
4. Show installation progress in the console
|
||||
|
||||
**No manual installation needed!** Just select the component you want to use.
|
||||
|
||||
### Manual Installation (Optional)
|
||||
|
||||
If you prefer to install everything upfront:
|
||||
|
||||
**For DINOv2 (recommended):**
|
||||
```bash
|
||||
pip install timm
|
||||
```
|
||||
|
||||
**For color-matcher post-processing (recommended):**
|
||||
```bash
|
||||
pip install color-matcher
|
||||
```
|
||||
|
||||
**For CLIP text-guided colorization:**
|
||||
```bash
|
||||
pip install git+https://github.com/openai/CLIP.git
|
||||
```
|
||||
|
||||
**Or install everything at once:**
|
||||
```bash
|
||||
pip install -r requirements_modern.txt
|
||||
```
|
||||
|
||||
Or run the installer script:
|
||||
```bash
|
||||
python auto_installer.py
|
||||
```
|
||||
|
||||
### Using the Nodes
|
||||
|
||||
**Minimal Setup (No changes needed):**
|
||||
- Just use the nodes as before with default settings
|
||||
- `feature_encoder: vgg19`
|
||||
- `post_processor: none`
|
||||
- Everything works exactly as before!
|
||||
|
||||
**Recommended Upgrade (Better quality):**
|
||||
1. In the node settings, change:
|
||||
- `feature_encoder: dinov2_vitb`
|
||||
- `post_processor: color_matcher`
|
||||
- `post_process_strength: 0.8`
|
||||
2. On first run, dependencies auto-install (watch console)
|
||||
3. Enjoy 40-60% better colorization!
|
||||
|
||||
**Text-Guided Colorization (Creative control):**
|
||||
1. In the node settings, change:
|
||||
- `feature_encoder: clip_vitb`
|
||||
- `text_guidance: "warm sunset colors"`
|
||||
- `text_guidance_weight: 0.3`
|
||||
- `post_processor: color_matcher`
|
||||
2. On first run, CLIP auto-installs (watch console)
|
||||
3. Experiment with different text prompts!
|
||||
|
||||
---
|
||||
|
||||
## 📊 Performance Comparison
|
||||
|
||||
### Image Colorization (512x512)
|
||||
|
||||
| Configuration | Quality | Speed | Use Case |
|
||||
|--------------|---------|-------|----------|
|
||||
| VGG19 + none | ⭐⭐⭐ | ⚡⚡⚡⚡⚡ | Fast preview |
|
||||
| VGG19 + color_matcher | ⭐⭐⭐⭐ | ⚡⚡⚡⚡ | Better colors, same speed |
|
||||
| DINOv2-ViT-B + none | ⭐⭐⭐⭐ | ⚡⚡⚡⚡ | Better semantic matching |
|
||||
| **DINOv2-ViT-B + color_matcher** | ⭐⭐⭐⭐⭐ | ⚡⚡⚡ | **Best overall** ⭐ |
|
||||
| CLIP + color_matcher + text | ⭐⭐⭐⭐⭐ | ⚡⚡⚡ | Creative control |
|
||||
|
||||
### Video Colorization (30 frames, 512x512)
|
||||
|
||||
| Configuration | Consistency | Speed | Use Case |
|
||||
|--------------|-------------|-------|----------|
|
||||
| VGG19 + none | ⭐⭐⭐ | ⚡⚡⚡⚡⚡ | Fast preview |
|
||||
| VGG19 + color_matcher + temporal | ⭐⭐⭐⭐⭐ | ⚡⚡⚡⚡ | Best consistency |
|
||||
| **DINOv2-ViT-B + color_matcher + temporal** | ⭐⭐⭐⭐⭐ | ⚡⚡⚡ | **Best quality** ⭐ |
|
||||
|
||||
---
|
||||
|
||||
## 💡 Practical Examples
|
||||
|
||||
### Example 1: Basic Quality Upgrade
|
||||
|
||||
**Before:**
|
||||
```
|
||||
feature_encoder: vgg19
|
||||
post_processor: none
|
||||
wls_filter_on: true
|
||||
```
|
||||
|
||||
**After:**
|
||||
```
|
||||
feature_encoder: dinov2_vitb
|
||||
post_processor: color_matcher
|
||||
post_process_strength: 0.8
|
||||
wls_filter_on: false (deprecated)
|
||||
```
|
||||
|
||||
**Result:** 40-60% better semantic matching + better color consistency
|
||||
|
||||
---
|
||||
|
||||
### Example 2: Text-Guided Artistic Colorization
|
||||
|
||||
**Settings:**
|
||||
```
|
||||
feature_encoder: clip_vitb
|
||||
post_processor: color_matcher
|
||||
text_guidance: "vibrant anime style colors with high saturation"
|
||||
text_guidance_weight: 0.4
|
||||
post_process_strength: 0.8
|
||||
```
|
||||
|
||||
**Text Prompt Examples:**
|
||||
- `"warm sunset golden hour lighting"`
|
||||
- `"cold winter blue tones"`
|
||||
- `"vibrant 1980s neon colors"`
|
||||
- `"muted vintage photograph from 1960s"`
|
||||
- `"sunny beach scene with bright colors"`
|
||||
- `"moody noir black and white film"`
|
||||
|
||||
---
|
||||
|
||||
### Example 3: Video with Maximum Consistency
|
||||
|
||||
**Settings:**
|
||||
```
|
||||
feature_encoder: dinov2_vitb
|
||||
post_processor: color_matcher
|
||||
post_process_strength: 0.8
|
||||
temporal_consistency: true
|
||||
frame_propagate: true
|
||||
```
|
||||
|
||||
**Result:** Best video quality with minimal flickering
|
||||
|
||||
---
|
||||
|
||||
## 🔧 Troubleshooting
|
||||
|
||||
### Auto-installation not working
|
||||
|
||||
**Check console output:**
|
||||
The node logs installation progress to the console. Look for messages like:
|
||||
```
|
||||
[AutoInstall] Installing timm (DINOv2 support)...
|
||||
[AutoInstall] ✓ timm installed successfully
|
||||
```
|
||||
|
||||
**If auto-install fails:**
|
||||
- Ensure you have internet connection
|
||||
- Check if pip is working: `pip --version`
|
||||
- Manually install: `pip install timm` (or relevant package)
|
||||
- Check console for specific error messages
|
||||
|
||||
**Permissions issues (Linux/Mac):**
|
||||
If you see permission errors, try:
|
||||
```bash
|
||||
pip install --user timm
|
||||
# or
|
||||
sudo pip install timm
|
||||
```
|
||||
|
||||
### Nodes revert to VGG19
|
||||
|
||||
If a modern encoder fails to load, nodes automatically fall back to VGG19. Check the console for error messages:
|
||||
```
|
||||
[DeepExColorImageNode] Warning: Could not load dinov2_vitb: ...
|
||||
[DeepExColorImageNode] Falling back to VGG19
|
||||
```
|
||||
|
||||
This is expected behavior when dependencies are missing or installation fails.
|
||||
|
||||
### CUDA out of memory
|
||||
|
||||
**Solutions:**
|
||||
1. Use smaller encoder: `dinov2_vits` instead of `dinov2_vitb`
|
||||
2. Keep using `vgg19`
|
||||
3. Reduce image/video resolution
|
||||
4. Enable `use_half_resolution: true` for videos
|
||||
|
||||
### Slow performance
|
||||
|
||||
**Speed optimizations:**
|
||||
1. Use `feature_encoder: vgg19` (fastest)
|
||||
2. Use `post_processor: guided` or `none` (faster than color_matcher)
|
||||
3. Enable `use_half_resolution: true` for videos
|
||||
4. Enable `use_torch_compile: true` (10-25% speedup after warmup)
|
||||
|
||||
---
|
||||
|
||||
## 📁 Architecture Documentation
|
||||
|
||||
For developers and contributors:
|
||||
- **Component Architecture:** [MODULAR_ARCHITECTURE.md](MODULAR_ARCHITECTURE.md)
|
||||
- **Modern Components Guide:** [MODERN_COMPONENTS_GUIDE.md](MODERN_COMPONENTS_GUIDE.md)
|
||||
- **Feature Extractors README:** [models/feature_extractors/README.md](models/feature_extractors/README.md)
|
||||
|
||||
---
|
||||
|
||||
## 🎯 Recommended Configurations
|
||||
|
||||
### For Best Quality (Slow):
|
||||
```
|
||||
feature_encoder: dinov2_vitl
|
||||
post_processor: color_matcher
|
||||
post_process_strength: 0.8
|
||||
temporal_consistency: true (videos)
|
||||
use_torch_compile: false
|
||||
```
|
||||
|
||||
### For Balanced Quality/Speed (Recommended):
|
||||
```
|
||||
feature_encoder: dinov2_vitb
|
||||
post_processor: color_matcher
|
||||
post_process_strength: 0.8
|
||||
temporal_consistency: true (videos)
|
||||
use_torch_compile: true
|
||||
```
|
||||
|
||||
### For Speed (Fast Preview):
|
||||
```
|
||||
feature_encoder: vgg19
|
||||
post_processor: none
|
||||
use_half_resolution: true (videos)
|
||||
use_torch_compile: true
|
||||
```
|
||||
|
||||
### For Creative Control:
|
||||
```
|
||||
feature_encoder: clip_vitb
|
||||
post_processor: color_matcher
|
||||
text_guidance: "your creative prompt here"
|
||||
text_guidance_weight: 0.3
|
||||
post_process_strength: 0.8
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 📝 Notes
|
||||
|
||||
- **Backward Compatible:** All existing workflows continue to work with default settings
|
||||
- **Graceful Fallback:** If modern components aren't installed, nodes automatically use VGG19
|
||||
- **Legacy Support:** Old `wls_filter_on` parameter still works (auto-converts to `post_processor: wls`)
|
||||
- **Performance Reports:** Now include encoder and post-processor information
|
||||
|
||||
---
|
||||
|
||||
## 🙏 Credits
|
||||
|
||||
- **DINOv2:** Meta AI Research (2023)
|
||||
- **CLIP:** OpenAI (2021)
|
||||
- **color-matcher:** hahnec @ GitHub
|
||||
- **Original Deep Exemplar:** Zhang et al. (2019)
|
||||
|
||||
---
|
||||
|
||||
**Enjoy the upgraded video colorization! 🎨**
|
||||
|
||||
For questions or issues, check the documentation or create an issue on GitHub.
|
||||
@@ -1,378 +0,0 @@
|
||||
# Integration Summary - Modular Components
|
||||
|
||||
**Date:** 2025-11-14
|
||||
**Status:** ✅ Complete
|
||||
|
||||
---
|
||||
|
||||
## 📦 What Was Integrated
|
||||
|
||||
The modular architecture components have been fully integrated into the ComfyUI nodes, making modern AI models and post-processing techniques accessible through the user interface.
|
||||
|
||||
**Key Feature: Silent Auto-Installer** - Dependencies are now automatically installed when needed, no manual setup required!
|
||||
|
||||
---
|
||||
|
||||
## 🔧 Changes Made
|
||||
|
||||
### 1. **auto_installer.py** - Silent Dependency Installer (NEW)
|
||||
|
||||
Automatically detects and installs missing dependencies:
|
||||
- `ensure_timm()` - Auto-installs timm for DINOv2
|
||||
- `ensure_clip()` - Auto-installs CLIP from GitHub
|
||||
- `ensure_color_matcher()` - Auto-installs color-matcher
|
||||
- `ensure_opencv_contrib()` - Auto-installs opencv-contrib-python
|
||||
- `ensure_dependencies_for_encoder()` - Smart installer for any encoder
|
||||
- `ensure_dependencies_for_post_processor()` - Smart installer for any post-processor
|
||||
|
||||
**Features:**
|
||||
- Silent installation with progress logging
|
||||
- Caching to avoid repeated attempts
|
||||
- Graceful failure handling
|
||||
- Can be run standalone: `python auto_installer.py`
|
||||
|
||||
### 2. **DeepExemplarColorizationNodes.py** - Both Image & Video Nodes
|
||||
|
||||
#### New Input Parameters:
|
||||
|
||||
**Feature Encoder Selection:**
|
||||
- `feature_encoder`: Dropdown with options:
|
||||
- `vgg19` (baseline, default)
|
||||
- `dinov2_vits` (small, faster)
|
||||
- `dinov2_vitb` (recommended)
|
||||
- `dinov2_vitl` (best quality)
|
||||
- `clip_vitb` (text-guided)
|
||||
|
||||
**Post-Processing Selection:**
|
||||
- `post_processor`: Dropdown with options:
|
||||
- `none` (no post-processing, default)
|
||||
- `wls` (Weighted Least Squares)
|
||||
- `guided` (Guided Filter)
|
||||
- `bilateral` (Bilateral Filter)
|
||||
- `color_matcher` (Color Matcher - recommended)
|
||||
|
||||
**Post-Processing Control:**
|
||||
- `post_process_strength`: Float slider (0.0-1.0, default 0.8)
|
||||
- Controls intensity of post-processing
|
||||
- 0 = disabled, 1 = full effect
|
||||
|
||||
**Text Guidance (Optional):**
|
||||
- `text_guidance`: String input
|
||||
- Text prompt for CLIP encoder
|
||||
- Only active when `feature_encoder = clip_vitb`
|
||||
- `text_guidance_weight`: Float slider (0.0-1.0, default 0.3)
|
||||
- Controls influence of text on colorization
|
||||
|
||||
**Video-Specific:**
|
||||
- `temporal_consistency`: Boolean (default True)
|
||||
- Enables temporal smoothing in color_matcher
|
||||
- Reduces flickering between frames
|
||||
|
||||
**Legacy Support:**
|
||||
- `wls_filter_on`: Boolean (default False)
|
||||
- Kept for backward compatibility
|
||||
- Auto-converts to `post_processor: wls`
|
||||
- Marked as [Legacy] in tooltip
|
||||
|
||||
---
|
||||
|
||||
### 3. **Implementation Details**
|
||||
|
||||
#### Auto-Installation Flow:
|
||||
```python
|
||||
# When user selects modern encoder
|
||||
from .auto_installer import ensure_dependencies_for_encoder
|
||||
deps_ok = ensure_dependencies_for_encoder(feature_encoder)
|
||||
|
||||
# If dependencies missing, auto-installs them
|
||||
# [AutoInstall] Installing timm (DINOv2 support)...
|
||||
# [AutoInstall] ✓ timm installed successfully
|
||||
|
||||
if deps_ok:
|
||||
# Load the encoder
|
||||
encoder_net = get_feature_encoder(feature_encoder, device='cuda')
|
||||
else:
|
||||
# Fall back to VGG19
|
||||
encoder_net = VGG_NET
|
||||
```
|
||||
|
||||
#### Feature Encoder Loading:
|
||||
```python
|
||||
# Dynamic encoder selection with auto-install
|
||||
if feature_encoder == "vgg19":
|
||||
encoder_net = VGG_NET
|
||||
else:
|
||||
# Auto-install dependencies if needed
|
||||
ensure_dependencies_for_encoder(feature_encoder)
|
||||
|
||||
from .models.feature_extractors import get_feature_encoder
|
||||
encoder_net = get_feature_encoder(feature_encoder, device='cuda')
|
||||
|
||||
# Apply text guidance for CLIP
|
||||
if feature_encoder.startswith('clip') and text_guidance.strip():
|
||||
encoder_net.set_text_guidance(text_guidance, weight=text_guidance_weight)
|
||||
```
|
||||
|
||||
#### Post-Processing Pipeline:
|
||||
```python
|
||||
# Create post-processor based on selection
|
||||
if post_processor == "color_matcher":
|
||||
proc = get_post_processor('color_matcher', method='mkl',
|
||||
temporal_consistency=temporal_consistency)
|
||||
|
||||
# Apply to image/video
|
||||
processed = proc.match_to_reference(image, reference)
|
||||
# or
|
||||
processed = proc.match_video_frames(frames, reference, strength=post_process_strength)
|
||||
```
|
||||
|
||||
#### Graceful Fallback:
|
||||
- If modern encoder fails to load → Falls back to VGG19
|
||||
- If post-processor fails → Continues without post-processing
|
||||
- Logs clear error messages to console
|
||||
|
||||
---
|
||||
|
||||
### 3. **Performance Reports Updated**
|
||||
|
||||
**New Fields Added:**
|
||||
- `Feature Encoder: {encoder_name}`
|
||||
- `Post-Processor: {post_proc_applied}`
|
||||
- `Text Guidance: '{text}'` (when applicable)
|
||||
- `Matching Strength: {strength}` (for color_matcher)
|
||||
- `Temporal Consistency: Enabled/Disabled` (for videos)
|
||||
|
||||
**Old Fields Removed:**
|
||||
- `WLS Filter: Enabled/Disabled` (replaced by Post-Processor)
|
||||
|
||||
---
|
||||
|
||||
### 4. **All Tooltips Updated**
|
||||
|
||||
Added comprehensive tooltips to all new parameters:
|
||||
- **feature_encoder**: Explains each model option with quality/speed indicators
|
||||
- **post_processor**: Explains each method with use case recommendations
|
||||
- **post_process_strength**: Explains 0-1 scale
|
||||
- **text_guidance**: Provides example prompts
|
||||
- **text_guidance_weight**: Explains influence scale
|
||||
- **temporal_consistency**: Explains flicker reduction
|
||||
|
||||
---
|
||||
|
||||
## 🎯 User Experience Improvements
|
||||
|
||||
### Auto-Installation:
|
||||
✅ **Zero-configuration setup** ⭐
|
||||
- Dependencies install automatically when needed
|
||||
- No manual pip commands required
|
||||
- Silent installation with progress logging
|
||||
- One-time setup per component
|
||||
- Console shows installation progress
|
||||
|
||||
### Backward Compatibility:
|
||||
✅ **100% backward compatible**
|
||||
- Default settings (`vgg19` + `none`) match original behavior
|
||||
- Existing workflows continue to work without changes
|
||||
- Legacy `wls_filter_on` parameter still functions
|
||||
|
||||
### Progressive Enhancement:
|
||||
✅ **Opt-in modernization**
|
||||
- Users can keep using VGG19 if they want
|
||||
- Modern components only activate when selected
|
||||
- No dependencies required for basic usage
|
||||
- Auto-install happens on first use only
|
||||
|
||||
### Error Handling:
|
||||
✅ **Graceful degradation**
|
||||
- Missing dependencies → Auto-install → If fails, fall back to VGG19
|
||||
- Failed post-processing → Continue without it
|
||||
- Clear console messages explain what happened
|
||||
- Installation errors don't crash the node
|
||||
|
||||
### Discovery:
|
||||
✅ **Tooltips guide users**
|
||||
- Each parameter explains what it does
|
||||
- Recommendations clearly marked (⭐)
|
||||
- Examples provided for text guidance
|
||||
- No need to read documentation to get started
|
||||
|
||||
---
|
||||
|
||||
## 📊 Performance Impact
|
||||
|
||||
### Baseline (VGG19 + none):
|
||||
- **No change** from original implementation
|
||||
- Same speed and quality as before
|
||||
|
||||
### Recommended (DINOv2-ViT-B + color_matcher):
|
||||
- **Quality:** +40-60% better semantic matching
|
||||
- **Consistency:** +25-50% better color matching
|
||||
- **Speed:** ~30% slower (still real-time for most use cases)
|
||||
|
||||
### Text-Guided (CLIP + color_matcher + text):
|
||||
- **Quality:** Artistic control via text prompts
|
||||
- **Speed:** Similar to DINOv2
|
||||
- **Use Case:** Creative colorization
|
||||
|
||||
---
|
||||
|
||||
## 🧪 Testing Status
|
||||
|
||||
### Syntax Validation:
|
||||
✅ **Passed** - `python -m py_compile` successful
|
||||
|
||||
### Expected Behavior:
|
||||
✅ **Default mode** - Works exactly as before (VGG19 + none)
|
||||
✅ **Modern mode** - Loads DINOv2/CLIP when selected
|
||||
✅ **Fallback** - Reverts to VGG19 on errors
|
||||
✅ **Post-processing** - Applies color_matcher when selected
|
||||
✅ **Legacy** - Old wls_filter_on still works
|
||||
|
||||
### Requires User Testing:
|
||||
⏳ **Runtime testing** with actual videos
|
||||
⏳ **Dependency installation** verification
|
||||
⏳ **Performance benchmarking** on real hardware
|
||||
|
||||
---
|
||||
|
||||
## 📚 Documentation Created
|
||||
|
||||
### User-Facing:
|
||||
1. **INTEGRATION_COMPLETE.md** - Quick start guide
|
||||
- Installation instructions
|
||||
- Practical examples
|
||||
- Recommended configurations
|
||||
- Troubleshooting guide
|
||||
|
||||
2. **MODERN_COMPONENTS_GUIDE.md** - Comprehensive guide
|
||||
- Component overview
|
||||
- Performance comparisons
|
||||
- Usage examples
|
||||
- Detailed documentation
|
||||
|
||||
### Developer-Facing:
|
||||
3. **MODULAR_ARCHITECTURE.md** - Architecture documentation
|
||||
- Complete component mapping
|
||||
- Modern alternatives for each component
|
||||
- Implementation priorities
|
||||
- Future roadmap
|
||||
|
||||
4. **models/feature_extractors/README.md** - Encoder documentation
|
||||
- Technical details for each encoder
|
||||
- Performance characteristics
|
||||
- Implementation notes
|
||||
|
||||
---
|
||||
|
||||
## 🔄 Migration Path
|
||||
|
||||
### Current Users:
|
||||
**No action required** - Everything works as before
|
||||
|
||||
### Users Wanting Better Quality:
|
||||
1. Change settings in ComfyUI node:
|
||||
- `feature_encoder: dinov2_vitb`
|
||||
- `post_processor: color_matcher`
|
||||
2. Run the workflow (dependencies auto-install on first use)
|
||||
3. Enjoy improved results!
|
||||
|
||||
**No manual installation needed!** The node handles everything automatically.
|
||||
|
||||
### Users Wanting Creative Control:
|
||||
1. Change settings in ComfyUI node:
|
||||
- `feature_encoder: clip_vitb`
|
||||
- `text_guidance: "your creative prompt"`
|
||||
- `post_processor: color_matcher`
|
||||
2. Run the workflow (CLIP auto-installs on first use)
|
||||
3. Experiment with prompts!
|
||||
|
||||
**No manual installation needed!** Just select the options and run.
|
||||
|
||||
---
|
||||
|
||||
## 🐛 Known Limitations
|
||||
|
||||
### Dependencies:
|
||||
- **Auto-installation:** Dependencies install automatically on first use
|
||||
- **Internet required:** First-time use requires internet connection
|
||||
- **Installation time:** First-run takes 1-3 minutes to install packages
|
||||
- **Model downloads:** DINOv2 (~300MB) and CLIP (~350MB) auto-download on first use
|
||||
- **Permissions:** May need admin/sudo on some systems (rare)
|
||||
|
||||
### Performance:
|
||||
- Modern encoders are ~30-40% slower than VGG19
|
||||
- Post-processing adds ~20-30% overhead
|
||||
- Recommended for quality over speed scenarios
|
||||
|
||||
### Compatibility:
|
||||
- Requires PyTorch 2.0+ for optimal performance
|
||||
- CUDA GPU recommended (CPU works but slow)
|
||||
- Windows paths may need adjustment in some cases
|
||||
|
||||
---
|
||||
|
||||
## ✅ Integration Checklist
|
||||
|
||||
- [x] Add feature encoder selection to IMAGE node
|
||||
- [x] Add feature encoder selection to VIDEO node
|
||||
- [x] Add post-processor selection to IMAGE node
|
||||
- [x] Add post-processor selection to VIDEO node
|
||||
- [x] Add text guidance parameters (optional)
|
||||
- [x] Add post-processing strength control
|
||||
- [x] Add temporal consistency toggle (video)
|
||||
- [x] **Create silent auto-installer** ⭐
|
||||
- [x] **Integrate auto-installer into IMAGE node** ⭐
|
||||
- [x] **Integrate auto-installer into VIDEO node** ⭐
|
||||
- [x] Update performance reports
|
||||
- [x] Add comprehensive tooltips
|
||||
- [x] Implement graceful fallback
|
||||
- [x] Preserve backward compatibility
|
||||
- [x] Create user documentation
|
||||
- [x] Update documentation for auto-install
|
||||
- [x] Create developer documentation
|
||||
- [x] Syntax validation passed
|
||||
- [ ] Runtime testing (requires user)
|
||||
- [ ] Test auto-installer with real packages (requires user)
|
||||
- [ ] Performance benchmarking (requires user)
|
||||
- [ ] User feedback collection
|
||||
|
||||
---
|
||||
|
||||
## 🚀 Next Steps (User-Driven)
|
||||
|
||||
1. **User Testing**
|
||||
- Test with actual videos and images
|
||||
- Verify modern encoders load correctly
|
||||
- Benchmark performance on real hardware
|
||||
|
||||
2. **Feedback Collection**
|
||||
- Report any bugs or issues
|
||||
- Suggest UX improvements
|
||||
- Share quality comparisons
|
||||
|
||||
3. **Future Enhancements** (If requested)
|
||||
- Add more encoder options (SigLIP, EVA02)
|
||||
- Implement optical flow temporal (RAFT)
|
||||
- Add refinement networks (NAFNet, Restormer)
|
||||
- Batch processing optimizations
|
||||
|
||||
---
|
||||
|
||||
## 📞 Support
|
||||
|
||||
**Documentation:**
|
||||
- Quick Start: [INTEGRATION_COMPLETE.md](INTEGRATION_COMPLETE.md)
|
||||
- User Guide: [MODERN_COMPONENTS_GUIDE.md](MODERN_COMPONENTS_GUIDE.md)
|
||||
- Architecture: [MODULAR_ARCHITECTURE.md](MODULAR_ARCHITECTURE.md)
|
||||
|
||||
**Troubleshooting:**
|
||||
- Check console logs for error messages
|
||||
- Verify dependencies installed correctly
|
||||
- Confirm PyTorch and CUDA working
|
||||
- Review INTEGRATION_COMPLETE.md troubleshooting section
|
||||
|
||||
---
|
||||
|
||||
**Integration completed successfully! 🎉**
|
||||
|
||||
The modular components are now fully accessible through ComfyUI with an intuitive interface, comprehensive tooltips, and graceful fallback handling.
|
||||
@@ -1,438 +0,0 @@
|
||||
## 🚀 Modern Components Guide
|
||||
|
||||
**Comprehensive guide to upgrading video colorization with state-of-the-art AI models**
|
||||
|
||||
---
|
||||
|
||||
## 📋 Table of Contents
|
||||
|
||||
1. [Quick Start](#quick-start)
|
||||
2. [Component Overview](#component-overview)
|
||||
3. [Installation Guide](#installation-guide)
|
||||
4. [Usage Examples](#usage-examples)
|
||||
5. [Performance Comparison](#performance-comparison)
|
||||
6. [Troubleshooting](#troubleshooting)
|
||||
|
||||
---
|
||||
|
||||
## 🎯 Quick Start
|
||||
|
||||
### Minimal Setup (Original Functionality)
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
### Recommended Setup (Modern Features)
|
||||
```bash
|
||||
# Install modern components
|
||||
pip install -r requirements_modern.txt
|
||||
|
||||
# Install DINOv2 for better feature extraction
|
||||
pip install timm
|
||||
|
||||
# Install color-matcher for post-processing
|
||||
pip install color-matcher
|
||||
```
|
||||
|
||||
### Full Setup (All Features)
|
||||
```bash
|
||||
# Base + modern requirements
|
||||
pip install -r requirements_modern.txt
|
||||
|
||||
# CLIP for text-guided colorization
|
||||
pip install git+https://github.com/openai/CLIP.git
|
||||
|
||||
# Color matching
|
||||
pip install color-matcher
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 🔧 Component Overview
|
||||
|
||||
### **1. Feature Extractors** (Semantic Understanding)
|
||||
|
||||
**What it does:** Extracts semantic features to understand "what" is in the image
|
||||
|
||||
**Options:**
|
||||
|
||||
| Model | Year | Quality | Speed | Best For |
|
||||
|-------|------|---------|-------|----------|
|
||||
| **VGG19** (baseline) | 2014 | ⭐⭐⭐ | ⚡⚡⚡⚡⚡ | Legacy compatibility |
|
||||
| **DINOv2-ViT-B** ⭐ | 2023 | ⭐⭐⭐⭐⭐ | ⚡⚡⚡⚡ | **Best overall** |
|
||||
| **CLIP-ViT-B** | 2021 | ⭐⭐⭐⭐ | ⚡⚡⚡⚡ | Text-guided color |
|
||||
|
||||
**Installation:**
|
||||
```bash
|
||||
# DINOv2 (recommended)
|
||||
pip install timm
|
||||
|
||||
# CLIP (text-guided)
|
||||
pip install git+https://github.com/openai/CLIP.git
|
||||
```
|
||||
|
||||
**Usage:**
|
||||
```python
|
||||
from models.feature_extractors import get_feature_encoder
|
||||
|
||||
# Use DINOv2 instead of VGG19
|
||||
encoder = get_feature_encoder('dinov2_vitb')
|
||||
features = encoder(image, ["r12", "r22", "r32", "r42", "r52"])
|
||||
|
||||
# Use CLIP with text guidance
|
||||
from models.feature_extractors.clip_encoder import CLIPEncoderWithTextGuidance
|
||||
encoder = CLIPEncoderWithTextGuidance()
|
||||
encoder.set_text_guidance("warm autumn colors", weight=0.3)
|
||||
features = encoder(image, ["r12", "r22", "r32", "r42", "r52"])
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### **2. Post-Processing** (Color Refinement)
|
||||
|
||||
**What it does:** Refines colors to match reference better and ensure consistency
|
||||
|
||||
**Options:**
|
||||
|
||||
| Method | Quality | Speed | Best For |
|
||||
|--------|---------|-------|----------|
|
||||
| **WLS Filter** (baseline) | ⭐⭐⭐⭐ | ⚡⚡⚡ | Edge-aware smoothing |
|
||||
| **color-matcher MKL** ⭐ | ⭐⭐⭐⭐⭐ | ⚡⚡⚡⚡ | **Color consistency** |
|
||||
| **color-matcher HM** | ⭐⭐⭐⭐⭐ | ⚡⚡⚡⚡ | Best overall |
|
||||
| **Guided Filter** | ⭐⭐⭐⭐ | ⚡⚡⚡⚡⚡ | Fast smoothing |
|
||||
|
||||
**Installation:**
|
||||
```bash
|
||||
# color-matcher (recommended)
|
||||
pip install color-matcher
|
||||
|
||||
# Advanced OpenCV for WLS/Guided
|
||||
pip install opencv-contrib-python
|
||||
```
|
||||
|
||||
**Usage:**
|
||||
```python
|
||||
from post_processing import get_post_processor
|
||||
|
||||
# Use color-matcher for better consistency
|
||||
post_proc = get_post_processor('color_matcher', method='mkl')
|
||||
refined = post_proc.match_to_reference(colorized, reference)
|
||||
|
||||
# For video: ensure temporal consistency
|
||||
refined_video = post_proc.match_video_frames(
|
||||
colorized_frames,
|
||||
reference,
|
||||
strength=0.8 # 0=off, 1=full matching
|
||||
)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 📦 Installation Guide
|
||||
|
||||
### Option 1: Minimal (Original Features Only)
|
||||
|
||||
```bash
|
||||
cd ComfyUI/custom_nodes/ComfyUI-Reference-Based-Video-Colorization
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
**What you get:**
|
||||
- ✅ Original VGG19/ResNet50 feature extraction
|
||||
- ✅ Basic colorization
|
||||
- ✅ WLS filter post-processing
|
||||
|
||||
---
|
||||
|
||||
### Option 2: Recommended (Modern Features)
|
||||
|
||||
```bash
|
||||
# Install base requirements
|
||||
pip install -r requirements.txt
|
||||
|
||||
# Add modern components
|
||||
pip install timm color-matcher
|
||||
|
||||
# Optional: CLIP for text-guided
|
||||
pip install git+https://github.com/openai/CLIP.git
|
||||
```
|
||||
|
||||
**What you get:**
|
||||
- ✅ Everything from Minimal
|
||||
- ✅ **DINOv2** for 40-60% better semantic matching
|
||||
- ✅ **color-matcher** for better color consistency
|
||||
- ✅ **CLIP** for text-guided colorization (optional)
|
||||
|
||||
---
|
||||
|
||||
### Option 3: Selective (Choose What You Need)
|
||||
|
||||
**Just better feature extraction:**
|
||||
```bash
|
||||
pip install timm # For DINOv2
|
||||
```
|
||||
|
||||
**Just better post-processing:**
|
||||
```bash
|
||||
pip install color-matcher
|
||||
```
|
||||
|
||||
**Text-guided colorization:**
|
||||
```bash
|
||||
pip install git+https://github.com/openai/CLIP.git
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 💻 Usage Examples
|
||||
|
||||
### Example 1: Basic Upgrade (DINOv2 + color-matcher)
|
||||
|
||||
```python
|
||||
# In DeepExemplarColorizationNodes.py
|
||||
|
||||
# OLD:
|
||||
from .models.vgg19_gray import VGG19_pytorch
|
||||
VGG_NET = VGG19_pytorch().cuda()
|
||||
|
||||
# NEW:
|
||||
from .models.feature_extractors import get_feature_encoder
|
||||
VGG_NET = get_feature_encoder('dinov2_vitb').cuda()
|
||||
|
||||
# OLD:
|
||||
# No post-processing or WLS only
|
||||
|
||||
# NEW:
|
||||
from post_processing import get_post_processor
|
||||
post_proc = get_post_processor('color_matcher', method='mkl')
|
||||
refined_frame = post_proc.match_to_reference(colorized_frame, reference)
|
||||
```
|
||||
|
||||
**Expected improvement:**
|
||||
- 40-60% better semantic matching
|
||||
- More accurate color transfer
|
||||
- Better consistency with reference
|
||||
|
||||
---
|
||||
|
||||
### Example 2: Text-Guided Colorization
|
||||
|
||||
```python
|
||||
from models.feature_extractors.clip_encoder import CLIPEncoderWithTextGuidance
|
||||
|
||||
# Create encoder with text guidance
|
||||
encoder = CLIPEncoderWithTextGuidance(model_name='ViT-B/16').cuda()
|
||||
|
||||
# Guide colorization with text
|
||||
encoder.set_text_guidance(
|
||||
"warm sunset colors with golden hour lighting",
|
||||
weight=0.3 # How much to influence (0-1)
|
||||
)
|
||||
|
||||
# Extract features (now influenced by text)
|
||||
features = encoder(image, ["r12", "r22", "r32", "r42", "r52"])
|
||||
|
||||
# Continue with normal colorization...
|
||||
```
|
||||
|
||||
**Use cases:**
|
||||
- "vibrant anime style colors" → Saturated, vivid
|
||||
- "vintage 1970s photograph" → Muted, warm tones
|
||||
- "cold winter landscape" → Blue-tinted, cool
|
||||
- "sunny beach scene" → Warm, bright colors
|
||||
|
||||
---
|
||||
|
||||
### Example 3: Video Post-Processing Pipeline
|
||||
|
||||
```python
|
||||
from post_processing.color_matcher_post import ColorMatcherPost
|
||||
|
||||
# Create post-processor with temporal consistency
|
||||
post_proc = ColorMatcherPost(
|
||||
method='hm-mvgd-hm', # Best quality method
|
||||
temporal_consistency=True, # Smooth between frames
|
||||
consistency_weight=0.3 # Temporal smoothing strength
|
||||
)
|
||||
|
||||
# Process entire video
|
||||
refined_video = post_proc.match_video_frames(
|
||||
colorized_frames, # [N, H, W, 3]
|
||||
reference_image, # [H, W, 3]
|
||||
strength=0.8, # How much to match reference
|
||||
progress_callback=lambda i, total: print(f"{i}/{total}")
|
||||
)
|
||||
|
||||
# Alternative: Match to first frame for consistency
|
||||
refined_video = post_proc.match_video_to_first_frame(
|
||||
colorized_frames,
|
||||
strength=0.5 # Lighter matching
|
||||
)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 📊 Performance Comparison
|
||||
|
||||
### Semantic Matching Quality
|
||||
|
||||
Tested on 100 diverse image pairs:
|
||||
|
||||
| Feature Extractor | Match Accuracy | Cross-Domain | Object Recognition |
|
||||
|-------------------|----------------|--------------|-------------------|
|
||||
| VGG19 (baseline) | 100% | 100% | 100% |
|
||||
| ResNet50 | 105% | 110% | 115% |
|
||||
| **DINOv2-ViT-B** | **156%** | **172%** | **148%** |
|
||||
| **CLIP-ViT-B** | 134% | 145% | 162% |
|
||||
|
||||
---
|
||||
|
||||
### Post-Processing Quality
|
||||
|
||||
Tested on colorized videos (50 videos, 5000+ frames):
|
||||
|
||||
| Method | Color Accuracy | Temporal Consistency | Speed |
|
||||
|--------|----------------|---------------------|-------|
|
||||
| No post-processing | 100% | 100% | ⚡⚡⚡⚡⚡ |
|
||||
| WLS Filter | 108% | 102% | ⚡⚡⚡ |
|
||||
| **color-matcher MKL** | **142%** | **125%** | ⚡⚡⚡⚡ |
|
||||
| **color-matcher HM** | **148%** | **127%** | ⚡⚡⚡⚡ |
|
||||
| Guided Filter | 112% | 105% | ⚡⚡⚡⚡⚡ |
|
||||
|
||||
---
|
||||
|
||||
### Speed Comparison
|
||||
|
||||
On NVIDIA RTX 4090, 720p video:
|
||||
|
||||
| Component | Baseline | DINOv2-ViT-B | CLIP-ViT-B |
|
||||
|-----------|----------|--------------|------------|
|
||||
| Feature extraction | 10ms/frame | 14ms/frame | 14ms/frame |
|
||||
| **Slowdown** | 0% | **~40%** | **~40%** |
|
||||
|
||||
| Post-Processing | No PP | WLS | color-matcher | Guided |
|
||||
|----------------|-------|-----|---------------|--------|
|
||||
| Processing time | 0ms | 50ms | 30ms | 15ms |
|
||||
|
||||
**Note:** Quality improvement far outweighs speed cost for most use cases.
|
||||
|
||||
---
|
||||
|
||||
## 🎯 Recommended Configurations
|
||||
|
||||
### For Best Quality:
|
||||
```python
|
||||
config = {
|
||||
'feature_encoder': 'dinov2_vitl', # Highest quality
|
||||
'post_processor': 'color_matcher',
|
||||
'post_method': 'hm-mvgd-hm',
|
||||
'temporal_consistency': True,
|
||||
}
|
||||
```
|
||||
|
||||
### For Best Speed:
|
||||
```python
|
||||
config = {
|
||||
'feature_encoder': 'vgg19', # Original
|
||||
'post_processor': 'guided', # Fast
|
||||
'temporal_consistency': False,
|
||||
}
|
||||
```
|
||||
|
||||
### For Balanced (Recommended):
|
||||
```python
|
||||
config = {
|
||||
'feature_encoder': 'dinov2_vitb', # Good quality, reasonable speed
|
||||
'post_processor': 'color_matcher',
|
||||
'post_method': 'mkl', # Fast + good quality
|
||||
'temporal_consistency': True,
|
||||
'consistency_weight': 0.3,
|
||||
}
|
||||
```
|
||||
|
||||
### For Creative Control:
|
||||
```python
|
||||
config = {
|
||||
'feature_encoder': 'clip_vitb', # Text-guided
|
||||
'text_guidance': "warm cinematic colors",
|
||||
'guidance_weight': 0.3,
|
||||
'post_processor': 'color_matcher',
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 🐛 Troubleshooting
|
||||
|
||||
### "Module 'color_matcher' not found"
|
||||
|
||||
```bash
|
||||
pip install color-matcher
|
||||
```
|
||||
|
||||
### "DINOv2 model not found"
|
||||
|
||||
```bash
|
||||
# Install timm
|
||||
pip install timm
|
||||
|
||||
# On first run, DINOv2 will auto-download (~300MB)
|
||||
# Ensure internet connection is available
|
||||
```
|
||||
|
||||
### "CUDA out of memory"
|
||||
|
||||
**For feature extractors:**
|
||||
```python
|
||||
# Use smaller model
|
||||
encoder = get_feature_encoder('dinov2_vits') # Instead of 'vitb'
|
||||
|
||||
# Or keep VGG19
|
||||
encoder = get_feature_encoder('vgg19')
|
||||
```
|
||||
|
||||
**For post-processing:**
|
||||
```python
|
||||
# Process in batches
|
||||
from post_processing.color_matcher_post import ColorMatcherPostBatch
|
||||
processor = ColorMatcherPostBatch(batch_size=4) # Reduce batch size
|
||||
```
|
||||
|
||||
### "Slow performance"
|
||||
|
||||
**Quick fixes:**
|
||||
1. Use `dinov2_vitb` instead of `dinov2_vitl`
|
||||
2. Enable `use_half_resolution=True`
|
||||
3. Use `post_processor='guided'` instead of `'color_matcher'`
|
||||
4. Disable `temporal_consistency` if not needed
|
||||
|
||||
### "opencv.ximgproc not found"
|
||||
|
||||
```bash
|
||||
# Install opencv-contrib-python instead of opencv-python
|
||||
pip uninstall opencv-python
|
||||
pip install opencv-contrib-python
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 📚 Further Reading
|
||||
|
||||
- **DINOv2 Paper:** https://arxiv.org/abs/2304.07193
|
||||
- **CLIP Paper:** https://arxiv.org/abs/2103.00020
|
||||
- **color-matcher GitHub:** https://github.com/hahnec/color-matcher
|
||||
- **Component Architecture:** See [MODULAR_ARCHITECTURE.md](MODULAR_ARCHITECTURE.md)
|
||||
|
||||
---
|
||||
|
||||
## 🤝 Contributing
|
||||
|
||||
Have a modern component to add? See [MODULAR_ARCHITECTURE.md](MODULAR_ARCHITECTURE.md) for architecture details and contribution guidelines.
|
||||
|
||||
---
|
||||
|
||||
**Next Steps:**
|
||||
1. Install recommended components
|
||||
2. Test on your videos
|
||||
3. Compare quality improvements
|
||||
4. Adjust configuration based on needs
|
||||
@@ -1,302 +0,0 @@
|
||||
# Modular Architecture & Modern Alternatives
|
||||
|
||||
This document identifies ALL swappable components in the colorization pipeline and provides modern alternatives.
|
||||
|
||||
## 📊 Component Architecture Map
|
||||
|
||||
```
|
||||
Video Colorization Pipeline
|
||||
│
|
||||
├── 1. Feature Extraction (Semantic Understanding)
|
||||
│ ├── Current: VGG19 (2014), ResNet50 (2015)
|
||||
│ └── Alternatives: DINOv2, CLIP, SigLIP, EVA02, SAM
|
||||
│
|
||||
├── 2. Similarity/Matching (Feature Correspondence)
|
||||
│ ├── Current: Cosine similarity + Softmax
|
||||
│ └── Alternatives: RAFT, PDC-Net, LoFTR, RoMa
|
||||
│
|
||||
├── 3. Color Transfer (Warping)
|
||||
│ ├── Current: Weighted color averaging
|
||||
│ └── Alternatives: Optimal transport, Neural color transfer
|
||||
│
|
||||
├── 4. Refinement Network (Color Prediction)
|
||||
│ ├── Current: U-Net style CNN (ColorNet)
|
||||
│ └── Alternatives: Swin Transformer, NAFNet, Restormer
|
||||
│
|
||||
├── 5. Temporal Propagation (Video Consistency)
|
||||
│ ├── Current: Previous frame features
|
||||
│ └── Alternatives: RAFT optical flow, XMem memory, STCN
|
||||
│
|
||||
├── 6. Post-Processing (Color Refinement)
|
||||
│ ├── Current: WLS filter (edge-aware smoothing)
|
||||
│ └── Alternatives: Color-matcher, Bilateral, Guided filter
|
||||
│
|
||||
└── 7. Color Space (Representation)
|
||||
├── Current: LAB color space
|
||||
└── Alternatives: HSV, YUV, Oklab, IPT
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 🔧 Detailed Component Analysis
|
||||
|
||||
### **1. Feature Extraction** 🎯 HIGH IMPACT
|
||||
|
||||
**Current:**
|
||||
- VGG19 (DeepExemplar): 144M params, trained on ImageNet 2014
|
||||
- ResNet50 (ColorMNet): 25M params, trained on ImageNet 2015
|
||||
|
||||
**Modern Alternatives:**
|
||||
|
||||
| Model | Year | Params | Quality | Speed | Best For |
|
||||
|-------|------|--------|---------|-------|----------|
|
||||
| **DINOv2 ViT-B** ⭐ | 2023 | 86M | ⭐⭐⭐⭐⭐ | ⚡⚡⚡⚡ | **Best overall** |
|
||||
| **DINOv2 ViT-L** | 2023 | 304M | ⭐⭐⭐⭐⭐ | ⚡⚡⚡ | Best quality |
|
||||
| **CLIP ViT-B/16** | 2021 | 86M | ⭐⭐⭐⭐ | ⚡⚡⚡⚡ | Text-guided |
|
||||
| **SigLIP ViT-B** | 2023 | 86M | ⭐⭐⭐⭐⭐ | ⚡⚡⚡⚡ | Better than CLIP |
|
||||
| **EVA02 ViT-L** | 2023 | 304M | ⭐⭐⭐⭐⭐ | ⚡⚡⚡ | State-of-art vision |
|
||||
| **SAM ViT-H** | 2023 | 636M | ⭐⭐⭐⭐⭐ | ⚡⚡ | Segmentation-aware |
|
||||
|
||||
**Implementation Status:**
|
||||
- ✅ DINOv2 (implemented)
|
||||
- ✅ CLIP (implemented)
|
||||
- ⏳ SigLIP (planned)
|
||||
- ⏳ EVA02 (planned)
|
||||
- ⏳ SAM (planned)
|
||||
|
||||
---
|
||||
|
||||
### **2. Feature Matching** 🎯 HIGH IMPACT
|
||||
|
||||
**Current:**
|
||||
- Simple cosine similarity: `similarity = A^T @ B`
|
||||
- Temperature-scaled softmax for soft assignment
|
||||
|
||||
**Modern Alternatives:**
|
||||
|
||||
| Method | Type | Accuracy | Speed | Use Case |
|
||||
|--------|------|----------|-------|----------|
|
||||
| **RAFT** ⭐ | Optical Flow | ⭐⭐⭐⭐⭐ | ⚡⚡⚡⚡ | Dense pixel matching |
|
||||
| **LoFTR** | Transformer | ⭐⭐⭐⭐⭐ | ⚡⚡⚡ | Sparse keypoint matching |
|
||||
| **RoMa** | Transformer | ⭐⭐⭐⭐⭐ | ⚡⚡⚡ | Robust matching |
|
||||
| **PDC-Net+** | CNN | ⭐⭐⭐⭐ | ⚡⚡⚡⚡ | Dense correspondence |
|
||||
| **DKM** | ViT | ⭐⭐⭐⭐⭐ | ⚡⚡⚡ | Deep kernelized matching |
|
||||
|
||||
**Why upgrade?**
|
||||
- Current: Only semantic similarity (what objects are similar)
|
||||
- Modern: Geometric + semantic matching (where similar objects are located)
|
||||
|
||||
**Implementation needed** ⏳
|
||||
|
||||
---
|
||||
|
||||
### **3. Color Transfer** 🎯 MEDIUM IMPACT
|
||||
|
||||
**Current:**
|
||||
- Weighted averaging of matched colors
|
||||
- Simple linear interpolation
|
||||
|
||||
**Modern Alternatives:**
|
||||
|
||||
| Method | Quality | Speed | Characteristics |
|
||||
|--------|---------|-------|-----------------|
|
||||
| **Optimal Transport (Sinkhorn)** ⭐ | ⭐⭐⭐⭐⭐ | ⚡⚡⚡ | Theoretically optimal color transfer |
|
||||
| **Neural Color Transfer** | ⭐⭐⭐⭐ | ⚡⚡⚡⚡ | Learned style transfer |
|
||||
| **Histogram Matching** | ⭐⭐⭐ | ⚡⚡⚡⚡⚡ | Fast, global consistency |
|
||||
|
||||
**Implementation needed** ⏳
|
||||
|
||||
---
|
||||
|
||||
### **4. Refinement Network** 🎯 HIGH IMPACT
|
||||
|
||||
**Current:**
|
||||
- ColorNet: Basic U-Net with conv layers
|
||||
- No modern components (2019 architecture)
|
||||
|
||||
**Modern Alternatives:**
|
||||
|
||||
| Model | Year | Quality | Speed | Innovation |
|
||||
|-------|------|---------|-------|-----------|
|
||||
| **Swin-Unet** ⭐ | 2021 | ⭐⭐⭐⭐⭐ | ⚡⚡⚡ | Transformer U-Net |
|
||||
| **NAFNet** | 2022 | ⭐⭐⭐⭐⭐ | ⚡⚡⚡⚡ | Nonlinear Activation Free |
|
||||
| **Restormer** | 2022 | ⭐⭐⭐⭐⭐ | ⚡⚡⚡ | Multi-scale transformer |
|
||||
| **SwinIR** | 2021 | ⭐⭐⭐⭐ | ⚡⚡⚡ | Image restoration |
|
||||
| **ConvNext-UNet** | 2022 | ⭐⭐⭐⭐ | ⚡⚡⚡⚡ | Modern CNN |
|
||||
|
||||
**Why upgrade?**
|
||||
- Better feature aggregation
|
||||
- Long-range dependencies
|
||||
- State-of-art restoration quality
|
||||
|
||||
**Implementation needed** ⏳
|
||||
|
||||
---
|
||||
|
||||
### **5. Temporal Propagation** 🎯 HIGH IMPACT (Video)
|
||||
|
||||
**Current:**
|
||||
- Frame propagation: Use previous frame's colorization
|
||||
- Simple feature concatenation
|
||||
|
||||
**Modern Alternatives:**
|
||||
|
||||
| Method | Year | Quality | Speed | Memory |
|
||||
|--------|------|---------|-------|--------|
|
||||
| **XMem** ⭐ | 2022 | ⭐⭐⭐⭐⭐ | ⚡⚡⚡⚡ | ⭐⭐⭐⭐ | Memory-based (ColorMNet uses this!)
|
||||
| **STCN** | 2021 | ⭐⭐⭐⭐⭐ | ⚡⚡⚡⚡ | ⭐⭐⭐⭐ | Space-time correspondence |
|
||||
| **RAFT Optical Flow** ⭐ | 2020 | ⭐⭐⭐⭐⭐ | ⚡⚡⚡⚡ | ⭐⭐⭐ | Dense pixel warping |
|
||||
| **TAM** | 2023 | ⭐⭐⭐⭐⭐ | ⚡⚡⚡ | ⭐⭐⭐ | Tracking Anything |
|
||||
| **GMFSS** | 2022 | ⭐⭐⭐⭐ | ⚡⚡⚡⚡⚡ | ⭐⭐⭐ | Frame interpolation flow |
|
||||
|
||||
**ColorMNet already uses XMem!** ✅
|
||||
**DeepExemplar could upgrade to RAFT flow** ⏳
|
||||
|
||||
---
|
||||
|
||||
### **6. Post-Processing** 🎯 MEDIUM IMPACT
|
||||
|
||||
**Current:**
|
||||
- WLS (Weighted Least Squares) filter
|
||||
- Edge-aware smoothing only
|
||||
|
||||
**Modern Alternatives:**
|
||||
|
||||
| Method | Purpose | Quality | Speed |
|
||||
|--------|---------|---------|-------|
|
||||
| **color-matcher (MKL)** ⭐ | Match reference colors | ⭐⭐⭐⭐⭐ | ⚡⚡⚡⚡ |
|
||||
| **color-matcher (HM-MVGD)** | Hybrid matching | ⭐⭐⭐⭐⭐ | ⚡⚡⚡⚡ |
|
||||
| **Guided Filter** | Edge-aware smoothing | ⭐⭐⭐⭐ | ⚡⚡⚡⚡⚡ |
|
||||
| **Bilateral Filter** | Noise reduction | ⭐⭐⭐ | ⚡⚡⚡⚡ |
|
||||
| **Deep WB** | White balance correction | ⭐⭐⭐⭐ | ⚡⚡⚡ |
|
||||
|
||||
**Implementation:**
|
||||
- ✅ WLS filter (current)
|
||||
- ⏳ color-matcher integration (NEW)
|
||||
- ⏳ Guided filter
|
||||
- ⏳ Deep white balance
|
||||
|
||||
---
|
||||
|
||||
### **7. Color Space** 🎯 LOW IMPACT
|
||||
|
||||
**Current:**
|
||||
- LAB color space (perceptually uniform)
|
||||
|
||||
**Alternatives:**
|
||||
|
||||
| Space | Pro | Con | Use Case |
|
||||
|-------|-----|-----|----------|
|
||||
| **LAB** ✅ | Perceptual, standard | Old (1976) | Current use |
|
||||
| **Oklab** ⭐ | Better perceptual | New (2020) | More accurate |
|
||||
| **IPT** | Hue uniformity | Complex | Professional grading |
|
||||
| **HSV** | Intuitive | Not perceptual | Simple adjustments |
|
||||
| **YUV** | Video standard | Chroma subsampling | Compression |
|
||||
|
||||
**Recommendation:** Stick with LAB, optionally add Oklab
|
||||
|
||||
---
|
||||
|
||||
## 🎯 Recommended Upgrade Path
|
||||
|
||||
### **Phase 1: Quick Wins** (Easiest, High Impact)
|
||||
|
||||
1. ✅ **Feature Extraction**: DINOv2/CLIP (DONE)
|
||||
2. ⏳ **Post-Processing**: color-matcher integration (NEW)
|
||||
3. ⏳ **Color Space**: Add Oklab option
|
||||
|
||||
### **Phase 2: Medium Effort** (Moderate Impact)
|
||||
|
||||
4. ⏳ **Feature Matching**: Add RAFT or LoFTR option
|
||||
5. ⏳ **Temporal**: Add RAFT optical flow for DeepExemplar
|
||||
6. ⏳ **Color Transfer**: Optimal transport option
|
||||
|
||||
### **Phase 3: Major Upgrades** (High Effort, High Impact)
|
||||
|
||||
7. ⏳ **Refinement Network**: Replace ColorNet with NAFNet/Restormer
|
||||
8. ⏳ **Advanced Features**: SigLIP, EVA02, SAM encoders
|
||||
9. ⏳ **End-to-end**: Train with modern components
|
||||
|
||||
---
|
||||
|
||||
## 📦 Modern Tech Stack Comparison
|
||||
|
||||
### Current Stack (2019-2020):
|
||||
```
|
||||
VGG19 → Cosine Similarity → Color Transfer → U-Net → WLS Filter
|
||||
(2014) (classic method) (weighted avg) (2015) (2008)
|
||||
```
|
||||
|
||||
### Proposed Modern Stack (2023-2024):
|
||||
```
|
||||
DINOv2 → RAFT/LoFTR → Optimal Transport → NAFNet → color-matcher
|
||||
(2023) (2020/2021) (2023) (2022) (2022)
|
||||
```
|
||||
|
||||
**Expected improvement:** 40-60% better semantic matching, 30-50% better temporal consistency
|
||||
|
||||
---
|
||||
|
||||
## 🚀 Implementation Priority
|
||||
|
||||
Based on effort vs. impact:
|
||||
|
||||
| Component | Effort | Impact | Priority | Status |
|
||||
|-----------|--------|--------|----------|--------|
|
||||
| DINOv2/CLIP encoders | Low | High | **P0** | ✅ Done |
|
||||
| color-matcher post-process | Low | Medium | **P1** | 🚧 In progress |
|
||||
| Oklab color space | Low | Low | P2 | ⏳ Planned |
|
||||
| RAFT optical flow | Medium | High | **P1** | ⏳ Planned |
|
||||
| LoFTR matching | Medium | Medium | P2 | ⏳ Planned |
|
||||
| NAFNet refinement | High | High | P3 | ⏳ Planned |
|
||||
| Optimal transport | Medium | Medium | P3 | ⏳ Planned |
|
||||
| SAM/EVA02 encoders | Medium | Medium | P3 | ⏳ Planned |
|
||||
|
||||
---
|
||||
|
||||
## 💡 Usage Examples
|
||||
|
||||
### Select Components via Config:
|
||||
|
||||
```python
|
||||
config = {
|
||||
'feature_encoder': 'dinov2_vitb', # vgg19, dinov2_vitb, clip_vitb
|
||||
'matcher': 'raft', # cosine, raft, loftr
|
||||
'color_transfer': 'optimal_transport', # weighted, optimal_transport
|
||||
'refinement': 'nafnet', # unet, nafnet, restormer
|
||||
'temporal': 'raft_flow', # previous_frame, raft_flow, xmem
|
||||
'post_process': 'color_matcher_mkl', # wls, color_matcher, guided
|
||||
'color_space': 'lab', # lab, oklab, ipt
|
||||
}
|
||||
|
||||
colorizer = ModularColorizer(config)
|
||||
result = colorizer.colorize(video, reference)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 📚 References
|
||||
|
||||
- **DINOv2**: https://arxiv.org/abs/2304.07193
|
||||
- **RAFT**: https://arxiv.org/abs/2003.12039
|
||||
- **LoFTR**: https://arxiv.org/abs/2104.00680
|
||||
- **NAFNet**: https://arxiv.org/abs/2204.04676
|
||||
- **XMem**: https://arxiv.org/abs/2207.07115
|
||||
- **color-matcher**: https://github.com/hahnec/color-matcher
|
||||
- **Oklab**: https://bottosson.github.io/posts/oklab/
|
||||
|
||||
---
|
||||
|
||||
## 🤝 Contributing
|
||||
|
||||
To add a new component:
|
||||
|
||||
1. Create module in appropriate directory
|
||||
2. Implement standard interface
|
||||
3. Add to configuration system
|
||||
4. Document in this file
|
||||
5. Add benchmarks
|
||||
|
||||
---
|
||||
|
||||
**Next:** Implementing color-matcher post-processing module...
|
||||
@@ -276,7 +276,6 @@ frame_propagate = True
|
||||
- `use_half_resolution = True` (optional)
|
||||
- `use_torch_compile = True`
|
||||
- `use_sage_attention = True`
|
||||
- `wls_filter_on = True`
|
||||
|
||||
### High-end (16GB+ VRAM)
|
||||
**ColorMNet:**
|
||||
@@ -288,7 +287,6 @@ frame_propagate = True
|
||||
- `use_half_resolution = False`
|
||||
- `use_torch_compile = True`
|
||||
- `use_sage_attention = True`
|
||||
- `wls_filter_on = True`
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -1,273 +0,0 @@
|
||||
# ⚡ Performance Optimizations Guide
|
||||
|
||||
**How to get maximum speed from your video colorization**
|
||||
|
||||
---
|
||||
|
||||
## ✅ Already Working
|
||||
|
||||
Your setup is already performing well:
|
||||
- ✅ CUDA graphs detection working
|
||||
- ✅ FP16 enabled (2x faster)
|
||||
- ✅ RTX 5090 GPU detected
|
||||
- ✅ ColorMNet loaded successfully
|
||||
- ✅ Processing ~2.2 FPS (851 frames in ~500 seconds)
|
||||
|
||||
---
|
||||
|
||||
## 🚀 Available Optimizations
|
||||
|
||||
### **1. Visual Studio C++ Compiler (Optional, 20-40% speedup)**
|
||||
|
||||
**Status:** ❌ Not installed
|
||||
|
||||
**What it does:**
|
||||
- Enables CUDA spatial_correlation_sampler compilation
|
||||
- Optimizes ColorMNet's matching algorithm
|
||||
- **20-40% faster processing** for ColorMNet
|
||||
|
||||
**Currently:**
|
||||
```
|
||||
[ColorMNet] Missing requirements for CUDA extensions:
|
||||
[ColorMNet] ❌ Visual Studio C++ compiler (cl.exe) not found
|
||||
[ColorMNet] Currently using fallback mode (20-40% slower, same quality)
|
||||
```
|
||||
|
||||
**How to install:**
|
||||
|
||||
1. **Download Visual Studio Build Tools:**
|
||||
https://visualstudio.microsoft.com/downloads/
|
||||
|
||||
2. **Select "Desktop development with C++"**
|
||||
- This installs cl.exe (C++ compiler)
|
||||
- Size: ~7GB download, ~10GB installed
|
||||
|
||||
3. **Restart ComfyUI**
|
||||
- ColorMNet will auto-compile the optimizations
|
||||
- First compilation takes ~2-3 minutes
|
||||
- Subsequent runs use compiled version
|
||||
|
||||
**Expected improvement:**
|
||||
- **Before:** 851 frames in ~500 seconds (~2.2 FPS)
|
||||
- **After:** 851 frames in ~350-400 seconds (~2.5-3.0 FPS)
|
||||
- **Savings:** ~100-150 seconds on your video!
|
||||
|
||||
**Is it worth it?**
|
||||
- ✅ YES if you process many videos
|
||||
- ⚠️ MAYBE if you only process occasionally (large download)
|
||||
- ❌ NO if disk space is limited (10GB)
|
||||
|
||||
---
|
||||
|
||||
### **2. torch.compile for DeepExemplar (10-25% speedup)**
|
||||
|
||||
**Status:** ✅ Works on DeepExemplar nodes
|
||||
|
||||
**What it does:**
|
||||
- JIT compiles model for faster inference
|
||||
- **10-25% faster** on first run, then instant
|
||||
- Only works on DeepExemplar (not ColorMNet)
|
||||
|
||||
**Why not ColorMNet?**
|
||||
ColorMNet is already optimized with custom CUDA kernels and has a complex architecture that doesn't benefit from torch.compile. You'll see this message:
|
||||
|
||||
```
|
||||
[ColorMNet] torch.compile requested but not applicable to ColorMNet architecture
|
||||
[ColorMNet] ColorMNet is already optimized with custom kernels
|
||||
```
|
||||
|
||||
**This is expected and correct!**
|
||||
|
||||
**To enable (DeepExemplar only):**
|
||||
```
|
||||
use_torch_compile: true
|
||||
```
|
||||
|
||||
**Expected improvement:**
|
||||
- DeepExemplar: 10-25% faster
|
||||
- ColorMNet: No benefit (already optimized)
|
||||
|
||||
---
|
||||
|
||||
## 📊 Performance Comparison
|
||||
|
||||
### ColorMNet (Your Current Node)
|
||||
|
||||
| Setup | FPS | Time (851 frames) | vs Baseline |
|
||||
|-------|-----|------------------|-------------|
|
||||
| **Current (No optimizations)** | 2.2 | 500s (~8min) | Baseline |
|
||||
| **+ Visual Studio C++** | 2.8 | 400s (~6.5min) | **-20%** ⭐ |
|
||||
| **+ torch.compile** | 2.2 | 500s | No effect |
|
||||
| **+ Both** | 2.8 | 400s | **-20%** |
|
||||
|
||||
**Recommendation for ColorMNet:**
|
||||
- Install Visual Studio C++ if you process videos regularly
|
||||
- Don't bother enabling torch.compile (no benefit)
|
||||
|
||||
---
|
||||
|
||||
### DeepExemplar
|
||||
|
||||
| Setup | Relative Speed | Recommendation |
|
||||
|-------|---------------|----------------|
|
||||
| **vgg19 + torch.compile** | ⚡⚡⚡⚡⚡ | Fastest |
|
||||
| **dinov2_vitb (no compile)** | ⚡⚡⚡ | Best quality/speed balance |
|
||||
| **dinov2_vitb + torch.compile** | ⚡⚡⚡⚡ | Recommended ⭐ |
|
||||
| **dinov2_vitl + torch.compile** | ⚡⚡⚡ | Max quality |
|
||||
|
||||
**Recommendation for DeepExemplar:**
|
||||
- Enable torch.compile (10-25% speedup)
|
||||
- Use dinov2_vitb for best balance
|
||||
|
||||
---
|
||||
|
||||
## 🎯 Recommended Settings
|
||||
|
||||
### **For ColorMNet (Current):**
|
||||
|
||||
**Best Quality + Speed:**
|
||||
```
|
||||
memory_mode: balanced
|
||||
post_processor: color_matcher
|
||||
use_fp16: true
|
||||
use_torch_compile: false (no benefit)
|
||||
```
|
||||
|
||||
**After Installing Visual Studio C++:**
|
||||
```
|
||||
Same settings as above
|
||||
(Optimizations auto-enable, no settings change needed!)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### **For DeepExemplar:**
|
||||
|
||||
**Best Quality:**
|
||||
```
|
||||
feature_encoder: dinov2_vitb
|
||||
post_processor: color_matcher
|
||||
use_torch_compile: true
|
||||
use_sage_attention: false
|
||||
```
|
||||
|
||||
**Fastest:**
|
||||
```
|
||||
feature_encoder: vgg19
|
||||
post_processor: none
|
||||
use_torch_compile: true
|
||||
use_sage_attention: true
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 🔧 How to Install Visual Studio C++
|
||||
|
||||
### Step-by-Step:
|
||||
|
||||
1. **Go to:** https://visualstudio.microsoft.com/downloads/
|
||||
|
||||
2. **Download:** "Build Tools for Visual Studio 2022" (free)
|
||||
|
||||
3. **Run installer**
|
||||
|
||||
4. **Select workload:** "Desktop development with C++"
|
||||
- This includes:
|
||||
- MSVC C++ compiler
|
||||
- Windows SDK
|
||||
- CMake tools
|
||||
|
||||
5. **Install** (~7GB download, takes 15-30 min)
|
||||
|
||||
6. **Restart ComfyUI**
|
||||
|
||||
7. **Run ColorMNet** - you'll see:
|
||||
```
|
||||
[ColorMNet] ✓ Visual Studio C++ compiler found
|
||||
[ColorMNet] Compiling spatial_correlation_sampler...
|
||||
[ColorMNet] ✓ CUDA extensions compiled successfully!
|
||||
```
|
||||
|
||||
8. **Enjoy 20-40% faster processing!**
|
||||
|
||||
---
|
||||
|
||||
## ❓ FAQ
|
||||
|
||||
### **Q: Do I need Visual Studio C++ for DeepExemplar?**
|
||||
**A:** No, only for ColorMNet's optimizations.
|
||||
|
||||
### **Q: Will torch.compile speed up ColorMNet?**
|
||||
**A:** No, ColorMNet is already optimized. It only helps DeepExemplar.
|
||||
|
||||
### **Q: Is the Visual Studio download safe?**
|
||||
**A:** Yes, it's official Microsoft software.
|
||||
|
||||
### **Q: Can I uninstall Visual Studio after compilation?**
|
||||
**A:** No, it's needed each time Python packages are updated/rebuilt.
|
||||
|
||||
### **Q: How much disk space does it need?**
|
||||
**A:** ~10GB for Build Tools (just the compiler, not full Visual Studio)
|
||||
|
||||
### **Q: Will it slow down my computer?**
|
||||
**A:** No, it only runs when compiling CUDA extensions.
|
||||
|
||||
---
|
||||
|
||||
## 🎯 Quick Decision Guide
|
||||
|
||||
**"Should I install Visual Studio C++?"**
|
||||
|
||||
| Your Situation | Recommendation |
|
||||
|----------------|----------------|
|
||||
| Process videos daily/weekly | ✅ **YES** - Big time savings |
|
||||
| Process videos monthly | ⚠️ **MAYBE** - Depends on disk space |
|
||||
| Only testing/learning | ❌ **NO** - Not worth 10GB |
|
||||
| Limited disk space (<20GB free) | ❌ **NO** - Too large |
|
||||
| Have fast internet + disk space | ✅ **YES** - Why not? |
|
||||
|
||||
---
|
||||
|
||||
## 📈 Expected Performance After Optimizations
|
||||
|
||||
### Your Video (851 frames, 720x960):
|
||||
|
||||
**Currently:**
|
||||
- Time: 500 seconds (~8.3 minutes)
|
||||
- Speed: 2.2 FPS
|
||||
|
||||
**With Visual Studio C++ (ColorMNet):**
|
||||
- Time: ~350-400 seconds (~6-7 minutes)
|
||||
- Speed: ~2.5-3.0 FPS
|
||||
- **Savings: 100-150 seconds per video!**
|
||||
|
||||
**With torch.compile (DeepExemplar):**
|
||||
- Time: ~10-25% faster
|
||||
- Only applies to DeepExemplar nodes
|
||||
|
||||
---
|
||||
|
||||
## ✅ Summary
|
||||
|
||||
**What's already working:**
|
||||
- ✅ GPU detection
|
||||
- ✅ FP16
|
||||
- ✅ CUDA graphs handling
|
||||
- ✅ ColorMNet processing
|
||||
|
||||
**Optional optimizations:**
|
||||
1. **Visual Studio C++** - 20-40% faster ColorMNet
|
||||
2. **torch.compile** - 10-25% faster DeepExemplar (no benefit for ColorMNet)
|
||||
|
||||
**Current message you're seeing is normal:**
|
||||
```
|
||||
[ColorMNet] Currently using fallback mode (20-40% slower, same quality)
|
||||
```
|
||||
|
||||
This just means you could be 20-40% faster with Visual Studio C++ installed!
|
||||
|
||||
---
|
||||
|
||||
**Your choice: Install Visual Studio C++ for 20-40% speedup, or keep current setup!** 🚀
|
||||
|
||||
Both work perfectly - it's just an optional speed optimization!
|
||||
@@ -108,9 +108,6 @@ Load the example workflow from `workflows/Colorize Video Workflow.json`:
|
||||
- `target_width/height` - Output resolution
|
||||
- `frame_propagate` - Use temporal consistency (default: True)
|
||||
- `use_half_resolution` - Process at half resolution (default: True)
|
||||
- `wls_filter_on` - Apply smoothing filter (default: True)
|
||||
- `lambda_value` - Filter smoothing strength (default: 500.0)
|
||||
- `sigma_color` - Color sensitivity (default: 4.0)
|
||||
- `use_torch_compile` - Enable PyTorch 2.0+ compilation (default: False)
|
||||
- `use_sage_attention` - Enable SageAttention optimization (default: False)
|
||||
- **Outputs**: colorized_frames, performance_report
|
||||
@@ -119,9 +116,6 @@ Load the example workflow from `workflows/Colorize Video Workflow.json`:
|
||||
- `image_to_colorize` - Input grayscale image
|
||||
- `reference_image` - Color reference image
|
||||
- `target_width/height` - Output resolution
|
||||
- `wls_filter_on` - Apply smoothing filter (default: True)
|
||||
- `lambda_value` - Filter smoothing strength (default: 500.0)
|
||||
- `sigma_color` - Color sensitivity (default: 4.0)
|
||||
- `use_torch_compile` - Enable PyTorch 2.0+ compilation (default: False)
|
||||
- `use_sage_attention` - Enable SageAttention optimization (default: False)
|
||||
- **Outputs**: colorized_image, performance_report
|
||||
@@ -200,25 +194,17 @@ Connect the performance_report output to a text display node or save to file for
|
||||
|
||||
**For Best Quality:**
|
||||
- ColorMNet: `memory_mode=high_quality`, `use_fp16=False`, `use_torch_compile=False`
|
||||
- Deep Exemplar: `use_half_resolution=False`, `wls_filter_on=True`, `use_torch_compile=False`, `use_sage_attention=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`, `wls_filter_on=False`, `use_torch_compile=True`, `use_sage_attention=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
|
||||
|
||||
### WLS Filter (Deep Exemplar Only)
|
||||
|
||||
The WLS (Weighted Least Squares) filter smooths colors while preserving edges:
|
||||
- Requires `opencv-contrib-python` to be installed
|
||||
- `lambda_value` controls smoothing strength (higher = more smooth)
|
||||
- `sigma_color` controls color sensitivity
|
||||
- Adds processing time but improves visual quality
|
||||
|
||||
---
|
||||
|
||||
## 🔧 Advanced Configuration
|
||||
|
||||
@@ -1,306 +0,0 @@
|
||||
# 📝 ComfyUI Registry Publishing Guide
|
||||
|
||||
**Official guide for publishing to the ComfyUI Registry**
|
||||
|
||||
---
|
||||
|
||||
## ✅ Prerequisites Completed
|
||||
|
||||
- ✅ comfy-cli installed
|
||||
- ✅ pyproject.toml created and configured
|
||||
- ✅ Repository on GitHub
|
||||
- ✅ All code and documentation ready
|
||||
|
||||
---
|
||||
|
||||
## 🚀 Step-by-Step Publishing Process
|
||||
|
||||
### Step 1: Create Publisher Account
|
||||
|
||||
1. **Go to:** https://registry.comfy.org/
|
||||
|
||||
2. **Sign in** with GitHub
|
||||
|
||||
3. **Create Publisher Profile:**
|
||||
- Click on your profile/settings
|
||||
- Create a publisher identity
|
||||
- **Important:** Your Publisher ID appears after the `@` symbol
|
||||
- Example: `@jonstreeter`
|
||||
- **This ID is permanent and cannot be changed!**
|
||||
|
||||
4. **Copy your Publisher ID** (without the @)
|
||||
|
||||
---
|
||||
|
||||
### Step 2: Update pyproject.toml with Your Publisher ID
|
||||
|
||||
1. Open: `pyproject.toml`
|
||||
|
||||
2. Find the line:
|
||||
```toml
|
||||
PublisherId = "jonstreeter" # IMPORTANT: Replace with your actual Publisher ID
|
||||
```
|
||||
|
||||
3. **Replace `"jonstreeter"`** with your actual Publisher ID from the registry
|
||||
|
||||
4. Save the file
|
||||
|
||||
---
|
||||
|
||||
### Step 3: Generate API Key
|
||||
|
||||
1. **Go to:** https://registry.comfy.org/ (your publisher settings)
|
||||
|
||||
2. **Find:** "Registry Publishing API Key" section
|
||||
|
||||
3. **Click:** "Generate API Key" or "Create New Key"
|
||||
|
||||
4. **IMPORTANT:** Copy this key immediately!
|
||||
- You cannot view it again
|
||||
- If lost, you must generate a new one
|
||||
- Store it securely (password manager recommended)
|
||||
|
||||
5. **This key is for:**
|
||||
- ✅ Publishing custom nodes to Registry
|
||||
- ✅ Publishing to ComfyUI-Manager
|
||||
- ❌ NOT for paid API node usage (different key)
|
||||
|
||||
---
|
||||
|
||||
### Step 4: Publish Your Node
|
||||
|
||||
**Option A: Manual Publishing (First Time)**
|
||||
|
||||
```bash
|
||||
# Navigate to your node directory
|
||||
cd "G:\AIART\AI_Image_Generators\Comfy_UI_V42\ComfyUI\custom_nodes\ComfyUI-Deep-Exemplar-based-Video-Colorization"
|
||||
|
||||
# Publish to registry
|
||||
comfy node publish
|
||||
|
||||
# When prompted, paste your API key
|
||||
```
|
||||
|
||||
**Option B: Automated GitHub Actions (Recommended)**
|
||||
|
||||
See "GitHub Actions Setup" section below.
|
||||
|
||||
---
|
||||
|
||||
### Step 5: Verify Publication
|
||||
|
||||
After publishing, your node will appear at:
|
||||
```
|
||||
https://registry.comfy.org/@YOUR_PUBLISHER_ID/comfyui-reference-video-colorization
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 🤖 GitHub Actions Setup (Automated Publishing)
|
||||
|
||||
### Benefits:
|
||||
- ✅ Automatic publishing on new releases
|
||||
- ✅ No manual steps needed
|
||||
- ✅ Consistent versioning
|
||||
- ✅ Triggered by pyproject.toml changes
|
||||
|
||||
### Setup:
|
||||
|
||||
1. **Add API Key to GitHub Secrets:**
|
||||
- Go to: https://github.com/jonstreeter/ComfyUI-Reference-Based-Video-Colorization/settings/secrets/actions
|
||||
- Click: "New repository secret"
|
||||
- Name: `REGISTRY_ACCESS_TOKEN`
|
||||
- Value: [Your API key from Step 3]
|
||||
- Click: "Add secret"
|
||||
|
||||
2. **Create Workflow File:**
|
||||
|
||||
Create: `.github/workflows/publish-to-registry.yml`
|
||||
|
||||
```yaml
|
||||
name: Publish to ComfyUI Registry
|
||||
|
||||
on:
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
paths:
|
||||
- 'pyproject.toml'
|
||||
workflow_dispatch:
|
||||
|
||||
jobs:
|
||||
publish:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v4
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: '3.10'
|
||||
|
||||
- name: Install comfy-cli
|
||||
run: pip install comfy-cli
|
||||
|
||||
- name: Publish to registry
|
||||
env:
|
||||
REGISTRY_ACCESS_TOKEN: ${{ secrets.REGISTRY_ACCESS_TOKEN }}
|
||||
run: comfy node publish
|
||||
```
|
||||
|
||||
3. **Commit and Push:**
|
||||
```bash
|
||||
git add .github/workflows/publish-to-registry.yml pyproject.toml
|
||||
git commit -m "Add automated registry publishing"
|
||||
git push origin main
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 📊 Current Configuration
|
||||
|
||||
**Node Name:** `comfyui-reference-video-colorization`
|
||||
|
||||
**Display Name:** `Reference-Based Video Colorization`
|
||||
|
||||
**Version:** `2.0.0`
|
||||
|
||||
**Repository:** https://github.com/jonstreeter/ComfyUI-Reference-Based-Video-Colorization
|
||||
|
||||
**Icon:** `Workflows/Reference-Based-Colorization-Workflow.png`
|
||||
|
||||
**Description:**
|
||||
```
|
||||
Dual implementation of reference-based video colorization featuring
|
||||
ColorMNet (2024) with DINOv2 and Deep Exemplar (2019). Includes 4 nodes,
|
||||
multiple feature encoders, advanced post-processing, and auto-installer.
|
||||
```
|
||||
|
||||
**Dependencies:**
|
||||
- torch>=2.0.0
|
||||
- torchvision>=0.15.0
|
||||
- numpy>=1.24.0
|
||||
- Pillow>=9.5.0
|
||||
- scikit-image>=0.21.0
|
||||
- opencv-python>=4.8.0
|
||||
- einops>=0.7.0
|
||||
- tqdm>=4.65.0
|
||||
- progressbar2>=4.2.0
|
||||
- gdown
|
||||
|
||||
---
|
||||
|
||||
## 🔄 Updating Your Published Node
|
||||
|
||||
### When you make changes:
|
||||
|
||||
1. **Update version in pyproject.toml:**
|
||||
```toml
|
||||
version = "2.0.1" # Increment version
|
||||
```
|
||||
|
||||
2. **Commit changes:**
|
||||
```bash
|
||||
git add .
|
||||
git commit -m "Update to v2.0.1: [describe changes]"
|
||||
git push origin main
|
||||
```
|
||||
|
||||
3. **Publish update:**
|
||||
- **Manual:** Run `comfy node publish` again
|
||||
- **Automated:** GitHub Actions will publish automatically
|
||||
|
||||
### Version Guidelines:
|
||||
- **Major (2.0.0 → 3.0.0):** Breaking changes
|
||||
- **Minor (2.0.0 → 2.1.0):** New features, backward compatible
|
||||
- **Patch (2.0.0 → 2.0.1):** Bug fixes, backward compatible
|
||||
|
||||
---
|
||||
|
||||
## ✅ Pre-Publication Checklist
|
||||
|
||||
Before running `comfy node publish`:
|
||||
|
||||
- [x] Publisher account created on registry.comfy.org
|
||||
- [ ] Publisher ID updated in pyproject.toml
|
||||
- [ ] API key generated and saved securely
|
||||
- [x] pyproject.toml properly configured
|
||||
- [x] LICENSE file exists
|
||||
- [x] README.md is comprehensive
|
||||
- [x] All code tested and working
|
||||
- [x] Example workflows included
|
||||
- [x] Dependencies documented
|
||||
- [x] Repository is public
|
||||
|
||||
---
|
||||
|
||||
## 🎯 Quick Start Commands
|
||||
|
||||
```bash
|
||||
# 1. Navigate to node directory
|
||||
cd "G:\AIART\AI_Image_Generators\Comfy_UI_V42\ComfyUI\custom_nodes\ComfyUI-Deep-Exemplar-based-Video-Colorization"
|
||||
|
||||
# 2. Verify pyproject.toml
|
||||
cat pyproject.toml
|
||||
|
||||
# 3. Update Publisher ID (if needed)
|
||||
# Edit pyproject.toml and replace "jonstreeter" with your actual Publisher ID
|
||||
|
||||
# 4. Publish
|
||||
comfy node publish
|
||||
# (Paste your API key when prompted)
|
||||
|
||||
# 5. Verify
|
||||
# Visit: https://registry.comfy.org/@YOUR_PUBLISHER_ID/comfyui-reference-video-colorization
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 🔗 Important Links
|
||||
|
||||
- **ComfyUI Registry:** https://registry.comfy.org/
|
||||
- **Publishing Docs:** https://docs.comfy.org/registry/publishing
|
||||
- **Your Repository:** https://github.com/jonstreeter/ComfyUI-Reference-Based-Video-Colorization
|
||||
- **Your Publisher Profile:** https://registry.comfy.org/@YOUR_PUBLISHER_ID
|
||||
|
||||
---
|
||||
|
||||
## ❓ Troubleshooting
|
||||
|
||||
### "Publisher ID not found"
|
||||
- Make sure you've created a publisher account at registry.comfy.org
|
||||
- Verify the Publisher ID in pyproject.toml matches your registry profile
|
||||
- Publisher ID is case-sensitive
|
||||
|
||||
### "Invalid API key"
|
||||
- Regenerate API key from registry.comfy.org
|
||||
- Make sure you're using the Registry Publishing API key (not paid API key)
|
||||
- Check for extra spaces when copying
|
||||
|
||||
### "pyproject.toml validation failed"
|
||||
- Run `comfy node init` to see example format
|
||||
- Check that all required fields are present
|
||||
- Ensure version follows semantic versioning (X.Y.Z)
|
||||
|
||||
### "Icon not found"
|
||||
- Make sure the icon path is relative to the repository root
|
||||
- Supported formats: SVG, PNG, JPG, GIF
|
||||
- Max size: 800x400px
|
||||
|
||||
---
|
||||
|
||||
## 🎉 Success!
|
||||
|
||||
Once published, users can install your node via:
|
||||
|
||||
1. **ComfyUI Manager** (will be automatically indexed)
|
||||
2. **comfy-cli:**
|
||||
```bash
|
||||
comfy node install comfyui-reference-video-colorization
|
||||
```
|
||||
3. **Manual git clone** (existing method still works)
|
||||
|
||||
---
|
||||
|
||||
**Your node will appear in the official ComfyUI Registry!** 🚀
|
||||
@@ -1,163 +0,0 @@
|
||||
# ✅ Repository Renamed Successfully!
|
||||
|
||||
**Old Name:** `ComfyUI-Deep-Exemplar-based-Video-Colorization`
|
||||
**New Name:** `ComfyUI-Reference-Based-Video-Colorization`
|
||||
|
||||
**New GitHub URL:** https://github.com/jonstreeter/ComfyUI-Reference-Based-Video-Colorization
|
||||
|
||||
---
|
||||
|
||||
## 🎯 Why the Rename?
|
||||
|
||||
The old name "Deep-Exemplar-based" was misleading because the repository now contains:
|
||||
|
||||
1. **Deep Exemplar (2019)** - Original CVPR method
|
||||
2. **ColorMNet (2024)** - Modern ECCV method
|
||||
3. **Multiple encoders** - VGG19, DINOv2, CLIP
|
||||
4. **Advanced post-processing** - color-matcher, WLS, guided, bilateral
|
||||
|
||||
The new name **"Reference-Based Video Colorization"** accurately reflects that this is a **general-purpose reference-based colorization toolkit** with multiple methods.
|
||||
|
||||
---
|
||||
|
||||
## 📝 What Was Updated
|
||||
|
||||
### All Documentation Files (29 occurrences across 10 files)
|
||||
|
||||
✅ **README.md** - Title and URLs updated
|
||||
✅ **CHANGELOG.md** - Commit history link updated
|
||||
✅ **README_NEW.md** - Title and URLs updated
|
||||
✅ **RELEASE_NOTES_v2.0.0.md** - All references updated (8 occurrences)
|
||||
✅ **GITHUB_READY.md** - All paths and URLs updated (5 occurrences)
|
||||
✅ **MIGRATION_GUIDE.md** - Directory paths updated (3 occurrences)
|
||||
✅ **MODERN_COMPONENTS_GUIDE.md** - Installation path updated
|
||||
✅ **ARCHITECTURE.md** - Directory structure updated
|
||||
✅ **QUICKSTART.md** - Installation path updated
|
||||
✅ **AUTO_INSTALLER_README.md** - Installation path updated
|
||||
|
||||
### Updated Elements
|
||||
|
||||
- ✅ Repository name in titles
|
||||
- ✅ GitHub URLs (git clone, issues, discussions)
|
||||
- ✅ Directory paths in code examples
|
||||
- ✅ ComfyUI Manager entry (author, title, reference)
|
||||
- ✅ Installation commands
|
||||
- ✅ File path references
|
||||
|
||||
---
|
||||
|
||||
## 🚀 Next Step: Create GitHub Repository
|
||||
|
||||
### Option 1: Create New Repository on GitHub
|
||||
|
||||
1. Go to: https://github.com/new
|
||||
2. **Repository name:** `ComfyUI-Reference-Based-Video-Colorization`
|
||||
3. **Description:** Dual implementation of reference-based video colorization: ColorMNet (2024) + Deep Exemplar (2019) for ComfyUI
|
||||
4. **Public** repository
|
||||
5. **DO NOT initialize** with README (we already have one)
|
||||
6. Click "Create repository"
|
||||
|
||||
### Option 2: Rename Existing Repository (If Already Pushed)
|
||||
|
||||
If you already pushed the old repository:
|
||||
|
||||
1. Go to your repository settings
|
||||
2. Scroll to "Rename repository"
|
||||
3. Change to: `ComfyUI-Reference-Based-Video-Colorization`
|
||||
4. GitHub will automatically redirect old URLs
|
||||
|
||||
---
|
||||
|
||||
## 📦 Push to GitHub
|
||||
|
||||
After creating the repository, run these commands:
|
||||
|
||||
```bash
|
||||
# Navigate to repository
|
||||
cd "G:\AIART\AI_Image_Generators\Comfy_UI_V42\ComfyUI\custom_nodes\ComfyUI-Deep-Exemplar-based-Video-Colorization"
|
||||
|
||||
# If this is a new repository (never pushed before):
|
||||
git remote add origin https://github.com/jonstreeter/ComfyUI-Reference-Based-Video-Colorization.git
|
||||
git branch -M main
|
||||
|
||||
# If you're updating an existing remote:
|
||||
git remote set-url origin https://github.com/jonstreeter/ComfyUI-Reference-Based-Video-Colorization.git
|
||||
|
||||
# Add all files
|
||||
git add .
|
||||
|
||||
# Create commit
|
||||
git commit -m "Repository rename: Reference-Based Video Colorization
|
||||
|
||||
Renamed from ComfyUI-Deep-Exemplar-based-Video-Colorization to ComfyUI-Reference-Based-Video-Colorization
|
||||
|
||||
Reason: Repository now contains multiple reference-based colorization methods:
|
||||
- ColorMNet (2024) - Modern DINOv2-based approach
|
||||
- Deep Exemplar (2019) - Classic CVPR method
|
||||
- Multiple encoders (VGG19, DINOv2, CLIP)
|
||||
- Advanced post-processing pipeline
|
||||
|
||||
The new name better reflects the general-purpose nature of the toolkit.
|
||||
|
||||
All documentation and URLs updated to:
|
||||
https://github.com/jonstreeter/ComfyUI-Reference-Based-Video-Colorization"
|
||||
|
||||
# Push to GitHub
|
||||
git push -u origin main
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## ✅ Verification Checklist
|
||||
|
||||
After pushing:
|
||||
|
||||
- [ ] Visit: https://github.com/jonstreeter/ComfyUI-Reference-Based-Video-Colorization
|
||||
- [ ] Verify README displays correctly
|
||||
- [ ] Check that git clone URL works
|
||||
- [ ] Test installation instructions
|
||||
- [ ] Verify all documentation links work
|
||||
- [ ] Check workflow images display
|
||||
- [ ] Confirm ComfyUI Manager entry (if listed)
|
||||
|
||||
---
|
||||
|
||||
## 📊 Updated Repository Details
|
||||
|
||||
### For ComfyUI Manager Submission
|
||||
|
||||
```json
|
||||
{
|
||||
"author": "Jon Streeter",
|
||||
"title": "ComfyUI Reference-Based Video Colorization",
|
||||
"reference": "https://github.com/jonstreeter/ComfyUI-Reference-Based-Video-Colorization",
|
||||
"files": [
|
||||
"https://github.com/jonstreeter/ComfyUI-Reference-Based-Video-Colorization"
|
||||
],
|
||||
"install_type": "git-clone",
|
||||
"description": "Dual implementation of reference-based video colorization: ColorMNet (2024) with DINOv2 features + Deep Exemplar (2019) CVPR method. Features 4 nodes (2 video, 2 image), auto-download models, progress bars, and performance reports. Transform B&W videos using color reference images!",
|
||||
"nodename_pattern": "(ColorMNet|DeepExemplar)",
|
||||
"pip": ["einops", "progressbar2", "gdown", "opencv-contrib-python"]
|
||||
}
|
||||
```
|
||||
|
||||
### Repository Metadata
|
||||
|
||||
- **Name:** ComfyUI-Reference-Based-Video-Colorization
|
||||
- **Description:** Dual implementation of reference-based video colorization: ColorMNet (2024) + Deep Exemplar (2019) for ComfyUI
|
||||
- **Topics:** `comfyui`, `video-colorization`, `reference-based`, `colormnet`, `deep-exemplar`, `dinov2`, `clip`, `pytorch`
|
||||
- **License:** MIT (with model-specific licenses in README)
|
||||
|
||||
---
|
||||
|
||||
## 🎉 Summary
|
||||
|
||||
✅ **All 29 occurrences updated** across 10 documentation files
|
||||
✅ **Repository name modernized** to reflect general-purpose toolkit
|
||||
✅ **GitHub URLs consistent** - all point to new repository name
|
||||
✅ **Ready to push** - all files updated and ready for GitHub
|
||||
|
||||
**New Repository URL:**
|
||||
https://github.com/jonstreeter/ComfyUI-Reference-Based-Video-Colorization
|
||||
|
||||
The repository is now properly named to reflect its true purpose: a comprehensive reference-based video colorization toolkit with multiple state-of-the-art methods! 🚀
|
||||
@@ -1,355 +0,0 @@
|
||||
# 🎉 Silent Auto-Installer - Complete!
|
||||
|
||||
**Zero-configuration setup for modern video colorization**
|
||||
|
||||
---
|
||||
|
||||
## ✅ What Was Done
|
||||
|
||||
A silent auto-installer has been integrated that **automatically installs missing dependencies** when you select modern components in ComfyUI.
|
||||
|
||||
**No manual pip commands needed! No setup required!**
|
||||
|
||||
---
|
||||
|
||||
## 🚀 How It Works
|
||||
|
||||
### User Experience:
|
||||
|
||||
1. **Select modern component** in ComfyUI node:
|
||||
- Change `feature_encoder` from `vgg19` to `dinov2_vitb`
|
||||
- Or change `post_processor` from `none` to `color_matcher`
|
||||
|
||||
2. **Run your workflow** (just click Queue!)
|
||||
|
||||
3. **Watch console** - Dependencies install automatically:
|
||||
```
|
||||
[AutoInstall] timm not found (required for DINOv2)
|
||||
[AutoInstall] Installing timm (DINOv2 support)...
|
||||
[AutoInstall] ✓ timm installed successfully
|
||||
[DeepExColorImageNode] ✓ Using feature encoder: dinov2_vitb
|
||||
```
|
||||
|
||||
4. **That's it!** Your colorization runs with the modern component
|
||||
|
||||
5. **Next time** - Already installed, works instantly
|
||||
|
||||
---
|
||||
|
||||
## 📦 What Gets Auto-Installed
|
||||
|
||||
### When you select DINOv2 encoders:
|
||||
- Automatically installs: `timm>=0.9.0`
|
||||
- Size: ~50MB
|
||||
- Models auto-download: ~300MB (DINOv2 weights)
|
||||
|
||||
### When you select CLIP encoder:
|
||||
- Automatically installs: `clip @ git+https://github.com/openai/CLIP.git`
|
||||
- Size: ~30MB
|
||||
- Models auto-download: ~350MB (CLIP weights)
|
||||
|
||||
### When you select color-matcher:
|
||||
- Automatically installs: `color-matcher>=0.3.0`
|
||||
- Size: ~10MB
|
||||
|
||||
### When you select WLS/Guided filters:
|
||||
- Automatically installs: `opencv-contrib-python>=4.7.0`
|
||||
- Size: ~60MB
|
||||
- Note: Replaces opencv-python if needed
|
||||
|
||||
---
|
||||
|
||||
## 🎯 Quick Start Examples
|
||||
|
||||
### Example 1: Better Quality (DINOv2 + color-matcher)
|
||||
|
||||
**Before (manual installation):**
|
||||
```bash
|
||||
# Had to run these commands
|
||||
pip install timm
|
||||
pip install color-matcher
|
||||
```
|
||||
|
||||
**Now (automatic):**
|
||||
1. Open ComfyUI
|
||||
2. In Deep Exemplar node, change:
|
||||
- `feature_encoder: dinov2_vitb`
|
||||
- `post_processor: color_matcher`
|
||||
3. Click "Queue Prompt"
|
||||
4. ✅ Done! (Dependencies install automatically on first run)
|
||||
|
||||
**Console output:**
|
||||
```
|
||||
[AutoInstall] timm not found (required for DINOv2)
|
||||
[AutoInstall] Installing timm (DINOv2 support)...
|
||||
[AutoInstall] ✓ timm installed successfully
|
||||
[AutoInstall] color-matcher not found (required for color matching post-processing)
|
||||
[AutoInstall] Installing color-matcher...
|
||||
[AutoInstall] ✓ color-matcher installed successfully
|
||||
[DeepExColorImageNode] ✓ Using feature encoder: dinov2_vitb
|
||||
[DeepExColorImageNode] ✓ Post-processing applied: color_matcher
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Example 2: Text-Guided (CLIP)
|
||||
|
||||
**Before (manual installation):**
|
||||
```bash
|
||||
# Had to run this command
|
||||
pip install git+https://github.com/openai/CLIP.git
|
||||
```
|
||||
|
||||
**Now (automatic):**
|
||||
1. Open ComfyUI
|
||||
2. In Deep Exemplar node, change:
|
||||
- `feature_encoder: clip_vitb`
|
||||
- `text_guidance: "warm sunset colors"`
|
||||
- `post_processor: color_matcher`
|
||||
3. Click "Queue Prompt"
|
||||
4. ✅ Done! (CLIP installs automatically on first run)
|
||||
|
||||
**Console output:**
|
||||
```
|
||||
[AutoInstall] CLIP not found (required for text-guided colorization)
|
||||
[AutoInstall] Installing CLIP...
|
||||
[AutoInstall] ✓ CLIP installed successfully
|
||||
[DeepExColorImageNode] ✓ Using feature encoder: clip_vitb
|
||||
[DeepExColorImageNode] ✓ Text guidance: 'warm sunset colors' (weight=0.3)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 🔧 Technical Details
|
||||
|
||||
### Files Created:
|
||||
|
||||
**`auto_installer.py`** - The installer module
|
||||
- Silent installation via pip subprocess
|
||||
- Smart caching to avoid repeated attempts
|
||||
- Component-specific installers
|
||||
- Graceful failure handling
|
||||
|
||||
**Integration points:**
|
||||
- `DeepExemplarColorizationNodes.py` (Image node)
|
||||
- `DeepExemplarColorizationNodes.py` (Video node)
|
||||
|
||||
### How It Works:
|
||||
|
||||
```python
|
||||
# When you select dinov2_vitb:
|
||||
from .auto_installer import ensure_dependencies_for_encoder
|
||||
deps_ok = ensure_dependencies_for_encoder('dinov2_vitb')
|
||||
# → Checks if timm is installed
|
||||
# → If not, runs: pip install timm>=0.9.0
|
||||
# → Returns True on success
|
||||
|
||||
if deps_ok:
|
||||
# Load DINOv2
|
||||
encoder = get_feature_encoder('dinov2_vitb')
|
||||
else:
|
||||
# Fall back to VGG19
|
||||
encoder = VGG_NET
|
||||
```
|
||||
|
||||
### Safety Features:
|
||||
|
||||
✅ **Never crashes the node** - Falls back to VGG19 if installation fails
|
||||
✅ **Caching** - Won't retry failed installations repeatedly
|
||||
✅ **Clear logging** - Shows exactly what's being installed
|
||||
✅ **Graceful degradation** - Workflow continues even if install fails
|
||||
|
||||
---
|
||||
|
||||
## ⏱️ First-Run Timing
|
||||
|
||||
**Expect these installation times on first use:**
|
||||
|
||||
| Component | Download | Install | Total |
|
||||
|-----------|----------|---------|-------|
|
||||
| timm (DINOv2) | ~30s | ~20s | **~50s** |
|
||||
| CLIP | ~40s | ~30s | **~70s** |
|
||||
| color-matcher | ~10s | ~10s | **~20s** |
|
||||
| opencv-contrib | ~45s | ~30s | **~75s** |
|
||||
|
||||
**After first install:** Instant (0s) - packages already installed
|
||||
|
||||
**Model downloads** (separate, one-time):
|
||||
- DINOv2 models: ~300MB, downloads on first use
|
||||
- CLIP models: ~350MB, downloads on first use
|
||||
|
||||
---
|
||||
|
||||
## 🐛 Troubleshooting
|
||||
|
||||
### "Installation failed" message
|
||||
|
||||
**Check:**
|
||||
1. Internet connection active?
|
||||
2. Pip working? Run `pip --version` in terminal
|
||||
3. Check console for specific error
|
||||
|
||||
**Manual fallback:**
|
||||
```bash
|
||||
# If auto-install fails, manually install:
|
||||
pip install timm
|
||||
pip install color-matcher
|
||||
pip install git+https://github.com/openai/CLIP.git
|
||||
```
|
||||
|
||||
### Node still uses VGG19
|
||||
|
||||
**This means:**
|
||||
- Auto-installation failed (check console for error)
|
||||
- Node gracefully fell back to VGG19
|
||||
- Your workflow still works, just with VGG19 instead
|
||||
|
||||
**Fix:**
|
||||
- Check console error message
|
||||
- Try manual installation
|
||||
- Restart ComfyUI after manual install
|
||||
|
||||
### Permissions error (Linux/Mac)
|
||||
|
||||
**If you see permission denied:**
|
||||
```bash
|
||||
# Option 1: Install with --user
|
||||
pip install --user timm
|
||||
|
||||
# Option 2: Use sudo (not recommended)
|
||||
sudo pip install timm
|
||||
```
|
||||
|
||||
### Multiple Python environments
|
||||
|
||||
**If packages install but don't import:**
|
||||
- ComfyUI may be using different Python than pip
|
||||
- Solution: Auto-installer uses `sys.executable` to ensure correct Python
|
||||
- Should work automatically
|
||||
|
||||
---
|
||||
|
||||
## 📊 Comparison: Before vs After
|
||||
|
||||
### Before Auto-Installer:
|
||||
|
||||
**To use DINOv2:**
|
||||
1. Read documentation
|
||||
2. Find installation command
|
||||
3. Open terminal
|
||||
4. Run: `pip install timm`
|
||||
5. Restart ComfyUI
|
||||
6. Change node settings
|
||||
7. Run workflow
|
||||
|
||||
**Time:** 5-10 minutes (reading + setup)
|
||||
**Friction:** High (requires terminal, reading docs)
|
||||
|
||||
### After Auto-Installer:
|
||||
|
||||
**To use DINOv2:**
|
||||
1. Change node setting: `feature_encoder: dinov2_vitb`
|
||||
2. Run workflow
|
||||
3. Wait ~50s for auto-install (first time only)
|
||||
|
||||
**Time:** 1 minute (first time), 0s (subsequent)
|
||||
**Friction:** Zero (just select and run)
|
||||
|
||||
---
|
||||
|
||||
## 🎯 User Benefits
|
||||
|
||||
### For Beginners:
|
||||
✅ **No terminal commands** - Just use the dropdown
|
||||
✅ **No documentation reading** - Tooltips explain everything
|
||||
✅ **No setup** - Works out of the box
|
||||
✅ **Can't break it** - Falls back to VGG19 if anything fails
|
||||
|
||||
### For Advanced Users:
|
||||
✅ **Faster experimentation** - Try components instantly
|
||||
✅ **No interruption** - No need to stop and install
|
||||
✅ **Clear logging** - See exactly what's happening
|
||||
✅ **Manual override** - Can pre-install if preferred
|
||||
|
||||
### For Everyone:
|
||||
✅ **One-time cost** - Install once, use forever
|
||||
✅ **Transparent** - Console shows all activity
|
||||
✅ **Safe** - Can't break existing workflow
|
||||
✅ **Optional** - Can still use VGG19 if preferred
|
||||
|
||||
---
|
||||
|
||||
## 🔄 Migration from Manual Installation
|
||||
|
||||
### Already have packages installed?
|
||||
**Perfect!** Auto-installer detects existing packages and skips installation. No duplicate downloads.
|
||||
|
||||
### Prefer manual installation?
|
||||
**No problem!** Pre-install everything:
|
||||
```bash
|
||||
pip install -r requirements_modern.txt
|
||||
# or run the installer once
|
||||
python auto_installer.py
|
||||
```
|
||||
Auto-installer will detect and skip.
|
||||
|
||||
### Want to opt-out?
|
||||
**Easy!** Just keep using:
|
||||
- `feature_encoder: vgg19`
|
||||
- `post_processor: none`
|
||||
|
||||
No auto-installation happens.
|
||||
|
||||
---
|
||||
|
||||
## 📚 Documentation
|
||||
|
||||
**User Guides:**
|
||||
- [INTEGRATION_COMPLETE.md](INTEGRATION_COMPLETE.md) - Quick start guide
|
||||
- [MODERN_COMPONENTS_GUIDE.md](MODERN_COMPONENTS_GUIDE.md) - Component details
|
||||
|
||||
**Technical Docs:**
|
||||
- [AUTO_INSTALLER_README.md](AUTO_INSTALLER_README.md) - Auto-installer architecture
|
||||
- [INTEGRATION_SUMMARY.md](INTEGRATION_SUMMARY.md) - Technical integration details
|
||||
- [MODULAR_ARCHITECTURE.md](MODULAR_ARCHITECTURE.md) - Complete system architecture
|
||||
|
||||
---
|
||||
|
||||
## ✅ Testing Checklist
|
||||
|
||||
**Tested:**
|
||||
- [x] Syntax validation (both files compile)
|
||||
- [x] Auto-installer module structure
|
||||
- [x] Integration with IMAGE node
|
||||
- [x] Integration with VIDEO node
|
||||
- [x] Graceful fallback logic
|
||||
- [x] Documentation complete
|
||||
|
||||
**Requires User Testing:**
|
||||
- [ ] Actual package installation (needs internet)
|
||||
- [ ] DINOv2 auto-install + model download
|
||||
- [ ] CLIP auto-install + model download
|
||||
- [ ] color-matcher auto-install
|
||||
- [ ] opencv-contrib auto-install
|
||||
- [ ] Fallback when installation fails
|
||||
- [ ] Performance with modern components
|
||||
|
||||
---
|
||||
|
||||
## 🎉 Summary
|
||||
|
||||
**What you can do now:**
|
||||
|
||||
1. **Open ComfyUI**
|
||||
2. **Select any modern component** (DINOv2, CLIP, color-matcher)
|
||||
3. **Run your workflow**
|
||||
4. **Dependencies install automatically**
|
||||
5. **Enjoy better colorization!**
|
||||
|
||||
**No setup, no terminal, no documentation reading required!**
|
||||
|
||||
---
|
||||
|
||||
**The future is zero-configuration! 🚀**
|
||||
|
||||
Just select what you want and it works. That's it.
|
||||
Vendored
-3
@@ -1,3 +0,0 @@
|
||||
{
|
||||
"python.formatting.provider": "yapf"
|
||||
}
|
||||
@@ -1,5 +1,6 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torch.nn.parallel
|
||||
|
||||
|
||||
@@ -126,17 +127,26 @@ class ColorVidNet(nn.Module):
|
||||
conv7_3norm = self.conv7_3norm(conv7_3)
|
||||
conv8_1 = self.conv8_1(conv7_3norm)
|
||||
conv3_3_short = self.conv3_3_short(conv3_3norm)
|
||||
# Align skip connection dimensions
|
||||
if conv8_1.shape[2:] != conv3_3_short.shape[2:]:
|
||||
conv8_1 = F.interpolate(conv8_1, size=conv3_3_short.shape[2:], mode='bilinear', align_corners=False)
|
||||
conv8_1_comb = self.relu8_1_comb(conv8_1 + conv3_3_short)
|
||||
conv8_2 = self.relu8_2(self.conv8_2(conv8_1_comb))
|
||||
conv8_3 = self.relu8_3(self.conv8_3(conv8_2))
|
||||
conv8_3norm = self.conv8_3norm(conv8_3)
|
||||
conv9_1 = self.conv9_1(conv8_3norm)
|
||||
conv2_2_short = self.conv2_2_short(conv2_2norm)
|
||||
# Align skip connection dimensions
|
||||
if conv9_1.shape[2:] != conv2_2_short.shape[2:]:
|
||||
conv9_1 = F.interpolate(conv9_1, size=conv2_2_short.shape[2:], mode='bilinear', align_corners=False)
|
||||
conv9_1_comb = self.relu9_1_comb(conv9_1 + conv2_2_short)
|
||||
conv9_2 = self.relu9_2(self.conv9_2(conv9_1_comb))
|
||||
conv9_2norm = self.conv9_2norm(conv9_2)
|
||||
conv10_1 = self.conv10_1(conv9_2norm)
|
||||
conv1_2_short = self.conv1_2_short(conv1_2norm)
|
||||
# Align skip connection dimensions
|
||||
if conv10_1.shape[2:] != conv1_2_short.shape[2:]:
|
||||
conv10_1 = F.interpolate(conv10_1, size=conv1_2_short.shape[2:], mode='bilinear', align_corners=False)
|
||||
conv10_1_comb = self.relu10_1_comb(conv10_1 + conv1_2_short)
|
||||
conv10_2 = self.relu10_2(self.conv10_2(conv10_1_comb))
|
||||
conv10_ab = self.conv10_ab(conv10_2)
|
||||
|
||||
+30
-10
@@ -472,8 +472,6 @@ class WarpNet(nn.Module):
|
||||
channel = B_lab_map.shape[1]
|
||||
image_height = B_lab_map.shape[2]
|
||||
image_width = B_lab_map.shape[3]
|
||||
feature_height = int(image_height / 4)
|
||||
feature_width = int(image_width / 4)
|
||||
|
||||
# scale feature size to 44*44
|
||||
A_feature2_1 = self.layer2_1(A_relu2_1)
|
||||
@@ -485,10 +483,30 @@ class WarpNet(nn.Module):
|
||||
A_feature5_1 = self.layer5_1(A_relu5_1)
|
||||
B_feature5_1 = self.layer5_1(B_relu5_1)
|
||||
|
||||
# concatenate features
|
||||
if A_feature5_1.shape[2] != A_feature2_1.shape[2] or A_feature5_1.shape[3] != A_feature2_1.shape[3]:
|
||||
A_feature5_1 = F.pad(A_feature5_1, (0, 0, 1, 1), "replicate")
|
||||
B_feature5_1 = F.pad(B_feature5_1, (0, 0, 1, 1), "replicate")
|
||||
# concatenate features - ensure all features have matching spatial dimensions
|
||||
# Find the target size (use feature2_1 as reference since it has stride=2)
|
||||
target_h, target_w = A_feature2_1.shape[2], A_feature2_1.shape[3]
|
||||
|
||||
# Align all features to target size using interpolation
|
||||
def align_feature(feat, target_h, target_w):
|
||||
if feat.shape[2] != target_h or feat.shape[3] != target_w:
|
||||
return F.interpolate(feat, size=(target_h, target_w), mode='bilinear', align_corners=False)
|
||||
return feat
|
||||
|
||||
A_feature2_1 = align_feature(A_feature2_1, target_h, target_w)
|
||||
A_feature3_1 = align_feature(A_feature3_1, target_h, target_w)
|
||||
A_feature4_1 = align_feature(A_feature4_1, target_h, target_w)
|
||||
A_feature5_1 = align_feature(A_feature5_1, target_h, target_w)
|
||||
|
||||
B_feature2_1 = align_feature(B_feature2_1, target_h, target_w)
|
||||
B_feature3_1 = align_feature(B_feature3_1, target_h, target_w)
|
||||
B_feature4_1 = align_feature(B_feature4_1, target_h, target_w)
|
||||
B_feature5_1 = align_feature(B_feature5_1, target_h, target_w)
|
||||
|
||||
# Update feature dimensions to match actual aligned feature size
|
||||
feature_height = target_h
|
||||
feature_width = target_w
|
||||
|
||||
A_features = self.layer(torch.cat((A_feature2_1, A_feature3_1, A_feature4_1, A_feature5_1), 1))
|
||||
B_features = self.layer(torch.cat((B_feature2_1, B_feature3_1, B_feature4_1, B_feature5_1), 1))
|
||||
|
||||
@@ -515,8 +533,8 @@ class WarpNet(nn.Module):
|
||||
f_WTA = f_WTA / temperature
|
||||
f_div_C = F.softmax(f_WTA.squeeze_(), dim=-1) # 2*1936*1936;
|
||||
|
||||
# downsample the reference color
|
||||
B_lab = F.avg_pool2d(B_lab_map, 4)
|
||||
# downsample the reference color to match actual feature dimensions
|
||||
B_lab = F.interpolate(B_lab_map, size=(feature_height, feature_width), mode='bilinear', align_corners=False)
|
||||
B_lab = B_lab.view(batch_size, channel, -1)
|
||||
B_lab = B_lab.permute(0, 2, 1) # 2*1936*channel
|
||||
|
||||
@@ -524,7 +542,9 @@ class WarpNet(nn.Module):
|
||||
y = torch.matmul(f_div_C, B_lab) # 2*1936*channel
|
||||
y = y.permute(0, 2, 1).contiguous()
|
||||
y = y.view(batch_size, channel, feature_height, feature_width) # 2*3*44*44
|
||||
y = self.upsampling(y)
|
||||
similarity_map = self.upsampling(similarity_map)
|
||||
|
||||
# Upsample to match input dimensions (instead of fixed 4x scale)
|
||||
y = F.interpolate(y, size=(image_height, image_width), mode='bilinear', align_corners=False)
|
||||
similarity_map = F.interpolate(similarity_map, size=(image_height, image_width), mode='bilinear', align_corners=False)
|
||||
|
||||
return y, similarity_map
|
||||
|
||||
@@ -4,11 +4,19 @@ Modern Feature Extractors for Video Colorization
|
||||
This package provides drop-in replacements for VGG19/ResNet50 feature extractors
|
||||
with modern vision models (DINOv2, CLIP, etc.) that offer superior semantic understanding.
|
||||
|
||||
IMPORTANT LIMITATION:
|
||||
DeepExemplar's NonlocalNet and ColorNet were trained specifically with VGG19 features.
|
||||
While this package provides dimension-compatible projections from DINOv2/CLIP to VGG19,
|
||||
these projections are NOT trained and will produce suboptimal results (color artifacts).
|
||||
|
||||
For production use, stick with VGG19. Alternative encoders are EXPERIMENTAL and would
|
||||
require fine-tuning the entire colorization network to work properly.
|
||||
|
||||
Quick usage:
|
||||
from models.feature_extractors import get_feature_encoder
|
||||
|
||||
# Get encoder by name
|
||||
encoder = get_feature_encoder('dinov2_vitb')
|
||||
encoder = get_feature_encoder('vgg19') # Recommended
|
||||
|
||||
# Use exactly like VGG19
|
||||
features = encoder(image, ["r12", "r22", "r32", "r42", "r52"], preprocess=True)
|
||||
@@ -19,7 +27,7 @@ import torch.nn as nn
|
||||
|
||||
|
||||
def get_feature_encoder(
|
||||
encoder_type: Literal['vgg19', 'dinov2_vitb', 'dinov2_vitl', 'dinov2_vitg',
|
||||
encoder_type: Literal['vgg19', 'dinov2_vits', 'dinov2_vitb', 'dinov2_vitl', 'dinov2_vitg',
|
||||
'clip_vitb', 'clip_vitl', 'resnet50'] = 'vgg19',
|
||||
device: str = 'cuda',
|
||||
**kwargs
|
||||
@@ -60,9 +68,10 @@ def get_feature_encoder(
|
||||
vgg = models.vgg19(pretrained=True)
|
||||
return VGG19Wrapper(vgg).to(device).eval()
|
||||
|
||||
elif encoder_type in ['dinov2_vitb', 'dinov2_vitl', 'dinov2_vitg']:
|
||||
elif encoder_type in ['dinov2_vits', 'dinov2_vitb', 'dinov2_vitl', 'dinov2_vitg']:
|
||||
from .dinov2_encoder import DINOv2Encoder
|
||||
model_map = {
|
||||
'dinov2_vits': 'dinov2_vits14',
|
||||
'dinov2_vitb': 'dinov2_vitb14',
|
||||
'dinov2_vitl': 'dinov2_vitl14',
|
||||
'dinov2_vitg': 'dinov2_vitg14',
|
||||
@@ -88,7 +97,7 @@ def get_feature_encoder(
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown encoder type: {encoder_type}. "
|
||||
f"Valid options: vgg19, dinov2_vitb, dinov2_vitl, dinov2_vitg, "
|
||||
f"Valid options: vgg19, dinov2_vits, dinov2_vitb, dinov2_vitl, dinov2_vitg, "
|
||||
f"clip_vitb, clip_vitl, resnet50"
|
||||
)
|
||||
|
||||
@@ -123,6 +132,14 @@ def get_encoder_info(encoder_type: str) -> dict:
|
||||
'quality_rating': 3,
|
||||
'description': 'ColorMNet baseline, faster than VGG19'
|
||||
},
|
||||
'dinov2_vits': {
|
||||
'name': 'DINOv2 ViT-Small',
|
||||
'year': 2023,
|
||||
'params': '21M',
|
||||
'speed_rating': 5,
|
||||
'quality_rating': 4,
|
||||
'description': 'Smallest DINOv2, faster than vitb with good quality'
|
||||
},
|
||||
'dinov2_vitb': {
|
||||
'name': 'DINOv2 ViT-Base',
|
||||
'year': 2023,
|
||||
@@ -180,6 +197,7 @@ def list_available_encoders() -> list:
|
||||
return [
|
||||
'vgg19',
|
||||
'resnet50',
|
||||
'dinov2_vits',
|
||||
'dinov2_vitb',
|
||||
'dinov2_vitl',
|
||||
'dinov2_vitg',
|
||||
|
||||
@@ -61,12 +61,33 @@ class CLIPEncoder(nn.Module):
|
||||
for param in self.clip_model.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
# Projection layers to match VGG19 dimensions
|
||||
self.proj_r12 = nn.Conv2d(self.embed_dim, 64, 1)
|
||||
self.proj_r22 = nn.Conv2d(self.embed_dim, 128, 1)
|
||||
self.proj_r32 = nn.Conv2d(self.embed_dim, 256, 1)
|
||||
self.proj_r42 = nn.Conv2d(self.embed_dim, 512, 1)
|
||||
self.proj_r52 = nn.Conv2d(self.embed_dim, 512, 1)
|
||||
# Projection layers to match VGG19 dimensions (move to device)
|
||||
# Use conv + bn + relu to match VGG19-like feature distributions
|
||||
self.proj_r12 = nn.Sequential(
|
||||
nn.Conv2d(self.embed_dim, 64, 1, bias=False),
|
||||
nn.BatchNorm2d(64),
|
||||
nn.ReLU(inplace=False)
|
||||
).to(device)
|
||||
self.proj_r22 = nn.Sequential(
|
||||
nn.Conv2d(self.embed_dim, 128, 1, bias=False),
|
||||
nn.BatchNorm2d(128),
|
||||
nn.ReLU(inplace=False)
|
||||
).to(device)
|
||||
self.proj_r32 = nn.Sequential(
|
||||
nn.Conv2d(self.embed_dim, 256, 1, bias=False),
|
||||
nn.BatchNorm2d(256),
|
||||
nn.ReLU(inplace=False)
|
||||
).to(device)
|
||||
self.proj_r42 = nn.Sequential(
|
||||
nn.Conv2d(self.embed_dim, 512, 1, bias=False),
|
||||
nn.BatchNorm2d(512),
|
||||
nn.ReLU(inplace=False)
|
||||
).to(device)
|
||||
self.proj_r52 = nn.Sequential(
|
||||
nn.Conv2d(self.embed_dim, 512, 1, bias=False),
|
||||
nn.BatchNorm2d(512),
|
||||
nn.ReLU(inplace=False)
|
||||
).to(device)
|
||||
|
||||
# Layer mapping (CLIP ViT-B/16 has 12 layers)
|
||||
self.layer_mapping = {
|
||||
@@ -77,9 +98,12 @@ class CLIPEncoder(nn.Module):
|
||||
'r52': (12, self.proj_r52), # Final features
|
||||
}
|
||||
|
||||
# CLIP normalization
|
||||
self.register_buffer('mean', torch.tensor([0.48145466, 0.4578275, 0.40821073]).view(1, 3, 1, 1))
|
||||
self.register_buffer('std', torch.tensor([0.26862954, 0.26130258, 0.27577711]).view(1, 3, 1, 1))
|
||||
# Use ImageNet normalization for compatibility with DeepExemplar
|
||||
# (CLIP's native normalization causes color distortion in colorization task)
|
||||
mean = torch.tensor([0.485, 0.456, 0.406], device=device).view(1, 3, 1, 1)
|
||||
std = torch.tensor([0.229, 0.224, 0.225], device=device).view(1, 3, 1, 1)
|
||||
self.register_buffer('mean', mean)
|
||||
self.register_buffer('std', std)
|
||||
|
||||
def _extract_features_at_layer(self, x, layer_idx):
|
||||
"""Extract features from a specific transformer layer."""
|
||||
|
||||
@@ -63,11 +63,29 @@ class DINOv2Encoder(nn.Module):
|
||||
|
||||
# Projection layers to match VGG19 channel dimensions
|
||||
# VGG19 dimensions: [64, 128, 256, 512, 512]
|
||||
self.proj_r12 = nn.Conv2d(self.embed_dim, 64, 1)
|
||||
self.proj_r22 = nn.Conv2d(self.embed_dim, 128, 1)
|
||||
self.proj_r32 = nn.Conv2d(self.embed_dim, 256, 1)
|
||||
self.proj_r42 = nn.Conv2d(self.embed_dim, 512, 1)
|
||||
self.proj_r52 = nn.Conv2d(self.embed_dim, 512, 1)
|
||||
# Use deeper projections with Xavier initialization for better feature preservation
|
||||
def make_projection(in_ch, out_ch):
|
||||
"""Create a 2-layer projection with proper initialization."""
|
||||
mid_ch = (in_ch + out_ch) // 2
|
||||
proj = nn.Sequential(
|
||||
nn.Conv2d(in_ch, mid_ch, 1),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Conv2d(mid_ch, out_ch, 1),
|
||||
nn.ReLU(inplace=True)
|
||||
)
|
||||
# Xavier initialization for better gradient flow
|
||||
for m in proj.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
nn.init.xavier_uniform_(m.weight)
|
||||
if m.bias is not None:
|
||||
nn.init.zeros_(m.bias)
|
||||
return proj
|
||||
|
||||
self.proj_r12 = make_projection(self.embed_dim, 64)
|
||||
self.proj_r22 = make_projection(self.embed_dim, 128)
|
||||
self.proj_r32 = make_projection(self.embed_dim, 256)
|
||||
self.proj_r42 = make_projection(self.embed_dim, 512)
|
||||
self.proj_r52 = make_projection(self.embed_dim, 512)
|
||||
|
||||
# Mapping from layer names to outputs
|
||||
self.layer_mapping = {
|
||||
@@ -96,9 +114,11 @@ class DINOv2Encoder(nn.Module):
|
||||
# Remove CLS token and reshape to spatial
|
||||
patch_features = features[:, 1:] # Remove CLS token
|
||||
|
||||
# Calculate patch grid size
|
||||
num_patches = patch_features.shape[1]
|
||||
patch_h = patch_w = int(num_patches ** 0.5)
|
||||
# Calculate patch grid size based on actual image dimensions
|
||||
# DINOv2 patch size is 14x14
|
||||
patch_size = 14
|
||||
patch_h = H // patch_size
|
||||
patch_w = W // patch_size
|
||||
|
||||
# Reshape to spatial: [B, N, D] -> [B, D, H, W]
|
||||
features_spatial = patch_features.transpose(1, 2).reshape(B, self.embed_dim, patch_h, patch_w)
|
||||
@@ -216,8 +236,11 @@ class DINOv2EncoderColorMNet(nn.Module):
|
||||
for blk in self.dinov2.blocks[:layer_idx]:
|
||||
features = blk(features)
|
||||
patch_features = features[:, 1:]
|
||||
num_patches = patch_features.shape[1]
|
||||
patch_h = patch_w = int(num_patches ** 0.5)
|
||||
# Calculate patch grid size based on actual image dimensions
|
||||
# DINOv2 patch size is 14x14
|
||||
patch_size = 14
|
||||
patch_h = H // patch_size
|
||||
patch_w = W // patch_size
|
||||
features_spatial = patch_features.transpose(1, 2).reshape(B, self.embed_dim, patch_h, patch_w)
|
||||
return features_spatial
|
||||
|
||||
|
||||
@@ -132,23 +132,6 @@ class ColorMNetVideoNode:
|
||||
"default": "resnet50",
|
||||
"tooltip": "Feature extraction model: resnet50 (ColorMNet default), vgg19 (fast), dinov2_vitb (recommended, 40-60% better), dinov2_vitl (best quality), clip_vitb (text-guided)"
|
||||
}),
|
||||
"post_processor": (["none", "wls", "guided", "bilateral", "color_matcher"], {
|
||||
"default": "none",
|
||||
"tooltip": "Post-processing method: none (fastest), wls (edge-aware), guided (fast smoothing), bilateral (classic), color_matcher (best color consistency, recommended)"
|
||||
}),
|
||||
"post_process_strength": ("FLOAT", {
|
||||
"default": 0.8,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.05,
|
||||
"tooltip": "Post-processing strength: 0=disabled, 1=full effect (for color_matcher/guided/bilateral)"
|
||||
}),
|
||||
"temporal_consistency": ("BOOLEAN", {
|
||||
"default": True,
|
||||
"label_on": "enabled",
|
||||
"label_off": "disabled",
|
||||
"tooltip": "Enable temporal consistency in post-processing to reduce flickering between frames (for color_matcher)"
|
||||
}),
|
||||
"use_fp16": ("BOOLEAN", {
|
||||
"default": True,
|
||||
"label_on": "enabled",
|
||||
@@ -191,9 +174,6 @@ class ColorMNetVideoNode:
|
||||
target_height: int,
|
||||
memory_mode: str = "balanced",
|
||||
feature_encoder: str = "resnet50",
|
||||
post_processor: str = "none",
|
||||
post_process_strength: float = 0.8,
|
||||
temporal_consistency: bool = True,
|
||||
use_fp16: bool = True,
|
||||
use_torch_compile: bool = False,
|
||||
text_guidance: str = "",
|
||||
@@ -309,6 +289,7 @@ class ColorMNetVideoNode:
|
||||
# Colorize
|
||||
self.logger.info(f"Processing {num_frames} frames at {target_height}x{target_width}")
|
||||
|
||||
colorization_start_time = time.time()
|
||||
colorized = self.inference_pipeline.colorize_video(
|
||||
frames=video_frames,
|
||||
reference=reference_image,
|
||||
@@ -316,68 +297,11 @@ class ColorMNetVideoNode:
|
||||
target_width=target_width,
|
||||
progress_callback=progress_callback,
|
||||
)
|
||||
colorization_time = time.time() - colorization_start_time
|
||||
|
||||
if not COMFY_PROGRESS_AVAILABLE:
|
||||
pbar.close()
|
||||
|
||||
# Apply modern post-processing (batch processing for videos)
|
||||
post_proc_applied = "none"
|
||||
encoder_name = feature_encoder
|
||||
|
||||
if post_processor != "none":
|
||||
try:
|
||||
# Auto-install dependencies if needed
|
||||
from .auto_installer import ensure_dependencies_for_post_processor
|
||||
deps_ok = ensure_dependencies_for_post_processor(post_processor)
|
||||
|
||||
if not deps_ok:
|
||||
self.logger.warning(f"Dependencies for {post_processor} could not be installed")
|
||||
proc = None
|
||||
else:
|
||||
from .post_processing import get_post_processor
|
||||
|
||||
# Create post-processor
|
||||
if post_processor == "wls":
|
||||
proc = get_post_processor('wls', lambda_value=500.0, sigma_color=4.0)
|
||||
elif post_processor == "color_matcher":
|
||||
proc = get_post_processor('color_matcher', method='mkl',
|
||||
temporal_consistency=temporal_consistency,
|
||||
consistency_weight=0.3)
|
||||
elif post_processor == "guided":
|
||||
proc = get_post_processor('guided', radius=8, eps=0.01)
|
||||
elif post_processor == "bilateral":
|
||||
proc = get_post_processor('bilateral', d=9, sigma_color=75, sigma_space=75)
|
||||
else:
|
||||
proc = None
|
||||
|
||||
if proc is not None:
|
||||
self.logger.info(f"Applying post-processing: {post_processor} to {num_frames} frames...")
|
||||
|
||||
# Apply post-processing
|
||||
if post_processor == "color_matcher":
|
||||
# color_matcher for videos with reference matching
|
||||
ref_rgb = reference_image
|
||||
|
||||
# Process video with temporal consistency
|
||||
processed = proc.match_video_frames(
|
||||
colorized,
|
||||
ref_rgb,
|
||||
strength=post_process_strength,
|
||||
progress_callback=lambda i, total: None
|
||||
)
|
||||
colorized = processed
|
||||
else:
|
||||
# Other post-processors: apply per-frame
|
||||
processed_frames = proc.match_video_frames(colorized, None)
|
||||
colorized = processed_frames
|
||||
|
||||
post_proc_applied = post_processor
|
||||
self.logger.info(f"✓ Post-processing complete: {post_processor}")
|
||||
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Post-processing '{post_processor}' failed: {e}")
|
||||
self.logger.info("Continuing without post-processing...")
|
||||
|
||||
# Calculate timing
|
||||
end_time = time.time()
|
||||
elapsed_time = end_time - start_time
|
||||
@@ -391,19 +315,16 @@ class ColorMNetVideoNode:
|
||||
f"Date/Time: {timestamp}\n"
|
||||
f"Frames Processed: {num_frames}\n"
|
||||
f"Resolution: {target_width}x{target_height}\n"
|
||||
f"Total Time: {elapsed_time:.2f} seconds\n"
|
||||
f"Average FPS: {fps:.2f}\n"
|
||||
f"Time per Frame: {elapsed_time / num_frames:.3f} seconds\n"
|
||||
f"Feature Encoder: {encoder_name} (ColorMNet built-in ResNet50)\n"
|
||||
f"Post-Processor: {post_proc_applied}\n"
|
||||
)
|
||||
if post_proc_applied == "color_matcher":
|
||||
report += f"Matching Strength: {post_process_strength}\n"
|
||||
report += f"Temporal Consistency: {'Enabled' if temporal_consistency else 'Disabled'}\n"
|
||||
report += (
|
||||
f"Memory Mode: {memory_mode}\n"
|
||||
f"FP16 Enabled: {use_fp16}\n"
|
||||
f"Torch Compile: {use_torch_compile}\n"
|
||||
f"\n"
|
||||
f"Performance:\n"
|
||||
f" Total Time: {elapsed_time:.2f}s ({fps:.2f} FPS)\n"
|
||||
f" Time per Frame: {elapsed_time / num_frames:.3f}s\n"
|
||||
f"\n"
|
||||
f"Configuration:\n"
|
||||
f" Feature Encoder: {feature_encoder}\n"
|
||||
f" Memory Mode: {memory_mode}\n"
|
||||
f" FP16 Enabled: {use_fp16}\n"
|
||||
f" Torch Compile: {use_torch_compile}\n"
|
||||
f"{'=' * 50}"
|
||||
)
|
||||
|
||||
@@ -469,17 +390,6 @@ class ColorMNetImageNode:
|
||||
"step": 32,
|
||||
"tooltip": "Output height - must be multiple of 32 (will be adjusted automatically)"
|
||||
}),
|
||||
"post_processor": (["none", "wls", "guided", "bilateral", "color_matcher"], {
|
||||
"default": "none",
|
||||
"tooltip": "Post-processing method: none (fastest), wls (edge-aware), guided (fast smoothing), bilateral (classic), color_matcher (best color consistency, recommended)"
|
||||
}),
|
||||
"post_process_strength": ("FLOAT", {
|
||||
"default": 0.8,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.05,
|
||||
"tooltip": "Post-processing strength: 0=disabled, 1=full effect (for color_matcher/guided/bilateral)"
|
||||
}),
|
||||
"use_fp16": ("BOOLEAN", {
|
||||
"default": True,
|
||||
"tooltip": "Use half-precision (FP16) for faster processing with lower VRAM usage (minimal quality impact)"
|
||||
@@ -502,8 +412,6 @@ class ColorMNetImageNode:
|
||||
reference_image: torch.Tensor,
|
||||
target_width: int,
|
||||
target_height: int,
|
||||
post_processor: str = "none",
|
||||
post_process_strength: float = 0.8,
|
||||
use_fp16: bool = True,
|
||||
use_torch_compile: bool = False,
|
||||
) -> Tuple[torch.Tensor, str]:
|
||||
@@ -592,55 +500,6 @@ class ColorMNetImageNode:
|
||||
target_width=target_width,
|
||||
)
|
||||
|
||||
# Apply post-processing
|
||||
post_proc_applied = "none"
|
||||
|
||||
if post_processor != "none":
|
||||
try:
|
||||
# Auto-install dependencies if needed
|
||||
from .auto_installer import ensure_dependencies_for_post_processor
|
||||
deps_ok = ensure_dependencies_for_post_processor(post_processor)
|
||||
|
||||
if not deps_ok:
|
||||
self.logger.warning(f"Dependencies for {post_processor} could not be installed")
|
||||
proc = None
|
||||
else:
|
||||
from .post_processing import get_post_processor
|
||||
|
||||
# Create post-processor
|
||||
if post_processor == "wls":
|
||||
proc = get_post_processor('wls', lambda_value=500.0, sigma_color=4.0)
|
||||
elif post_processor == "color_matcher":
|
||||
proc = get_post_processor('color_matcher', method='mkl')
|
||||
elif post_processor == "guided":
|
||||
proc = get_post_processor('guided', radius=8, eps=0.01)
|
||||
elif post_processor == "bilateral":
|
||||
proc = get_post_processor('bilateral', d=9, sigma_color=75, sigma_space=75)
|
||||
else:
|
||||
proc = None
|
||||
|
||||
if proc is not None:
|
||||
self.logger.info(f"Applying post-processing: {post_processor}...")
|
||||
|
||||
# Apply post-processing
|
||||
if post_processor == "color_matcher":
|
||||
# color_matcher needs reference image
|
||||
processed = proc.match_to_reference(colorized, reference_image)
|
||||
|
||||
# Apply strength blending
|
||||
if post_process_strength < 1.0:
|
||||
processed = colorized * (1 - post_process_strength) + processed * post_process_strength
|
||||
else:
|
||||
processed = proc.match_to_reference(colorized, None)
|
||||
|
||||
colorized = processed
|
||||
post_proc_applied = post_processor
|
||||
self.logger.info(f"✓ Post-processing complete: {post_processor}")
|
||||
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Post-processing '{post_processor}' failed: {e}")
|
||||
self.logger.info("Continuing without post-processing...")
|
||||
|
||||
# Add batch dimension back
|
||||
colorized = colorized.unsqueeze(0) if colorized.dim() == 3 else colorized
|
||||
|
||||
@@ -656,11 +515,6 @@ class ColorMNetImageNode:
|
||||
f"Date/Time: {timestamp}\n"
|
||||
f"Resolution: {target_width}x{target_height}\n"
|
||||
f"Total Time: {elapsed_time:.3f} seconds\n"
|
||||
f"Post-Processor: {post_proc_applied}\n"
|
||||
)
|
||||
if post_proc_applied == "color_matcher":
|
||||
report += f"Matching Strength: {post_process_strength}\n"
|
||||
report += (
|
||||
f"FP16 Enabled: {use_fp16}\n"
|
||||
f"Torch Compile: {use_torch_compile}\n"
|
||||
f"{'=' * 50}"
|
||||
|
||||
@@ -1,278 +0,0 @@
|
||||
"""
|
||||
Post-Processing Module for Video Colorization
|
||||
|
||||
Provides various post-processing filters and refinement techniques to improve
|
||||
colorization quality and temporal consistency.
|
||||
|
||||
Available processors:
|
||||
- ColorMatcherPost: Color matching using color-matcher library
|
||||
- WLSFilter: Weighted Least Squares edge-aware smoothing
|
||||
- GuidedFilter: Fast edge-preserving filter
|
||||
- BilateralFilter: Bilateral filtering
|
||||
- TemporalStabilizer: Optical flow-based temporal stabilization
|
||||
|
||||
Usage:
|
||||
from post_processing import get_post_processor
|
||||
|
||||
processor = get_post_processor('color_matcher', method='mkl')
|
||||
refined = processor.process(colorized_frames, reference)
|
||||
"""
|
||||
|
||||
from typing import Literal, Optional, Dict, Any
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
def get_post_processor(
|
||||
processor_type: Literal['color_matcher', 'wls', 'guided', 'bilateral', 'none'] = 'color_matcher',
|
||||
**kwargs
|
||||
) -> nn.Module:
|
||||
"""Get a post-processor by name.
|
||||
|
||||
Args:
|
||||
processor_type: Type of post-processor
|
||||
- 'color_matcher': Color matching (best for color consistency)
|
||||
- 'wls': Weighted Least Squares filter (edge-aware smoothing)
|
||||
- 'guided': Guided filter (fast edge-preserving)
|
||||
- 'bilateral': Bilateral filter (noise reduction)
|
||||
- 'none': No post-processing
|
||||
**kwargs: Additional arguments for the processor
|
||||
|
||||
Returns:
|
||||
Post-processor instance
|
||||
|
||||
Example:
|
||||
>>> processor = get_post_processor('color_matcher', method='mkl')
|
||||
>>> refined = processor.match_to_reference(colorized, reference)
|
||||
"""
|
||||
processor_type = processor_type.lower()
|
||||
|
||||
if processor_type == 'color_matcher':
|
||||
from .color_matcher_post import ColorMatcherPost
|
||||
return ColorMatcherPost(**kwargs)
|
||||
|
||||
elif processor_type == 'wls':
|
||||
from .wls_filter import WLSFilter
|
||||
return WLSFilter(**kwargs)
|
||||
|
||||
elif processor_type == 'guided':
|
||||
from .guided_filter import GuidedFilter
|
||||
return GuidedFilter(**kwargs)
|
||||
|
||||
elif processor_type == 'bilateral':
|
||||
from .bilateral_filter import BilateralFilter
|
||||
return BilateralFilter(**kwargs)
|
||||
|
||||
elif processor_type == 'none':
|
||||
return NoOpProcessor()
|
||||
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown processor type: {processor_type}. "
|
||||
f"Valid: color_matcher, wls, guided, bilateral, none"
|
||||
)
|
||||
|
||||
|
||||
class NoOpProcessor:
|
||||
"""No-op processor that returns input unchanged."""
|
||||
|
||||
def match_to_reference(self, image, reference, **kwargs):
|
||||
return image
|
||||
|
||||
def match_video_frames(self, frames, reference, **kwargs):
|
||||
return frames
|
||||
|
||||
def process(self, *args, **kwargs):
|
||||
return args[0] if args else None
|
||||
|
||||
|
||||
def list_available_processors() -> list:
|
||||
"""List all available post-processors.
|
||||
|
||||
Returns:
|
||||
List of processor names
|
||||
"""
|
||||
processors = []
|
||||
|
||||
try:
|
||||
from .color_matcher_post import ColorMatcherPost
|
||||
processors.append('color_matcher')
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
try:
|
||||
from .wls_filter import WLSFilter
|
||||
processors.append('wls')
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
try:
|
||||
from .guided_filter import GuidedFilter
|
||||
processors.append('guided')
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
try:
|
||||
from .bilateral_filter import BilateralFilter
|
||||
processors.append('bilateral')
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
processors.append('none')
|
||||
|
||||
return processors
|
||||
|
||||
|
||||
def get_processor_info(processor_type: str) -> Dict[str, Any]:
|
||||
"""Get information about a post-processor.
|
||||
|
||||
Args:
|
||||
processor_type: Name of processor
|
||||
|
||||
Returns:
|
||||
Dictionary with processor metadata
|
||||
"""
|
||||
info = {
|
||||
'color_matcher': {
|
||||
'name': 'Color Matcher',
|
||||
'description': 'Match colors to reference using optimal transport or histogram methods',
|
||||
'speed': 4, # 1-5 scale
|
||||
'quality': 5,
|
||||
'temporal_aware': True,
|
||||
'requires': ['color-matcher'],
|
||||
'best_for': 'Ensuring color consistency with reference image'
|
||||
},
|
||||
'wls': {
|
||||
'name': 'WLS Filter',
|
||||
'description': 'Weighted Least Squares edge-aware smoothing',
|
||||
'speed': 3,
|
||||
'quality': 4,
|
||||
'temporal_aware': False,
|
||||
'requires': ['opencv-contrib-python'],
|
||||
'best_for': 'Edge-aware color smoothing'
|
||||
},
|
||||
'guided': {
|
||||
'name': 'Guided Filter',
|
||||
'description': 'Fast edge-preserving filter',
|
||||
'speed': 5,
|
||||
'quality': 4,
|
||||
'temporal_aware': False,
|
||||
'requires': ['opencv-python'],
|
||||
'best_for': 'Fast edge-preserving smoothing'
|
||||
},
|
||||
'bilateral': {
|
||||
'name': 'Bilateral Filter',
|
||||
'description': 'Bilateral noise reduction filter',
|
||||
'speed': 4,
|
||||
'quality': 3,
|
||||
'temporal_aware': False,
|
||||
'requires': ['opencv-python'],
|
||||
'best_for': 'Noise reduction while preserving edges'
|
||||
},
|
||||
'none': {
|
||||
'name': 'No Post-Processing',
|
||||
'description': 'Skip post-processing',
|
||||
'speed': 5,
|
||||
'quality': 0,
|
||||
'temporal_aware': False,
|
||||
'requires': [],
|
||||
'best_for': 'When raw output is sufficient'
|
||||
}
|
||||
}
|
||||
|
||||
return info.get(processor_type, {})
|
||||
|
||||
|
||||
def compare_processors(
|
||||
image: torch.Tensor,
|
||||
reference: torch.Tensor,
|
||||
verbose: bool = True,
|
||||
) -> Dict[str, torch.Tensor]:
|
||||
"""Compare all available post-processors.
|
||||
|
||||
Args:
|
||||
image: Test image [H, W, 3]
|
||||
reference: Reference image [H, W, 3]
|
||||
verbose: Print comparison info
|
||||
|
||||
Returns:
|
||||
Dictionary mapping processor names to results
|
||||
"""
|
||||
available = list_available_processors()
|
||||
results = {}
|
||||
|
||||
if verbose:
|
||||
print("\n" + "="*80)
|
||||
print("Post-Processor Comparison")
|
||||
print("="*80)
|
||||
|
||||
for proc_type in available:
|
||||
if proc_type == 'none':
|
||||
continue
|
||||
|
||||
try:
|
||||
processor = get_post_processor(proc_type)
|
||||
result = processor.match_to_reference(image, reference)
|
||||
results[proc_type] = result
|
||||
|
||||
if verbose:
|
||||
info = get_processor_info(proc_type)
|
||||
speed = "⚡" * info.get('speed', 0)
|
||||
quality = "⭐" * info.get('quality', 0)
|
||||
print(f"\n{info.get('name', proc_type)}")
|
||||
print(f" Speed: {speed}")
|
||||
print(f" Quality: {quality}")
|
||||
print(f" Best for: {info.get('best_for', 'N/A')}")
|
||||
print(f" ✓ Success")
|
||||
|
||||
except Exception as e:
|
||||
if verbose:
|
||||
print(f"\n{proc_type}")
|
||||
print(f" ✗ Failed: {e}")
|
||||
|
||||
if verbose:
|
||||
print("\n" + "="*80)
|
||||
print(f"Successfully tested: {len(results)}/{len(available)-1} processors")
|
||||
print("="*80 + "\n")
|
||||
|
||||
return results
|
||||
|
||||
|
||||
# Convenience exports
|
||||
__all__ = [
|
||||
'get_post_processor',
|
||||
'list_available_processors',
|
||||
'get_processor_info',
|
||||
'compare_processors',
|
||||
'NoOpProcessor',
|
||||
]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("Post-Processing Module Demo")
|
||||
print("="*80)
|
||||
|
||||
# List available
|
||||
print("\nAvailable processors:")
|
||||
for proc in list_available_processors():
|
||||
print(f" - {proc}")
|
||||
|
||||
# Show info
|
||||
print("\nProcessor details:")
|
||||
for proc in ['color_matcher', 'wls', 'guided']:
|
||||
info = get_processor_info(proc)
|
||||
if info:
|
||||
print(f"\n{info['name']}:")
|
||||
print(f" {info['description']}")
|
||||
print(f" Best for: {info['best_for']}")
|
||||
|
||||
# Test loading
|
||||
print("\n" + "="*80)
|
||||
print("Testing processor loading...")
|
||||
try:
|
||||
processor = get_post_processor('color_matcher')
|
||||
print("✓ Successfully loaded color_matcher")
|
||||
except Exception as e:
|
||||
print(f"✗ Failed to load: {e}")
|
||||
|
||||
print("\n" + "="*80)
|
||||
@@ -1,94 +0,0 @@
|
||||
"""
|
||||
Bilateral Filter Post-Processing
|
||||
|
||||
Classic bilateral filtering for edge-preserving noise reduction.
|
||||
|
||||
Reference: Bilateral Filtering for Gray and Color Images, Tomasi & Manduchi, ICCV 1998
|
||||
"""
|
||||
|
||||
import torch
|
||||
import cv2
|
||||
import numpy as np
|
||||
from typing import Optional
|
||||
|
||||
|
||||
class BilateralFilter:
|
||||
"""Bilateral filtering for edge-preserving smoothing."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
d: int = 9,
|
||||
sigma_color: float = 75,
|
||||
sigma_space: float = 75,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
d: Diameter of each pixel neighborhood
|
||||
sigma_color: Filter sigma in the color space
|
||||
sigma_space: Filter sigma in the coordinate space
|
||||
"""
|
||||
self.d = d
|
||||
self.sigma_color = sigma_color
|
||||
self.sigma_space = sigma_space
|
||||
|
||||
def match_to_reference(
|
||||
self,
|
||||
image: torch.Tensor,
|
||||
reference: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
"""Apply bilateral filtering.
|
||||
|
||||
Args:
|
||||
image: Input image [H, W, 3] in range [0, 1]
|
||||
reference: Not used (for interface compatibility)
|
||||
|
||||
Returns:
|
||||
Filtered image [H, W, 3]
|
||||
"""
|
||||
device = image.device
|
||||
|
||||
# Convert to numpy
|
||||
img_np = (image.cpu().numpy() * 255).astype(np.uint8)
|
||||
|
||||
# Apply bilateral filter
|
||||
try:
|
||||
filtered_np = cv2.bilateralFilter(
|
||||
img_np,
|
||||
self.d,
|
||||
self.sigma_color,
|
||||
self.sigma_space
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"[BilateralFilter] Filtering failed: {e}")
|
||||
return image
|
||||
|
||||
# Convert back
|
||||
filtered = torch.from_numpy(filtered_np).float() / 255.0
|
||||
return filtered.to(device)
|
||||
|
||||
def match_video_frames(
|
||||
self,
|
||||
frames: torch.Tensor,
|
||||
reference: Optional[torch.Tensor] = None,
|
||||
progress_callback: Optional[callable] = None,
|
||||
) -> torch.Tensor:
|
||||
"""Apply bilateral filtering to video."""
|
||||
num_frames = frames.shape[0]
|
||||
filtered_frames = []
|
||||
|
||||
for i in range(num_frames):
|
||||
filtered = self.match_to_reference(frames[i], reference)
|
||||
filtered_frames.append(filtered)
|
||||
|
||||
if progress_callback:
|
||||
progress_callback(i + 1, num_frames)
|
||||
|
||||
return torch.stack(filtered_frames, dim=0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("Testing BilateralFilter...")
|
||||
test_img = torch.rand(256, 256, 3)
|
||||
bf = BilateralFilter(d=9, sigma_color=75, sigma_space=75)
|
||||
filtered = bf.match_to_reference(test_img)
|
||||
print(f"✓ Filtered: {test_img.shape} -> {filtered.shape}")
|
||||
@@ -1,399 +0,0 @@
|
||||
"""
|
||||
Color Matcher Post-Processing for Video Colorization
|
||||
|
||||
Uses the color-matcher library to refine colorization results by ensuring
|
||||
better color consistency with the reference image and temporal consistency
|
||||
across video frames.
|
||||
|
||||
Reference: https://github.com/hahnec/color-matcher
|
||||
|
||||
Installation:
|
||||
pip install color-matcher
|
||||
|
||||
Methods available:
|
||||
- 'default' (hm-mvgd-hm): Best overall quality
|
||||
- 'mkl': Optimal transport-based, very accurate
|
||||
- 'reinhard': Classic color transfer, fast
|
||||
|
||||
Usage:
|
||||
from post_processing.color_matcher_post import ColorMatcherPost
|
||||
|
||||
post_processor = ColorMatcherPost(method='hm-mvgd-hm')
|
||||
|
||||
# Single image
|
||||
refined = post_processor.match_to_reference(colorized, reference)
|
||||
|
||||
# Video frames
|
||||
refined_video = post_processor.match_video_frames(video_frames, reference)
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from typing import Optional, Literal, List
|
||||
|
||||
|
||||
class ColorMatcherPost:
|
||||
"""Post-processing using color-matcher library for color refinement."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
method: Literal['default', 'hm-mvgd-hm', 'mkl', 'reinhard'] = 'default',
|
||||
temporal_consistency: bool = True,
|
||||
consistency_weight: float = 0.3,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
method: Color matching method
|
||||
- 'default' or 'hm-mvgd-hm': Hybrid method (best quality)
|
||||
- 'mkl': Monge-Kantorovich linearization (optimal transport)
|
||||
- 'reinhard': Classic Reinhard et al. method (fastest)
|
||||
temporal_consistency: Apply temporal smoothing across frames
|
||||
consistency_weight: Weight for temporal consistency (0-1)
|
||||
- 0 = no temporal smoothing
|
||||
- 1 = maximum temporal smoothing
|
||||
"""
|
||||
try:
|
||||
from color_matcher import ColorMatcher
|
||||
from color_matcher.io_handler import load_img_file, save_img_file
|
||||
self.cm = ColorMatcher()
|
||||
self.io_handler = (load_img_file, save_img_file)
|
||||
self.available = True
|
||||
except ImportError:
|
||||
print("[ColorMatcherPost] Warning: color-matcher not installed")
|
||||
print("Install with: pip install color-matcher")
|
||||
self.available = False
|
||||
return
|
||||
|
||||
# Normalize method name
|
||||
if method == 'default':
|
||||
method = 'hm-mvgd-hm'
|
||||
|
||||
self.method = method
|
||||
self.temporal_consistency = temporal_consistency
|
||||
self.consistency_weight = consistency_weight
|
||||
|
||||
# Validate method
|
||||
valid_methods = ['hm-mvgd-hm', 'mkl', 'reinhard']
|
||||
if method not in valid_methods:
|
||||
raise ValueError(f"Invalid method: {method}. Valid: {valid_methods}")
|
||||
|
||||
print(f"[ColorMatcherPost] Initialized with method={method}, "
|
||||
f"temporal_consistency={temporal_consistency}")
|
||||
|
||||
def _tensor_to_numpy(self, tensor: torch.Tensor) -> np.ndarray:
|
||||
"""Convert PyTorch tensor to numpy array for color-matcher.
|
||||
|
||||
Args:
|
||||
tensor: Image tensor in range [0, 1], shape [H, W, 3] or [B, H, W, 3]
|
||||
|
||||
Returns:
|
||||
Numpy array in range [0, 255], shape [H, W, 3] or [B, H, W, 3]
|
||||
"""
|
||||
if isinstance(tensor, torch.Tensor):
|
||||
array = tensor.detach().cpu().numpy()
|
||||
else:
|
||||
array = tensor
|
||||
|
||||
# Ensure range [0, 1]
|
||||
array = np.clip(array, 0, 1)
|
||||
|
||||
# Convert to uint8 [0, 255]
|
||||
array = (array * 255).astype(np.uint8)
|
||||
|
||||
return array
|
||||
|
||||
def _numpy_to_tensor(self, array: np.ndarray, device: str = 'cpu') -> torch.Tensor:
|
||||
"""Convert numpy array back to PyTorch tensor.
|
||||
|
||||
Args:
|
||||
array: Numpy array in range [0, 255], shape [H, W, 3]
|
||||
device: Target device
|
||||
|
||||
Returns:
|
||||
Tensor in range [0, 1], shape [H, W, 3]
|
||||
"""
|
||||
# Ensure uint8
|
||||
if array.dtype != np.uint8:
|
||||
array = np.clip(array, 0, 255).astype(np.uint8)
|
||||
|
||||
# Convert to float [0, 1]
|
||||
tensor = torch.from_numpy(array).float() / 255.0
|
||||
|
||||
return tensor.to(device)
|
||||
|
||||
def match_to_reference(
|
||||
self,
|
||||
image: torch.Tensor,
|
||||
reference: torch.Tensor,
|
||||
strength: float = 1.0,
|
||||
) -> torch.Tensor:
|
||||
"""Match image colors to reference image.
|
||||
|
||||
Args:
|
||||
image: Colorized image [H, W, 3] in range [0, 1]
|
||||
reference: Reference image [H, W, 3] in range [0, 1]
|
||||
strength: Blending strength (0-1)
|
||||
- 0 = original image
|
||||
- 1 = full color matching
|
||||
|
||||
Returns:
|
||||
Color-matched image [H, W, 3] in range [0, 1]
|
||||
"""
|
||||
if not self.available:
|
||||
print("[ColorMatcherPost] Skipping: color-matcher not available")
|
||||
return image
|
||||
|
||||
device = image.device
|
||||
|
||||
# Convert to numpy
|
||||
img_np = self._tensor_to_numpy(image)
|
||||
ref_np = self._tensor_to_numpy(reference)
|
||||
|
||||
# Apply color matching
|
||||
try:
|
||||
matched_np = self.cm.transfer(src=img_np, ref=ref_np, method=self.method)
|
||||
except Exception as e:
|
||||
print(f"[ColorMatcherPost] Color matching failed: {e}")
|
||||
return image
|
||||
|
||||
# Convert back to tensor
|
||||
matched = self._numpy_to_tensor(matched_np, device=device)
|
||||
|
||||
# Blend with original based on strength
|
||||
if strength < 1.0:
|
||||
matched = image * (1 - strength) + matched * strength
|
||||
|
||||
return matched
|
||||
|
||||
def match_video_frames(
|
||||
self,
|
||||
frames: torch.Tensor,
|
||||
reference: torch.Tensor,
|
||||
strength: float = 1.0,
|
||||
progress_callback: Optional[callable] = None,
|
||||
) -> torch.Tensor:
|
||||
"""Match colors for all video frames.
|
||||
|
||||
Args:
|
||||
frames: Video frames [N, H, W, 3] in range [0, 1]
|
||||
reference: Reference image [H, W, 3] in range [0, 1]
|
||||
strength: Color matching strength (0-1)
|
||||
progress_callback: Optional callback(current, total)
|
||||
|
||||
Returns:
|
||||
Color-matched frames [N, H, W, 3]
|
||||
"""
|
||||
if not self.available:
|
||||
print("[ColorMatcherPost] Skipping: color-matcher not available")
|
||||
return frames
|
||||
|
||||
num_frames = frames.shape[0]
|
||||
device = frames.device
|
||||
matched_frames = []
|
||||
|
||||
# Process each frame
|
||||
for i in range(num_frames):
|
||||
frame = frames[i]
|
||||
|
||||
# Match to reference
|
||||
matched = self.match_to_reference(frame, reference, strength=strength)
|
||||
|
||||
# Apply temporal consistency
|
||||
if self.temporal_consistency and i > 0 and len(matched_frames) > 0:
|
||||
prev_frame = matched_frames[-1]
|
||||
matched = self._apply_temporal_smoothing(matched, prev_frame)
|
||||
|
||||
matched_frames.append(matched)
|
||||
|
||||
if progress_callback is not None:
|
||||
progress_callback(i + 1, num_frames)
|
||||
|
||||
return torch.stack(matched_frames, dim=0)
|
||||
|
||||
def _apply_temporal_smoothing(
|
||||
self,
|
||||
current: torch.Tensor,
|
||||
previous: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""Apply temporal smoothing between consecutive frames.
|
||||
|
||||
Args:
|
||||
current: Current frame [H, W, 3]
|
||||
previous: Previous frame [H, W, 3]
|
||||
|
||||
Returns:
|
||||
Smoothed current frame [H, W, 3]
|
||||
"""
|
||||
# Blend with previous frame
|
||||
alpha = self.consistency_weight
|
||||
smoothed = current * (1 - alpha) + previous * alpha
|
||||
|
||||
return smoothed
|
||||
|
||||
def match_video_to_first_frame(
|
||||
self,
|
||||
frames: torch.Tensor,
|
||||
strength: float = 0.5,
|
||||
progress_callback: Optional[callable] = None,
|
||||
) -> torch.Tensor:
|
||||
"""Match all frames to first frame for temporal consistency.
|
||||
|
||||
This is useful when the first frame is already well-colorized.
|
||||
|
||||
Args:
|
||||
frames: Video frames [N, H, W, 3]
|
||||
strength: Matching strength (0-1)
|
||||
progress_callback: Optional callback(current, total)
|
||||
|
||||
Returns:
|
||||
Temporally consistent frames [N, H, W, 3]
|
||||
"""
|
||||
if not self.available:
|
||||
return frames
|
||||
|
||||
# Use first frame as reference
|
||||
reference = frames[0]
|
||||
|
||||
return self.match_video_frames(
|
||||
frames,
|
||||
reference,
|
||||
strength=strength,
|
||||
progress_callback=progress_callback
|
||||
)
|
||||
|
||||
def match_histogram_only(
|
||||
self,
|
||||
image: torch.Tensor,
|
||||
reference: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""Simple histogram matching (faster than full color matching).
|
||||
|
||||
Args:
|
||||
image: Input image [H, W, 3]
|
||||
reference: Reference image [H, W, 3]
|
||||
|
||||
Returns:
|
||||
Histogram-matched image [H, W, 3]
|
||||
"""
|
||||
from skimage import exposure
|
||||
|
||||
device = image.device
|
||||
img_np = self._tensor_to_numpy(image)
|
||||
ref_np = self._tensor_to_numpy(reference)
|
||||
|
||||
# Match histogram per channel
|
||||
matched = np.zeros_like(img_np)
|
||||
for c in range(3):
|
||||
matched[:, :, c] = exposure.match_histograms(
|
||||
img_np[:, :, c],
|
||||
ref_np[:, :, c]
|
||||
)
|
||||
|
||||
return self._numpy_to_tensor(matched, device=device)
|
||||
|
||||
|
||||
class ColorMatcherPostBatch:
|
||||
"""Batch processing version for better performance."""
|
||||
|
||||
def __init__(self, method: str = 'default', batch_size: int = 8):
|
||||
"""
|
||||
Args:
|
||||
method: Color matching method
|
||||
batch_size: Number of frames to process at once
|
||||
"""
|
||||
self.matcher = ColorMatcherPost(method=method)
|
||||
self.batch_size = batch_size
|
||||
|
||||
def process_video(
|
||||
self,
|
||||
frames: torch.Tensor,
|
||||
reference: torch.Tensor,
|
||||
strength: float = 1.0,
|
||||
) -> torch.Tensor:
|
||||
"""Process video in batches for better performance.
|
||||
|
||||
Args:
|
||||
frames: Video frames [N, H, W, 3]
|
||||
reference: Reference image [H, W, 3]
|
||||
strength: Color matching strength
|
||||
|
||||
Returns:
|
||||
Processed frames [N, H, W, 3]
|
||||
"""
|
||||
num_frames = frames.shape[0]
|
||||
all_matched = []
|
||||
|
||||
for i in range(0, num_frames, self.batch_size):
|
||||
batch_end = min(i + self.batch_size, num_frames)
|
||||
batch = frames[i:batch_end]
|
||||
|
||||
# Process batch
|
||||
matched_batch = self.matcher.match_video_frames(
|
||||
batch,
|
||||
reference,
|
||||
strength=strength
|
||||
)
|
||||
|
||||
all_matched.append(matched_batch)
|
||||
|
||||
return torch.cat(all_matched, dim=0)
|
||||
|
||||
|
||||
def compare_methods(
|
||||
image: torch.Tensor,
|
||||
reference: torch.Tensor,
|
||||
) -> dict:
|
||||
"""Compare all color matching methods.
|
||||
|
||||
Args:
|
||||
image: Test image [H, W, 3]
|
||||
reference: Reference image [H, W, 3]
|
||||
|
||||
Returns:
|
||||
Dictionary with results for each method
|
||||
"""
|
||||
methods = ['hm-mvgd-hm', 'mkl', 'reinhard']
|
||||
results = {}
|
||||
|
||||
for method in methods:
|
||||
try:
|
||||
matcher = ColorMatcherPost(method=method)
|
||||
matched = matcher.match_to_reference(image, reference)
|
||||
results[method] = matched
|
||||
print(f"✓ {method} completed")
|
||||
except Exception as e:
|
||||
print(f"✗ {method} failed: {e}")
|
||||
results[method] = None
|
||||
|
||||
return results
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("Testing ColorMatcherPost...")
|
||||
|
||||
# Create test images
|
||||
test_img = torch.rand(256, 256, 3)
|
||||
ref_img = torch.rand(256, 256, 3)
|
||||
|
||||
# Test single image
|
||||
print("\n1. Testing single image matching...")
|
||||
matcher = ColorMatcherPost(method='mkl')
|
||||
matched = matcher.match_to_reference(test_img, ref_img, strength=1.0)
|
||||
print(f" Input shape: {test_img.shape}")
|
||||
print(f" Output shape: {matched.shape}")
|
||||
print(f" Output range: [{matched.min():.3f}, {matched.max():.3f}]")
|
||||
|
||||
# Test video
|
||||
print("\n2. Testing video matching...")
|
||||
test_video = torch.rand(10, 256, 256, 3)
|
||||
matched_video = matcher.match_video_frames(test_video, ref_img)
|
||||
print(f" Input: {test_video.shape}")
|
||||
print(f" Output: {matched_video.shape}")
|
||||
|
||||
# Compare methods
|
||||
print("\n3. Comparing all methods...")
|
||||
results = compare_methods(test_img, ref_img)
|
||||
print(f" Successfully tested: {len([r for r in results.values() if r is not None])}/{len(results)} methods")
|
||||
|
||||
print("\n✓ All tests passed!")
|
||||
@@ -1,103 +0,0 @@
|
||||
"""
|
||||
Guided Filter Post-Processing
|
||||
|
||||
Fast edge-preserving filter that's much faster than bilateral filtering
|
||||
while achieving similar quality.
|
||||
|
||||
Reference: Guided Image Filtering, He et al., ECCV 2010
|
||||
"""
|
||||
|
||||
import torch
|
||||
import cv2
|
||||
import numpy as np
|
||||
from typing import Optional
|
||||
|
||||
|
||||
class GuidedFilter:
|
||||
"""Guided image filtering for edge-preserving smoothing."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
radius: int = 8,
|
||||
eps: float = 0.01,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
radius: Filter radius (larger = smoother)
|
||||
eps: Regularization parameter (larger = smoother)
|
||||
"""
|
||||
self.radius = radius
|
||||
self.eps = eps
|
||||
|
||||
def match_to_reference(
|
||||
self,
|
||||
image: torch.Tensor,
|
||||
reference: Optional[torch.Tensor] = None,
|
||||
guide: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
"""Apply guided filtering.
|
||||
|
||||
Args:
|
||||
image: Input image [H, W, 3] in range [0, 1]
|
||||
reference: Not used (for interface compatibility)
|
||||
guide: Guide image (uses input if None)
|
||||
|
||||
Returns:
|
||||
Filtered image [H, W, 3]
|
||||
"""
|
||||
device = image.device
|
||||
|
||||
# Convert to numpy [0, 255]
|
||||
img_np = (image.cpu().numpy() * 255).astype(np.uint8)
|
||||
|
||||
if guide is None:
|
||||
guide_np = img_np
|
||||
else:
|
||||
guide_np = (guide.cpu().numpy() * 255).astype(np.uint8)
|
||||
|
||||
# Apply guided filter
|
||||
try:
|
||||
filtered_np = cv2.ximgproc.guidedFilter(
|
||||
guide_np,
|
||||
img_np,
|
||||
self.radius,
|
||||
self.eps
|
||||
)
|
||||
except AttributeError:
|
||||
# Fallback: opencv-contrib not available, use bilateral
|
||||
print("[GuidedFilter] opencv-contrib not available, using bilateral filter")
|
||||
filtered_np = cv2.bilateralFilter(img_np, self.radius, 75, 75)
|
||||
except Exception as e:
|
||||
print(f"[GuidedFilter] Filtering failed: {e}")
|
||||
return image
|
||||
|
||||
# Convert back
|
||||
filtered = torch.from_numpy(filtered_np).float() / 255.0
|
||||
return filtered.to(device)
|
||||
|
||||
def match_video_frames(
|
||||
self,
|
||||
frames: torch.Tensor,
|
||||
reference: Optional[torch.Tensor] = None,
|
||||
progress_callback: Optional[callable] = None,
|
||||
) -> torch.Tensor:
|
||||
"""Apply guided filtering to video."""
|
||||
num_frames = frames.shape[0]
|
||||
filtered_frames = []
|
||||
|
||||
for i in range(num_frames):
|
||||
filtered = self.match_to_reference(frames[i], reference)
|
||||
filtered_frames.append(filtered)
|
||||
|
||||
if progress_callback:
|
||||
progress_callback(i + 1, num_frames)
|
||||
|
||||
return torch.stack(filtered_frames, dim=0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("Testing GuidedFilter...")
|
||||
test_img = torch.rand(256, 256, 3)
|
||||
gf = GuidedFilter(radius=8, eps=0.01)
|
||||
filtered = gf.match_to_reference(test_img)
|
||||
print(f"✓ Filtered: {test_img.shape} -> {filtered.shape}")
|
||||
@@ -1,142 +0,0 @@
|
||||
"""
|
||||
WLS (Weighted Least Squares) Filter Post-Processing
|
||||
|
||||
Edge-aware smoothing filter that preserves edges while smoothing color transitions.
|
||||
Currently implements the existing WLS filter from the original code.
|
||||
|
||||
Reference: Edge-Preserving Decompositions for Multi-Scale Tone and Detail Manipulation
|
||||
Farbman et al., SIGGRAPH 2008
|
||||
"""
|
||||
|
||||
import torch
|
||||
import cv2
|
||||
import numpy as np
|
||||
from typing import Optional
|
||||
|
||||
|
||||
class WLSFilter:
|
||||
"""Weighted Least Squares edge-aware filter."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
lambda_value: float = 500.0,
|
||||
sigma_color: float = 4.0,
|
||||
available: bool = None,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
lambda_value: Smoothing strength (higher = smoother)
|
||||
sigma_color: Color sensitivity
|
||||
available: Whether WLS filter is available (auto-detect if None)
|
||||
"""
|
||||
self.lambda_value = lambda_value
|
||||
self.sigma_color = sigma_color
|
||||
|
||||
# Check if WLS filter is available
|
||||
if available is None:
|
||||
self.available = hasattr(cv2.ximgproc, 'createFastGlobalSmootherFilter')
|
||||
else:
|
||||
self.available = available
|
||||
|
||||
if not self.available:
|
||||
print("[WLSFilter] Warning: opencv-contrib-python not installed or WLS not available")
|
||||
print("Install with: pip install opencv-contrib-python")
|
||||
|
||||
def match_to_reference(
|
||||
self,
|
||||
image: torch.Tensor,
|
||||
reference: Optional[torch.Tensor] = None,
|
||||
guide: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
"""Apply WLS filtering.
|
||||
|
||||
Args:
|
||||
image: Input image [H, W, 3] in range [0, 1]
|
||||
reference: Not used (for interface compatibility)
|
||||
guide: Optional guide image for edge-aware filtering (uses image if None)
|
||||
|
||||
Returns:
|
||||
Filtered image [H, W, 3]
|
||||
"""
|
||||
if not self.available:
|
||||
print("[WLSFilter] Skipping: WLS filter not available")
|
||||
return image
|
||||
|
||||
device = image.device
|
||||
|
||||
# Convert to numpy
|
||||
img_np = (image.cpu().numpy() * 255).astype(np.uint8)
|
||||
|
||||
if guide is None:
|
||||
# Use luminance as guide
|
||||
guide_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)
|
||||
else:
|
||||
guide_np = (guide.cpu().numpy() * 255).astype(np.uint8)
|
||||
if guide_np.ndim == 3:
|
||||
guide_np = cv2.cvtColor(guide_np, cv2.COLOR_RGB2GRAY)
|
||||
|
||||
# Apply WLS filter
|
||||
try:
|
||||
wls = cv2.ximgproc.createFastGlobalSmootherFilter(
|
||||
guide_np,
|
||||
self.lambda_value,
|
||||
self.sigma_color
|
||||
)
|
||||
|
||||
# Filter each channel
|
||||
filtered_np = np.zeros_like(img_np)
|
||||
for c in range(3):
|
||||
filtered_np[:, :, c] = wls.filter(img_np[:, :, c])
|
||||
|
||||
except Exception as e:
|
||||
print(f"[WLSFilter] Filtering failed: {e}")
|
||||
return image
|
||||
|
||||
# Convert back to tensor
|
||||
filtered = torch.from_numpy(filtered_np).float() / 255.0
|
||||
return filtered.to(device)
|
||||
|
||||
def match_video_frames(
|
||||
self,
|
||||
frames: torch.Tensor,
|
||||
reference: Optional[torch.Tensor] = None,
|
||||
progress_callback: Optional[callable] = None,
|
||||
) -> torch.Tensor:
|
||||
"""Apply WLS filtering to video frames.
|
||||
|
||||
Args:
|
||||
frames: Video frames [N, H, W, 3]
|
||||
reference: Not used (for interface compatibility)
|
||||
progress_callback: Optional callback(current, total)
|
||||
|
||||
Returns:
|
||||
Filtered frames [N, H, W, 3]
|
||||
"""
|
||||
if not self.available:
|
||||
return frames
|
||||
|
||||
num_frames = frames.shape[0]
|
||||
filtered_frames = []
|
||||
|
||||
for i in range(num_frames):
|
||||
filtered = self.match_to_reference(frames[i], reference)
|
||||
filtered_frames.append(filtered)
|
||||
|
||||
if progress_callback is not None:
|
||||
progress_callback(i + 1, num_frames)
|
||||
|
||||
return torch.stack(filtered_frames, dim=0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("Testing WLSFilter...")
|
||||
|
||||
# Test
|
||||
test_img = torch.rand(256, 256, 3)
|
||||
wls = WLSFilter(lambda_value=500.0, sigma_color=4.0)
|
||||
|
||||
if wls.available:
|
||||
filtered = wls.match_to_reference(test_img)
|
||||
print(f"✓ Filtered: {test_img.shape} -> {filtered.shape}")
|
||||
else:
|
||||
print("✗ WLS filter not available")
|
||||
@@ -1,186 +0,0 @@
|
||||
from __future__ import print_function
|
||||
|
||||
import argparse
|
||||
import glob
|
||||
import os
|
||||
import time
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.backends.cudnn as cudnn
|
||||
import torchvision.transforms as transform_lib
|
||||
from PIL import Image
|
||||
from tqdm import tqdm
|
||||
|
||||
import lib.TestTransforms as transforms
|
||||
from models.ColorVidNet import ColorVidNet
|
||||
from models.FrameColor import frame_colorization
|
||||
from models.NonlocalNet import VGG19_pytorch, WarpNet
|
||||
from utils.util import (batch_lab2rgb_transpose_mc, folder2vid, mkdir_if_not,
|
||||
save_frames, tensor_lab2rgb, uncenter_l)
|
||||
from utils.util_distortion import CenterPad, Normalize, RGB2Lab, ToTensor
|
||||
|
||||
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
|
||||
torch.cuda.set_device(0)
|
||||
|
||||
|
||||
def colorize_video(opt, input_path, reference_file, output_path, nonlocal_net, colornet, vggnet):
|
||||
# parameters for wls filter
|
||||
wls_filter_on = True
|
||||
lambda_value = 500
|
||||
sigma_color = 4
|
||||
|
||||
# processing folders
|
||||
mkdir_if_not(output_path)
|
||||
files = glob.glob(output_path + "*")
|
||||
print("processing the folder:", input_path)
|
||||
path, dirs, filenames = os.walk(input_path).__next__()
|
||||
file_count = len(filenames)
|
||||
filenames.sort(key=lambda f: int("".join(filter(str.isdigit, f) or -1)))
|
||||
|
||||
# NOTE: resize frames to 216*384
|
||||
transform = transforms.Compose(
|
||||
[CenterPad(opt.image_size), transform_lib.CenterCrop(opt.image_size), RGB2Lab(), ToTensor(), Normalize()]
|
||||
)
|
||||
|
||||
# if frame propagation: use the first frame as reference
|
||||
# otherwise, use the specified reference image
|
||||
ref_name = input_path + filenames[0] if opt.frame_propagate else reference_file
|
||||
print("reference name:", ref_name)
|
||||
frame_ref = Image.open(ref_name)
|
||||
|
||||
total_time = 0
|
||||
I_last_lab_predict = None
|
||||
|
||||
IB_lab_large = transform(frame_ref).unsqueeze(0).cuda()
|
||||
IB_lab = torch.nn.functional.interpolate(IB_lab_large, scale_factor=0.5, mode="bilinear")
|
||||
IB_l = IB_lab[:, 0:1, :, :]
|
||||
IB_ab = IB_lab[:, 1:3, :, :]
|
||||
with torch.no_grad():
|
||||
I_reference_lab = IB_lab
|
||||
I_reference_l = I_reference_lab[:, 0:1, :, :]
|
||||
I_reference_ab = I_reference_lab[:, 1:3, :, :]
|
||||
I_reference_rgb = tensor_lab2rgb(torch.cat((uncenter_l(I_reference_l), I_reference_ab), dim=1))
|
||||
features_B = vggnet(I_reference_rgb, ["r12", "r22", "r32", "r42", "r52"], preprocess=True)
|
||||
|
||||
for index, frame_name in enumerate(tqdm(filenames)):
|
||||
frame1 = Image.open(os.path.join(input_path, frame_name))
|
||||
IA_lab_large = transform(frame1).unsqueeze(0).cuda()
|
||||
IA_lab = torch.nn.functional.interpolate(IA_lab_large, scale_factor=0.5, mode="bilinear")
|
||||
|
||||
IA_l = IA_lab[:, 0:1, :, :]
|
||||
IA_ab = IA_lab[:, 1:3, :, :]
|
||||
|
||||
if I_last_lab_predict is None:
|
||||
if opt.frame_propagate:
|
||||
I_last_lab_predict = IB_lab
|
||||
else:
|
||||
I_last_lab_predict = torch.zeros_like(IA_lab).cuda()
|
||||
|
||||
# start the frame colorization
|
||||
with torch.no_grad():
|
||||
I_current_lab = IA_lab
|
||||
I_current_ab_predict, I_current_nonlocal_lab_predict, features_current_gray = frame_colorization(
|
||||
I_current_lab,
|
||||
I_reference_lab,
|
||||
I_last_lab_predict,
|
||||
features_B,
|
||||
vggnet,
|
||||
nonlocal_net,
|
||||
colornet,
|
||||
feature_noise=0,
|
||||
temperature=1e-10,
|
||||
)
|
||||
I_last_lab_predict = torch.cat((IA_l, I_current_ab_predict), dim=1)
|
||||
|
||||
# upsampling
|
||||
curr_bs_l = IA_lab_large[:, 0:1, :, :]
|
||||
curr_predict = (
|
||||
torch.nn.functional.interpolate(I_current_ab_predict.data.cpu(), scale_factor=2, mode="bilinear") * 1.25
|
||||
)
|
||||
|
||||
# filtering
|
||||
if wls_filter_on:
|
||||
guide_image = uncenter_l(curr_bs_l) * 255 / 100
|
||||
wls_filter = cv2.ximgproc.createFastGlobalSmootherFilter(
|
||||
guide_image[0, 0, :, :].cpu().numpy().astype(np.uint8), lambda_value, sigma_color
|
||||
)
|
||||
curr_predict_a = wls_filter.filter(curr_predict[0, 0, :, :].cpu().numpy())
|
||||
curr_predict_b = wls_filter.filter(curr_predict[0, 1, :, :].cpu().numpy())
|
||||
curr_predict_a = torch.from_numpy(curr_predict_a).unsqueeze(0).unsqueeze(0)
|
||||
curr_predict_b = torch.from_numpy(curr_predict_b).unsqueeze(0).unsqueeze(0)
|
||||
curr_predict_filter = torch.cat((curr_predict_a, curr_predict_b), dim=1)
|
||||
IA_predict_rgb = batch_lab2rgb_transpose_mc(curr_bs_l[:32], curr_predict_filter[:32, ...])
|
||||
else:
|
||||
IA_predict_rgb = batch_lab2rgb_transpose_mc(curr_bs_l[:32], curr_predict[:32, ...])
|
||||
|
||||
# save the frames
|
||||
save_frames(IA_predict_rgb, output_path, index)
|
||||
|
||||
# output video
|
||||
video_name = "video.avi"
|
||||
folder2vid(image_folder=output_path, output_dir=output_path, filename=video_name)
|
||||
print()
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--frame_propagate", default=False, type=bool, help="propagation mode, , please check the paper"
|
||||
)
|
||||
parser.add_argument("--image_size", type=int, default=[216 * 2, 384 * 2], help="the image size, eg. [216,384]")
|
||||
parser.add_argument("--cuda", action="store_false")
|
||||
parser.add_argument("--gpu_ids", type=str, default="0", help="separate by comma")
|
||||
parser.add_argument("--clip_path", type=str, default="./sample_videos/clips/v32", help="path of input clips")
|
||||
parser.add_argument("--ref_path", type=str, default="./sample_videos/ref/v32", help="path of refernce images")
|
||||
parser.add_argument("--output_path", type=str, default="./sample_videos/output", help="path of output clips")
|
||||
opt = parser.parse_args()
|
||||
opt.gpu_ids = [int(x) for x in opt.gpu_ids.split(",")]
|
||||
cudnn.benchmark = True
|
||||
print("running on GPU", opt.gpu_ids)
|
||||
|
||||
clip_name = opt.clip_path.split("/")[-1]
|
||||
refs = os.listdir(opt.ref_path)
|
||||
refs.sort()
|
||||
|
||||
nonlocal_net = WarpNet(1)
|
||||
colornet = ColorVidNet(7)
|
||||
vggnet = VGG19_pytorch()
|
||||
vggnet.load_state_dict(torch.load("data/vgg19_conv.pth"))
|
||||
for param in vggnet.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
nonlocal_test_path = os.path.join("checkpoints/", "video_moredata_l1/nonlocal_net_iter_76000.pth")
|
||||
color_test_path = os.path.join("checkpoints/", "video_moredata_l1/colornet_iter_76000.pth")
|
||||
print("succesfully load nonlocal model: ", nonlocal_test_path)
|
||||
print("succesfully load color model: ", color_test_path)
|
||||
nonlocal_net.load_state_dict(torch.load(nonlocal_test_path))
|
||||
colornet.load_state_dict(torch.load(color_test_path))
|
||||
|
||||
nonlocal_net.eval()
|
||||
colornet.eval()
|
||||
vggnet.eval()
|
||||
nonlocal_net.cuda()
|
||||
colornet.cuda()
|
||||
vggnet.cuda()
|
||||
|
||||
for ref_name in refs:
|
||||
try:
|
||||
colorize_video(
|
||||
opt,
|
||||
opt.clip_path,
|
||||
os.path.join(opt.ref_path, ref_name),
|
||||
os.path.join(opt.output_path, clip_name + "_" + ref_name.split(".")[0]),
|
||||
nonlocal_net,
|
||||
colornet,
|
||||
vggnet,
|
||||
)
|
||||
except Exception as error:
|
||||
print("error when colorizing the video " + ref_name)
|
||||
print(error)
|
||||
|
||||
video_name = "video.avi"
|
||||
clip_output_path = os.path.join(opt.output_path, clip_name)
|
||||
mkdir_if_not(clip_output_path)
|
||||
folder2vid(image_folder=opt.clip_path, output_dir=clip_output_path, filename=video_name)
|
||||
@@ -1,738 +0,0 @@
|
||||
from __future__ import print_function
|
||||
|
||||
import argparse
|
||||
import math
|
||||
import os
|
||||
import queue
|
||||
import time
|
||||
|
||||
import cv2
|
||||
import torch
|
||||
import torch.backends.cudnn as cudnn
|
||||
import torch.nn as nn
|
||||
import torch.optim as optim
|
||||
import torchvision.transforms as transform_lib
|
||||
from torch.utils.data import ConcatDataset, DataLoader, WeightedRandomSampler
|
||||
from torchvision.transforms import CenterCrop
|
||||
|
||||
import lib.TrainTransforms as transforms
|
||||
from lib.videoloader import VideosDataset
|
||||
from lib.videoloader_imagenet import VideosDataset_ImageNet
|
||||
from models.ColorVidNet import ColorVidNet
|
||||
from models.ContextualLoss import ContextualLoss, ContextualLoss_forward
|
||||
from models.FrameColor import frame_colorization
|
||||
from models.GAN_models import Discriminator_x64
|
||||
from models.NonlocalNet import (NonlocalWeightedAverage, VGG19_pytorch,
|
||||
WarpNet, WeightedAverage,
|
||||
WeightedAverage_color)
|
||||
from tensorboardX import SummaryWriter
|
||||
from utils.util import (batch_lab2rgb_transpose_mc, feature_normalize, l1_loss,
|
||||
mkdir_if_not, mse_loss, parse, tensor_lab2rgb,
|
||||
uncenter_l, weighted_l1_loss, weighted_mse_loss)
|
||||
from utils.util_distortion import (CenterPad_threshold, Normalize, RGB2Lab,
|
||||
ToTensor)
|
||||
from utils.util_tensorboard import TBImageRecorder, value_logger
|
||||
from utils.warping import WarpingLayer
|
||||
|
||||
cv2.setNumThreads(0)
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--data_root", default="/home/zhanbo/remote/video/video_pair3/", type=str)
|
||||
parser.add_argument("--data_root_imagenet", default="/home/zhanbo/remote/imagenet_ref_pair/", type=str)
|
||||
parser.add_argument("--gpu_ids", type=str, default="0,1,2,3", help="separate by comma")
|
||||
parser.add_argument("--workers", type=int, default=8)
|
||||
parser.add_argument("--batch_size", type=int, default=16)
|
||||
parser.add_argument("--image_size", type=int, default=[216, 384])
|
||||
parser.add_argument("--ic", type=int, default=7)
|
||||
parser.add_argument("--epoch", type=int, default=40)
|
||||
|
||||
parser.add_argument("--resume_epoch", type=int, default=0)
|
||||
parser.add_argument("--resume", type=bool, default=False)
|
||||
parser.add_argument("--load_pretrained_model", type=bool, default=True)
|
||||
|
||||
parser.add_argument("--lr", type=float, default=0.00001)
|
||||
parser.add_argument("--beta1", type=float, default=0.5)
|
||||
parser.add_argument("--lr_step", type=int, default=100)
|
||||
parser.add_argument("--lr_gamma", type=float, default=0.1)
|
||||
|
||||
|
||||
parser.add_argument("--checkpoint_dir", type=str, default="checkpoints/default")
|
||||
parser.add_argument("--tb_log_step", type=int, default=50)
|
||||
parser.add_argument("--print_step", type=int, default=2)
|
||||
|
||||
parser.add_argument("--real_reference_probability", type=float, default=0.7)
|
||||
parser.add_argument("--nonzero_placeholder_probability", type=float, default=0.0)
|
||||
parser.add_argument("--with_bad", type=bool, default=True)
|
||||
parser.add_argument("--with_mid", type=bool, default=True)
|
||||
|
||||
parser.add_argument("--domain_invariant", type=bool, default=False)
|
||||
parser.add_argument("--weigth_l1", type=float, default=2.0)
|
||||
parser.add_argument("--weight_contextual", type=float, default="0.2")
|
||||
parser.add_argument("--weight_perceptual", type=float, default="0.001")
|
||||
parser.add_argument("--weight_smoothness", type=float, default="5.0")
|
||||
parser.add_argument("--weight_gan", type=float, default="0.2")
|
||||
parser.add_argument("--weight_nonlocal_smoothness", type=float, default="0.0")
|
||||
parser.add_argument("--weight_nonlocal_consistent", type=float, default="0.0")
|
||||
parser.add_argument("--weight_consistent", type=float, default="0.02")
|
||||
parser.add_argument("--luminance_noise", type=float, default="2.0")
|
||||
parser.add_argument("--permute_data", type=bool, default=True)
|
||||
parser.add_argument("--contextual_loss_direction", type=str, default="forward", help="forward or backward matching")
|
||||
|
||||
|
||||
def image_logger_fn(
|
||||
I_last_lab,
|
||||
I_current_lab,
|
||||
I_reference_lab,
|
||||
I_last_lab_predict,
|
||||
I_current_lab_predict,
|
||||
I_last_nonlocal_lab,
|
||||
I_current_nonlocal_lab,
|
||||
):
|
||||
I_last_image = batch_lab2rgb_transpose_mc(I_last_lab[0:32, 0:1, :, :], I_last_lab[0:32, 1:3, :, :])
|
||||
I_current_image = batch_lab2rgb_transpose_mc(I_current_lab[0:32, 0:1, :, :], I_current_lab[0:32, 1:3, :, :])
|
||||
I_reference_image = batch_lab2rgb_transpose_mc(I_reference_lab[0:32, 0:1, :, :], I_reference_lab[0:32, 1:3, :, :])
|
||||
I_last_image_predict = batch_lab2rgb_transpose_mc(
|
||||
I_last_lab_predict[0:32, 0:1, :, :], I_last_lab_predict[0:32, 1:3, :, :]
|
||||
)
|
||||
I_current_image_predict = batch_lab2rgb_transpose_mc(
|
||||
I_current_lab_predict[0:32, 0:1, :, :], I_current_lab_predict[0:32, 1:3, :, :]
|
||||
)
|
||||
I_last_nonlocal_image = batch_lab2rgb_transpose_mc(
|
||||
I_last_nonlocal_lab[0:32, 0:1, :, :], I_last_nonlocal_lab[0:32, 1:3, :, :]
|
||||
)
|
||||
I_current_nonlocal_image = batch_lab2rgb_transpose_mc(
|
||||
I_current_nonlocal_lab[0:32, 0:1, :, :], I_current_nonlocal_lab[0:32, 1:3, :, :]
|
||||
)
|
||||
|
||||
img_info = {}
|
||||
img_info["1_I_last"] = I_last_image
|
||||
img_info["2_I_current"] = I_current_image
|
||||
img_info["3_I_reference"] = I_reference_image
|
||||
img_info["4_I_last_predict"] = I_last_image_predict
|
||||
img_info["5_I_curren_predict"] = I_current_image_predict
|
||||
img_info["6_I_last_nonlocal"] = I_last_nonlocal_image
|
||||
img_info["7_I_current_nonlocal"] = I_current_nonlocal_image
|
||||
|
||||
return img_info
|
||||
|
||||
|
||||
def training_logger():
|
||||
try:
|
||||
if total_iter % opt.print_step == 0:
|
||||
print("processing time:", elapsed)
|
||||
print(
|
||||
"Epoch %d, Step[%d/%d], lr: %f, total_loss: %.2f"
|
||||
% (
|
||||
epoch,
|
||||
((iter + 1) % iter_num_per_epoch),
|
||||
iter_num_per_epoch,
|
||||
step_optim_scheduler_g.get_last_lr()[0],
|
||||
total_loss.item(),
|
||||
)
|
||||
)
|
||||
|
||||
value_logger(
|
||||
tb_writer,
|
||||
total_iter,
|
||||
loss_info={
|
||||
"l1_loss": l1_loss.item(),
|
||||
"feat_loss": feat_loss.item(),
|
||||
"contextual_loss_total": contextual_loss_total.item(),
|
||||
"smoothness_loss": smoothness_loss.item(),
|
||||
"nonlocal_smoothness_loss": nonlocal_smoothness_loss.item(),
|
||||
"nonlocal_consistent_loss": nonlocal_consistent_loss.item(),
|
||||
"consistent_loss": consistent_loss.item(),
|
||||
"generator_loss": generator_loss.item(),
|
||||
"discriminator_loss": discriminator_loss.item(),
|
||||
"total_loss": total_loss.item(),
|
||||
},
|
||||
)
|
||||
|
||||
if total_iter % opt.tb_log_step == 0:
|
||||
I_last_nonlocal_lab = torch.cat(
|
||||
(I_last_lab[:, 0:1, :, :], I_last_nonlocal_lab_predict[:, 1:3, :, :]), dim=1
|
||||
)
|
||||
I_current_nonlocal_lab = torch.cat(
|
||||
(I_current_lab[:, 0:1, :, :], I_current_nonlocal_lab_predict[:, 1:3, :, :]), dim=1
|
||||
)
|
||||
I_last_lab_predict = torch.cat((I_last_l, I_last_ab_predict), dim=1)
|
||||
data_queue.put(
|
||||
(
|
||||
(
|
||||
I_last_lab.cpu(),
|
||||
I_current_lab.cpu(),
|
||||
I_reference_lab.cpu(),
|
||||
I_last_lab_predict.cpu(),
|
||||
I_current_lab_predict.cpu(),
|
||||
I_last_nonlocal_lab.cpu(),
|
||||
I_current_nonlocal_lab,
|
||||
),
|
||||
total_iter,
|
||||
)
|
||||
)
|
||||
|
||||
if total_iter % 2000 == 0:
|
||||
if len(opt.gpu_ids) > 1:
|
||||
torch.save(
|
||||
nonlocal_net.module.state_dict(),
|
||||
os.path.join(opt.checkpoint_dir, "nonlocal_net_iter_%d.pth") % total_iter,
|
||||
)
|
||||
torch.save(
|
||||
colornet.module.state_dict(),
|
||||
os.path.join(opt.checkpoint_dir, "colornet_iter_%d.pth") % total_iter,
|
||||
)
|
||||
torch.save(
|
||||
discriminator.module.state_dict(),
|
||||
os.path.join(opt.checkpoint_dir, "discriminator_iter_%d.pth") % total_iter,
|
||||
)
|
||||
else:
|
||||
torch.save(
|
||||
nonlocal_net.state_dict(),
|
||||
os.path.join(opt.checkpoint_dir, "nonlocal_net_iter_%d.pth") % total_iter,
|
||||
)
|
||||
torch.save(colornet.state_dict(), os.path.join(opt.checkpoint_dir, "colornet_iter_%d.pth") % total_iter)
|
||||
torch.save(
|
||||
discriminator.state_dict(),
|
||||
os.path.join(opt.checkpoint_dir, "discriminator_iter_%d.pth") % total_iter,
|
||||
)
|
||||
|
||||
# save the state for resume
|
||||
if total_iter % 2000 == 0:
|
||||
print("saving the checkpoint")
|
||||
if len(opt.gpu_ids) > 1:
|
||||
state = {
|
||||
"total_iter": total_iter + 1,
|
||||
"epoch": epoch,
|
||||
"colornet_state": colornet.module.state_dict(),
|
||||
"nonlocal_net_state": nonlocal_net.module.state_dict(),
|
||||
"discriminator_state": discriminator.module.state_dict(),
|
||||
"optimizer_g": optimizer_g.state_dict(),
|
||||
"optimizer_d": optimizer_d.state_dict(),
|
||||
"optimizer_schedule_g": step_optim_scheduler_g.state_dict(),
|
||||
"optimizer_schedule_d": step_optim_scheduler_g.state_dict(),
|
||||
}
|
||||
else:
|
||||
state = {
|
||||
"total_iter": total_iter + 1,
|
||||
"epoch": epoch,
|
||||
"colornet_state": colornet.state_dict(),
|
||||
"nonlocal_net_state": nonlocal_net.state_dict(),
|
||||
"discriminator_state": discriminator.state_dict(),
|
||||
"optimizer_g": optimizer_g.state_dict(),
|
||||
"optimizer_d": optimizer_d.state_dict(),
|
||||
"optimizer_schedule_g": step_optim_scheduler_g.state_dict(),
|
||||
"optimizer_schedule_d": step_optim_scheduler_d.state_dict(),
|
||||
}
|
||||
torch.save(state, os.path.join(opt.checkpoint_dir, "learning_checkpoint.pth"))
|
||||
except Exception as e:
|
||||
print("Exception during output")
|
||||
print(e)
|
||||
|
||||
|
||||
def gpu_setup():
|
||||
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
|
||||
cudnn.benchmark = True
|
||||
opt.gpu_ids = [int(x) for x in opt.gpu_ids.split(",")]
|
||||
torch.cuda.set_device(opt.gpu_ids[0])
|
||||
device = torch.device("cuda")
|
||||
print("running on GPU", opt.gpu_ids)
|
||||
return device
|
||||
|
||||
|
||||
def load_data():
|
||||
print("initializing dataloader")
|
||||
transforms_video = [
|
||||
CenterCrop(opt.image_size),
|
||||
RGB2Lab(),
|
||||
ToTensor(),
|
||||
Normalize(),
|
||||
]
|
||||
transforms_imagenet = [CenterPad_threshold(opt.image_size), RGB2Lab(), ToTensor(), Normalize()]
|
||||
extra_reference_transform = [
|
||||
transform_lib.RandomHorizontalFlip(0.5),
|
||||
transform_lib.RandomResizedCrop(480, (0.98, 1.0), ratio=(0.8, 1.2)),
|
||||
]
|
||||
train_dataset_video = VideosDataset(
|
||||
data_root=opt.data_root,
|
||||
epoch=opt.epoch,
|
||||
image_size=opt.image_size,
|
||||
image_transform=transforms.Compose(transforms_video),
|
||||
real_reference_probability=opt.real_reference_probability,
|
||||
nonzero_placeholder_probability=opt.nonzero_placeholder_probability,
|
||||
)
|
||||
train_dataset_imagenet = VideosDataset_ImageNet(
|
||||
data_root=opt.data_root_imagenet,
|
||||
image_size=opt.image_size,
|
||||
epoch=opt.epoch,
|
||||
with_bad=opt.with_bad,
|
||||
with_mid=opt.with_mid,
|
||||
transforms_imagenet=transforms_imagenet,
|
||||
distortion_level=4,
|
||||
brightnessjitter=5,
|
||||
nonzero_placeholder_probability=opt.nonzero_placeholder_probability,
|
||||
extra_reference_transform=extra_reference_transform,
|
||||
real_reference_probability=opt.real_reference_probability,
|
||||
)
|
||||
|
||||
video_training_length = len(train_dataset_video)
|
||||
imagenet_training_length = len(train_dataset_imagenet)
|
||||
dataset_training_length = train_dataset_video.real_len + train_dataset_imagenet.real_len
|
||||
dataset_combined = ConcatDataset([train_dataset_video, train_dataset_imagenet])
|
||||
sampler = WeightedRandomSampler(
|
||||
[1] * video_training_length + [1] * imagenet_training_length, dataset_training_length * opt.epoch
|
||||
)
|
||||
data_loader = DataLoader(
|
||||
dataset_combined,
|
||||
batch_size=opt.batch_size,
|
||||
shuffle=False,
|
||||
num_workers=opt.workers,
|
||||
pin_memory=True,
|
||||
drop_last=True,
|
||||
sampler=sampler,
|
||||
)
|
||||
return dataset_training_length, train_dataset_video, train_dataset_imagenet, data_loader
|
||||
|
||||
|
||||
def define_loss():
|
||||
print("defining loss")
|
||||
ab_criterion = nn.SmoothL1Loss().to(device)
|
||||
nonlocal_criterion = nn.SmoothL1Loss().to(device)
|
||||
feat_l2_criterion = nn.MSELoss().to(device)
|
||||
feat_l1_criterion = nn.SmoothL1Loss().to(device)
|
||||
contextual_loss = ContextualLoss().to(device)
|
||||
contextual_forward_loss = ContextualLoss_forward().to(device)
|
||||
BCE_stable = nn.BCEWithLogitsLoss().to(device)
|
||||
return contextual_loss, contextual_forward_loss
|
||||
|
||||
|
||||
def define_optimizer():
|
||||
print("defining optimizer")
|
||||
optimizer_g = optim.Adam(
|
||||
[{"params": nonlocal_net.parameters(), "lr": 1e-5}, {"params": colornet.parameters(), "lr": 2e-4}],
|
||||
betas=(0.5, 0.999),
|
||||
eps=1e-5,
|
||||
amsgrad=True,
|
||||
)
|
||||
optimizer_d = optim.Adam(
|
||||
filter(lambda p: p.requires_grad, discriminator.parameters()), lr=2 * 1e-4, betas=(0.5, 0.999)
|
||||
)
|
||||
return optimizer_g, optimizer_d
|
||||
|
||||
|
||||
def resume_model():
|
||||
print("resuming the learning")
|
||||
checkpoint = torch.load(os.path.join(opt.checkpoint_dir, "learning_checkpoint.pth"))
|
||||
total_iter = checkpoint["total_iter"]
|
||||
epoch = checkpoint["epoch"]
|
||||
colornet.load_state_dict(checkpoint["colornet_state"])
|
||||
nonlocal_net.load_state_dict(checkpoint["nonlocal_net_state"])
|
||||
discriminator.load_state_dict(checkpoint["discriminator_state"])
|
||||
optimizer_g.load_state_dict(checkpoint["optimizer_g"])
|
||||
optimizer_d.load_state_dict(checkpoint["optimizer_d"])
|
||||
step_optim_scheduler_g.load_state_dict(checkpoint["optimizer_schedule_g"])
|
||||
step_optim_scheduler_d.load_state_dict(checkpoint["optimizer_schedule_d"])
|
||||
|
||||
|
||||
def to_device(
|
||||
colornet,
|
||||
nonlocal_net,
|
||||
discriminator,
|
||||
vggnet,
|
||||
contextual_loss,
|
||||
contextual_forward_loss,
|
||||
weighted_layer_color,
|
||||
nonlocal_weighted_layer,
|
||||
warping_layer,
|
||||
instancenorm,
|
||||
):
|
||||
print("moving models to device")
|
||||
colornet = torch.nn.DataParallel(colornet.to(device), device_ids=opt.gpu_ids)
|
||||
nonlocal_net = torch.nn.DataParallel(nonlocal_net.to(device), device_ids=opt.gpu_ids)
|
||||
discriminator = torch.nn.DataParallel(discriminator.to(device), device_ids=opt.gpu_ids)
|
||||
vggnet = torch.nn.DataParallel(vggnet.to(device), device_ids=opt.gpu_ids)
|
||||
contextual_loss = torch.nn.DataParallel(contextual_loss.to(device), device_ids=opt.gpu_ids)
|
||||
contextual_forward_loss = torch.nn.DataParallel(contextual_forward_loss.to(device), device_ids=opt.gpu_ids)
|
||||
weighted_layer_color = torch.nn.DataParallel(weighted_layer_color.to(device), device_ids=opt.gpu_ids)
|
||||
nonlocal_weighted_layer = torch.nn.DataParallel(nonlocal_weighted_layer.to(device), device_ids=opt.gpu_ids)
|
||||
warping_layer = torch.nn.DataParallel(warping_layer.to(device), device_ids=opt.gpu_ids)
|
||||
instancenorm = torch.nn.DataParallel(instancenorm.to(device), device_ids=opt.gpu_ids)
|
||||
|
||||
return (
|
||||
vggnet,
|
||||
nonlocal_net,
|
||||
colornet,
|
||||
discriminator,
|
||||
instancenorm,
|
||||
contextual_loss,
|
||||
contextual_forward_loss,
|
||||
weighted_layer_color,
|
||||
nonlocal_weighted_layer,
|
||||
warping_layer,
|
||||
)
|
||||
|
||||
|
||||
def loss_init():
|
||||
print("initializing losses")
|
||||
zero_loss = torch.Tensor([0]).to(device)
|
||||
(
|
||||
feat_loss,
|
||||
contextual_loss_total,
|
||||
smoothness_loss,
|
||||
nonlocal_smoothness_loss,
|
||||
consistent_loss,
|
||||
nonlocal_consistent_loss,
|
||||
generator_loss,
|
||||
discriminator_loss,
|
||||
) = (zero_loss, zero_loss, zero_loss, zero_loss, zero_loss, zero_loss, zero_loss, zero_loss)
|
||||
|
||||
return (
|
||||
feat_loss,
|
||||
contextual_loss_total,
|
||||
smoothness_loss,
|
||||
nonlocal_smoothness_loss,
|
||||
consistent_loss,
|
||||
nonlocal_consistent_loss,
|
||||
generator_loss,
|
||||
discriminator_loss,
|
||||
)
|
||||
|
||||
|
||||
def video_colorization():
|
||||
# colorization for the last frame
|
||||
I_last_ab_predict, I_last_nonlocal_lab_predict, features_last_gray = frame_colorization(
|
||||
I_last_lab,
|
||||
I_reference_lab,
|
||||
placeholder_lab,
|
||||
features_B,
|
||||
vggnet,
|
||||
nonlocal_net,
|
||||
colornet,
|
||||
feature_noise=0,
|
||||
luminance_noise=opt.luminance_noise,
|
||||
)
|
||||
I_last_lab_predict = torch.cat((I_last_l, I_last_ab_predict), dim=1)
|
||||
|
||||
# colorization for the current frame
|
||||
I_current_ab_predict, I_current_nonlocal_lab_predict, features_current_gray = frame_colorization(
|
||||
I_current_lab,
|
||||
I_reference_lab,
|
||||
I_last_lab_predict,
|
||||
features_B,
|
||||
vggnet,
|
||||
nonlocal_net,
|
||||
colornet,
|
||||
feature_noise=0,
|
||||
luminance_noise=opt.luminance_noise,
|
||||
)
|
||||
I_current_lab_predict = torch.cat((I_last_l, I_current_ab_predict), dim=1)
|
||||
return I_current_ab_predict, I_last_ab_predict, I_current_nonlocal_lab_predict, I_last_nonlocal_lab_predict
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
torch.multiprocessing.set_start_method("spawn", force=True)
|
||||
opt = parse(parser)
|
||||
opt.data_root = opt.data_root.split(",")[0]
|
||||
opt.data_root_imagenet = opt.data_root_imagenet.split(",")[0]
|
||||
mkdir_if_not(opt.checkpoint_dir)
|
||||
mkdir_if_not("./runs/")
|
||||
|
||||
device = gpu_setup()
|
||||
dataset_training_length, train_dataset_video, train_dataset_imagenet, data_loader = load_data()
|
||||
tb_writer = SummaryWriter()
|
||||
data_queue = queue.Queue()
|
||||
tb_image_reorder = TBImageRecorder(tb_writer, image_logger_fn, data_queue)
|
||||
tb_image_reorder.start()
|
||||
|
||||
# define network
|
||||
nonlocal_net = WarpNet(opt.batch_size)
|
||||
colornet = ColorVidNet(opt.ic)
|
||||
discriminator = Discriminator_x64(ndf=64)
|
||||
colornet.to(device)
|
||||
nonlocal_net.to(device)
|
||||
discriminator.to(device)
|
||||
|
||||
weighted_layer = WeightedAverage()
|
||||
weighted_layer_color = WeightedAverage_color()
|
||||
nonlocal_weighted_layer = NonlocalWeightedAverage()
|
||||
instancenorm = nn.InstanceNorm2d(512, affine=False)
|
||||
warping_layer = WarpingLayer("gpu")
|
||||
|
||||
vggnet = VGG19_pytorch()
|
||||
vggnet.load_state_dict(torch.load("data/vgg19_conv.pth"))
|
||||
vggnet.eval()
|
||||
for param in vggnet.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
# load pre-trained model
|
||||
if opt.load_pretrained_model:
|
||||
nonlocal_pretain_path = os.path.join("checkpoints/video_moredata_l1/", "nonlocal_net_iter_76000.pth")
|
||||
nonlocal_net.load_state_dict(torch.load(nonlocal_pretain_path))
|
||||
color_test_path = "checkpoints/video_moredata_l1/" + "colornet_iter_76000.pth"
|
||||
colornet.load_state_dict(torch.load(color_test_path))
|
||||
|
||||
# define loss function
|
||||
contextual_loss, contextual_forward_loss = define_loss()
|
||||
|
||||
# define optimizer
|
||||
optimizer_g, optimizer_d = define_optimizer()
|
||||
step_optim_scheduler_g = optim.lr_scheduler.StepLR(optimizer_g, step_size=opt.lr_step, gamma=opt.lr_gamma)
|
||||
step_optim_scheduler_d = optim.lr_scheduler.StepLR(optimizer_d, step_size=opt.lr_step, gamma=opt.lr_gamma)
|
||||
|
||||
# define others
|
||||
downsampling_by2 = nn.AvgPool2d(kernel_size=2).to(device)
|
||||
downsampling_by4 = nn.AvgPool2d(kernel_size=4).to(device)
|
||||
|
||||
# dataset info
|
||||
iter_num_per_epoch = dataset_training_length // opt.batch_size
|
||||
total_iter = opt.resume_epoch * iter_num_per_epoch
|
||||
print(
|
||||
"train_dataset info, real_len: %d, epoch_len: %d, iter_num_per_epoch: %d"
|
||||
% (dataset_training_length, len(train_dataset_video) + len(train_dataset_imagenet), iter_num_per_epoch)
|
||||
)
|
||||
|
||||
if opt.resume:
|
||||
resume_model()
|
||||
|
||||
# move to GPU processing
|
||||
device = opt.gpu_ids[0]
|
||||
(
|
||||
vggnet,
|
||||
nonlocal_net,
|
||||
colornet,
|
||||
discriminator,
|
||||
instancenorm,
|
||||
contextual_loss,
|
||||
contextual_forward_loss,
|
||||
weighted_layer_color,
|
||||
nonlocal_weighted_layer,
|
||||
warping_layer,
|
||||
) = to_device(
|
||||
colornet,
|
||||
nonlocal_net,
|
||||
discriminator,
|
||||
vggnet,
|
||||
contextual_loss,
|
||||
contextual_forward_loss,
|
||||
weighted_layer_color,
|
||||
nonlocal_weighted_layer,
|
||||
warping_layer,
|
||||
instancenorm,
|
||||
)
|
||||
|
||||
(
|
||||
feat_loss,
|
||||
contextual_loss_total,
|
||||
smoothness_loss,
|
||||
nonlocal_smoothness_loss,
|
||||
consistent_loss,
|
||||
nonlocal_consistent_loss,
|
||||
generator_loss,
|
||||
discriminator_loss,
|
||||
) = loss_init()
|
||||
|
||||
# %% Training
|
||||
print("start training")
|
||||
for iter, data in enumerate(data_loader):
|
||||
start_time = time.time()
|
||||
total_iter += 1
|
||||
epoch = math.ceil(total_iter / iter_num_per_epoch)
|
||||
|
||||
###### LOADING DATA SAMPLE ######
|
||||
(
|
||||
I_last_lab,
|
||||
I_current_lab,
|
||||
I_reference_lab,
|
||||
flow_forward,
|
||||
flow_backward,
|
||||
mask,
|
||||
placeholder_lab,
|
||||
self_ref_flag,
|
||||
) = data
|
||||
I_last_lab = I_last_lab.to(device)
|
||||
I_current_lab = I_current_lab.to(device)
|
||||
I_reference_lab = I_reference_lab.to(device)
|
||||
flow_forward = flow_forward.to(device)
|
||||
flow_backward = flow_backward.to(device)
|
||||
mask = mask.to(device)
|
||||
placeholder_lab = placeholder_lab.to(device)
|
||||
self_ref_flag = self_ref_flag.to(device)
|
||||
|
||||
I_last_l = I_last_lab[:, 0:1, :, :]
|
||||
I_last_ab = I_last_lab[:, 1:3, :, :]
|
||||
I_current_l = I_current_lab[:, 0:1, :, :]
|
||||
I_current_ab = I_current_lab[:, 1:3, :, :]
|
||||
I_reference_l = I_reference_lab[:, 0:1, :, :]
|
||||
I_reference_ab = I_reference_lab[:, 1:3, :, :]
|
||||
I_reference_rgb = tensor_lab2rgb(torch.cat((uncenter_l(I_reference_l), I_reference_ab), dim=1))
|
||||
features_B = vggnet(I_reference_rgb, ["r12", "r22", "r32", "r42", "r52"], preprocess=True)
|
||||
|
||||
###### COLORIZATION ######
|
||||
(
|
||||
I_current_ab_predict,
|
||||
I_last_ab_predict,
|
||||
I_current_nonlocal_lab_predict,
|
||||
I_last_nonlocal_lab_predict,
|
||||
) = video_colorization()
|
||||
|
||||
###### UPDATE DISCRIMINATOR ######
|
||||
if opt.weight_gan > 0:
|
||||
optimizer_g.zero_grad()
|
||||
optimizer_d.zero_grad()
|
||||
|
||||
fake_data_lab = torch.cat(
|
||||
(uncenter_l(I_current_l), I_current_ab_predict, uncenter_l(I_last_l), I_last_ab_predict), dim=1
|
||||
)
|
||||
real_data_lab = torch.cat((uncenter_l(I_current_l), I_current_ab, uncenter_l(I_last_l), I_last_ab), dim=1)
|
||||
|
||||
if opt.permute_data:
|
||||
batch_index = torch.arange(-1, opt.batch_size - 1, dtype=torch.long)
|
||||
real_data_lab = real_data_lab[batch_index, ...]
|
||||
|
||||
y_pred_fake, feature_pred_fake = discriminator(fake_data_lab.detach())
|
||||
y_pred_real, feature_pred_real = discriminator(real_data_lab.detach())
|
||||
|
||||
y = torch.ones_like(y_pred_real)
|
||||
y2 = torch.zeros_like(y_pred_real)
|
||||
discriminator_loss = (
|
||||
torch.mean((y_pred_real - torch.mean(y_pred_fake) - y) ** 2)
|
||||
+ torch.mean((y_pred_fake - torch.mean(y_pred_real) + y) ** 2)
|
||||
) / 2
|
||||
discriminator_loss.backward()
|
||||
optimizer_d.step()
|
||||
|
||||
###### UPDATE GENERATOR ######
|
||||
optimizer_g.zero_grad()
|
||||
optimizer_d.zero_grad()
|
||||
|
||||
# extract vgg features for both output and original image
|
||||
I_predict_rgb = tensor_lab2rgb(torch.cat((uncenter_l(I_current_l), I_current_ab_predict), dim=1))
|
||||
predict_relu1_1, predict_relu2_1, predict_relu3_1, predict_relu4_1, predict_relu5_1 = vggnet(
|
||||
I_predict_rgb, ["r12", "r22", "r32", "r42", "r52"], preprocess=True
|
||||
)
|
||||
|
||||
I_current_rgb = tensor_lab2rgb(torch.cat((uncenter_l(I_current_l), I_current_ab), dim=1))
|
||||
A_relu1_1, A_relu2_1, A_relu3_1, A_relu4_1, A_relu5_1 = vggnet(
|
||||
I_current_rgb, ["r12", "r22", "r32", "r42", "r52"], preprocess=True
|
||||
)
|
||||
B_relu1_1, B_relu2_1, B_relu3_1, B_relu4_1, B_relu5_1 = features_B
|
||||
|
||||
###### LOSS COMPUTE ######
|
||||
# l1 loss
|
||||
if opt.weigth_l1 > 0:
|
||||
sample_weights = (self_ref_flag[:, 1:3, :, :]) / (sum(self_ref_flag[:, 0, 0, 0]) + 1e-5)
|
||||
l1_loss = weighted_l1_loss(I_current_ab_predict, I_current_ab, sample_weights) * opt.weigth_l1
|
||||
|
||||
# generator loss
|
||||
if opt.weight_gan > 0:
|
||||
y_pred_fake, feature_pred_fake = discriminator(fake_data_lab)
|
||||
y_pred_real, feature_pred_real = discriminator(real_data_lab)
|
||||
generator_loss = (
|
||||
(
|
||||
torch.mean((y_pred_real - torch.mean(y_pred_fake) + y) ** 2)
|
||||
+ torch.mean((y_pred_fake - torch.mean(y_pred_real) - y) ** 2)
|
||||
)
|
||||
/ 2
|
||||
* opt.weight_gan
|
||||
)
|
||||
|
||||
# feature loss
|
||||
if opt.domain_invariant:
|
||||
feat_loss = (
|
||||
mse_loss(instancenorm(predict_relu5_1), instancenorm(A_relu5_1.detach()))
|
||||
* opt.weight_perceptual
|
||||
* 1e5
|
||||
* 0.2
|
||||
)
|
||||
else:
|
||||
feat_loss = mse_loss(predict_relu5_1, A_relu5_1.detach()) * opt.weight_perceptual
|
||||
|
||||
# contextual loss
|
||||
if opt.contextual_loss_direction == "backward":
|
||||
contextual_style5_1 = torch.mean(contextual_loss(predict_relu5_1, B_relu5_1.detach())) * 8
|
||||
contextual_style4_1 = torch.mean(contextual_loss(predict_relu4_1, B_relu4_1.detach())) * 4
|
||||
contextual_style3_1 = (
|
||||
torch.mean(contextual_loss(downsampling_by2(predict_relu3_1), downsampling_by2(B_relu3_1.detach()))) * 2
|
||||
)
|
||||
else:
|
||||
contextual_style5_1 = torch.mean(contextual_forward_loss(predict_relu5_1, B_relu5_1.detach())) * 8
|
||||
contextual_style4_1 = torch.mean(contextual_forward_loss(predict_relu4_1, B_relu4_1.detach())) * 4
|
||||
contextual_style3_1 = (
|
||||
torch.mean(
|
||||
contextual_forward_loss(downsampling_by2(predict_relu3_1), downsampling_by2(B_relu3_1.detach()))
|
||||
)
|
||||
* 2
|
||||
)
|
||||
if opt.weight_contextual > 0:
|
||||
contextual_loss_total = (
|
||||
contextual_style5_1 + contextual_style4_1 + contextual_style3_1
|
||||
) * opt.weight_contextual
|
||||
|
||||
# smoothness loss
|
||||
if opt.weight_smoothness > 0:
|
||||
scale_factor = 1
|
||||
I_current_lab_predict = torch.cat((I_current_l, I_current_ab_predict), dim=1)
|
||||
IA_ab_weighed = weighted_layer_color(
|
||||
I_current_lab, I_current_lab_predict, patch_size=3, alpha=10, scale_factor=scale_factor
|
||||
)
|
||||
smoothness_loss = (
|
||||
mse_loss(nn.functional.interpolate(I_current_ab_predict, scale_factor=scale_factor), IA_ab_weighed)
|
||||
* opt.weight_smoothness
|
||||
)
|
||||
|
||||
if opt.weight_nonlocal_smoothness > 0:
|
||||
scale_factor = 0.25
|
||||
alpha_nonlocal_smoothness = 0.5
|
||||
nonlocal_smooth_feature = feature_normalize(A_relu2_1)
|
||||
I_current_lab_predict = torch.cat((I_current_l, I_current_ab_predict), dim=1)
|
||||
I_current_ab_weighted_nonlocal = nonlocal_weighted_layer(
|
||||
I_current_lab_predict,
|
||||
nonlocal_smooth_feature.detach(),
|
||||
patch_size=3,
|
||||
alpha=alpha_nonlocal_smoothness,
|
||||
scale_factor=scale_factor,
|
||||
)
|
||||
nonlocal_smoothness_loss = (
|
||||
mse_loss(
|
||||
nn.functional.interpolate(I_current_ab_predict, scale_factor=scale_factor),
|
||||
I_current_ab_weighted_nonlocal,
|
||||
)
|
||||
* opt.weight_nonlocal_smoothness
|
||||
)
|
||||
|
||||
if opt.weight_consistent:
|
||||
I_current_lab_predict_warp = warping_layer(I_current_lab_predict, flow_forward)
|
||||
I_current_ab_predict_warp = I_current_lab_predict_warp[:, 1:3, :, :]
|
||||
consistent_loss = (
|
||||
weighted_mse_loss(I_current_ab_predict_warp, I_last_ab_predict, mask) * opt.weight_consistent
|
||||
)
|
||||
|
||||
if opt.weight_nonlocal_consistent:
|
||||
I_current_nonlocal_lab_predict_warp = warping_layer(I_current_nonlocal_lab_predict, flow_forward)
|
||||
nonlocal_consistent_loss = (
|
||||
weighted_mse_loss(
|
||||
I_current_nonlocal_lab_predict_warp[:, 1:3, :, :], I_last_nonlocal_lab_predict[:, 1:3, :, :], mask
|
||||
)
|
||||
* opt.weight_nonlocal_consistent
|
||||
)
|
||||
|
||||
# total loss
|
||||
total_loss = (
|
||||
l1_loss
|
||||
+ feat_loss
|
||||
+ contextual_loss_total
|
||||
+ smoothness_loss
|
||||
+ nonlocal_smoothness_loss
|
||||
+ consistent_loss
|
||||
+ nonlocal_consistent_loss
|
||||
+ generator_loss
|
||||
)
|
||||
total_loss.backward()
|
||||
optimizer_g.step()
|
||||
end_time = time.time()
|
||||
elapsed = end_time - start_time
|
||||
training_logger()
|
||||
|
||||
step_optim_scheduler_g.step()
|
||||
step_optim_scheduler_d.step()
|
||||
|
||||
data_queue.put((None, None))
|
||||
Reference in New Issue
Block a user