feat(app): migrate demo to Gradio by @AK391 (#179)

README.md

@@ -3,8 +3,9 @@ title: DALL·E mini
 emoji: 🥑
 colorFrom: yellow
 colorTo: green
-sdk: streamlit
-app_file: app/streamlit/app.py
+sdk: gradio
+sdk_version: 3.0b6
+app_file: app/gradio/app.py
 pinned: True
 license: apache-2.0
 ---

app/gradio/app.py

@@ -0,0 +1,53 @@
+#!/usr/bin/env python
+# coding: utf-8
+import os
+
+import gradio as gr
+from backend import get_images_from_backend
+
+block = gr.Blocks(css=".container { max-width: 800px; margin: auto; }")
+backend_url = os.environ["BACKEND_SERVER"] + "/generate"
+
+
+def infer(prompt):
+    response = get_images_from_backend(prompt, backend_url)
+    return response["images"]
+
+
+with block:
+    gr.Markdown("<h1><center>DALL·E mini</center></h1>")
+    gr.Markdown(
+        "DALL·E mini is an AI model that generates images from any prompt you give!"
+    )
+    with gr.Group():
+        with gr.Box():
+            with gr.Row().style(mobile_collapse=False, equal_height=True):
+
+                text = gr.Textbox(
+                    label="Enter your prompt", show_label=False, max_lines=1
+                ).style(
+                    border=(True, False, True, True),
+                    margin=False,
+                    rounded=(True, False, False, True),
+                    container=False,
+                )
+                btn = gr.Button("Run").style(
+                    margin=False,
+                    rounded=(False, True, True, False),
+                )
+        gallery = gr.Gallery(label="Generated images", show_label=False).style(
+            grid=[3], height="auto"
+        )
+        btn.click(infer, inputs=text, outputs=gallery)
+
+    gr.Markdown(
+        """___
+   <p style='text-align: center'>
+   Created by Boris Dayma et al. 2021-2022
+   <br/>
+   <a href="https://github.com/borisdayma/dalle-mini" target="_blank">GitHub</a> | <a href="https://wandb.ai/dalle-mini/dalle-mini/reports/DALL-E-mini--Vmlldzo4NjIxODA" target="_blank">Project Report</a>
+   </p>"""
+    )
+
+
+block.launch()

app/gradio/app_gradio.py

@@ -1,179 +0,0 @@
-#!/usr/bin/env python
-# coding: utf-8
-
-# Uncomment to run on cpu
-# import os
-# os.environ["JAX_PLATFORM_NAME"] = "cpu"
-
-import random
-
-import gradio as gr
-import jax
-import numpy as np
-from flax.jax_utils import replicate
-from flax.training.common_utils import shard
-from PIL import Image, ImageDraw, ImageFont
-
-# ## CLIP Scoring
-from transformers import BartTokenizer, CLIPProcessor, FlaxCLIPModel
-from vqgan_jax.modeling_flax_vqgan import VQModel
-
-from dalle_mini.model import CustomFlaxBartForConditionalGeneration
-
-DALLE_REPO = "flax-community/dalle-mini"
-DALLE_COMMIT_ID = "4d34126d0df8bc4a692ae933e3b902a1fa8b6114"
-
-VQGAN_REPO = "flax-community/vqgan_f16_16384"
-VQGAN_COMMIT_ID = "90cc46addd2dd8f5be21586a9a23e1b95aa506a9"
-
-tokenizer = BartTokenizer.from_pretrained(DALLE_REPO, revision=DALLE_COMMIT_ID)
-model = CustomFlaxBartForConditionalGeneration.from_pretrained(
-    DALLE_REPO, revision=DALLE_COMMIT_ID
-)
-vqgan = VQModel.from_pretrained(VQGAN_REPO, revision=VQGAN_COMMIT_ID)
-
-
-def captioned_strip(images, caption=None, rows=1):
-    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 // rows, h * rows + increased_h))
-    for i, img_ in enumerate(images):
-        img.paste(img_, (i // rows * w, increased_h + (i % rows) * 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 custom_to_pil(x):
-    x = np.clip(x, 0.0, 1.0)
-    x = (255 * x).astype(np.uint8)
-    x = Image.fromarray(x)
-    if not x.mode == "RGB":
-        x = x.convert("RGB")
-    return x
-
-
-def generate(input, rng, params):
-    return model.generate(
-        **input,
-        max_length=257,
-        num_beams=1,
-        do_sample=True,
-        prng_key=rng,
-        eos_token_id=50000,
-        pad_token_id=50000,
-        params=params,
-    )
-
-
-def get_images(indices, params):
-    return vqgan.decode_code(indices, params=params)
-
-
-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 = FlaxCLIPModel.from_pretrained("openai/clip-vit-base-patch32")
-print("Initialize FlaxCLIPModel")
-processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
-print("Initialize CLIPProcessor")
-
-
-def hallucinate(prompt, num_images=64):
-    prompt = [prompt] * jax.device_count()
-    inputs = tokenizer(
-        prompt,
-        return_tensors="jax",
-        padding="max_length",
-        truncation=True,
-        max_length=128,
-    ).data
-    inputs = shard(inputs)
-
-    all_images = []
-    for i in range(num_images // jax.device_count()):
-        key = random.randint(0, 1e7)
-        rng = jax.random.PRNGKey(key)
-        rngs = jax.random.split(rng, jax.local_device_count())
-        indices = p_generate(inputs, rngs, bart_params).sequences
-        indices = indices[:, :, 1:]
-
-        images = p_get_images(indices, vqgan_params)
-        images = np.squeeze(np.asarray(images), 1)
-        for image in images:
-            all_images.append(custom_to_pil(image))
-    return all_images
-
-
-def clip_top_k(prompt, images, k=8):
-    inputs = processor(text=prompt, images=images, return_tensors="np", padding=True)
-    outputs = clip(**inputs)
-    logits = outputs.logits_per_text
-    scores = np.array(logits[0]).argsort()[-k:][::-1]
-    return [images[score] for score in scores]
-
-
-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 top_k_predictions(prompt, num_candidates=32, k=8):
-    images = hallucinate(prompt, num_images=num_candidates)
-    images = clip_top_k(prompt, images, k=k)
-    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 = captioned_strip(images)
-    output_title = f"""
-    <b>{prompt}</b>
-    """
-    return (output_title, predictions)
-
-
-outputs = [
-    gr.outputs.HTML(label=""),  # To be used as title
-    gr.outputs.Image(label=""),
-]
-
-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?")],
-    outputs=outputs,
-    title="DALL·E mini",
-    description=description,
-    article="<p style='text-align: center'> Created by Boris Dayma et al. 2021 | <a href='https://github.com/borisdayma/dalle-mini'>GitHub</a> | <a href='https://wandb.ai/dalle-mini/dalle-mini/reports/DALL-E-mini--Vmlldzo4NjIxODA'>Report</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)

app/gradio/backend.py

@@ -0,0 +1,33 @@
+# Client requests to Dalle-Mini Backend server
+
+import base64
+from io import BytesIO
+
+import requests
+from PIL import Image
+
+
+class ServiceError(Exception):
+    def __init__(self, status_code):
+        self.status_code = status_code
+
+
+def get_images_from_backend(prompt, backend_url):
+    r = requests.post(backend_url, json={"prompt": prompt})
+    if r.status_code == 200:
+        json = r.json()
+        images = json["images"]
+        images = [Image.open(BytesIO(base64.b64decode(img))) for img in images]
+        version = json.get("version", "unknown")
+        return {"images": images, "version": version}
+    else:
+        raise ServiceError(r.status_code)
+
+
+def get_model_version(url):
+    r = requests.get(url)
+    if r.status_code == 200:
+        version = r.json()["version"]
+        return version
+    else:
+        raise ServiceError(r.status_code)

app/gradio/requirements.txt

@@ -1,4 +0,0 @@
-# Requirements for huggingface spaces
-gradio>=2.2.3
-flax
-transformers

app/streamlit/app.py

@@ -1,8 +1,6 @@
 #!/usr/bin/env python
 # coding: utf-8
 
-from datetime import datetime
-
 import streamlit as st
 from backend import ServiceError, get_images_from_backend