diff --git a/DeepExemplarImageColorNode.py b/DeepExemplarImageColorNode.py deleted file mode 100644 index a00ae39..0000000 --- a/DeepExemplarImageColorNode.py +++ /dev/null @@ -1,238 +0,0 @@ -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