| | import torch |
| | from torch import nn |
| | import torch.optim as optim |
| | import torch.nn.functional as F |
| | from torch.utils.data.dataloader import DataLoader |
| | from torchvision import transforms |
| | from torchvision import utils as vutils |
| |
|
| | from models import Generator |
| | from utils import copy_G_params, load_params |
| |
|
| |
|
| |
|
| | def get_early_features(net, noise): |
| | with torch.no_grad(): |
| | feat_4 = net._init(noise) |
| | feat_8 = net._upsample_8(feat_4) |
| | feat_16 = net._upsample_16(feat_8) |
| | feat_32 = net._upsample_32(feat_16) |
| | feat_64 = net._upsample_64(feat_32) |
| | return feat_8, feat_16, feat_32, feat_64 |
| |
|
| | def get_late_features(net, feat_64, feat_8, feat_16, feat_32): |
| | with torch.no_grad(): |
| | feat_128 = net._upsample_128(feat_64) |
| | feat_128 = net._sle_128(feat_8, feat_128) |
| |
|
| | feat_256 = net._upsample_256(feat_128) |
| | feat_256 = net._sle_256(feat_16, feat_256) |
| |
|
| | feat_512 = net._upsample_512(feat_256) |
| | feat_512 = net._sle_512(feat_32, feat_512) |
| |
|
| | feat_1024 = net._upsample_1024(feat_512) |
| |
|
| | return net._out_1024(feat_1024) |
| |
|
| | def style_mix(model_name_or_path, bs, device): |
| | _in_channels = 256 |
| | im_size = 1024 |
| |
|
| | netG = Generator(in_channels=_in_channels, out_channels=3) |
| | netG = netG.from_pretrained(model_name_or_path, in_channels=256, out_channels=3) |
| | _ = netG.to(device) |
| | _ = netG.eval() |
| |
|
| | avg_param_G = copy_G_params(netG) |
| | load_params(netG, avg_param_G) |
| |
|
| | noise_a = torch.randn(bs, 256, 1, 1, device=device).to(device) |
| | noise_b = torch.randn(bs, 256, 1, 1, device=device).to(device) |
| |
|
| | feat_8_a, feat_16_a, feat_32_a, feat_64_a = get_early_features(netG, noise_a) |
| | feat_8_b, feat_16_b, feat_32_b, feat_64_b = get_early_features(netG, noise_b) |
| |
|
| | images_b = get_late_features(netG, feat_64_b, feat_8_b, feat_16_b, feat_32_b) |
| | images_a = get_late_features(netG, feat_64_a, feat_8_a, feat_16_a, feat_32_a) |
| |
|
| | imgs = [ torch.ones(1, 3, im_size, im_size) ] |
| |
|
| | imgs.append(images_b.cpu()) |
| | for i in range(bs): |
| | imgs.append(images_a[i].unsqueeze(0).cpu()) |
| | gimgs = get_late_features(netG, feat_64_a[i].unsqueeze(0).repeat(bs, 1, 1, 1), feat_8_b, feat_16_b, feat_32_b) |
| | imgs.append(gimgs.cpu()) |
| |
|
| | imgs = torch.cat(imgs) |
| | |
| |
|
| | return imgs |
| |
|
