From 3eef088d35283b7c118536bf0bfbb8006e391b56 Mon Sep 17 00:00:00 2001 From: Blonbar <117039186+jonstreeter@users.noreply.github.com> Date: Tue, 11 Mar 2025 11:36:22 -0500 Subject: [PATCH] Add files via upload current snapshot - still buggy --- DeepExemplarImageColorNode.py | 238 ++++++++++++++++++++++ DeepExemplarVideoColorNode.py | 365 ++++++++++++++++++++++++++++++++++ __init__.py | 15 +- requirements.txt | 2 +- 4 files changed, 618 insertions(+), 2 deletions(-) create mode 100644 DeepExemplarImageColorNode.py create mode 100644 DeepExemplarVideoColorNode.py diff --git a/DeepExemplarImageColorNode.py b/DeepExemplarImageColorNode.py new file mode 100644 index 0000000..a00ae39 --- /dev/null +++ b/DeepExemplarImageColorNode.py @@ -0,0 +1,238 @@ +import os +import sys +import torch +import torch.nn.functional as F +import torchvision.transforms as T +import numpy as np +from PIL import Image +from typing import Any, List + +############################################################################### +# Optional: fix relative paths for "data/vgg19_gray.pth" if used +############################################################################### +current_dir = os.path.dirname(os.path.abspath(__file__)) +orig_torch_load = torch.load +def custom_torch_load(path, *args, **kwargs): + if path == "data/vgg19_gray.pth": + path = os.path.join(current_dir, "data", "vgg19_gray.pth") + return orig_torch_load(path, *args, **kwargs) +torch.load = custom_torch_load +############################################################################### + +# Import from your local code +import importlib.abc +import importlib.machinery + +class CustomImportFinder(importlib.abc.MetaPathFinder): + def __init__(self, base_path): + self.base_path = base_path + def find_spec(self, fullname, path, target=None): + if fullname.startswith('utils.') or fullname == 'utils': + parts = fullname.split('.') + rel_path = os.path.join(*parts) + ".py" + full_path = os.path.join(self.base_path, rel_path) + if os.path.exists(full_path): + return importlib.machinery.ModuleSpec( + name=fullname, + loader=importlib.machinery.SourceFileLoader(fullname, full_path), + origin=full_path + ) + return None + +sys.meta_path.insert(0, CustomImportFinder(current_dir)) + +from models.NonlocalNet import WarpNet, VGG19_pytorch +from models.ColorVidNet import ColorVidNet +from models.FrameColor import frame_colorization +from utils.util import batch_lab2rgb_transpose_mc, tensor_lab2rgb, uncenter_l +from utils.util_distortion import CenterPad, Normalize, RGB2Lab, ToTensor + +import cv2 +try: + _ = cv2.ximgproc +except AttributeError: + print("[DeepExemplarImageColorNode] Warning: ximgproc not found. WLS filter may be unavailable.") + +# ComfyUI node registry +NODE_CLASS_MAPPINGS = {} + +# Paths to your model checkpoints +ROOT_DIR = os.path.dirname(__file__) +NONLOCAL_TEST_PATH = os.path.join(ROOT_DIR, "checkpoints", "video_moredata_l1", "nonlocal_net_iter_76000.pth") +COLOR_TEST_PATH = os.path.join(ROOT_DIR, "checkpoints", "video_moredata_l1", "colornet_iter_76000.pth") +VGG_PATH = os.path.join(ROOT_DIR, "data", "vgg19_conv.pth") + +# Global references +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("[DeepExemplarImageColorNode] Loading models...") + nonlocal_net = WarpNet(1) + color_net = ColorVidNet(7) + vgg_net = VGG19_pytorch() + + nonlocal_net.load_state_dict(torch.load(NONLOCAL_TEST_PATH, map_location="cuda")) + color_net.load_state_dict(torch.load(COLOR_TEST_PATH, map_location="cuda")) + vgg_net.load_state_dict(torch.load(VGG_PATH, map_location="cuda")) + + nonlocal_net.cuda().eval() + color_net.cuda().eval() + vgg_net.cuda().eval() + + MODELS_LOADED = True + NONLOCAL_NET = nonlocal_net + COLOR_NET = color_net + VGG_NET = vgg_net + print("[DeepExemplarImageColorNode] Models loaded.") + +############################################################################### +# Utilities to convert Torch Tensors <-> PIL +############################################################################### +def tensor_to_pil(image_tensor: torch.Tensor) -> Image.Image: + """ + Expects shape [C,H,W]. Converts to PIL (RGB). + Clamps to [0..1], then scales to [0..255]. + """ + image_tensor = image_tensor.detach().cpu().clamp(0,1) + np_img = (image_tensor.numpy() * 255).astype(np.uint8) + np_img = np.transpose(np_img, (1,2,0)) # [H,W,C] + pil_img = Image.fromarray(np_img, mode="RGB") + return pil_img + +def pil_to_tensor(pil_img: Image.Image) -> torch.Tensor: + """ + Converts a PIL (RGB) to Torch float tensor [C,H,W], scaled 0..1. + """ + return T.ToTensor()(pil_img) + +############################################################################### +# Transforms from test.py +############################################################################### +class ComposeTransforms: + def __init__(self, steps): + self.steps = steps + def __call__(self, img): + for step in self.steps: + img = step(img) + return img + +def build_transform(image_size=(432,768)): + return ComposeTransforms([ + CenterPad(image_size), + T.CenterCrop(image_size), + RGB2Lab(), + ToTensor(), + Normalize(), + ]) + +############################################################################### +# The Node +############################################################################### +class DeepExemplarImageColorNode: + """ + Accepts a single image (Torch Tensor or PIL) + optional boolean WLS, etc. + Outputs a single colorized image (Torch Tensor). + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "input_image": ("IMAGE",), + "wls_filter_on": ("BOOLEAN", { + "default": True, + "forceInput": False, + "widget": "checkbox" + }), + "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" + + def colorize_image(self, + input_image: Any, + wls_filter_on: bool, + lambda_value: float, + sigma_color: float) -> (torch.Tensor,): + + load_models_if_needed() + + # 1) Convert input to PIL + pil_img = None + if isinstance(input_image, torch.Tensor): + # shape [C,H,W] + pil_img = tensor_to_pil(input_image) + elif isinstance(input_image, Image.Image): + pil_img = input_image + else: + print("[DeepExemplarImageColorNode] Unsupported image type. Returning empty.") + return (torch.zeros((3,64,64)),) + + # 2) Prepare transforms + transform = build_transform((432,768)) + input_lab = transform(pil_img).unsqueeze(0).cuda() # shape [1,3,H,W] + input_small = F.interpolate(input_lab, scale_factor=0.5, mode="bilinear") + + # We'll treat the reference as itself for a single image (no external ref). + ref_tensor_small = input_small.clone() + + # Precompute VGG features + with torch.no_grad(): + I_l = input_small[:,0:1,:,:] + I_ab = input_small[:,1:3,:,:] + I_rgb = tensor_lab2rgb(torch.cat((uncenter_l(I_l), I_ab), dim=1)) + features_B = VGG_NET(I_rgb, ["r12","r22","r32","r42","r52"], preprocess=True) + + # colorize + with torch.no_grad(): + I_current_ab_predict, _, _ = frame_colorization( + I_current_lab=input_small, + I_reference_lab=ref_tensor_small, + I_last_lab=torch.zeros_like(input_small), + features_B=features_B, + vggnet=VGG_NET, + nonlocal_net=NONLOCAL_NET, + color_net=COLOR_NET, + feature_noise=0, + temperature=1e-10, + ) + + # Upscale ab + ab_up = F.interpolate(I_current_ab_predict, scale_factor=2.0, mode="bilinear") * 1.25 + + # optional WLS + if wls_filter_on: + guide_image = uncenter_l(input_lab[:,0:1,:,:]) * 255/100.0 + guide_numpy = guide_image[0,0].detach().cpu().numpy().astype(np.uint8) + wls_filter = cv2.ximgproc.createFastGlobalSmootherFilter( + guide_numpy, lambda_value, sigma_color + ) + a_channel = ab_up[0,0].detach().cpu().numpy() + b_channel = ab_up[0,1].detach().cpu().numpy() + a_filtered = wls_filter.filter(a_channel) + b_filtered = wls_filter.filter(b_channel) + a_filt_t = torch.from_numpy(a_filtered).unsqueeze(0).unsqueeze(0) + b_filt_t = torch.from_numpy(b_filtered).unsqueeze(0).unsqueeze(0) + ab_up = torch.cat([a_filt_t, b_filt_t], dim=1).to(ab_up.device).unsqueeze(0) + + # Convert final to RGB + out_rgb = batch_lab2rgb_transpose_mc(input_lab[:,0:1,:,:], ab_up[:1]) + np_rgb = out_rgb[0].astype(np.uint8) + color_pil = Image.fromarray(np_rgb, mode="RGB") + + # Convert back to Torch + output_tensor = pil_to_tensor(color_pil) # shape [C,H,W], float 0..1 + return (output_tensor,) + + +NODE_CLASS_MAPPINGS["DeepExemplarImageColorNode"] = DeepExemplarImageColorNode diff --git a/DeepExemplarVideoColorNode.py b/DeepExemplarVideoColorNode.py new file mode 100644 index 0000000..b9a2192 --- /dev/null +++ b/DeepExemplarVideoColorNode.py @@ -0,0 +1,365 @@ +import os +import sys +import torch +import torch.nn.functional as F +import numpy as np +from typing import Any, List +from tqdm import tqdm + +import cv2 +from PIL import Image + +############################################################################### +# Global references for the colorization nets +############################################################################### +MODELS_LOADED = False +NONLOCAL_NET = None +COLOR_NET = None +VGG_NET = None + +############################################################################### +# Paths from install.py +############################################################################### +current_dir = os.path.dirname(os.path.abspath(__file__)) +nonlocal_ckpt = os.path.join(current_dir, "checkpoints", "video_moredata_l1", "nonlocal_net_iter_76000.pth") +color_ckpt = os.path.join(current_dir, "checkpoints", "video_moredata_l1", "colornet_iter_76000.pth") +vgg_ckpt = os.path.join(current_dir, "data", "vgg19_conv.pth") + +############################################################################### +# Helper: get height,width from PIL or Torch +############################################################################### +def get_hw(image_or_tensor): + """Return (height, width) from a PIL image or Torch [C,H,W] or [N,C,H,W].""" + if isinstance(image_or_tensor, Image.Image): + return image_or_tensor.height, image_or_tensor.width + elif isinstance(image_or_tensor, torch.Tensor): + # shape [C,H,W] => (H,W) + # shape [N,C,H,W] => pick the first + if image_or_tensor.dim()==4: + if image_or_tensor.size(0)==0: + return 0,0 + return image_or_tensor.shape[-2], image_or_tensor.shape[-1] + elif image_or_tensor.dim()==3: + return image_or_tensor.shape[-2], image_or_tensor.shape[-1] + else: + return 0,0 + else: + return 0,0 + +############################################################################### +# CenterPadToRef: Pad or Crop to (ref_h, ref_w) +############################################################################### +import numpy as np +import torchvision.transforms as Tlib + +class CenterPadToRef: + """ + If the image is smaller in a dimension => pad + If bigger => center-crop + Expects a PIL image as input. + """ + def __init__(self, ref_h, ref_w): + self.ref_h = ref_h + self.ref_w = ref_w + def __call__(self, pil_img): + arr = np.array(pil_img) + h_old, w_old = arr.shape[:2] + out_h, out_w = self.ref_h, self.ref_w + # If smaller => pad, if bigger => crop + # We'll do a naive approach + if h_old==out_h and w_old==out_w: + return pil_img + # create an output + c_old = 1 if arr.ndim==3 and arr.shape[2]==1 else 3 + out_arr = np.zeros((out_h, out_w, c_old), dtype=arr.dtype) + # if bigger => crop + if h_old>out_h: + start_h = (h_old - out_h)//2 + arr = arr[start_h:start_h+out_h,:,:] + h_old = out_h + if w_old>out_w: + start_w = (w_old - out_w)//2 + arr = arr[:, start_w:start_w+out_w, :] + w_old = out_w + # if smaller => pad + off_h = (out_h - h_old)//2 if out_h>h_old else 0 + off_w = (out_w - w_old)//2 if out_w>w_old else 0 + h_end = off_h + h_old + w_end = off_w + w_old + out_arr[off_h:h_end, off_w:w_end] = arr + return Image.fromarray(out_arr) + +############################################################################### +# Convert PIL or Torch to Lab tensor (L => L-50) +############################################################################### +from skimage import color +def convert_any_image_to_lab_tensor(item): + """ + If PIL => do np.array => rgb2lab => (L-50) + If Torch => shape [C,H,W] or [N,C,H,W], convert to CPU np => rgb2lab => (L-50) + """ + if isinstance(item, Image.Image): + arr = np.array(item).astype(np.float32) + if arr.ndim==2: + arr = np.stack([arr,arr,arr], axis=-1) + arr /= 255.0 + lab = color.rgb2lab(arr) + lab_t = torch.from_numpy(lab.transpose(2,0,1)) + lab_t[0:1,:,:] = lab_t[0:1,:,:]-50.0 + return lab_t + elif isinstance(item, torch.Tensor): + # if shape [N,C,H,W], pick first + if item.dim()==4: + if item.size(0)==0: + return None + item = item[0] + # now [C,H,W] + arr = item.detach().cpu().float() + # guess scale + if arr.max()>2.0: + arr = arr/255.0 + np_arr = arr.numpy() + np_arr = np_arr.transpose(1,2,0) # [H,W,C] + lab = color.rgb2lab(np_arr) + lab_t = torch.from_numpy(lab.transpose(2,0,1)) + lab_t[0:1,:,:] = lab_t[0:1,:,:]-50.0 + return lab_t + else: + return None + +############################################################################### +# final step: batch_lab2rgb_transpose_mc, optional WLS +############################################################################### +from utils.util import batch_lab2rgb_transpose_mc, tensor_lab2rgb, uncenter_l +from models.NonlocalNet import WarpNet, VGG19_pytorch +from models.ColorVidNet import ColorVidNet +from models.FrameColor import frame_colorization + +############################################################################### +# The Node +############################################################################### +NODE_CLASS_MAPPINGS = {} + +class DeepExemplarVideoColorNodeTestPyLike: + """ + Node that: + - checks if reference_image is PIL or Torch + - obtains (ref_h, ref_w) + - optional center pad/crop each frame to (ref_h, ref_w) + - convert to Lab => downsample => colorize => upsample => optional WLS => final batch_lab2rgb + - frame_propagate => first frame is reference + - WLS defaults: lambda=500, sigma=4 + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "video_frames": ("IMAGE",), + "reference_image": ("IMAGE",), + "frame_propagate": ("BOOLEAN", { + "default": False, + "widget": "checkbox" + }), + "use_center_pad": ("BOOLEAN", { + "default": False, + "widget": "checkbox" + }), + "use_center_crop": ("BOOLEAN", { + "default": False, + "widget": "checkbox" + }), + "wls_filter_on": ("BOOLEAN", { + "default": True, + "widget": "checkbox" + }), + "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 = "Video/DeepExemplar" + + def load_models_if_needed(self): + global MODELS_LOADED, NONLOCAL_NET, COLOR_NET, VGG_NET + if MODELS_LOADED: + return + print("[DeepExemplarVideoColorNodeTestPyLike] Loading models from local checkpoints...") + + nonlocal_net = WarpNet(1) + color_net = ColorVidNet(7) + vgg_net = VGG19_pytorch() + + 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")) + + nonlocal_net.cuda().eval() + color_net.cuda().eval() + vgg_net.cuda().eval() + + NONLOCAL_NET = nonlocal_net + COLOR_NET = color_net + VGG_NET = vgg_net + MODELS_LOADED = True + print("[DeepExemplarVideoColorNodeTestPyLike] Models loaded.") + + def colorize_video(self, + video_frames: Any, + reference_image: Any, + frame_propagate: bool, + use_center_pad: bool, + use_center_crop: bool, + wls_filter_on: bool, + lambda_value: float, + sigma_color: float + ) -> (List[np.ndarray],): + + self.load_models_if_needed() + + # if frame_propagate => first frame is reference + # unify frames => list + frames_list = [] + if isinstance(video_frames, list): + frames_list = video_frames + elif isinstance(video_frames, torch.Tensor): + # shape [N,C,H,W] or [C,H,W] + if video_frames.dim()==4: + frames_list = [video_frames[i] for i in range(video_frames.size(0))] + else: + frames_list = [video_frames] + elif isinstance(video_frames, Image.Image): + frames_list = [video_frames] + + if frame_propagate and len(frames_list)>0: + reference_image = frames_list[0] + + # get (ref_h, ref_w) + ref_h, ref_w = get_hw(reference_image) + if ref_h<1 or ref_w<1: + print("[DeepExemplarVideoColorNodeTestPyLike] reference_image is zero dimension.") + return ([],) + + # build optional pad/crop + pad_transform = CenterPadToRef(ref_h, ref_w) if use_center_pad else None + import torchvision.transforms as Tlib + crop_transform = Tlib.CenterCrop((ref_h, ref_w)) if use_center_crop else None + + # Convert reference => lab + # if it's PIL or Torch + # if pad/crop => must be PIL + # so if reference_image is Torch, we convert to PIL if needed + ref_pil = reference_image + if isinstance(reference_image, torch.Tensor): + # convert to PIL + ref_pil = self.torch_to_pil(reference_image) + + if pad_transform: + ref_pil = pad_transform(ref_pil) + if crop_transform: + ref_pil = crop_transform(ref_pil) + + ref_lab_large = convert_any_image_to_lab_tensor(ref_pil).unsqueeze(0).cuda() + ref_lab_small = F.interpolate(ref_lab_large, scale_factor=0.5, mode="bilinear", align_corners=False) + + with torch.no_grad(): + ref_l = ref_lab_small[:,0:1,:,:]+50.0 + ref_ab = ref_lab_small[:,1:3,:,:] + ref_unnorm = torch.cat([ref_l, ref_ab], dim=1) + ref_rgb = tensor_lab2rgb(ref_unnorm) + features_B = VGG_NET(ref_rgb, ["r12","r22","r32","r42","r52"], preprocess=True) + + last_lab_predict = ref_lab_small.clone() if frame_propagate else None + + out_frames = [] + for idx, item in enumerate(tqdm(frames_list, desc="[DeepExemplar Node]")): + # unify item => PIL + if isinstance(item, torch.Tensor): + pil_frame = self.torch_to_pil(item) + else: + pil_frame = item + + # pad/crop if toggles + if pad_transform: + pil_frame = pad_transform(pil_frame) + if crop_transform: + pil_frame = crop_transform(pil_frame) + + # convert to lab + lab_large = convert_any_image_to_lab_tensor(pil_frame) + if lab_large is None: + print(f"[DeepExemplar] skipping frame {idx}, invalid.") + continue + lab_large = lab_large.unsqueeze(0).cuda() + lab_small = F.interpolate(lab_large, scale_factor=0.5, mode="bilinear", align_corners=False) + + IA_l = lab_small[:,0:1,:,:] + if last_lab_predict is None: + last_lab_predict = torch.zeros_like(lab_small) + + with torch.no_grad(): + I_current_ab_predict, _, _ = frame_colorization( + lab_small, + ref_lab_small, + last_lab_predict, + features_B, + VGG_NET, + NONLOCAL_NET, + COLOR_NET, + False, + 1e-10 + ) + last_lab_predict = torch.cat([IA_l, I_current_ab_predict], dim=1) + + ab_up = F.interpolate(I_current_ab_predict, scale_factor=2.0, mode="bilinear", align_corners=False)*1.25 + + if wls_filter_on: + guide_l = (lab_large[:,0:1,:,:]+50.0).clamp(0,100)*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).cuda() + b_t = torch.from_numpy(b_filt).unsqueeze(0).unsqueeze(0).cuda() + ab_up = torch.cat([a_t,b_t], dim=1) + + out_rgb = batch_lab2rgb_transpose_mc(lab_large[:,0:1,:,:]+50.0, ab_up) + np_rgb = out_rgb[0].astype(np.uint8) + out_frames.append(np_rgb) + + return (out_frames,) + + def torch_to_pil(self, tensor: torch.Tensor): + """ + Convert a Torch [C,H,W] or [N,C,H,W] in 0..1 or 0..255 to a PIL image. + We'll just pick the first if [N,C,H,W]. + """ + if tensor.dim()==4: + if tensor.size(0)==0: + # return blank + arr = np.zeros((1,1,3), dtype=np.uint8) + return Image.fromarray(arr) + tensor = tensor[0] + arr = tensor.detach().cpu().float() + if arr.max()>2.0: + arr = arr/255.0 + arr_np = arr.numpy() # shape [C,H,W] + arr_np = arr_np.transpose(1,2,0) # [H,W,C] + arr_np = (arr_np*255).clip(0,255).astype(np.uint8) + pil_img = Image.fromarray(arr_np) + return pil_img + + +NODE_CLASS_MAPPINGS["DeepExemplarVideoColorNodeTestPyLike"] = DeepExemplarVideoColorNodeTestPyLike diff --git a/__init__.py b/__init__.py index 8b13789..a10732c 100644 --- a/__init__.py +++ b/__init__.py @@ -1 +1,14 @@ - +import os +import sys + +# 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) + +from .DeepExemplarImageColorNode import NODE_CLASS_MAPPINGS as IMG_MAP +from .DeepExemplarVideoColorNode import NODE_CLASS_MAPPINGS as VID_MAP + +NODE_CLASS_MAPPINGS = {} +NODE_CLASS_MAPPINGS.update(IMG_MAP) +NODE_CLASS_MAPPINGS.update(VID_MAP) \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 4b3a54e..961023e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,7 +3,6 @@ tqdm>=4.23.0 scipy==1.2 torchsummary>=1.3 matplotlib>=2.2.2 -opencv_contrib_python>=4.1.0.25 torchvision>=0.3.0 scikit_image>=0.15.0 torchviz>=0.0.1 @@ -17,3 +16,4 @@ future numba pypng opencv-python>=4.5 +opencv-contrib-python \ No newline at end of file