Spaces:
Runtime error
Runtime error
File size: 2,347 Bytes
5e78353 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
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)
# vutils.save_image(imgs.add(1).mul(0.5), 'style_mix/style_mix_2.jpg', nrow=bs+1)
return imgs
|