107 lines
4.1 KiB
Python
107 lines
4.1 KiB
Python
import os
|
|
import torch
|
|
import cv2
|
|
import numpy as np
|
|
from .ddcolor.ddcolor_arch import DDColor
|
|
import torch.nn.functional as F
|
|
import comfy.model_management
|
|
from huggingface_hub import snapshot_download
|
|
|
|
script_directory = os.path.dirname(os.path.abspath(__file__))
|
|
|
|
class DDColor_Colorize:
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"image": ("IMAGE", ),
|
|
"model_input_size": ("INT", {"default": 512,"min": 32, "max": 4096, "step": 32}),
|
|
"checkpoint": (
|
|
[
|
|
"ddcolor_paper_tiny.pth",
|
|
"ddcolor_paper.pth",
|
|
"ddcolor_modelscope.pth",
|
|
"ddcolor_artistic.pth",
|
|
], {
|
|
"default": "ddcolor_paper_tiny.pth"
|
|
}),
|
|
|
|
|
|
},
|
|
|
|
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES =("colorized_image",)
|
|
FUNCTION = "process"
|
|
|
|
CATEGORY = "DDColor"
|
|
@torch.no_grad()
|
|
def process(self, image, model_input_size, checkpoint):
|
|
self.device = comfy.model_management.get_torch_device()
|
|
batch_size = image.shape[0]
|
|
self.input_size = model_input_size
|
|
self.checkpoint = checkpoint
|
|
self.checkpoint_folder = os.path.join(script_directory, f"checkpoints")
|
|
self.checkpoint_path = os.path.join(script_directory, f"checkpoints/{checkpoint}")
|
|
|
|
if not os.path.isfile(self.checkpoint_path):
|
|
try:
|
|
snapshot_download(repo_id="piddnad/DDColor-models", allow_patterns=[self.checkpoint], local_dir=self.checkpoint_folder, local_dir_use_symlinks=False)
|
|
except:
|
|
raise FileNotFoundError("Checkpoint load failed.")
|
|
if not hasattr(self, "model") or not hasattr(self, "ddcolor_model") or self.model is None or self.checkpoint != self.ddcolor_model:
|
|
|
|
self.ddcolor_model = self.checkpoint
|
|
if self.ddcolor_model == "ddcolor_paper_tiny.pth":
|
|
encoder="convnext-t"
|
|
else:
|
|
encoder="convnext-l"
|
|
self.model = DDColor(
|
|
encoder_name=encoder,
|
|
decoder_name="MultiScaleColorDecoder",
|
|
input_size=[self.input_size, self.input_size],
|
|
num_output_channels=2,
|
|
last_norm="Spectral",
|
|
do_normalize=False,
|
|
num_queries=100,
|
|
num_scales=3,
|
|
dec_layers=9,
|
|
).to(self.device)
|
|
self.model.load_state_dict(torch.load(self.checkpoint_path, map_location=torch.device("cpu"))["params"], strict=False)
|
|
self.model.eval()
|
|
|
|
out=[]
|
|
for i in range(batch_size):
|
|
self.height, self.width = image.shape[1:3]
|
|
img = image[i].numpy().astype(np.float32)
|
|
|
|
orig_l = cv2.cvtColor(img, cv2.COLOR_RGB2LAB)[:, :, :1] # (h, w, 1)
|
|
|
|
# resize rgb image -> lab -> get grey -> rgb
|
|
img = cv2.resize(img, (self.input_size, self.input_size))
|
|
img_l = cv2.cvtColor(img, cv2.COLOR_RGB2LAB)[:, :, :1]
|
|
img_gray_lab = np.concatenate((img_l, np.zeros_like(img_l), np.zeros_like(img_l)), axis=-1)
|
|
img_gray_rgb = cv2.cvtColor(img_gray_lab, cv2.COLOR_LAB2RGB)
|
|
|
|
tensor_gray_rgb = torch.from_numpy(img_gray_rgb.transpose((2, 0, 1))).float().unsqueeze(0).to(self.device)
|
|
output_ab = self.model(tensor_gray_rgb).cpu() # (1, 2, self.height, self.width)
|
|
|
|
# resize ab -> concat original l -> rgb
|
|
output_ab_resize = F.interpolate(output_ab, size=(self.height, self.width))[0].float().numpy().transpose(1, 2, 0)
|
|
output_lab = np.concatenate((orig_l, output_ab_resize), axis=-1)
|
|
|
|
output_rgb = cv2.cvtColor(output_lab, cv2.COLOR_LAB2RGB)
|
|
output_img = torch.from_numpy(output_rgb).float() # CHW format and add batch dimension
|
|
out.append(output_img)
|
|
|
|
batch_out = torch.stack(out, dim=0)
|
|
return(batch_out,)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"DDColor_Colorize": DDColor_Colorize,
|
|
}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"DDColor_Colorize": "DDColor_Colorize",
|
|
} |