File size: 5,588 Bytes
6b0d541
 
 
 
0ca6514
 
6b0d541
 
 
 
 
 
0ca6514
6b0d541
0ca6514
6b0d541
0ca6514
6b0d541
 
11ae595
7f962d6
6b0d541
a0b5dc7
 
 
6b0d541
 
0ca6514
a0b5dc7
6b0d541
0ca6514
 
6b0d541
0ca6514
 
6b0d541
7f962d6
0ca6514
 
 
7f962d6
6b0d541
0ca6514
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6b0d541
0ca6514
 
6b0d541
 
 
 
 
0ca6514
6b0d541
 
 
 
 
 
 
 
 
 
 
 
0ca6514
6b0d541
 
 
0ca6514
6b0d541
 
 
 
 
 
 
 
 
 
 
0ca6514
6b0d541
 
0ca6514
 
 
 
 
 
 
6b0d541
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0ca6514
6b0d541
 
 
 
 
 
 
0ca6514
85eab14
6b0d541
 
0ca6514
6b0d541
0ca6514
6b0d541
 
 
0ca6514
 
 
 
6b0d541
 
0ca6514
85eab14
 
 
 
 
0ca6514
85eab14
 
a0b5dc7
85eab14
 
 
139d801
85eab14
0ca6514
85eab14
0ca6514
 
85eab14
 
 
139d801
85eab14
0ca6514
 
 
 
 
85eab14
139d801
0ca6514
 
 
 
 
 
85eab14
 
 
adcb063
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
#!/usr/bin/env python
# coding: utf-8

# Uncomment to run on cpu
# import os
# os.environ["JAX_PLATFORM_NAME"] = "cpu"

import random

import jax
import flax.linen as nn
from flax.training.common_utils import shard
from flax.jax_utils import replicate

from transformers import BartTokenizer

from PIL import Image, ImageDraw, ImageFont
import numpy as np

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 PIL import Image, ImageDraw, ImageFont


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)