feat: update gradio app

app/gradio/app_gradio.py

@@ -18,12 +18,16 @@ from PIL import Image
 import numpy as np
 import matplotlib.pyplot as plt
 
-
 from vqgan_jax.modeling_flax_vqgan import VQModel
 from dalle_mini.model import CustomFlaxBartForConditionalGeneration
 
+# ## CLIP Scoring
+from transformers import CLIPProcessor, FlaxCLIPModel
+
 import gradio as gr
 
+from dalle_mini.helpers import captioned_strip
+
 
 DALLE_REPO = 'flax-community/dalle-mini'
 DALLE_COMMIT_ID = '4d34126d0df8bc4a692ae933e3b902a1fa8b6114'
@@ -58,34 +62,12 @@ def generate(input, rng, params):
 def get_images(indices, params):
     return vqgan.decode_code(indices, params=params)
 
-def plot_images(images):
-    fig = plt.figure(figsize=(40, 20))
-    columns = 4
-    rows = 2
-    plt.subplots_adjust(hspace=0, wspace=0)
-
-    for i in range(1, columns*rows +1):
-        fig.add_subplot(rows, columns, i)
-        plt.imshow(images[i-1])
-    plt.gca().axes.get_yaxis().set_visible(False)
-    plt.show()
-    
-def stack_reconstructions(images):
-    w, h = images[0].size[0], images[0].size[1]
-    img = Image.new("RGB", (len(images)*w, h))
-    for i, img_ in enumerate(images):
-        img.paste(img_, (i*w,0))
-    return img
-
 p_generate = jax.pmap(generate, "batch")
 p_get_images = jax.pmap(get_images, "batch")
 
 bart_params = replicate(model.params)
 vqgan_params = replicate(vqgan.params)
 
-# ## CLIP Scoring
-from transformers import CLIPProcessor, FlaxCLIPModel
-
 clip = FlaxCLIPModel.from_pretrained("openai/clip-vit-base-patch32")
 print("Initialize FlaxCLIPModel")
 processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
@@ -137,7 +119,7 @@ def top_k_predictions(prompt, num_candidates=32, k=8):
 
 def run_inference(prompt, num_images=32, num_preds=8):
     images = top_k_predictions(prompt, num_candidates=num_images, k=num_preds)
-    predictions = compose_predictions(images)
+    predictions = captioned_strip(images)
     output_title = f"""
     <b>{prompt}</b>
     """
@@ -152,7 +134,7 @@ description = """
 DALL·E-mini is an AI model that generates images from any prompt you give! Generate images from text:
 """
 gr.Interface(run_inference, 
-    inputs=[gr.inputs.Textbox(label='What do you want to see?')], #, gr.inputs.Slider(1,64,1,8, label='Candidates to generate'), gr.inputs.Slider(1,8,1,1, label='Best predictions to show')], 
+    inputs=[gr.inputs.Textbox(label='What do you want to see?')],
     outputs=outputs, 
     title='DALL·E mini',
     description=description,

app/gradio/app_gradio_ngrok.py

@@ -7,25 +7,15 @@ import numpy as np
 import matplotlib.pyplot as plt
 from io import BytesIO
 import base64
+import os
 
 import gradio as gr
 
-# If we use streamlit, this would be exported as a streamlit secret
-import os
-backend_url = os.environ["BACKEND_SERVER"]
+from dalle_mini.helpers import captioned_strip
 
-def compose_predictions(images, caption=None):
-    increased_h = 0 if caption is None else 48
-    w, h = images[0].size[0], images[0].size[1]
-    img = Image.new("RGB", (len(images)*w, h + increased_h))
-    for i, img_ in enumerate(images):
-        img.paste(img_, (i*w, increased_h))
 
-    if caption is not None:
-        draw = ImageDraw.Draw(img)
-        font = ImageFont.truetype("/usr/share/fonts/truetype/liberation2/LiberationMono-Bold.ttf", 40)
-        draw.text((20, 3), caption, (255,255,255), font=font)
-    return img
+backend_url = os.environ["BACKEND_SERVER"]
+
 
 class ServiceError(Exception):
     def __init__(self, status_code):
@@ -46,7 +36,7 @@ def get_images_from_ngrok(prompt):
 def run_inference(prompt):
     try:
         images = get_images_from_ngrok(prompt)
-        predictions = compose_predictions(images)
+        predictions = captioned_strip(images)
         output_title = f"""
         <p style="font-size:22px; font-style:bold">Best predictions</p>
         <p>We asked our model to generate 128 candidates for your prompt:</p>

app/sample_images/readme.txt

@@ -1 +0,0 @@
-These images were generated by one of our checkpoints, as responses to the prompt "snowy mountains by the sea".

app/ui_gradio.py

@@ -1,91 +0,0 @@
-#!/usr/bin/env python
-# coding: utf-8
-
-from PIL import Image
-import gradio as gr
-
-def compose_predictions(images, caption=None):
-    increased_h = 0 if caption is None else 48
-    w, h = images[0].size[0], images[0].size[1]
-    img = Image.new("RGB", (len(images)*w, h + increased_h))
-    for i, img_ in enumerate(images):
-        img.paste(img_, (i*w, increased_h))
-
-    if caption is not None:
-        draw = ImageDraw.Draw(img)
-        font = ImageFont.truetype("/usr/share/fonts/truetype/liberation2/LiberationMono-Bold.ttf", 40)
-        draw.text((20, 3), caption, (255,255,255), font=font)
-    return img
-
-def compose_predictions_grid(images):
-    cols = 4
-    rows = len(images) // cols
-    w, h = images[0].size[0], images[0].size[1]
-    img = Image.new("RGB", (w * cols, h * rows))
-    for i, img_ in enumerate(images):
-        row = i // cols
-        col = i % cols
-        img.paste(img_, (w * col, h * row))
-    return img
-
-def top_k_predictions_real(prompt, num_candidates=32, k=8):
-    images = hallucinate(prompt, num_images=num_candidates)
-    images = clip_top_k(prompt, images, k=num_preds)
-    return images
-
-def top_k_predictions(prompt, num_candidates=32, k=8):
-    images = []
-    for i in range(k):
-        image = Image.open(f"sample_images/image_{i}.jpg")
-        images.append(image)
-    return images
-
-def run_inference(prompt, num_images=32, num_preds=8):
-    images = top_k_predictions(prompt, num_candidates=num_images, k=num_preds)
-    predictions = compose_predictions(images)
-    output_title = f"""
-    <p style="font-size:22px; font-style:bold">Best predictions</p>
-    <p>We asked our model to generate 32 candidates for your prompt:</p>
-
-    <pre>
-
-    <b>{prompt}</b>
-    </pre>
-    <p>We then used a pre-trained <a href="https://huggingface.co/openai/clip-vit-base-patch32">CLIP model</a> to score them according to the
-    similarity of the text and the image representations.</p>
-
-    <p>This is the result:</p>
-    """
-    output_description = """
-    <p>Read more about the process <a href="https://wandb.ai/dalle-mini/dalle-mini/reports/DALL-E-mini--Vmlldzo4NjIxODA">in our report</a>.<p>
-    <p style='text-align: center'>Created with <a href="https://github.com/borisdayma/dalle-mini">DALLE·mini</a></p>
-    """
-    return (output_title, predictions, output_description)
-
-outputs = [
-    gr.outputs.HTML(label=""),      # To be used as title
-    gr.outputs.Image(label=''),
-    gr.outputs.HTML(label=""),      # Additional text that appears in the screenshot
-]
-
-description = """
-Welcome to our demo of DALL·E-mini. This project was created on TPU v3-8s during the 🤗 Flax / JAX Community Week.
-It reproduces the essential characteristics of OpenAI's DALL·E, at a fraction of the size.
-
-Please, write what you would like the model to generate, or select one of the examples below.
-"""
-gr.Interface(run_inference, 
-    inputs=[gr.inputs.Textbox(label='Prompt')], #, gr.inputs.Slider(1,64,1,8, label='Candidates to generate'), gr.inputs.Slider(1,8,1,1, label='Best predictions to show')], 
-    outputs=outputs, 
-    title='DALL·E mini',
-    description=description,
-    article="<p style='text-align: center'> DALLE·mini by Boris Dayma et al. | <a href='https://github.com/borisdayma/dalle-mini'>GitHub</a></p>",
-    layout='vertical',
-    theme='huggingface',
-    examples=[['an armchair in the shape of an avocado'], ['snowy mountains by the sea']],
-    allow_flagging=False,
-    live=False,
-    server_port=8999
-).launch(
-    share=True     # Creates temporary public link if true
-)