299 lines
14 KiB
Python
299 lines
14 KiB
Python
"""
|
|
DeepExemplarColorizationNodes.py
|
|
|
|
This module provides two ComfyUI nodes for Deep Exemplar-based Colorization:
|
|
- DeepExColorImageNode: Colorizes a grayscale image using a reference color image.
|
|
- DeepExColorVideoNode: Colorizes a sequence of grayscale video frames using a reference color image.
|
|
|
|
Both nodes follow the test.py logic from the base Deep Exemplar project.
|
|
"""
|
|
|
|
import sys
|
|
import subprocess
|
|
import importlib
|
|
import cv2
|
|
|
|
# Check for cv2.ximgproc.createFastGlobalSmootherFilter and attempt installation if needed.
|
|
if not hasattr(cv2.ximgproc, 'createFastGlobalSmootherFilter'):
|
|
print("cv2.ximgproc.createFastGlobalSmootherFilter not found. Attempting to install opencv-contrib-python...")
|
|
try:
|
|
subprocess.check_call([sys.executable, "-m", "pip", "install", "--upgrade", "opencv-contrib-python"])
|
|
importlib.reload(cv2)
|
|
if hasattr(cv2.ximgproc, 'createFastGlobalSmootherFilter'):
|
|
print("Installation successful and function is now available.")
|
|
else:
|
|
print("Installation attempted, but function is still unavailable. A restart might be required.")
|
|
except Exception as e:
|
|
print(f"Failed to install opencv-contrib-python: {e}")
|
|
|
|
WLS_FILTER_AVAILABLE = hasattr(cv2.ximgproc, 'createFastGlobalSmootherFilter')
|
|
|
|
import os
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import numpy as np
|
|
from PIL import Image
|
|
import torchvision.transforms as Tlib
|
|
from tqdm import tqdm as console_tqdm
|
|
|
|
# Import ComfyUI's internal ProgressBar for GUI updates.
|
|
from comfy.utils import ProgressBar
|
|
|
|
from utils.util_distortion import CenterPad, Normalize, RGB2Lab, ToTensor
|
|
from utils.util import batch_lab2rgb_transpose_mc, uncenter_l, tensor_lab2rgb
|
|
from models.FrameColor import frame_colorization
|
|
|
|
MODELS_LOADED = False
|
|
NONLOCAL_NET = None
|
|
COLOR_NET = None
|
|
VGG_NET = None
|
|
|
|
def load_models_if_needed():
|
|
global MODELS_LOADED, NONLOCAL_NET, COLOR_NET, VGG_NET
|
|
if MODELS_LOADED:
|
|
return
|
|
|
|
print("[DeepExemplar] Loading model checkpoints...")
|
|
script_dir = os.path.dirname(os.path.abspath(__file__))
|
|
nonlocal_ckpt = os.path.join(script_dir, "checkpoints", "video_moredata_l1", "nonlocal_net_iter_76000.pth")
|
|
color_ckpt = os.path.join(script_dir, "checkpoints", "video_moredata_l1", "colornet_iter_76000.pth")
|
|
vgg_ckpt = os.path.join(script_dir, "data", "vgg19_conv.pth")
|
|
from models.NonlocalNet import WarpNet, VGG19_pytorch
|
|
from models.ColorVidNet import ColorVidNet
|
|
NONLOCAL_NET = WarpNet(1).cuda().eval()
|
|
COLOR_NET = ColorVidNet(7).cuda().eval()
|
|
VGG_NET = VGG19_pytorch().cuda().eval()
|
|
NONLOCAL_NET.load_state_dict(torch.load(nonlocal_ckpt, map_location="cuda"))
|
|
COLOR_NET.load_state_dict(torch.load(color_ckpt, map_location="cuda"))
|
|
VGG_NET.load_state_dict(torch.load(vgg_ckpt, map_location="cuda"))
|
|
MODELS_LOADED = True
|
|
print("[DeepExemplar] Models loaded successfully.")
|
|
|
|
def adjust_target_size(input_h: int, input_w: int, min_h: int = 64, min_w: int = 64, downsample_multiple: int = 32) -> (int, int):
|
|
final_h = max(input_h, min_h)
|
|
final_w = max(input_w, min_w)
|
|
remainder_h = final_h % downsample_multiple
|
|
if remainder_h != 0:
|
|
down_dist = remainder_h
|
|
up_dist = downsample_multiple - remainder_h
|
|
if up_dist < down_dist:
|
|
final_h += up_dist
|
|
else:
|
|
if final_h - down_dist >= min_h:
|
|
final_h -= down_dist
|
|
else:
|
|
final_h += up_dist
|
|
remainder_w = final_w % downsample_multiple
|
|
if remainder_w != 0:
|
|
down_dist = remainder_w
|
|
up_dist = downsample_multiple - remainder_w
|
|
if up_dist < down_dist:
|
|
final_w += up_dist
|
|
else:
|
|
if final_w - down_dist >= min_w:
|
|
final_w -= down_dist
|
|
else:
|
|
final_w += up_dist
|
|
return final_h, final_w
|
|
|
|
def build_test_py_transform(image_size=(432,768)):
|
|
return Tlib.Compose([
|
|
CenterPad(image_size),
|
|
Tlib.CenterCrop(image_size),
|
|
RGB2Lab(),
|
|
ToTensor(),
|
|
Normalize(),
|
|
])
|
|
|
|
def tensor_to_pil(tensor_img: torch.Tensor) -> Image.Image:
|
|
if tensor_img.dim() == 4 and tensor_img.size(0) == 1:
|
|
tensor_img = tensor_img.squeeze(0)
|
|
if tensor_img.dim() == 3 and tensor_img.shape[-1] in (1, 3):
|
|
tensor_img = tensor_img.permute(2, 0, 1)
|
|
tensor_img = tensor_img.detach().cpu().clamp(0,1)
|
|
np_img = (tensor_img.numpy()*255).astype(np.uint8)
|
|
np_img = np.transpose(np_img, (1,2,0))
|
|
if np_img.shape[2] == 1:
|
|
return Image.fromarray(np_img[:,:,0], mode="L")
|
|
else:
|
|
return Image.fromarray(np_img, mode="RGB")
|
|
|
|
class DeepExColorImageNode:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image_to_colorize": ("IMAGE",),
|
|
"reference_image": ("IMAGE",),
|
|
"target_height": ("INT", {"default":432, "min":16, "max":4096}),
|
|
"target_width": ("INT", {"default":768, "min":16, "max":4096}),
|
|
"wls_filter_on": ("BOOLEAN", {"default":True}),
|
|
"lambda_value": ("FLOAT", {"default":500.0, "min":0, "max":2000}),
|
|
"sigma_color": ("FLOAT", {"default":4.0, "min":0, "max":50}),
|
|
}
|
|
}
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "colorize_image"
|
|
CATEGORY = "DeepExemplar/Image"
|
|
|
|
def colorize_image(self, image_to_colorize: torch.Tensor, reference_image: torch.Tensor,
|
|
target_height: int, target_width: int, wls_filter_on: bool,
|
|
lambda_value: float, sigma_color: float) -> (list,):
|
|
load_models_if_needed()
|
|
final_h, final_w = adjust_target_size(target_height, target_width, 64, 64, 32)
|
|
print(f"[DeepExColorImageNode] Adjusted size: ({target_height}, {target_width}) -> ({final_h}, {final_w})")
|
|
transform = build_test_py_transform((final_h, final_w))
|
|
src_pil = tensor_to_pil(image_to_colorize)
|
|
ref_pil = tensor_to_pil(reference_image)
|
|
src_lab_full = transform(src_pil).unsqueeze(0).cuda()
|
|
ref_lab_full = transform(ref_pil).unsqueeze(0).cuda()
|
|
src_half = F.interpolate(src_lab_full, scale_factor=0.5, mode="bilinear", align_corners=False)
|
|
ref_half = F.interpolate(ref_lab_full, scale_factor=0.5, mode="bilinear", align_corners=False)
|
|
I_ref_l = ref_half[:, 0:1, :, :]
|
|
I_ref_ab = ref_half[:, 1:3, :, :]
|
|
I_ref_rgb = tensor_lab2rgb(torch.cat((uncenter_l(I_ref_l), I_ref_ab), dim=1)).to(ref_half.device)
|
|
with torch.no_grad():
|
|
features_B = VGG_NET(I_ref_rgb, ["r12","r22","r32","r42","r52"], preprocess=True)
|
|
I_ab_predict, _, _ = frame_colorization(
|
|
src_half, ref_half, torch.zeros_like(src_half),
|
|
features_B, VGG_NET, NONLOCAL_NET, COLOR_NET,
|
|
feature_noise=0, temperature=1e-10
|
|
)
|
|
ab_up = F.interpolate(I_ab_predict, scale_factor=2.0, mode="bilinear", align_corners=False) * 1.25
|
|
if wls_filter_on and WLS_FILTER_AVAILABLE:
|
|
guide_l = uncenter_l(src_lab_full[:, 0:1, :, :]) * (255.0 / 100.0)
|
|
guide_np = guide_l[0, 0].detach().cpu().numpy().astype(np.uint8)
|
|
wlsf = cv2.ximgproc.createFastGlobalSmootherFilter(guide_np, lambda_value, sigma_color)
|
|
a_np = ab_up[0, 0].detach().cpu().numpy()
|
|
b_np = ab_up[0, 1].detach().cpu().numpy()
|
|
a_filt = wlsf.filter(a_np)
|
|
b_filt = wlsf.filter(b_np)
|
|
a_t = torch.from_numpy(a_filt).unsqueeze(0).unsqueeze(0).to(ab_up.device)
|
|
b_t = torch.from_numpy(b_filt).unsqueeze(0).unsqueeze(0).to(ab_up.device)
|
|
ab_up = torch.cat([a_t, b_t], dim=1)
|
|
else:
|
|
if wls_filter_on and not WLS_FILTER_AVAILABLE:
|
|
print("[DeepExColorImageNode] Warning: 'createFastGlobalSmootherFilter' not available; skipping filtering.")
|
|
out_np = batch_lab2rgb_transpose_mc(src_lab_full[:, 0:1, :, :], ab_up)
|
|
out_tensor = torch.from_numpy(out_np).float().div(255.0)
|
|
return ([out_tensor],)
|
|
|
|
class DeepExColorVideoNode:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"video_frames": ("IMAGE",),
|
|
"reference_image": ("IMAGE",),
|
|
"frame_propagate": ("BOOLEAN", {"default":True}),
|
|
"use_half_resolution": ("BOOLEAN", {"default":True}),
|
|
"target_height": ("INT", {"default":432, "min":16, "max":4096}),
|
|
"target_width": ("INT", {"default":768, "min":16, "max":4096}),
|
|
"wls_filter_on": ("BOOLEAN", {"default":True}),
|
|
"lambda_value": ("FLOAT", {"default":500.0, "min":0, "max":2000}),
|
|
"sigma_color": ("FLOAT", {"default":4.0, "min":0, "max":50}),
|
|
}
|
|
}
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "colorize_video"
|
|
CATEGORY = "DeepExemplar/Video"
|
|
|
|
def colorize_video(self, video_frames: any, reference_image: torch.Tensor,
|
|
frame_propagate: bool, use_half_resolution: bool,
|
|
target_height: int, target_width: int,
|
|
wls_filter_on: bool, lambda_value: float, sigma_color: float) -> (list,):
|
|
load_models_if_needed()
|
|
final_h, final_w = adjust_target_size(target_height, target_width, 64, 64, 32)
|
|
print(f"[DeepExColorVideoNode] Adjusted size: ({target_height}, {target_width}) -> ({final_h}, {final_w})")
|
|
if wls_filter_on and not WLS_FILTER_AVAILABLE:
|
|
print("[DeepExColorVideoNode] Warning: 'createFastGlobalSmootherFilter' not available; skipping filtering.")
|
|
if isinstance(video_frames, list):
|
|
frames_list = video_frames
|
|
elif isinstance(video_frames, torch.Tensor):
|
|
if video_frames.dim() == 4:
|
|
frames_list = [video_frames[i] for i in range(video_frames.size(0))]
|
|
else:
|
|
frames_list = [video_frames]
|
|
else:
|
|
frames_list = [video_frames]
|
|
print(f"[DeepExColorVideoNode] Number of frames: {len(frames_list)}")
|
|
transform = build_test_py_transform((final_h, final_w))
|
|
ref_pil = tensor_to_pil(reference_image)
|
|
ref_lab_full = transform(ref_pil).unsqueeze(0).cuda()
|
|
if use_half_resolution:
|
|
ref_half = F.interpolate(ref_lab_full, scale_factor=0.5, mode="bilinear", align_corners=False)
|
|
else:
|
|
ref_half = ref_lab_full.clone()
|
|
I_ref_l = ref_half[:, 0:1, :, :]
|
|
I_ref_ab = ref_half[:, 1:3, :, :]
|
|
I_ref_rgb = tensor_lab2rgb(torch.cat((uncenter_l(I_ref_l), I_ref_ab), dim=1)).to(ref_half.device)
|
|
with torch.no_grad():
|
|
features_B = VGG_NET(I_ref_rgb, ["r12", "r22", "r32", "r42", "r52"], preprocess=True)
|
|
out_frames = []
|
|
last_lab = None
|
|
num_frames = len(frames_list)
|
|
gui_pbar = ProgressBar(num_frames)
|
|
console_pbar = console_tqdm(total=num_frames, desc="[DeepExColorVideoNode] Console Progress")
|
|
for i, frm in enumerate(frames_list):
|
|
frame_pil = tensor_to_pil(frm)
|
|
frm_lab_full = transform(frame_pil).unsqueeze(0).cuda()
|
|
if use_half_resolution:
|
|
frm_half = F.interpolate(frm_lab_full, scale_factor=0.5, mode="bilinear", align_corners=False)
|
|
else:
|
|
frm_half = frm_lab_full.clone()
|
|
if i == 0:
|
|
with torch.no_grad():
|
|
I_ab, _, _ = frame_colorization(
|
|
frm_half, ref_half, torch.zeros_like(frm_half),
|
|
features_B, VGG_NET, NONLOCAL_NET, COLOR_NET,
|
|
feature_noise=0, temperature=1e-10
|
|
)
|
|
if frame_propagate:
|
|
last_lab = torch.cat([frm_half[:, 0:1, :, :], I_ab], dim=1)
|
|
else:
|
|
if frame_propagate and last_lab is not None:
|
|
with torch.no_grad():
|
|
I_ab, _, _ = frame_colorization(
|
|
frm_half, ref_half, last_lab,
|
|
features_B, VGG_NET, NONLOCAL_NET, COLOR_NET,
|
|
feature_noise=0, temperature=1e-10
|
|
)
|
|
last_lab = torch.cat([frm_half[:, 0:1, :, :], I_ab], dim=1)
|
|
else:
|
|
with torch.no_grad():
|
|
I_ab, _, _ = frame_colorization(
|
|
frm_half, ref_half, torch.zeros_like(frm_half),
|
|
features_B, VGG_NET, NONLOCAL_NET, COLOR_NET,
|
|
feature_noise=0, temperature=1e-10
|
|
)
|
|
last_lab = None
|
|
if use_half_resolution:
|
|
ab_up = F.interpolate(I_ab, scale_factor=2.0, mode="bilinear", align_corners=False)*1.25
|
|
else:
|
|
ab_up = I_ab.clone()
|
|
if wls_filter_on and WLS_FILTER_AVAILABLE:
|
|
guide_l = uncenter_l(frm_lab_full[:, 0:1, :, :])*(255.0/100.0)
|
|
guide_np = guide_l[0,0].detach().cpu().numpy().astype(np.uint8)
|
|
wlsf = cv2.ximgproc.createFastGlobalSmootherFilter(guide_np, lambda_value, sigma_color)
|
|
a_np = ab_up[0,0].detach().cpu().numpy()
|
|
b_np = ab_up[0,1].detach().cpu().numpy()
|
|
a_filt = wlsf.filter(a_np)
|
|
b_filt = wlsf.filter(b_np)
|
|
a_t = torch.from_numpy(a_filt).unsqueeze(0).unsqueeze(0).to(ab_up.device)
|
|
b_t = torch.from_numpy(b_filt).unsqueeze(0).unsqueeze(0).to(ab_up.device)
|
|
ab_up = torch.cat([a_t, b_t], dim=1)
|
|
out_np = batch_lab2rgb_transpose_mc(frm_lab_full[:, 0:1, :, :], ab_up)
|
|
out_t = torch.from_numpy(out_np).float().div(255.0)
|
|
out_frames.append(out_t)
|
|
gui_pbar.update_absolute(i+1, num_frames)
|
|
console_pbar.update(1)
|
|
if hasattr(self, 'set_progress'):
|
|
self.set_progress(gui_pbar.progress)
|
|
console_pbar.close()
|
|
return (out_frames,)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"DeepExColorImageNode": DeepExColorImageNode,
|
|
"DeepExColorVideoNode": DeepExColorVideoNode
|
|
}
|