Update app.py

app.py

@@ -1,11 +1,13 @@
 import torch
 torch.hub.download_url_to_file('http://mirror.io.community/blob/vqgan/vqgan_imagenet_f16_16384.yaml', 'vqgan_imagenet_f16_16384.yaml')
 torch.hub.download_url_to_file('http://mirror.io.community/blob/vqgan/vqgan_imagenet_f16_16384.ckpt', 'vqgan_imagenet_f16_16384.ckpt')
+# import torch
+# torch.hub.download_url_to_file('http://mirror.io.community/blob/vqgan/vqgan_imagenet_f16_16384.yaml', 'vqgan_imagenet_f16_16384.yaml')
+# torch.hub.download_url_to_file('http://mirror.io.community/blob/vqgan/vqgan_imagenet_f16_16384.ckpt', 'vqgan_imagenet_f16_16384.ckpt')
 import argparse
 import math
 from pathlib import Path
 import sys
-
 sys.path.insert(1, './taming-transformers')
 #from IPython import display
 from base64 import b64encode
@@ -18,7 +20,6 @@ from torch.nn import functional as F
 from torchvision import transforms
 from torchvision.transforms import functional as TF
 from tqdm.notebook import tqdm
-
 from CLIP import clip
 import kornia.augmentation as K
 import numpy as np
@@ -26,17 +27,12 @@ import imageio
 from PIL import ImageFile, Image
 ImageFile.LOAD_TRUNCATED_IMAGES = True
 import gradio as gr
-
 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])
@@ -45,44 +41,31 @@ def ramp(ratio, width):
         out[i] = cur
         cur += ratio
     return torch.cat([-out[1:].flip([0]), out])[1:-1]
-
-
 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
-
-
 class ClampWithGrad(torch.autograd.Function):
     @staticmethod
     def forward(ctx, input, min, max):
@@ -90,51 +73,38 @@ class ClampWithGrad(torch.autograd.Function):
         ctx.max = max
         ctx.save_for_backward(input)
         return input.clamp(min, max)
-
     @staticmethod
     def backward(ctx, grad_in):
         input, = ctx.saved_tensors
         return grad_in * (grad_in * (input - input.clamp(ctx.min, ctx.max)) >= 0), None, None
-
-
 clamp_with_grad = ClampWithGrad.apply
-
-
 def vector_quantize(x, codebook):
     d = x.pow(2).sum(dim=-1, keepdim=True) + codebook.pow(2).sum(dim=1) - 2 * x @ codebook.T
     indices = d.argmin(-1)
     x_q = F.one_hot(indices, codebook.shape[0]).to(d.dtype) @ codebook
     return replace_grad(x_q, x)
-
-
 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()
-
-
 def parse_prompt(prompt):
     vals = prompt.rsplit(':', 2)
     vals = vals + ['', '1', '-inf'][len(vals):]
     return vals[0], float(vals[1]), float(vals[2])
-
-
 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.RandomVerticalFlip(p=0.5),
@@ -151,7 +121,6 @@ class MakeCutouts(nn.Module):
         self.noise_fac = 0.1
         self.av_pool = nn.AdaptiveAvgPool2d((self.cut_size, self.cut_size))
         self.max_pool = nn.AdaptiveMaxPool2d((self.cut_size, self.cut_size))
-
     def forward(self, input):
         sideY, sideX = input.shape[2:4]
         max_size = min(sideX, sideY)
@@ -159,13 +128,11 @@ class MakeCutouts(nn.Module):
         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)))
-
             # cutout = transforms.Resize(size=(self.cut_size, self.cut_size))(input)
             
             cutout = (self.av_pool(input) + self.max_pool(input))/2
@@ -175,8 +142,6 @@ class MakeCutouts(nn.Module):
             facs = batch.new_empty([self.cutn, 1, 1, 1]).uniform_(0, self.noise_fac)
             batch = batch + facs * torch.randn_like(batch)
         return batch
-
-
 def load_vqgan_model(config_path, checkpoint_path):
     config = OmegaConf.load(config_path)
     if config.model.target == 'taming.models.vqgan.VQModel':
@@ -196,8 +161,6 @@ def load_vqgan_model(config_path, checkpoint_path):
         raise ValueError(f'unknown model type: {config.model.target}')
     del model.loss
     return model
-
-
 def resize_image(image, out_size):
     ratio = image.size[0] / image.size[1]
     area = min(image.size[0] * image.size[1], out_size[0] * out_size[1])
@@ -205,82 +168,66 @@ def resize_image(image, out_size):
     return image.resize(size, Image.LANCZOS)
 
 
+model_name = "vqgan_imagenet_f16_16384" 
+images_interval =  50
+width =  256
+height = 256
+init_image = ""
+seed = 42
+args = argparse.Namespace(
+    noise_prompt_seeds=[],
+    noise_prompt_weights=[],
+    size=[width, height],
+    init_image=init_image,
+    init_weight=0.,
+    clip_model='ViT-B/32',
+    vqgan_config=f'{model_name}.yaml',
+    vqgan_checkpoint=f'{model_name}.ckpt',
+    step_size=0.1,
+    cutn=1,
+    cut_pow=1.,
+    display_freq=images_interval,
+    seed=seed,
+)
+device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+print('Using device:', device)
+model = load_vqgan_model(args.vqgan_config, args.vqgan_checkpoint).to(device)
+perceptor = clip.load(args.clip_model, jit=False)[0].eval().requires_grad_(False).to(device)
+
 def inference(text):
     texts = text
-    width =  128
-    height = 128
-    model = "vqgan_imagenet_f16_16384" 
-    images_interval =  50
-    init_image = ""
     target_images = ""
-    seed = 42
     max_iterations = 100
-
     model_names={"vqgan_imagenet_f16_16384": 'ImageNet 16384',"vqgan_imagenet_f16_1024":"ImageNet 1024", 'vqgan_openimages_f16_8192':'OpenImages 8912',
                     "wikiart_1024":"WikiArt 1024", "wikiart_16384":"WikiArt 16384", "coco":"COCO-Stuff", "faceshq":"FacesHQ", "sflckr":"S-FLCKR"}
-    name_model = model_names[model]     
-
-    if seed == -1:
-        seed = None
-    if init_image == "None":
-        init_image = None
+    name_model = model_names[model_name]
     if target_images == "None" or not target_images:
         target_images = []
     else:
         target_images = target_images.split("|")
         target_images = [image.strip() for image in target_images]
-
     texts = [phrase.strip() for phrase in texts.split("|")]
     if texts == ['']:
         texts = []
-
-
-    args = argparse.Namespace(
-        prompts=texts,
-        image_prompts=target_images,
-        noise_prompt_seeds=[],
-        noise_prompt_weights=[],
-        size=[width, height],
-        init_image=init_image,
-        init_weight=0.,
-        clip_model='ViT-B/32',
-        vqgan_config=f'{model}.yaml',
-        vqgan_checkpoint=f'{model}.ckpt',
-        step_size=0.1,
-        cutn=1,
-        cut_pow=1.,
-        display_freq=images_interval,
-        seed=seed,
-    )
     from urllib.request import urlopen
-
-    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
-    print('Using device:', device)
     if texts:
         print('Using texts:', texts)
     if target_images:
         print('Using image prompts:', target_images)
-    if args.seed is None:
+    if args.seed is None or args.seed == -1:
         seed = torch.seed()
     else:
         seed = args.seed
     torch.manual_seed(seed)
     print('Using seed:', seed)
-
-    model = load_vqgan_model(args.vqgan_config, args.vqgan_checkpoint).to(device)
-    perceptor = clip.load(args.clip_model, jit=False)[0].eval().requires_grad_(False).to(device)
     # clock=deepcopy(perceptor.visual.positional_embedding.data)
     # perceptor.visual.positional_embedding.data = clock/clock.max()
     # perceptor.visual.positional_embedding.data=clamp_with_grad(clock,0,1)
-
     cut_size = perceptor.visual.input_resolution
-
     f = 2**(model.decoder.num_resolutions - 1)
     make_cutouts = MakeCutouts(cut_size, args.cutn, cut_pow=args.cut_pow)
-
     toksX, toksY = args.size[0] // f, args.size[1] // f
     sideX, sideY = toksX * f, toksY * f
-
     if args.vqgan_checkpoint == 'vqgan_openimages_f16_8192.ckpt':
         e_dim = 256
         n_toks = model.quantize.n_embed
@@ -293,10 +240,8 @@ def inference(text):
         z_max = model.quantize.embedding.weight.max(dim=0).values[None, :, None, None]
     # z_min = model.quantize.embedding.weight.min(dim=0).values[None, :, None, None]
     # z_max = model.quantize.embedding.weight.max(dim=0).values[None, :, None, None]
-
     # normalize_imagenet = transforms.Normalize(mean=[0.485, 0.456, 0.406],
     #                                            std=[0.229, 0.224, 0.225])
-
     if args.init_image:
         if 'http' in args.init_image:
             img = Image.open(urlopen(args.init_image))
@@ -318,20 +263,14 @@ def inference(text):
     z_orig = z.clone()
     z.requires_grad_(True)
     opt = optim.Adam([z], lr=args.step_size)
-
     normalize = transforms.Normalize(mean=[0.48145466, 0.4578275, 0.40821073],
                                     std=[0.26862954, 0.26130258, 0.27577711])
-
-
-
     pMs = []
-
-    for prompt in args.prompts:
+    for prompt in texts:
         txt, weight, stop = parse_prompt(prompt)
         embed = perceptor.encode_text(clip.tokenize(txt).to(device)).float()
         pMs.append(Prompt(embed, weight, stop).to(device))
-
-    for prompt in args.image_prompts:
+    for prompt in target_images:
         path, weight, stop = parse_prompt(prompt)
         img = Image.open(path)
         pil_image = img.convert('RGB')
@@ -339,19 +278,16 @@ def inference(text):
         batch = make_cutouts(TF.to_tensor(img).unsqueeze(0).to(device))
         embed = perceptor.encode_image(normalize(batch)).float()
         pMs.append(Prompt(embed, weight, stop).to(device))
-
     for seed, weight in zip(args.noise_prompt_seeds, args.noise_prompt_weights):
         gen = torch.Generator().manual_seed(seed)
         embed = torch.empty([1, perceptor.visual.output_dim]).normal_(generator=gen)
         pMs.append(Prompt(embed, weight).to(device))
-
     def synth(z):
         if args.vqgan_checkpoint == 'vqgan_openimages_f16_8192.ckpt':
             z_q = vector_quantize(z.movedim(1, 3), model.quantize.embed.weight).movedim(3, 1)
         else:
             z_q = vector_quantize(z.movedim(1, 3), model.quantize.embedding.weight).movedim(3, 1)
         return clamp_with_grad(model.decode(z_q).add(1).div(2), 0, 1)
-
     @torch.no_grad()
     def checkin(i, losses):
         losses_str = ', '.join(f'{loss.item():g}' for loss in losses)
@@ -359,14 +295,12 @@ def inference(text):
         out = synth(z)
         #TF.to_pil_image(out[0].cpu()).save('progress.png')
         #display.display(display.Image('progress.png'))
-
     def ascend_txt():
         # global i
         out = synth(z)
         iii = perceptor.encode_image(normalize(make_cutouts(out))).float()
         
         result = []
-
         if args.init_weight:
             # result.append(F.mse_loss(z, z_orig) * args.init_weight / 2)
             result.append(F.mse_loss(z, torch.zeros_like(z_orig)) * ((1/torch.tensor(i*2 + 1))*args.init_weight) / 2)
@@ -375,9 +309,7 @@ def inference(text):
         img = np.array(out.mul(255).clamp(0, 255)[0].cpu().detach().numpy().astype(np.uint8))[:,:,:]
         img = np.transpose(img, (1, 2, 0))
         #imageio.imwrite('./steps/' + str(i) + '.png', np.array(img))
-
         return result, np.array(img)
-
     def train(i):
         opt.zero_grad()
         lossAll, image = ascend_txt()
@@ -390,7 +322,6 @@ def inference(text):
         with torch.no_grad():
             z.copy_(z.maximum(z_min).minimum(z_max))
         return image
-
     i = 0
     try:
         with tqdm() as pbar:
@@ -404,12 +335,29 @@ def inference(text):
         pass
     return image
 
+inferences_running = 0
+
+def throttled_inference(text):
+    global inferences_running
+    current = inferences_running
+    if current >= 2:
+        print(f"Rejected inference when we already had {current} running")
+        return None
+
+    print(f"Inference starting when we already had {current} running")
+    inferences_running += 1
+    try:
+        return inference(text)
+    finally:
+        print("Inference finished")
+        inferences_running -= 1
+
+
 title = "VQGAN + CLIP"
 description = "Gradio demo for VQGAN + CLIP. To use it, simply add your text, or click one of the examples to load them. Read more at the links below. Please click submit only once"
 article = "<p style='text-align: center'>Originally made by Katherine Crowson (https://github.com/crowsonkb, https://twitter.com/RiversHaveWings). The original BigGAN+CLIP method was by https://twitter.com/advadnoun. Added some explanations and modifications by Eleiber#8347, pooling trick by Crimeacs#8222 (https://twitter.com/EarthML1) and the GUI was made with the help of Abulafia#3734. | <a href='https://colab.research.google.com/drive/1ZAus_gn2RhTZWzOWUpPERNC0Q8OhZRTZ'>Colab</a> | <a href='https://github.com/CompVis/taming-transformers'>Taming Transformers Github Repo</a> | <a href='https://github.com/openai/CLIP'>CLIP Github Repo</a></p>"
-
 gr.Interface(
-    inference, 
+    throttled_inference, 
     gr.inputs.Textbox(label="Input"), 
     gr.outputs.Image(type="numpy", label="Output"),
     title=title,