Delete DeepExemplarColor.py
del
This commit is contained in:
Blonbar
@@ -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": <int>,
|
||||
...
|
||||
}
|
||||
"""
|
||||
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
|
||||
Reference in New Issue
Block a user