Add files via upload

Functional

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
+}

__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/<filename>.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)
+###############################################################################
+# 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

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()