def __init__(self, opt, sw, sh):
super().__init__()
kw = 3
pw = int(np.ceil((kw - 1.0) / 2))
ndf = opt.ngf # 64
self.sw = sw
self.sh = sh
self.opt = opt
norm_layer = get_nonspade_norm_layer(opt, opt.norm_E)
self.layer1 = norm_layer(nn.Conv2d(3, ndf, kw, stride=2, padding=pw))
self.layer2 = norm_layer(
nn.Conv2d(ndf * 1, ndf * 2, kw, stride=2, padding=pw))
self.layer3 = norm_layer(
nn.Conv2d(ndf * 2, ndf * 4, kw, stride=2, padding=pw))
self.layer4 = norm_layer(
nn.Conv2d(ndf * 4, ndf * 8, kw, stride=2, padding=pw))
self.layer5 = norm_layer(
nn.Conv2d(ndf * 8, ndf * 16, kw, stride=2, padding=pw))
if opt.crop_size >= 256:
self.layer6 = norm_layer(
nn.Conv2d(ndf * 8, ndf * 8, kw, stride=2, padding=pw))
self.so = s0 = 4
self.adaptivepool = nn.AdaptiveAvgPool2d(1)
self.fc = self.fc = nn.Conv2d(ndf * 16, ndf * 16 * self.sw * self.sh,
1, 1, 0)
self.actvn = nn.LeakyReLU(0.2, False)
self.opt = opt
def __init__(self, opt, input_nc, netS=None, params_free_layers=0, first_layer_free=False):
super().__init__()
self.opt = opt
norm_layer = get_nonspade_norm_layer(opt, opt.norm_F)
activation = nn.LeakyReLU(0.2)
nf = opt.ngf
nf_max = 1024
self.netS = netS if netS is not None else opt.netS
self.unet = 'unet' in self.netS
self.decode = 'decoder' in self.netS or self.unet
self.n_downsample_S = n_downsample_S = opt.n_downsample_G
self.params_free_layers = params_free_layers if params_free_layers != -1 else n_downsample_S
self.first_layer_free = first_layer_free
ch = [min(nf_max, nf * (2 ** i)) for i in range(n_downsample_S + 1)]
if not first_layer_free:
layer = [nn.Conv2d(input_nc, nf, kernel_size=3, padding=1), activation]
self.conv_first = nn.Sequential(*layer)
# downsample
for i in range(n_downsample_S):
layer = [nn.Conv2d(ch[i], ch[i+1], kernel_size=3, stride=2, padding=1), activation]
if i >= params_free_layers or 'decoder' in netS:
setattr(self, 'down_%d' % i, nn.Sequential(*layer))
# upsample
if self.decode:
for i in reversed(range(n_downsample_S)):
ch_i = ch[i+1] * (2 if self.unet and i != n_downsample_S -1 else 1)
# layer = [nn.ConvTranspose2d(ch_i, ch[i], kernel_size=3, stride=2, padding=1, output_padding=1), activation]
layer = [nn.Upsample(scale_factor=2), nn.Conv2d(ch_i, ch[i], kernel_size=3, padding=1), activation]
if i >= params_free_layers:
setattr(self, 'up_%d' % i, nn.Sequential(*layer))
def __init__(self, opt, ic, oc, size):
super().__init__()
self.opt = opt
self.downsample = True if size == 256 else False
nf = opt.ngf
opt.spade_ic = ic
if opt.warp_reverseG_s:
self.backbone_0 = SPADEResnetBlock(4 * nf, 4 * nf, opt)
else:
self.backbone_0 = SPADEResnetBlock(4 * nf, 8 * nf, opt)
self.backbone_1 = SPADEResnetBlock(8 * nf, 8 * nf, opt)
self.backbone_2 = SPADEResnetBlock(8 * nf, 8 * nf, opt)
self.backbone_3 = SPADEResnetBlock(8 * nf, 4 * nf, opt)
self.backbone_4 = SPADEResnetBlock(4 * nf, 2 * nf, opt)
self.backbone_5 = SPADEResnetBlock(2 * nf, nf, opt)
del opt.spade_ic
if self.downsample:
kw = 3
pw = int(np.ceil((kw - 1.0) / 2))
ndf = opt.ngf
norm_layer = get_nonspade_norm_layer(opt, opt.norm_E)
self.layer1 = norm_layer(nn.Conv2d(ic, ndf, kw, stride=1, padding=pw))
self.layer2 = norm_layer(nn.Conv2d(ndf * 1, ndf * 2, 4, stride=2, padding=pw))
self.layer3 = norm_layer(nn.Conv2d(ndf * 2, ndf * 4, kw, stride=1, padding=pw))
self.layer4 = norm_layer(nn.Conv2d(ndf * 4, ndf * 4, 4, stride=2, padding=pw))
self.up = nn.Upsample(scale_factor=2)
self.actvn = nn.LeakyReLU(0.2, False)
self.conv_img = nn.Conv2d(nf, oc, 3, padding=1)
def __init__(self, opt):
super().__init__()
self.opt = opt
kw = 4
padw = int(np.ceil((kw - 1.0) / 2))
nf = opt.ndf
input_nc = self.compute_D_input_nc(opt)
norm_layer = get_nonspade_norm_layer(opt, opt.norm_D)
sequence = [[nn.Conv2d(input_nc, nf, kernel_size=kw, stride=2, padding=padw),
nn.LeakyReLU(0.2, False)]]
for n in range(1, opt.n_layers_D):
nf_prev = nf
nf = min(nf * 2, 512)
stride = 1 if n == opt.n_layers_D - 1 else 2
sequence += [[norm_layer(nn.Conv2d(nf_prev, nf, kernel_size=kw,
stride=stride, padding=padw)),
nn.LeakyReLU(0.2, False)
]]
sequence += [[nn.Conv2d(nf, 1, kernel_size=kw, stride=1, padding=padw)]]
# We divide the layers into groups to extract intermediate layer outputs
for n in range(len(sequence)):
self.add_module('model' + str(n), nn.Sequential(*sequence[n]))
def __init__(self, opt):
super().__init__()
self.opt = opt
RealFakeDiscriminator.modify_commandline_options(self.opt)
n = math.log2(opt['RF_size'])
assert n == round(n)
assert n >= 3
n = int(n)
in_c = opt['RF_in']
nc = opt['RF_nc']
norm_layer = get_nonspade_norm_layer(opt, opt['RF_norm'])
activation = nn.LeakyReLU(0.2, inplace=True)
model = [nn.Conv2d(in_c, nc, 4, 2, 1, bias=False), activation]
for i in range(n - 3):
model += [
norm_layer(
nn.Conv2d(nc * 2**(i),
nc * 2**(i + 1),
4,
2,
1,
bias=False)), activation
]
model += [nn.Conv2d(nc * 2**(n - 3), 1, 4, 1, 0, bias=False)]
self.model = nn.Sequential(*model)
def __init__(self, opt):
super().__init__()
kw = 3
pw = int(np.ceil((kw - 1.0) / 2))
ndf = opt.ngf
norm_layer = get_nonspade_norm_layer(opt, opt.norm_E)
self.layer1 = norm_layer(nn.Conv2d(3, ndf, kw, stride=2, padding=pw))
self.layer2 = norm_layer(
nn.Conv2d(ndf * 1, ndf * 2, kw, stride=2, padding=pw))
self.layer3 = norm_layer(
nn.Conv2d(ndf * 2, ndf * 4, kw, stride=2, padding=pw))
self.layer4 = norm_layer(
nn.Conv2d(ndf * 4, ndf * 8, kw, stride=2, padding=pw))
self.layer5 = norm_layer(
nn.Conv2d(ndf * 8, ndf * 8, kw, stride=2, padding=pw))
if opt.crop_size >= 256:
self.layer6 = norm_layer(
nn.Conv2d(ndf * 8, ndf * 8, kw, stride=2, padding=pw))
self.so = s0 = 4
self.fc_mu = nn.Linear(ndf * 8 * s0 * s0, opt.z_dim)
self.fc_var = nn.Linear(ndf * 8 * s0 * s0, opt.z_dim)
self.actvn = nn.LeakyReLU(0.2, False)
self.opt = opt
def __init__(self, opt):
super().__init__()
self.opt = opt
kw = 4
padw = int(np.ceil((kw - 1.0) / 2))
nf = opt.ndf
norm_layer = get_nonspade_norm_layer(opt, opt.norm_D)
self.emb = nn.Conv2d(opt.label_nc, opt.mask_emb_dim, 1, bias=False)
if self.opt.embed_attributes:
attr_dim = opt.attr_emb_dim
self.attr_emb = nn.Conv2d(opt.attr_nc, opt.attr_emb_dim, 1)
else:
attr_dim = 0
self.attr_emb = None
if self.opt.embed_captions:
sent_dim = opt.attr_emb_dim # Same as attributes
self.text_emb = SentenceSemanticAttention(
opt.label_nc,
query_dim=128,
output_dim=sent_dim,
num_attention_heads=opt.attention_heads)
else:
sent_dim = 0
self.text_emb = None
sequence = [[
nn.Conv2d(opt.output_nc + opt.mask_emb_dim + attr_dim + sent_dim,
nf,
kernel_size=kw,
stride=2,
padding=padw),
nn.LeakyReLU(0.2, False)
]]
for n in range(1, opt.n_layers_D):
nf_prev = nf
nf = min(nf * 2, 512)
stride = 1 if n == opt.n_layers_D - 1 else 2
sequence += [[
norm_layer(
nn.Conv2d(nf_prev,
nf,
kernel_size=kw,
stride=stride,
padding=padw)),
nn.LeakyReLU(0.2, False)
]]
sequence += [[
nn.Conv2d(nf, 1, kernel_size=kw, stride=1, padding=padw)
]]
# We divide the layers into groups to extract intermediate layer outputs
for n in range(len(sequence)):
self.add_module('model' + str(n), nn.Sequential(*sequence[n]))
def __init__(self, opt):
super().__init__()
self.opt = opt
nf = opt.ngf
self.sw, self.sh = self.compute_latent_vector_size(opt)
if opt.use_vae:
# In case of VAE, we will sample from random z vector
self.fc = nn.Linear(opt.z_dim, 16*nf*self.sw*self.sh)
else:
# Otherwise, we make the network deterministic by starting with
# downsampled segmentation map instead of random z
self.fc = nn.Conv2d(self.opt.semantic_nc, 16*nf*self.sw*self.sh, 3, padding=1)
# semantic map encoder to incorporate context
# currently only supports 'normal' mode
norm_layer = get_nonspade_norm_layer(opt, opt.norm_G)
self.labelenc1 = nn.Sequential(norm_layer(nn.Conv2d(self.opt.semantic_nc, nf, 3, padding=1), opt), nn.LeakyReLU(0.2, True)) # 256
self.labelenc2 = nn.Sequential(norm_layer(nn.Conv2d(nf, nf, 3, padding=1, stride=2), opt), nn.LeakyReLU(0.2, True)) # 128
self.labelenc3 = nn.Sequential(norm_layer(nn.Conv2d(nf, nf, 3, padding=1, stride=2), opt), nn.LeakyReLU(0.2, True)) # 64
self.labelenc4 = nn.Sequential(norm_layer(nn.Conv2d(nf, nf, 3, padding=1, stride=2), opt), nn.LeakyReLU(0.2, True)) # 32
self.labelenc5 = nn.Sequential(norm_layer(nn.Conv2d(nf, nf, 3, padding=1, stride=2), opt), nn.LeakyReLU(0.2, True)) # 16
self.labelenc6 = nn.Sequential(norm_layer(nn.Conv2d(nf, nf, 3, padding=1, stride=2), opt), nn.LeakyReLU(0.2, True)) # 8
self.labelenc7 = nn.Sequential(norm_layer(nn.Conv2d(nf, nf, 3, padding=1, stride=2), opt), nn.LeakyReLU(0.2, True)) # 4
# lateral for fpn
self.labellat1 = nn.Sequential(norm_layer(nn.Conv2d(nf, nf, 1), opt), nn.LeakyReLU(0.2, True))#16
self.labellat2 = nn.Sequential(norm_layer(nn.Conv2d(nf, nf, 1), opt), nn.LeakyReLU(0.2, True))#32
self.labellat3 = nn.Sequential(norm_layer(nn.Conv2d(nf, nf, 1), opt), nn.LeakyReLU(0.2, True))#64
self.labellat4 = nn.Sequential(norm_layer(nn.Conv2d(nf, nf, 1), opt), nn.LeakyReLU(0.2, True))#128
self.labellat5 = nn.Sequential(norm_layer(nn.Conv2d(nf, nf, 1), opt), nn.LeakyReLU(0.2, True))#256
self.labellat6 = nn.Sequential(norm_layer(nn.Conv2d(nf, nf, 1), opt), nn.LeakyReLU(0.2, True))
# semantic map decoder
self.labeldec1 = nn.Sequential(norm_layer(nn.Conv2d(nf, nf, 3, padding=1), opt), nn.LeakyReLU(0.2, True))
self.labeldec2 = nn.Sequential(norm_layer(nn.Conv2d(nf, nf, 3, padding=1), opt), nn.LeakyReLU(0.2, True))
self.labeldec3 = nn.Sequential(norm_layer(nn.Conv2d(nf, nf, 3, padding=1), opt), nn.LeakyReLU(0.2, True))
self.labeldec4 = nn.Sequential(norm_layer(nn.Conv2d(nf, nf, 3, padding=1), opt), nn.LeakyReLU(0.2, True))
self.labeldec5 = nn.Sequential(norm_layer(nn.Conv2d(nf, nf, 3, padding=1), opt), nn.LeakyReLU(0.2, True))
self.labeldec6 = nn.Sequential(norm_layer(nn.Conv2d(nf, nf, 3, padding=1), opt), nn.LeakyReLU(0.2, True))
# image generator
self.head_0 = DepthsepCCBlock(16*nf, 16*nf, opt, opt.semantic_nc + nf)
self.G_middle_0 = DepthsepCCBlock(16*nf, 16*nf, opt, opt.semantic_nc + nf)
self.G_middle_1 = DepthsepCCBlock(16*nf, 16*nf, opt, opt.semantic_nc + nf)
self.up_0 = DepthsepCCBlock(16*nf, 8*nf, opt, opt.semantic_nc + nf)
self.up_1 = DepthsepCCBlock(8*nf, 4*nf, opt, opt.semantic_nc + nf)
self.up_2 = DepthsepCCBlock(4*nf, 2*nf, opt, opt.semantic_nc + nf)
self.up_3 = DepthsepCCBlock(2*nf, 1*nf, opt, opt.semantic_nc + nf)
final_nc = nf
self.conv_img = nn.Conv2d(final_nc, 3, 3, padding=1)
self.up = nn.Upsample(scale_factor=2)
def __init__(self, opt):
super(RotateGenerator, self).__init__()
input_nc = 3
norm_layer = get_nonspade_norm_layer(opt, opt.norm_G)
activation = nn.ReLU(False)
# initial conv
self.first_layer = nn.Sequential(nn.ReflectionPad2d(opt.resnet_initial_kernel_size // 2),
norm_layer(nn.Conv2d(input_nc, opt.ngf,
kernel_size=opt.resnet_initial_kernel_size,
padding=0)),
activation)
# downsample
downsample_model = []
mult = 1
for i in range(opt.resnet_n_downsample):
downsample_model += [norm_layer(nn.Conv2d(opt.ngf * mult, opt.ngf * mult * 2,
kernel_size=3, stride=2, padding=1)),
activation]
mult *= 2
self.downsample_layers = nn.Sequential(*downsample_model)
# resnet blocks
resnet_model = []
for i in range(opt.resnet_n_blocks):
resnet_model += [ResnetBlock(opt.ngf * mult,
norm_layer=norm_layer,
activation=activation,
kernel_size=opt.resnet_kernel_size)]
self.resnet_layers = nn.Sequential(*resnet_model)
# upsample
upsample_model = []
for i in range(opt.resnet_n_downsample):
nc_in = int(opt.ngf * mult)
nc_out = int((opt.ngf * mult) / 2)
upsample_model += [norm_layer(nn.ConvTranspose2d(nc_in, nc_out,
kernel_size=3, stride=2,
padding=1, output_padding=1)),
activation]
mult = mult // 2
self.upsample_layers = nn.Sequential(*upsample_model)
# final output conv
self.final_layer = nn.Sequential(nn.ReflectionPad2d(3),
nn.Conv2d(nc_out, opt.output_nc, kernel_size=7, padding=0),
nn.Tanh())
def __init__(self, opt, input_nc=None):
super().__init__()
self.opt = opt
kw = 4
padw = int(np.ceil((kw - 1.0) / 2))
nf = opt.ndf
if input_nc is None:
input_nc = self.compute_D_input_nc(opt)
branch = []
sizes = (input_nc - 3, 3)
original_nf = nf
for input_nc in sizes:
nf = original_nf
norm_layer = get_nonspade_norm_layer(opt, opt.norm_D)
sequence = [[
nn.Conv2d(input_nc, nf, kernel_size=kw, stride=2,
padding=padw),
nn.LeakyReLU(0.2, False)
]]
for n in range(1, opt.n_layers_D):
nf_prev = nf
nf = min(nf * 2, 512)
stride = 1 if n == opt.n_layers_D - 1 else 2
sequence += [[
norm_layer(
nn.Conv2d(nf_prev,
nf,
kernel_size=kw,
stride=stride,
padding=padw)),
nn.LeakyReLU(0.2, False)
]]
branch.append(sequence)
sem_sequence = nn.ModuleList()
for n in range(len(branch[0])):
sem_sequence.append(nn.Sequential(*branch[0][n]))
self.sem_sequence = nn.Sequential(*sem_sequence)
sequence = branch[1]
sequence += [[
nn.Conv2d(nf, 1, kernel_size=kw, stride=1, padding=padw)
]]
# We divide the layers into groups to extract intermediate layer outputs
self.img_sequence = nn.ModuleList()
for n in range(len(sequence)):
self.img_sequence.append(nn.Sequential(*sequence[n]))
def __init__(self, opt):
super().__init__()
input_nc = opt.label_nc + (1 if opt.contain_dontcare_label else 0) + (0 if opt.no_instance else 1)
norm_layer = get_nonspade_norm_layer(opt, opt.norm_G)
activation = nn.ReLU(False)
model = []
## initial conv
model += [nn.ReflectionPad2d(opt.resnet_initial_kernel_size // 2),
norm_layer(nn.Conv2d(input_nc, opt.ngf,
kernel_size=opt.resnet_initial_kernel_size,
padding=0)),
activation]
### downsample
mult = 1
for i in range(opt.resnet_n_downsample):
model += [norm_layer(nn.Conv2d(opt.ngf * mult, opt.ngf * mult * 2,
kernel_size=3, stride=2, padding=1)),
activation]
mult *= 2
### resnet blocks
for i in range(opt.resnet_n_blocks):
model += [ResnetBlock(opt.ngf * mult,
norm_layer=norm_layer,
activation=activation,
kernel_size=opt.resnet_kernel_size)]
### upsample
for i in range(opt.resnet_n_downsample):
nc_in = int(opt.ngf * mult)
nc_out = int((opt.ngf * mult) / 2)
model += [norm_layer(nn.ConvTranspose2d(nc_in, nc_out,
kernel_size=3, stride=2,
padding=1, output_padding=1)),
activation]
mult = mult // 2
# final output conv
model += [nn.ReflectionPad2d(3),
nn.Conv2d(nc_out, opt.output_nc, kernel_size=7, padding=0),
nn.Tanh()]
self.model = nn.Sequential(*model)
def __init__(self, fin, fout, opt):
super().__init__()
# Attributes
self.learned_shortcut = (fin != fout)
fmiddle = min(fin, fout)
norm_layer = get_nonspade_norm_layer(opt, 'spectralsync_batch')
# create conv layers
self.conv_0 = norm_layer(
nn.Conv2d(fin, fmiddle, kernel_size=3, padding=1))
self.conv_1 = norm_layer(
nn.Conv2d(fmiddle, fout, kernel_size=3, padding=1))
if self.learned_shortcut:
self.conv_s = spectral_norm(
nn.Conv2d(fin, fout, kernel_size=1, bias=False))
def define_D(opt,
input_nc,
ndf,
n_layers_D,
norm='spectralinstance',
subarch='n_layers',
num_D=1,
getIntermFeat=False,
stride=2,
gpu_ids=[]):
norm_layer = get_nonspade_norm_layer(opt, norm_type=norm)
if opt.which_model_netD == 'multiscale':
netD = MultiscaleDiscriminator(opt, input_nc, ndf, n_layers_D,
norm_layer, subarch, num_D,
getIntermFeat, stride, gpu_ids)
elif opt.which_model_netD == 'n_layers':
netD = NLayerDiscriminator(input_nc, ndf, n_layers_D, norm_layer,
getIntermFeat)
elif opt.which_model_netD == 'syncframe':
if subarch == 'sync':
netD = SyncDiscriminator(opt, 64, ndf, 256, 2, 5)
elif subarch == 'frame':
netD = FrameDiscriminator(opt, 256, input_nc, ndf, 6)
elif subarch == 'syncframe':
netD = SepDiscriminator(opt, ndf, 256, 1920)
elif opt.which_model_netD == 'sepfea':
if subarch == 'mouth':
netD = AudioFeaDiscriminator(input_nc,
ndf,
n_layers_D,
norm_layer,
getIntermFeat,
stride,
nf_final=256)
else:
netD = SepFeaDiscriminator(opt, ndf, n_layers_D, norm_layer,
getIntermFeat, stride)
else:
raise ('unknown type discriminator %s!' % opt.which_model_netD)
if opt.isTrain and opt.print_D: netD.print_network()
if len(gpu_ids) > 0:
assert (torch.cuda.is_available())
netD.cuda()
netD.init_weights(opt.init_type, opt.init_variance)
return netD
def __init__(self, opt, sw, sh):
super().__init__()
kw = 3
pw = int(np.ceil((kw - 1.0) / 2))
ndf = opt.ngf # 64
self.sw = sw
self.sh = sh
self.opt = opt
norm_layer = get_nonspade_norm_layer(opt, opt.norm_E)
self.layer1 = norm_layer(nn.Conv2d(3, ndf, kw, stride=2, padding=pw))
self.layer2 = norm_layer(nn.Conv2d(ndf * 1, ndf * 2, kw, stride=2, padding=pw))
self.layer3 = norm_layer(nn.Conv2d(ndf * 2, ndf * 4, kw, stride=2, padding=pw))
self.layer4 = norm_layer(nn.Conv2d(ndf * 4, ndf * 8, kw, stride=2, padding=pw))
self.layer5 = norm_layer(nn.Conv2d(ndf * 8, ndf * 16, kw, stride=2, padding=pw))
self.actvn = nn.LeakyReLU(0.2, False)
self.opt = opt
def __init__(self, opt):
super().__init__()
FeatGenerator.modify_commandline_options(opt)
input_nc = opt['FG_c']
keep_conv = opt['FG_keep']
norm_layer = get_nonspade_norm_layer(opt, opt['norm_FG'])
activation = nn.ReLU(True)
model = []
for i in range(keep_conv):
model += [
ResnetBlock(input_nc,
norm_layer=norm_layer,
activation=activation,
kernel_size=opt['FG_resnet_kernel_size'])
]
model += [nn.Conv2d(input_nc, input_nc, 1, 1, 0)]
self.model = nn.Sequential(*model)
def __init__(self, opt, n_frames_G):
super().__init__()
self.opt = opt
input_nc = (opt.label_nc if opt.label_nc != 0 else opt.input_nc) * n_frames_G
input_nc += opt.output_nc * (n_frames_G - 1)
nf = opt.nff
n_blocks = opt.n_blocks_F
n_downsample_F = opt.n_downsample_F
self.flow_multiplier = opt.flow_multiplier
nf_max = 1024
ch = [min(nf_max, nf * (2 ** i)) for i in range(n_downsample_F + 1)]
norm = opt.norm_F
norm_layer = get_nonspade_norm_layer(opt, norm)
activation = nn.LeakyReLU(0.2)
down_flow = [norm_layer(nn.Conv2d(input_nc, nf, kernel_size=3, padding=1)), activation]
for i in range(n_downsample_F):
down_flow += [norm_layer(nn.Conv2d(ch[i], ch[i+1], kernel_size=3, padding=1, stride=2)), activation]
### resnet blocks
res_flow = []
ch_r = min(nf_max, nf * (2**n_downsample_F))
for i in range(n_blocks):
res_flow += [SPADEResnetBlock(ch_r, ch_r, norm=norm)]
### upsample
up_flow = []
for i in reversed(range(n_downsample_F)):
up_flow += [nn.Upsample(scale_factor=2), norm_layer(nn.Conv2d(ch[i+1], ch[i], kernel_size=3, padding=1)), activation]
conv_flow = [nn.Conv2d(nf, 2, kernel_size=3, padding=1)]
conv_mask = [nn.Conv2d(nf, 1, kernel_size=3, padding=1), nn.Sigmoid()]
self.down_flow = nn.Sequential(*down_flow)
self.res_flow = nn.Sequential(*res_flow)
self.up_flow = nn.Sequential(*up_flow)
self.conv_flow = nn.Sequential(*conv_flow)
self.conv_mask = nn.Sequential(*conv_mask)
def __init__(self, opt):
# TODO: kernel=4, concat noise, or change architecture to vgg feature pyramid
super().__init__()
self.opt = opt
kw = 3
pw = int(np.ceil((kw - 1.0) / 2))
ndf = opt.ngf
norm_layer = get_nonspade_norm_layer(opt, opt.norm_E)
self.layer1 = norm_layer(nn.Conv2d(opt.spade_ic, ndf, kw, stride=1, padding=pw))
self.layer2 = norm_layer(nn.Conv2d(ndf * 1, ndf * 2, opt.adaptor_kernel, stride=2, padding=pw))
self.layer3 = norm_layer(nn.Conv2d(ndf * 2, ndf * 4, kw, stride=1, padding=pw))
if opt.warp_stride == 2:
self.layer4 = norm_layer(nn.Conv2d(ndf * 4, ndf * 8, kw, stride=1, padding=pw))
else:
self.layer4 = norm_layer(nn.Conv2d(ndf * 4, ndf * 8, opt.adaptor_kernel, stride=2, padding=pw))
self.layer5 = norm_layer(nn.Conv2d(ndf * 8, ndf * 8, kw, stride=1, padding=pw))
self.actvn = nn.LeakyReLU(0.2, False)
self.opt = opt
nf = opt.ngf
self.head_0 = SPADEResnetBlock(8 * nf, 8 * nf, opt, use_se=opt.adaptor_se)
if opt.adaptor_nonlocal:
self.attn = Attention(8 * nf, False)
self.G_middle_0 = SPADEResnetBlock(8 * nf, 8 * nf, opt, use_se=opt.adaptor_se)
self.G_middle_1 = SPADEResnetBlock(8 * nf, 4 * nf, opt, use_se=opt.adaptor_se)
if opt.adaptor_res_deeper:
self.deeper0 = SPADEResnetBlock(4 * nf, 4 * nf, opt)
if opt.dilation_conv:
self.deeper1 = SPADEResnetBlock(4 * nf, 4 * nf, opt, dilation=2)
self.deeper2 = SPADEResnetBlock(4 * nf, 4 * nf, opt, dilation=4)
self.degridding0 = norm_layer(nn.Conv2d(ndf * 4, ndf * 4, 3, stride=1, padding=2, dilation=2))
self.degridding1 = norm_layer(nn.Conv2d(ndf * 4, ndf * 4, 3, stride=1, padding=1))
else:
self.deeper1 = SPADEResnetBlock(4 * nf, 4 * nf, opt)
self.deeper2 = SPADEResnetBlock(4 * nf, 4 * nf, opt)
def __init__(self, opt):
super().__init__()
self.opt = opt
NLayerDiscriminator.modify_commandline_options(self.opt)
kw = 4
padw = int(np.ceil((kw - 1.0) / 2))
nf = opt['ndf']
input_nc = opt['output_nc']
norm_layer = get_nonspade_norm_layer(opt, opt['norm_D'])
sequence = [[
nn.Conv2d(input_nc, nf, kernel_size=kw, stride=2, padding=padw),
nn.LeakyReLU(0.2, False)
]]
for n in range(1, opt['n_layers_D']):
nf_prev = nf
nf = min(nf * 2, 512)
stride = 1 if n == opt['n_layers_D'] - 1 else 2
sequence += [[
norm_layer(
nn.Conv2d(nf_prev,
nf,
kernel_size=kw,
stride=stride,
padding=padw)),
nn.LeakyReLU(0.2, False)
]]
sequence += [[
nn.Conv2d(nf, 1, kernel_size=kw, stride=1, padding=padw)
]]
# We divide the layers into groups to extract intermediate layer outputs
for n in range(len(sequence)):
self.add_module('model' + str(n), nn.Sequential(*sequence[n]))
def __init__(self, opt, stage1=False):
super().__init__()
self.opt = opt
self.stage1 = stage1
kw = 4
#padw = int(np.ceil((kw - 1.0) / 2))
padw = int((kw - 1.0) / 2)
nf = opt.ndf
input_nc = self.compute_D_input_nc(opt)
norm_layer = get_nonspade_norm_layer(opt, opt.norm_D)
sequence = [[
nn.Conv2d(input_nc, nf, kernel_size=kw, stride=2, padding=padw),
nn.LeakyReLU(0.2, False)
]]
for n in range(1, opt.n_layers_D):
nf_prev = nf
nf = min(nf * 2, 512)
stride = 1 if n == opt.n_layers_D - 1 else 2
if (((not stage1) and opt.use_attention) or
(stage1
and opt.use_attention_st1)) and n == opt.n_layers_D - 1:
self.attn = Attention(nf_prev, 'spectral' in opt.norm_D)
if n == opt.n_layers_D - 1 and (not stage1):
dec = []
nc_dec = nf_prev
for _ in range(opt.n_layers_D - 1):
dec += [
nn.Upsample(scale_factor=2),
norm_layer(
nn.Conv2d(nc_dec,
int(nc_dec // 2),
kernel_size=3,
stride=1,
padding=1)),
nn.LeakyReLU(0.2, False)
]
nc_dec = int(nc_dec // 2)
dec += [
nn.Conv2d(nc_dec,
opt.semantic_nc,
kernel_size=3,
stride=1,
padding=1)
]
self.dec = nn.Sequential(*dec)
sequence += [[
norm_layer(
nn.Conv2d(nf_prev,
nf,
kernel_size=kw,
stride=stride,
padding=padw)),
nn.LeakyReLU(0.2, False)
]]
sequence += [[
nn.Conv2d(nf, 1, kernel_size=kw, stride=1, padding=padw)
]]
if opt.D_cam > 0:
mult = min(2**(opt.n_layers_D - 1), 8)
if opt.eqlr_sn:
self.gap_fc = equal_lr(nn.Linear(opt.ndf * mult, 1,
bias=False))
self.gmp_fc = equal_lr(nn.Linear(opt.ndf * mult, 1,
bias=False))
else:
self.gap_fc = nn.utils.spectral_norm(
nn.Linear(opt.ndf * mult, 1, bias=False))
self.gmp_fc = nn.utils.spectral_norm(
nn.Linear(opt.ndf * mult, 1, bias=False))
self.conv1x1 = nn.Conv2d(opt.ndf * mult * 2,
opt.ndf * mult,
kernel_size=1,
stride=1,
bias=True)
self.leaky_relu = nn.LeakyReLU(0.2, True)
# We divide the layers into groups to extract intermediate layer outputs
for n in range(len(sequence)):
self.add_module('model' + str(n), nn.Sequential(*sequence[n]))
def __init__(self, opt, n_addition_layer=3):
super().__init__()
assert n_addition_layer > 0
self.opt = opt
max_dim = 512
self.n_add_layer = n_addition_layer
self.backboneNet = Vgg16backbone(block_num=min(5,
2 + n_addition_layer),
requires_grad=False)
activation = nn.LeakyReLU(0.2, False)
norm_layer = get_nonspade_norm_layer(opt, opt.norm_D)
nf = opt.ndf
input_nc = self.compute_D_input_nc(opt)
sequence = [
nn.Conv2d(input_nc, nf, kernel_size=3, stride=2, padding=1),
activation,
]
model_1 = nn.Sequential(*sequence)
sequence = [
norm_layer(
nn.Conv2d(nf, nf * 2, kernel_size=3, stride=2, padding=1)),
activation,
]
model_2 = nn.Sequential(*sequence)
sequence = [
norm_layer(
nn.Conv2d(nf * 2, nf * 4, kernel_size=3, stride=2, padding=1)),
activation,
]
model_3 = nn.Sequential(*sequence)
self.common = nn.ModuleList([model_1, model_2, model_3])
for i in range(n_addition_layer):
nf = max_dim if i > 0 else max_dim // 2
sequence = [
norm_layer(
nn.Conv2d(nf, nf, kernel_size=3, stride=1, padding=1)),
activation,
]
model = nn.Sequential(*sequence)
setattr(self, 'pre' + str(i + 1), model)
nf = 512 if i == 0 else 1024
sequence = [
norm_layer(
nn.Conv2d(nf, 512, kernel_size=3, stride=1, padding=1)),
activation,
]
model = nn.Sequential(*sequence)
setattr(self, 'gcb' + str(i + 1), model)
for i in range(n_addition_layer):
if i != n_addition_layer - 1:
nf = 512
sequence = [
InceptionBlock(nf,
512,
stride=2,
norm_layer=norm_layer,
activation=activation),
]
model = nn.Sequential(*sequence)
setattr(self, 'conv' + str(i + 1) + '_line', model)
nf = 512
sequence = [
norm_layer(
nn.Conv2d(nf, nf, kernel_size=3, stride=1, padding=1)),
activation,
nn.Conv2d(nf, 1, kernel_size=5, stride=1, padding=2),
nn.Upsample(scale_factor=2**i),
]
model = nn.Sequential(*sequence)
setattr(self, 'conv' + str(i + 1) + '_out', model)
def __init__(self, opt):
super().__init__()
input_nc = 3
# print("xxxxx")
# print(opt.norm_G)
norm_layer = get_nonspade_norm_layer(opt, opt.norm_G)
activation = nn.ReLU(False)
model = []
# initial conv
model += [
nn.ReflectionPad2d(opt.resnet_initial_kernel_size // 2),
norm_layer(
nn.Conv2d(
input_nc,
opt.ngf,
kernel_size=opt.resnet_initial_kernel_size,
padding=0,
)
),
activation,
]
# downsample
mult = 1
for i in range(opt.resnet_n_downsample):
model += [
norm_layer(
nn.Conv2d(
opt.ngf * mult,
opt.ngf * mult * 2,
kernel_size=3,
stride=2,
padding=1,
)
),
activation,
]
mult *= 2
# resnet blocks
for i in range(opt.resnet_n_blocks):
model += [
ResnetBlock(
opt.ngf * mult,
norm_layer=norm_layer,
activation=activation,
kernel_size=opt.resnet_kernel_size,
)
]
# upsample
for i in range(opt.resnet_n_downsample):
nc_in = int(opt.ngf * mult)
nc_out = int((opt.ngf * mult) / 2)
model += [
norm_layer(
nn.ConvTranspose2d(
nc_in,
nc_out,
kernel_size=3,
stride=2,
padding=1,
output_padding=1,
)
),
activation,
]
mult = mult // 2
# final output conv
model += [
nn.ReflectionPad2d(3),
nn.Conv2d(nc_out, opt.output_nc, kernel_size=7, padding=0),
nn.Tanh(),
]
self.model = nn.Sequential(*model)
def __init__(self, opt):
super().__init__()
kw = 3
self.opt = opt
pw = int(np.ceil((kw - 1.0) / 2))
if opt.dataset_mode == 'cityscapes':
ndf = 350
elif opt.dataset_mode == 'ade20k':
ndf = 151 * 4
elif opt.dataset_mode == 'deepfashion':
ndf = 256
norm_layer = get_nonspade_norm_layer(opt, opt.norm_E)
self.layer1 = norm_layer(
nn.Conv2d(self.opt.semantic_nc * 3,
ndf,
kw,
stride=2,
padding=pw,
groups=self.opt.semantic_nc))
self.layer2 = norm_layer(
nn.Conv2d(ndf * 1,
ndf * 2,
kw,
stride=2,
padding=pw,
groups=self.opt.semantic_nc))
self.layer3 = norm_layer(
nn.Conv2d(ndf * 2,
ndf * 4,
kw,
stride=2,
padding=pw,
groups=self.opt.semantic_nc))
self.layer4 = norm_layer(
nn.Conv2d(ndf * 4,
ndf * 8,
kw,
stride=2,
padding=pw,
groups=self.opt.semantic_nc))
self.layer5 = norm_layer(
nn.Conv2d(ndf * 8,
ndf * 8,
kw,
stride=2,
padding=pw,
groups=self.opt.semantic_nc))
if opt.crop_size >= 256:
self.layer6 = norm_layer(
nn.Conv2d(ndf * 8,
ndf * 8,
kw,
stride=2,
padding=pw,
groups=self.opt.semantic_nc))
self.so = s0 = 4
self.fc_mu = nn.Conv2d(ndf * 8,
8 * self.opt.semantic_nc,
stride=1,
kernel_size=3,
padding=1,
groups=self.opt.semantic_nc)
self.fc_var = nn.Conv2d(ndf * 8,
8 * self.opt.semantic_nc,
stride=1,
kernel_size=3,
padding=1,
groups=self.opt.semantic_nc)
self.actvn = nn.LeakyReLU(0.2, False)
def __init__(self, opt):
super().__init__()
Pix2PixHDGenerator.modify_commandline_options(opt)
input_nc = opt['input_nc']
norm_layer = get_nonspade_norm_layer(opt, opt['norm_G'])
activation = nn.ReLU(False)
model = []
# initial conv
model += [
nn.ReflectionPad2d(opt['resnet_initial_kernel_size'] // 2),
norm_layer(
nn.Conv2d(input_nc,
opt['ngf'],
kernel_size=opt['resnet_initial_kernel_size'],
padding=0)), activation
]
# downsample
mult = 1
for i in range(opt['resnet_n_downsample']):
model += [
norm_layer(
nn.Conv2d(opt['ngf'] * mult,
opt['ngf'] * mult * 2,
kernel_size=3,
stride=2,
padding=1)), activation
]
mult *= 2
# resnet blocks
for i in range(opt['resnet_n_blocks']):
model += [
ResnetBlock(opt['ngf'] * mult,
norm_layer=norm_layer,
activation=activation,
kernel_size=opt['resnet_kernel_size'])
]
# upsample
for i in range(opt['resnet_n_upsample']):
nc_in = int(opt['ngf'] * mult)
nc_out = int((opt['ngf'] * mult) / 2)
model += [
norm_layer(
nn.ConvTranspose2d(nc_in,
nc_out,
kernel_size=3,
stride=2,
padding=1,
output_padding=1)), activation
]
# model += [
# nn.Upsample(scale_factor=2, mode='bilinear'),
# nn.ReflectionPad2d(1),
# norm_layer(
# nn.Conv2d(nc_in,
# nc_out,
# kernel_size=3,
# stride=1,
# padding=0))
# ]
mult = mult // 2
# final output conv
model += [
nn.ReflectionPad2d(3),
nn.Conv2d(nc_out, opt['output_nc'], kernel_size=7, padding=0),
nn.Tanh()
]
self.model = nn.Sequential(*model)
def __init__(self, opt):
super().__init__()
self.opt = opt
output_nc = 3
label_nc = opt.label_nc
input_nc = label_nc + (1 if opt.contain_dontcare_label else
0) + (0 if opt.no_instance else 1)
if opt.mix_input_gen:
input_nc += 4
norm_layer = get_nonspade_norm_layer(opt, 'instance')
activation = nn.ReLU(False)
# initial block
self.init_block = nn.Sequential(*[
nn.ReflectionPad2d(opt.resnet_initial_kernel_size // 2),
norm_layer(
nn.Conv2d(input_nc,
opt.ngf,
kernel_size=opt.resnet_initial_kernel_size,
padding=0)), activation
])
# Downsampling blocks
self.downlayers = nn.ModuleList()
mult = 1
for i in range(opt.resnet_n_downsample):
self.downlayers.append(
nn.Sequential(*[
norm_layer(
nn.Conv2d(opt.ngf * mult,
opt.ngf * mult * 2,
kernel_size=3,
stride=2,
padding=1)), activation
]))
mult *= 2
# Semantic core blocks
self.resnet_core = nn.ModuleList()
if opt.wide:
self.resnet_core += [
ResnetBlock(opt.ngf * mult,
dim2=opt.ngf * mult * 2,
norm_layer=norm_layer,
activation=activation,
kernel_size=opt.resnet_kernel_size)
]
mult *= 2
else:
self.resnet_core += [
ResnetBlock(opt.ngf * mult,
norm_layer=norm_layer,
activation=activation,
kernel_size=opt.resnet_kernel_size)
]
for i in range(opt.resnet_n_blocks - 1):
self.resnet_core += [
ResnetBlock(opt.ngf * mult,
norm_layer=norm_layer,
activation=activation,
kernel_size=opt.resnet_kernel_size,
dilation=2)
]
self.spade_core = nn.ModuleList()
for i in range(opt.spade_n_blocks - 1):
self.spade_core += [
SPADEResnetBlock(opt.ngf * mult,
opt.ngf * mult,
opt,
dilation=2)
]
if opt.wide:
self.spade_core += [
SPADEResnetBlock(
opt.ngf * mult *
(2 if not self.opt.no_skip_connections else 1),
opt.ngf * mult // 2, opt)
]
mult //= 2
else:
self.spade_core += [
SPADEResnetBlock(
opt.ngf * mult *
(2 if not self.opt.no_skip_connections else 1),
opt.ngf * mult, opt)
]
# Upsampling blocks
self.uplayers = nn.ModuleList()
for i in range(opt.resnet_n_downsample):
self.uplayers.append(
SPADEResnetBlock(
mult * opt.ngf *
(3 if not self.opt.no_skip_connections else 2) // 2,
opt.ngf * mult // 2, opt))
mult //= 2
final_nc = opt.ngf
self.conv_img = nn.Conv2d(
(input_nc +
final_nc) if not self.opt.no_skip_connections else final_nc,
output_nc,
3,
padding=1)
self.up = nn.Upsample(scale_factor=2)
def __init__(self, opt):
super().__init__()
self.opt = opt
nf = opt.ndf
input_nc = 3
label_nc = opt.label_nc + (1 if opt.contain_dontcare_label else 0) \
+ (0 if opt.no_instance else 1)
norm_layer = get_nonspade_norm_layer(opt, opt.norm_D)
# bottom-up pathway
self.enc1 = nn.Sequential(
norm_layer(
nn.Conv2d(input_nc, nf, kernel_size=3, stride=2, padding=1),
opt), nn.LeakyReLU(0.2, True))
self.enc2 = nn.Sequential(
norm_layer(
nn.Conv2d(nf, nf * 2, kernel_size=3, stride=2, padding=1),
opt), nn.LeakyReLU(0.2, True))
self.enc3 = nn.Sequential(
norm_layer(
nn.Conv2d(nf * 2, nf * 4, kernel_size=3, stride=2, padding=1),
opt), nn.LeakyReLU(0.2, True))
self.enc4 = nn.Sequential(
norm_layer(
nn.Conv2d(nf * 4, nf * 8, kernel_size=3, stride=2, padding=1),
opt), nn.LeakyReLU(0.2, True))
self.enc5 = nn.Sequential(
norm_layer(
nn.Conv2d(nf * 8, nf * 8, kernel_size=3, stride=2, padding=1),
opt), nn.LeakyReLU(0.2, True))
# top-down pathway
self.lat2 = nn.Sequential(
norm_layer(nn.Conv2d(nf * 2, nf * 4, kernel_size=1), opt),
nn.LeakyReLU(0.2, True))
self.lat3 = nn.Sequential(
norm_layer(nn.Conv2d(nf * 4, nf * 4, kernel_size=1), opt),
nn.LeakyReLU(0.2, True))
self.lat4 = nn.Sequential(
norm_layer(nn.Conv2d(nf * 8, nf * 4, kernel_size=1), opt),
nn.LeakyReLU(0.2, True))
self.lat5 = nn.Sequential(
norm_layer(nn.Conv2d(nf * 8, nf * 4, kernel_size=1), opt),
nn.LeakyReLU(0.2, True))
# upsampling
self.up = nn.Upsample(scale_factor=2, mode='bilinear')
# final layers
self.final2 = nn.Sequential(
norm_layer(nn.Conv2d(nf * 4, nf * 2, kernel_size=3, padding=1),
opt), nn.LeakyReLU(0.2, True))
self.final3 = nn.Sequential(
norm_layer(nn.Conv2d(nf * 4, nf * 2, kernel_size=3, padding=1),
opt), nn.LeakyReLU(0.2, True))
self.final4 = nn.Sequential(
norm_layer(nn.Conv2d(nf * 4, nf * 2, kernel_size=3, padding=1),
opt), nn.LeakyReLU(0.2, True))
# true/false prediction and semantic alignment prediction
self.tf = nn.Conv2d(nf * 2, 1, kernel_size=1)
self.seg = nn.Conv2d(nf * 2, nf * 2, kernel_size=1)
self.embedding = nn.Conv2d(label_nc, nf * 2, kernel_size=1)
