Spaces:
Runtime error
Runtime error
| import argparse | |
| import math | |
| import os | |
| import pickle | |
| import torch | |
| import torchvision | |
| from torch import optim | |
| from tqdm import tqdm | |
| from StyleCLIP.criteria.clip_loss import CLIPLoss | |
| from StyleCLIP.models.stylegan2.model import Generator | |
| import clip | |
| from StyleCLIP.utils import ensure_checkpoint_exists | |
| def get_lr(t, initial_lr, rampdown=0.25, rampup=0.05): | |
| lr_ramp = min(1, (1 - t) / rampdown) | |
| lr_ramp = 0.5 - 0.5 * math.cos(lr_ramp * math.pi) | |
| lr_ramp = lr_ramp * min(1, t / rampup) | |
| return initial_lr * lr_ramp | |
| def main(args, use_old_G): | |
| ensure_checkpoint_exists(args.ckpt) | |
| text_inputs = torch.cat([clip.tokenize(args.description)]).cuda() | |
| os.makedirs(args.results_dir, exist_ok=True) | |
| new_generator_path = f'/disk2/danielroich/Sandbox/stylegan2_ada_pytorch/checkpoints/model_{args.run_id}_{args.image_name}.pt' | |
| old_generator_path = '/disk2/danielroich/Sandbox/pretrained_models/ffhq.pkl' | |
| if not use_old_G: | |
| with open(new_generator_path, 'rb') as f: | |
| G = torch.load(f).cuda().eval() | |
| else: | |
| with open(old_generator_path, 'rb') as f: | |
| G = pickle.load(f)['G_ema'].cuda().eval() | |
| if args.latent_path: | |
| latent_code_init = torch.load(args.latent_path).cuda() | |
| elif args.mode == "edit": | |
| latent_code_init_not_trunc = torch.randn(1, 512).cuda() | |
| with torch.no_grad(): | |
| latent_code_init = G.mapping(latent_code_init_not_trunc, None) | |
| latent = latent_code_init.detach().clone() | |
| latent.requires_grad = True | |
| clip_loss = CLIPLoss(args) | |
| optimizer = optim.Adam([latent], lr=args.lr) | |
| pbar = tqdm(range(args.step)) | |
| for i in pbar: | |
| t = i / args.step | |
| lr = get_lr(t, args.lr) | |
| optimizer.param_groups[0]["lr"] = lr | |
| img_gen = G.synthesis(latent, noise_mode='const') | |
| c_loss = clip_loss(img_gen, text_inputs) | |
| if args.mode == "edit": | |
| l2_loss = ((latent_code_init - latent) ** 2).sum() | |
| loss = c_loss + args.l2_lambda * l2_loss | |
| else: | |
| loss = c_loss | |
| optimizer.zero_grad() | |
| loss.backward() | |
| optimizer.step() | |
| pbar.set_description( | |
| ( | |
| f"loss: {loss.item():.4f};" | |
| ) | |
| ) | |
| if args.save_intermediate_image_every > 0 and i % args.save_intermediate_image_every == 0: | |
| with torch.no_grad(): | |
| img_gen = G.synthesis(latent, noise_mode='const') | |
| torchvision.utils.save_image(img_gen, | |
| f"/disk2/danielroich/Sandbox/StyleCLIP/results/inference_results/{str(i).zfill(5)}.png", | |
| normalize=True, range=(-1, 1)) | |
| if args.mode == "edit": | |
| with torch.no_grad(): | |
| img_orig = G.synthesis(latent_code_init, noise_mode='const') | |
| final_result = torch.cat([img_orig, img_gen]) | |
| else: | |
| final_result = img_gen | |
| return final_result | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--description", type=str, default="a person with purple hair", | |
| help="the text that guides the editing/generation") | |
| parser.add_argument("--ckpt", type=str, default="../pretrained_models/stylegan2-ffhq-config-f.pt", | |
| help="pretrained StyleGAN2 weights") | |
| parser.add_argument("--stylegan_size", type=int, default=1024, help="StyleGAN resolution") | |
| parser.add_argument("--lr_rampup", type=float, default=0.05) | |
| parser.add_argument("--lr", type=float, default=0.1) | |
| parser.add_argument("--step", type=int, default=300, help="number of optimization steps") | |
| parser.add_argument("--mode", type=str, default="edit", choices=["edit", "free_generation"], | |
| help="choose between edit an image an generate a free one") | |
| parser.add_argument("--l2_lambda", type=float, default=0.008, | |
| help="weight of the latent distance (used for editing only)") | |
| parser.add_argument("--latent_path", type=str, default=None, | |
| help="starts the optimization from the given latent code if provided. Otherwose, starts from" | |
| "the mean latent in a free generation, and from a random one in editing. " | |
| "Expects a .pt format") | |
| parser.add_argument("--truncation", type=float, default=0.7, | |
| help="used only for the initial latent vector, and only when a latent code path is" | |
| "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") | |
| args = parser.parse_args() | |
| result_image = main(args) | |
| torchvision.utils.save_image(result_image.detach().cpu(), os.path.join(args.results_dir, "final_result.jpg"), | |
| normalize=True, scale_each=True, range=(-1, 1)) | |