67 lines
1.9 KiB
Python
67 lines
1.9 KiB
Python
import torch
|
|
from torch import nn
|
|
|
|
from . import networks
|
|
|
|
|
|
class AesFA_test(nn.Module):
|
|
def __init__(
|
|
self, alpha_in, alpha_out, style_kernel, input_nc, nf, output_nc, freq_ratio
|
|
):
|
|
super(AesFA_test, self).__init__()
|
|
|
|
self.netE = networks.define_network(
|
|
"Encoder",
|
|
alpha_in,
|
|
alpha_out,
|
|
style_kernel,
|
|
input_nc,
|
|
nf,
|
|
output_nc,
|
|
freq_ratio,
|
|
)
|
|
self.netS = networks.define_network(
|
|
"Encoder",
|
|
alpha_in,
|
|
alpha_out,
|
|
style_kernel,
|
|
input_nc,
|
|
nf,
|
|
output_nc,
|
|
freq_ratio,
|
|
)
|
|
self.netG = networks.define_network(
|
|
"Generator",
|
|
alpha_in,
|
|
alpha_out,
|
|
style_kernel,
|
|
input_nc,
|
|
nf,
|
|
output_nc,
|
|
freq_ratio,
|
|
)
|
|
|
|
def forward(self, real_A, real_B):
|
|
with torch.no_grad():
|
|
content_A = self.netE.forward_test(real_A, "content")
|
|
style_B = self.netS.forward_test(real_B, "style")
|
|
# if freq:
|
|
# trs_AtoB, trs_AtoB_high, trs_AtoB_low = self.netG(content_A, style_B)
|
|
# end = time.time()
|
|
# during = end - start
|
|
# return trs_AtoB, trs_AtoB_high, trs_AtoB_low, during
|
|
# else:
|
|
trs_AtoB = self.netG.forward_test(content_A, style_B)
|
|
return trs_AtoB
|
|
|
|
def style_blending(self, real_A, real_B_1, real_B_2):
|
|
with torch.no_grad():
|
|
content_A = self.netE.forward_test(real_A, "content")
|
|
style_B1_h = self.netS.forward_test(real_B_1, "style")[0]
|
|
style_B2_l = self.netS.forward_test(real_B_2, "style")[1]
|
|
style_B = style_B1_h, style_B2_l
|
|
|
|
trs_AtoB = self.netG.forward_test(content_A, style_B)
|
|
|
|
return trs_AtoB
|