init

.gitignore

@@ -0,0 +1,2 @@
+**/__pycache__
+.DS_Store

__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",
+}

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

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,)

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)

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,)