huggan/FastGan

1import torch2from torch import nn3import torch.optim as optim4import torch.nn.functional as F5from torch.utils.data.dataloader import DataLoader6from torchvision import transforms7from torchvision import utils as vutils8 9from models import Generator10from utils import copy_G_params, load_params11 12 13 14def get_early_features(net, noise):15    with torch.no_grad():16        feat_4 = net._init(noise)17        feat_8 = net._upsample_8(feat_4)18        feat_16 = net._upsample_16(feat_8)19        feat_32 = net._upsample_32(feat_16)20        feat_64 = net._upsample_64(feat_32)21    return feat_8, feat_16, feat_32, feat_6422 23def get_late_features(net, feat_64, feat_8, feat_16, feat_32):24    with torch.no_grad():25        feat_128 = net._upsample_128(feat_64)26        feat_128 = net._sle_128(feat_8, feat_128)27 28        feat_256 = net._upsample_256(feat_128)29        feat_256 = net._sle_256(feat_16, feat_256)30 31        feat_512 = net._upsample_512(feat_256)32        feat_512 = net._sle_512(feat_32, feat_512)33 34        feat_1024 = net._upsample_1024(feat_512)35 36    return net._out_1024(feat_1024)37 38def style_mix(model_name_or_path, bs, device):39    _in_channels = 25640    im_size = 102441 42    netG = Generator(in_channels=_in_channels, out_channels=3)43    netG = netG.from_pretrained(model_name_or_path, in_channels=256, out_channels=3)44    _ = netG.to(device)45    _ = netG.eval()46 47    avg_param_G = copy_G_params(netG)48    load_params(netG, avg_param_G)49 50    noise_a = torch.randn(bs, 256, 1, 1, device=device).to(device)51    noise_b = torch.randn(bs, 256, 1, 1, device=device).to(device)52 53    feat_8_a, feat_16_a, feat_32_a, feat_64_a = get_early_features(netG, noise_a)54    feat_8_b, feat_16_b, feat_32_b, feat_64_b = get_early_features(netG, noise_b)55 56    images_b = get_late_features(netG, feat_64_b, feat_8_b, feat_16_b, feat_32_b)57    images_a = get_late_features(netG, feat_64_a, feat_8_a, feat_16_a, feat_32_a)58 59    imgs = [ torch.ones(1, 3, im_size, im_size) ]60 61    imgs.append(images_b.cpu())62    for i in range(bs):63        imgs.append(images_a[i].unsqueeze(0).cpu())64        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)65        imgs.append(gimgs.cpu())66 67    imgs = torch.cat(imgs)68    # vutils.save_image(imgs.add(1).mul(0.5), 'style_mix/style_mix_2.jpg', nrow=bs+1)69 70    return imgs71