first

app.py

@@ -0,0 +1,225 @@
+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()

requirements.txt

@@ -0,0 +1,7 @@
+numpy
+torch
+einops 
+timm 
+torchvision 
+functorch  
+torch