def draw_concat(Gs, reals, NoiseAmp, in_s, mode, inject_level, opt):
if len(Gs) > 0:
if mode == 'rand':
count = 0
for G,real_curr,real_next,noise_amp in zip(Gs,reals,reals[1:],NoiseAmp):
print("cnt " + str(count))
if count == 0:
z = generate_noise2([1, 3, real_curr.shape[2], real_curr.shape[3]], device=opt.device)
G_z = in_s
else:
z = generate_noise2([1, opt.nc_z,real_curr.shape[2], real_curr.shape[3]], device=opt.device)
G_z = G_z[:,:,0:real_curr.shape[2],0:real_curr.shape[3]]
if count == inject_level:
print("Inject")
z_in = noise_amp*z + real_curr.cuda()
else:
z_in = noise_amp*z+G_z
if count > opt.switch_scale:
G_z = G(z_in.detach())
else:
G_z = G(z_in.detach(), G_z)
G_z = imresize2(G_z.detach(),1/opt.scale_factor,opt)
G_z = G_z[:,:,0:real_next.shape[2],0:real_next.shape[3]]
count += 1
return G_z
def draw_concat(Gs, reals, NoiseAmp, in_s, mode, opt):
if len(Gs) > 0:
if mode == 'rand':
count = 0
for G, real_curr, real_next, noise_amp in zip(
Gs, reals, reals[1:], NoiseAmp):
G = G.cuda()
if count == 0:
z = generate_noise2(
[1, 3, real_curr.shape[2], real_curr.shape[3]],
device=opt.device)
G_z = in_s
else:
z = generate_noise2(
[1, opt.nc_z, real_curr.shape[2], real_curr.shape[3]],
device=opt.device)
G_z = G_z[:, :, 0:real_curr.shape[2], 0:real_curr.shape[3]]
z_in = noise_amp * z + G_z
if count > opt.switch_scale:
G_z = G(z_in.detach())
else:
G_z = G(z_in.detach(), G_z)
G_z = imresize2(G_z.detach(), 1 / opt.scale_factor, opt)
G_z = G_z[:, :, 0:real_next.shape[2], 0:real_next.shape[3]]
count += 1
if mode == 'rec':
count = 0
for G, real_curr, real_next, noise_amp in zip(
Gs, reals, reals[1:], NoiseAmp):
G = G.cuda()
if count == 0:
size = list(real_curr.size())
#print(size)
G_z = generate_noise2(size, device=opt.device)
G_z = G_z[:, :, 0:real_curr.shape[2], 0:real_curr.shape[3]]
if count > opt.switch_scale:
G_z = G(G_z)
else:
G_z = G(G_z, G_z)
G_z = imresize2(G_z.detach(), 1 / opt.scale_factor, opt)
G_z = G_z[:, :, 0:real_next.shape[2], 0:real_next.shape[3]]
count += 1
return G_z
noise_amount_a = 0
noise_cnt_a = 0
noise_amount_b = 0
noise_cnt_b = 0
i = 0
for epoch in range(n_iters):
for data_a in data_loader_a:
real_a = data_a[len(Gs_a)].cuda()
real_b = data_b[len(Gs_b)].cuda()
# Create fake from noise
noise_ = generate_noise2([opt.bs, opt.nc_z, opt.nzx, opt.nzy],
device=opt.device)
if Gs_a == []:
noise_a = noise_
prev_a = torch.full([opt.bs, opt.nc_z, opt.nzx, opt.nzy],
0,
device=opt.device)
else:
prev_a = draw_concat(Gs_a, list(data_a), NoiseAmp_a, in_s,
'rand', opt)
noise_a = opt.noise_amp_a * noise_ + prev_a
noise_ = generate_noise2([opt.bs, opt.nc_z, opt.nzx, opt.nzy],
device=opt.device)
if Gs_b == []:
for ss in range(il, opt.stop_scale + 1):
print("inject " + str(il) + " - switch " + str(ss))
opt.nfc = min(
opt.nfc_init * pow(2, math.floor(opt.stop_scale / 4)), 128)
opt.min_nfc = min(
opt.min_nfc_init * pow(2, math.floor(opt.stop_scale / 4)), 128)
pad_image = int(((opt.ker_size - 1) * opt.num_layer) / 2)
opt.nzx = size_arr[len(Gs_a)]
opt.nzy = size_arr[len(Gs_a)]
real_a = data_a[len(Gs_a)].cuda()
real_b = data_b[len(Gs_b)].cuda()
noise_ = generate_noise2([1, opt.nc_z, opt.nzx, opt.nzy],
device=opt.device)
prev_a = draw_concat(Gs_a,
Gs_b,
list(data_a),
NoiseAmp_a,
in_s,
'rand',
il,
opt,
switch_level=ss)
noise_a = opt.noise_amp_a * noise_ + prev_a
noise_ = generate_noise2([1, opt.nc_z, opt.nzx, opt.nzy],
device=opt.device)
NoiseAmp_b = []
nfc_prev = 0
scale_num = 0
r_loss = nn.MSELoss()
dataset_a = Video_dataset(opt.video_dir, opt.num_images, opt.vid_ext, opt)
data_loader_a = DataLoader(dataset_a, shuffle=True, batch_size=1)
fixed_noise = []
size_arr = []
for ii in range(0, opt.stop_scale + 1, 1):
scale = math.pow(opt.scale_factor, opt.stop_scale - ii)
size_arr.append(math.ceil(scale * opt.img_size))
fixed_noise.append(generate_noise2([1, 3, size_arr[-1], size_arr[-1]], device=opt.device))
# print(size_arr)
opt.switch_scale = opt.stop_scale - opt.switch_res
opt.nzx = size_arr[0]
opt.nzy = size_arr[0]
in_s = torch.full([1, opt.nc_z, opt.nzx, opt.nzy], 0, device=opt.device)
scale_num = opt.stop_scale
if opt.a2b:
Gs_a, reals_a, NoiseAmp_a, Gs_b, reals_b, NoiseAmp_b = load_trained_pyramid_mix(opt)
else:
Gs_b, reals_b, NoiseAmp_b , Gs_b, reals_b, NoiseAmp_b = load_trained_pyramid_mix(opt)
