From a4a1edda8eb3fd8db69f02c90603eb12f6c8a159 Mon Sep 17 00:00:00 2001 From: logtd Date: Tue, 10 Sep 2024 20:44:28 -0600 Subject: [PATCH] init --- .gitignore | 2 + __init__.py | 16 ++++ modules/neural_texture_embedding.py | 118 ++++++++++++++++++++++++++ nodes/load_dilight_controlnet_node.py | 55 ++++++++++++ nodes/load_dilight_image_node.py | 72 ++++++++++++++++ nodes/prepare_dilight_cond_node.py | 25 ++++++ 6 files changed, 288 insertions(+) create mode 100644 .gitignore create mode 100644 __init__.py create mode 100644 modules/neural_texture_embedding.py create mode 100644 nodes/load_dilight_controlnet_node.py create mode 100644 nodes/load_dilight_image_node.py create mode 100644 nodes/prepare_dilight_cond_node.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..bd1c29a --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +**/__pycache__ +.DS_Store diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..2f0a6ad --- /dev/null +++ b/__init__.py @@ -0,0 +1,16 @@ +from .nodes.load_dilight_controlnet_node import LoadDiLightControlNetNode +from .nodes.prepare_dilight_cond_node import PrepareDiLightCondNode +from .nodes.load_dilight_image_node import LoadDiLightImageNode + + +NODE_CLASS_MAPPINGS = { + "LoadDiLightControlNet": LoadDiLightControlNetNode, + "PrepareDiLightCond": PrepareDiLightCondNode, + "LoadDiLightImage": LoadDiLightImageNode, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "LoadDiLightControlNet": "Load DiLight ControlNet", + "PrepareDiLightCond": "Prepare DiLight CN Image", + "LoadDiLightImage": "Load DiLight Image", +} diff --git a/modules/neural_texture_embedding.py b/modules/neural_texture_embedding.py new file mode 100644 index 0000000..b438a1b --- /dev/null +++ b/modules/neural_texture_embedding.py @@ -0,0 +1,118 @@ +from typing import Tuple + +import torch +import torch.nn.functional as F +from torch import nn + +import comfy.ops +ops = comfy.ops.disable_weight_init + + +class ResBlock(nn.Module): + def __init__(self, dim): + super().__init__() + self.conv = nn.Sequential( + ops.Conv2d(dim, dim, 3, 1, 1), + ops.GroupNorm(num_groups=8, num_channels=dim), + nn.SiLU(inplace=True), + ops.Conv2d(dim, dim, 3, 1, 1), + ) + + def forward(self, x): + return x + self.conv(x) + + +class NeuralTextureEncoder(nn.Module): + def __init__(self, in_dim=3, out_dim=16, dims=(32, 64, 128), groups=8): + super().__init__() + self.model = nn.Sequential( + ops.Conv2d(in_dim, dims[0], kernel_size=3, padding=1), + nn.SiLU(inplace=True), + + # down 1 + ops.Conv2d(dims[0], dims[1], kernel_size=3, padding=1, stride=2), + ops.GroupNorm(num_groups=groups, num_channels=dims[1]), + nn.SiLU(inplace=True), + + # down 2 + ops.Conv2d(dims[1], dims[2], kernel_size=3, padding=1, stride=2), + ops.GroupNorm(num_groups=groups, num_channels=dims[2]), + nn.SiLU(inplace=True), + + # res blocks + ResBlock(dims[2]), + ResBlock(dims[2]), + ResBlock(dims[2]), + ResBlock(dims[2]), + + # up 1 + ops.ConvTranspose2d(dims[2], dims[1], kernel_size=4, padding=1, stride=2), + ops.GroupNorm(num_groups=groups, num_channels=dims[1]), + nn.SiLU(inplace=True), + + # up 2 + ops.ConvTranspose2d(dims[1], dims[0], kernel_size=4, padding=1, stride=2), + ops.GroupNorm(num_groups=groups, num_channels=dims[0]), + nn.SiLU(inplace=True), + + # out + ops.Conv2d(dims[0], out_dim, kernel_size=3, padding=1), + ) + self.gradient_checkpointing = False + + def forward(self, x): + x = self.model(x) + return x + + +class NeuralTextureEmbedding(nn.Module): + def __init__( + self, + conditioning_embedding_channels: int, + conditioning_channels: int = 3, + block_out_channels: Tuple[int] = (16, 32, 96, 256), + shading_hint_channels: int = 12, # diffuse + 3 * ggx + ): + super().__init__() + self.conditioning_channels = conditioning_channels + self.shading_hint_channels = shading_hint_channels + + self.conv_in = ops.Conv2d(shading_hint_channels, block_out_channels[0], kernel_size=3, padding=1) + self.neural_texture_encoder = NeuralTextureEncoder(in_dim=conditioning_channels, out_dim=shading_hint_channels) + + self.blocks = nn.ModuleList([]) + + for i in range(len(block_out_channels) - 1): + channel_in = block_out_channels[i] + channel_out = block_out_channels[i + 1] + self.blocks.append(ops.Conv2d(channel_in, channel_in, kernel_size=3, padding=1)) + self.blocks.append(ops.Conv2d(channel_in, channel_out, kernel_size=3, padding=1, stride=2)) + + self.conv_out = ops.Conv2d( + block_out_channels[-1], + conditioning_embedding_channels, + kernel_size=3, + padding=1 + ) + + def forward(self, all_conditioning, emb=None, context=None): + # conditioning: [BS, 4 + 12, 512, 512] # RGB ref image + shading hint (diffuse + 3 * ggx) + conditioning, shading_hint = torch.split( + all_conditioning, + [self.conditioning_channels, self.shading_hint_channels], + dim=1 + ) + embedding = self.neural_texture_encoder(conditioning) # [BS, 15, 512, 512] + + # multiply shading hint to each channel + embedding = embedding * shading_hint + embedding = self.conv_in(embedding) + embedding = F.silu(embedding) + + for block in self.blocks: + embedding = block(embedding) + embedding = F.silu(embedding) + + embedding = self.conv_out(embedding) + + return embedding diff --git a/nodes/load_dilight_controlnet_node.py b/nodes/load_dilight_controlnet_node.py new file mode 100644 index 0000000..df4ac04 --- /dev/null +++ b/nodes/load_dilight_controlnet_node.py @@ -0,0 +1,55 @@ + +import torch + +import folder_paths +import comfy.utils +from comfy.cldm.cldm import ControlNet +from nodes import ControlNetLoader + +from ..modules.neural_texture_embedding import NeuralTextureEmbedding + + +class DiLightControlnet(ControlNet): + def set_input_hint_block(self, input_hint_block): + self.input_hint_block = input_hint_block + + def to(self, device): + self.input_hint_block.to(device) + + return super().to(device) + + +class LoadDiLightControlNetNode: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "control_net_name": (folder_paths.get_filename_list("controlnet"), ), + }} + RETURN_TYPES = ("CONTROL_NET",) + FUNCTION = "apply" + + CATEGORY = "dilight" + + def apply(self, control_net_name): + controlnet, = ControlNetLoader().load_controlnet(control_net_name) + controlnet_path = folder_paths.get_full_path("controlnet", control_net_name) + controlnet_data = comfy.utils.load_torch_file(controlnet_path, safe_load=True) + state_dict = {} + for key in list(controlnet_data.keys()): + if 'controlnet_cond_embedding.' in key: + new_key = key.replace('controlnet_cond_embedding.', '') + state_dict[new_key] = controlnet_data[key] + + # https://huggingface.co/dilightnet/DiLightNet/blob/main/config.json + embedder = NeuralTextureEmbedding( + conditioning_embedding_channels=320, + conditioning_channels=4, + block_out_channels=[16, 32, 96, 256], + shading_hint_channels=12 + ).to(torch.float16) + embedder.load_state_dict(state_dict) + + controlnet.control_model.__class__ = DiLightControlnet + controlnet.control_model.set_input_hint_block(embedder) + + return (controlnet,) diff --git a/nodes/load_dilight_image_node.py b/nodes/load_dilight_image_node.py new file mode 100644 index 0000000..1eec180 --- /dev/null +++ b/nodes/load_dilight_image_node.py @@ -0,0 +1,72 @@ +import os + +import torch +import numpy as np +from PIL import Image, ImageOps, ImageSequence + +import folder_paths +import node_helpers + + +class LoadDiLightImageNode: + @classmethod + def INPUT_TYPES(s): + input_dir = folder_paths.get_input_directory() + files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))] + return {"required": + {"image": (sorted(files), {"image_upload": True})}, + } + + CATEGORY = "image" + + RETURN_TYPES = ("IMAGE", "MASK") + FUNCTION = "load_image" + def load_image(self, image): + image_path = folder_paths.get_annotated_filepath(image) + + img = node_helpers.pillow(Image.open, image_path) + + output_images = [] + output_masks = [] + w, h = None, None + + excluded_formats = ['MPO'] + + for i in ImageSequence.Iterator(img): + i = node_helpers.pillow(ImageOps.exif_transpose, i) + + if i.mode == 'I': + i = i.point(lambda i: i * (1 / 255)) + if 'A' in i.getbands(): + image = i.convert("RGBA") + else: + image = i.convert("RGB") + + if len(output_images) == 0: + w = image.size[0] + h = image.size[1] + + if image.size[0] != w or image.size[1] != h: + continue + + image = np.array(image).astype(np.float32) / 255.0 + image = torch.from_numpy(image)[None,] + if 'A' in i.getbands(): + mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0 + mask = 1. - torch.from_numpy(mask) + else: + mask = torch.zeros((64,64), dtype=torch.float32, device="cpu") + output_images.append(image) + output_masks.append(mask.unsqueeze(0)) + + if len(output_images) > 1 and img.format not in excluded_formats: + output_image = torch.cat(output_images, dim=0) + output_mask = torch.cat(output_masks, dim=0) + else: + output_image = output_images[0] + output_mask = output_masks[0] + + if output_image.shape[-1] == 4: + output_image = output_image[..., :3] * output_image[..., 3:] + + return (output_image, output_mask) diff --git a/nodes/prepare_dilight_cond_node.py b/nodes/prepare_dilight_cond_node.py new file mode 100644 index 0000000..cbaa84a --- /dev/null +++ b/nodes/prepare_dilight_cond_node.py @@ -0,0 +1,25 @@ +import torch + + +class PrepareDiLightCondNode: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "mask": ("MASK",), + "image": ("IMAGE",), + "diff_image": ("IMAGE",), + "ggx5": ("IMAGE",), + "ggx13": ("IMAGE",), + "ggx34": ("IMAGE",), + }} + RETURN_TYPES = ("IMAGE",) + FUNCTION = "prepare" + + CATEGORY = "dilight" + + def prepare(self, mask, image, diff_image, ggx5, ggx13, ggx34): + # shape: mask + [1, 3, 512, 512] * 4 -> [1, 16, 512, 512] + mask = mask.unsqueeze(-1) + cond = torch.cat([diff_image, ggx5, ggx13, ggx34], dim=3) * mask + cond = torch.cat([mask, image, cond], dim=3) + return (cond,)