nodes/jonstreeter-ComfyUI-Reference-Based-Video-Colorization
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
main
jonstreeter-ComfyUI-Referen…/nodes.py
T
jonstreeterandClaude 0d831af1b9 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>
2025-11-18 14:52:13 -06:00

542 lines
21 KiB
Python
Raw Permalink 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)"
}),
"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",
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}")
colorization_start_time = time.time()
colorized = self.inference_pipeline.colorize_video(
frames=video_frames,
reference=reference_image,
target_height=target_height,
target_width=target_width,
progress_callback=progress_callback,
)
colorization_time = time.time() - colorization_start_time
if not COMFY_PROGRESS_AVAILABLE:
pbar.close()
# 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"\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}"
)
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)"
}),
"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,
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,
)
# 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"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