From 65db494f8bace9734e6e7788768c43b26326efd4 Mon Sep 17 00:00:00 2001 From: Kijai <40791699+kijai@users.noreply.github.com> Date: Wed, 17 Jan 2024 15:28:22 +0200 Subject: [PATCH] Initial commit --- README.md | 4 + __init__.py | 3 + checkpoints/place_checkpoints_here.txt | 1 + ddcolor/__init__.py | 0 ddcolor/ddcolor_arch.py | 382 ++++++++++++++++++ ddcolor/ddcolor_arch_utils/__int__.py | 0 ddcolor/ddcolor_arch_utils/convnext.py | 155 +++++++ .../ddcolor_arch_utils/position_encoding.py | 52 +++ ddcolor/ddcolor_arch_utils/transformer.py | 368 +++++++++++++++++ .../ddcolor_arch_utils/transformer_utils.py | 192 +++++++++ ddcolor/ddcolor_arch_utils/unet.py | 208 ++++++++++ nodes.py | 107 +++++ requirements.txt | 5 + 13 files changed, 1477 insertions(+) create mode 100644 README.md create mode 100644 __init__.py create mode 100644 checkpoints/place_checkpoints_here.txt create mode 100644 ddcolor/__init__.py create mode 100644 ddcolor/ddcolor_arch.py create mode 100644 ddcolor/ddcolor_arch_utils/__int__.py create mode 100644 ddcolor/ddcolor_arch_utils/convnext.py create mode 100644 ddcolor/ddcolor_arch_utils/position_encoding.py create mode 100644 ddcolor/ddcolor_arch_utils/transformer.py create mode 100644 ddcolor/ddcolor_arch_utils/transformer_utils.py create mode 100644 ddcolor/ddcolor_arch_utils/unet.py create mode 100644 nodes.py create mode 100644 requirements.txt diff --git a/README.md b/README.md new file mode 100644 index 0000000..4d53daa --- /dev/null +++ b/README.md @@ -0,0 +1,4 @@ +#ComfyUI-DDColor + +Node to use DDColor (https://github.com/piddnad/DDColor) in ComfyUI + diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..2e96bd6 --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] \ No newline at end of file diff --git a/checkpoints/place_checkpoints_here.txt b/checkpoints/place_checkpoints_here.txt new file mode 100644 index 0000000..f448309 --- /dev/null +++ b/checkpoints/place_checkpoints_here.txt @@ -0,0 +1 @@ +https://huggingface.co/piddnad/DDColor-models \ No newline at end of file diff --git a/ddcolor/__init__.py b/ddcolor/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/ddcolor/ddcolor_arch.py b/ddcolor/ddcolor_arch.py new file mode 100644 index 0000000..cca2441 --- /dev/null +++ b/ddcolor/ddcolor_arch.py @@ -0,0 +1,382 @@ +import torch +import torch.nn as nn + +from ..ddcolor.ddcolor_arch_utils.unet import Hook, CustomPixelShuffle_ICNR, UnetBlockWide, NormType, custom_conv_layer +from ..ddcolor.ddcolor_arch_utils.convnext import ConvNeXt +from ..ddcolor.ddcolor_arch_utils.transformer_utils import SelfAttentionLayer, CrossAttentionLayer, FFNLayer, MLP +from ..ddcolor.ddcolor_arch_utils.position_encoding import PositionEmbeddingSine +from ..ddcolor.ddcolor_arch_utils.transformer import Transformer + +class DDColor(nn.Module): + + def __init__(self, + encoder_name='convnext-l', + decoder_name='MultiScaleColorDecoder', + num_input_channels=3, + input_size=(256, 256), + nf=512, + num_output_channels=3, + last_norm='Weight', + do_normalize=False, + num_queries=256, + num_scales=3, + dec_layers=9, + encoder_from_pretrain=False): + super().__init__() + + self.encoder = Encoder(encoder_name, ['norm0', 'norm1', 'norm2', 'norm3'], from_pretrain=encoder_from_pretrain) + self.encoder.eval() + test_input = torch.randn(1, num_input_channels, *input_size) + self.encoder(test_input) + + self.decoder = Decoder( + self.encoder.hooks, + nf=nf, + last_norm=last_norm, + num_queries=num_queries, + num_scales=num_scales, + dec_layers=dec_layers, + decoder_name=decoder_name + ) + self.refine_net = nn.Sequential(custom_conv_layer(num_queries + 3, num_output_channels, ks=1, use_activ=False, norm_type=NormType.Spectral)) + + self.do_normalize = do_normalize + self.register_buffer('mean', torch.Tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)) + self.register_buffer('std', torch.Tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)) + + def normalize(self, img): + return (img - self.mean) / self.std + + def denormalize(self, img): + return img * self.std + self.mean + + def forward(self, x): + if x.shape[1] == 3: + x = self.normalize(x) + + self.encoder(x) + out_feat = self.decoder() + coarse_input = torch.cat([out_feat, x], dim=1) + out = self.refine_net(coarse_input) + + if self.do_normalize: + out = self.denormalize(out) + return out + + +class Decoder(nn.Module): + + def __init__(self, + hooks, + nf=512, + blur=True, + last_norm='Weight', + num_queries=256, + num_scales=3, + dec_layers=9, + decoder_name='MultiScaleColorDecoder'): + super().__init__() + self.hooks = hooks + self.nf = nf + self.blur = blur + self.last_norm = getattr(NormType, last_norm) + self.decoder_name = decoder_name + + self.layers = self.make_layers() + embed_dim = nf // 2 + + self.last_shuf = CustomPixelShuffle_ICNR(embed_dim, embed_dim, blur=self.blur, norm_type=self.last_norm, scale=4) + + if self.decoder_name == 'MultiScaleColorDecoder': + self.color_decoder = MultiScaleColorDecoder( + in_channels=[512, 512, 256], + num_queries=num_queries, + num_scales=num_scales, + dec_layers=dec_layers, + ) + else: + self.color_decoder = SingleColorDecoder( + in_channels=hooks[-1].feature.shape[1], + num_queries=num_queries, + ) + + + def forward(self): + encode_feat = self.hooks[-1].feature + out0 = self.layers[0](encode_feat) + out1 = self.layers[1](out0) + out2 = self.layers[2](out1) + out3 = self.last_shuf(out2) + + if self.decoder_name == 'MultiScaleColorDecoder': + out = self.color_decoder([out0, out1, out2], out3) + else: + out = self.color_decoder(out3, encode_feat) + + return out + + def make_layers(self): + decoder_layers = [] + + e_in_c = self.hooks[-1].feature.shape[1] + in_c = e_in_c + + out_c = self.nf + setup_hooks = self.hooks[-2::-1] + for layer_index, hook in enumerate(setup_hooks): + feature_c = hook.feature.shape[1] + if layer_index == len(setup_hooks) - 1: + out_c = out_c // 2 + decoder_layers.append( + UnetBlockWide( + in_c, feature_c, out_c, hook, blur=self.blur, self_attention=False, norm_type=NormType.Spectral)) + in_c = out_c + return nn.Sequential(*decoder_layers) + + +class Encoder(nn.Module): + + def __init__(self, encoder_name, hook_names, from_pretrain, **kwargs): + super().__init__() + + if encoder_name == 'convnext-t' or encoder_name == 'convnext': + self.arch = ConvNeXt() + elif encoder_name == 'convnext-s': + self.arch = ConvNeXt(depths=[3, 3, 27, 3], dims=[96, 192, 384, 768]) + elif encoder_name == 'convnext-b': + self.arch = ConvNeXt(depths=[3, 3, 27, 3], dims=[128, 256, 512, 1024]) + elif encoder_name == 'convnext-l': + self.arch = ConvNeXt(depths=[3, 3, 27, 3], dims=[192, 384, 768, 1536]) + else: + raise NotImplementedError + + self.encoder_name = encoder_name + self.hook_names = hook_names + self.hooks = self.setup_hooks() + + if from_pretrain: + self.load_pretrain_model() + + def setup_hooks(self): + hooks = [Hook(self.arch._modules[name]) for name in self.hook_names] + return hooks + + def forward(self, x): + return self.arch(x) + + def load_pretrain_model(self): + if self.encoder_name == 'convnext-t' or self.encoder_name == 'convnext': + self.load('pretrain/convnext_tiny_22k_224.pth') + elif self.encoder_name == 'convnext-s': + self.load('pretrain/convnext_small_22k_224.pth') + elif self.encoder_name == 'convnext-b': + self.load('pretrain/convnext_base_22k_224.pth') + elif self.encoder_name == 'convnext-l': + self.load('pretrain/convnext_large_22k_224.pth') + else: + raise NotImplementedError + print('Loaded pretrained convnext model.') + + def load(self, path): + from basicsr.utils import get_root_logger + logger = get_root_logger() + if not path: + logger.info("No checkpoint found. Initializing model from scratch") + return + logger.info("[Encoder] Loading from {} ...".format(path)) + checkpoint = torch.load(path, map_location=torch.device("cpu")) + checkpoint_state_dict = checkpoint['model'] if 'model' in checkpoint.keys() else checkpoint + incompatible = self.arch.load_state_dict(checkpoint_state_dict, strict=False) + + if incompatible.missing_keys: + msg = "Some model parameters or buffers are not found in the checkpoint:\n" + msg += str(incompatible.missing_keys) + logger.warning(msg) + if incompatible.unexpected_keys: + msg = "The checkpoint state_dict contains keys that are not used by the model:\n" + msg += str(incompatible.unexpected_keys) + logger.warning(msg) + + +class MultiScaleColorDecoder(nn.Module): + + def __init__( + self, + in_channels, + hidden_dim=256, + num_queries=100, + nheads=8, + dim_feedforward=2048, + dec_layers=9, + pre_norm=False, + color_embed_dim=256, + enforce_input_project=True, + num_scales=3 + ): + super().__init__() + + # positional encoding + N_steps = hidden_dim // 2 + self.pe_layer = PositionEmbeddingSine(N_steps, normalize=True) + + # define Transformer decoder here + self.num_heads = nheads + self.num_layers = dec_layers + self.transformer_self_attention_layers = nn.ModuleList() + self.transformer_cross_attention_layers = nn.ModuleList() + self.transformer_ffn_layers = nn.ModuleList() + + for _ in range(self.num_layers): + self.transformer_self_attention_layers.append( + SelfAttentionLayer( + d_model=hidden_dim, + nhead=nheads, + dropout=0.0, + normalize_before=pre_norm, + ) + ) + self.transformer_cross_attention_layers.append( + CrossAttentionLayer( + d_model=hidden_dim, + nhead=nheads, + dropout=0.0, + normalize_before=pre_norm, + ) + ) + self.transformer_ffn_layers.append( + FFNLayer( + d_model=hidden_dim, + dim_feedforward=dim_feedforward, + dropout=0.0, + normalize_before=pre_norm, + ) + ) + + self.decoder_norm = nn.LayerNorm(hidden_dim) + + self.num_queries = num_queries + # learnable color query features + self.query_feat = nn.Embedding(num_queries, hidden_dim) + # learnable color query p.e. + self.query_embed = nn.Embedding(num_queries, hidden_dim) + + # level embedding + self.num_feature_levels = num_scales + self.level_embed = nn.Embedding(self.num_feature_levels, hidden_dim) + + # input projections + self.input_proj = nn.ModuleList() + for i in range(self.num_feature_levels): + if in_channels[i] != hidden_dim or enforce_input_project: + self.input_proj.append(nn.Conv2d(in_channels[i], hidden_dim, kernel_size=1)) + nn.init.kaiming_uniform_(self.input_proj[-1].weight, a=1) + if self.input_proj[-1].bias is not None: + nn.init.constant_(self.input_proj[-1].bias, 0) + else: + self.input_proj.append(nn.Sequential()) + + # output FFNs + self.color_embed = MLP(hidden_dim, hidden_dim, color_embed_dim, 3) + + def forward(self, x, img_features): + # x is a list of multi-scale feature + assert len(x) == self.num_feature_levels + src = [] + pos = [] + + for i in range(self.num_feature_levels): + pos.append(self.pe_layer(x[i], None).flatten(2)) + src.append(self.input_proj[i](x[i]).flatten(2) + self.level_embed.weight[i][None, :, None]) + + # flatten NxCxHxW to HWxNxC + pos[-1] = pos[-1].permute(2, 0, 1) + src[-1] = src[-1].permute(2, 0, 1) + + _, bs, _ = src[0].shape + + # QxNxC + query_embed = self.query_embed.weight.unsqueeze(1).repeat(1, bs, 1) + output = self.query_feat.weight.unsqueeze(1).repeat(1, bs, 1) + + for i in range(self.num_layers): + level_index = i % self.num_feature_levels + # attention: cross-attention first + output = self.transformer_cross_attention_layers[i]( + output, src[level_index], + memory_mask=None, + memory_key_padding_mask=None, + pos=pos[level_index], query_pos=query_embed + ) + output = self.transformer_self_attention_layers[i]( + output, tgt_mask=None, + tgt_key_padding_mask=None, + query_pos=query_embed + ) + # FFN + output = self.transformer_ffn_layers[i]( + output + ) + + decoder_output = self.decoder_norm(output) + decoder_output = decoder_output.transpose(0, 1) # [N, bs, C] -> [bs, N, C] + color_embed = self.color_embed(decoder_output) + out = torch.einsum("bqc,bchw->bqhw", color_embed, img_features) + + return out + + +class SingleColorDecoder(nn.Module): + + def __init__( + self, + in_channels=768, + hidden_dim=256, + num_queries=256, # 100 + nheads=8, + dropout=0.1, + dim_feedforward=2048, + enc_layers=0, + dec_layers=6, + pre_norm=False, + deep_supervision=True, + enforce_input_project=True, + ): + + super().__init__() + + N_steps = hidden_dim // 2 + self.pe_layer = PositionEmbeddingSine(N_steps, normalize=True) + + transformer = Transformer( + d_model=hidden_dim, + dropout=dropout, + nhead=nheads, + dim_feedforward=dim_feedforward, + num_encoder_layers=enc_layers, + num_decoder_layers=dec_layers, + normalize_before=pre_norm, + return_intermediate_dec=deep_supervision, + ) + self.num_queries = num_queries + self.transformer = transformer + + self.query_embed = nn.Embedding(num_queries, hidden_dim) + + if in_channels != hidden_dim or enforce_input_project: + self.input_proj = nn.Conv2d(in_channels, hidden_dim, kernel_size=1) + nn.init.kaiming_uniform_(self.input_proj.weight, a=1) + if self.input_proj.bias is not None: + nn.init.constant_(self.input_proj.bias, 0) + else: + self.input_proj = nn.Sequential() + + + def forward(self, img_features, encode_feat): + pos = self.pe_layer(encode_feat) + src = encode_feat + mask = None + hs, memory = self.transformer(self.input_proj(src), mask, self.query_embed.weight, pos) + color_embed = hs[-1] + color_preds = torch.einsum('bqc,bchw->bqhw', color_embed, img_features) + return color_preds + diff --git a/ddcolor/ddcolor_arch_utils/__int__.py b/ddcolor/ddcolor_arch_utils/__int__.py new file mode 100644 index 0000000..e69de29 diff --git a/ddcolor/ddcolor_arch_utils/convnext.py b/ddcolor/ddcolor_arch_utils/convnext.py new file mode 100644 index 0000000..e44fd6b --- /dev/null +++ b/ddcolor/ddcolor_arch_utils/convnext.py @@ -0,0 +1,155 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. + +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + + +import torch +import torch.nn as nn +import torch.nn.functional as F +from timm.models.layers import trunc_normal_, DropPath + +class Block(nn.Module): + r""" ConvNeXt Block. There are two equivalent implementations: + (1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W) + (2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back + We use (2) as we find it slightly faster in PyTorch + + Args: + dim (int): Number of input channels. + drop_path (float): Stochastic depth rate. Default: 0.0 + layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6. + """ + def __init__(self, dim, drop_path=0., layer_scale_init_value=1e-6): + super().__init__() + self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim) # depthwise conv + self.norm = LayerNorm(dim, eps=1e-6) + self.pwconv1 = nn.Linear(dim, 4 * dim) # pointwise/1x1 convs, implemented with linear layers + self.act = nn.GELU() + self.pwconv2 = nn.Linear(4 * dim, dim) + self.gamma = nn.Parameter(layer_scale_init_value * torch.ones((dim)), + requires_grad=True) if layer_scale_init_value > 0 else None + self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() + + def forward(self, x): + input = x + x = self.dwconv(x) + x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C) + x = self.norm(x) + x = self.pwconv1(x) + x = self.act(x) + x = self.pwconv2(x) + if self.gamma is not None: + x = self.gamma * x + x = x.permute(0, 3, 1, 2) # (N, H, W, C) -> (N, C, H, W) + + x = input + self.drop_path(x) + return x + +class ConvNeXt(nn.Module): + r""" ConvNeXt + A PyTorch impl of : `A ConvNet for the 2020s` - + https://arxiv.org/pdf/2201.03545.pdf + Args: + in_chans (int): Number of input image channels. Default: 3 + num_classes (int): Number of classes for classification head. Default: 1000 + depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3] + dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768] + drop_path_rate (float): Stochastic depth rate. Default: 0. + layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6. + head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1. + """ + def __init__(self, in_chans=3, num_classes=1000, + depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], drop_path_rate=0., + layer_scale_init_value=1e-6, head_init_scale=1., + ): + super().__init__() + + self.downsample_layers = nn.ModuleList() # stem and 3 intermediate downsampling conv layers + stem = nn.Sequential( + nn.Conv2d(in_chans, dims[0], kernel_size=4, stride=4), + LayerNorm(dims[0], eps=1e-6, data_format="channels_first") + ) + self.downsample_layers.append(stem) + for i in range(3): + downsample_layer = nn.Sequential( + LayerNorm(dims[i], eps=1e-6, data_format="channels_first"), + nn.Conv2d(dims[i], dims[i+1], kernel_size=2, stride=2), + ) + self.downsample_layers.append(downsample_layer) + + self.stages = nn.ModuleList() # 4 feature resolution stages, each consisting of multiple residual blocks + dp_rates=[x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] + cur = 0 + for i in range(4): + stage = nn.Sequential( + *[Block(dim=dims[i], drop_path=dp_rates[cur + j], + layer_scale_init_value=layer_scale_init_value) for j in range(depths[i])] + ) + self.stages.append(stage) + cur += depths[i] + + # add norm layers for each output + out_indices = (0, 1, 2, 3) + for i in out_indices: + layer = LayerNorm(dims[i], eps=1e-6, data_format="channels_first") + # layer = nn.Identity() + layer_name = f'norm{i}' + self.add_module(layer_name, layer) + + self.norm = nn.LayerNorm(dims[-1], eps=1e-6) # final norm layer + # self.head_cls = nn.Linear(dims[-1], 4) + + self.apply(self._init_weights) + # self.head_cls.weight.data.mul_(head_init_scale) + # self.head_cls.bias.data.mul_(head_init_scale) + + def _init_weights(self, m): + if isinstance(m, (nn.Conv2d, nn.Linear)): + trunc_normal_(m.weight, std=.02) + nn.init.constant_(m.bias, 0) + + def forward_features(self, x): + for i in range(4): + x = self.downsample_layers[i](x) + x = self.stages[i](x) + + # add extra norm + norm_layer = getattr(self, f'norm{i}') + # x = norm_layer(x) + norm_layer(x) + + return self.norm(x.mean([-2, -1])) # global average pooling, (N, C, H, W) -> (N, C) + + def forward(self, x): + x = self.forward_features(x) + # x = self.head_cls(x) + return x + +class LayerNorm(nn.Module): + r""" LayerNorm that supports two data formats: channels_last (default) or channels_first. + The ordering of the dimensions in the inputs. channels_last corresponds to inputs with + shape (batch_size, height, width, channels) while channels_first corresponds to inputs + with shape (batch_size, channels, height, width). + """ + def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"): + super().__init__() + self.weight = nn.Parameter(torch.ones(normalized_shape)) + self.bias = nn.Parameter(torch.zeros(normalized_shape)) + self.eps = eps + self.data_format = data_format + if self.data_format not in ["channels_last", "channels_first"]: + raise NotImplementedError + self.normalized_shape = (normalized_shape, ) + + def forward(self, x): + if self.data_format == "channels_last": # B H W C + return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps) + elif self.data_format == "channels_first": # B C H W + u = x.mean(1, keepdim=True) + s = (x - u).pow(2).mean(1, keepdim=True) + x = (x - u) / torch.sqrt(s + self.eps) + x = self.weight[:, None, None] * x + self.bias[:, None, None] + return x diff --git a/ddcolor/ddcolor_arch_utils/position_encoding.py b/ddcolor/ddcolor_arch_utils/position_encoding.py new file mode 100644 index 0000000..9fa7b9e --- /dev/null +++ b/ddcolor/ddcolor_arch_utils/position_encoding.py @@ -0,0 +1,52 @@ +# Copyright (c) Facebook, Inc. and its affiliates. +# Modified from: https://github.com/facebookresearch/detr/blob/master/models/position_encoding.py +""" +Various positional encodings for the transformer. +""" +import math + +import torch +from torch import nn + + +class PositionEmbeddingSine(nn.Module): + """ + This is a more standard version of the position embedding, very similar to the one + used by the Attention is all you need paper, generalized to work on images. + """ + + def __init__(self, num_pos_feats=64, temperature=10000, normalize=False, scale=None): + super().__init__() + self.num_pos_feats = num_pos_feats + self.temperature = temperature + self.normalize = normalize + if scale is not None and normalize is False: + raise ValueError("normalize should be True if scale is passed") + if scale is None: + scale = 2 * math.pi + self.scale = scale + + def forward(self, x, mask=None): + if mask is None: + mask = torch.zeros((x.size(0), x.size(2), x.size(3)), device=x.device, dtype=torch.bool) + not_mask = ~mask + y_embed = not_mask.cumsum(1, dtype=torch.float32) + x_embed = not_mask.cumsum(2, dtype=torch.float32) + if self.normalize: + eps = 1e-6 + y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale + x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale + + dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device) + dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats) + + pos_x = x_embed[:, :, :, None] / dim_t + pos_y = y_embed[:, :, :, None] / dim_t + pos_x = torch.stack( + (pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4 + ).flatten(3) + pos_y = torch.stack( + (pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4 + ).flatten(3) + pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2) + return pos \ No newline at end of file diff --git a/ddcolor/ddcolor_arch_utils/transformer.py b/ddcolor/ddcolor_arch_utils/transformer.py new file mode 100644 index 0000000..f462a10 --- /dev/null +++ b/ddcolor/ddcolor_arch_utils/transformer.py @@ -0,0 +1,368 @@ +# Copyright (c) Facebook, Inc. and its affiliates. +# Modified from: https://github.com/facebookresearch/detr/blob/master/models/transformer.py +""" +Transformer class. +Copy-paste from torch.nn.Transformer with modifications: + * positional encodings are passed in MHattention + * extra LN at the end of encoder is removed + * decoder returns a stack of activations from all decoding layers +""" +import copy +from typing import List, Optional + +import torch +import torch.nn.functional as F +from torch import Tensor, nn + + +class Transformer(nn.Module): + def __init__( + self, + d_model=512, + nhead=8, + num_encoder_layers=6, + num_decoder_layers=6, + dim_feedforward=2048, + dropout=0.1, + activation="relu", + normalize_before=False, + return_intermediate_dec=False, + ): + super().__init__() + + encoder_layer = TransformerEncoderLayer( + d_model, nhead, dim_feedforward, dropout, activation, normalize_before + ) + encoder_norm = nn.LayerNorm(d_model) if normalize_before else None + self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm) + + decoder_layer = TransformerDecoderLayer( + d_model, nhead, dim_feedforward, dropout, activation, normalize_before + ) + decoder_norm = nn.LayerNorm(d_model) + self.decoder = TransformerDecoder( + decoder_layer, + num_decoder_layers, + decoder_norm, + return_intermediate=return_intermediate_dec, + ) + + self._reset_parameters() + + self.d_model = d_model + self.nhead = nhead + + def _reset_parameters(self): + for p in self.parameters(): + if p.dim() > 1: + nn.init.xavier_uniform_(p) + + def forward(self, src, mask, query_embed, pos_embed): + # flatten NxCxHxW to HWxNxC + bs, c, h, w = src.shape + src = src.flatten(2).permute(2, 0, 1) + pos_embed = pos_embed.flatten(2).permute(2, 0, 1) + query_embed = query_embed.unsqueeze(1).repeat(1, bs, 1) + if mask is not None: + mask = mask.flatten(1) + + tgt = torch.zeros_like(query_embed) + memory = self.encoder(src, src_key_padding_mask=mask, pos=pos_embed) + hs = self.decoder( + tgt, memory, memory_key_padding_mask=mask, pos=pos_embed, query_pos=query_embed + ) + return hs.transpose(1, 2), memory.permute(1, 2, 0).view(bs, c, h, w) + + +class TransformerEncoder(nn.Module): + def __init__(self, encoder_layer, num_layers, norm=None): + super().__init__() + self.layers = _get_clones(encoder_layer, num_layers) + self.num_layers = num_layers + self.norm = norm + + def forward( + self, + src, + mask: Optional[Tensor] = None, + src_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None, + ): + output = src + + for layer in self.layers: + output = layer( + output, src_mask=mask, src_key_padding_mask=src_key_padding_mask, pos=pos + ) + + if self.norm is not None: + output = self.norm(output) + + return output + + +class TransformerDecoder(nn.Module): + def __init__(self, decoder_layer, num_layers, norm=None, return_intermediate=False): + super().__init__() + self.layers = _get_clones(decoder_layer, num_layers) + self.num_layers = num_layers + self.norm = norm + self.return_intermediate = return_intermediate + + def forward( + self, + tgt, + memory, + tgt_mask: Optional[Tensor] = None, + memory_mask: Optional[Tensor] = None, + tgt_key_padding_mask: Optional[Tensor] = None, + memory_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None, + query_pos: Optional[Tensor] = None, + ): + output = tgt + + intermediate = [] + + for layer in self.layers: + output = layer( + output, + memory, + tgt_mask=tgt_mask, + memory_mask=memory_mask, + tgt_key_padding_mask=tgt_key_padding_mask, + memory_key_padding_mask=memory_key_padding_mask, + pos=pos, + query_pos=query_pos, + ) + if self.return_intermediate: + intermediate.append(self.norm(output)) + + if self.norm is not None: + output = self.norm(output) + if self.return_intermediate: + intermediate.pop() + intermediate.append(output) + + if self.return_intermediate: + return torch.stack(intermediate) + + return output.unsqueeze(0) + + +class TransformerEncoderLayer(nn.Module): + def __init__( + self, + d_model, + nhead, + dim_feedforward=2048, + dropout=0.1, + activation="relu", + normalize_before=False, + ): + super().__init__() + self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) + # Implementation of Feedforward model + self.linear1 = nn.Linear(d_model, dim_feedforward) + self.dropout = nn.Dropout(dropout) + self.linear2 = nn.Linear(dim_feedforward, d_model) + + self.norm1 = nn.LayerNorm(d_model) + self.norm2 = nn.LayerNorm(d_model) + self.dropout1 = nn.Dropout(dropout) + self.dropout2 = nn.Dropout(dropout) + + self.activation = _get_activation_fn(activation) + self.normalize_before = normalize_before + + def with_pos_embed(self, tensor, pos: Optional[Tensor]): + return tensor if pos is None else tensor + pos + + def forward_post( + self, + src, + src_mask: Optional[Tensor] = None, + src_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None, + ): + q = k = self.with_pos_embed(src, pos) + src2 = self.self_attn( + q, k, value=src, attn_mask=src_mask, key_padding_mask=src_key_padding_mask + )[0] + src = src + self.dropout1(src2) + src = self.norm1(src) + src2 = self.linear2(self.dropout(self.activation(self.linear1(src)))) + src = src + self.dropout2(src2) + src = self.norm2(src) + return src + + def forward_pre( + self, + src, + src_mask: Optional[Tensor] = None, + src_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None, + ): + src2 = self.norm1(src) + q = k = self.with_pos_embed(src2, pos) + src2 = self.self_attn( + q, k, value=src2, attn_mask=src_mask, key_padding_mask=src_key_padding_mask + )[0] + src = src + self.dropout1(src2) + src2 = self.norm2(src) + src2 = self.linear2(self.dropout(self.activation(self.linear1(src2)))) + src = src + self.dropout2(src2) + return src + + def forward( + self, + src, + src_mask: Optional[Tensor] = None, + src_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None, + ): + if self.normalize_before: + return self.forward_pre(src, src_mask, src_key_padding_mask, pos) + return self.forward_post(src, src_mask, src_key_padding_mask, pos) + + +class TransformerDecoderLayer(nn.Module): + def __init__( + self, + d_model, + nhead, + dim_feedforward=2048, + dropout=0.1, + activation="relu", + normalize_before=False, + ): + super().__init__() + self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) + self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) + # Implementation of Feedforward model + self.linear1 = nn.Linear(d_model, dim_feedforward) + self.dropout = nn.Dropout(dropout) + self.linear2 = nn.Linear(dim_feedforward, d_model) + + self.norm1 = nn.LayerNorm(d_model) + self.norm2 = nn.LayerNorm(d_model) + self.norm3 = nn.LayerNorm(d_model) + self.dropout1 = nn.Dropout(dropout) + self.dropout2 = nn.Dropout(dropout) + self.dropout3 = nn.Dropout(dropout) + + self.activation = _get_activation_fn(activation) + self.normalize_before = normalize_before + + def with_pos_embed(self, tensor, pos: Optional[Tensor]): + return tensor if pos is None else tensor + pos + + def forward_post( + self, + tgt, + memory, + tgt_mask: Optional[Tensor] = None, + memory_mask: Optional[Tensor] = None, + tgt_key_padding_mask: Optional[Tensor] = None, + memory_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None, + query_pos: Optional[Tensor] = None, + ): + q = k = self.with_pos_embed(tgt, query_pos) + tgt2 = self.self_attn( + q, k, value=tgt, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask + )[0] + tgt = tgt + self.dropout1(tgt2) + tgt = self.norm1(tgt) + tgt2 = self.multihead_attn( + query=self.with_pos_embed(tgt, query_pos), + key=self.with_pos_embed(memory, pos), + value=memory, + attn_mask=memory_mask, + key_padding_mask=memory_key_padding_mask, + )[0] + tgt = tgt + self.dropout2(tgt2) + tgt = self.norm2(tgt) + tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt)))) + tgt = tgt + self.dropout3(tgt2) + tgt = self.norm3(tgt) + return tgt + + def forward_pre( + self, + tgt, + memory, + tgt_mask: Optional[Tensor] = None, + memory_mask: Optional[Tensor] = None, + tgt_key_padding_mask: Optional[Tensor] = None, + memory_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None, + query_pos: Optional[Tensor] = None, + ): + tgt2 = self.norm1(tgt) + q = k = self.with_pos_embed(tgt2, query_pos) + tgt2 = self.self_attn( + q, k, value=tgt2, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask + )[0] + tgt = tgt + self.dropout1(tgt2) + tgt2 = self.norm2(tgt) + tgt2 = self.multihead_attn( + query=self.with_pos_embed(tgt2, query_pos), + key=self.with_pos_embed(memory, pos), + value=memory, + attn_mask=memory_mask, + key_padding_mask=memory_key_padding_mask, + )[0] + tgt = tgt + self.dropout2(tgt2) + tgt2 = self.norm3(tgt) + tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2)))) + tgt = tgt + self.dropout3(tgt2) + return tgt + + def forward( + self, + tgt, + memory, + tgt_mask: Optional[Tensor] = None, + memory_mask: Optional[Tensor] = None, + tgt_key_padding_mask: Optional[Tensor] = None, + memory_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None, + query_pos: Optional[Tensor] = None, + ): + if self.normalize_before: + return self.forward_pre( + tgt, + memory, + tgt_mask, + memory_mask, + tgt_key_padding_mask, + memory_key_padding_mask, + pos, + query_pos, + ) + return self.forward_post( + tgt, + memory, + tgt_mask, + memory_mask, + tgt_key_padding_mask, + memory_key_padding_mask, + pos, + query_pos, + ) + + +def _get_clones(module, N): + return nn.ModuleList([copy.deepcopy(module) for i in range(N)]) + + +def _get_activation_fn(activation): + """Return an activation function given a string""" + if activation == "relu": + return F.relu + if activation == "gelu": + return F.gelu + if activation == "glu": + return F.glu + raise RuntimeError(f"activation should be relu/gelu, not {activation}.") \ No newline at end of file diff --git a/ddcolor/ddcolor_arch_utils/transformer_utils.py b/ddcolor/ddcolor_arch_utils/transformer_utils.py new file mode 100644 index 0000000..ea38f39 --- /dev/null +++ b/ddcolor/ddcolor_arch_utils/transformer_utils.py @@ -0,0 +1,192 @@ +from typing import Optional +from torch import nn, Tensor +from torch.nn import functional as F + +class SelfAttentionLayer(nn.Module): + + def __init__(self, d_model, nhead, dropout=0.0, + activation="relu", normalize_before=False): + super().__init__() + self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) + + self.norm = nn.LayerNorm(d_model) + self.dropout = nn.Dropout(dropout) + + self.activation = _get_activation_fn(activation) + self.normalize_before = normalize_before + + self._reset_parameters() + + def _reset_parameters(self): + for p in self.parameters(): + if p.dim() > 1: + nn.init.xavier_uniform_(p) + + def with_pos_embed(self, tensor, pos: Optional[Tensor]): + return tensor if pos is None else tensor + pos + + def forward_post(self, tgt, + tgt_mask: Optional[Tensor] = None, + tgt_key_padding_mask: Optional[Tensor] = None, + query_pos: Optional[Tensor] = None): + q = k = self.with_pos_embed(tgt, query_pos) + tgt2 = self.self_attn(q, k, value=tgt, attn_mask=tgt_mask, + key_padding_mask=tgt_key_padding_mask)[0] + tgt = tgt + self.dropout(tgt2) + tgt = self.norm(tgt) + + return tgt + + def forward_pre(self, tgt, + tgt_mask: Optional[Tensor] = None, + tgt_key_padding_mask: Optional[Tensor] = None, + query_pos: Optional[Tensor] = None): + tgt2 = self.norm(tgt) + q = k = self.with_pos_embed(tgt2, query_pos) + tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask, + key_padding_mask=tgt_key_padding_mask)[0] + tgt = tgt + self.dropout(tgt2) + + return tgt + + def forward(self, tgt, + tgt_mask: Optional[Tensor] = None, + tgt_key_padding_mask: Optional[Tensor] = None, + query_pos: Optional[Tensor] = None): + if self.normalize_before: + return self.forward_pre(tgt, tgt_mask, + tgt_key_padding_mask, query_pos) + return self.forward_post(tgt, tgt_mask, + tgt_key_padding_mask, query_pos) + + +class CrossAttentionLayer(nn.Module): + + def __init__(self, d_model, nhead, dropout=0.0, + activation="relu", normalize_before=False): + super().__init__() + self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) + + self.norm = nn.LayerNorm(d_model) + self.dropout = nn.Dropout(dropout) + + self.activation = _get_activation_fn(activation) + self.normalize_before = normalize_before + + self._reset_parameters() + + def _reset_parameters(self): + for p in self.parameters(): + if p.dim() > 1: + nn.init.xavier_uniform_(p) + + def with_pos_embed(self, tensor, pos: Optional[Tensor]): + return tensor if pos is None else tensor + pos + + def forward_post(self, tgt, memory, + memory_mask: Optional[Tensor] = None, + memory_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None, + query_pos: Optional[Tensor] = None): + tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt, query_pos), + key=self.with_pos_embed(memory, pos), + value=memory, attn_mask=memory_mask, + key_padding_mask=memory_key_padding_mask)[0] + tgt = tgt + self.dropout(tgt2) + tgt = self.norm(tgt) + + return tgt + + def forward_pre(self, tgt, memory, + memory_mask: Optional[Tensor] = None, + memory_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None, + query_pos: Optional[Tensor] = None): + tgt2 = self.norm(tgt) + tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt2, query_pos), + key=self.with_pos_embed(memory, pos), + value=memory, attn_mask=memory_mask, + key_padding_mask=memory_key_padding_mask)[0] + tgt = tgt + self.dropout(tgt2) + + return tgt + + def forward(self, tgt, memory, + memory_mask: Optional[Tensor] = None, + memory_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None, + query_pos: Optional[Tensor] = None): + if self.normalize_before: + return self.forward_pre(tgt, memory, memory_mask, + memory_key_padding_mask, pos, query_pos) + return self.forward_post(tgt, memory, memory_mask, + memory_key_padding_mask, pos, query_pos) + + +class FFNLayer(nn.Module): + + def __init__(self, d_model, dim_feedforward=2048, dropout=0.0, + activation="relu", normalize_before=False): + super().__init__() + # Implementation of Feedforward model + self.linear1 = nn.Linear(d_model, dim_feedforward) + self.dropout = nn.Dropout(dropout) + self.linear2 = nn.Linear(dim_feedforward, d_model) + + self.norm = nn.LayerNorm(d_model) + + self.activation = _get_activation_fn(activation) + self.normalize_before = normalize_before + + self._reset_parameters() + + def _reset_parameters(self): + for p in self.parameters(): + if p.dim() > 1: + nn.init.xavier_uniform_(p) + + def with_pos_embed(self, tensor, pos: Optional[Tensor]): + return tensor if pos is None else tensor + pos + + def forward_post(self, tgt): + tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt)))) + tgt = tgt + self.dropout(tgt2) + tgt = self.norm(tgt) + return tgt + + def forward_pre(self, tgt): + tgt2 = self.norm(tgt) + tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2)))) + tgt = tgt + self.dropout(tgt2) + return tgt + + def forward(self, tgt): + if self.normalize_before: + return self.forward_pre(tgt) + return self.forward_post(tgt) + + +def _get_activation_fn(activation): + """Return an activation function given a string""" + if activation == "relu": + return F.relu + if activation == "gelu": + return F.gelu + if activation == "glu": + return F.glu + raise RuntimeError(F"activation should be relu/gelu, not {activation}.") + + +class MLP(nn.Module): + """ Very simple multi-layer perceptron (also called FFN)""" + + def __init__(self, input_dim, hidden_dim, output_dim, num_layers): + super().__init__() + self.num_layers = num_layers + h = [hidden_dim] * (num_layers - 1) + self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim])) + + def forward(self, x): + for i, layer in enumerate(self.layers): + x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x) + return x \ No newline at end of file diff --git a/ddcolor/ddcolor_arch_utils/unet.py b/ddcolor/ddcolor_arch_utils/unet.py new file mode 100644 index 0000000..4610ba5 --- /dev/null +++ b/ddcolor/ddcolor_arch_utils/unet.py @@ -0,0 +1,208 @@ +from enum import Enum +import torch +import torch.nn as nn +from torch.nn import functional as F +import collections + + +NormType = Enum('NormType', 'Batch BatchZero Weight Spectral') + + +class Hook: + feature = None + + def __init__(self, module): + self.hook = module.register_forward_hook(self.hook_fn) + + def hook_fn(self, module, input, output): + if isinstance(output, torch.Tensor): + self.feature = output + elif isinstance(output, collections.OrderedDict): + self.feature = output['out'] + + def remove(self): + self.hook.remove() + + +class SelfAttention(nn.Module): + "Self attention layer for nd." + + def __init__(self, n_channels: int): + super().__init__() + self.query = conv1d(n_channels, n_channels // 8) + self.key = conv1d(n_channels, n_channels // 8) + self.value = conv1d(n_channels, n_channels) + self.gamma = nn.Parameter(torch.tensor([0.])) + + def forward(self, x): + #Notation from https://arxiv.org/pdf/1805.08318.pdf + size = x.size() + x = x.view(*size[:2], -1) + f, g, h = self.query(x), self.key(x), self.value(x) + beta = F.softmax(torch.bmm(f.permute(0, 2, 1).contiguous(), g), dim=1) + o = self.gamma * torch.bmm(h, beta) + x + return o.view(*size).contiguous() + + +def batchnorm_2d(nf: int, norm_type: NormType = NormType.Batch): + "A batchnorm2d layer with `nf` features initialized depending on `norm_type`." + bn = nn.BatchNorm2d(nf) + with torch.no_grad(): + bn.bias.fill_(1e-3) + bn.weight.fill_(0. if norm_type == NormType.BatchZero else 1.) + return bn + + +def init_default(m: nn.Module, func=nn.init.kaiming_normal_) -> None: + "Initialize `m` weights with `func` and set `bias` to 0." + if func: + if hasattr(m, 'weight'): func(m.weight) + if hasattr(m, 'bias') and hasattr(m.bias, 'data'): m.bias.data.fill_(0.) + return m + + +def icnr(x, scale=2, init=nn.init.kaiming_normal_): + "ICNR init of `x`, with `scale` and `init` function." + ni, nf, h, w = x.shape + ni2 = int(ni / (scale**2)) + k = init(torch.zeros([ni2, nf, h, w])).transpose(0, 1) + k = k.contiguous().view(ni2, nf, -1) + k = k.repeat(1, 1, scale**2) + k = k.contiguous().view([nf, ni, h, w]).transpose(0, 1) + x.data.copy_(k) + + +def conv1d(ni: int, no: int, ks: int = 1, stride: int = 1, padding: int = 0, bias: bool = False): + "Create and initialize a `nn.Conv1d` layer with spectral normalization." + conv = nn.Conv1d(ni, no, ks, stride=stride, padding=padding, bias=bias) + nn.init.kaiming_normal_(conv.weight) + if bias: conv.bias.data.zero_() + return nn.utils.spectral_norm(conv) + + +def custom_conv_layer( + ni: int, + nf: int, + ks: int = 3, + stride: int = 1, + padding: int = None, + bias: bool = None, + is_1d: bool = False, + norm_type=NormType.Batch, + use_activ: bool = True, + transpose: bool = False, + init=nn.init.kaiming_normal_, + self_attention: bool = False, + extra_bn: bool = False, +): + "Create a sequence of convolutional (`ni` to `nf`), ReLU (if `use_activ`) and batchnorm (if `bn`) layers." + if padding is None: + padding = (ks - 1) // 2 if not transpose else 0 + bn = norm_type in (NormType.Batch, NormType.BatchZero) or extra_bn == True + if bias is None: + bias = not bn + conv_func = nn.ConvTranspose2d if transpose else nn.Conv1d if is_1d else nn.Conv2d + conv = init_default( + conv_func(ni, nf, kernel_size=ks, bias=bias, stride=stride, padding=padding), + init, + ) + + if norm_type == NormType.Weight: + conv = nn.utils.weight_norm(conv) + elif norm_type == NormType.Spectral: + conv = nn.utils.spectral_norm(conv) + layers = [conv] + if use_activ: + layers.append(nn.ReLU(True)) + if bn: + layers.append((nn.BatchNorm1d if is_1d else nn.BatchNorm2d)(nf)) + if self_attention: + layers.append(SelfAttention(nf)) + return nn.Sequential(*layers) + + +def conv_layer(ni: int, + nf: int, + ks: int = 3, + stride: int = 1, + padding: int = None, + bias: bool = None, + is_1d: bool = False, + norm_type=NormType.Batch, + use_activ: bool = True, + transpose: bool = False, + init=nn.init.kaiming_normal_, + self_attention: bool = False): + "Create a sequence of convolutional (`ni` to `nf`), ReLU (if `use_activ`) and batchnorm (if `bn`) layers." + if padding is None: padding = (ks - 1) // 2 if not transpose else 0 + bn = norm_type in (NormType.Batch, NormType.BatchZero) + if bias is None: bias = not bn + conv_func = nn.ConvTranspose2d if transpose else nn.Conv1d if is_1d else nn.Conv2d + conv = init_default(conv_func(ni, nf, kernel_size=ks, bias=bias, stride=stride, padding=padding), init) + if norm_type == NormType.Weight: conv = nn.utils.weight_norm(conv) + elif norm_type == NormType.Spectral: conv = nn.utils.spectral_norm(conv) + layers = [conv] + if use_activ: layers.append(nn.ReLU(True)) + if bn: layers.append((nn.BatchNorm1d if is_1d else nn.BatchNorm2d)(nf)) + if self_attention: layers.append(SelfAttention(nf)) + return nn.Sequential(*layers) + + +def _conv(ni: int, nf: int, ks: int = 3, stride: int = 1, **kwargs): + return conv_layer(ni, nf, ks=ks, stride=stride, norm_type=NormType.Spectral, **kwargs) + + +class CustomPixelShuffle_ICNR(nn.Module): + "Upsample by `scale` from `ni` filters to `nf` (default `ni`), using `nn.PixelShuffle`, `icnr` init, and `weight_norm`." + + def __init__(self, + ni: int, + nf: int = None, + scale: int = 2, + blur: bool = True, + norm_type=NormType.Spectral, + extra_bn=False): + super().__init__() + self.conv = custom_conv_layer( + ni, nf * (scale**2), ks=1, use_activ=False, norm_type=norm_type, extra_bn=extra_bn) + icnr(self.conv[0].weight) + self.shuf = nn.PixelShuffle(scale) + self.do_blur = blur + # Blurring over (h*w) kernel + # "Super-Resolution using Convolutional Neural Networks without Any Checkerboard Artifacts" + # - https://arxiv.org/abs/1806.02658 + self.pad = nn.ReplicationPad2d((1, 0, 1, 0)) + self.blur = nn.AvgPool2d(2, stride=1) + self.relu = nn.ReLU(True) + + def forward(self, x): + x = self.shuf(self.relu(self.conv(x))) + return self.blur(self.pad(x)) if self.do_blur else x + + +class UnetBlockWide(nn.Module): + "A quasi-UNet block, using `PixelShuffle_ICNR upsampling`." + + def __init__(self, + up_in_c: int, + x_in_c: int, + n_out: int, + hook, + blur: bool = False, + self_attention: bool = False, + norm_type=NormType.Spectral): + super().__init__() + + self.hook = hook + up_out = n_out + self.shuf = CustomPixelShuffle_ICNR(up_in_c, up_out, blur=blur, norm_type=norm_type, extra_bn=True) + self.bn = batchnorm_2d(x_in_c) + ni = up_out + x_in_c + self.conv = custom_conv_layer(ni, n_out, norm_type=norm_type, self_attention=self_attention, extra_bn=True) + self.relu = nn.ReLU() + + def forward(self, up_in): + s = self.hook.feature + up_out = self.shuf(up_in) + cat_x = self.relu(torch.cat([up_out, self.bn(s)], dim=1)) + return self.conv(cat_x) diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..8b8a743 --- /dev/null +++ b/nodes.py @@ -0,0 +1,107 @@ +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": 0, "max": 0xffffffffffffffff, "step": 1}), + "checkpoint": ( + [ + "ddcolor_paper_tiny.pth", + "ddcolor_paper.pth", + "ddcolor_modelscope.pth", + "ddcolor_artistic.pth", + ], { + "default": "ddcolor_paper.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", +} \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..25c7e86 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,5 @@ +numpy>=1.24.3 +opencv_python>=4.7.0.72 +Pillow>=10.1.0 +timm>=0.9.2 +huggingface_hub