Created app.py

app.py

@@ -0,0 +1,66 @@
+import torch
+from tqdm import tqdm
+from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
+import torchvision
+import BigGAN_utils.utils as utils
+import clip
+import torch.nn.functional as F
+from DiffAugment_pytorch import DiffAugment
+import numpy as np
+from fusedream_utils import FuseDreamBaseGenerator, get_G, save_image
+import gradio as gr
+
+
+def Text_To_Img(prompt, INIT_ITERS, OPT_ITERS, NUM_BASIS, MODEL, seed):
+  import sys
+  sys.argv = [''] ### workaround to deal with the argparse in Jupyter
+
+  utils.seed_rng(seed)
+  if MODEL == "biggan-256":
+    G, config = get_G(256) 
+  elif MODEL == "biggan-512":
+    G, config = get_G(512) 
+  else:
+    raise Exception('Model not supported')
+
+  generator = FuseDreamBaseGenerator(G, config, 10)
+  z_cllt, y_cllt = generator.generate_basis(prompt, init_iters=INIT_ITERS, num_basis=NUM_BASIS)
+  z_cllt_save = torch.cat(z_cllt).cpu().numpy()
+  y_cllt_save = torch.cat(y_cllt).cpu().numpy()
+  img, z, y = generator.optimize_clip_score(z_cllt, y_cllt, sentence, latent_noise=False, augment=True, opt_iters=OPT_ITERS, optimize_y=True)
+
+  score = generator.measureAugCLIP(z, y, sentence, augment=True, num_samples=20)
+
+  return img, score
+
+
+demo = gr.Interface(
+    fn = Text_To_Img,
+    inputs = [
+              gr.inputs.Textbox(lines=1, placeholder="Enter your text prompt here.", default="", label="Prompt", optional=False),
+              gr.inputs.Slider(100, 10000, step=20, default=500, label="Init_Iters"),
+              gr.inputs.Slider(100, 10000, step=20, default=500, label="Opt_Iters"), 
+              gr.inputs.Slider(0, 1000, step=5, default=5, label="Num_Basis"), 
+              gr.inputs.Radio(["biggan-256", "biggan-512"], default="biggan-512", label="Model"), 
+              gr.inputs.Slider(0, 1000, step=0, default=0, label="Seed")
+              ],
+
+    outputs = [
+               gr.outputs.Image(type="auto", label="Generated Image"), 
+               gr.outputs.Textbox(type="str", label="AugCLIP Score")
+               ],
+
+    #live=True,
+    #examples = [
+    #            ["A beautiful morning sunrise in the fields", 1000, 1000, 5, "biggan-512", 0],
+    #            ["A black pyramid", 500, 500, 5, "biggan-512", 1729],
+    #            ["Two dogs sitting on a bench in a park", 500, 500, 5, "biggan-256", 1234],
+    #            ],
+    #theme = "dark",
+    title = "Text to Image Generation",
+    description = "This was achieved using a combination of BIGGAN as the image generator and AugCLIP as the image-prompt correlator. Below are some examples. Try them for yourself!",
+)
+
+
+if __name__ == "__main__":
+  app, local_url, share_url = demo.launch(debug = True)