Spaces:
Runtime error
Runtime error
| 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() | |