| import math
|
| import torch
|
| import torch.nn.functional as F
|
| from math import log2
|
| from torch import nn, einsum
|
| from kornia.filters import filter2d
|
| from einops import reduce, rearrange, repeat
|
|
|
|
|
| def exists(val):
|
| return val is not None
|
|
|
|
|
| def is_power_of_two(val):
|
| return log2(val).is_integer()
|
|
|
|
|
| def default(val, d):
|
| return val if exists(val) else d
|
|
|
|
|
| def get_1d_dct(i, freq, L):
|
| result = math.cos(math.pi * freq * (i + 0.5) / L) / math.sqrt(L)
|
| return result * (1 if freq == 0 else math.sqrt(2))
|
|
|
|
|
| def get_dct_weights(width, channel, fidx_u, fidx_v):
|
| dct_weights = torch.zeros(1, channel, width, width)
|
| c_part = channel // len(fidx_u)
|
|
|
| for i, (u_x, v_y) in enumerate(zip(fidx_u, fidx_v)):
|
| for x in range(width):
|
| for y in range(width):
|
| coor_value = get_1d_dct(x, u_x, width) * get_1d_dct(y, v_y, width)
|
| dct_weights[:, i * c_part : (i + 1) * c_part, x, y] = coor_value
|
|
|
| return dct_weights
|
|
|
|
|
| class Blur(nn.Module):
|
| def __init__(self):
|
| super().__init__()
|
| f = torch.Tensor([1, 2, 1])
|
| self.register_buffer("f", f)
|
|
|
| def forward(self, x):
|
| f = self.f
|
| f = f[None, None, :] * f[None, :, None]
|
| return filter2d(x, f, normalized=True)
|
|
|
|
|
| class ChanNorm(nn.Module):
|
| def __init__(self, dim, eps=1e-5):
|
| super().__init__()
|
| self.eps = eps
|
| self.g = nn.Parameter(torch.ones(1, dim, 1, 1))
|
| self.b = nn.Parameter(torch.zeros(1, dim, 1, 1))
|
|
|
| def forward(self, x):
|
| var = torch.var(x, dim=1, unbiased=False, keepdim=True)
|
| mean = torch.mean(x, dim=1, keepdim=True)
|
| return (x - mean) / (var + self.eps).sqrt() * self.g + self.b
|
|
|
|
|
| def Conv2dSame(dim_in, dim_out, kernel_size, bias=True):
|
| pad_left = kernel_size // 2
|
| pad_right = (pad_left - 1) if (kernel_size % 2) == 0 else pad_left
|
|
|
| return nn.Sequential(
|
| nn.ZeroPad2d((pad_left, pad_right, pad_left, pad_right)),
|
| nn.Conv2d(dim_in, dim_out, kernel_size, bias=bias),
|
| )
|
|
|
|
|
| class DepthWiseConv2d(nn.Module):
|
| def __init__(self, dim_in, dim_out, kernel_size, padding=0, stride=1, bias=True):
|
| super().__init__()
|
| self.net = nn.Sequential(
|
| nn.Conv2d(
|
| dim_in,
|
| dim_in,
|
| kernel_size=kernel_size,
|
| padding=padding,
|
| groups=dim_in,
|
| stride=stride,
|
| bias=bias,
|
| ),
|
| nn.Conv2d(dim_in, dim_out, kernel_size=1, bias=bias),
|
| )
|
|
|
| def forward(self, x):
|
| return self.net(x)
|
|
|
|
|
| class FCANet(nn.Module):
|
| def __init__(self, *, chan_in, chan_out, reduction=4, width):
|
| super().__init__()
|
|
|
| freq_w, freq_h = ([0] * 8), list(
|
| range(8)
|
| )
|
| dct_weights = get_dct_weights(
|
| width, chan_in, [*freq_w, *freq_h], [*freq_h, *freq_w]
|
| )
|
| self.register_buffer("dct_weights", dct_weights)
|
|
|
| chan_intermediate = max(3, chan_out // reduction)
|
|
|
| self.net = nn.Sequential(
|
| nn.Conv2d(chan_in, chan_intermediate, 1),
|
| nn.LeakyReLU(0.1),
|
| nn.Conv2d(chan_intermediate, chan_out, 1),
|
| nn.Sigmoid(),
|
| )
|
|
|
| def forward(self, x):
|
| x = reduce(
|
| x * self.dct_weights, "b c (h h1) (w w1) -> b c h1 w1", "sum", h1=1, w1=1
|
| )
|
| return self.net(x)
|
|
|
|
|
| class Generator(nn.Module):
|
| def __init__(
|
| self,
|
| *,
|
| image_size,
|
| latent_dim=256,
|
| fmap_max=512,
|
| fmap_inverse_coef=12,
|
| transparent=False,
|
| greyscale=False,
|
| attn_res_layers=[],
|
| freq_chan_attn=False,
|
| syncbatchnorm=False,
|
| antialias=False,
|
| ):
|
| super().__init__()
|
| resolution = log2(image_size)
|
| assert is_power_of_two(image_size), "image size must be a power of 2"
|
|
|
|
|
| norm_class = nn.SyncBatchNorm if syncbatchnorm else nn.BatchNorm2d
|
| Blur = nn.Identity if not antialias else Blur
|
|
|
| if transparent:
|
| init_channel = 4
|
| elif greyscale:
|
| init_channel = 1
|
| else:
|
| init_channel = 3
|
|
|
| self.latent_dim = latent_dim
|
|
|
| fmap_max = default(fmap_max, latent_dim)
|
|
|
| self.initial_conv = nn.Sequential(
|
| nn.ConvTranspose2d(latent_dim, latent_dim * 2, 4),
|
| norm_class(latent_dim * 2),
|
| nn.GLU(dim=1),
|
| )
|
|
|
| num_layers = int(resolution) - 2
|
| features = list(
|
| map(lambda n: (n, 2 ** (fmap_inverse_coef - n)), range(2, num_layers + 2))
|
| )
|
| features = list(map(lambda n: (n[0], min(n[1], fmap_max)), features))
|
| features = list(map(lambda n: 3 if n[0] >= 8 else n[1], features))
|
| features = [latent_dim, *features]
|
|
|
| in_out_features = list(zip(features[:-1], features[1:]))
|
|
|
| self.res_layers = range(2, num_layers + 2)
|
| self.layers = nn.ModuleList([])
|
| self.res_to_feature_map = dict(zip(self.res_layers, in_out_features))
|
|
|
| self.sle_map = ((3, 7), (4, 8), (5, 9), (6, 10))
|
| self.sle_map = list(
|
| filter(lambda t: t[0] <= resolution and t[1] <= resolution, self.sle_map)
|
| )
|
| self.sle_map = dict(self.sle_map)
|
|
|
| self.num_layers_spatial_res = 1
|
|
|
| for res, (chan_in, chan_out) in zip(self.res_layers, in_out_features):
|
| image_width = 2**res
|
|
|
| attn = None
|
| if image_width in attn_res_layers:
|
| attn = PreNorm(chan_in, LinearAttention(chan_in))
|
|
|
| sle = None
|
| if res in self.sle_map:
|
| residual_layer = self.sle_map[res]
|
| sle_chan_out = self.res_to_feature_map[residual_layer - 1][-1]
|
|
|
| if freq_chan_attn:
|
| sle = FCANet(
|
| chan_in=chan_out, chan_out=sle_chan_out, width=2 ** (res + 1)
|
| )
|
| else:
|
| sle = GlobalContext(chan_in=chan_out, chan_out=sle_chan_out)
|
|
|
| layer = nn.ModuleList(
|
| [
|
| nn.Sequential(
|
| PixelShuffleUpsample(chan_in),
|
| Blur(),
|
| Conv2dSame(chan_in, chan_out * 2, 4),
|
| Noise(),
|
| norm_class(chan_out * 2),
|
| nn.GLU(dim=1),
|
| ),
|
| sle,
|
| attn,
|
| ]
|
| )
|
| self.layers.append(layer)
|
|
|
| self.out_conv = nn.Conv2d(features[-1], init_channel, 3, padding=1)
|
|
|
| def forward(self, x):
|
| x = rearrange(x, "b c -> b c () ()")
|
| x = self.initial_conv(x)
|
| x = F.normalize(x, dim=1)
|
|
|
| residuals = dict()
|
|
|
| for res, (up, sle, attn) in zip(self.res_layers, self.layers):
|
| if exists(attn):
|
| x = attn(x) + x
|
|
|
| x = up(x)
|
|
|
| if exists(sle):
|
| out_res = self.sle_map[res]
|
| residual = sle(x)
|
| residuals[out_res] = residual
|
|
|
| next_res = res + 1
|
| if next_res in residuals:
|
| x = x * residuals[next_res]
|
|
|
| return self.out_conv(x)
|
|
|
|
|
| class GlobalContext(nn.Module):
|
| def __init__(self, *, chan_in, chan_out):
|
| super().__init__()
|
| self.to_k = nn.Conv2d(chan_in, 1, 1)
|
| chan_intermediate = max(3, chan_out // 2)
|
|
|
| self.net = nn.Sequential(
|
| nn.Conv2d(chan_in, chan_intermediate, 1),
|
| nn.LeakyReLU(0.1),
|
| nn.Conv2d(chan_intermediate, chan_out, 1),
|
| nn.Sigmoid(),
|
| )
|
|
|
| def forward(self, x):
|
| context = self.to_k(x)
|
| context = context.flatten(2).softmax(dim=-1)
|
| out = einsum("b i n, b c n -> b c i", context, x.flatten(2))
|
| out = out.unsqueeze(-1)
|
| return self.net(out)
|
|
|
|
|
| class LinearAttention(nn.Module):
|
| def __init__(self, dim, dim_head=64, heads=8, kernel_size=3):
|
| super().__init__()
|
| self.scale = dim_head**-0.5
|
| self.heads = heads
|
| self.dim_head = dim_head
|
| inner_dim = dim_head * heads
|
|
|
| self.kernel_size = kernel_size
|
| self.nonlin = nn.GELU()
|
|
|
| self.to_lin_q = nn.Conv2d(dim, inner_dim, 1, bias=False)
|
| self.to_lin_kv = DepthWiseConv2d(dim, inner_dim * 2, 3, padding=1, bias=False)
|
|
|
| self.to_q = nn.Conv2d(dim, inner_dim, 1, bias=False)
|
| self.to_kv = nn.Conv2d(dim, inner_dim * 2, 1, bias=False)
|
|
|
| self.to_out = nn.Conv2d(inner_dim * 2, dim, 1)
|
|
|
| def forward(self, fmap):
|
| h, x, y = self.heads, *fmap.shape[-2:]
|
|
|
|
|
|
|
| lin_q, lin_k, lin_v = (
|
| self.to_lin_q(fmap),
|
| *self.to_lin_kv(fmap).chunk(2, dim=1),
|
| )
|
| lin_q, lin_k, lin_v = map(
|
| lambda t: rearrange(t, "b (h c) x y -> (b h) (x y) c", h=h),
|
| (lin_q, lin_k, lin_v),
|
| )
|
|
|
| lin_q = lin_q.softmax(dim=-1)
|
| lin_k = lin_k.softmax(dim=-2)
|
|
|
| lin_q = lin_q * self.scale
|
|
|
| context = einsum("b n d, b n e -> b d e", lin_k, lin_v)
|
| lin_out = einsum("b n d, b d e -> b n e", lin_q, context)
|
| lin_out = rearrange(lin_out, "(b h) (x y) d -> b (h d) x y", h=h, x=x, y=y)
|
|
|
|
|
|
|
| q, k, v = (self.to_q(fmap), *self.to_kv(fmap).chunk(2, dim=1))
|
| q, k, v = map(
|
| lambda t: rearrange(t, "b (h c) x y -> (b h) c x y", h=h), (q, k, v)
|
| )
|
|
|
| k = F.unfold(k, kernel_size=self.kernel_size, padding=self.kernel_size // 2)
|
| v = F.unfold(v, kernel_size=self.kernel_size, padding=self.kernel_size // 2)
|
|
|
| k, v = map(
|
| lambda t: rearrange(t, "b (d j) n -> b n j d", d=self.dim_head), (k, v)
|
| )
|
|
|
| q = rearrange(q, "b c ... -> b (...) c") * self.scale
|
|
|
| sim = einsum("b i d, b i j d -> b i j", q, k)
|
| sim = sim - sim.amax(dim=-1, keepdim=True).detach()
|
|
|
| attn = sim.softmax(dim=-1)
|
|
|
| full_out = einsum("b i j, b i j d -> b i d", attn, v)
|
| full_out = rearrange(full_out, "(b h) (x y) d -> b (h d) x y", h=h, x=x, y=y)
|
|
|
|
|
|
|
| lin_out = self.nonlin(lin_out)
|
| out = torch.cat((lin_out, full_out), dim=1)
|
| return self.to_out(out)
|
|
|
|
|
| class Noise(nn.Module):
|
| def __init__(self):
|
| super().__init__()
|
| self.weight = nn.Parameter(torch.zeros(1))
|
|
|
| def forward(self, x, noise=None):
|
| b, _, h, w, device = *x.shape, x.device
|
|
|
| if not exists(noise):
|
| noise = torch.randn(b, 1, h, w, device=device)
|
|
|
| return x + self.weight * noise
|
|
|
|
|
| class PixelShuffleUpsample(nn.Module):
|
| def __init__(self, dim, dim_out=None):
|
| super().__init__()
|
| dim_out = default(dim_out, dim)
|
| conv = nn.Conv2d(dim, dim_out * 4, 1)
|
|
|
| self.net = nn.Sequential(conv, nn.SiLU(), nn.PixelShuffle(2))
|
|
|
| self.init_conv_(conv)
|
|
|
| def init_conv_(self, conv):
|
| o, i, h, w = conv.weight.shape
|
| conv_weight = torch.empty(o // 4, i, h, w)
|
| nn.init.kaiming_uniform_(conv_weight)
|
| conv_weight = repeat(conv_weight, "o ... -> (o 4) ...")
|
|
|
| conv.weight.data.copy_(conv_weight)
|
| nn.init.zeros_(conv.bias.data)
|
|
|
| def forward(self, x):
|
| return self.net(x)
|
|
|
|
|
| class PreNorm(nn.Module):
|
| def __init__(self, dim, fn):
|
| super().__init__()
|
| self.fn = fn
|
| self.norm = ChanNorm(dim)
|
|
|
| def forward(self, x):
|
| return self.fn(self.norm(x))
|
|
|
