lucidrains.py

import json, time, pickle, argparse
from math import floor, log2, sqrt
from random import random
from shutil import rmtree
from functools import partial
from datetime import datetime
import multiprocessing
from PIL import Image
from pathlib import Path
from retry.api import retry_call
from tqdm import tqdm

import numpy as np
from scipy import linalg

import torch
from torch import nn
from torch.utils import data
import torch.nn.functional as F

from torch_optimizer import DiffGrad
from torch.autograd import grad as torch_grad

import torchvision
from torchvision import transforms

from vector_quantize_pytorch import VectorQuantize
from linear_attention_transformer import ImageLinearAttention
from contrastive_learner import ContrastiveLearner, RandomApply

from kornia import augmentation as augs
from kornia import filters

import validation
from validation.inception import InceptionV3
from validation import lpips

import wandb

try:
    from apex import amp

    APEX_AVAILABLE = True
except:
    APEX_AVAILABLE = False

assert torch.cuda.is_available(), "You need to have an Nvidia GPU with CUDA installed."

num_cores = multiprocessing.cpu_count()

# constants

EXTS = ["jpg", "png"]
EPS = 1e-8

# helper classes


class NanException(Exception):
    pass


class EMA:
    def __init__(self, beta):
        super().__init__()
        self.beta = beta

    def update_average(self, old, new):
        if old is None:
            return new
        return old * self.beta + (1 - self.beta) * new


class Flatten(nn.Module):
    def forward(self, x):
        return x.reshape(x.shape[0], -1)


class Residual(nn.Module):
    def __init__(self, fn):
        super().__init__()
        self.fn = fn

    def forward(self, x):
        return self.fn(x) + x


class Rezero(nn.Module):
    def __init__(self, fn):
        super().__init__()
        self.fn = fn
        self.g = nn.Parameter(torch.zeros(1))

    def forward(self, x):
        return self.fn(x) * self.g


class PermuteToFrom(nn.Module):
    def __init__(self, fn):
        super().__init__()
        self.fn = fn

    def forward(self, x):
        x = x.permute(0, 2, 3, 1)
        out, loss = self.fn(x)
        out = out.permute(0, 3, 1, 2)
        return out, loss


# helpers


def default(value, d):
    return d if value is None else value


def cycle(iterable):
    while True:
        for i in iterable:
            yield i


def cast_list(el):
    return el if isinstance(el, list) else [el]


def is_empty(t):
    if isinstance(t, torch.Tensor):
        return t.nelement() == 0
    return t is None


def raise_if_nan(t):
    if torch.isnan(t):
        raise NanException


def loss_backwards(fp16, loss, optimizer, **kwargs):
    if fp16:
        with amp.scale_loss(loss, optimizer) as scaled_loss:
            scaled_loss.backward(**kwargs)
    else:
        loss.backward(**kwargs)


def gradient_penalty(images, output, weight=10):
    batch_size = images.shape[0]
    gradients = torch_grad(
        outputs=output,
        inputs=images,
        grad_outputs=torch.ones(output.size()).cuda(),
        create_graph=True,
        retain_graph=True,
        only_inputs=True,
    )[0]

    gradients = gradients.view(batch_size, -1)
    return weight * ((gradients.norm(2, dim=1) - 1) ** 2).mean()


def noise(n, latent_dim):
    return torch.randn(n, latent_dim).cuda()


def noise_list(n, layers, latent_dim):
    return [(noise(n, latent_dim), layers)]


def mixed_list(n, layers, latent_dim):
    tt = int(torch.rand(()).numpy() * layers)
    return noise_list(n, tt, latent_dim) + noise_list(n, layers - tt, latent_dim)


def latent_to_w(style_vectorizer, latent_descr):
    return [(style_vectorizer(z), num_layers) for z, num_layers in latent_descr]


def image_noise(n, im_size):
    return torch.FloatTensor(n, im_size, im_size, 1).uniform_(0.0, 1.0).cuda()


def leaky_relu(p=0.2):
    return nn.LeakyReLU(p, inplace=True)


def evaluate_in_chunks(max_batch_size, model, *args):
    split_args = list(zip(*list(map(lambda x: x.split(max_batch_size, dim=0), args))))
    chunked_outputs = [model(*i) for i in split_args]
    if len(chunked_outputs) == 1:
        return chunked_outputs[0]
    return torch.cat(chunked_outputs, dim=0)


def styles_def_to_tensor(styles_def):
    return torch.cat([t[:, None, :].expand(-1, n, -1) for t, n in styles_def], dim=1)


def set_requires_grad(model, bool):
    for p in model.parameters():
        p.requires_grad = bool


# dataset


def convert_rgb_to_transparent(image):
    if image.mode == "RGB":
        return image.convert("RGBA")
    return image


def convert_transparent_to_rgb(image):
    if image.mode == "RGBA":
        return image.convert("RGB")
    return image


class expand_greyscale(object):
    def __init__(self, num_channels):
        self.num_channels = num_channels

    def __call__(self, tensor):
        return tensor.expand(self.num_channels, -1, -1)


def resize_to_minimum_size(min_size, image):
    if max(*image.size) < min_size:
        return torchvision.transforms.functional.resize(image, min_size)
    return image


class Dataset(data.Dataset):
    def __init__(self, folder, image_size, transparent=False):
        super().__init__()
        self.folder = folder
        self.image_size = image_size
        self.paths = [p for ext in EXTS for p in Path(f"{folder}").glob(f"**/*.{ext}")]

        convert_image_fn = convert_transparent_to_rgb if not transparent else convert_rgb_to_transparent
        num_channels = 3 if not transparent else 4

        self.transform = transforms.Compose(
            [
                transforms.Lambda(convert_image_fn),
                transforms.Lambda(partial(resize_to_minimum_size, image_size)),
                transforms.RandomHorizontalFlip(),
                transforms.Resize(image_size),
                transforms.CenterCrop(image_size),
                transforms.ToTensor(),
                transforms.Lambda(expand_greyscale(num_channels)),
            ]
        )

    def __len__(self):
        return len(self.paths)

    def __getitem__(self, index):
        path = self.paths[index]
        img = Image.open(path)
        return self.transform(img)


# stylegan2 classes


class StyleVectorizer(nn.Module):
    def __init__(self, emb, depth):
        super().__init__()

        layers = []
        for i in range(depth):
            layers.extend([nn.Linear(emb, emb), leaky_relu()])

        self.net = nn.Sequential(*layers)

    def forward(self, x):
        return self.net(x)


class RGBBlock(nn.Module):
    def __init__(self, latent_dim, input_channel, upsample, rgba=False):
        super().__init__()
        self.input_channel = input_channel
        self.to_style = nn.Linear(latent_dim, input_channel)

        out_filters = 3 if not rgba else 4
        self.conv = Conv2DMod(input_channel, out_filters, 1, demod=False)

        self.upsample = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False) if upsample else None

    def forward(self, x, prev_rgb, istyle):
        b, c, h, w = x.shape
        style = self.to_style(istyle)
        x = self.conv(x, style)

        if prev_rgb is not None:
            x = x + prev_rgb

        if self.upsample is not None:
            x = self.upsample(x)

        return x


class Conv2DMod(nn.Module):
    def __init__(self, in_chan, out_chan, kernel, demod=True, stride=1, dilation=1, **kwargs):
        super().__init__()
        self.filters = out_chan
        self.demod = demod
        self.kernel = kernel
        self.stride = stride
        self.dilation = dilation
        self.weight = nn.Parameter(torch.randn((out_chan, in_chan, kernel, kernel)))
        nn.init.kaiming_normal_(self.weight, a=0, mode="fan_in", nonlinearity="leaky_relu")

    def _get_same_padding(self, size, kernel, dilation, stride):
        return ((size - 1) * (stride - 1) + dilation * (kernel - 1)) // 2

    def forward(self, x, y):
        b, c, h, w = x.shape

        w1 = y[:, None, :, None, None]
        w2 = self.weight[None, :, :, :, :]
        weights = w2 * (w1 + 1)

        if self.demod:
            d = torch.rsqrt((weights ** 2).sum(dim=(2, 3, 4), keepdim=True) + EPS)
            weights = weights * d

        x = x.reshape(1, -1, h, w)

        _, _, *ws = weights.shape
        weights = weights.reshape(b * self.filters, *ws)

        padding = self._get_same_padding(h, self.kernel, self.dilation, self.stride)
        x = F.conv2d(x, weights, padding=padding, groups=b)

        x = x.reshape(-1, self.filters, h, w)
        return x


class GeneratorBlock(nn.Module):
    def __init__(self, latent_dim, input_channels, filters, upsample=True, upsample_rgb=True, rgba=False):
        super().__init__()
        self.upsample = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False) if upsample else None

        self.to_style1 = nn.Linear(latent_dim, input_channels)
        self.to_noise1 = nn.Linear(1, filters)
        self.conv1 = Conv2DMod(input_channels, filters, 3)

        self.to_style2 = nn.Linear(latent_dim, filters)
        self.to_noise2 = nn.Linear(1, filters)
        self.conv2 = Conv2DMod(filters, filters, 3)

        self.activation = leaky_relu()
        self.to_rgb = RGBBlock(latent_dim, filters, upsample_rgb, rgba)

    def forward(self, x, prev_rgb, istyle, inoise):
        if self.upsample is not None:
            x = self.upsample(x)

        inoise = inoise[:, : x.shape[2], : x.shape[3], :]
        noise1 = self.to_noise1(inoise).permute((0, 3, 2, 1))
        noise2 = self.to_noise2(inoise).permute((0, 3, 2, 1))

        style1 = self.to_style1(istyle)
        x = self.conv1(x, style1)
        x = self.activation(x + noise1)

        style2 = self.to_style2(istyle)
        x = self.conv2(x, style2)
        x = self.activation(x + noise2)

        rgb = self.to_rgb(x, prev_rgb, istyle)
        return x, rgb


class DiscriminatorBlock(nn.Module):
    def __init__(self, input_channels, filters, downsample=True):
        super().__init__()
        self.conv_res = nn.Conv2d(input_channels, filters, 1)

        self.net = nn.Sequential(
            nn.Conv2d(input_channels, filters, 3, padding=1),
            leaky_relu(),
            nn.Conv2d(filters, filters, 3, padding=1),
            leaky_relu(),
        )

        self.downsample = nn.Conv2d(filters, filters, 3, padding=1, stride=2) if downsample else None

    def forward(self, x):
        res = self.conv_res(x)
        x = self.net(x)
        x = x + res
        if self.downsample is not None:
            x = self.downsample(x)
        return x


class Generator(nn.Module):
    def __init__(self, image_size, latent_dim, network_capacity=16, transparent=False, attn_layers=[]):
        super().__init__()
        self.image_size = image_size
        self.latent_dim = latent_dim
        self.num_layers = int(log2(image_size) - 1)

        init_channels = 4 * network_capacity
        self.initial_block = nn.Parameter(torch.randn((init_channels, 4, 4)))
        filters = [init_channels] + [network_capacity * (2 ** (i + 1)) for i in range(self.num_layers)][::-1]
        in_out_pairs = zip(filters[0:-1], filters[1:])

        self.blocks = nn.ModuleList([])
        self.attns = nn.ModuleList([])

        for ind, (in_chan, out_chan) in enumerate(in_out_pairs):
            not_first = ind != 0
            not_last = ind != (self.num_layers - 1)
            num_layer = self.num_layers - ind

            attn_fn = (
                nn.Sequential(*[Residual(Rezero(ImageLinearAttention(in_chan))) for _ in range(2)])
                if num_layer in attn_layers
                else None
            )

            self.attns.append(attn_fn)

            block = GeneratorBlock(
                latent_dim, in_chan, out_chan, upsample=not_first, upsample_rgb=not_last, rgba=transparent
            )
            self.blocks.append(block)

    def forward(self, styles, input_noise):
        batch_size = styles.shape[0]
        image_size = self.image_size
        x = self.initial_block.expand(batch_size, -1, -1, -1)
        styles = styles.transpose(0, 1)

        rgb = None
        for style, block, attn in zip(styles, self.blocks, self.attns):
            if attn is not None:
                x = attn(x)
            x, rgb = block(x, rgb, style, input_noise)

        return rgb


class Discriminator(nn.Module):
    def __init__(
        self, image_size, network_capacity=16, fq_layers=[], fq_dict_size=256, attn_layers=[], transparent=False
    ):
        super().__init__()
        num_layers = int(log2(image_size) - 1)
        num_init_filters = 3 if not transparent else 4

        blocks = []
        filters = [num_init_filters] + [(network_capacity) * (2 ** i) for i in range(num_layers + 1)]
        chan_in_out = list(zip(filters[0:-1], filters[1:]))

        blocks = []
        quantize_blocks = []
        attn_blocks = []

        for ind, (in_chan, out_chan) in enumerate(chan_in_out):
            num_layer = ind + 1
            is_not_last = ind != (len(chan_in_out) - 1)

            block = DiscriminatorBlock(in_chan, out_chan, downsample=is_not_last)
            blocks.append(block)

            attn_fn = (
                nn.Sequential(*[Residual(Rezero(ImageLinearAttention(out_chan))) for _ in range(2)])
                if num_layer in attn_layers
                else None
            )

            attn_blocks.append(attn_fn)

            quantize_fn = PermuteToFrom(VectorQuantize(out_chan, fq_dict_size)) if num_layer in fq_layers else None
            quantize_blocks.append(quantize_fn)

        self.blocks = nn.ModuleList(blocks)
        self.attn_blocks = nn.ModuleList(attn_blocks)
        self.quantize_blocks = nn.ModuleList(quantize_blocks)

        latent_dim = 2 * 2 * filters[-1]

        self.flatten = Flatten()
        self.to_logit = nn.Linear(latent_dim, 1)

    def forward(self, x):
        b, *_ = x.shape

        quantize_loss = torch.zeros(1).to(x)

        for (block, attn_block, q_block) in zip(self.blocks, self.attn_blocks, self.quantize_blocks):
            x = block(x)

            if attn_block is not None:
                x = attn_block(x)

            if q_block is not None:
                x, loss = q_block(x)
                quantize_loss += loss

        x = self.flatten(x)
        x = self.to_logit(x)
        return x.squeeze(), quantize_loss


class StyleGAN2(nn.Module):
    def __init__(
        self,
        image_size,
        latent_dim=512,
        style_depth=8,
        network_capacity=16,
        transparent=False,
        fp16=False,
        cl_reg=False,
        augment_fn=None,
        steps=1,
        lr=1e-4,
        fq_layers=[],
        fq_dict_size=256,
        attn_layers=[],
    ):
        super().__init__()
        self.lr = lr
        self.steps = steps
        self.ema_updater = EMA(0.995)

        self.S = StyleVectorizer(latent_dim, style_depth)
        self.G = Generator(image_size, latent_dim, network_capacity, transparent=transparent, attn_layers=attn_layers)
        self.D = Discriminator(
            image_size,
            network_capacity,
            fq_layers=fq_layers,
            fq_dict_size=fq_dict_size,
            attn_layers=attn_layers,
            transparent=transparent,
        )

        self.SE = StyleVectorizer(latent_dim, style_depth)
        self.GE = Generator(image_size, latent_dim, network_capacity, transparent=transparent, attn_layers=attn_layers)

        set_requires_grad(self.SE, False)
        set_requires_grad(self.GE, False)

        generator_params = list(self.G.parameters()) + list(self.S.parameters())
        self.G_opt = DiffGrad(generator_params, lr=self.lr, betas=(0.5, 0.9))
        self.D_opt = DiffGrad(self.D.parameters(), lr=self.lr, betas=(0.5, 0.9))

        self._init_weights()
        self.reset_parameter_averaging()

        self.cuda()

        if fp16:
            (self.S, self.G, self.D, self.SE, self.GE), (self.G_opt, self.D_opt) = amp.initialize(
                [self.S, self.G, self.D, self.SE, self.GE], [self.G_opt, self.D_opt], opt_level="O2"
            )

        # experimental contrastive loss discriminator regularization
        if augment_fn is not None:
            self.augment_fn = augment_fn
        else:
            self.augment_fn = nn.Sequential(
                nn.ReflectionPad2d(int((sqrt(2) - 1) * image_size / 4)),
                RandomApply(augs.ColorJitter(0.8, 0.8, 0.8, 0.2), p=0.7),
                augs.RandomGrayscale(p=0.2),
                augs.RandomHorizontalFlip(),
                RandomApply(augs.RandomAffine(degrees=0, translate=(0.25, 0.25), shear=(15, 15)), p=0.3),
                RandomApply(
                    nn.Sequential(augs.RandomRotation(180), augs.CenterCrop(size=(image_size, image_size))), p=0.2
                ),
                augs.RandomResizedCrop(size=(image_size, image_size)),
                RandomApply(filters.GaussianBlur2d((3, 3), (1.5, 1.5)), p=0.1),
                RandomApply(augs.RandomErasing(), p=0.1),
            )

        self.D_cl = (
            ContrastiveLearner(self.D, image_size, augment_fn=self.augment_fn, fp16=fp16, hidden_layer="flatten")
            if cl_reg
            else None
        )

        # self.S, self.G, self.D, self.SE, self.GE = (
        #     nn.DataParallel(self.S),
        #     nn.DataParallel(self.G),
        #     nn.DataParallel(self.D),
        #     nn.DataParallel(self.SE),
        #     nn.DataParallel(self.GE),
        # )

    def _init_weights(self):
        for m in self.modules():
            if type(m) in {nn.Conv2d, nn.Linear}:
                nn.init.kaiming_normal_(m.weight, a=0, mode="fan_in", nonlinearity="leaky_relu")

        for block in self.G.blocks:
            nn.init.zeros_(block.to_noise1.weight)
            nn.init.zeros_(block.to_noise2.weight)
            nn.init.zeros_(block.to_noise1.bias)
            nn.init.zeros_(block.to_noise2.bias)

    def EMA(self):
        def update_moving_average(ma_model, current_model):
            for current_params, ma_params in zip(current_model.parameters(), ma_model.parameters()):
                old_weight, up_weight = ma_params.data, current_params.data
                ma_params.data = self.ema_updater.update_average(old_weight, up_weight)

        update_moving_average(self.SE, self.S)
        update_moving_average(self.GE, self.G)

    def reset_parameter_averaging(self):
        self.SE.load_state_dict(self.S.state_dict())
        self.GE.load_state_dict(self.G.state_dict())

    def forward(self, x):
        return x


class Trainer:
    def __init__(
        self,
        name,
        results_dir,
        models_dir,
        image_size,
        network_capacity,
        transparent=False,
        batch_size=4,
        mixed_prob=0.9,
        gradient_accumulate_every=1,
        lr=2e-4,
        num_workers=None,
        save_every=1000,
        trunc_psi=0.6,
        fp16=False,
        cl_reg=False,
        fq_layers=[],
        fq_dict_size=256,
        attn_layers=[],
        fid_n_sample=5000,
        ppl_n_sample=2500,
        *args,
        **kwargs,
    ):
        self.GAN_params = [args, kwargs]
        self.GAN = None

        self.name = name
        self.results_dir = Path(results_dir)
        self.models_dir = Path(models_dir)
        self.config_path = self.models_dir / name / ".config.json"

        assert log2(image_size).is_integer(), "image size must be a power of 2 (64, 128, 256, 512, 1024)"
        self.image_size = image_size
        self.network_capacity = network_capacity
        self.transparent = transparent
        self.fq_layers = cast_list(fq_layers)
        self.fq_dict_size = fq_dict_size

        self.attn_layers = cast_list(attn_layers)

        self.lr = lr
        self.batch_size = batch_size
        self.num_workers = num_workers
        self.mixed_prob = mixed_prob

        self.save_every = save_every
        self.steps = 0

        self.av = None
        self.trunc_psi = trunc_psi

        self.pl_mean = 0

        self.fid_n_sample = fid_n_sample
        self.ppl_n_sample = ppl_n_sample

        self.gradient_accumulate_every = gradient_accumulate_every

        assert not fp16 or fp16 and APEX_AVAILABLE, "Apex is not available for you to use mixed precision training"
        self.fp16 = fp16

        self.cl_reg = cl_reg

        self.d_loss = 0
        self.g_loss = 0
        self.last_gp_loss = 0
        self.last_cr_loss = 0
        self.q_loss = 0

        self.pl_length_ma = EMA(0.99)
        self.init_folders()

        self.loader = None

    def init_GAN(self):
        args, kwargs = self.GAN_params
        self.GAN = StyleGAN2(
            lr=self.lr,
            image_size=self.image_size,
            network_capacity=self.network_capacity,
            transparent=self.transparent,
            fq_layers=self.fq_layers,
            fq_dict_size=self.fq_dict_size,
            attn_layers=self.attn_layers,
            fp16=self.fp16,
            cl_reg=self.cl_reg,
            *args,
            **kwargs,
        )

    def write_config(self):
        self.config_path.write_text(json.dumps(self.config()))

    def load_config(self):
        config = self.config() if not self.config_path.exists() else json.loads(self.config_path.read_text())
        self.image_size = config["image_size"]
        self.network_capacity = config["network_capacity"]
        self.transparent = config["transparent"]
        self.fq_layers = config["fq_layers"]
        self.fq_dict_size = config["fq_dict_size"]
        self.attn_layers = config.pop("attn_layers", [])
        del self.GAN
        self.init_GAN()

    def config(self):
        return {
            "image_size": self.image_size,
            "network_capacity": self.network_capacity,
            "transparent": self.transparent,
            "fq_layers": self.fq_layers,
            "fq_dict_size": self.fq_dict_size,
            "attn_layers": self.attn_layers,
        }

    def set_data_src(self, folder):
        self.dataset = Dataset(folder, self.image_size, transparent=self.transparent)
        self.loader = cycle(
            data.DataLoader(
                self.dataset,
                num_workers=default(self.num_workers, num_cores),
                batch_size=self.batch_size,
                drop_last=True,
                shuffle=True,
                pin_memory=True,
            )
        )
        validation.get_dataset_inception_features(self.loader, self.name, self.image_size)

    def train(self):
        assert (
            self.loader is not None
        ), "You must first initialize the data source with `.set_data_src(<folder of images>)`"

        if self.GAN is None:
            self.init_GAN()

        self.GAN.train()
        total_disc_loss = torch.tensor(0.0).cuda()
        total_gen_loss = torch.tensor(0.0).cuda()

        batch_size = self.batch_size

        image_size = self.GAN.G.image_size
        latent_dim = self.GAN.G.latent_dim
        num_layers = self.GAN.G.num_layers

        apply_gradient_penalty = self.steps % 4 == 0
        apply_path_penalty = self.steps % 32 == 0
        apply_cl_reg_to_generated = self.steps > 20000

        log_dict = {"Divergence": 0, "Quantize": 0, "Generator": 0}
        if apply_gradient_penalty:
            log_dict["R1"] = 0
        if apply_path_penalty:
            log_dict["Path Length"] = 0

        backwards = partial(loss_backwards, self.fp16)

        if self.GAN.D_cl is not None:
            self.GAN.D_opt.zero_grad()

            if apply_cl_reg_to_generated:
                for i in range(self.gradient_accumulate_every):
                    get_latents_fn = mixed_list if random() < self.mixed_prob else noise_list
                    style = get_latents_fn(batch_size, num_layers, latent_dim)
                    noise = image_noise(batch_size, image_size)

                    w_space = latent_to_w(self.GAN.S, style)
                    w_styles = styles_def_to_tensor(w_space)

                    generated_images = self.GAN.G(w_styles, noise)
                    self.GAN.D_cl(generated_images.clone().detach(), accumulate=True)

            for i in range(self.gradient_accumulate_every):
                image_batch = next(self.loader).cuda()
                self.GAN.D_cl(image_batch, accumulate=True)

            loss = self.GAN.D_cl.calculate_loss()
            self.last_cr_loss = loss.clone().detach().item()
            log_dict["Consistency"] = self.last_cr_loss
            backwards(loss, self.GAN.D_opt)

            self.GAN.D_opt.step()

        # train discriminator

        avg_pl_length = self.pl_mean
        self.GAN.D_opt.zero_grad()

        for i in range(self.gradient_accumulate_every):
            get_latents_fn = mixed_list if random() < self.mixed_prob else noise_list
            style = get_latents_fn(batch_size, num_layers, latent_dim)
            noise = image_noise(batch_size, image_size)

            w_space = latent_to_w(self.GAN.S, style)
            w_styles = styles_def_to_tensor(w_space)

            generated_images = self.GAN.G(w_styles, noise)
            fake_output, fake_q_loss = self.GAN.D(generated_images.clone().detach())

            image_batch = next(self.loader).cuda()
            image_batch.requires_grad_()
            real_output, real_q_loss = self.GAN.D(image_batch)

            divergence = (F.relu(1 + real_output) + F.relu(1 - fake_output)).mean()
            disc_loss = divergence
            log_dict["Divergence"] += divergence / self.gradient_accumulate_every

            quantize_loss = (fake_q_loss + real_q_loss).mean()
            self.q_loss = float(quantize_loss.detach().item())
            log_dict["Quantize"] += self.q_loss / self.gradient_accumulate_every

            disc_loss = disc_loss + quantize_loss

            if apply_gradient_penalty:
                gp = gradient_penalty(image_batch, real_output)
                self.last_gp_loss = gp.clone().detach().item()
                disc_loss = disc_loss + gp
                log_dict["R1"] += gp / self.gradient_accumulate_every

            disc_loss = disc_loss / self.gradient_accumulate_every
            disc_loss.register_hook(raise_if_nan)
            backwards(disc_loss, self.GAN.D_opt)

            total_disc_loss += divergence.detach().item() / self.gradient_accumulate_every

        self.d_loss = float(total_disc_loss)
        log_dict["Discriminator"] = self.d_loss
        self.GAN.D_opt.step()

        # train generator

        self.GAN.G_opt.zero_grad()
        for i in range(self.gradient_accumulate_every):
            style = get_latents_fn(batch_size, num_layers, latent_dim)
            noise = image_noise(batch_size, image_size)

            w_space = latent_to_w(self.GAN.S, style)
            w_styles = styles_def_to_tensor(w_space)

            generated_images = self.GAN.G(w_styles, noise)
            fake_output, _ = self.GAN.D(generated_images)
            loss = fake_output.mean()
            gen_loss = loss
            log_dict["Generator"] += gen_loss / self.gradient_accumulate_every

            if apply_path_penalty:
                std = 0.1 / (w_styles.std(dim=0, keepdim=True) + EPS)
                w_styles_2 = w_styles + torch.randn(w_styles.shape).cuda() / (std + EPS)
                pl_images = self.GAN.G(w_styles_2, noise)
                pl_lengths = ((pl_images - generated_images) ** 2).mean(dim=(1, 2, 3))
                avg_pl_length = np.mean(pl_lengths.detach().cpu().numpy())

                if not is_empty(self.pl_mean):
                    pl_loss = ((pl_lengths - self.pl_mean) ** 2).mean()
                    log_dict["Path Length"] += pl_loss / self.gradient_accumulate_every
                    if not torch.isnan(pl_loss):
                        gen_loss = gen_loss + pl_loss

            gen_loss = gen_loss / self.gradient_accumulate_every
            gen_loss.register_hook(raise_if_nan)
            backwards(gen_loss, self.GAN.G_opt)

            total_gen_loss += loss.detach().item() / self.gradient_accumulate_every

        self.g_loss = float(total_gen_loss)
        self.GAN.G_opt.step()

        # calculate moving averages
        if apply_path_penalty and not np.isnan(avg_pl_length):
            self.pl_mean = self.pl_length_ma.update_average(self.pl_mean, avg_pl_length)
            log_dict["Mean Path Length"] = self.pl_mean

        if self.steps % 10 == 0 and self.steps > 20000:
            self.GAN.EMA()

        if self.steps <= 25000 and self.steps % 1000 == 2:
            self.GAN.reset_parameter_averaging()

        # save from NaN errors
        checkpoint_num = floor(self.steps / self.save_every)

        if any(torch.isnan(l) for l in (total_gen_loss, total_disc_loss)):
            print(f"NaN detected for generator or discriminator. Loading from checkpoint #{checkpoint_num}")
            self.load(checkpoint_num)
            raise NanException

        # periodically save results
        if self.steps % self.save_every == 0:
            self.save(checkpoint_num)

        if self.steps % 1000 == 0 or (self.steps % 100 == 0 and self.steps < 2500):
            self.evaluate(floor(self.steps / 1000))

        if self.steps % 1000 == 0:
            start_time = time.time()

            PBAR.set_description((f"Calculating FID..."))
            fid, density, coverage = self.calculate_fid()
            log_dict["Evaluation/FID"] = fid
            log_dict["Evaluation/Density"] = density
            log_dict["Evaluation/Coverage"] = coverage

            PBAR.set_description((f"Calculating PPL..."))
            ppl = self.calculate_ppl()

            PBAR.set_description(
                (
                    f"FID: {fid:.4f}; Density: {density:.4f}; Coverage: {coverage:.4f}; PPL: {ppl:.4f} in {time.time() - start_time:.1f}s"
                )
            )
            log_dict["Evaluation/PPL"] = ppl

        wandb.log(log_dict)

        self.steps += 1
        self.av = None

    @torch.no_grad()
    def evaluate(self, num=0, num_image_tiles=8):
        self.GAN.eval()
        ext = "jpg" if not self.transparent else "png"
        num_rows = num_image_tiles

        latent_dim = self.GAN.G.latent_dim
        image_size = self.GAN.G.image_size
        num_layers = self.GAN.G.num_layers

        # latents and noise
        latents = noise_list(num_rows ** 2, num_layers, latent_dim)
        n = image_noise(num_rows ** 2, image_size)

        # regular
        generated_images = self.generate_truncated(self.GAN.S, self.GAN.G, latents, n, trunc_psi=self.trunc_psi)
        grid = torchvision.utils.make_grid(generated_images, nrow=num_rows)
        wandb.log({"Generated Images": [wandb.Image(grid, caption=f"Step {num}")]})

        # moving averages
        generated_images = self.generate_truncated(self.GAN.SE, self.GAN.GE, latents, n, trunc_psi=self.trunc_psi)
        grid = torchvision.utils.make_grid(generated_images, nrow=num_rows)
        wandb.log({"Generated Images EMA": [wandb.Image(grid, caption=f"Step {num}")]})

        # mixing regularities
        def tile(a, dim, n_tile):
            init_dim = a.size(dim)
            repeat_idx = [1] * a.dim()
            repeat_idx[dim] = n_tile
            a = a.repeat(*(repeat_idx))
            order_index = torch.LongTensor(
                np.concatenate([init_dim * np.arange(n_tile) + i for i in range(init_dim)])
            ).cuda()
            return torch.index_select(a, dim, order_index)

        nn = noise(num_rows, latent_dim)
        tmp1 = tile(nn, 0, num_rows)
        tmp2 = nn.repeat(num_rows, 1)

        tt = int(num_layers / 2)
        mixed_latents = [(tmp1, tt), (tmp2, num_layers - tt)]

        generated_images = self.generate_truncated(self.GAN.SE, self.GAN.GE, mixed_latents, n, trunc_psi=self.trunc_psi)
        grid = torchvision.utils.make_grid(generated_images, nrow=num_rows)
        wandb.log({"Style Mixing": [wandb.Image(grid, caption=f"Step {num}")]})

    @torch.no_grad()
    def calculate_fid(self):
        self.GAN.eval()

        inception = InceptionV3([3], normalize_input=False, init_weights=False)
        inception = inception.eval().to(next(self.GAN.parameters()).device)

        latent_dim = self.GAN.G.latent_dim
        image_size = self.GAN.G.image_size
        num_layers = self.GAN.G.num_layers

        features = []
        for _ in range(floor(self.fid_n_sample / self.batch_size) + 1):
            latents = noise_list(self.batch_size, num_layers, latent_dim)
            n = image_noise(self.batch_size, image_size)
            imgs = self.generate_truncated(self.GAN.SE, self.GAN.GE, latents, n, trunc_psi=self.trunc_psi)
            feat = inception(imgs)[0].view(imgs.shape[0], -1)
            features.append(feat.to("cpu"))
        features = torch.cat(features, 0).numpy()

        del inception

        sample_mean = np.mean(features, 0)
        sample_cov = np.cov(features, rowvar=False)

        with open(f"inception_{self.name}_stats.pkl", "rb") as f:
            embeds = pickle.load(f)
            real_mean = embeds["mean"]
            real_cov = embeds["cov"]

        cov_sqrt, _ = linalg.sqrtm(sample_cov @ real_cov, disp=False)
        if not np.isfinite(cov_sqrt).all():
            print("product of cov matrices is singular")
            offset = np.eye(sample_cov.shape[0]) * 1e-6
            cov_sqrt = linalg.sqrtm((sample_cov + offset) @ (real_cov + offset))
        if np.iscomplexobj(cov_sqrt):
            if not np.allclose(np.diagonal(cov_sqrt).imag, 0, atol=1e-3):
                raise ValueError(f"Imaginary component {np.max(np.abs(cov_sqrt.imag))}")
            cov_sqrt = cov_sqrt.real

        mean_diff = sample_mean - real_mean
        mean_norm = mean_diff @ mean_diff
        trace = np.trace(sample_cov) + np.trace(real_cov) - 2 * np.trace(cov_sqrt)
        inception_distance = mean_norm + trace

        with open(f"inception_{self.name}_features.pkl", "rb") as f:
            embeds = pickle.load(f)
            real_feats = embeds["features"]
        _, _, density, coverage = validation.prdc(real_feats[:80000], features[:80000])

        return inception_distance, density, coverage

    @torch.no_grad()
    def calculate_ppl(self):
        self.GAN.eval()
        latent_dim = self.GAN.G.latent_dim
        image_size = self.GAN.G.image_size
        num_layers = self.GAN.G.num_layers
        eps = 1e-4

        def lerp(a, b, t):
            return a + (b - a) * t

        percept = lpips.PerceptualLoss(
            model="net-lin", net="vgg", use_gpu=True, gpu_ids=[next(self.GAN.parameters()).device.index]
        )

        distances = []
        for _ in range(floor(self.fid_n_sample / self.batch_size) + 1):
            noise = image_noise(self.batch_size * 2, image_size)
            inputs = noise_list(self.batch_size * 2, num_layers, latent_dim)
            lerp_t = torch.rand(self.batch_size).cuda()

            # print(lerp_t.shape)

            w_space = []
            for tensor, num_layers in inputs:
                av = torch.from_numpy(self.av).cuda()
                tmp = self.trunc_psi * (self.GAN.SE(tensor) - av) + av
                w_space.append((tmp, num_layers))
            latent = styles_def_to_tensor(w_space)

            latent_t0, latent_t1 = latent[::2], latent[1::2]
            latent_e0 = lerp(latent_t0, latent_t1, lerp_t[:, None, None])
            latent_e1 = lerp(latent_t0, latent_t1, lerp_t[:, None, None] + eps)
            latent_e = torch.stack([latent_e0, latent_e1], 1).view(*latent.shape)

            image = evaluate_in_chunks(self.batch_size, self.GAN.GE, latent_e, noise)

            factor = image.shape[2] // 256
            if factor > 1:
                image = F.interpolate(image, size=(256, 256), mode="bilinear", align_corners=False)

            dist = percept(image[::2], image[1::2]).view(image.shape[0] // 2) / (eps ** 2)
            distances.append(dist.to("cpu").numpy())
        distances = np.concatenate(distances, 0)

        lo = np.percentile(distances, 1, interpolation="lower")
        hi = np.percentile(distances, 99, interpolation="higher")
        filtered_dist = np.extract(np.logical_and(lo <= distances, distances <= hi), distances)
        path_length = filtered_dist.mean()

        del percept, inputs, lerp_t, image, dist

        return path_length

    @torch.no_grad()
    def generate_truncated(self, S, G, style, noi, trunc_psi=0.75, num_image_tiles=8):
        latent_dim = G.latent_dim

        if self.av is None:
            z = noise(2 ** 14, latent_dim)
            samples = evaluate_in_chunks(self.batch_size, S, z).cpu().numpy()
            self.av = np.mean(samples, axis=0)
            self.av = np.expand_dims(self.av, axis=0)

        w_space = []
        for tensor, num_layers in style:
            tmp = S(tensor)
            av_torch = torch.from_numpy(self.av).cuda()
            tmp = trunc_psi * (tmp - av_torch) + av_torch
            w_space.append((tmp, num_layers))

        w_styles = styles_def_to_tensor(w_space)
        generated_images = evaluate_in_chunks(self.batch_size, G, w_styles, noi)
        return generated_images.clamp_(0.0, 1.0)

    def print_log(self):
        print(
            f"G: {self.g_loss:.2f} | D: {self.d_loss:.2f} | GP: {self.last_gp_loss:.2f} | PL: {self.pl_mean:.2f} | CR: {self.last_cr_loss:.2f} | Q: {self.q_loss:.2f}"
        )

    def model_name(self, num):
        return str(self.models_dir / self.name / f"model_{wandb.run.dir.split('/')[-1].split('-')[-1]}_{num}.pt")

    def init_folders(self):
        (self.results_dir / self.name).mkdir(parents=True, exist_ok=True)
        (self.models_dir / self.name).mkdir(parents=True, exist_ok=True)

    def clear(self):
        rmtree(f"./models/{self.name}", True)
        rmtree(f"./results/{self.name}", True)
        rmtree(str(self.config_path), True)
        self.init_folders()

    def save(self, num):
        torch.save(self.GAN.state_dict(), self.model_name(num))
        self.write_config()

    def load(self, num=-1):
        self.load_config()

        name = num
        if num == -1:
            file_paths = [p for p in Path(self.models_dir / self.name).glob("model_*.pt")]
            saved_nums = sorted(map(lambda x: int(x.stem.split("_")[-1]), file_paths))
            if len(saved_nums) == 0:
                return
            name = saved_nums[-1]
            print(f"continuing from previous epoch - {name}")
        self.steps = name * self.save_every
        self.GAN.load_state_dict(torch.load(self.model_name(name)))


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("data", type=str)
    parser.add_argument("name", type=str)
    parser.add_argument("--results_dir", type=str, default="/home/hans/neurout/")
    parser.add_argument("--models_dir", type=str, default="/home/hans/modelzoo/maua-sg2/")
    parser.add_argument("--new", type=bool, default=False)
    parser.add_argument("--load_from", type=str, default=-1)
    parser.add_argument("--image_size", type=int, default=256)
    parser.add_argument("--network_capacity", type=int, default=16)
    parser.add_argument("--transparent", type=bool, default=False)
    parser.add_argument("--batch_size", type=int, default=2)
    parser.add_argument("--gradient_accumulate_every", type=int, default=12)
    parser.add_argument("--num_train_steps", type=int, default=150000)
    parser.add_argument("--learning_rate", type=float, default=2e-4)
    parser.add_argument("--num_workers", type=int, default=None)
    parser.add_argument("--save_every", type=int, default=1000)
    parser.add_argument("--generate", type=bool, default=False)
    parser.add_argument("--num_image_tiles", type=int, default=8)
    parser.add_argument("--trunc_psi", type=float, default=1)
    parser.add_argument("--fp16", type=bool, default=False)
    parser.add_argument("--cl_reg", type=bool, default=True)
    parser.add_argument("--fq_layers", default=[])
    parser.add_argument("--fq_dict_size", default=256)
    parser.add_argument("--attn_layers", default=[])
    args = parser.parse_args()

    wandb.init(project=f"maua-stylegan", name="lucidrains-" + args.name)

    model = Trainer(
        args.name,
        args.results_dir,
        args.models_dir,
        batch_size=args.batch_size,
        gradient_accumulate_every=args.gradient_accumulate_every,
        image_size=args.image_size,
        network_capacity=args.network_capacity,
        transparent=args.transparent,
        lr=args.learning_rate,
        num_workers=args.num_workers,
        save_every=args.save_every,
        trunc_psi=args.trunc_psi,
        fp16=args.fp16,
        cl_reg=args.cl_reg,
        fq_layers=args.fq_layers,
        fq_dict_size=args.fq_dict_size,
        attn_layers=args.attn_layers,
    )

    if not args.new:
        model.load(args.load_from)
    else:
        model.clear()

    if args.generate:
        now = datetime.now()
        timestamp = now.strftime("%m-%d-%Y_%H-%M-%S")
        samples_name = f"generated-{timestamp}"
        model.evaluate(samples_name, args.num_image_tiles)
        print(f"sample images generated at {args.results_dir}/{args.name}/{args.samples_name}")
        exit()

    model.set_data_src(args.data)

    PBAR = tqdm(range(args.num_train_steps - model.steps), mininterval=10.0, desc=f"{args.name}<{args.data}>")
    for _ in PBAR:
        retry_call(model.train, tries=3, exceptions=NanException)
        if _ % 50 == 0:
            model.print_log()