import torch import torch.nn as nn import numpy as np import torch.nn.functional as F import gradio as gr from einops import rearrange, repeat import torchvision.transforms.functional as ttf from timm.models.convmixer import ConvMixer import functorch def img_to_patches(im, patch_h, patch_w): "B, C, H, W -> B, C, D, h_patch, w_patch" bs, c, h, w = im.shape im = im.unfold(-1, patch_h, patch_w).unfold(2, patch_h, patch_w) im = im.permute(0, 1, 2, 3, 5, 4) im = im.contiguous().view(bs, c, -1, patch_h, patch_w) return im def patches_to_img(patches, num_patch_h, num_patch_w): "B, C, D, h_patch, w_patch -> B, C, H, W" bs, c, d, h, w = patches.shape patches = patches.view(bs, c, num_patch_h, num_patch_w, h, w) # fold patches patches = torch.cat([patches[..., k, :, :] for k in range(num_patch_w)], dim=-1) x = torch.cat([patches[..., k, :, :] for k in range(num_patch_h)], dim=-2) return x def vmapped_rotate(x, angle, in_dims=1): "B, C, D, H, W -> B, C, D, H, W" rotate_ = functorch.vmap(ttf.rotate, in_dims=in_dims, out_dims=in_dims) return rotate_(x, angle=angle) class CollageOperator2d(nn.Module): def __init__(self, res, rh, rw, dh=None, dw=None, use_augmentations=False): """Collage Operator for two-dimensional data. Given a fractal code, it outputs the corresponding fixed-point. Args: res (int): Spatial resolutions of input (and output) data. rh (int): Height of range (target) square patches. rw (int): Width of range (target) square patches. dh (int, optional): Height of range domain (source) patches. Defaults to `res`. dw (int, optional): Width of range domain (source) patches. Defaults to `res`. use_augmentations (bool, optional): Use augmentations of domain square patches at each decoding iteration. Defaults to `False`. """ super().__init__() self.dh, self.dw = dh, dw if self.dh is None: self.dh = res if self.dw is None: self.dw = res # 5 refers to the 5 copies of domain patches generated with the current choice of augmentations: # 3 rotations (90, 180, 270), horizontal flips and vertical flips. self.n_aug_transforms = 9 if use_augmentations else 0 # precompute useful quantities related to the partitioning scheme into patches, given # the desired `dh`, `dw`, `rh`, `rw`. partition_info = self.collage_partition_info(res, self.dh, self.dw, rh, rw) self.n_dh, self.n_dw, self.n_rh, self.n_rw, self.h_factors, self.w_factors, self.n_domains, self.n_ranges = partition_info # At each step of the collage, all (source) domain patches are pooled down to the size of range (target) patches. # Notices how the pooling factors do not change if one decodes at higher resolutions, since both domain and range # patch sizes are multiplied by the same integer. self.pool = nn.AvgPool3d(kernel_size=(1, self.h_factors, self.w_factors), stride=(1, self.h_factors, self.w_factors)) def collage_operator(self, z, collage_weight, collage_bias): """Collage Operator (decoding). Performs the steps described in Def. 3.1, Figure 2.""" # Given the current iterate `z`, we split it into domain patches according to the partitioning scheme. domains = img_to_patches(z) # Pool domains (pre augmentation) to range patch sizes. pooled_domains = self.pool(domains) # If needed, produce additional candidate domain patches as augmentations of existing domains. # Auxiliary learned feature maps / patches are also introduced here. if self.n_aug_transforms > 1: pooled_domains = self.generate_candidates(pooled_domains) pooled_domains = repeat(pooled_domains, 'b c d h w -> b c d r h w', r=self.num_ranges) # Apply the affine maps to domain patches range_domains = torch.einsum('bcdrhw, bcdr -> bcrhw', pooled_domains, collage_weight) range_domains = range_domains + collage_bias[..., None, None] # Reconstruct data by "composing" the output patches back together (collage!). z = patches_to_img(range_domains) return z def decode_step(self, z, weight, bias, superres_factor, return_patches=False): """Single Collage Operator step. Performs the steps described in: https://arxiv.org/pdf/2204.07673.pdf (Def. 3.1, Figure 2). """ # Given the current iterate `z`, we split it into `n_domains` domain patches. domains = img_to_patches(z, patch_h=self.dh * superres_factor, patch_w=self.dw * superres_factor) # Pool domains (pre augmentation) for compatibility with range patches. pooled_domains = self.pool(domains) # If needed, produce additional candidate domain patches as augmentations of existing domains. if self.n_aug_transforms > 1: pooled_domains = self.generate_candidates(pooled_domains) pooled_domains = repeat(pooled_domains, 'b c d h w -> b c d r h w', r=self.n_ranges) # Apply the affine maps to domain patches range_domains = torch.einsum('bcdrhw, bcdr -> bcrhw', pooled_domains, weight) range_domains = range_domains + bias[:, :, :, None, None] # Reconstruct data by "composing" the output patches back together (collage!). z = patches_to_img(range_domains, self.n_rh, self.n_rw) if return_patches: return z, (domains, pooled_domains, range_domains) return z def generate_candidates(self, domains): domains = domains.permute(0,2,1,3,4) rotations = [vmapped_rotate(domains, angle=angle) for angle in (90, 180, 270)] hflips = ttf.hflip(domains) vflips = ttf.vflip(domains) br_shift = ttf.adjust_brightness(domains, 0.5) cr_shift = ttf.adjust_contrast(domains, 0.5) hue_shift = ttf.adjust_hue(domains, 0.5) sat_shift = ttf.adjust_saturation(domains, 0.5) domains = torch.cat([domains, *rotations, hflips, vflips, br_shift, cr_shift, hue_shift, sat_shift], dim=1) return domains.permute(0,2,1,3,4) def forward(self, x, co_w, co_bias, decode_steps=20, superres_factor=1): B, C, H, W = x.shape # It does not matter which initial condition is chosen, so long as the dimensions match. # The fixed-point of a Collage Operator is uniquely determined* by the fractal code # *: and auxiliary learned patches, if any. z = torch.randn(B, C, H * superres_factor, W * superres_factor).to(x.device) for _ in range(decode_steps): z = self.decode_step(z, co_w, co_bias, superres_factor) return z def collage_partition_info(self, input_res, dh, dw, rh, rw): """ Computes auxiliary information for the collage (number of source and target domains, and relative size factors) """ height = width = input_res n_dh, n_dw = height // dh, width // dw n_domains = n_dh * n_dw # Adjust number of domain patches to include augmentations n_domains = n_domains + n_domains * self.n_aug_transforms # (3 rotations, hflip, vlip) h_factors, w_factors = dh // rh, dw // rw n_rh, n_rw = input_res // rh, input_res // rw n_ranges = n_rh * n_rw return n_dh, n_dw, n_rh, n_rw, h_factors, w_factors, n_domains, n_ranges class NeuralCollageOperator2d(nn.Module): def __init__(self, out_res, out_channels, rh, rw, dh=None, dw=None, net=None, use_augmentations=False): super().__init__() self.co = CollageOperator2d(out_res, rh, rw, dh, dw, use_augmentations) # In a Collage Operator, the affine map requires a single scalar weight # for each pair of domain and range patches, and a single scalar bias for each range. # `net` learns to output these weights based on the objective. self.co_w_dim = self.co.n_domains * self.co.n_ranges * out_channels self.co_bias_dim = self.co.n_ranges * out_channels tot_out_dim = self.co_w_dim + self.co_bias_dim # Does not need to be a ConvMixer: for deep generative Neural Collages `net` can be e.g, a VDVAE. if net is None: net = ConvMixer(dim=32, depth=8, kernel_size=9, patch_size=7, num_classes=tot_out_dim) self.net = net self.softmax = nn.Softmax(dim=-1) self.tanh = nn.Tanh() def forward(self, x, decode_steps=10, superres_factor=1, return_co_code=False): B, C, H, W = x.shape co_code = self.net(x) # B, C, co_w_dim + co_mix_dim + co_bias_dim co_w, co_bias = torch.split(co_code, [self.co_w_dim, self.co_bias_dim], dim=-1) co_w = co_w.view(B, C, self.co.n_domains, self.co.n_ranges) # No restrictions on co_w, thus no guarantee of contractiveness. # In the full jax version of Neural Collages we enforce the constraint |co_w| < 1 (elementwise). co_bias = co_bias.view(B, C, self.co.n_ranges) co_bias = self.tanh(co_bias) z = self.co(x, co_w, co_bias, decode_steps=decode_steps, superres_factor=superres_factor) if return_co_code: return z, co_w, co_bias else: return z def fractalize(img, superresolution_factor=1): superresolution_factor = int(superresolution_factor) device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') im = np.asarray(img) im = torch.from_numpy(im).permute(2,0,1).to(device) co = NeuralCollageOperator2d(out_res=100, out_channels=3, rh=2, rw=2, dh=100, dw=100).to(device) opt = torch.optim.Adam(co.parameters(), lr=1e-2) objective = nn.MSELoss() norm_im = im.float().unsqueeze(0) / 255 for _ in range(200): recon = co(norm_im, decode_steps=10, return_co_code=False) loss = objective(recon, norm_im) loss.backward() opt.step() opt.zero_grad() fractal_img = co(norm_im, decode_steps=10, superres_factor=superresolution_factor)[0].permute(1,2,0).clamp(-1, 1) return fractal_img.cpu().detach().numpy() demo = gr.Interface( fn=fractalize, inputs=[gr.Image(shape=(100, 100), image_mode='RGB'), gr.Slider(1, 40, step=1)], outputs="image" ) demo.launch()