Spaces:
Running
on
T4
Running
on
T4
| import models.stylegan2.lpips as lpips | |
| from torch import autograd, optim | |
| from torchvision import transforms, utils | |
| from tqdm import tqdm | |
| import torch | |
| from scripts.align_all_parallel import align_face | |
| from utils.inference_utils import noise_regularize, noise_normalize_, get_lr, latent_noise, visualize | |
| def latent_optimization(frame, pspex, landmarkpredictor, step=500, device='cuda'): | |
| percept = lpips.PerceptualLoss( | |
| model="net-lin", net="vgg", use_gpu=device.startswith("cuda") | |
| ) | |
| transform = transforms.Compose([ | |
| transforms.ToTensor(), | |
| transforms.Normalize(mean=[0.5, 0.5, 0.5],std=[0.5,0.5,0.5]), | |
| ]) | |
| with torch.no_grad(): | |
| noise_sample = torch.randn(1000, 512, device=device) | |
| latent_out = pspex.decoder.style(noise_sample) | |
| latent_mean = latent_out.mean(0) | |
| latent_std = ((latent_out - latent_mean).pow(2).sum() / 1000) ** 0.5 | |
| y = transform(frame).unsqueeze(dim=0).to(device) | |
| I_ = align_face(frame, landmarkpredictor) | |
| I_ = transform(I_).unsqueeze(dim=0).to(device) | |
| wplus = pspex.encoder(I_) + pspex.latent_avg.unsqueeze(0) | |
| _, f = pspex.encoder(y, return_feat=True) | |
| latent_in = wplus.detach().clone() | |
| feat = [f[0].detach().clone(), f[1].detach().clone()] | |
| # wplus and f to optimize | |
| latent_in.requires_grad = True | |
| feat[0].requires_grad = True | |
| feat[1].requires_grad = True | |
| noises_single = pspex.decoder.make_noise() | |
| basic_height, basic_width = int(y.shape[2]*32/256), int(y.shape[3]*32/256) | |
| noises = [] | |
| for noise in noises_single: | |
| noises.append(noise.new_empty(y.shape[0], 1, max(basic_height, int(y.shape[2]*noise.shape[2]/256)), | |
| max(basic_width, int(y.shape[3]*noise.shape[2]/256))).normal_()) | |
| for noise in noises: | |
| noise.requires_grad = True | |
| init_lr=0.05 | |
| optimizer = optim.Adam(feat + noises, lr=init_lr) | |
| optimizer2 = optim.Adam([latent_in], lr=init_lr) | |
| noise_weight = 0.05 * 0.2 | |
| pbar = tqdm(range(step)) | |
| latent_path = [] | |
| for i in pbar: | |
| t = i / step | |
| lr = get_lr(t, init_lr) | |
| optimizer.param_groups[0]["lr"] = lr | |
| optimizer2.param_groups[0]["lr"] = get_lr(t, init_lr) | |
| noise_strength = latent_std * noise_weight * max(0, 1 - t / 0.75) ** 2 | |
| latent_n = latent_noise(latent_in, noise_strength.item()) | |
| y_hat, _ = pspex.decoder([latent_n], input_is_latent=True, randomize_noise=False, | |
| first_layer_feature=feat, noise=noises) | |
| batch, channel, height, width = y_hat.shape | |
| if height > y.shape[2]: | |
| factor = height // y.shape[2] | |
| y_hat = y_hat.reshape( | |
| batch, channel, height // factor, factor, width // factor, factor | |
| ) | |
| y_hat = y_hat.mean([3, 5]) | |
| p_loss = percept(y_hat, y).sum() | |
| n_loss = noise_regularize(noises) * 1e3 | |
| loss = p_loss + n_loss | |
| optimizer.zero_grad() | |
| optimizer2.zero_grad() | |
| loss.backward() | |
| optimizer.step() | |
| optimizer2.step() | |
| noise_normalize_(noises) | |
| ''' for visualization | |
| if (i + 1) % 100 == 0 or i == 0: | |
| viz = torch.cat((y_hat,y,y_hat-y), dim=3) | |
| visualize(torch.clamp(viz[0].cpu(),-1,1), 60) | |
| ''' | |
| pbar.set_description( | |
| ( | |
| f"perceptual: {p_loss.item():.4f}; noise regularize: {n_loss.item():.4f};" | |
| f" lr: {lr:.4f}" | |
| ) | |
| ) | |
| return latent_n, feat, noises, wplus, f |