nodes/jonstreeter-ComfyUI-Reference-Based-Video-Colorization
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Simplify DeepExemplar to VGG19-only, fix ColorVidNet skip connections

Browse Source 
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
2025-11-18 14:52:13 -06:00
co-authored by Claude
parent 9727e8479e
commit 0d831af1b9
32 changed files with 189 additions and 5669 deletions
Download Patch File Download Diff File Expand all files Collapse all files
models/ColorVidNet.py
+10
View File
@@ -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)
models/NonlocalNet.py
+30 -10
View File
@@ -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
models/feature_extractors/__init__.py
+22 -4
View File
@@ -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',
models/feature_extractors/clip_encoder.py
+33 -9
View File
@@ -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."""
models/feature_extractors/dinov2_encoder.py
+33 -10
View File
@@ -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