vtoonify/model/dualstylegan.py

import random
import torch
from torch import nn
from model.stylegan.model import ConvLayer, PixelNorm, EqualLinear, Generator

class AdaptiveInstanceNorm(nn.Module):
    def __init__(self, fin, style_dim=512):
        super().__init__()

        self.norm = nn.InstanceNorm2d(fin, affine=False)
        self.style = nn.Linear(style_dim, fin * 2)

        self.style.bias.data[:fin] = 1
        self.style.bias.data[fin:] = 0

    def forward(self, input, style):
        style = self.style(style).unsqueeze(2).unsqueeze(3)
        gamma, beta = style.chunk(2, 1)
        out = self.norm(input)
        out = gamma * out + beta
        return out

# modulative residual blocks (ModRes)
class AdaResBlock(nn.Module):
    def __init__(self, fin, style_dim=512, dilation=1): # modified
        super().__init__()

        self.conv = ConvLayer(fin, fin, 3, dilation=dilation) # modified
        self.conv2 = ConvLayer(fin, fin, 3, dilation=dilation) # modified
        self.norm = AdaptiveInstanceNorm(fin, style_dim)
        self.norm2 = AdaptiveInstanceNorm(fin, style_dim)
        
        # model initialization
        # the convolution filters are set to values close to 0 to produce negligible residual features
        self.conv[0].weight.data *= 0.01
        self.conv2[0].weight.data *= 0.01
        
    def forward(self, x, s, w=1):
        skip = x
        if w == 0:
            return skip
        out = self.conv(self.norm(x, s))
        out = self.conv2(self.norm2(out, s))
        out = out * w + skip
        return out

class DualStyleGAN(nn.Module):
    def __init__(self, size, style_dim, n_mlp, channel_multiplier=2, twoRes=True, res_index=6):
        super().__init__()
        
        layers = [PixelNorm()]
        for i in range(n_mlp-6):
            layers.append(EqualLinear(512, 512, lr_mul=0.01, activation="fused_lrelu"))
        # color transform blocks T_c
        self.style = nn.Sequential(*layers)
        # StyleGAN2
        self.generator = Generator(size, style_dim, n_mlp, channel_multiplier) 
        # The extrinsic style path
        self.res = nn.ModuleList()
        self.res_index = res_index//2 * 2
        self.res.append(AdaResBlock(self.generator.channels[2 ** 2])) # for conv1
        for i in range(3, self.generator.log_size + 1):
            out_channel = self.generator.channels[2 ** i]
            if i < 3 + self.res_index//2:
                # ModRes
                self.res.append(AdaResBlock(out_channel))
                self.res.append(AdaResBlock(out_channel))
            else:
                # structure transform block T_s
                self.res.append(EqualLinear(512, 512))
                # FC layer is initialized with identity matrices, meaning no changes to the input latent code
                self.res[-1].weight.data = torch.eye(512) * 512.0**0.5 + torch.randn(512, 512) * 0.01
                self.res.append(EqualLinear(512, 512)) 
                self.res[-1].weight.data = torch.eye(512) * 512.0**0.5 + torch.randn(512, 512) * 0.01   
        self.res.append(EqualLinear(512, 512)) # for to_rgb7
        self.res[-1].weight.data = torch.eye(512) * 512.0**0.5 + torch.randn(512, 512) * 0.01             
        self.size = self.generator.size
        self.style_dim = self.generator.style_dim
        self.log_size = self.generator.log_size
        self.num_layers = self.generator.num_layers
        self.n_latent = self.generator.n_latent
        self.channels = self.generator.channels
        
    def forward(
        self,
        styles, # intrinsic style code
        exstyles, # extrinsic style code
        return_latents=False, 
        return_feat=False,
        inject_index=None,
        truncation=1,
        truncation_latent=None,
        input_is_latent=False,
        noise=None,
        randomize_noise=True,
        z_plus_latent=False, # intrinsic style code is z+ or z
        use_res=True,        # whether to use the extrinsic style path
        fuse_index=18,       # layers > fuse_index do not use the extrinsic style path
        interp_weights=[1]*18, # weight vector for style combination of two paths
    ):

        if not input_is_latent:
            if not z_plus_latent:
                styles = [self.generator.style(s) for s in styles]
            else:
                styles = [self.generator.style(s.reshape(s.shape[0]*s.shape[1], s.shape[2])).reshape(s.shape) for s in styles]

        if noise is None:
            if randomize_noise:
                noise = [None] * self.generator.num_layers
            else:
                noise = [
                    getattr(self.generator.noises, f"noise_{i}") for i in range(self.generator.num_layers)
                ]

        if truncation < 1:
            style_t = []

            for style in styles:
                style_t.append(
                    truncation_latent + truncation * (style - truncation_latent)
                )

            styles = style_t
        
        if len(styles) < 2:
            inject_index = self.generator.n_latent

            if styles[0].ndim < 3:
                latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1)

            else:
                latent = styles[0]

        else:
            if inject_index is None:
                inject_index = random.randint(1, self.generator.n_latent - 1)

            if styles[0].ndim < 3:
                latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1)
                latent2 = styles[1].unsqueeze(1).repeat(1, self.generator.n_latent - inject_index, 1)

                latent = torch.cat([latent, latent2], 1)
            else:
                latent = torch.cat([styles[0][:,0:inject_index], styles[1][:,inject_index:]], 1)
            
        if use_res:
            if exstyles.ndim < 3:
                resstyles = self.style(exstyles).unsqueeze(1).repeat(1, self.generator.n_latent, 1)
                adastyles = exstyles.unsqueeze(1).repeat(1, self.generator.n_latent, 1)
            else:
                nB, nL, nD = exstyles.shape
                resstyles = self.style(exstyles.reshape(nB*nL, nD)).reshape(nB, nL, nD)
                adastyles = exstyles
        
        out = self.generator.input(latent)
        out = self.generator.conv1(out, latent[:, 0], noise=noise[0])
        if use_res and fuse_index > 0:
            out = self.res[0](out, resstyles[:, 0], interp_weights[0])
        
        skip = self.generator.to_rgb1(out, latent[:, 1])
        i = 1
        for conv1, conv2, noise1, noise2, to_rgb in zip(
            self.generator.convs[::2], self.generator.convs[1::2], noise[1::2], noise[2::2], self.generator.to_rgbs):
            if use_res and fuse_index >= i and i > self.res_index:
                out = conv1(out, interp_weights[i] * self.res[i](adastyles[:, i]) + 
                            (1-interp_weights[i]) * latent[:, i], noise=noise1)
            else:
                out = conv1(out, latent[:, i], noise=noise1)
            if use_res and fuse_index >= i and i <= self.res_index:
                out = self.res[i](out, resstyles[:, i], interp_weights[i])
            if use_res and fuse_index >= (i+1) and i > self.res_index:
                out = conv2(out, interp_weights[i+1] * self.res[i+1](adastyles[:, i+1]) + 
                            (1-interp_weights[i+1]) * latent[:, i+1], noise=noise2)
            else:
                out = conv2(out, latent[:, i + 1], noise=noise2)
            if use_res and fuse_index >= (i+1) and i <= self.res_index:
                out = self.res[i+1](out, resstyles[:, i+1], interp_weights[i+1])   
            if use_res and fuse_index >= (i+2) and i >= self.res_index-1:
                skip = to_rgb(out, interp_weights[i+2] * self.res[i+2](adastyles[:, i+2]) +
                              (1-interp_weights[i+2]) * latent[:, i + 2], skip)
            else:
                skip = to_rgb(out, latent[:, i + 2], skip)
            i += 2
            if i > self.res_index and return_feat:
                return out, skip

        image = skip

        if return_latents:
            return image, latent

        else:
            return image, None  
        
    def make_noise(self):
        return self.generator.make_noise()

    def mean_latent(self, n_latent):
        return self.generator.mean_latent(n_latent)

    def get_latent(self, input):
        return self.generator.style(input)