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>
304 lines
12 KiB
Python
304 lines
12 KiB
Python
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
import torchvision
|
|
from ..utils.util import feature_normalize
|
|
|
|
postpa = torchvision.transforms.Compose(
|
|
[
|
|
torchvision.transforms.Lambda(lambda x: x.mul_(1.0 / 255)),
|
|
torchvision.transforms.Normalize(
|
|
mean=[-0.40760392, -0.45795686, -0.48501961], std=[1, 1, 1] # add imagenet mean
|
|
),
|
|
torchvision.transforms.Lambda(lambda x: x[torch.LongTensor([2, 1, 0])]), # turn to RGB
|
|
]
|
|
)
|
|
postpb = torchvision.transforms.Compose([torchvision.transforms.ToPILImage()])
|
|
|
|
|
|
def post_processing(tensor):
|
|
t = postpa(tensor) # denormalize the image since the optimized tensor is the normalized one
|
|
t[t > 1] = 1
|
|
t[t < 0] = 0
|
|
img = postpb(t)
|
|
img = np.array(img)
|
|
return img
|
|
|
|
|
|
class ContextualLoss(nn.Module):
|
|
"""
|
|
input is Al, Bl, channel = 1, range ~ [0, 255]
|
|
"""
|
|
|
|
def __init__(self):
|
|
super(ContextualLoss, self).__init__()
|
|
return None
|
|
|
|
def forward(self, X_features, Y_features, h=0.1, feature_centering=True):
|
|
"""
|
|
X_features&Y_features are are feature vectors or feature 2d array
|
|
h: bandwidth
|
|
return the per-sample loss
|
|
"""
|
|
batch_size = X_features.shape[0]
|
|
feature_depth = X_features.shape[1]
|
|
feature_size = X_features.shape[2]
|
|
|
|
# to normalized feature vectors
|
|
if feature_centering:
|
|
X_features = X_features - Y_features.view(batch_size, feature_depth, -1).mean(dim=-1).unsqueeze(
|
|
dim=-1
|
|
).unsqueeze(dim=-1)
|
|
Y_features = Y_features - Y_features.view(batch_size, feature_depth, -1).mean(dim=-1).unsqueeze(
|
|
dim=-1
|
|
).unsqueeze(dim=-1)
|
|
X_features = feature_normalize(X_features).view(
|
|
batch_size, feature_depth, -1
|
|
) # batch_size * feature_depth * feature_size^2
|
|
Y_features = feature_normalize(Y_features).view(
|
|
batch_size, feature_depth, -1
|
|
) # batch_size * feature_depth * feature_size^2
|
|
|
|
# conine distance = 1 - similarity
|
|
X_features_permute = X_features.permute(0, 2, 1) # batch_size * feature_size^2 * feature_depth
|
|
d = 1 - torch.matmul(X_features_permute, Y_features) # batch_size * feature_size^2 * feature_size^2
|
|
|
|
# normalized distance: dij_bar
|
|
d_norm = d / (torch.min(d, dim=-1, keepdim=True)[0] + 1e-5) # batch_size * feature_size^2 * feature_size^2
|
|
|
|
# pairwise affinity
|
|
w = torch.exp((1 - d_norm) / h)
|
|
A_ij = w / torch.sum(w, dim=-1, keepdim=True)
|
|
|
|
# contextual loss per sample
|
|
CX = torch.mean(torch.max(A_ij, dim=1)[0], dim=-1)
|
|
return -torch.log(CX)
|
|
|
|
|
|
class ContextualLoss_forward(nn.Module):
|
|
"""
|
|
input is Al, Bl, channel = 1, range ~ [0, 255]
|
|
"""
|
|
|
|
def __init__(self):
|
|
super(ContextualLoss_forward, self).__init__()
|
|
return None
|
|
|
|
def forward(self, X_features, Y_features, h=0.1, feature_centering=True):
|
|
"""
|
|
X_features&Y_features are are feature vectors or feature 2d array
|
|
h: bandwidth
|
|
return the per-sample loss
|
|
"""
|
|
batch_size = X_features.shape[0]
|
|
feature_depth = X_features.shape[1]
|
|
feature_size = X_features.shape[2]
|
|
|
|
# to normalized feature vectors
|
|
if feature_centering:
|
|
X_features = X_features - Y_features.view(batch_size, feature_depth, -1).mean(dim=-1).unsqueeze(
|
|
dim=-1
|
|
).unsqueeze(dim=-1)
|
|
Y_features = Y_features - Y_features.view(batch_size, feature_depth, -1).mean(dim=-1).unsqueeze(
|
|
dim=-1
|
|
).unsqueeze(dim=-1)
|
|
X_features = feature_normalize(X_features).view(
|
|
batch_size, feature_depth, -1
|
|
) # batch_size * feature_depth * feature_size^2
|
|
Y_features = feature_normalize(Y_features).view(
|
|
batch_size, feature_depth, -1
|
|
) # batch_size * feature_depth * feature_size^2
|
|
|
|
# conine distance = 1 - similarity
|
|
X_features_permute = X_features.permute(0, 2, 1) # batch_size * feature_size^2 * feature_depth
|
|
d = 1 - torch.matmul(X_features_permute, Y_features) # batch_size * feature_size^2 * feature_size^2
|
|
|
|
# normalized distance: dij_bar
|
|
d_norm = d / (torch.min(d, dim=-1, keepdim=True)[0] + 1e-5) # batch_size * feature_size^2 * feature_size^2
|
|
|
|
# pairwise affinity
|
|
w = torch.exp((1 - d_norm) / h)
|
|
A_ij = w / torch.sum(w, dim=-1, keepdim=True)
|
|
|
|
# contextual loss per sample
|
|
CX = torch.mean(torch.max(A_ij, dim=-1)[0], dim=1)
|
|
return -torch.log(CX)
|
|
|
|
|
|
class ContextualLoss_complex(nn.Module):
|
|
"""
|
|
input is Al, Bl, channel = 1, range ~ [0, 255]
|
|
"""
|
|
|
|
def __init__(self):
|
|
super(ContextualLoss_complex, self).__init__()
|
|
return None
|
|
|
|
def forward(self, X_features, Y_features, h=0.1, patch_size=1, direction="forward"):
|
|
"""
|
|
X_features&Y_features are are feature vectors or feature 2d array
|
|
h: bandwidth
|
|
return the per-sample loss
|
|
"""
|
|
batch_size = X_features.shape[0]
|
|
feature_depth = X_features.shape[1]
|
|
feature_size = X_features.shape[2]
|
|
|
|
# to normalized feature vectors
|
|
X_features = X_features - Y_features.view(batch_size, feature_depth, -1).mean(dim=-1).unsqueeze(
|
|
dim=-1
|
|
).unsqueeze(dim=-1)
|
|
Y_features = Y_features - Y_features.view(batch_size, feature_depth, -1).mean(dim=-1).unsqueeze(
|
|
dim=-1
|
|
).unsqueeze(dim=-1)
|
|
X_features = feature_normalize(X_features) # batch_size * feature_depth * feature_size^2
|
|
Y_features = feature_normalize(Y_features) # batch_size * feature_depth * feature_size^2
|
|
|
|
# to normalized feature vectors
|
|
X_features = F.unfold(
|
|
X_features, kernel_size=(patch_size, patch_size), stride=(1, 1), padding=(patch_size // 2, patch_size // 2)
|
|
) # batch_size * feature_depth_new * feature_size^2
|
|
Y_features = F.unfold(
|
|
Y_features, kernel_size=(patch_size, patch_size), stride=(1, 1), padding=(patch_size // 2, patch_size // 2)
|
|
) # batch_size * feature_depth_new * feature_size^2
|
|
|
|
# conine distance = 1 - similarity
|
|
X_features_permute = X_features.permute(0, 2, 1) # batch_size * feature_size^2 * feature_depth
|
|
d = 1 - torch.matmul(X_features_permute, Y_features) # batch_size * feature_size^2 * feature_size^2
|
|
|
|
# normalized distance: dij_bar
|
|
d_norm = d / (torch.min(d, dim=-1, keepdim=True)[0] + 1e-5) # batch_size * feature_size^2 * feature_size^2
|
|
|
|
# pairwise affinity
|
|
w = torch.exp((1 - d_norm) / h)
|
|
A_ij = w / torch.sum(w, dim=-1, keepdim=True)
|
|
|
|
# contextual loss per sample
|
|
if direction == "forward":
|
|
CX = torch.mean(torch.max(A_ij, dim=-1)[0], dim=1)
|
|
else:
|
|
CX = torch.mean(torch.max(A_ij, dim=1)[0], dim=-1)
|
|
|
|
return -torch.log(CX)
|
|
|
|
|
|
class ChamferDistance_patch_loss(nn.Module):
|
|
"""
|
|
input is Al, Bl, channel = 1, range ~ [0, 255]
|
|
"""
|
|
|
|
def __init__(self):
|
|
super(ChamferDistance_patch_loss, self).__init__()
|
|
return None
|
|
|
|
def forward(self, X_features, Y_features, patch_size=3, image_x=None, image_y=None, h=0.1, Y_features_in=None):
|
|
"""
|
|
X_features&Y_features are are feature vectors or feature 2d array
|
|
h: bandwidth
|
|
return the per-sample loss
|
|
"""
|
|
batch_size = X_features.shape[0]
|
|
feature_depth = X_features.shape[1]
|
|
feature_size = X_features.shape[2]
|
|
|
|
# to normalized feature vectors
|
|
X_features = F.unfold(
|
|
X_features, kernel_size=(patch_size, patch_size), stride=(1, 1), padding=(patch_size // 2, patch_size // 2)
|
|
) # batch_size, feature_depth_new * feature_size^2
|
|
Y_features = F.unfold(
|
|
Y_features, kernel_size=(patch_size, patch_size), stride=(1, 1), padding=(patch_size // 2, patch_size // 2)
|
|
) # batch_size, feature_depth_new * feature_size^2
|
|
|
|
if image_x is not None and image_y is not None:
|
|
image_x = torch.nn.functional.interpolate(image_x, size=(feature_size, feature_size), mode="bilinear").view(
|
|
batch_size, 3, -1
|
|
)
|
|
image_y = torch.nn.functional.interpolate(image_y, size=(feature_size, feature_size), mode="bilinear").view(
|
|
batch_size, 3, -1
|
|
)
|
|
|
|
X_features_permute = X_features.permute(0, 2, 1) # batch_size * feature_size^2 * feature_depth
|
|
similarity_matrix = torch.matmul(X_features_permute, Y_features) # batch_size * feature_size^2 * feature_size^2
|
|
NN_index = similarity_matrix.max(dim=-1, keepdim=True)[1].squeeze()
|
|
|
|
if Y_features_in is not None:
|
|
loss = torch.mean((X_features - Y_features_in.detach()) ** 2)
|
|
Y_features_in = Y_features_in.detach()
|
|
else:
|
|
loss = torch.mean((X_features - Y_features[:, :, NN_index].detach()) ** 2)
|
|
Y_features_in = Y_features[:, :, NN_index].detach()
|
|
|
|
# re-arrange image
|
|
if image_x is not None and image_y is not None:
|
|
image_y_rearrange = image_y[:, :, NN_index]
|
|
image_y_rearrange = image_y_rearrange.view(batch_size, 3, feature_size, feature_size)
|
|
image_x = image_x.view(batch_size, 3, feature_size, feature_size)
|
|
image_y = image_y.view(batch_size, 3, feature_size, feature_size)
|
|
|
|
return loss
|
|
|
|
|
|
class ChamferDistance_loss(nn.Module):
|
|
"""
|
|
input is Al, Bl, channel = 1, range ~ [0, 255]
|
|
"""
|
|
|
|
def __init__(self):
|
|
super(ChamferDistance_loss, self).__init__()
|
|
return None
|
|
|
|
def forward(self, X_features, Y_features, image_x, image_y, h=0.1, Y_features_in=None):
|
|
"""
|
|
X_features&Y_features are are feature vectors or feature 2d array
|
|
h: bandwidth
|
|
return the per-sample loss
|
|
"""
|
|
batch_size = X_features.shape[0]
|
|
feature_depth = X_features.shape[1]
|
|
feature_size = X_features.shape[2]
|
|
|
|
# to normalized feature vectors
|
|
X_features = feature_normalize(X_features).view(
|
|
batch_size, feature_depth, -1
|
|
) # batch_size * feature_depth * feature_size^2
|
|
Y_features = feature_normalize(Y_features).view(
|
|
batch_size, feature_depth, -1
|
|
) # batch_size * feature_depth * feature_size^2
|
|
image_x = torch.nn.functional.interpolate(image_x, size=(feature_size, feature_size), mode="bilinear").view(
|
|
batch_size, 3, -1
|
|
)
|
|
image_y = torch.nn.functional.interpolate(image_y, size=(feature_size, feature_size), mode="bilinear").view(
|
|
batch_size, 3, -1
|
|
)
|
|
|
|
X_features_permute = X_features.permute(0, 2, 1) # batch_size * feature_size^2 * feature_depth
|
|
similarity_matrix = torch.matmul(X_features_permute, Y_features) # batch_size * feature_size^2 * feature_size^2
|
|
NN_index = similarity_matrix.max(dim=-1, keepdim=True)[1].squeeze()
|
|
if Y_features_in is not None:
|
|
loss = torch.mean((X_features - Y_features_in.detach()) ** 2)
|
|
Y_features_in = Y_features_in.detach()
|
|
else:
|
|
loss = torch.mean((X_features - Y_features[:, :, NN_index].detach()) ** 2)
|
|
Y_features_in = Y_features[:, :, NN_index].detach()
|
|
|
|
# re-arrange image
|
|
image_y_rearrange = image_y[:, :, NN_index]
|
|
image_y_rearrange = image_y_rearrange.view(batch_size, 3, feature_size, feature_size)
|
|
image_x = image_x.view(batch_size, 3, feature_size, feature_size)
|
|
image_y = image_y.view(batch_size, 3, feature_size, feature_size)
|
|
|
|
return loss, Y_features_in, X_features
|
|
|
|
|
|
if __name__ == "__main__":
|
|
contextual_loss = ContextualLoss()
|
|
batch_size = 32
|
|
feature_depth = 8
|
|
feature_size = 16
|
|
X_features = torch.zeros(batch_size, feature_depth, feature_size, feature_size)
|
|
Y_features = torch.zeros(batch_size, feature_depth, feature_size, feature_size)
|
|
|
|
cx_loss = contextual_loss(X_features, Y_features, 1)
|
|
print(cx_loss)
|