go

app.py

@@ -1,10 +1,89 @@
-from dalle_pytorch import DiscreteVAE, OpenAIDiscreteVAE, VQGanVAE, DALLE
-from dalle_pytorch.tokenizer import tokenizer, HugTokenizer, YttmTokenizer, ChineseTokenizer
+from tqdm import tqdm
+import numpy as np
 
+from pathlib import Path
+import json
+
+# torch
+
+import torch
+
+from einops import repeat
+
+# vision imports
+
+from PIL import Image
+
+# dalle related classes and utils
+
+from dalle_pytorch import VQGanVAE, DALLE
+from dalle_pytorch.tokenizer import tokenizer
+
+from io import BytesIO
 import gradio as gr
 
-def greet(name):
-    return "Hello " + name + "!!"
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
+# load DALL-E
+
+def exists(val):
+    return val is not None
+
+models = json.load(open("model_paths.json"))
+
+
+vae = VQGanVAE(None, None)
+
+dalles = {}
+
+for name, model_path in models.items():
+    assert Path(model_path).exists(), 'trained DALL-E '+model_path+' must exist'
+    load_obj = torch.load(model_path)
+    dalle_params, _, weights = load_obj.pop('hparams'), load_obj.pop('vae_params'), load_obj.pop('weights')
+    dalle_params.pop('vae', None) # cleanup later
+    
+    dalle = DALLE(vae = vae, **dalle_params).cuda()
+    
+    dalle.load_state_dict(weights)
+    dalles[name] = dalle
+
+
+
+batch_size = 4
+
+top_k = 0.9
+
+# generate images
+
+image_size = vae.image_size
+
+def generate(text):
+    text_input = text
+    num_images = 4
+    dalle_name = "weird_car"
+    dalle = dalles[dalle_name]
+
+    text = tokenizer.tokenize([text_input], dalle.text_seq_len).cuda()
+
+    text = repeat(text, '() n -> b n', b = num_images)
+
+    outputs = []
+
+    for text_chunk in tqdm(text.split(batch_size), desc = f'generating images for - {text}'):
+        output = dalle.generate_images(text_chunk, filter_thres = top_k)
+        outputs.append(output)
+
+    outputs = torch.cat(outputs)
+
+    response = []
+
+    for image in tqdm(outputs, desc = 'saving images'):
+        np_image = np.moveaxis(image.cpu().numpy(), 0, -1)
+        formatted = (np_image * 255).astype('uint8')
+
+        img = Image.fromarray(formatted)
+        response.append(img)
+
+    return response
+
+iface = gr.Interface(fn=generate, inputs="text", outputs=gr.outputs.Carousel("image"))
 iface.launch()

model_paths.json

@@ -0,0 +1 @@
+{"weird_car":"weird_car_model_continue.pt"}

requirements.txt

@@ -1 +1,6 @@
 dalle-pytorch
+numpy
+tqdm
+torch
+torchvision
+einops