File size: 6,176 Bytes
33719d5 |
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 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 |
# original implementation: https://github.com/odegeasslbc/FastGAN-pytorch/blob/main/models.py
#
# modified by Axel Sauer for "Projected GANs Converge Faster"
#
import torch.nn as nn
from blocks import (InitLayer, UpBlockBig, UpBlockBigCond, UpBlockSmall, UpBlockSmallCond, SEBlock, conv2d)
from huggingface_hub import PyTorchModelHubMixin
def normalize_second_moment(x, dim=1, eps=1e-8):
return x * (x.square().mean(dim=dim, keepdim=True) + eps).rsqrt()
class DummyMapping(nn.Module):
def __init__(self):
super().__init__()
def forward(self, z, c, **kwargs):
return z.unsqueeze(1) # to fit the StyleGAN API
class FastganSynthesis(nn.Module):
def __init__(self, ngf=128, z_dim=256, nc=3, img_resolution=256, lite=False):
super().__init__()
self.img_resolution = img_resolution
self.z_dim = z_dim
# channel multiplier
nfc_multi = {2: 16, 4:16, 8:8, 16:4, 32:2, 64:2, 128:1, 256:0.5,
512:0.25, 1024:0.125}
nfc = {}
for k, v in nfc_multi.items():
nfc[k] = int(v*ngf)
# layers
self.init = InitLayer(z_dim, channel=nfc[2], sz=4)
UpBlock = UpBlockSmall if lite else UpBlockBig
self.feat_8 = UpBlock(nfc[4], nfc[8])
self.feat_16 = UpBlock(nfc[8], nfc[16])
self.feat_32 = UpBlock(nfc[16], nfc[32])
self.feat_64 = UpBlock(nfc[32], nfc[64])
self.feat_128 = UpBlock(nfc[64], nfc[128])
self.feat_256 = UpBlock(nfc[128], nfc[256])
self.se_64 = SEBlock(nfc[4], nfc[64])
self.se_128 = SEBlock(nfc[8], nfc[128])
self.se_256 = SEBlock(nfc[16], nfc[256])
self.to_big = conv2d(nfc[img_resolution], nc, 3, 1, 1, bias=True)
if img_resolution > 256:
self.feat_512 = UpBlock(nfc[256], nfc[512])
self.se_512 = SEBlock(nfc[32], nfc[512])
if img_resolution > 512:
self.feat_1024 = UpBlock(nfc[512], nfc[1024])
def forward(self, input, c, **kwargs):
# map noise to hypersphere as in "Progressive Growing of GANS"
input = normalize_second_moment(input[:, 0])
feat_4 = self.init(input)
feat_8 = self.feat_8(feat_4)
feat_16 = self.feat_16(feat_8)
feat_32 = self.feat_32(feat_16)
feat_64 = self.se_64(feat_4, self.feat_64(feat_32))
feat_128 = self.se_128(feat_8, self.feat_128(feat_64))
if self.img_resolution >= 128:
feat_last = feat_128
if self.img_resolution >= 256:
feat_last = self.se_256(feat_16, self.feat_256(feat_last))
if self.img_resolution >= 512:
feat_last = self.se_512(feat_32, self.feat_512(feat_last))
if self.img_resolution >= 1024:
feat_last = self.feat_1024(feat_last)
return self.to_big(feat_last)
class FastganSynthesisCond(nn.Module):
def __init__(self, ngf=64, z_dim=256, nc=3, img_resolution=256, num_classes=1000, lite=False):
super().__init__()
self.z_dim = z_dim
nfc_multi = {2: 16, 4:16, 8:8, 16:4, 32:2, 64:2, 128:1, 256:0.5,
512:0.25, 1024:0.125, 2048:0.125}
nfc = {}
for k, v in nfc_multi.items():
nfc[k] = int(v*ngf)
self.img_resolution = img_resolution
self.init = InitLayer(z_dim, channel=nfc[2], sz=4)
UpBlock = UpBlockSmallCond if lite else UpBlockBigCond
self.feat_8 = UpBlock(nfc[4], nfc[8], z_dim)
self.feat_16 = UpBlock(nfc[8], nfc[16], z_dim)
self.feat_32 = UpBlock(nfc[16], nfc[32], z_dim)
self.feat_64 = UpBlock(nfc[32], nfc[64], z_dim)
self.feat_128 = UpBlock(nfc[64], nfc[128], z_dim)
self.feat_256 = UpBlock(nfc[128], nfc[256], z_dim)
self.se_64 = SEBlock(nfc[4], nfc[64])
self.se_128 = SEBlock(nfc[8], nfc[128])
self.se_256 = SEBlock(nfc[16], nfc[256])
self.to_big = conv2d(nfc[img_resolution], nc, 3, 1, 1, bias=True)
if img_resolution > 256:
self.feat_512 = UpBlock(nfc[256], nfc[512])
self.se_512 = SEBlock(nfc[32], nfc[512])
if img_resolution > 512:
self.feat_1024 = UpBlock(nfc[512], nfc[1024])
self.embed = nn.Embedding(num_classes, z_dim)
def forward(self, input, c, update_emas=False):
c = self.embed(c.argmax(1))
# map noise to hypersphere as in "Progressive Growing of GANS"
input = normalize_second_moment(input[:, 0])
feat_4 = self.init(input)
feat_8 = self.feat_8(feat_4, c)
feat_16 = self.feat_16(feat_8, c)
feat_32 = self.feat_32(feat_16, c)
feat_64 = self.se_64(feat_4, self.feat_64(feat_32, c))
feat_128 = self.se_128(feat_8, self.feat_128(feat_64, c))
if self.img_resolution >= 128:
feat_last = feat_128
if self.img_resolution >= 256:
feat_last = self.se_256(feat_16, self.feat_256(feat_last, c))
if self.img_resolution >= 512:
feat_last = self.se_512(feat_32, self.feat_512(feat_last, c))
if self.img_resolution >= 1024:
feat_last = self.feat_1024(feat_last, c)
return self.to_big(feat_last)
class MyGenerator(nn.Module, PyTorchModelHubMixin):
def __init__(
self,
z_dim=256,
c_dim=0,
w_dim=0,
img_resolution=256,
img_channels=3,
ngf=128,
cond=0,
mapping_kwargs={},
synthesis_kwargs={}
):
super().__init__()
#self.config = kwargs.pop("config", None)
self.z_dim = z_dim
self.c_dim = c_dim
self.w_dim = w_dim
self.img_resolution = img_resolution
self.img_channels = img_channels
# Mapping and Synthesis Networks
self.mapping = DummyMapping() # to fit the StyleGAN API
Synthesis = FastganSynthesisCond if cond else FastganSynthesis
self.synthesis = Synthesis(ngf=ngf, z_dim=z_dim, nc=img_channels, img_resolution=img_resolution, **synthesis_kwargs)
def forward(self, z, c, **kwargs):
w = self.mapping(z, c)
img = self.synthesis(w, c)
return img
|