models/style_based_pix2pixII_model.py

import copy
import clip
import os
import random
import torch
import torch.nn as nn
import torch.nn.functional as F
from enum import Enum
from PIL import Image
from torch import autograd
from .base_model import BaseModel
from models.modules import networks
from models.modules.stylegan2.model import Generator, Discriminator, StyledConv, ToRGB, EqualLinear, ResBlock, ConvLayer, PixelNorm
from models.modules.stylegan2.op import conv2d_gradfix
from models.modules.stylegan2.non_leaking import augment
from models.modules.vit.losses import LossG


class TrainingPhase(Enum):
    ENCODER = 1
    BASE_MODEL = 2
    CLIP_MAPPING = 3
    FEW_SHOT = 4


class CLIPFeats2Wplus(nn.Module):

    def __init__(self, n_tokens=16, embedding_dim=512):
        super().__init__()

        self.position_embedding = nn.Parameter(embedding_dim ** -0.5 * torch.randn(n_tokens, embedding_dim))
        self.transformer = nn.TransformerEncoder(nn.TransformerEncoderLayer(d_model=embedding_dim, nhead=8, norm_first=True), num_layers=4)

    def forward(self, x):
        x_in = x.view(x.shape[0], 1, x.shape[1]) + self.position_embedding
        return F.leaky_relu(self.transformer(x_in.permute(1, 0, 2)), negative_slope=0.2)


class Stylizer(nn.Module):

    def __init__(self, ngf=64, phase=TrainingPhase.ENCODER, model_weights=None):
        super(Stylizer, self).__init__()

        # encoder
        self.encoder = nn.Sequential(
            ConvLayer(3, ngf, 3),          # 512
            ResBlock(ngf * 1, ngf * 1),    # 256
            ResBlock(ngf * 1, ngf * 2),    # 128
            ResBlock(ngf * 2, ngf * 4),    # 64
            ResBlock(ngf * 4, ngf * 8),    # 32
            ConvLayer(ngf * 8, ngf * 8, 3) # 32
        )

        # mapping network
        self.mapping_z = nn.Sequential(*([ PixelNorm() ] + [ EqualLinear(512, 512, activation='fused_lrelu', lr_mul=0.01) for _ in range(8) ]))

        # style-based decoder
        channels = {
            32 : ngf * 8,
            64 : ngf * 8,
            128: ngf * 4,
            256: ngf * 2,
            512: ngf * 1
        }
        self.decoder0 = StyledConv(channels[32], channels[32], 3, 512)
        self.to_rgb0 = ToRGB(channels[32], 512, upsample=False)
        for i in range(4):
            ichan = channels[2 ** (i + 5)]
            ochan = channels[2 ** (i + 6)]
            setattr(self, f'decoder{i + 1}a', StyledConv(ichan, ochan, 3, 512, upsample=True))
            setattr(self, f'decoder{i + 1}b', StyledConv(ochan, ochan, 3, 512))
            setattr(self, f'to_rgb{i + 1}', ToRGB(ochan, 512))
        self.n_latent = 10

        # random style for testing
        self.test_z = torch.randn(1, 512)

        # load pretrained model weights
        if phase == TrainingPhase.ENCODER:
            # load pretrained stylegan2 and freeze these params
            for param in self.mapping_z.parameters():
                param.requires_grad = False
            for i in range(4):
                for key in [f'decoder{i + 1}a', f'decoder{i + 1}b', f'to_rgb{i + 1}']:
                    for param in getattr(self, key).parameters():
                        param.requires_grad = False
            self.load_state_dict(self._convert_stylegan2_dict(model_weights), strict=False)
        elif phase == TrainingPhase.BASE_MODEL:
            # load pretrained encoder and stylegan2 decoder
            self.load_state_dict(model_weights)
        elif phase == TrainingPhase.CLIP_MAPPING:
            self.clip_mapper = CLIPFeats2Wplus(n_tokens=self.n_latent)
            # load pretraned base model and freeze all params except clip mapper
            self.load_state_dict(model_weights, strict=False)
            params = dict(self.named_parameters())
            for k in params.keys():
                if 'clip_mapper' in k:
                    print(f'{k} not freezed !')
                    continue
                params[k].requires_grad = False
        elif phase == TrainingPhase.FEW_SHOT:
            self.clip_mapper = CLIPFeats2Wplus(n_tokens=self.n_latent)
            # load pretrained base model and freeze encoder & mapping
            self.load_state_dict(model_weights)
            self.encoder.requires_grad_(False)
            self.mapping_z.requires_grad_(False)
            self.clip_mapper.requires_grad_(False)

    def _convert_stylegan2_dict(self, src):
        res = {}
        for k, v in src.items():
            if k.startswith('style.'):
                res[k.replace('style.', 'mapping_z.')] = v
            else:
                name, idx = k.split('.')[:2]
                if name == 'convs':
                    idx = int(idx)
                    if idx >= 6:
                        res[k.replace(f'{name}.{idx}.', f'decoder{idx // 2 - 2}{chr(97 + idx % 2)}.')] = v
                elif name == 'to_rgbs':
                    idx = int(idx)
                    if idx >= 3:
                        res[k.replace(f'{name}.{idx}.', f'to_rgb{idx - 2}.')] = v
        return res

    def get_styles(self, x, **kwargs):
        if len(kwargs) == 0:
            return self.mapping_z(self.test_z.to(x.device).repeat(x.shape[0], 1)).repeat(self.n_latent, 1, 1)
        elif 'mixing' in kwargs and kwargs['mixing']:
            w0 = self.mapping_z(torch.randn(x.shape[0], 512, device=x.device))
            w1 = self.mapping_z(torch.randn(x.shape[0], 512, device=x.device))
            inject_index = random.randint(1, self.n_latent - 1)
            return torch.cat([
                w0.repeat(inject_index, 1, 1),
                w1.repeat(self.n_latent - inject_index, 1, 1)
            ])
        elif 'z' in kwargs:
            return self.mapping_z(kwargs['z']).repeat(self.n_latent, 1, 1)
        elif 'clip_feats' in kwargs:
            return self.clip_mapper(kwargs['clip_feats'])
        else:
            z = torch.randn(x.shape[0], 512, device=x.device)
            return self.mapping_z(z).repeat(self.n_latent, 1, 1)

    def forward(self, x, **kwargs):
        # encode
        feat = self.encoder(x)

        # get style code
        styles = self.get_styles(x, **kwargs)

        # style-based generate
        feat = self.decoder0(feat, styles[0])
        out = self.to_rgb0(feat, styles[1])
        for i in range(4):
            feat = getattr(self, f'decoder{i + 1}a')(feat, styles[i * 2 + 1])
            feat = getattr(self, f'decoder{i + 1}b')(feat, styles[i * 2 + 2])
            out = getattr(self, f'to_rgb{i + 1}')(feat, styles[i * 2 + 3], out)

        return F.hardtanh(out)


class StyleBasedPix2PixIIModel(BaseModel):
    """
    This class implements the Style-Based Pix2Pix model version II.
    """

    def __init__(self, config, DDP_device=None):
        BaseModel.__init__(self, config, DDP_device=DDP_device)

        self.d_reg_freq = 16
        self.lambda_r1 = 10
        self.step = 0
        self.phase = TrainingPhase(config['training']['phase'])

        # specify the training losses you want to print out. The training/test scripts will call <BaseModel.get_current_losses>
        if self.phase == TrainingPhase.ENCODER:
            self.loss_names = ['G', 'G_L1', 'G_Feat']
        elif self.phase == TrainingPhase.BASE_MODEL:
            self.loss_names = ['G', 'G_ST', 'G_GAN', 'D']
        elif self.phase == TrainingPhase.CLIP_MAPPING:
            self.loss_names = ['G', 'G_L1', 'G_Feat']
        elif self.phase == TrainingPhase.FEW_SHOT:
            self.loss_names = ['G', 'G_ST', 'G_CLIP', 'G_PROJ']

        # specify the images you want to save/display. The training/test scripts will call <BaseModel.get_current_visuals>
        self.visual_names = ['real_A', 'real_B', 'fake_B']

        # specify the models you want to save to the disk. The training/test scripts will call <BaseModel.save_networks> and <BaseModel.load_networks>.
        if self.isTrain:
            self.model_names = ['G', 'G_ema', 'D']
        else:  # during test time, only load Gs
            self.model_names = ['G_ema']

        self.data_aug_prob = config['training']['data_aug_prob']

        min_feats_size = tuple(config['model']['min_feats_size'])

        def __init_net(model):
            return networks.init_net(model, init_type='none', init_gain=0.0, gpu_ids=self.gpu_ids,
                DDP_device=self.DDP_device, find_unused_parameters=config['training']['find_unused_parameters'])

        if self.phase == TrainingPhase.ENCODER: # train a encoder for stylegan2
            # load and init pretrained stylegan2
            model_dict = torch.load(config['training']['pretrained_model'], map_location='cpu')
            self.stylegan2 = Generator(512, 512, 8)
            self.stylegan2.load_state_dict(model_dict['g'])
            self.stylegan2 = __init_net(self.stylegan2)
            self.stylegan2.eval()
            self.stylegan2.requires_grad_(False)

            # init netG
            self.netG = Stylizer(ngf=config['model']['ngf'], phase=self.phase, model_weights=model_dict['g'])
            self.netG = __init_net(self.netG)

            # init netD
            self.netD = Discriminator(min(min_feats_size) * 128, min_feats_size)
            self.netD.load_state_dict(model_dict['d'])
            self.netD = __init_net(self.netD)
            self.netD.eval()
            self.netD.requires_grad_(False)

        elif self.phase == TrainingPhase.BASE_MODEL: # finetune the whole model
            model_dict = torch.load(config['training']['pretrained_model'], map_location='cpu')

            # init netG
            self.netG = Stylizer(ngf=config['model']['ngf'], phase=self.phase, model_weights=model_dict['G_ema_model'])
            self.netG = __init_net(self.netG)

            # init netD
            self.netD = Discriminator(min(min_feats_size) * 128, min_feats_size)
            self.netD.load_state_dict(model_dict['D_model'])
            self.netD = __init_net(self.netD)

        elif self.phase == TrainingPhase.CLIP_MAPPING or self.phase == TrainingPhase.FEW_SHOT: # train the clip mapper or zero/one shot finetune
            # init CLIP
            self.clip_model, self.pil_to_tensor = clip.load('ViT-B/32', device=self.device)
            self.clip_model.eval()
            self.clip_model.requires_grad_(False)

            model_dict = torch.load(config['training']['pretrained_model'], map_location='cpu')

            # init netG
            self.netG = Stylizer(ngf=config['model']['ngf'], phase=self.phase, model_weights=model_dict['G_ema_model'])
            self.netG = __init_net(self.netG)

            # init netD
            self.netD = Discriminator(min(min_feats_size) * 128, min_feats_size)
            self.netD.load_state_dict(model_dict['D_model'])
            self.netD = __init_net(self.netD)
            self.netD.eval()
            self.netD.requires_grad_(False)

        if self.phase == TrainingPhase.FEW_SHOT: # set hook to get clip vit tokens
            def clip_vit_hook(model, feat_in, feat_out):
                self.clip_vit_tokens = feat_out[1:].permute(1, 0, 2).float() # remove cls token
            self.clip_model.visual.transformer.resblocks[3].register_forward_hook(clip_vit_hook)

        # create netG ema
        self.netG_ema = copy.deepcopy(self.netG)
        self.netG_ema.eval()
        self.netG_ema.requires_grad_(False)
        self.ema(self.netG_ema, self.netG, 0.0)

        # CLIP mean & std
        self.clip_mean = torch.tensor((0.48145466, 0.4578275, 0.40821073), device=self.device).view(1, 3, 1, 1)
        self.clip_std = torch.tensor((0.26862954, 0.26130258, 0.27577711), device=self.device).view(1, 3, 1, 1)

        if self.isTrain:
            # define loss functions
            if self.phase == TrainingPhase.ENCODER:
                self.criterionL1 = nn.L1Loss()
            elif self.phase == TrainingPhase.BASE_MODEL:
                self.criterionStyleGAN = networks.GANLoss('wgangp').to(self.device)
                self.vitLoss = LossG(self.device)
            elif self.phase == TrainingPhase.CLIP_MAPPING:
                self.criterionL1 = nn.L1Loss()
            elif self.phase == TrainingPhase.FEW_SHOT:
                self.criterionL1 = nn.L1Loss()
                self.vitLoss = LossG(self.device)
                self.cosineSim = nn.CosineSimilarity(dim=1)

            # initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>.
            self.optimizer_G = torch.optim.Adam(self.netG.parameters(), lr=config['training']['lr'], betas=(config['training']['beta1'], 0.999))
            d_reg_ratio = self.d_reg_freq / (self.d_reg_freq + 1)
            self.optimizer_D = torch.optim.Adam(self.netD.parameters(), lr=config['training']['lr'] * d_reg_ratio, betas=(config['training']['beta1'] ** d_reg_ratio, 0.999 ** d_reg_ratio))
            self.optimizers.append(self.optimizer_G)
            self.optimizers.append(self.optimizer_D)

    def ema(self, tgt, src, decay=0.999):
        param_tgt = dict(tgt.named_parameters())
        param_src = dict(src.named_parameters())
        for key in param_tgt.keys():
            param_tgt[key].data.mul_(decay).add_(param_src[key].data, alpha=1.0 - decay)

    def preprocess_clip_image(self, x, size):
        x = x * 0.5 + 0.5
        x = F.interpolate(x, (size, size), mode='bilinear', antialias=True, align_corners=False)
        return (x - self.clip_mean) / self.clip_std

    def set_input(self, input):
        if self.phase == TrainingPhase.ENCODER:
            # sample via stylegan2
            self.z = torch.randn(self.config['dataset']['batch_size'], 512, device=self.device)
            with torch.no_grad():
                self.real_A = F.hardtanh(self.stylegan2.forward([self.z])[0])
                self.real_B = self.real_A.clone()
        elif self.phase == TrainingPhase.BASE_MODEL:
            if self.config['common']['phase'] == 'test':
                self.real_A = input['test_A'].to(self.device)
                self.real_B = input['test_B'].to(self.device)
                self.image_paths = input['test_A_path']
            else:
                self.real_A = input['unpaired_A'].to(self.device)
                self.real_B = input['unpaired_B'].to(self.device)
                self.image_paths = input['unpaired_A_path']
        elif self.phase == TrainingPhase.CLIP_MAPPING:
            self.real_A = input['unpaired_A'].to(self.device)
            with torch.no_grad():
                self.real_B = self.netG_ema(self.real_A, mixing=random.random() < self.config['training']['style_mixing_prob'])
                self.clip_feats = self.clip_model.encode_image(self.preprocess_clip_image(self.real_B, self.clip_model.visual.input_resolution))
                self.clip_feats /= self.clip_feats.norm(dim=1, keepdim=True)
        elif self.phase == TrainingPhase.FEW_SHOT:
            self.real_A = input['unpaired_A'].to(self.device)
            self.real_B = self.real_A
            if not hasattr(self, 'clip_feats'):
                with torch.no_grad():
                    if os.path.isfile(self.config['training']['image_prompt']):
                        image = self.pil_to_tensor(Image.open(self.config['training']['image_prompt'])).unsqueeze(0).to(self.device)
                        self.clip_feats = self.clip_model.encode_image(image)
                        ref_tokens = self.clip_vit_tokens
                        ref_tokens /= ref_tokens.norm(dim=2, keepdim=True)
                        D = ref_tokens.shape[2]
                        ref_tokens = ref_tokens.reshape(-1, D).permute(1, 0)
                        U, _, _ = torch.linalg.svd(ref_tokens, full_matrices=False)
                        self.UUT = U @ U.permute(1, 0)
                        self.use_image_prompt = True
                    else:
                        text = clip.tokenize(self.config['training']['text_prompt']).to(self.device)
                        self.clip_feats = self.clip_model.encode_text(text)
                        self.use_image_prompt = False
                        # get source text prompt feature
                        text = clip.tokenize(self.config['training']['src_text_prompt']).to(self.device)
                        self.src_clip_feats = self.clip_model.encode_text(text)
                        self.src_clip_feats /= self.src_clip_feats.norm(dim=1, keepdim=True)
                        self.src_clip_feats = self.src_clip_feats.repeat(self.config['dataset']['batch_size'], 1)
                    self.clip_feats /= self.clip_feats.norm(dim=1, keepdim=True)
                    self.clip_feats = self.clip_feats.repeat(self.config['dataset']['batch_size'], 1)
            # get direction in clip space
            with torch.no_grad():
                self.real_A_clip_feats = self.clip_model.encode_image(self.preprocess_clip_image(self.real_A, self.clip_model.visual.input_resolution))
                self.real_A_clip_feats /= self.real_A_clip_feats.norm(dim=1, keepdim=True)
                if self.use_image_prompt:
                    self.src_clip_feats = self.real_A_clip_feats.mean(dim=0, keepdim=True).repeat(self.config['dataset']['batch_size'], 1)
                self.clip_feats_dir = self.clip_feats - self.src_clip_feats

    def forward(self, use_ema=False):
        if self.phase == TrainingPhase.ENCODER:
            if use_ema:
                self.fake_B = self.netG_ema(self.real_A, z=self.z)
            else:
                self.fake_B = self.netG(self.real_A, z=self.z)
        elif self.phase == TrainingPhase.BASE_MODEL:
            if not self.isTrain:
                self.fake_B = self.netG_ema(self.real_A, mixing=False)
            elif use_ema:
                self.fake_B = self.netG_ema(self.real_A, mixing=random.random() < self.config['training']['style_mixing_prob'])
            else:
                self.fake_B = self.netG(self.real_A, mixing=random.random() < self.config['training']['style_mixing_prob'])
        elif self.phase == TrainingPhase.CLIP_MAPPING or self.phase == TrainingPhase.FEW_SHOT:
            if use_ema:
                self.fake_B = self.netG_ema(self.real_A, clip_feats=self.clip_feats)
            else:
                self.fake_B = self.netG(self.real_A, clip_feats=self.clip_feats)

    def backward_D_r1(self):
        self.real_B.requires_grad = True
        if self.data_aug_prob == 0.0:
            real_aug = self.real_B
        else:
            real_aug, _ = augment(self.real_B, self.data_aug_prob)
        real_pred = self.netD(real_aug)
        with conv2d_gradfix.no_weight_gradients():
            grad, = autograd.grad(outputs=real_pred.sum(), inputs=real_aug, create_graph=True)
        r1_loss = grad.pow(2).reshape(grad.shape[0], -1).sum(1).mean()
        (r1_loss * self.lambda_r1 / 2 * self.d_reg_freq + 0 * real_pred[0]).backward()

    def backward_D(self, backward=True):
        if self.data_aug_prob == 0.0:
            loss_fake = self.criterionStyleGAN(self.netD(self.fake_B.detach()), False)
            loss_real = self.criterionStyleGAN(self.netD(self.real_B), True)
        else:
            fake_aug, _ = augment(self.fake_B.detach(), self.data_aug_prob)
            real_aug, _ = augment(self.real_B, self.data_aug_prob)
            loss_fake = self.criterionStyleGAN(self.netD(fake_aug), False)
            loss_real = self.criterionStyleGAN(self.netD(real_aug), True)
        self.loss_D = (loss_fake + loss_real) * 0.5

        if backward:
            self.loss_D.backward()

    def backward_G(self, backward=True):
        self.loss_G = 0

        if self.phase == TrainingPhase.ENCODER or self.phase == TrainingPhase.CLIP_MAPPING:
            self.loss_G_L1 = self.criterionL1(self.fake_B, self.real_B)
            with torch.no_grad():
                real_feats = self.netD(self.real_B, rtn_feats=True)
            fake_feats = self.netD(self.fake_B, rtn_feats=True)
            self.loss_G_Feat = sum([ self.criterionL1(fake, real) for fake, real in zip(fake_feats, real_feats) ])

            self.loss_G += self.loss_G_L1 * self.config['training']['lambda_L1']
            self.loss_G += self.loss_G_Feat * self.config['training']['lambda_Feat']

        elif self.phase == TrainingPhase.BASE_MODEL:
            self.loss_G_ST = self.vitLoss.calculate_global_ssim_loss(self.fake_B * 0.5 + 0.5, self.real_A * 0.5 + 0.5)
            if self.data_aug_prob == 0.0:
                self.loss_G_GAN = self.criterionStyleGAN(self.netD(self.fake_B), True)
            else:
                fake_aug, _ = augment(self.fake_B, self.data_aug_prob)
                self.loss_G_GAN = self.criterionStyleGAN(self.netD(fake_aug), True)

            self.loss_G += self.loss_G_ST * self.config['training']['lambda_ST']
            self.loss_G += self.loss_G_GAN * self.config['training']['lambda_GAN']

        elif self.phase == TrainingPhase.FEW_SHOT:
            self.loss_G_ST = self.vitLoss.calculate_global_ssim_loss(self.fake_B * 0.5 + 0.5, self.real_A * 0.5 + 0.5)
            fake_clip_feats = self.clip_model.encode_image(self.preprocess_clip_image(self.fake_B, self.clip_model.visual.input_resolution))
            fake_clip_feats = fake_clip_feats / fake_clip_feats.norm(dim=1, keepdim=True)
            fake_clip_feats_dir = fake_clip_feats - self.real_A_clip_feats
            self.loss_G_CLIP = (1.0 - self.cosineSim(fake_clip_feats_dir, self.clip_feats_dir)).mean()
            if self.use_image_prompt:
                fake_tokens = self.clip_vit_tokens
                fake_tokens = fake_tokens / fake_tokens.norm(dim=2, keepdim=True)
                D = fake_tokens.shape[2]
                fake_tokens = fake_tokens.reshape(-1, D).permute(1, 0)
                self.loss_G_PROJ = self.criterionL1(self.UUT @ fake_tokens, fake_tokens)
            else:
                self.loss_G_PROJ = 0.0

            self.loss_G += self.loss_G_ST * self.config['training']['lambda_ST']
            self.loss_G += self.loss_G_CLIP * self.config['training']['lambda_CLIP']
            self.loss_G += self.loss_G_PROJ * self.config['training']['lambda_PROJ']

        if backward:
            self.loss_G.backward()

    def optimize_parameters(self):
        # forward
        self.forward()

        if not self.phase == TrainingPhase.BASE_MODEL:
            # only G
            self.optimizer_G.zero_grad()
            self.backward_G()
            self.optimizer_G.step()

            # update G_ema
            self.ema(self.netG_ema, self.netG, decay=self.config['training']['ema'])

        else:
            # G
            self.set_requires_grad([self.netD], False)
            self.optimizer_G.zero_grad()
            self.backward_G()
            self.optimizer_G.step()
            # D
            self.set_requires_grad([self.netD], True)
            self.optimizer_D.zero_grad()
            self.backward_D()
            self.optimizer_D.step()

            # update G_ema
            self.ema(self.netG_ema, self.netG, decay=self.config['training']['ema'])

            # r1 reg
            self.step += 1
            if self.step % self.d_reg_freq == 0:
                self.optimizer_D.zero_grad()
                self.backward_D_r1()
                self.optimizer_D.step()

    def eval_step(self):
        self.forward(use_ema=True)
        self.backward_G(False)
        if self.phase == TrainingPhase.BASE_MODEL:
            self.backward_D(False)
        self.step += 1

    def trace_jit(self, input):
        self.netG = self.netG.module.cpu()
        traced_script_module = torch.jit.trace(self.netG, input)
        dummy_output = self.netG_ema(input)
        dummy_output_traced = traced_script_module(input)
        return traced_script_module, dummy_output, dummy_output_traced