Create app.py

app.py

@@ -0,0 +1,262 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd import Variable
+import torch.optim as optim
+
+from imstack.core import ImStack
+from tqdm.notebook import tqdm
+
+import kornia.augmentation as K
+from CLIP import clip
+from torchvision import transforms
+
+from PIL import Image
+import numpy as np
+import math
+
+from matplotlib import pyplot as plt
+from fastprogress.fastprogress import master_bar, progress_bar
+from IPython.display import HTML
+from base64 import b64encode
+
+import warnings
+warnings.filterwarnings('ignore') # Some pytorch functions give warnings about behaviour changes that I don't want to see over and over again :)
+
+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 
+
+#Load CLOOB model
+import sys
+sys.path.append('./cloob-training')
+from cloob_training import model_pt, pretrained
+
+config = pretrained.get_config('cloob_laion_400m_vit_b_16_16_epochs')
+cloob = model_pt.get_pt_model(config)
+checkpoint = pretrained.download_checkpoint(config)
+cloob.load_state_dict(model_pt.get_pt_params(config, checkpoint))
+cloob.eval().requires_grad_(False).to(device)
+print('done')
+
+# Load fastai model
+
+import gradio as gr
+from fastai.vision.all import *
+from os.path import exists
+import requests
+
+model_fn = 'quick_224px'
+url = 'https://huggingface.co/johnowhitaker/sketchy_unet_rn34/resolve/main/quick_224px'
+
+if not exists(model_fn):
+  print('starting download')
+  with requests.get(url, stream=True) as r:
+    r.raise_for_status()
+    with open(model_fn, 'wb') as f:
+      for chunk in r.iter_content(chunk_size=8192):
+        f.write(chunk)
+  print('done')
+else:
+  print('file exists')
+
+def get_x(item):return None
+def get_y(item):return None
+sketch_model = load_learner(model_fn)
+
+# Cutouts
+cutn=16
+cut_pow=1
+make_cutouts = MakeCutouts(cloob.config['image_encoder']['image_size'], cutn, cut_pow)
+
+def process_im(image_path, 
+               sketchify_first=True,
+               prompt='A watercolor painting of a face',
+               lr=0.03,
+               n_iter=10
+               ):
+  
+  n_iter = int(n_iter)
+
+  pil_im = None
+
+  if sketchify_first:
+    pred = sketch_model.predict(image_path)
+    np_im = pred[0].permute(1, 2, 0).numpy()
+    pil_im = Image.fromarray(np_im.astype(np.uint8))
+  else:
+    pil_im = Image.open(image_path).resize((540, 540))
+
+
+  prompt_texts = [prompt]
+  weight_decay=1e-4
+  
+  out_size=540
+  base_size=8
+  n_layers=5
+  scale=3
+  layer_decay = 0.3
+
+
+  # The prompts
+  p_prompts = []
+  for pr in prompt_texts:
+    embed = cloob.text_encoder(cloob.tokenize(pr).to(device)).float()
+    p_prompts.append(Prompt(embed, 1, float('-inf')).to(device)) # 1 is the weight
+
+  # Some negative prompts
+  n_prompts = []
+  for pr in ["Random noise", 'saturated rainbow RGB deep dream']:
+    embed = cloob.text_encoder(cloob.tokenize(pr).to(device)).float()
+    n_prompts.append(Prompt(embed, 0.5, float('-inf')).to(device)) # 0.5 is the weight
+
+  # The ImageStack - trying a different scale and n_layers
+  ims = ImStack(base_size=base_size, 
+                scale=scale, 
+                n_layers=n_layers, 
+                out_size=out_size, 
+                decay=layer_decay,
+                init_image = pil_im)
+
+  # desaturate starting image
+  desat =  0.6#@param
+
+  if desat != 1:
+    for i in range(n_layers):
+      ims.layers[i] = ims.layers[i].detach()*desat
+      ims.layers[i].requires_grad = True
+
+
+  optimizer = optim.Adam(ims.layers, lr=lr, weight_decay=weight_decay)
+  losses = []
+
+  for i in tqdm(range(n_iter)):
+    optimizer.zero_grad()
+
+    im = ims()
+    batch = cloob.normalize(make_cutouts(im))
+    iii = cloob.image_encoder(batch).float()
+    
+    l = 0
+    for prompt in p_prompts:
+      l += prompt(iii)
+    for prompt in n_prompts:
+      l -= prompt(iii)
+
+    losses.append(float(l.detach().cpu()))
+    l.backward() # Backprop
+    optimizer.step() # Update
+
+  return ims.to_pil()
+
+from gradio.inputs import Checkbox
+iface = gr.Interface(fn=process_im, 
+  inputs=[
+          gr.inputs.Image(label="Input Image", shape=(512, 512), type="filepath"),
+          gr.inputs.Checkbox(label='Sketchify First', default=True),
+          gr.inputs.Textbox(default="A charcoal and watercolor sketch of a person", label="Prompt"),
+          gr.inputs.Number(default=0.03, label='LR'),
+          gr.inputs.Number(default=10, label='num_steps'),
+
+          ], 
+  outputs=[gr.outputs.Image(type="pil", label="Model Output")],
+  title = 'Sketchy ImStack + CLOOB', description = "Stylize an image with ImStack+CLOOB after a Sketchy Unet", 
+  article = "More info on datasciencecastnet.home.blog"
+)
+iface.launch(enable_queue=True)