init

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