def optimize_latents():
print("Optimizing Latents.")
synthesizer = StyleGANGenerator(args.model_type).model.synthesis
latent_optimizer = LatentOptimizer(synthesizer, args.vgg_layer)
# Optimize only the dlatents.
for param in latent_optimizer.parameters():
param.requires_grad_(False)
if args.video or args.save_optimized_image:
# Hook, saves an image during optimization to be used to create video.
generated_image_hook = GeneratedImageHook(
latent_optimizer.post_synthesis_processing, args.save_frequency)
reference_image = load_images([args.image_path])
reference_image = torch.from_numpy(reference_image).cuda()
reference_image = latent_optimizer.vgg_processing(reference_image)
reference_features = latent_optimizer.vgg16(reference_image).detach()
reference_image = reference_image.detach()
if args.use_latent_finder:
image_to_latent = ImageToLatent().cuda()
image_to_latent.load_state_dict(torch.load(args.image_to_latent_path))
image_to_latent.eval()
latents_to_be_optimized = image_to_latent(reference_image)
latents_to_be_optimized = latents_to_be_optimized.detach().cuda(
).requires_grad_(True)
else:
latents_to_be_optimized = torch.zeros(
(1, 18, 512)).cuda().requires_grad_(True)
criterion = LatentLoss()
optimizer = torch.optim.SGD([latents_to_be_optimized],
lr=args.learning_rate)
progress_bar = tqdm(range(args.iterations))
for step in progress_bar:
optimizer.zero_grad()
generated_image_features = latent_optimizer(latents_to_be_optimized)
loss = criterion(generated_image_features, reference_features)
loss.backward()
loss = loss.item()
optimizer.step()
progress_bar.set_description("Step: {}, Loss: {}".format(step, loss))
optimized_dlatents = latents_to_be_optimized.detach().cpu().numpy()
np.save(args.dlatent_path, optimized_dlatents)
if args.video:
images_to_video(generated_image_hook.get_images(), args.video_path)
if args.save_optimized_image:
save_image(generated_image_hook.last_image, args.optimized_image_path)
def optimize_latents():
print("Optimizing Latents.")
generator = StyledGenerator(512).to(device)
generator.load_state_dict(torch.load(args.path)['generator'])
generator.eval()
latent_optimizer = LatentOptimizer(generator, args.vgg_layer)
mean_style = get_mean_style(generator, device)
total = np.zeros((83 * 3, 512))
# Optimize only the dlatents.
for param in latent_optimizer.parameters():
param.requires_grad_(False)
if args.video or args.save_optimized_image:
# Hook, saves an image during optimization to be used to create video.
generated_image_hook = GeneratedImageHook(
latent_optimizer.post_synthesis_processing, args.save_frequency)
for i in range(3 * 83): #3 for each pictrue
iid = i % 3
path = int(i / 3)
iterations = int(200 * iid + 300)
image_path = './data/' + str(path) + '.jpg'
print(image_path)
reference_image = load_images([image_path])
reference_image = torch.from_numpy(reference_image).to(device)
reference_image = latent_optimizer.vgg_processing(
reference_image) #normalize
reference_features = latent_optimizer.vgg16(
reference_image).detach() #vgg
reference_image = reference_image.detach()
if args.use_latent_finder:
image_to_latent = ImageToLatent().cuda()
image_to_latent.load_state_dict(
torch.load(args.image_to_latent_path))
image_to_latent.eval()
latents_to_be_optimized = image_to_latent(reference_image)
latents_to_be_optimized = latents_to_be_optimized.detach().cuda(
).requires_grad_(True)
else:
latents_to_be_optimized = torch.zeros(
(1, 512)).cuda().requires_grad_(True)
criterion = LatentLoss()
optimizer = torch.optim.SGD([latents_to_be_optimized],
lr=args.learning_rate)
progress_bar = tqdm(range(iterations))
for step in progress_bar:
optimizer.zero_grad()
generated_image_features = latent_optimizer(
latents_to_be_optimized, mean_style, i)
#print(latents_to_be_optimized)
loss = criterion(generated_image_features, reference_features)
loss.backward()
loss = loss.item()
optimizer.step()
progress_bar.set_description("Step: {}, Loss: {}".format(
step, loss))
optimized_dlatents = latents_to_be_optimized.detach().cpu().numpy()
total[i] = optimized_dlatents[0]
np.save(args.dlatent_path, total)
def optimize_latents():
print("Optimizing Latents.")
generator = StyledGenerator(512).to(device)
generator.load_state_dict(torch.load(args.path)['generator'])
generator.eval()
latent_optimizer = LatentOptimizer(generator, args.vgg_layer)
mean_style = get_mean_style(generator, device)
# Optimize only the dlatents.
for param in latent_optimizer.parameters():
param.requires_grad_(False)
if args.video or args.save_optimized_image:
# Hook, saves an image during optimization to be used to create video.
generated_image_hook = GeneratedImageHook(
latent_optimizer.post_synthesis_processing, args.save_frequency)
reference_image = load_images([args.image_path])
reference_image = torch.from_numpy(reference_image).to(device)
reference_image = latent_optimizer.vgg_processing(
reference_image) #normalize
utils.save_image(reference_image,
'./reference.png',
nrow=1,
normalize=True,
range=(-1, 1))
reference_features = latent_optimizer.vgg16(reference_image).detach() #vgg
reference_image = reference_image.detach()
if args.use_latent_finder:
image_to_latent = ImageToLatent().cuda()
image_to_latent.load_state_dict(torch.load(args.image_to_latent_path))
image_to_latent.eval()
latents_to_be_optimized = image_to_latent(reference_image)
latents_to_be_optimized = latents_to_be_optimized.detach().cuda(
).requires_grad_(True)
else:
latents_to_be_optimized = torch.zeros(
(1, 512)).cuda().requires_grad_(True)
#print(latents_to_be_optimized.mean(),latents_to_be_optimized.std())
criterion = LatentLoss()
optimizer = torch.optim.RMSprop([latents_to_be_optimized],
lr=args.learning_rate,
weight_decay=0.02)
#print(latents_to_be_optimized.mean(),latents_to_be_optimized.std())
progress_bar = tqdm(range(args.iterations))
for step in progress_bar:
#print(latents_to_be_optimized)
optimizer.zero_grad()
generated_image_features = latent_optimizer(latents_to_be_optimized,
mean_style)
loss = criterion(generated_image_features, reference_features)
loss.backward()
loss = loss.item()
optimizer.step()
# if step==args.iterations:
# break
# with torch.no_grad():
# latents_to_be_optimized.add_(-latents_to_be_optimized.mean()+3e-2*torch.randn(1).to('cuda'))
# latents_to_be_optimized.div_(latents_to_be_optimized.std()+3e-2*torch.randn(1).to('cuda'))
#print(latents_to_be_optimized.mean(),latents_to_be_optimized.std())
progress_bar.set_description("Step: {}, Loss: {}".format(step, loss))
print(latents_to_be_optimized)
#latents_to_be_optimized=latent_optimizer.normalize(latents_to_be_optimized)
#print(latents_to_be_optimized.mean(),latents_to_be_optimized.std())
optimized_dlatents = latents_to_be_optimized.detach().cpu().numpy()
np.save(args.dlatent_path, optimized_dlatents)
if args.video:
images_to_video(generated_image_hook.get_images(), args.video_path)
if args.save_optimized_image:
save_image(generated_image_hook.last_image, args.optimized_image_path)