jonstreeterandClaude
6dda556567
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>
68 lines
2.1 KiB
Python
68 lines
2.1 KiB
Python
import torch
|
|
from ..utils.util import *
|
|
|
|
|
|
def warp_color(IA_l, IB_lab, features_B, vggnet, nonlocal_net, colornet, feature_noise=0, temperature=0.01):
|
|
IA_rgb_from_gray = gray2rgb_batch(IA_l)
|
|
with torch.no_grad():
|
|
A_relu1_1, A_relu2_1, A_relu3_1, A_relu4_1, A_relu5_1 = vggnet(
|
|
IA_rgb_from_gray, ["r12", "r22", "r32", "r42", "r52"], preprocess=True
|
|
)
|
|
B_relu1_1, B_relu2_1, B_relu3_1, B_relu4_1, B_relu5_1 = features_B
|
|
|
|
# NOTE: output the feature before normalization
|
|
features_A = [A_relu1_1, A_relu2_1, A_relu3_1, A_relu4_1, A_relu5_1]
|
|
|
|
A_relu2_1 = feature_normalize(A_relu2_1)
|
|
A_relu3_1 = feature_normalize(A_relu3_1)
|
|
A_relu4_1 = feature_normalize(A_relu4_1)
|
|
A_relu5_1 = feature_normalize(A_relu5_1)
|
|
B_relu2_1 = feature_normalize(B_relu2_1)
|
|
B_relu3_1 = feature_normalize(B_relu3_1)
|
|
B_relu4_1 = feature_normalize(B_relu4_1)
|
|
B_relu5_1 = feature_normalize(B_relu5_1)
|
|
|
|
nonlocal_BA_lab, similarity_map = nonlocal_net(
|
|
IB_lab,
|
|
A_relu2_1,
|
|
A_relu3_1,
|
|
A_relu4_1,
|
|
A_relu5_1,
|
|
B_relu2_1,
|
|
B_relu3_1,
|
|
B_relu4_1,
|
|
B_relu5_1,
|
|
temperature=temperature,
|
|
)
|
|
|
|
return nonlocal_BA_lab, similarity_map, features_A
|
|
|
|
|
|
def frame_colorization(
|
|
IA_lab,
|
|
IB_lab,
|
|
IA_last_lab,
|
|
features_B,
|
|
vggnet,
|
|
nonlocal_net,
|
|
colornet,
|
|
joint_training=True,
|
|
feature_noise=0,
|
|
luminance_noise=0,
|
|
temperature=0.01,
|
|
):
|
|
|
|
IA_l = IA_lab[:, 0:1, :, :]
|
|
if luminance_noise:
|
|
IA_l = IA_l + torch.randn_like(IA_l, requires_grad=False) * luminance_noise
|
|
|
|
with torch.autograd.set_grad_enabled(joint_training):
|
|
nonlocal_BA_lab, similarity_map, features_A_gray = warp_color(
|
|
IA_l, IB_lab, features_B, vggnet, nonlocal_net, colornet, feature_noise, temperature=temperature
|
|
)
|
|
nonlocal_BA_ab = nonlocal_BA_lab[:, 1:3, :, :]
|
|
color_input = torch.cat((IA_l, nonlocal_BA_ab, similarity_map, IA_last_lab), dim=1)
|
|
IA_ab_predict = colornet(color_input)
|
|
|
|
return IA_ab_predict, nonlocal_BA_lab, features_A_gray
|