add CAST
This commit is contained in:
FuouM
@@ -1,3 +1,6 @@
|
||||
*.pth
|
||||
|
||||
|
||||
# Byte-compiled / optimized / DLL files
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
|
||||
@@ -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 |
|
||||
|:-:|-|-|
|
||||
| | X | X | X
|
||||
| Neural Neighbor (2020-2022) - Slowest | CAST (2022) - Fast | X |
|
||||
|:-:|:-:|-|
|
||||
| |  | X |
|
||||
|
||||
### Neural Neighbor
|
||||
|
||||
@@ -32,3 +32,46 @@ See more information on the original repository. No model is needed.
|
||||
|
||||
|
||||
|
||||
### 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)
|
||||
|
||||
|
||||
|
||||
## 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}
|
||||
}
|
||||
```
|
||||
+3
-1
@@ -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",
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -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"
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 1.1 MiB |
@@ -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
|
||||
@@ -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
|
||||
@@ -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 <opt.n_epochs> epochs
|
||||
and linearly decay the rate to zero over the next <opt.n_epochs_decay> 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 <init_func>
|
||||
|
||||
|
||||
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 <n_layers_D> (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 <init_net>. 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)
|
||||
+6
-1
@@ -1 +1,6 @@
|
||||
torch
|
||||
torch
|
||||
torchvision
|
||||
Pillow
|
||||
numpy
|
||||
kornia
|
||||
scipy
|
||||
@@ -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:
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 585 KiB |
@@ -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
|
||||
}
|
||||
Reference in New Issue
Block a user