Cropinky commited on
Commit
33719d5
1 Parent(s): a80f735

initial upload hana houses

Browse files
Files changed (5) hide show
  1. app.py +30 -0
  2. blocks.py +325 -0
  3. image_generator.py +153 -0
  4. networks_fastgan.py +179 -0
  5. requirements.txt +4 -0
app.py ADDED
@@ -0,0 +1,30 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from huggingface_hub import PyTorchModelHubMixin
3
+ import torch
4
+ import matplotlib.pyplot as plt
5
+ import torchvision
6
+ from networks_fastgan import MyGenerator
7
+ import click
8
+ import PIL
9
+ from image_generator import generate_images
10
+
11
+ def image_generation(model, number_of_images=1):
12
+ img = generate_images(model)
13
+ #return f"generating {number_of_images} images from {model}"
14
+ return img
15
+ if __name__ == "__main__":
16
+ description = "TODO: when generating only 1 image use an esrgan to increase its resolution \n TODO: allow generation of multiple images TODO: walk through input space video i have exams now c u in 2 weeks (:"
17
+ inputs = gr.inputs.Radio([ "Impressionism", "Abstract Expressionism", "Cubism", "Pop Art", "Color Field", "Hana Hanak houses", "Hana Hanak houses - abstract expressionism", "Hana Hanak houses - color field"])
18
+ outputs = gr.outputs.Image(label="Generated Image", type="pil")
19
+ #outputs = "text"
20
+ title = "Projected GAN for painting generation v0.2"
21
+ article = "<p style='text-align: center'><a href='https://github.com/autonomousvision/projected_gan'>Official projected GAN github repo + paper</a></p>"
22
+
23
+
24
+
25
+ gr.Interface(image_generation, inputs, outputs, title=title, article = article, description = description,
26
+ analytics_enabled=False).launch(debug=True)
27
+
28
+ app, local_url, share_url = iface.launch()
29
+
30
+
blocks.py ADDED
@@ -0,0 +1,325 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import functools
2
+ import torch
3
+ import torch.nn as nn
4
+ import torch.nn.functional as F
5
+ from torch.nn.utils import spectral_norm
6
+
7
+
8
+ ### single layers
9
+
10
+
11
+ def conv2d(*args, **kwargs):
12
+ return spectral_norm(nn.Conv2d(*args, **kwargs))
13
+
14
+
15
+ def convTranspose2d(*args, **kwargs):
16
+ return spectral_norm(nn.ConvTranspose2d(*args, **kwargs))
17
+
18
+
19
+ def embedding(*args, **kwargs):
20
+ return spectral_norm(nn.Embedding(*args, **kwargs))
21
+
22
+
23
+ def linear(*args, **kwargs):
24
+ return spectral_norm(nn.Linear(*args, **kwargs))
25
+
26
+
27
+ def NormLayer(c, mode='batch'):
28
+ if mode == 'group':
29
+ return nn.GroupNorm(c//2, c)
30
+ elif mode == 'batch':
31
+ return nn.BatchNorm2d(c)
32
+
33
+
34
+ ### Activations
35
+
36
+
37
+ class GLU(nn.Module):
38
+ def forward(self, x):
39
+ nc = x.size(1)
40
+ assert nc % 2 == 0, 'channels dont divide 2!'
41
+ nc = int(nc/2)
42
+ return x[:, :nc] * torch.sigmoid(x[:, nc:])
43
+
44
+
45
+ class Swish(nn.Module):
46
+ def forward(self, feat):
47
+ return feat * torch.sigmoid(feat)
48
+
49
+
50
+ ### Upblocks
51
+
52
+
53
+ class InitLayer(nn.Module):
54
+ def __init__(self, nz, channel, sz=4):
55
+ super().__init__()
56
+
57
+ self.init = nn.Sequential(
58
+ convTranspose2d(nz, channel*2, sz, 1, 0, bias=False),
59
+ NormLayer(channel*2),
60
+ GLU(),
61
+ )
62
+
63
+ def forward(self, noise):
64
+ noise = noise.view(noise.shape[0], -1, 1, 1)
65
+ return self.init(noise)
66
+
67
+
68
+ def UpBlockSmall(in_planes, out_planes):
69
+ block = nn.Sequential(
70
+ nn.Upsample(scale_factor=2, mode='nearest'),
71
+ conv2d(in_planes, out_planes*2, 3, 1, 1, bias=False),
72
+ NormLayer(out_planes*2), GLU())
73
+ return block
74
+
75
+
76
+ class UpBlockSmallCond(nn.Module):
77
+ def __init__(self, in_planes, out_planes, z_dim):
78
+ super().__init__()
79
+ self.in_planes = in_planes
80
+ self.out_planes = out_planes
81
+ self.up = nn.Upsample(scale_factor=2, mode='nearest')
82
+ self.conv = conv2d(in_planes, out_planes*2, 3, 1, 1, bias=False)
83
+
84
+ which_bn = functools.partial(CCBN, which_linear=linear, input_size=z_dim)
85
+ self.bn = which_bn(2*out_planes)
86
+ self.act = GLU()
87
+
88
+ def forward(self, x, c):
89
+ x = self.up(x)
90
+ x = self.conv(x)
91
+ x = self.bn(x, c)
92
+ x = self.act(x)
93
+ return x
94
+
95
+
96
+ def UpBlockBig(in_planes, out_planes):
97
+ block = nn.Sequential(
98
+ nn.Upsample(scale_factor=2, mode='nearest'),
99
+ conv2d(in_planes, out_planes*2, 3, 1, 1, bias=False),
100
+ NoiseInjection(),
101
+ NormLayer(out_planes*2), GLU(),
102
+ conv2d(out_planes, out_planes*2, 3, 1, 1, bias=False),
103
+ NoiseInjection(),
104
+ NormLayer(out_planes*2), GLU()
105
+ )
106
+ return block
107
+
108
+
109
+ class UpBlockBigCond(nn.Module):
110
+ def __init__(self, in_planes, out_planes, z_dim):
111
+ super().__init__()
112
+ self.in_planes = in_planes
113
+ self.out_planes = out_planes
114
+ self.up = nn.Upsample(scale_factor=2, mode='nearest')
115
+ self.conv1 = conv2d(in_planes, out_planes*2, 3, 1, 1, bias=False)
116
+ self.conv2 = conv2d(out_planes, out_planes*2, 3, 1, 1, bias=False)
117
+
118
+ which_bn = functools.partial(CCBN, which_linear=linear, input_size=z_dim)
119
+ self.bn1 = which_bn(2*out_planes)
120
+ self.bn2 = which_bn(2*out_planes)
121
+ self.act = GLU()
122
+ self.noise = NoiseInjection()
123
+
124
+ def forward(self, x, c):
125
+ # block 1
126
+ x = self.up(x)
127
+ x = self.conv1(x)
128
+ x = self.noise(x)
129
+ x = self.bn1(x, c)
130
+ x = self.act(x)
131
+
132
+ # block 2
133
+ x = self.conv2(x)
134
+ x = self.noise(x)
135
+ x = self.bn2(x, c)
136
+ x = self.act(x)
137
+
138
+ return x
139
+
140
+
141
+ class SEBlock(nn.Module):
142
+ def __init__(self, ch_in, ch_out):
143
+ super().__init__()
144
+ self.main = nn.Sequential(
145
+ nn.AdaptiveAvgPool2d(4),
146
+ conv2d(ch_in, ch_out, 4, 1, 0, bias=False),
147
+ Swish(),
148
+ conv2d(ch_out, ch_out, 1, 1, 0, bias=False),
149
+ nn.Sigmoid(),
150
+ )
151
+
152
+ def forward(self, feat_small, feat_big):
153
+ return feat_big * self.main(feat_small)
154
+
155
+
156
+ ### Downblocks
157
+
158
+
159
+ class SeparableConv2d(nn.Module):
160
+ def __init__(self, in_channels, out_channels, kernel_size, bias=False):
161
+ super(SeparableConv2d, self).__init__()
162
+ self.depthwise = conv2d(in_channels, in_channels, kernel_size=kernel_size,
163
+ groups=in_channels, bias=bias, padding=1)
164
+ self.pointwise = conv2d(in_channels, out_channels,
165
+ kernel_size=1, bias=bias)
166
+
167
+ def forward(self, x):
168
+ out = self.depthwise(x)
169
+ out = self.pointwise(out)
170
+ return out
171
+
172
+
173
+ class DownBlock(nn.Module):
174
+ def __init__(self, in_planes, out_planes, separable=False):
175
+ super().__init__()
176
+ if not separable:
177
+ self.main = nn.Sequential(
178
+ conv2d(in_planes, out_planes, 4, 2, 1),
179
+ NormLayer(out_planes),
180
+ nn.LeakyReLU(0.2, inplace=True),
181
+ )
182
+ else:
183
+ self.main = nn.Sequential(
184
+ SeparableConv2d(in_planes, out_planes, 3),
185
+ NormLayer(out_planes),
186
+ nn.LeakyReLU(0.2, inplace=True),
187
+ nn.AvgPool2d(2, 2),
188
+ )
189
+
190
+ def forward(self, feat):
191
+ return self.main(feat)
192
+
193
+
194
+ class DownBlockPatch(nn.Module):
195
+ def __init__(self, in_planes, out_planes, separable=False):
196
+ super().__init__()
197
+ self.main = nn.Sequential(
198
+ DownBlock(in_planes, out_planes, separable),
199
+ conv2d(out_planes, out_planes, 1, 1, 0, bias=False),
200
+ NormLayer(out_planes),
201
+ nn.LeakyReLU(0.2, inplace=True),
202
+ )
203
+
204
+ def forward(self, feat):
205
+ return self.main(feat)
206
+
207
+
208
+ ### CSM
209
+
210
+
211
+ class ResidualConvUnit(nn.Module):
212
+ def __init__(self, cin, activation, bn):
213
+ super().__init__()
214
+ self.conv = nn.Conv2d(cin, cin, kernel_size=3, stride=1, padding=1, bias=True)
215
+ self.skip_add = nn.quantized.FloatFunctional()
216
+
217
+ def forward(self, x):
218
+ return self.skip_add.add(self.conv(x), x)
219
+
220
+
221
+ class FeatureFusionBlock(nn.Module):
222
+ def __init__(self, features, activation, deconv=False, bn=False, expand=False, align_corners=True, lowest=False):
223
+ super().__init__()
224
+
225
+ self.deconv = deconv
226
+ self.align_corners = align_corners
227
+
228
+ self.expand = expand
229
+ out_features = features
230
+ if self.expand==True:
231
+ out_features = features//2
232
+
233
+ self.out_conv = nn.Conv2d(features, out_features, kernel_size=1, stride=1, padding=0, bias=True, groups=1)
234
+ self.skip_add = nn.quantized.FloatFunctional()
235
+
236
+ def forward(self, *xs):
237
+ output = xs[0]
238
+
239
+ if len(xs) == 2:
240
+ output = self.skip_add.add(output, xs[1])
241
+
242
+ output = nn.functional.interpolate(
243
+ output, scale_factor=2, mode="bilinear", align_corners=self.align_corners
244
+ )
245
+
246
+ output = self.out_conv(output)
247
+
248
+ return output
249
+
250
+
251
+ ### Misc
252
+
253
+
254
+ class NoiseInjection(nn.Module):
255
+ def __init__(self):
256
+ super().__init__()
257
+ self.weight = nn.Parameter(torch.zeros(1), requires_grad=True)
258
+
259
+ def forward(self, feat, noise=None):
260
+ if noise is None:
261
+ batch, _, height, width = feat.shape
262
+ noise = torch.randn(batch, 1, height, width).to(feat.device)
263
+
264
+ return feat + self.weight * noise
265
+
266
+
267
+ class CCBN(nn.Module):
268
+ ''' conditional batchnorm '''
269
+ def __init__(self, output_size, input_size, which_linear, eps=1e-5, momentum=0.1):
270
+ super().__init__()
271
+ self.output_size, self.input_size = output_size, input_size
272
+
273
+ # Prepare gain and bias layers
274
+ self.gain = which_linear(input_size, output_size)
275
+ self.bias = which_linear(input_size, output_size)
276
+
277
+ # epsilon to avoid dividing by 0
278
+ self.eps = eps
279
+ # Momentum
280
+ self.momentum = momentum
281
+
282
+ self.register_buffer('stored_mean', torch.zeros(output_size))
283
+ self.register_buffer('stored_var', torch.ones(output_size))
284
+
285
+ def forward(self, x, y):
286
+ # Calculate class-conditional gains and biases
287
+ gain = (1 + self.gain(y)).view(y.size(0), -1, 1, 1)
288
+ bias = self.bias(y).view(y.size(0), -1, 1, 1)
289
+ out = F.batch_norm(x, self.stored_mean, self.stored_var, None, None,
290
+ self.training, 0.1, self.eps)
291
+ return out * gain + bias
292
+
293
+
294
+ class Interpolate(nn.Module):
295
+ """Interpolation module."""
296
+
297
+ def __init__(self, size, mode='bilinear', align_corners=False):
298
+ """Init.
299
+ Args:
300
+ scale_factor (float): scaling
301
+ mode (str): interpolation mode
302
+ """
303
+ super(Interpolate, self).__init__()
304
+
305
+ self.interp = nn.functional.interpolate
306
+ self.size = size
307
+ self.mode = mode
308
+ self.align_corners = align_corners
309
+
310
+ def forward(self, x):
311
+ """Forward pass.
312
+ Args:
313
+ x (tensor): input
314
+ Returns:
315
+ tensor: interpolated data
316
+ """
317
+
318
+ x = self.interp(
319
+ x,
320
+ size=self.size,
321
+ mode=self.mode,
322
+ align_corners=self.align_corners,
323
+ )
324
+
325
+ return x
image_generator.py ADDED
@@ -0,0 +1,153 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
+ #
3
+ # NVIDIA CORPORATION and its licensors retain all intellectual property
4
+ # and proprietary rights in and to this software, related documentation
5
+ # and any modifications thereto. Any use, reproduction, disclosure or
6
+ # distribution of this software and related documentation without an express
7
+ # license agreement from NVIDIA CORPORATION is strictly prohibited.
8
+
9
+ """Generate images using pretrained network pickle."""
10
+
11
+ from ast import parse
12
+ import os
13
+ from pyexpat import model
14
+ import re
15
+ from typing import List, Optional, Tuple, Union
16
+ import numpy as np
17
+ import PIL.Image
18
+ import torch
19
+ from networks_fastgan import MyGenerator
20
+ import random
21
+ #----------------------------------------------------------------------------
22
+
23
+ def parse_range(s: Union[str, List]) -> List[int]:
24
+ '''Parse a comma separated list of numbers or ranges and return a list of ints.
25
+
26
+ Example: '1,2,5-10' returns [1, 2, 5, 6, 7]
27
+ '''
28
+ if isinstance(s, list): return s
29
+ ranges = []
30
+ range_re = re.compile(r'^(\d+)-(\d+)$')
31
+ for p in s.split(','):
32
+ m = range_re.match(p)
33
+ if m:
34
+ ranges.extend(range(int(m.group(1)), int(m.group(2))+1))
35
+ else:
36
+ ranges.append(int(p))
37
+ return ranges
38
+
39
+ #----------------------------------------------------------------------------
40
+
41
+ def parse_vec2(s: Union[str, Tuple[float, float]]) -> Tuple[float, float]:
42
+ '''Parse a floating point 2-vector of syntax 'a,b'.
43
+
44
+ Example:
45
+ '0,1' returns (0,1)
46
+ '''
47
+ if isinstance(s, tuple): return s
48
+ parts = s.split(',')
49
+ if len(parts) == 2:
50
+ return (float(parts[0]), float(parts[1]))
51
+ raise ValueError(f'cannot parse 2-vector {s}')
52
+
53
+ #----------------------------------------------------------------------------
54
+
55
+ def make_transform(translate: Tuple[float,float], angle: float):
56
+ m = np.eye(3)
57
+ s = np.sin(angle/360.0*np.pi*2)
58
+ c = np.cos(angle/360.0*np.pi*2)
59
+ m[0][0] = c
60
+ m[0][1] = s
61
+ m[0][2] = translate[0]
62
+ m[1][0] = -s
63
+ m[1][1] = c
64
+ m[1][2] = translate[1]
65
+ return m
66
+
67
+ #----------------------------------------------------------------------------
68
+
69
+ def generate_images(
70
+ model_path,
71
+ seeds = "10-12",
72
+ truncation_psi = 1.0,
73
+ noise_mode = "const",
74
+ outdir = "out",
75
+ translate = "0,0",
76
+ rotate = 0,
77
+ number_of_images = 16
78
+ ):
79
+ model_owner = "huggan"
80
+ #inputs = gr.inputs.Radio(["Abstract Expressionism", "Impressionism", "Cubism", "Minimalism", "Pop Art", "Color Field", "Hana Hanak houses"])
81
+ model_path_dict = {
82
+ 'Impressionism' : 'projected_gan_impressionism',
83
+ 'Cubism' : 'projected_gan_cubism',
84
+ 'Abstract Expressionism' : 'projected_gan_abstract_expressionism',
85
+ 'Pop Art' : 'projected_gan_popart',
86
+ 'Minimalism' : 'projected_gan_minimalism',
87
+ 'Color Field' : 'projected_gan_color_field',
88
+ 'Hana Hanak houses' : 'projected_gan_Hana_Hanak',
89
+ 'Hana Hanak houses - abstract expressionism' : 'projected_gan_abstract_expressionism_hana',
90
+ 'Hana Hanak houses - color field' : 'projected_gan_color_field_hana',
91
+
92
+ }
93
+
94
+ model_path = model_owner + "/" + model_path_dict[model_path]
95
+ print(model_path)
96
+ print(seeds)
97
+ seeds=[random.randint(1,230)]
98
+ #seeds =f"{seeds}-{seeds+number_of_images-1}"
99
+ #seeds = parse_range(seeds)
100
+ device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
101
+ G = MyGenerator.from_pretrained(model_path)
102
+ os.makedirs(outdir, exist_ok=True)
103
+ # Labels.
104
+ label = torch.zeros([1, G.c_dim], device=device)
105
+ """
106
+ if G.c_dim != 0:
107
+ if class_idx is None:
108
+ raise click.ClickException('Must specify class label with --class when using a conditional network')
109
+ label[:, class_idx] = 1
110
+ else:
111
+ if class_idx is not None:
112
+ print ('warn: --class=lbl ignored when running on an unconditional network')
113
+ """
114
+
115
+ print(f"z dimenzija mi je: {G.z_dim}")
116
+ # Generate images.
117
+
118
+ #imgs_row = np.array()
119
+ #imgs_complete = np.array()
120
+ for seed_idx, seed in enumerate(seeds):
121
+ print('Generating image for seed %d (%d/%d) ...' % (seed, seed_idx, len(seeds)))
122
+ z = torch.from_numpy(np.random.RandomState(seed).randn(1, G.z_dim)).to(device).float()
123
+ # Construct an inverse rotation/translation matrix and pass to the generator. The
124
+ # generator expects this matrix as an inverse to avoid potentially failing numerical
125
+ # operations in the network.
126
+ if hasattr(G.synthesis, 'input'):
127
+ m = make_transform(translate, rotate)
128
+ m = np.linalg.inv(m)
129
+ G.synthesis.input.transform.copy_(torch.from_numpy(m))
130
+
131
+ img = G(z, label, truncation_psi=truncation_psi, noise_mode=noise_mode)
132
+ img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
133
+ print(seed_idx)
134
+ #first image
135
+ if seed_idx == 0:
136
+ imgs_row = img[0].cpu().numpy()
137
+ else:
138
+ imgs_row = np.hstack((imgs_row, img[0].cpu().numpy()))
139
+ #img = PIL.Image.fromarray(img[0].cpu().numpy(), 'RGB')
140
+ #PIL.Image.fromarray(img[0].cpu().numpy(), 'RGB').save(f'{outdir}/seed{seed:04d}.png')
141
+ #napravi vsplit i toe to ka
142
+ imgs_complete = np.vstack(np.hsplit(imgs_row, 4))
143
+ #cv2.imshow("lalaxd", imgs_complete)
144
+ #cv2.waitKey()
145
+ return PIL.Image.fromarray(imgs_complete, 'RGB')
146
+
147
+
148
+ #----------------------------------------------------------------------------
149
+
150
+ if __name__ == "__main__":
151
+ generate_images() # pylint: disable=no-value-for-parameter
152
+
153
+ #----------------------------------------------------------------------------
networks_fastgan.py ADDED
@@ -0,0 +1,179 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # original implementation: https://github.com/odegeasslbc/FastGAN-pytorch/blob/main/models.py
2
+ #
3
+ # modified by Axel Sauer for "Projected GANs Converge Faster"
4
+ #
5
+ import torch.nn as nn
6
+ from blocks import (InitLayer, UpBlockBig, UpBlockBigCond, UpBlockSmall, UpBlockSmallCond, SEBlock, conv2d)
7
+ from huggingface_hub import PyTorchModelHubMixin
8
+
9
+ def normalize_second_moment(x, dim=1, eps=1e-8):
10
+ return x * (x.square().mean(dim=dim, keepdim=True) + eps).rsqrt()
11
+
12
+
13
+ class DummyMapping(nn.Module):
14
+ def __init__(self):
15
+ super().__init__()
16
+
17
+ def forward(self, z, c, **kwargs):
18
+ return z.unsqueeze(1) # to fit the StyleGAN API
19
+
20
+
21
+ class FastganSynthesis(nn.Module):
22
+ def __init__(self, ngf=128, z_dim=256, nc=3, img_resolution=256, lite=False):
23
+ super().__init__()
24
+ self.img_resolution = img_resolution
25
+ self.z_dim = z_dim
26
+
27
+ # channel multiplier
28
+ nfc_multi = {2: 16, 4:16, 8:8, 16:4, 32:2, 64:2, 128:1, 256:0.5,
29
+ 512:0.25, 1024:0.125}
30
+ nfc = {}
31
+ for k, v in nfc_multi.items():
32
+ nfc[k] = int(v*ngf)
33
+
34
+ # layers
35
+ self.init = InitLayer(z_dim, channel=nfc[2], sz=4)
36
+
37
+ UpBlock = UpBlockSmall if lite else UpBlockBig
38
+
39
+ self.feat_8 = UpBlock(nfc[4], nfc[8])
40
+ self.feat_16 = UpBlock(nfc[8], nfc[16])
41
+ self.feat_32 = UpBlock(nfc[16], nfc[32])
42
+ self.feat_64 = UpBlock(nfc[32], nfc[64])
43
+ self.feat_128 = UpBlock(nfc[64], nfc[128])
44
+ self.feat_256 = UpBlock(nfc[128], nfc[256])
45
+
46
+ self.se_64 = SEBlock(nfc[4], nfc[64])
47
+ self.se_128 = SEBlock(nfc[8], nfc[128])
48
+ self.se_256 = SEBlock(nfc[16], nfc[256])
49
+
50
+ self.to_big = conv2d(nfc[img_resolution], nc, 3, 1, 1, bias=True)
51
+
52
+ if img_resolution > 256:
53
+ self.feat_512 = UpBlock(nfc[256], nfc[512])
54
+ self.se_512 = SEBlock(nfc[32], nfc[512])
55
+ if img_resolution > 512:
56
+ self.feat_1024 = UpBlock(nfc[512], nfc[1024])
57
+
58
+ def forward(self, input, c, **kwargs):
59
+ # map noise to hypersphere as in "Progressive Growing of GANS"
60
+ input = normalize_second_moment(input[:, 0])
61
+
62
+ feat_4 = self.init(input)
63
+ feat_8 = self.feat_8(feat_4)
64
+ feat_16 = self.feat_16(feat_8)
65
+ feat_32 = self.feat_32(feat_16)
66
+ feat_64 = self.se_64(feat_4, self.feat_64(feat_32))
67
+ feat_128 = self.se_128(feat_8, self.feat_128(feat_64))
68
+
69
+ if self.img_resolution >= 128:
70
+ feat_last = feat_128
71
+
72
+ if self.img_resolution >= 256:
73
+ feat_last = self.se_256(feat_16, self.feat_256(feat_last))
74
+
75
+ if self.img_resolution >= 512:
76
+ feat_last = self.se_512(feat_32, self.feat_512(feat_last))
77
+
78
+ if self.img_resolution >= 1024:
79
+ feat_last = self.feat_1024(feat_last)
80
+
81
+ return self.to_big(feat_last)
82
+
83
+
84
+ class FastganSynthesisCond(nn.Module):
85
+ def __init__(self, ngf=64, z_dim=256, nc=3, img_resolution=256, num_classes=1000, lite=False):
86
+ super().__init__()
87
+
88
+ self.z_dim = z_dim
89
+ nfc_multi = {2: 16, 4:16, 8:8, 16:4, 32:2, 64:2, 128:1, 256:0.5,
90
+ 512:0.25, 1024:0.125, 2048:0.125}
91
+ nfc = {}
92
+ for k, v in nfc_multi.items():
93
+ nfc[k] = int(v*ngf)
94
+
95
+ self.img_resolution = img_resolution
96
+
97
+ self.init = InitLayer(z_dim, channel=nfc[2], sz=4)
98
+
99
+ UpBlock = UpBlockSmallCond if lite else UpBlockBigCond
100
+
101
+ self.feat_8 = UpBlock(nfc[4], nfc[8], z_dim)
102
+ self.feat_16 = UpBlock(nfc[8], nfc[16], z_dim)
103
+ self.feat_32 = UpBlock(nfc[16], nfc[32], z_dim)
104
+ self.feat_64 = UpBlock(nfc[32], nfc[64], z_dim)
105
+ self.feat_128 = UpBlock(nfc[64], nfc[128], z_dim)
106
+ self.feat_256 = UpBlock(nfc[128], nfc[256], z_dim)
107
+
108
+ self.se_64 = SEBlock(nfc[4], nfc[64])
109
+ self.se_128 = SEBlock(nfc[8], nfc[128])
110
+ self.se_256 = SEBlock(nfc[16], nfc[256])
111
+
112
+ self.to_big = conv2d(nfc[img_resolution], nc, 3, 1, 1, bias=True)
113
+
114
+ if img_resolution > 256:
115
+ self.feat_512 = UpBlock(nfc[256], nfc[512])
116
+ self.se_512 = SEBlock(nfc[32], nfc[512])
117
+ if img_resolution > 512:
118
+ self.feat_1024 = UpBlock(nfc[512], nfc[1024])
119
+
120
+ self.embed = nn.Embedding(num_classes, z_dim)
121
+
122
+ def forward(self, input, c, update_emas=False):
123
+ c = self.embed(c.argmax(1))
124
+
125
+ # map noise to hypersphere as in "Progressive Growing of GANS"
126
+ input = normalize_second_moment(input[:, 0])
127
+
128
+ feat_4 = self.init(input)
129
+ feat_8 = self.feat_8(feat_4, c)
130
+ feat_16 = self.feat_16(feat_8, c)
131
+ feat_32 = self.feat_32(feat_16, c)
132
+ feat_64 = self.se_64(feat_4, self.feat_64(feat_32, c))
133
+ feat_128 = self.se_128(feat_8, self.feat_128(feat_64, c))
134
+
135
+ if self.img_resolution >= 128:
136
+ feat_last = feat_128
137
+
138
+ if self.img_resolution >= 256:
139
+ feat_last = self.se_256(feat_16, self.feat_256(feat_last, c))
140
+
141
+ if self.img_resolution >= 512:
142
+ feat_last = self.se_512(feat_32, self.feat_512(feat_last, c))
143
+
144
+ if self.img_resolution >= 1024:
145
+ feat_last = self.feat_1024(feat_last, c)
146
+
147
+ return self.to_big(feat_last)
148
+
149
+
150
+ class MyGenerator(nn.Module, PyTorchModelHubMixin):
151
+ def __init__(
152
+ self,
153
+ z_dim=256,
154
+ c_dim=0,
155
+ w_dim=0,
156
+ img_resolution=256,
157
+ img_channels=3,
158
+ ngf=128,
159
+ cond=0,
160
+ mapping_kwargs={},
161
+ synthesis_kwargs={}
162
+ ):
163
+ super().__init__()
164
+ #self.config = kwargs.pop("config", None)
165
+ self.z_dim = z_dim
166
+ self.c_dim = c_dim
167
+ self.w_dim = w_dim
168
+ self.img_resolution = img_resolution
169
+ self.img_channels = img_channels
170
+
171
+ # Mapping and Synthesis Networks
172
+ self.mapping = DummyMapping() # to fit the StyleGAN API
173
+ Synthesis = FastganSynthesisCond if cond else FastganSynthesis
174
+ self.synthesis = Synthesis(ngf=ngf, z_dim=z_dim, nc=img_channels, img_resolution=img_resolution, **synthesis_kwargs)
175
+
176
+ def forward(self, z, c, **kwargs):
177
+ w = self.mapping(z, c)
178
+ img = self.synthesis(w, c)
179
+ return img
requirements.txt ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ gradio
2
+ torchvision
3
+ matplotlib
4
+ torch