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>
208 lines
6.0 KiB
Python
208 lines
6.0 KiB
Python
from functools import reduce
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
class LambdaBase(nn.Sequential):
|
|
def __init__(self, fn, *args):
|
|
super(LambdaBase, self).__init__(*args)
|
|
self.lambda_func = fn
|
|
|
|
def forward_prepare(self, input):
|
|
output = []
|
|
for module in self._modules.values():
|
|
output.append(module(input))
|
|
|
|
return output if output else input
|
|
|
|
|
|
class Lambda(LambdaBase):
|
|
def forward(self, input):
|
|
return self.lambda_func(self.forward_prepare(input))
|
|
|
|
|
|
class LambdaMap(LambdaBase):
|
|
def forward(self, input):
|
|
return list(map(self.lambda_func, self.forward_prepare(input)))
|
|
|
|
|
|
class LambdaReduce(LambdaBase):
|
|
def forward(self, input):
|
|
return reduce(self.lambda_func, self.forward_prepare(input))
|
|
|
|
|
|
layer_names = [
|
|
"conv1_1",
|
|
"relu1_1",
|
|
"conv1_2",
|
|
"relu1_2",
|
|
"pool1",
|
|
"conv2_1",
|
|
"relu2_1",
|
|
"conv2_2",
|
|
"relu2_2",
|
|
"pool2",
|
|
"conv3_1",
|
|
"relu3_1",
|
|
"conv3_2",
|
|
"relu3_2",
|
|
"conv3_3",
|
|
"relu3_3",
|
|
"conv3_4",
|
|
"relu3_4",
|
|
"pool3",
|
|
"conv4_1",
|
|
"relu4_1",
|
|
"conv4_2",
|
|
"relu4_2",
|
|
"conv4_3",
|
|
"relu4_3",
|
|
"conv4_4",
|
|
"relu4_4",
|
|
"pool4",
|
|
"conv5_1",
|
|
"relu5_1",
|
|
"conv5_2",
|
|
"relu5_2",
|
|
"conv5_3",
|
|
"relu5_3",
|
|
"conv5_4",
|
|
"relu5_4",
|
|
"pool5",
|
|
"view1",
|
|
"fc6",
|
|
"fc6_relu",
|
|
"fc7",
|
|
"fc7_relu",
|
|
"fc8",
|
|
]
|
|
|
|
model = nn.Sequential( # Sequential,
|
|
nn.Conv2d(3, 64, (3, 3), (1, 1), (1, 1)),
|
|
nn.ReLU(),
|
|
nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1)),
|
|
nn.ReLU(),
|
|
nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True),
|
|
nn.Conv2d(64, 128, (3, 3), (1, 1), (1, 1)),
|
|
nn.ReLU(),
|
|
nn.Conv2d(128, 128, (3, 3), (1, 1), (1, 1)),
|
|
nn.ReLU(),
|
|
nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True),
|
|
nn.Conv2d(128, 256, (3, 3), (1, 1), (1, 1)),
|
|
nn.ReLU(),
|
|
nn.Conv2d(256, 256, (3, 3), (1, 1), (1, 1)),
|
|
nn.ReLU(),
|
|
nn.Conv2d(256, 256, (3, 3), (1, 1), (1, 1)),
|
|
nn.ReLU(),
|
|
nn.Conv2d(256, 256, (3, 3), (1, 1), (1, 1)),
|
|
nn.ReLU(),
|
|
nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True),
|
|
nn.Conv2d(256, 512, (3, 3), (1, 1), (1, 1)),
|
|
nn.ReLU(),
|
|
nn.Conv2d(512, 512, (3, 3), (1, 1), (1, 1)),
|
|
nn.ReLU(),
|
|
nn.Conv2d(512, 512, (3, 3), (1, 1), (1, 1)),
|
|
nn.ReLU(),
|
|
nn.Conv2d(512, 512, (3, 3), (1, 1), (1, 1)),
|
|
nn.ReLU(),
|
|
nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True),
|
|
nn.Conv2d(512, 512, (3, 3), (1, 1), (1, 1)),
|
|
nn.ReLU(),
|
|
nn.Conv2d(512, 512, (3, 3), (1, 1), (1, 1)),
|
|
nn.ReLU(),
|
|
nn.Conv2d(512, 512, (3, 3), (1, 1), (1, 1)),
|
|
nn.ReLU(),
|
|
nn.Conv2d(512, 512, (3, 3), (1, 1), (1, 1)),
|
|
nn.ReLU(),
|
|
nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True),
|
|
Lambda(lambda x: x.view(x.size(0), -1)), # View,
|
|
nn.Sequential(Lambda(lambda x: x.view(1, -1) if 1 == len(x.size()) else x), nn.Linear(25088, 4096)), # Linear,
|
|
nn.ReLU(),
|
|
nn.Sequential(Lambda(lambda x: x.view(1, -1) if 1 == len(x.size()) else x), nn.Linear(4096, 4096)), # Linear,
|
|
nn.ReLU(),
|
|
nn.Sequential(Lambda(lambda x: x.view(1, -1) if 1 == len(x.size()) else x), nn.Linear(4096, 1000)), # Linear,
|
|
)
|
|
|
|
|
|
# Lazy loading: Create VGG network only when needed
|
|
_vgg19_gray_net = None
|
|
|
|
def _get_vgg19_gray_net():
|
|
global _vgg19_gray_net
|
|
if _vgg19_gray_net is None:
|
|
import os
|
|
# Get the path to the model file relative to this script
|
|
script_dir = os.path.dirname(os.path.abspath(__file__))
|
|
model_path = os.path.join(script_dir, "..", "data", "vgg19_conv.pth")
|
|
|
|
model.load_state_dict(torch.load(model_path, map_location='cpu'))
|
|
_vgg19_gray_net = torch.nn.Sequential()
|
|
for (name, layer) in model._modules.items():
|
|
_vgg19_gray_net.add_module(layer_names[int(name)], model[int(name)])
|
|
|
|
for param in _vgg19_gray_net.parameters():
|
|
param.requires_grad = False
|
|
_vgg19_gray_net.eval()
|
|
|
|
return _vgg19_gray_net
|
|
|
|
|
|
class vgg19_gray(torch.nn.Module):
|
|
def __init__(self, requires_grad=False):
|
|
super(vgg19_gray, self).__init__()
|
|
vgg_pretrained_features = _get_vgg19_gray_net()
|
|
self.slice1 = torch.nn.Sequential()
|
|
self.slice2 = torch.nn.Sequential()
|
|
self.slice3 = torch.nn.Sequential()
|
|
for x in range(12):
|
|
self.slice1.add_module(layer_names[x], vgg_pretrained_features[x])
|
|
for x in range(12, 21):
|
|
self.slice2.add_module(layer_names[x], vgg_pretrained_features[x])
|
|
for x in range(21, 30):
|
|
self.slice3.add_module(layer_names[x], vgg_pretrained_features[x])
|
|
if not requires_grad:
|
|
for param in self.parameters():
|
|
param.requires_grad = False
|
|
|
|
def forward(self, X):
|
|
h = self.slice1(X)
|
|
h_relu3_1 = h
|
|
h = self.slice2(h)
|
|
h_relu4_1 = h
|
|
h = self.slice3(h)
|
|
h_relu5_1 = h
|
|
return h_relu3_1, h_relu4_1, h_relu5_1
|
|
|
|
|
|
class vgg19_gray_new(torch.nn.Module):
|
|
def __init__(self, requires_grad=False):
|
|
super(vgg19_gray_new, self).__init__()
|
|
vgg_pretrained_features = _get_vgg19_gray_net()
|
|
self.slice0 = torch.nn.Sequential()
|
|
self.slice1 = torch.nn.Sequential()
|
|
self.slice2 = torch.nn.Sequential()
|
|
self.slice3 = torch.nn.Sequential()
|
|
for x in range(7):
|
|
self.slice0.add_module(layer_names[x], vgg_pretrained_features[x])
|
|
for x in range(7, 12):
|
|
self.slice1.add_module(layer_names[x], vgg_pretrained_features[x])
|
|
for x in range(12, 21):
|
|
self.slice2.add_module(layer_names[x], vgg_pretrained_features[x])
|
|
for x in range(21, 30):
|
|
self.slice3.add_module(layer_names[x], vgg_pretrained_features[x])
|
|
if not requires_grad:
|
|
for param in self.parameters():
|
|
param.requires_grad = False
|
|
|
|
def forward(self, X):
|
|
h = self.slice0(X)
|
|
h_relu2_1 = h
|
|
h = self.slice1(h)
|
|
h_relu3_1 = h
|
|
h = self.slice2(h)
|
|
h_relu4_1 = h
|
|
h = self.slice3(h)
|
|
h_relu5_1 = h
|
|
return h_relu2_1, h_relu3_1, h_relu4_1, h_relu5_1
|