Skip to content
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 9 additions & 1 deletion optimization/run_optimization.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import argparse
import math
import os
import numpy as np

import torch
import torchvision
Expand Down Expand Up @@ -80,14 +81,20 @@ def main(args):
with torch.no_grad():
img_gen, _ = g_ema([latent], input_is_latent=True, randomize_noise=False)

torchvision.utils.save_image(img_gen, f"results/{str(i).zfill(5)}.png", normalize=True, range=(-1, 1))
torchvision.utils.save_image(img_gen, f"{args.results_dir}/{str(i).zfill(5)}.png", normalize=True, range=(-1, 1))

if(args.save_vector):
np.savez(f'{args.results_dir}/w-{str(i).zfill(5)}.npz', w=latent.unsqueeze(0).detach().cpu().numpy())


if args.mode == "edit":
with torch.no_grad():
img_orig, _ = g_ema([latent_code_init], input_is_latent=True, randomize_noise=False)

final_result = torch.cat([img_orig, img_gen])
else:
if(args.save_vector):
np.savez(f'{args.results_dir}/projected_w.npz', w=latent.unsqueeze(0).detach().cpu().numpy())
final_result = img_gen

return final_result
Expand All @@ -111,6 +118,7 @@ def main(args):
"not provided")
parser.add_argument("--save_intermediate_image_every", type=int, default=20, help="if > 0 then saves intermidate results during the optimization")
parser.add_argument("--results_dir", type=str, default="results")
parser.add_argument('--save_vector', action='store_true')

args = parser.parse_args()

Expand Down