Update app.py

app.py

@@ -19,8 +19,194 @@ from fastprogress.fastprogress import master_bar, progress_bar
 from IPython.display import HTML
 from base64 import b64encode
 
+# Definitions
+device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+
+def sinc(x):
+    return torch.where(x != 0, torch.sin(math.pi * x) / (math.pi * x), x.new_ones([]))
+ 
+ 
+def lanczos(x, a):
+    cond = torch.logical_and(-a < x, x < a)
+    out = torch.where(cond, sinc(x) * sinc(x/a), x.new_zeros([]))
+    return out / out.sum()
+ 
+ 
+def ramp(ratio, width):
+    n = math.ceil(width / ratio + 1)
+    out = torch.empty([n])
+    cur = 0
+    for i in range(out.shape[0]):
+        out[i] = cur
+        cur += ratio
+    return torch.cat([-out[1:].flip([0]), out])[1:-1]
+
+class Prompt(nn.Module):
+    def __init__(self, embed, weight=1., stop=float('-inf')):
+        super().__init__()
+        self.register_buffer('embed', embed)
+        self.register_buffer('weight', torch.as_tensor(weight))
+        self.register_buffer('stop', torch.as_tensor(stop))
+ 
+    def forward(self, input):
+        input_normed = F.normalize(input.unsqueeze(1), dim=2)
+        embed_normed = F.normalize(self.embed.unsqueeze(0), dim=2)
+        dists = input_normed.sub(embed_normed).norm(dim=2).div(2).arcsin().pow(2).mul(2)
+        dists = dists * self.weight.sign()
+        return self.weight.abs() * replace_grad(dists, torch.maximum(dists, self.stop)).mean()
+
+class MakeCutouts(nn.Module):
+    def __init__(self, cut_size, cutn, cut_pow=1.):
+        super().__init__()
+        self.cut_size = cut_size
+        self.cutn = cutn
+        self.cut_pow = cut_pow
+        self.augs = nn.Sequential(
+            K.RandomHorizontalFlip(p=0.5),
+            K.RandomSharpness(0.3,p=0.4),
+            K.RandomAffine(degrees=30, translate=0.1, p=0.8, padding_mode='border'),
+            K.RandomPerspective(0.2,p=0.4),
+            K.ColorJitter(hue=0.01, saturation=0.01, p=0.7))
+        self.noise_fac = 0.1 
+
+    def forward(self, input):
+        sideY, sideX = input.shape[2:4]
+        max_size = min(sideX, sideY)
+        min_size = min(sideX, sideY, self.cut_size)
+        cutouts = []
+        for _ in range(self.cutn):
+            size = int(torch.rand([])**self.cut_pow * (max_size - min_size) + min_size)
+            offsetx = torch.randint(0, sideX - size + 1, ())
+            offsety = torch.randint(0, sideY - size + 1, ())
+            cutout = input[:, :, offsety:offsety + size, offsetx:offsetx + size]
+            cutouts.append(resample(cutout, (self.cut_size, self.cut_size)))
+        batch = self.augs(torch.cat(cutouts, dim=0))
+        if self.noise_fac:
+            facs = batch.new_empty([self.cutn, 1, 1, 1]).uniform_(0, self.noise_fac)
+            batch = batch + facs * torch.randn_like(batch)
+        return batch
+
+def resample(input, size, align_corners=True):
+    n, c, h, w = input.shape
+    dh, dw = size
+ 
+    input = input.view([n * c, 1, h, w])
+ 
+    if dh < h:
+        kernel_h = lanczos(ramp(dh / h, 2), 2).to(input.device, input.dtype)
+        pad_h = (kernel_h.shape[0] - 1) // 2
+        input = F.pad(input, (0, 0, pad_h, pad_h), 'reflect')
+        input = F.conv2d(input, kernel_h[None, None, :, None])
+ 
+    if dw < w:
+        kernel_w = lanczos(ramp(dw / w, 2), 2).to(input.device, input.dtype)
+        pad_w = (kernel_w.shape[0] - 1) // 2
+        input = F.pad(input, (pad_w, pad_w, 0, 0), 'reflect')
+        input = F.conv2d(input, kernel_w[None, None, None, :])
+ 
+    input = input.view([n, c, h, w])
+    return F.interpolate(input, size, mode='bicubic', align_corners=align_corners)
+
+class ReplaceGrad(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x_forward, x_backward):
+        ctx.shape = x_backward.shape
+        return x_forward
+ 
+    @staticmethod
+    def backward(ctx, grad_in):
+        return None, grad_in.sum_to_size(ctx.shape)
+ 
+ 
+replace_grad = ReplaceGrad.apply 
+
+# Set up CLIP
+perceptor = clip.load('ViT-B/32', jit=False)[0].eval().requires_grad_(False).to(device)
+normalize = transforms.Normalize(mean=[0.48145466, 0.4578275, 0.40821073],
+                                 std=[0.26862954, 0.26130258, 0.27577711])
+cut_size = perceptor.visual.input_resolution
+cutn=64
+cut_pow=1
+make_cutouts = MakeCutouts(cut_size, cutn, cut_pow=cut_pow)
+
+# ImStack
+class ImStack(nn.Module):
+  """ This class represents an image as a series of stacked arrays, where each is 1/2
+  the resolution of the next. This is useful eg when trying to create an image to minimise
+  some loss - parameters in the early (small) layers can have an affect on the overall 
+  structure and shapes while those in later layers act as residuals and fill in fine detail.
+  """
+
+  def __init__(self, n_layers=3, base_size=32, scale=2,
+               init_image=None, out_size=256, decay=0.7):
+    """Constructs the Image Stack
+
+    Args:
+        TODO
+    """
+    super().__init__()
+    self.n_layers = n_layers
+    self.base_size = base_size
+    self.sig = nn.Sigmoid()
+    self.layers = []
+
+    for i in range(n_layers):
+        side = base_size * (scale**i)
+        tim = torch.randn((3, side, side)).to(device)*(decay**i)
+        self.layers.append(tim)
+
+    self.scalers = [nn.Upsample(scale_factor=out_size/(l.shape[1]), mode='bilinear', align_corners=False) for l in self.layers]
+    
+    self.preview_scalers = [nn.Upsample(scale_factor=224/(l.shape[1]), mode='bilinear', align_corners=False) for l in self.layers]
+    
+    if init_image != None: # Given a PIL image, decompose it into a stack
+      downscalers = [nn.Upsample(scale_factor=(l.shape[1]/out_size), mode='bilinear', align_corners=False) for l in self.layers]
+      final_side = base_size * (scale ** n_layers)
+      im = torch.tensor(np.array(init_image.resize((out_size, out_size)))/255).clip(1e-03, 1-1e-3) # Between 0 and 1 (non-inclusive)
+      im = im.permute(2, 0, 1).unsqueeze(0).to(device) # torch.log(im/(1-im))
+      for i in range(n_layers):self.layers[i] *= 0 # Sero out the layers
+      for i in range(n_layers):
+        side = base_size * (scale**i)
+        out = self.forward()
+        residual = (torch.logit(im) - torch.logit(out))
+        Image.fromarray((torch.logit(residual).detach().cpu().squeeze().permute([1, 2, 0]) * 255).numpy().astype(np.uint8)).save(f'residual{i}.png')
+        self.layers[i] = downscalers[i](residual).squeeze()
+    
+    for l in self.layers: l.requires_grad = True
+
+  def forward(self):
+    im = self.scalers[0](self.layers[0].unsqueeze(0))
+    for i in range(1, self.n_layers):
+      im += self.scalers[i](self.layers[i].unsqueeze(0))
+    return self.sig(im)
+
+  def preview(self, n_preview=2):
+    im = self.preview_scalers[0](self.layers[0].unsqueeze(0))
+    for i in range(1, n_preview):
+      im += self.preview_scalers[i](self.layers[i].unsqueeze(0))
+    return self.sig(im)
+  
+  def to_pil(self):
+    return Image.fromarray((self.forward().detach().cpu().squeeze().permute([1, 2, 0]) * 255).numpy().astype(np.uint8))
+
+  def preview_pil(self):
+    return Image.fromarray((self.preview().detach().cpu().squeeze().permute([1, 2, 0]) * 255).numpy().astype(np.uint8))
+
+  def save(self, fn):
+    self.to_pil().save(fn)
+
+  def plot_layers(self):
+    fig, axs = plt.subplots(1, self.n_layers, figsize=(15, 5))
+    for i in range(self.n_layers):
+      im = (self.sig(self.layers[i].unsqueeze(0)).detach().cpu().squeeze().permute([1, 2, 0]) * 255).numpy().astype(np.uint8)
+      axs[i].imshow(im)
+      
+ 
+
 
 def generate(text, n_steps):
+  # Encode prompt
+  embed = perceptor.encode_text(clip.tokenize(text).to(device)).float()
   #todo
   return np.random.random((128, 128, 3)).astype(np.uint8)