From fbe4fc4eabfc155f6297e5b51c54285e1877c58a Mon Sep 17 00:00:00 2001 From: Blonbar <117039186+jonstreeter@users.noreply.github.com> Date: Sun, 16 Mar 2025 09:07:23 -0500 Subject: [PATCH] Add files via upload Functional --- DeepExemplarColorizationNodes.py | 298 +++++++++++++++++++++++++++++++ __init__.py | 55 +++++- install.py | 81 +++++++-- 3 files changed, 411 insertions(+), 23 deletions(-) create mode 100644 DeepExemplarColorizationNodes.py diff --git a/DeepExemplarColorizationNodes.py b/DeepExemplarColorizationNodes.py new file mode 100644 index 0000000..1c90f15 --- /dev/null +++ b/DeepExemplarColorizationNodes.py @@ -0,0 +1,298 @@ +""" +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 +} diff --git a/__init__.py b/__init__.py index a10732c..c16ad42 100644 --- a/__init__.py +++ b/__init__.py @@ -1,14 +1,51 @@ import os import sys +import torch +import types -# Add the parent directory of 'models' and 'utils' to the Python path -current_dir = os.path.dirname(os.path.abspath(__file__)) -if current_dir not in sys.path: - sys.path.append(current_dir) +############################################################################### +# 1) Patch torch.load for data/.pth so base code's relative paths +# are mapped to absolute paths in your local "data" folder. +############################################################################### +orig_torch_load = torch.load +REMAP_PATHS = { + "data/vgg19_gray.pth": os.path.join(os.path.dirname(__file__), "data", "vgg19_gray.pth"), +} +def custom_torch_load(path, *args, **kwargs): + if path in REMAP_PATHS: + path = REMAP_PATHS[path] + return orig_torch_load(path, *args, **kwargs) +torch.load = custom_torch_load -from .DeepExemplarImageColorNode import NODE_CLASS_MAPPINGS as IMG_MAP -from .DeepExemplarVideoColorNode import NODE_CLASS_MAPPINGS as VID_MAP +############################################################################### +# 2) Insert this custom node root directory into sys.path if not already present +############################################################################### +custom_node_root = os.path.dirname(os.path.abspath(__file__)) +if custom_node_root not in sys.path: + sys.path.insert(0, custom_node_root) -NODE_CLASS_MAPPINGS = {} -NODE_CLASS_MAPPINGS.update(IMG_MAP) -NODE_CLASS_MAPPINGS.update(VID_MAP) \ No newline at end of file +############################################################################### +# 3) Remove any pre-existing "utils" or "models" entries from sys.modules +############################################################################### +for modname in ["utils", "models"]: + if modname in sys.modules: + del sys.modules[modname] + +############################################################################### +# 4) Dynamically create "utils" and "models" packages, pointing them to our local directories. +############################################################################### +utils_dir = os.path.join(custom_node_root, "utils") +models_dir = os.path.join(custom_node_root, "models") +if os.path.isdir(utils_dir): + utils_module = types.ModuleType("utils") + utils_module.__path__ = [utils_dir] + sys.modules["utils"] = utils_module +if os.path.isdir(models_dir): + models_module = types.ModuleType("models") + models_module.__path__ = [models_dir] + sys.modules["models"] = models_module + +############################################################################### +# 5) Import and register node mappings. +############################################################################### +from .DeepExemplarColorizationNodes import NODE_CLASS_MAPPINGS diff --git a/install.py b/install.py index 850da66..00db141 100644 --- a/install.py +++ b/install.py @@ -1,40 +1,93 @@ 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 main(): +def ensure_opencv_contrib(): """ - Downloads `colorization_checkpoint.zip` from the official repo - and extracts it in this folder, - creating `checkpoints/video_moredata_l1` and `data/`. + Checks if cv2.ximgproc is available. If not, installs opencv-contrib-python. """ + 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) - # 1) Download if missing + # 3) Download the checkpoint zip if it's missing if not os.path.exists(zip_path): - print(f"Downloading from {CHECKPOINT_URL} ...") - urllib.request.urlretrieve(CHECKPOINT_URL, zip_path) - print("Download complete.") + 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 else: print(f"Zip file already exists: {zip_path}") - # 2) Extract + # 4) Extract the downloaded zip file print(f"Extracting {zip_path} into {root_dir} ...") - with zipfile.ZipFile(zip_path, "r") as zf: - zf.extractall(root_dir) - print("Extraction complete.") + try: + extract_zip(zip_path, root_dir) + print("Extraction complete.") + except Exception as e: + print(f"Error during extraction: {e}") + return - # Optionally remove the zip after extraction + # Optionally remove the zip file after successful extraction # os.remove(zip_path) - print("colorization_checkpoint files are installed in `checkpoints/` and `data/`.") + print("Checkpoint files are installed in the `checkpoints/` and `data/` directories.") if __name__ == "__main__": main()