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

Add files via upload

Browse Source 
working version
This commit is contained in:
Blonbar
2025-03-28 08:51:47 -05:00
committed by GitHub
parent bdde8605e1
commit ba47d1fe36
2 changed files with 34 additions and 87 deletions
Download Patch File Download Diff File Expand all files Collapse all files
DeepExemplarColorizationNodes.py
+20 -20
View File
@@ -22,9 +22,9 @@ if not hasattr(cv2.ximgproc, 'createFastGlobalSmootherFilter'):
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.")
print("Installation attempted, but function is still unavailable. Please install or upgrade opencv-contrib-python (e.g., 'pip install --upgrade opencv-contrib-python') and restart ComfyUI.")
except Exception as e:
print(f"Failed to install opencv-contrib-python: {e}")
print(f"Failed to install opencv-contrib-python: {e}\nPlease install opencv-contrib-python manually and restart ComfyUI.")
WLS_FILTER_AVAILABLE = hasattr(cv2.ximgproc, 'createFastGlobalSmootherFilter')
@@ -36,7 +36,7 @@ from PIL import Image
import torchvision.transforms as Tlib
from tqdm import tqdm as console_tqdm
# Import ComfyUI's internal ProgressBar for GUI updates.
# Import ComfyUI's internal ProgressBar for GUI progress updates.
from comfy.utils import ProgressBar
from utils.util_distortion import CenterPad, Normalize, RGB2Lab, ToTensor
@@ -52,7 +52,6 @@ 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")
@@ -125,8 +124,8 @@ class DeepExColorImageNode:
"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}),
"target_height": ("INT", {"default":432, "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}),
@@ -137,11 +136,12 @@ class DeepExColorImageNode:
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,
target_width: int, target_height: int, wls_filter_on: bool,
lambda_value: float, sigma_color: float) -> (list,):
load_models_if_needed()
# Note: adjust_target_size expects (height, width); here we pass target_height, target_width.
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})")
print(f"[DeepExColorImageNode] Adjusted size: ({target_width}, {target_height}) -> ({final_w}, {final_h})")
transform = build_test_py_transform((final_h, final_w))
src_pil = tensor_to_pil(image_to_colorize)
ref_pil = tensor_to_pil(reference_image)
@@ -153,7 +153,7 @@ class DeepExColorImageNode:
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)
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,
@@ -161,7 +161,7 @@ class DeepExColorImageNode:
)
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_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()
@@ -173,7 +173,7 @@ class DeepExColorImageNode:
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.")
print("[DeepExColorImageNode] Warning: 'createFastGlobalSmootherFilter' not available; skipping filtering. To fix, install or upgrade opencv-contrib-python (e.g., 'pip install --upgrade opencv-contrib-python') and restart ComfyUI.")
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],)
@@ -187,8 +187,8 @@ class DeepExColorVideoNode:
"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}),
"target_height": ("INT", {"default":432, "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}),
@@ -200,13 +200,13 @@ class DeepExColorVideoNode:
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,):
target_width: int, target_height: 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})")
print(f"[DeepExColorVideoNode] Adjusted size: ({target_width}, {target_height}) -> ({final_w}, {final_h})")
if wls_filter_on and not WLS_FILTER_AVAILABLE:
print("[DeepExColorVideoNode] Warning: 'createFastGlobalSmootherFilter' not available; skipping filtering.")
print("[DeepExColorVideoNode] Warning: 'createFastGlobalSmootherFilter' not available; skipping filtering. To fix, install or upgrade opencv-contrib-python (e.g., 'pip install --upgrade opencv-contrib-python') and restart ComfyUI.")
if isinstance(video_frames, list):
frames_list = video_frames
elif isinstance(video_frames, torch.Tensor):
@@ -268,15 +268,15 @@ class DeepExColorVideoNode:
)
last_lab = None
if use_half_resolution:
ab_up = F.interpolate(I_ab, scale_factor=2.0, mode="bilinear", align_corners=False)*1.25
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)
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_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)
install.py
+14 -67
View File
@@ -1,93 +1,40 @@
import os
import sys
import urllib.request
import zipfile
import subprocess
from tqdm import tqdm
# URL to download the checkpoints (colorization_checkpoint.zip)
CHECKPOINT_URL = (
"https://github.com/zhangmozhe/Deep-Exemplar-based-Video-Colorization/"
"releases/download/v1.0/colorization_checkpoint.zip"
)
ZIP_FILENAME = "colorization_checkpoint.zip"
def ensure_opencv_contrib():
def main():
"""
Checks if cv2.ximgproc is available. If not, installs opencv-contrib-python.
Downloads `colorization_checkpoint.zip` from the official repo
and extracts it in this folder,
creating `checkpoints/video_moredata_l1` and `data/`.
"""
try:
import cv2
_ = cv2.ximgproc # just a test import
print("[install.py] opencv-contrib-python is already installed.")
except (ImportError, AttributeError):
print("[install.py] Installing opencv-contrib-python for WLS filter support...")
subprocess.check_call([sys.executable, "-m", "pip", "install", "--upgrade", "opencv-contrib-python"])
print("[install.py] opencv-contrib-python installed successfully.")
def download_file(url: str, dest_path: str) -> None:
"""
Downloads a file from a URL to a local destination with a progress bar.
Args:
url (str): The URL from which to download the file.
dest_path (str): The file path to save the downloaded file.
"""
with tqdm(unit='B', unit_scale=True, miniters=1, desc=os.path.basename(dest_path)) as t:
def reporthook(block_num, block_size, total_size):
if total_size > 0:
t.total = total_size
t.update(block_size)
urllib.request.urlretrieve(url, dest_path, reporthook=reporthook)
def extract_zip(zip_path: str, extract_to: str) -> None:
"""
Extracts a zip file into the specified directory.
Args:
zip_path (str): The path to the zip file.
extract_to (str): The directory where files will be extracted.
"""
with zipfile.ZipFile(zip_path, "r") as zf:
zf.extractall(extract_to)
def main() -> None:
"""
Main function that downloads and extracts the checkpoint files.
This will create the necessary `checkpoints/` and `data/` directories.
"""
# 1) Ensure opencv-contrib-python is present (for cv2.ximgproc)
ensure_opencv_contrib()
# 2) Determine the root directory (where this script resides)
root_dir = os.path.dirname(os.path.abspath(__file__))
zip_path = os.path.join(root_dir, ZIP_FILENAME)
# 3) Download the checkpoint zip if it's missing
# 1) Download if missing
if not os.path.exists(zip_path):
print(f"Downloading checkpoints from {CHECKPOINT_URL} ...")
try:
download_file(CHECKPOINT_URL, zip_path)
print("Download complete.")
except Exception as e:
print(f"Error during download: {e}")
return
print(f"Downloading from {CHECKPOINT_URL} ...")
urllib.request.urlretrieve(CHECKPOINT_URL, zip_path)
print("Download complete.")
else:
print(f"Zip file already exists: {zip_path}")
# 4) Extract the downloaded zip file
# 2) Extract
print(f"Extracting {zip_path} into {root_dir} ...")
try:
extract_zip(zip_path, root_dir)
print("Extraction complete.")
except Exception as e:
print(f"Error during extraction: {e}")
return
with zipfile.ZipFile(zip_path, "r") as zf:
zf.extractall(root_dir)
print("Extraction complete.")
# Optionally remove the zip file after successful extraction
# Optionally remove the zip after extraction
# os.remove(zip_path)
print("Checkpoint files are installed in the `checkpoints/` and `data/` directories.")
print("colorization_checkpoint files are installed in `checkpoints/` and `data/`.")
if __name__ == "__main__":
main()