From 075ee71879b862c1e491ec28abff0a2d793dfe9e Mon Sep 17 00:00:00 2001 From: Blonbar <117039186+jonstreeter@users.noreply.github.com> Date: Fri, 28 Mar 2025 12:26:02 -0500 Subject: [PATCH] Delete DeepExemplarColor.py del --- DeepExemplarColor.py | 220 ------------------------------------------- 1 file changed, 220 deletions(-) delete mode 100644 DeepExemplarColor.py diff --git a/DeepExemplarColor.py b/DeepExemplarColor.py deleted file mode 100644 index 03659e6..0000000 --- a/DeepExemplarColor.py +++ /dev/null @@ -1,220 +0,0 @@ -import os -import torch -import torch.nn.functional as F -import numpy as np -import cv2 -from PIL import Image -from typing import Dict, Any - -########################################### -# ComfyUI Node Registration -########################################### -NODE_CLASS_MAPPINGS = {} - -########################################### -# Imports from your local code -########################################### -from .models.NonlocalNet import WarpNet, VGG19_pytorch -from .models.ColorVidNet import ColorVidNet -from .models.FrameColor import frame_colorization -# Add any other imports you need from .utils or .models - -########################################### -# Paths to 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") -VGG_GRAY_PATH = os.path.join(ROOT_DIR, "data", "vgg19_gray.pth") # if you use it - -########################################### -# Global Single Load (lazy load approach) -########################################### -MODELS_LOADED = False -NONLOCAL_NET = None -COLOR_NET = None -VGG_NET = None - -def load_models_if_needed(): - """Load colorization models only once. - If they are already loaded, do nothing.""" - global MODELS_LOADED, NONLOCAL_NET, COLOR_NET, VGG_NET - if MODELS_LOADED: - return - - # Initialize models - print("[DeepExemplar] Loading nonlocal, color, and VGG19 models...") - nonlocal_net = WarpNet(1) - colornet = ColorVidNet(7) - vggnet = VGG19_pytorch() - - # Load checkpoints - nonlocal_net.load_state_dict(torch.load(NONLOCAL_TEST_PATH, map_location="cuda")) - colornet.load_state_dict(torch.load(COLOR_TEST_PATH, map_location="cuda")) - vggnet.load_state_dict(torch.load(VGG_PATH, map_location="cuda")) - - nonlocal_net.cuda().eval() - colornet.cuda().eval() - vggnet.cuda().eval() - - NONLOCAL_NET = nonlocal_net - COLOR_NET = colornet - VGG_NET = vggnet - MODELS_LOADED = True - print("[DeepExemplar] Models loaded successfully.") - -########################################### -# Utility: Convert PIL->Lab Tensor (Simplified) -########################################### -def pil_to_lab_tensor(img: Image.Image, size=(432, 768)): - """ - Resizes to `size`, then does BGR->Lab via OpenCV. - Normalization is a simplified approach. - Replicate your exact 'TestTransforms' from the original code if needed. - """ - img = img.convert("RGB").resize((size[1], size[0]), Image.BICUBIC) - arr = np.array(img)[:, :, ::-1] # PIL to BGR - arr_lab = cv2.cvtColor(arr, cv2.COLOR_BGR2LAB).astype(np.float32) - - # Convert to tensor shape [1,3,H,W] - lab_t = torch.from_numpy(arr_lab).permute(2, 0, 1).unsqueeze(0) - # Simple scaling to ~[-1..+1] - L = (lab_t[:, 0:1, :, :] - 50.) / 50. - A = lab_t[:, 1:2, :, :] / 128. - B = lab_t[:, 2:3, :, :] / 128. - lab_t = torch.cat([L, A, B], dim=1).cuda() - return lab_t - -########################################### -# Utility: Convert Lab->PIL -########################################### -def lab_to_pil(lab_t: torch.Tensor): - """ - Inverse of pil_to_lab_tensor. shape [1,3,H,W], roughly in [-1..+1]. - Output: PIL.Image in RGB - """ - L = (lab_t[:, 0:1, :, :] * 50.) + 50. - A = lab_t[:, 1:2, :, :] * 128. - B = lab_t[:, 2:3, :, :] * 128. - - lab_255 = torch.cat([L, A, B], dim=1).clamp(0, 255) - np_lab = lab_255[0].permute(1, 2, 0).cpu().numpy().astype(np.uint8) - - bgr = cv2.cvtColor(np_lab, cv2.COLOR_Lab2BGR) - rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) - return Image.fromarray(rgb) - -########################################### -# The Node -########################################### -class DeepExemplarVideoColorNode: - """ - A ComfyUI node that colorizes a list of frames (VIDEO) using a single reference image. - The result is a new dictionary with colorized frames. - """ - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "video_input": ("VIDEO", ), - "reference_image": ("IMAGE", ), - "frame_propagate": ("BOOL", {"default": False}), - }, - } - - RETURN_TYPES = ("VIDEO",) - FUNCTION = "colorize_video" - CATEGORY = "Video/DeepExemplar" - - def colorize_video(self, video_input: Dict[str, Any], reference_image: Image.Image, frame_propagate: bool): - """ - ComfyUI calls this. 'video_input' is expected to have: - { - "frames": [PIL.Image, PIL.Image, ...], - "fps": int, - ... - } - 'reference_image' is a single PIL image. - 'frame_propagate' toggles using the last colorized frame as the next reference. - - Returns a dictionary with colorized frames, e.g.: - { - "frames": [...], - "fps": , - ... - } - """ - load_models_if_needed() - - frames = video_input.get("frames", []) - if not frames: - print("[DeepExemplar] No frames found in the input video.") - return (video_input,) # just return the input unmodified - - # If frame_propagate is True, we use the first frame as reference - if frame_propagate: - reference_image = frames[0] - - # Convert reference to Lab (half-res) - ref_lab_half = F.interpolate(pil_to_lab_tensor(reference_image), scale_factor=0.5, mode="bilinear") - # We should compute 'features_B' from the reference if fully replicating test.py logic: - # But for brevity, let's skip. In a real implementation, you'd replicate test.py's approach - # (converting reference to approximate RGB, passing through vgg, etc.) - features_B = None # or your real feature extraction code - - last_lab = None # for frame_propagation - out_frames = [] - - for idx, frame_pil in enumerate(frames): - # Convert frame to Lab half-res - frame_lab_half = F.interpolate(pil_to_lab_tensor(frame_pil), scale_factor=0.5, mode="bilinear") - - if last_lab is None: - if frame_propagate: - last_lab = ref_lab_half.clone() - else: - last_lab = torch.zeros_like(frame_lab_half) - - # Actual colorization logic - # We call frame_colorization(...). - # test.py calls: - # frame_colorization(I_current_lab, I_reference_lab, I_last_lab, features_B, vggnet, nonlocal_net, color_net, ...) - # We'll do a simplified approach: - I_current_ab_predict, *_ = frame_colorization( - frame_lab_half, # I_current_lab - ref_lab_half, # I_reference_lab - last_lab, # I_last_lab - features_B, # features_B - VGG_NET, - NONLOCAL_NET, - COLOR_NET, - feature_noise=0, - temperature=1e-10, - ) - - # Update last_lab - last_lab = torch.cat([frame_lab_half[:, 0:1, :, :], I_current_ab_predict], dim=1) - - # Upsample predicted ab back to full resolution - ab_up = F.interpolate(I_current_ab_predict, scale_factor=2.0, mode="bilinear") * 1.25 - - # Combine with full-res L - full_lab = pil_to_lab_tensor(frame_pil) # Full-res L - L_full = full_lab[:, 0:1, :, :] - colorized_lab = torch.cat([L_full, ab_up], dim=1) - - # Convert Lab->PIL - colorized_pil = lab_to_pil(colorized_lab) - out_frames.append(colorized_pil) - - # Return new dictionary - new_video = { - "frames": out_frames, - "fps": video_input.get("fps", 30), - # Copy other fields if needed - } - return (new_video,) - -# Register the node -NODE_CLASS_MAPPINGS["DeepExemplarVideoColorNode"] = DeepExemplarVideoColorNode