asdasdsa

app.py

@@ -1,18 +1,21 @@
 import gradio as gr
-
+from huggingface_hub import PyTorchModelHubMixin
+import torch
+import matplotlib.pyplot as plt
+import torchvision
+from networks_fastgan import Generator
 
 def image_generation(model, number_of_images=1):
-    Images = 0
+    G = Generator.from_pretrained("cropinky/projected_gan_impressionism")
     return f"generating {number_of_images} images from {model}"
-
 if __name__ == "__main__":
 
     inputs = gr.inputs.Radio(["Abstract Expressionism", "Impressionism", "Cubism", "Minimalism", "Pop Art", "Color Field", "Hana Hanak houses"])
-    outputs = gr.outputs.Image(label="Output Image")
-
+    #outputs = gr.outputs.Image(label="Output Image")
+    outputs = "text"
     title = "Projected GAN for painting generation"
-    description = "choose your artistic direction "
-    article = "<p style='text-align: center'><a href='https://github.com/autonomousvision/projected_gan'>Official projected GAn github repo + paper</a> | <a href='https://github.com/facebookresearch/pytorch_GAN_zoo/blob/master/models/DCGAN.py'>Github Repo</a></p>"
+    description = "Choose your artistic direction "
+    article = "<p style='text-align: center'><a href='https://github.com/autonomousvision/projected_gan'>Official projected GAN github repo + paper</a></p>"
 
 
 

blocks.py

@@ -0,0 +1,325 @@
+import functools
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.utils import spectral_norm
+
+
+### single layers
+
+
+def conv2d(*args, **kwargs):
+    return spectral_norm(nn.Conv2d(*args, **kwargs))
+
+
+def convTranspose2d(*args, **kwargs):
+    return spectral_norm(nn.ConvTranspose2d(*args, **kwargs))
+
+
+def embedding(*args, **kwargs):
+    return spectral_norm(nn.Embedding(*args, **kwargs))
+
+
+def linear(*args, **kwargs):
+    return spectral_norm(nn.Linear(*args, **kwargs))
+
+
+def NormLayer(c, mode='batch'):
+    if mode == 'group':
+        return nn.GroupNorm(c//2, c)
+    elif mode == 'batch':
+        return nn.BatchNorm2d(c)
+
+
+### Activations
+
+
+class GLU(nn.Module):
+    def forward(self, x):
+        nc = x.size(1)
+        assert nc % 2 == 0, 'channels dont divide 2!'
+        nc = int(nc/2)
+        return x[:, :nc] * torch.sigmoid(x[:, nc:])
+
+
+class Swish(nn.Module):
+    def forward(self, feat):
+        return feat * torch.sigmoid(feat)
+
+
+### Upblocks
+
+
+class InitLayer(nn.Module):
+    def __init__(self, nz, channel, sz=4):
+        super().__init__()
+
+        self.init = nn.Sequential(
+            convTranspose2d(nz, channel*2, sz, 1, 0, bias=False),
+            NormLayer(channel*2),
+            GLU(),
+        )
+
+    def forward(self, noise):
+        noise = noise.view(noise.shape[0], -1, 1, 1)
+        return self.init(noise)
+
+
+def UpBlockSmall(in_planes, out_planes):
+    block = nn.Sequential(
+        nn.Upsample(scale_factor=2, mode='nearest'),
+        conv2d(in_planes, out_planes*2, 3, 1, 1, bias=False),
+        NormLayer(out_planes*2), GLU())
+    return block
+
+
+class UpBlockSmallCond(nn.Module):
+    def __init__(self, in_planes, out_planes, z_dim):
+        super().__init__()
+        self.in_planes = in_planes
+        self.out_planes = out_planes
+        self.up = nn.Upsample(scale_factor=2, mode='nearest')
+        self.conv = conv2d(in_planes, out_planes*2, 3, 1, 1, bias=False)
+
+        which_bn = functools.partial(CCBN, which_linear=linear, input_size=z_dim)
+        self.bn = which_bn(2*out_planes)
+        self.act = GLU()
+
+    def forward(self, x, c):
+        x = self.up(x)
+        x = self.conv(x)
+        x = self.bn(x, c)
+        x = self.act(x)
+        return x
+
+
+def UpBlockBig(in_planes, out_planes):
+    block = nn.Sequential(
+        nn.Upsample(scale_factor=2, mode='nearest'),
+        conv2d(in_planes, out_planes*2, 3, 1, 1, bias=False),
+        NoiseInjection(),
+        NormLayer(out_planes*2), GLU(),
+        conv2d(out_planes, out_planes*2, 3, 1, 1, bias=False),
+        NoiseInjection(),
+        NormLayer(out_planes*2), GLU()
+        )
+    return block
+
+
+class UpBlockBigCond(nn.Module):
+    def __init__(self, in_planes, out_planes, z_dim):
+        super().__init__()
+        self.in_planes = in_planes
+        self.out_planes = out_planes
+        self.up = nn.Upsample(scale_factor=2, mode='nearest')
+        self.conv1 = conv2d(in_planes, out_planes*2, 3, 1, 1, bias=False)
+        self.conv2 = conv2d(out_planes, out_planes*2, 3, 1, 1, bias=False)
+
+        which_bn = functools.partial(CCBN, which_linear=linear, input_size=z_dim)
+        self.bn1 = which_bn(2*out_planes)
+        self.bn2 = which_bn(2*out_planes)
+        self.act = GLU()
+        self.noise = NoiseInjection()
+
+    def forward(self, x, c):
+        # block 1
+        x = self.up(x)
+        x = self.conv1(x)
+        x = self.noise(x)
+        x = self.bn1(x, c)
+        x = self.act(x)
+
+        # block 2
+        x = self.conv2(x)
+        x = self.noise(x)
+        x = self.bn2(x, c)
+        x = self.act(x)
+
+        return x
+
+
+class SEBlock(nn.Module):
+    def __init__(self, ch_in, ch_out):
+        super().__init__()
+        self.main = nn.Sequential(
+            nn.AdaptiveAvgPool2d(4),
+            conv2d(ch_in, ch_out, 4, 1, 0, bias=False),
+            Swish(),
+            conv2d(ch_out, ch_out, 1, 1, 0, bias=False),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, feat_small, feat_big):
+        return feat_big * self.main(feat_small)
+
+
+### Downblocks
+
+
+class SeparableConv2d(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, bias=False):
+        super(SeparableConv2d, self).__init__()
+        self.depthwise = conv2d(in_channels, in_channels, kernel_size=kernel_size,
+            groups=in_channels, bias=bias, padding=1)
+        self.pointwise = conv2d(in_channels, out_channels,
+            kernel_size=1, bias=bias)
+
+    def forward(self, x):
+        out = self.depthwise(x)
+        out = self.pointwise(out)
+        return out
+
+
+class DownBlock(nn.Module):
+    def __init__(self, in_planes, out_planes, separable=False):
+        super().__init__()
+        if not separable:
+            self.main = nn.Sequential(
+                conv2d(in_planes, out_planes, 4, 2, 1),
+                NormLayer(out_planes),
+                nn.LeakyReLU(0.2, inplace=True),
+            )
+        else:
+            self.main = nn.Sequential(
+                SeparableConv2d(in_planes, out_planes, 3),
+                NormLayer(out_planes),
+                nn.LeakyReLU(0.2, inplace=True),
+                nn.AvgPool2d(2, 2),
+            )
+
+    def forward(self, feat):
+        return self.main(feat)
+
+
+class DownBlockPatch(nn.Module):
+    def __init__(self, in_planes, out_planes, separable=False):
+        super().__init__()
+        self.main = nn.Sequential(
+            DownBlock(in_planes, out_planes, separable),
+            conv2d(out_planes, out_planes, 1, 1, 0, bias=False),
+            NormLayer(out_planes),
+            nn.LeakyReLU(0.2, inplace=True),
+        )
+
+    def forward(self, feat):
+        return self.main(feat)
+
+
+### CSM
+
+
+class ResidualConvUnit(nn.Module):
+    def __init__(self, cin, activation, bn):
+        super().__init__()
+        self.conv = nn.Conv2d(cin, cin, kernel_size=3, stride=1, padding=1, bias=True)
+        self.skip_add = nn.quantized.FloatFunctional()
+
+    def forward(self, x):
+        return self.skip_add.add(self.conv(x), x)
+
+
+class FeatureFusionBlock(nn.Module):
+    def __init__(self, features, activation, deconv=False, bn=False, expand=False, align_corners=True, lowest=False):
+        super().__init__()
+
+        self.deconv = deconv
+        self.align_corners = align_corners
+
+        self.expand = expand
+        out_features = features
+        if self.expand==True:
+            out_features = features//2
+
+        self.out_conv = nn.Conv2d(features, out_features, kernel_size=1, stride=1, padding=0, bias=True, groups=1)
+        self.skip_add = nn.quantized.FloatFunctional()
+
+    def forward(self, *xs):
+        output = xs[0]
+
+        if len(xs) == 2:
+            output = self.skip_add.add(output, xs[1])
+
+        output = nn.functional.interpolate(
+            output, scale_factor=2, mode="bilinear", align_corners=self.align_corners
+        )
+
+        output = self.out_conv(output)
+
+        return output
+
+
+### Misc
+
+
+class NoiseInjection(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.weight = nn.Parameter(torch.zeros(1), requires_grad=True)
+
+    def forward(self, feat, noise=None):
+        if noise is None:
+            batch, _, height, width = feat.shape
+            noise = torch.randn(batch, 1, height, width).to(feat.device)
+
+        return feat + self.weight * noise
+
+
+class CCBN(nn.Module):
+    ''' conditional batchnorm '''
+    def __init__(self, output_size, input_size, which_linear, eps=1e-5, momentum=0.1):
+        super().__init__()
+        self.output_size, self.input_size = output_size, input_size
+
+        # Prepare gain and bias layers
+        self.gain = which_linear(input_size, output_size)
+        self.bias = which_linear(input_size, output_size)
+
+        # epsilon to avoid dividing by 0
+        self.eps = eps
+        # Momentum
+        self.momentum = momentum
+
+        self.register_buffer('stored_mean', torch.zeros(output_size))
+        self.register_buffer('stored_var', torch.ones(output_size))
+
+    def forward(self, x, y):
+        # Calculate class-conditional gains and biases
+        gain = (1 + self.gain(y)).view(y.size(0), -1, 1, 1)
+        bias = self.bias(y).view(y.size(0), -1, 1, 1)
+        out = F.batch_norm(x, self.stored_mean, self.stored_var, None, None,
+                           self.training, 0.1, self.eps)
+        return out * gain + bias
+
+
+class Interpolate(nn.Module):
+    """Interpolation module."""
+
+    def __init__(self, size, mode='bilinear', align_corners=False):
+        """Init.
+        Args:
+            scale_factor (float): scaling
+            mode (str): interpolation mode
+        """
+        super(Interpolate, self).__init__()
+
+        self.interp = nn.functional.interpolate
+        self.size = size
+        self.mode = mode
+        self.align_corners = align_corners
+
+    def forward(self, x):
+        """Forward pass.
+        Args:
+            x (tensor): input
+        Returns:
+            tensor: interpolated data
+        """
+
+        x = self.interp(
+            x,
+            size=self.size,
+            mode=self.mode,
+            align_corners=self.align_corners,
+        )
+
+        return x

networks_fastgan.py

@@ -0,0 +1,179 @@
+# 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 Generator(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