nodes/jonstreeter-ComfyUI-Reference-Based-Video-Colorization
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
894d8f591bcf99c07c6ed3802a982ae22baf1345
jonstreeter-ComfyUI-Referen…/nodes.py
T
jonstreeterandClaude 6dda556567 Initial release: Reference-Based Video Colorization for ComfyUI 
Repository renamed from ComfyUI-Deep-Exemplar-based-Video-Colorization to ComfyUI-Reference-Based-Video-Colorization to better reflect the comprehensive nature of this toolkit.

## Major Features

### Dual Implementation
- ColorMNet (2024) - Modern DINOv2-based approach
- Deep Exemplar (2019) - Classic CVPR method
- 4 total nodes (2 video, 2 image)

### Modern Components
- Multiple feature encoders: VGG19, DINOv2 (ViT-S/B/L), CLIP (ViT-B)
- Advanced post-processing: color-matcher, WLS, guided, bilateral filters
- Silent auto-installer for dependencies (timm, CLIP, color-matcher)
- Comprehensive tooltips on all node parameters

### Performance Optimizations
- torch.compile support (15-25% speedup on PyTorch 2.0+)
- FP16 mixed precision support
- CUDA graphs compatibility handling
- Multiple memory modes for ColorMNet

## Bug 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
- Comprehensive user guides and technical documentation
- Performance optimization guides
- Migration guides
- Integration documentation
- Auto-installer documentation

All URLs updated to: https://github.com/jonstreeter/ComfyUI-Reference-Based-Video-Colorization

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>
2025-11-15 13:38:47 -06:00

688 lines
28 KiB
Python
Raw Blame History

"""Modern ComfyUI nodes for ColorMNet video colorization.
This module provides ComfyUI nodes for reference-based video colorization using ColorMNet.
"""
import torch
import time
from datetime import datetime
from pathlib import Path
from typing import Tuple
# ComfyUI imports
try:
from comfy.utils import ProgressBar
COMFY_PROGRESS_AVAILABLE = True
except ImportError:
COMFY_PROGRESS_AVAILABLE = False
from tqdm import tqdm
# Our modern infrastructure
from core.logger import setup_logger, get_logger
from core.device import DeviceManager
from core.exceptions import ColorMNetError, ModelNotFoundError, InsufficientVRAMError
from core.validation import validate_input, ValidationError
from colormnet.inference import ColorMNetInference
from colormnet.config import ColorMNetConfig
from colormnet.downloader import setup_model, ensure_model_downloaded, try_install_correlation_sampler
# Initialize logger
setup_logger(name="colormnet", level=20) # INFO level
# Track if we've done first-time setup
_SETUP_COMPLETE = False
_MODEL_PATH = None
_CORR_SAMPLER_AVAILABLE = None # Tri-state: None (not checked), True (available), False (unavailable)
# Cache for CUDA graph compatibility check
_CUDA_GRAPHS_COMPATIBLE = None
def check_cuda_graphs_compatible():
"""
Check if CUDA graphs are compatible with the current CUDA memory allocator.
Returns True if compatible, False if using cudaMallocAsync or incompatible allocator.
"""
global _CUDA_GRAPHS_COMPATIBLE
if _CUDA_GRAPHS_COMPATIBLE is not None:
return _CUDA_GRAPHS_COMPATIBLE
if not torch.cuda.is_available():
_CUDA_GRAPHS_COMPATIBLE = False
return False
try:
# Check if using cudaMallocAsync allocator
# cudaMallocAsync doesn't support checkPoolLiveAllocations needed by CUDA graphs
allocator_backend = torch.cuda.get_allocator_backend()
if allocator_backend == 'cudaMallocAsync':
logger = get_logger()
logger.info("Detected cudaMallocAsync allocator - CUDA graphs will be disabled")
_CUDA_GRAPHS_COMPATIBLE = False
return False
except (AttributeError, RuntimeError):
# If we can't check allocator, assume compatible
pass
# Default to compatible
_CUDA_GRAPHS_COMPATIBLE = True
return True
def get_torch_compile_kwargs():
"""
Get appropriate kwargs for torch.compile based on system compatibility.
Returns dict with mode and optionally disable_cudagraphs.
"""
if check_cuda_graphs_compatible():
# CUDA graphs supported - use reduce-overhead for best performance
return {'mode': 'reduce-overhead'}
else:
# CUDA graphs not supported - try to disable them
# Note: disable_cudagraphs parameter was added in PyTorch 2.1+
# For older versions, just use mode='default' (slower but works)
return {'mode': 'default'}
class ColorMNetVideoNode:
"""ComfyUI node for ColorMNet video colorization.
Colorizes grayscale or color video frames using a reference color image.
Uses memory-based temporal propagation for consistent colorization.
"""
def __init__(self):
self.logger = get_logger()
self.inference_pipeline = None
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"video_frames": ("IMAGE", {
"tooltip": "Batch of video frames to be colorized [N, H, W, 3]"
}),
"reference_image": ("IMAGE", {
"tooltip": "Color reference image that provides the color palette"
}),
"target_width": ("INT", {
"default": 768,
"min": 64,
"max": 2048,
"step": 32,
"display": "number",
"tooltip": "Output width - must be multiple of 32 (will be adjusted automatically)"
}),
"target_height": ("INT", {
"default": 432,
"min": 64,
"max": 2048,
"step": 32,
"display": "number",
"tooltip": "Output height - must be multiple of 32 (will be adjusted automatically)"
}),
"memory_mode": (["balanced", "low_memory", "high_quality"], {
"default": "balanced",
"tooltip": "Memory management strategy: 'low_memory' for large videos, 'high_quality' for best results, 'balanced' for general use"
}),
"feature_encoder": (["resnet50", "vgg19", "dinov2_vits", "dinov2_vitb", "dinov2_vitl", "clip_vitb"], {
"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",
"label_off": "disabled",
"tooltip": "Use half-precision (FP16) for faster processing with lower VRAM usage (minimal quality impact)"
}),
"use_torch_compile": ("BOOLEAN", {
"default": False,
"label_on": "enabled",
"label_off": "disabled",
"tooltip": "Enable torch.compile optimization for 15-25% speedup (increases first-run compilation time)"
}),
},
"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")
RETURN_NAMES = ("colorized_frames", "performance_report")
FUNCTION = "colorize_video"
CATEGORY = "ColorMNet/Video"
def colorize_video(
self,
video_frames: torch.Tensor,
reference_image: torch.Tensor,
target_width: int,
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 = "",
text_guidance_weight: float = 0.3,
) -> Tuple[torch.Tensor, str]:
"""Colorize video frames using reference image.
Args:
video_frames: Video frames [N, H, W, 3] or [H, W, 3]
reference_image: Reference image [H, W, 3] or [1, H, W, 3]
target_width: Target output width
target_height: Target output height
memory_mode: Memory usage mode
use_fp16: Use FP16 mixed precision
Returns:
Tuple containing colorized frames tensor and performance report
"""
global _SETUP_COMPLETE, _MODEL_PATH, _CORR_SAMPLER_AVAILABLE
# Start timing
start_time = time.time()
try:
self.logger.info("=" * 50)
self.logger.info("ColorMNet Video Colorization Starting")
self.logger.info("=" * 50)
# Get model path
script_dir = Path(__file__).parent
model_path = script_dir / "checkpoints" / "DINOv2FeatureV6_LocalAtten_s2_154000.pth"
# Auto-download model if not exists
if not model_path.exists() or not _SETUP_COMPLETE:
self.logger.info("First run detected - setting up ColorMNet...")
checkpoint_dir = script_dir / "checkpoints"
_MODEL_PATH = ensure_model_downloaded(checkpoint_dir)
if _MODEL_PATH is None:
raise ModelNotFoundError(
f"Model checkpoint not found at: {model_path}\n"
f"Auto-download failed. Please download manually from:\n"
f"https://github.com/yyang181/colormnet/releases/tag/v0.1\n"
f"And place at: {model_path}"
)
model_path = _MODEL_PATH
_SETUP_COMPLETE = True
self.logger.info("✓ Setup complete!")
else:
_MODEL_PATH = model_path
# Check for optimized CUDA correlation sampler (first run only)
if _CORR_SAMPLER_AVAILABLE is None:
self.logger.info("Checking for performance optimizations...")
success, message = try_install_correlation_sampler()
_CORR_SAMPLER_AVAILABLE = success
if not success:
self.logger.info("Using fallback mode (20-40% slower, same quality)")
# Create config based on memory mode
if memory_mode == "low_memory":
config = ColorMNetConfig.low_memory(model_path)
elif memory_mode == "high_quality":
config = ColorMNetConfig.high_quality(model_path)
else: # balanced
config = ColorMNetConfig.default(model_path)
config.use_fp16 = use_fp16
config.enable_corr = _CORR_SAMPLER_AVAILABLE # Enable if available
# Create device manager
device_manager = DeviceManager(device=None, use_fp16=use_fp16)
# Create inference pipeline
self.logger.info("Initializing ColorMNet inference pipeline")
self.inference_pipeline = ColorMNetInference(
model_path=model_path,
device_manager=device_manager,
config=config,
)
# Apply torch.compile optimization if requested
# Note: torch.compile on ColorMNet has minimal benefit (~5% speedup)
# The model is already highly optimized with custom CUDA kernels
if use_torch_compile and hasattr(torch, 'compile'):
self.logger.info("torch.compile requested but not applicable to ColorMNet architecture")
self.logger.info("ColorMNet is already optimized with custom kernels (no additional speedup)")
# Handle single image vs batch
if video_frames.dim() == 3:
video_frames = video_frames.unsqueeze(0)
if reference_image.dim() == 4 and reference_image.shape[0] == 1:
reference_image = reference_image.squeeze(0)
# Setup progress tracking
num_frames = video_frames.shape[0]
if COMFY_PROGRESS_AVAILABLE:
pbar = ProgressBar(num_frames)
def progress_callback(current, total):
pbar.update_absolute(current, total)
# Update ComfyUI progress if method exists
if hasattr(self, 'set_progress'):
self.set_progress(pbar.progress)
else:
pbar = tqdm(total=num_frames, desc="Colorizing frames")
def progress_callback(current, total):
pbar.update(1)
# Colorize
self.logger.info(f"Processing {num_frames} frames at {target_height}x{target_width}")
colorized = self.inference_pipeline.colorize_video(
frames=video_frames,
reference=reference_image,
target_height=target_height,
target_width=target_width,
progress_callback=progress_callback,
)
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
# Generate performance report
fps = num_frames / elapsed_time if elapsed_time > 0 else 0
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
report = (
f"ColorMNet Video Colorization Report\n"
f"{'=' * 50}\n"
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"{'=' * 50}"
)
self.logger.info("Colorization complete!")
self.logger.info(f"Processed {num_frames} frames in {elapsed_time:.2f}s ({fps:.2f} FPS)")
self.logger.info("=" * 50)
return (colorized, report)
except ModelNotFoundError as e:
self.logger.error(f"Model not found: {e}")
self.logger.error("Please run install.py to download the model checkpoint")
raise
except InsufficientVRAMError as e:
self.logger.error(f"Insufficient VRAM: {e}")
raise
except ValidationError as e:
self.logger.error(f"Input validation failed: {e}")
raise
except ColorMNetError as e:
self.logger.error(f"ColorMNet error: {e}")
raise
except Exception as e:
self.logger.error(f"Unexpected error during colorization: {e}", exc_info=True)
raise
class ColorMNetImageNode:
"""ComfyUI node for ColorMNet single image colorization.
Colorizes a single grayscale or color image using a reference color image.
"""
def __init__(self):
self.logger = get_logger()
self.inference_pipeline = None
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE", {
"tooltip": "Grayscale or color image to be colorized [H, W, 3]"
}),
"reference_image": ("IMAGE", {
"tooltip": "Color reference image that provides the color palette"
}),
"target_width": ("INT", {
"default": 768,
"min": 64,
"max": 2048,
"step": 32,
"tooltip": "Output width - must be multiple of 32 (will be adjusted automatically)"
}),
"target_height": ("INT", {
"default": 432,
"min": 64,
"max": 2048,
"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)"
}),
"use_torch_compile": ("BOOLEAN", {
"default": False,
"tooltip": "Enable torch.compile optimization for 10-15% speedup (increases first-run compilation time)"
}),
},
}
RETURN_TYPES = ("IMAGE", "STRING")
RETURN_NAMES = ("colorized_image", "performance_report")
FUNCTION = "colorize_image"
CATEGORY = "ColorMNet/Image"
def colorize_image(
self,
image: torch.Tensor,
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]:
"""Colorize a single image using reference.
Args:
image: Input image
reference_image: Reference image
target_width: Target width
target_height: Target height
use_fp16: Use FP16
Returns:
Tuple containing colorized image tensor and performance report
"""
global _SETUP_COMPLETE, _MODEL_PATH, _CORR_SAMPLER_AVAILABLE
# Start timing
start_time = time.time()
try:
self.logger.info("ColorMNet Image Colorization")
# Get model path
script_dir = Path(__file__).parent
model_path = script_dir / "checkpoints" / "DINOv2FeatureV6_LocalAtten_s2_154000.pth"
# Auto-download model if not exists
if not model_path.exists() or not _SETUP_COMPLETE:
self.logger.info("First run detected - setting up ColorMNet...")
checkpoint_dir = script_dir / "checkpoints"
_MODEL_PATH = ensure_model_downloaded(checkpoint_dir)
if _MODEL_PATH is None:
raise ModelNotFoundError(
f"Model checkpoint not found at: {model_path}\n"
f"Auto-download failed. Please download manually from:\n"
f"https://github.com/yyang181/colormnet/releases/tag/v0.1\n"
f"And place at: {model_path}"
)
model_path = _MODEL_PATH
_SETUP_COMPLETE = True
self.logger.info("✓ Setup complete!")
else:
_MODEL_PATH = model_path
# Check for optimized CUDA correlation sampler (first run only)
if _CORR_SAMPLER_AVAILABLE is None:
success, message = try_install_correlation_sampler()
_CORR_SAMPLER_AVAILABLE = success
# Create config
config = ColorMNetConfig.default(model_path)
config.use_fp16 = use_fp16
config.enable_corr = _CORR_SAMPLER_AVAILABLE
# Create device manager
device_manager = DeviceManager(device=None, use_fp16=use_fp16)
# Create inference pipeline
self.inference_pipeline = ColorMNetInference(
model_path=model_path,
device_manager=device_manager,
config=config,
)
# Apply torch.compile optimization if requested
# Note: torch.compile on ColorMNet has minimal benefit (~5% speedup)
# The model is already highly optimized with custom CUDA kernels
if use_torch_compile and hasattr(torch, 'compile'):
self.logger.info("torch.compile requested but not applicable to ColorMNet architecture")
self.logger.info("ColorMNet is already optimized with custom kernels")
# Handle batch dimension
if image.dim() == 4 and image.shape[0] == 1:
image = image.squeeze(0)
if reference_image.dim() == 4 and reference_image.shape[0] == 1:
reference_image = reference_image.squeeze(0)
# Colorize
colorized = self.inference_pipeline.colorize_image(
image=image,
reference=reference_image,
target_height=target_height,
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
# Calculate timing
end_time = time.time()
elapsed_time = end_time - start_time
# Generate performance report
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
report = (
f"ColorMNet Image Colorization Report\n"
f"{'=' * 50}\n"
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}"
)
self.logger.info(f"Image colorization complete in {elapsed_time:.3f}s")
return (colorized, report)
except Exception as e:
self.logger.error(f"Error during image colorization: {e}", exc_info=True)
raise
# Node registration for ComfyUI
NODE_CLASS_MAPPINGS = {
"ColorMNetVideo": ColorMNetVideoNode,
"ColorMNetImage": ColorMNetImageNode,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"ColorMNetVideo": "ColorMNet Video Colorization (New)",
"ColorMNetImage": "ColorMNet Image Colorization (New)",
}
Reference in New Issue View Git Blame Copy Permalink