diff --git a/.gitignore b/.gitignore index 82f9275..5ec8bfe 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,6 @@ +*.pth + + # Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] diff --git a/README.md b/README.md index 486198a..3a91792 100644 --- a/README.md +++ b/README.md @@ -3,16 +3,16 @@ Advanced Non-Diffusion-based Style Transfer in ComfyUI Click name to jump to workflow -1. [Neural Neighbor](#neural-neighbor). Paper: [Neural Neighbor Style Transfer](https://github.com/nkolkin13/NeuralNeighborStyleTransfer) - +1. [**Neural Neighbor**](#neural-neighbor). Paper: [Neural Neighbor Style Transfer](https://github.com/nkolkin13/NeuralNeighborStyleTransfer) +2. [**CAST**](#cast). Paper: [A Unified Arbitrary Style Transfer Framework via Adaptive Contrastive Learning](https://github.com/zyxElsa/CAST_pytorch) ## Workflows All nodes support batched input (i.e video) but is generally not recommended. One should generate 1 or 2 style frames (start and end), then use [ComfyUI-EbSynth](https://github.com/FuouM/ComfyUI-EbSynth) to propagate the style to the entire video. -| Neural Neighbor | X | X | -|:-:|-|-| -|![neural_neighbor](example_outputs/neural_neighbor.png) | X | X | X +| Neural Neighbor (2020-2022) - Slowest | CAST (2022) - Fast | X | +|:-:|:-:|-| +|![neural_neighbor](example_outputs/neural_neighbor.png) | ![cast](example_outputs/cast.png) | X | ### Neural Neighbor @@ -32,3 +32,46 @@ See more information on the original repository. No model is needed. ![wf_neural_neighbor](workflows/wf_neural_neighbor.png) +### CAST + +Download the CAST and UCAST models in the [Test](https://github.com/zyxElsa/CAST_pytorch?tab=readme-ov-file#test) section (Google Drive) (Unzip). + +Download the `vgg_normalized` model in the [Train](https://github.com/zyxElsa/CAST_pytorch?tab=readme-ov-file#train) section + +``` +models\ +│ .gitkeep +│ vgg_normalised.pth +─CAST_model +│ latest_net_AE.pth +│ latest_net_Dec_A.pth # Safe to delete +│ latest_net_Dec_B.pth +│ test_opt.txt # Safe to delete +│ +└─UCAST_model + latest_net_AE.pth + latest_net_Dec_A.pth # Safe to delete + latest_net_Dec_B.pth +``` + +[workflow_cast.json](workflows/workflow_cast.json) + +![wf_cast](workflows/wf_cast.png) + +## Credits + +``` +@inproceedings{zhang2020cast, +author = {Zhang, Yuxin and Tang, Fan and Dong, Weiming and Huang, Haibin and Ma, Chongyang and Lee, Tong-Yee and Xu, Changsheng}, +title = {Domain Enhanced Arbitrary Image Style Transfer via Contrastive Learning}, +booktitle = {ACM SIGGRAPH}, +year = {2022}} + +@article{zhang2023unified, + title={A Unified Arbitrary Style Transfer Framework via Adaptive Contrastive Learning}, + author={Zhang, Yuxin and Tang, Fan and Dong, Weiming and Huang, Haibin and Ma, Chongyang and Lee, Tong-Yee and Xu, Changsheng}, + journal={ACM Transactions on Graphics}, + year={2023}, + publisher={ACM New York, NY} +} +``` \ No newline at end of file diff --git a/__init__.py b/__init__.py index 480f43d..b6e46c9 100644 --- a/__init__.py +++ b/__init__.py @@ -1,11 +1,13 @@ -from .run import NeuralNeighbor +from .run import NeuralNeighbor, CAST NODE_CLASS_MAPPINGS = { "NeuralNeighbor": NeuralNeighbor, + "CAST": CAST, } NODE_DISPLAY_NAME_MAPPINGS = { "NeuralNeighbor": "Neural Neighbor", + "CAST": "CAST", } diff --git a/constants.py b/constants.py index e69de29..9d27d95 100644 --- a/constants.py +++ b/constants.py @@ -0,0 +1,6 @@ +CAST_TYPES = ["CAST", "UCAST"] +CAST_DEFAULT = "CAST" + +CAST_NET_AE_PATH = "latest_net_AE.pth" +CAST_NET_DEC_B_PATH = "latest_net_Dec_B.pth" +CAST_VGG_PATH = "models/vgg_normalised.pth" \ No newline at end of file diff --git a/example_outputs/cast.png b/example_outputs/cast.png new file mode 100644 index 0000000..979ded7 Binary files /dev/null and b/example_outputs/cast.png differ diff --git a/models/.gitkeep b/models/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/module_cast/cast_model.py b/module_cast/cast_model.py new file mode 100644 index 0000000..3ff3f96 --- /dev/null +++ b/module_cast/cast_model.py @@ -0,0 +1,41 @@ +import torch + +from . import net + + +def load_a_ckpt(ckpt_path: str): + state_dict = torch.load(ckpt_path) + if hasattr(state_dict, "_metadata"): + del state_dict._metadata + return state_dict + + +def inference_ucast( + src_img: torch.Tensor, + style_img: torch.Tensor, + ckpt_ae: str, + ckpt_decB: str, + vgg_path, + device, +): + vgg = net.vgg + vgg.load_state_dict(torch.load(vgg_path)) + vgg = torch.nn.Sequential(*list(vgg.children())[:31]) + + netAE = net.ADAIN_Encoder(vgg) + netDec_B = net.Decoder() + + netAE.load_state_dict(load_a_ckpt(ckpt_ae)) + netDec_B.load_state_dict(load_a_ckpt(ckpt_decB)) + + netAE = netAE.to(device).eval() + netDec_B = netDec_B.to(device).eval() + + real_A = src_img.permute(0, 3, 1, 2).to(device) + real_B = style_img.permute(0, 3, 1, 2).to(device) + + with torch.no_grad(): + real_A_feat = netAE.forward(real_A, real_B) + fake_B = netDec_B.forward(real_A_feat) + + return fake_B diff --git a/module_cast/net.py b/module_cast/net.py new file mode 100644 index 0000000..09cd30f --- /dev/null +++ b/module_cast/net.py @@ -0,0 +1,154 @@ +import torch.nn as nn + +vgg = nn.Sequential( + nn.Conv2d(3, 3, (1, 1)), + nn.ReflectionPad2d((1, 1, 1, 1)), + nn.Conv2d(3, 64, (3, 3)), + nn.ReLU(), # relu1-1 + nn.ReflectionPad2d((1, 1, 1, 1)), + nn.Conv2d(64, 64, (3, 3)), + nn.ReLU(), # relu1-2 + nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True), + nn.ReflectionPad2d((1, 1, 1, 1)), + nn.Conv2d(64, 128, (3, 3)), + nn.ReLU(), # relu2-1 + nn.ReflectionPad2d((1, 1, 1, 1)), + nn.Conv2d(128, 128, (3, 3)), + nn.ReLU(), # relu2-2 + nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True), + nn.ReflectionPad2d((1, 1, 1, 1)), + nn.Conv2d(128, 256, (3, 3)), + nn.ReLU(), # relu3-1 + nn.ReflectionPad2d((1, 1, 1, 1)), + nn.Conv2d(256, 256, (3, 3)), + nn.ReLU(), # relu3-2 + nn.ReflectionPad2d((1, 1, 1, 1)), + nn.Conv2d(256, 256, (3, 3)), + nn.ReLU(), # relu3-3 + nn.ReflectionPad2d((1, 1, 1, 1)), + nn.Conv2d(256, 256, (3, 3)), + nn.ReLU(), # relu3-4 + nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True), + nn.ReflectionPad2d((1, 1, 1, 1)), + nn.Conv2d(256, 512, (3, 3)), + nn.ReLU(), # relu4-1, this is the last layer used + nn.ReflectionPad2d((1, 1, 1, 1)), + nn.Conv2d(512, 512, (3, 3)), + nn.ReLU(), # relu4-2 + nn.ReflectionPad2d((1, 1, 1, 1)), + nn.Conv2d(512, 512, (3, 3)), + nn.ReLU(), # relu4-3 + nn.ReflectionPad2d((1, 1, 1, 1)), + nn.Conv2d(512, 512, (3, 3)), + nn.ReLU(), # relu4-4 + nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True), + nn.ReflectionPad2d((1, 1, 1, 1)), + nn.Conv2d(512, 512, (3, 3)), + nn.ReLU(), # relu5-1 + nn.ReflectionPad2d((1, 1, 1, 1)), + nn.Conv2d(512, 512, (3, 3)), + nn.ReLU(), # relu5-2 + nn.ReflectionPad2d((1, 1, 1, 1)), + nn.Conv2d(512, 512, (3, 3)), + nn.ReLU(), # relu5-3 + nn.ReflectionPad2d((1, 1, 1, 1)), + nn.Conv2d(512, 512, (3, 3)), + nn.ReLU(), # relu5-4 +) + + +class ADAIN_Encoder(nn.Module): + def __init__(self, encoder): + super(ADAIN_Encoder, self).__init__() + enc_layers = list(encoder.children()) + self.enc_1 = nn.Sequential(*enc_layers[:4]) # input -> relu1_1 64 + self.enc_2 = nn.Sequential(*enc_layers[4:11]) # relu1_1 -> relu2_1 128 + self.enc_3 = nn.Sequential(*enc_layers[11:18]) # relu2_1 -> relu3_1 256 + self.enc_4 = nn.Sequential(*enc_layers[18:31]) # relu3_1 -> relu4_1 512 + + self.mse_loss = nn.MSELoss() + + # fix the encoder + for name in ["enc_1", "enc_2", "enc_3", "enc_4"]: + for param in getattr(self, name).parameters(): + param.requires_grad = False + + # extract relu1_1, relu2_1, relu3_1, relu4_1 from input image + def encode_with_intermediate(self, input): + results = [input] + for i in range(4): + func = getattr(self, "enc_{:d}".format(i + 1)) + results.append(func(results[-1])) + return results[1:] + + def calc_mean_std(self, feat, eps=1e-5): + # eps is a small value added to the variance to avoid divide-by-zero. + size = feat.size() + assert len(size) == 4 + N, C = size[:2] # [N, C, H, W] + feat_var = feat.view(N, C, -1).var(dim=2) + eps + feat_std = feat_var.sqrt().view(N, C, 1, 1) + feat_mean = feat.view(N, C, -1).mean(dim=2).view(N, C, 1, 1) + return feat_mean, feat_std + + def adain(self, content_feat, style_feat): + assert content_feat.size()[:2] == style_feat.size()[:2] + size = content_feat.size() + style_mean, style_std = self.calc_mean_std(style_feat) + content_mean, content_std = self.calc_mean_std(content_feat) + + normalized_feat = ( + content_feat - content_mean.expand(size) + ) / content_std.expand(size) + return normalized_feat * style_std.expand(size) + style_mean.expand(size) + + def forward(self, content, style, encoded_only=False): + style_feats = self.encode_with_intermediate(style) + content_feats = self.encode_with_intermediate(content) + if encoded_only: + return content_feats[-1], style_feats[-1] + else: + adain_feat = self.adain(content_feats[-1], style_feats[-1]) + return adain_feat + + +class Decoder(nn.Module): + def __init__(self): + super(Decoder, self).__init__() + decoder = [ + nn.ReflectionPad2d((1, 1, 1, 1)), + nn.Conv2d(512, 256, (3, 3)), + nn.ReLU(), # 256 + nn.Upsample(scale_factor=2, mode="nearest"), + nn.ReflectionPad2d((1, 1, 1, 1)), + nn.Conv2d(256, 256, (3, 3)), + nn.ReLU(), + nn.ReflectionPad2d((1, 1, 1, 1)), + nn.Conv2d(256, 256, (3, 3)), + nn.ReLU(), + nn.ReflectionPad2d((1, 1, 1, 1)), + nn.Conv2d(256, 256, (3, 3)), + nn.ReLU(), + nn.ReflectionPad2d((1, 1, 1, 1)), + nn.Conv2d(256, 128, (3, 3)), + nn.ReLU(), # 128 + nn.Upsample(scale_factor=2, mode="nearest"), + nn.ReflectionPad2d((1, 1, 1, 1)), + nn.Conv2d(128, 128, (3, 3)), + nn.ReLU(), + nn.ReflectionPad2d((1, 1, 1, 1)), + nn.Conv2d(128, 64, (3, 3)), + nn.ReLU(), # 64 + nn.Upsample(scale_factor=2, mode="nearest"), + nn.ReflectionPad2d((1, 1, 1, 1)), + nn.Conv2d(64, 64, (3, 3)), + nn.ReLU(), + nn.ReflectionPad2d((1, 1, 1, 1)), + nn.Conv2d(64, 3, (3, 3)), + ] + self.decoder = nn.Sequential(*decoder) + + def forward(self, adain_feat): + fake_image = self.decoder(adain_feat) + + return fake_image diff --git a/module_cast/networks.py b/module_cast/networks.py new file mode 100644 index 0000000..5b0ab80 --- /dev/null +++ b/module_cast/networks.py @@ -0,0 +1,854 @@ +import functools + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.nn import init +from torch.optim import lr_scheduler + +############################################################################### +# Helper Functions +############################################################################### + + +def get_filter(filt_size=3): + if filt_size == 1: + a = np.array( + [ + 1.0, + ] + ) + elif filt_size == 2: + a = np.array([1.0, 1.0]) + elif filt_size == 3: + a = np.array([1.0, 2.0, 1.0]) + elif filt_size == 4: + a = np.array([1.0, 3.0, 3.0, 1.0]) + elif filt_size == 5: + a = np.array([1.0, 4.0, 6.0, 4.0, 1.0]) + elif filt_size == 6: + a = np.array([1.0, 5.0, 10.0, 10.0, 5.0, 1.0]) + elif filt_size == 7: + a = np.array([1.0, 6.0, 15.0, 20.0, 15.0, 6.0, 1.0]) + + filt = torch.Tensor(a[:, None] * a[None, :]) + filt = filt / torch.sum(filt) + + return filt + + +class Downsample(nn.Module): + def __init__(self, channels, pad_type="reflect", filt_size=3, stride=2, pad_off=0): + super(Downsample, self).__init__() + self.filt_size = filt_size + self.pad_off = pad_off + self.pad_sizes = [ + int(1.0 * (filt_size - 1) / 2), + int(np.ceil(1.0 * (filt_size - 1) / 2)), + int(1.0 * (filt_size - 1) / 2), + int(np.ceil(1.0 * (filt_size - 1) / 2)), + ] + self.pad_sizes = [pad_size + pad_off for pad_size in self.pad_sizes] + self.stride = stride + self.off = int((self.stride - 1) / 2.0) + self.channels = channels + + filt = get_filter(filt_size=self.filt_size) + self.register_buffer( + "filt", filt[None, None, :, :].repeat((self.channels, 1, 1, 1)) + ) + + self.pad = get_pad_layer(pad_type)(self.pad_sizes) + + def forward(self, inp): + if self.filt_size == 1: + if self.pad_off == 0: + return inp[:, :, :: self.stride, :: self.stride] + else: + return self.pad(inp)[:, :, :: self.stride, :: self.stride] + else: + return F.conv2d( + self.pad(inp), self.filt, stride=self.stride, groups=inp.shape[1] + ) + + +class Upsample2(nn.Module): + def __init__(self, scale_factor, mode="nearest"): + super().__init__() + self.factor = scale_factor + self.mode = mode + + def forward(self, x): + return torch.nn.functional.interpolate( + x, scale_factor=self.factor, mode=self.mode + ) + + +class Upsample(nn.Module): + def __init__(self, channels, pad_type="repl", filt_size=4, stride=2): + super(Upsample, self).__init__() + self.filt_size = filt_size + self.filt_odd = np.mod(filt_size, 2) == 1 + self.pad_size = int((filt_size - 1) / 2) + self.stride = stride + self.off = int((self.stride - 1) / 2.0) + self.channels = channels + + filt = get_filter(filt_size=self.filt_size) * (stride**2) + self.register_buffer( + "filt", filt[None, None, :, :].repeat((self.channels, 1, 1, 1)) + ) + + self.pad = get_pad_layer(pad_type)([1, 1, 1, 1]) + + def forward(self, inp): + ret_val = F.conv_transpose2d( + self.pad(inp), + self.filt, + stride=self.stride, + padding=1 + self.pad_size, + groups=inp.shape[1], + )[:, :, 1:, 1:] + if self.filt_odd: + return ret_val + else: + return ret_val[:, :, :-1, :-1] + + +def get_pad_layer(pad_type): + if pad_type in ["refl", "reflect"]: + PadLayer = nn.ReflectionPad2d + elif pad_type in ["repl", "replicate"]: + PadLayer = nn.ReplicationPad2d + elif pad_type == "zero": + PadLayer = nn.ZeroPad2d + else: + print("Pad type [%s] not recognized" % pad_type) + return PadLayer + + +class Identity(nn.Module): + def forward(self, x): + return x + + +def get_norm_layer(norm_type="instance"): + """Return a normalization layer + + Parameters: + norm_type (str) -- the name of the normalization layer: batch | instance | none + + For BatchNorm, we use learnable affine parameters and track running statistics (mean/stddev). + For InstanceNorm, we do not use learnable affine parameters. We do not track running statistics. + """ + if norm_type == "batch": + norm_layer = functools.partial( + nn.BatchNorm2d, affine=True, track_running_stats=True + ) + elif norm_type == "instance": + norm_layer = functools.partial( + nn.InstanceNorm2d, affine=False, track_running_stats=False + ) + elif norm_type == "none": + + def norm_layer(x): + return Identity() + else: + raise NotImplementedError("normalization layer [%s] is not found" % norm_type) + return norm_layer + + +def get_scheduler(optimizer, opt): + """Return a learning rate scheduler + + Parameters: + optimizer -- the optimizer of the network + opt (option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions． + opt.lr_policy is the name of learning rate policy: linear | step | plateau | cosine + + For 'linear', we keep the same learning rate for the first epochs + and linearly decay the rate to zero over the next epochs. + For other schedulers (step, plateau, and cosine), we use the default PyTorch schedulers. + See https://pytorch.org/docs/stable/optim.html for more details. + """ + if opt.lr_policy == "linear": + + def lambda_rule(epoch): + lr_l = 1.0 - max(0, epoch + opt.epoch_count - opt.n_epochs) / float( + opt.n_epochs_decay + 1 + ) + return lr_l + + scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule) + elif opt.lr_policy == "step": + scheduler = lr_scheduler.StepLR( + optimizer, step_size=opt.lr_decay_iters, gamma=0.1 + ) + elif opt.lr_policy == "plateau": + scheduler = lr_scheduler.ReduceLROnPlateau( + optimizer, mode="min", factor=0.2, threshold=0.01, patience=5 + ) + elif opt.lr_policy == "cosine": + scheduler = lr_scheduler.CosineAnnealingLR( + optimizer, T_max=opt.n_epochs, eta_min=0 + ) + else: + return NotImplementedError( + "learning rate policy [%s] is not implemented", opt.lr_policy + ) + return scheduler + + +def init_weights(net, init_type="kaiming", init_gain=0.02, debug=False): + """Initialize network weights. + + Parameters: + net (network) -- network to be initialized + init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal + init_gain (float) -- scaling factor for normal, xavier and orthogonal. + + We use 'normal' in the original pix2pix and CycleGAN paper. But xavier and kaiming might + work better for some applications. Feel free to try yourself. + """ + + def init_func(m): # define the initialization function + classname = m.__class__.__name__ + if hasattr(m, "weight") and ( + classname.find("Conv") != -1 or classname.find("Linear") != -1 + ): + if debug: + print(classname) + if init_type == "normal": + init.normal_(m.weight.data, 0.0, init_gain) + elif init_type == "xavier": + init.xavier_normal_(m.weight.data, gain=init_gain) + elif init_type == "kaiming": + init.kaiming_normal_(m.weight.data, a=0, mode="fan_in") + elif init_type == "orthogonal": + init.orthogonal_(m.weight.data, gain=init_gain) + else: + raise NotImplementedError( + "initialization method [%s] is not implemented" % init_type + ) + if hasattr(m, "bias") and m.bias is not None: + init.constant_(m.bias.data, 0.0) + elif ( + classname.find("BatchNorm2d") != -1 + ): # BatchNorm Layer's weight is not a matrix; only normal distribution applies. + init.normal_(m.weight.data, 1.0, init_gain) + init.constant_(m.bias.data, 0.0) + + net.apply(init_func) # apply the initialization function + + +def init_net( + net, + init_type="normal", + init_gain=0.02, + gpu_ids=[], + debug=False, + initialize_weights=True, +): + """Initialize a network: 1. register CPU/GPU device (with multi-GPU support); 2. initialize the network weights + Parameters: + net (network) -- the network to be initialized + init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal + gain (float) -- scaling factor for normal, xavier and orthogonal. + gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2 + + Return an initialized network. + """ + if len(gpu_ids) > 0: + assert torch.cuda.is_available() + net.to(gpu_ids[0]) + # if not amp: + # net = torch.nn.DataParallel(net, gpu_ids) # multi-GPUs for non-AMP training + if initialize_weights: + init_weights(net, init_type, init_gain=init_gain, debug=debug) + return net + + +def define_D( + input_nc, + ndf, + netD, + n_layers_D=3, + image_size=256, + feature_dim=256, + max_conv_dim=512, + norm="batch", + init_type="normal", + init_gain=0.02, + no_antialias=False, + gpu_ids=[], + opt=None, +): + """Create a discriminator + + Parameters: + input_nc (int) -- the number of channels in input images + ndf (int) -- the number of filters in the first conv layer + netD (str) -- the architecture's name: basic | n_layers | pixel + n_layers_D (int) -- the number of conv layers in the discriminator; effective when netD=='n_layers' + norm (str) -- the type of normalization layers used in the network. + init_type (str) -- the name of the initialization method. + init_gain (float) -- scaling factor for normal, xavier and orthogonal. + gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2 + + Returns a discriminator + + Our current implementation provides three types of discriminators: + [basic]: 'PatchGAN' classifier described in the original pix2pix paper. + It can classify whether 70x70 overlapping patches are real or fake. + Such a patch-level discriminator architecture has fewer parameters + than a full-image discriminator and can work on arbitrarily-sized images + in a fully convolutional fashion. + + [n_layers]: With this mode, you cna specify the number of conv layers in the discriminator + with the parameter (default=3 as used in [basic] (PatchGAN).) + + [pixel]: 1x1 PixelGAN discriminator can classify whether a pixel is real or not. + It encourages greater color diversity but has no effect on spatial statistics. + + The discriminator has been initialized by . It uses Leaky RELU for non-linearity. + """ + net = None + norm_layer = get_norm_layer(norm_type=norm) + + if netD == "basic": # default PatchGAN classifier + net = NLayerDiscriminator( + input_nc, + ndf, + n_layers=3, + image_size=image_size, + norm_layer=norm_layer, + no_antialias=no_antialias, + ) + elif netD == "n_layers": # more options + net = NLayerDiscriminator( + input_nc, + ndf, + n_layers_D, + norm_layer=norm_layer, + no_antialias=no_antialias, + ) + elif netD == "pixel": # classify if each pixel is real or fake + net = PixelDiscriminator(input_nc, ndf, norm_layer=norm_layer) + # elif 'stylegan2' in netD: + # net = StyleGAN2Discriminator(input_nc, ndf, n_layers_D, no_antialias=no_antialias, opt=opt) + # elif netD == 'NCE': # default PatchGAN classifier + # net = NCEDiscriminator(input_nc, ndf, n_layers=3, image_size = image_size, feature_dim = feature_dim, max_conv_dim = max_conv_dim, norm_layer=norm_layer, no_antialias=no_antialias,) + else: + raise NotImplementedError( + "Discriminator model name [%s] is not recognized" % netD + ) + + return init_net( + net, init_type, init_gain, gpu_ids, initialize_weights=("stylegan2" not in netD) + ) + + +############################################################################## +# Classes +############################################################################## + + +class GANLoss(nn.Module): + """Define different GAN objectives. + + The GANLoss class abstracts away the need to create the target label tensor + that has the same size as the input. + """ + + def __init__(self, gan_mode, target_real_label=1.0, target_fake_label=0.0): + """Initialize the GANLoss class. + + Parameters: + gan_mode (str) - - the type of GAN objective. It currently supports vanilla, lsgan, and wgangp. + target_real_label (bool) - - label for a real image + target_fake_label (bool) - - label of a fake image + + Note: Do not use sigmoid as the last layer of Discriminator. + LSGAN needs no sigmoid. vanilla GANs will handle it with BCEWithLogitsLoss. + """ + super(GANLoss, self).__init__() + self.register_buffer("real_label", torch.tensor(target_real_label)) + self.register_buffer("fake_label", torch.tensor(target_fake_label)) + self.gan_mode = gan_mode + if gan_mode == "lsgan": + self.loss = nn.MSELoss() + elif gan_mode == "vanilla": + self.loss = nn.BCEWithLogitsLoss() + elif gan_mode in ["wgangp", "nonsaturating"]: + self.loss = None + elif gan_mode == "hinge": + self.loss = None + else: + raise NotImplementedError("gan mode %s not implemented" % gan_mode) + + def get_target_tensor(self, prediction, target_is_real): + """Create label tensors with the same size as the input. + + Parameters: + prediction (tensor) - - tpyically the prediction from a discriminator + target_is_real (bool) - - if the ground truth label is for real images or fake images + + Returns: + A label tensor filled with ground truth label, and with the size of the input + """ + + if target_is_real: + target_tensor = self.real_label + else: + target_tensor = self.fake_label + return target_tensor.expand_as(prediction) + + def __call__(self, prediction, target_is_real): + """Calculate loss given Discriminator's output and grount truth labels. + + Parameters: + prediction (tensor) - - tpyically the prediction output from a discriminator + target_is_real (bool) - - if the ground truth label is for real images or fake images + + Returns: + the calculated loss. + """ + bs = prediction.size(0) + if self.gan_mode in ["lsgan", "vanilla"]: + target_tensor = self.get_target_tensor(prediction, target_is_real) + loss = self.loss(prediction, target_tensor) + elif self.gan_mode == "wgangp": + if target_is_real: + loss = -prediction.mean() + else: + loss = prediction.mean() + elif self.gan_mode == "nonsaturating": + if target_is_real: + loss = F.softplus(-prediction).view(bs, -1).mean(dim=1) + else: + loss = F.softplus(prediction).view(bs, -1).mean(dim=1) + elif self.gan_mode == "hinge": + if target_is_real: + minvalue = torch.min( + prediction - 1, torch.zeros(prediction.shape).to(prediction.device) + ) + loss = -torch.mean(minvalue) + else: + minvalue = torch.min( + -prediction - 1, torch.zeros(prediction.shape).to(prediction.device) + ) + loss = -torch.mean(minvalue) + return loss + + +def cal_gradient_penalty( + netD, real_data, fake_data, device, type="mixed", constant=1.0, lambda_gp=10.0 +): + """Calculate the gradient penalty loss, used in WGAN-GP paper https://arxiv.org/abs/1704.00028 + + Arguments: + netD (network) -- discriminator network + real_data (tensor array) -- real images + fake_data (tensor array) -- generated images from the generator + device (str) -- GPU / CPU: from torch.device('cuda:{}'.format(self.gpu_ids[0])) if self.gpu_ids else torch.device('cpu') + type (str) -- if we mix real and fake data or not [real | fake | mixed]. + constant (float) -- the constant used in formula ( | |gradient||_2 - constant)^2 + lambda_gp (float) -- weight for this loss + + Returns the gradient penalty loss + """ + if lambda_gp > 0.0: + if ( + type == "real" + ): # either use real images, fake images, or a linear interpolation of two. + interpolatesv = real_data + elif type == "fake": + interpolatesv = fake_data + elif type == "mixed": + alpha = torch.rand(real_data.shape[0], 1, device=device) + alpha = ( + alpha.expand( + real_data.shape[0], real_data.nelement() // real_data.shape[0] + ) + .contiguous() + .view(*real_data.shape) + ) + interpolatesv = alpha * real_data + ((1 - alpha) * fake_data) + else: + raise NotImplementedError("{} not implemented".format(type)) + interpolatesv.requires_grad_(True) + disc_interpolates = netD(interpolatesv) + gradients = torch.autograd.grad( + outputs=disc_interpolates, + inputs=interpolatesv, + grad_outputs=torch.ones(disc_interpolates.size()).to(device), + create_graph=True, + retain_graph=True, + only_inputs=True, + ) + gradients = gradients[0].view(real_data.size(0), -1) # flat the data + gradient_penalty = ( + ((gradients + 1e-16).norm(2, dim=1) - constant) ** 2 + ).mean() * lambda_gp # added eps + return gradient_penalty, gradients + else: + return 0.0, None + + +class Normalize(nn.Module): + def __init__(self, power=2): + super(Normalize, self).__init__() + self.power = power + + def forward(self, x): + norm = x.pow(self.power).sum(1, keepdim=True).pow(1.0 / self.power) + out = x.div(norm + 1e-7) + return out + + +################################################################################## +# Sequential Models +################################################################################## + + +# class ResBlocks(nn.Module): +# def __init__( +# self, num_blocks, dim, norm="inst", activation="relu", pad_type="zero", nz=0 +# ): +# super(ResBlocks, self).__init__() +# self.model = [] +# for i in range(num_blocks): +# self.model += [ +# ResBlock( +# dim, norm=norm, activation=activation, pad_type=pad_type, nz=nz +# ) +# ] +# self.model = nn.Sequential(*self.model) + +# def forward(self, x): +# return self.model(x) + + +################################################################################## +# Basic Blocks +################################################################################## +def cat_feature(x, y): + y_expand = y.view(y.size(0), y.size(1), 1, 1).expand( + y.size(0), y.size(1), x.size(2), x.size(3) + ) + x_cat = torch.cat([x, y_expand], 1) + return x_cat + + +class Conv2dBlock(nn.Module): + def __init__( + self, + input_dim, + output_dim, + kernel_size, + stride, + padding=0, + norm="none", + activation="relu", + pad_type="zero", + ): + super(Conv2dBlock, self).__init__() + self.use_bias = True + # initialize padding + if pad_type == "reflect": + self.pad = nn.ReflectionPad2d(padding) + elif pad_type == "zero": + self.pad = nn.ZeroPad2d(padding) + else: + assert 0, "Unsupported padding type: {}".format(pad_type) + + # initialize normalization + norm_dim = output_dim + if norm == "batch": + self.norm = nn.BatchNorm2d(norm_dim) + elif norm == "inst": + self.norm = nn.InstanceNorm2d(norm_dim, track_running_stats=False) + elif norm == "ln": + self.norm = LayerNorm(norm_dim) + elif norm == "none": + self.norm = None + else: + assert 0, "Unsupported normalization: {}".format(norm) + + # initialize activation + if activation == "relu": + self.activation = nn.ReLU(inplace=True) + elif activation == "lrelu": + self.activation = nn.LeakyReLU(0.2, inplace=True) + elif activation == "prelu": + self.activation = nn.PReLU() + elif activation == "selu": + self.activation = nn.SELU(inplace=True) + elif activation == "tanh": + self.activation = nn.Tanh() + elif activation == "none": + self.activation = None + else: + assert 0, "Unsupported activation: {}".format(activation) + + # initialize convolution + self.conv = nn.Conv2d( + input_dim, output_dim, kernel_size, stride, bias=self.use_bias + ) + + def forward(self, x): + x = self.conv(self.pad(x)) + if self.norm: + x = self.norm(x) + if self.activation: + x = self.activation(x) + return x + + +class LinearBlock(nn.Module): + def __init__(self, input_dim, output_dim, norm="none", activation="relu"): + super(LinearBlock, self).__init__() + use_bias = True + # initialize fully connected layer + self.fc = nn.Linear(input_dim, output_dim, bias=use_bias) + + # initialize normalization + norm_dim = output_dim + if norm == "batch": + self.norm = nn.BatchNorm1d(norm_dim) + elif norm == "inst": + self.norm = nn.InstanceNorm1d(norm_dim) + elif norm == "ln": + self.norm = LayerNorm(norm_dim) + elif norm == "none": + self.norm = None + else: + assert 0, "Unsupported normalization: {}".format(norm) + + # initialize activation + if activation == "relu": + self.activation = nn.ReLU(inplace=True) + elif activation == "lrelu": + self.activation = nn.LeakyReLU(0.2, inplace=True) + elif activation == "prelu": + self.activation = nn.PReLU() + elif activation == "selu": + self.activation = nn.SELU(inplace=True) + elif activation == "tanh": + self.activation = nn.Tanh() + elif activation == "none": + self.activation = None + else: + assert 0, "Unsupported activation: {}".format(activation) + + def forward(self, x): + out = self.fc(x) + if self.norm: + out = self.norm(out) + if self.activation: + out = self.activation(out) + return out + + +################################################################################## +# Normalization layers +################################################################################## + + +class LayerNorm(nn.Module): + def __init__(self, num_features, eps=1e-5, affine=True): + super(LayerNorm, self).__init__() + self.num_features = num_features + self.affine = affine + self.eps = eps + + if self.affine: + self.gamma = nn.Parameter(torch.Tensor(num_features).uniform_()) + self.beta = nn.Parameter(torch.zeros(num_features)) + + def forward(self, x): + shape = [-1] + [1] * (x.dim() - 1) + mean = x.view(x.size(0), -1).mean(1).view(*shape) + std = x.view(x.size(0), -1).std(1).view(*shape) + x = (x - mean) / (std + self.eps) + + if self.affine: + shape = [1, -1] + [1] * (x.dim() - 2) + x = x * self.gamma.view(*shape) + self.beta.view(*shape) + return x + + +class NLayerDiscriminator(nn.Module): + """Defines a PatchGAN discriminator""" + + def __init__( + self, + input_nc, + ndf=64, + n_layers=3, + image_size=256, + norm_layer=nn.BatchNorm2d, + no_antialias=False, + ): + """Construct a PatchGAN discriminator + + Parameters: + input_nc (int) -- the number of channels in input images + ndf (int) -- the number of filters in the last conv layer + n_layers (int) -- the number of conv layers in the discriminator + norm_layer -- normalization layer + """ + super(NLayerDiscriminator, self).__init__() + if ( + type(norm_layer) == functools.partial + ): # no need to use bias as BatchNorm2d has affine parameters + use_bias = norm_layer.func == nn.InstanceNorm2d + else: + use_bias = norm_layer == nn.InstanceNorm2d + + kw = 4 + padw = 1 + if no_antialias: + sequence = [ + nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), + nn.LeakyReLU(0.2, True), + ] + else: + sequence = [ + nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=1, padding=padw), + nn.LeakyReLU(0.2, True), + Downsample(ndf), + ] + nf_mult = 1 + nf_mult_prev = 1 + for n in range(1, n_layers): # gradually increase the number of filters + nf_mult_prev = nf_mult + nf_mult = min(2**n, 8) + if no_antialias: + sequence += [ + nn.Conv2d( + ndf * nf_mult_prev, + ndf * nf_mult, + kernel_size=kw, + stride=2, + padding=padw, + bias=use_bias, + ), + norm_layer(ndf * nf_mult), + nn.LeakyReLU(0.2, True), + ] + else: + sequence += [ + nn.Conv2d( + ndf * nf_mult_prev, + ndf * nf_mult, + kernel_size=kw, + stride=1, + padding=padw, + bias=use_bias, + ), + norm_layer(ndf * nf_mult), + nn.LeakyReLU(0.2, True), + Downsample(ndf * nf_mult), + ] + + nf_mult_prev = nf_mult + nf_mult = min(2**n_layers, 8) + sequence += [ + nn.Conv2d( + ndf * nf_mult_prev, + ndf * nf_mult, + kernel_size=kw, + stride=1, + padding=padw, + bias=use_bias, + ), + norm_layer(ndf * nf_mult), + nn.LeakyReLU(0.2, True), + ] + + sequence += [ + nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw) + ] # output 1 channel prediction map + self.model = nn.Sequential(*sequence) + + def forward(self, input): + """Standard forward.""" + logit = self.model(input) + return logit + + +class PixelDiscriminator(nn.Module): + """Defines a 1x1 PatchGAN discriminator (pixelGAN)""" + + def __init__(self, input_nc, ndf=64, norm_layer=nn.BatchNorm2d): + """Construct a 1x1 PatchGAN discriminator + + Parameters: + input_nc (int) -- the number of channels in input images + ndf (int) -- the number of filters in the last conv layer + norm_layer -- normalization layer + """ + super(PixelDiscriminator, self).__init__() + if ( + type(norm_layer) == functools.partial + ): # no need to use bias as BatchNorm2d has affine parameters + use_bias = norm_layer.func == nn.InstanceNorm2d + else: + use_bias = norm_layer == nn.InstanceNorm2d + + self.net = [ + nn.Conv2d(input_nc, ndf, kernel_size=1, stride=1, padding=0), + nn.LeakyReLU(0.2, True), + nn.Conv2d(ndf, ndf * 2, kernel_size=1, stride=1, padding=0, bias=use_bias), + norm_layer(ndf * 2), + nn.LeakyReLU(0.2, True), + nn.Conv2d(ndf * 2, 1, kernel_size=1, stride=1, padding=0, bias=use_bias), + ] + + self.net = nn.Sequential(*self.net) + + def forward(self, input): + """Standard forward.""" + return self.net(input) + + +class PatchDiscriminator(NLayerDiscriminator): + """Defines a PatchGAN discriminator""" + + def __init__( + self, + input_nc, + ndf=64, + n_layers=3, + norm_layer=nn.BatchNorm2d, + no_antialias=False, + ): + super().__init__(input_nc, ndf, 2, norm_layer, no_antialias) + + def forward(self, input): + B, C, H, W = input.size(0), input.size(1), input.size(2), input.size(3) + size = 16 + Y = H // size + X = W // size + input = input.view(B, C, Y, size, X, size) + input = ( + input.permute(0, 2, 4, 1, 3, 5).contiguous().view(B * Y * X, C, size, size) + ) + return super().forward(input) + + +class GroupedChannelNorm(nn.Module): + def __init__(self, num_groups): + super().__init__() + self.num_groups = num_groups + + def forward(self, x): + shape = list(x.shape) + new_shape = [shape[0], self.num_groups, shape[1] // self.num_groups] + shape[2:] + x = x.view(*new_shape) + mean = x.mean(dim=2, keepdim=True) + std = x.std(dim=2, keepdim=True) + x_norm = (x - mean) / (std + 1e-7) + return x_norm.view(*shape) diff --git a/requirements.txt b/requirements.txt index 08ed5ee..3c3eb54 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1 +1,6 @@ -torch \ No newline at end of file +torch +torchvision +Pillow +numpy +kornia +scipy \ No newline at end of file diff --git a/run.py b/run.py index 2a7c328..daa96a6 100644 --- a/run.py +++ b/run.py @@ -12,10 +12,18 @@ import torch.nn.functional as F import tqdm from comfy.utils import ProgressBar +from .module_cast.cast_model import inference_ucast + +from .constants import ( + CAST_DEFAULT, + CAST_NET_AE_PATH, + CAST_NET_DEC_B_PATH, + CAST_TYPES, + CAST_VGG_PATH, +) + from .module_neural_neighbor.neural_neighbor_model import inference_neural_neighbor -# from .constants import ( -# ) base_dir = Path(__file__).resolve().parent @@ -99,24 +107,52 @@ class NeuralNeighbor: return (torch.cat(result, dim=0),) -# def inference_realviformer(images: torch.Tensor, model: RealViformer) -> torch.Tensor: -# padded = False -# with torch.no_grad(): -# h, w = images.shape[-2:] -# ah, aw = h % 4, w % 4 -# padh = 0 if ah == 0 else 4 - ah -# padw = 0 if aw == 0 else 4 - aw -# if padh != 0 or padw != 0: -# padded = True -# images = F.pad( -# images.squeeze(0), pad=(padw, 0, padh, 0), mode="reflect" -# ).unsqueeze(0) -# outputs = model(images) +class CAST: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "src_img": ("IMAGE",), + "style_img": ("IMAGE",), + "model_arch": (CAST_TYPES, {"default": CAST_DEFAULT}), + }, + } -# if padded: -# outputs = outputs[..., padh * 4 :, padw * 4 :] + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("res_img",) + FUNCTION = "todo" + CATEGORY = "StyleTransferPlus" -# return outputs + def todo( + self, + src_img: torch.Tensor, + style_img: torch.Tensor, + model_arch: str, + ): + print(f"{src_img.shape=}") + print(f"{style_img.shape=}") + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + params = { + "style_img": style_img, + "device": device, + "ckpt_ae": f"{base_dir}/models/{model_arch}_model/{CAST_NET_AE_PATH}", + "ckpt_decB": f"{base_dir}/models/{model_arch}_model/{CAST_NET_DEC_B_PATH}", + "vgg_path": f"{base_dir}/{CAST_VGG_PATH}", + } + + num_frames = src_img.size(0) + pbar = ProgressBar(num_frames) + + result: list[torch.Tensor] = [] + with torch.inference_mode(mode=False): + for i in range(num_frames): + params["src_img"] = src_img[i].unsqueeze(0) + res_tensor = inference_ucast(**params) + result.append(res_tensor.permute(0, 2, 3, 1)) + pbar.update_absolute(i, num_frames) + + return (torch.cat(result, dim=0),) # class RealViFormerSR: diff --git a/workflows/wf_cast.png b/workflows/wf_cast.png new file mode 100644 index 0000000..ddf6382 Binary files /dev/null and b/workflows/wf_cast.png differ diff --git a/workflows/workflow_cast.json b/workflows/workflow_cast.json new file mode 100644 index 0000000..51d8bf8 --- /dev/null +++ b/workflows/workflow_cast.json @@ -0,0 +1,191 @@ +{ + "last_node_id": 6, + "last_link_id": 10, + "nodes": [ + { + "id": 2, + "type": "LoadImage", + "pos": [ + 309, + 187 + ], + "size": { + "0": 315, + "1": 314 + }, + "flags": {}, + "order": 0, + "mode": 0, + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 8 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "MASK", + "type": "MASK", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "006.jpg", + "image" + ] + }, + { + "id": 3, + "type": "LoadImage", + "pos": [ + 315, + 570 + ], + "size": { + "0": 315, + "1": 314 + }, + "flags": {}, + "order": 1, + "mode": 0, + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 9 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "MASK", + "type": "MASK", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "S4.jpg", + "image" + ] + }, + { + "id": 4, + "type": "SaveImage", + "pos": [ + 1050, + 327 + ], + "size": { + "0": 315, + "1": 270 + }, + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 10 + } + ], + "properties": {}, + "widgets_values": [ + "StyleTransfer" + ] + }, + { + "id": 6, + "type": "CAST", + "pos": [ + 767, + 326 + ], + "size": { + "0": 216.59999084472656, + "1": 78 + }, + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [ + { + "name": "src_img", + "type": "IMAGE", + "link": 8 + }, + { + "name": "style_img", + "type": "IMAGE", + "link": 9 + } + ], + "outputs": [ + { + "name": "res_img", + "type": "IMAGE", + "links": [ + 10 + ], + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "CAST" + }, + "widgets_values": [ + "CAST" + ] + } + ], + "links": [ + [ + 8, + 2, + 0, + 6, + 0, + "IMAGE" + ], + [ + 9, + 3, + 0, + 6, + 1, + "IMAGE" + ], + [ + 10, + 6, + 0, + 4, + 0, + "IMAGE" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 1.100000000000002, + "offset": [ + -106.49900776658754, + -117.32656624605369 + ] + } + }, + "version": 0.4 +} \ No newline at end of file