Tanishq Abraham commited on
Commit
7e6e1fe
2 Parent(s): 21a35bc 6b0d541

Merge pull request #37 from tmabraham/add-gradio-demo

Browse files

add gradio demo

Former-commit-id: a0a824c63b4e6198d111aa6bceb6afec9ae4ef0a

app/app_gradio.py ADDED
@@ -0,0 +1,194 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python
2
+ # coding: utf-8
3
+
4
+ # Uncomment to run on cpu
5
+ #import os
6
+ #os.environ["JAX_PLATFORM_NAME"] = "cpu"
7
+
8
+ import random
9
+
10
+ import jax
11
+ import flax.linen as nn
12
+ from flax.training.common_utils import shard
13
+ from flax.jax_utils import replicate, unreplicate
14
+
15
+ from transformers.models.bart.modeling_flax_bart import *
16
+ from transformers import BartTokenizer, FlaxBartForConditionalGeneration
17
+
18
+
19
+ import requests
20
+ from PIL import Image
21
+ import numpy as np
22
+ import matplotlib.pyplot as plt
23
+
24
+
25
+ from dalle_mini.vqgan_jax.modeling_flax_vqgan import VQModel
26
+
27
+ import gradio as gr
28
+
29
+
30
+ # TODO: set those args in a config file
31
+ OUTPUT_VOCAB_SIZE = 16384 + 1 # encoded image token space + 1 for bos
32
+ OUTPUT_LENGTH = 256 + 1 # number of encoded tokens + 1 for bos
33
+ BOS_TOKEN_ID = 16384
34
+ BASE_MODEL = 'flax-community/dalle-mini'
35
+
36
+ class CustomFlaxBartModule(FlaxBartModule):
37
+ def setup(self):
38
+ # we keep shared to easily load pre-trained weights
39
+ self.shared = nn.Embed(
40
+ self.config.vocab_size,
41
+ self.config.d_model,
42
+ embedding_init=jax.nn.initializers.normal(self.config.init_std, self.dtype),
43
+ dtype=self.dtype,
44
+ )
45
+ # a separate embedding is used for the decoder
46
+ self.decoder_embed = nn.Embed(
47
+ OUTPUT_VOCAB_SIZE,
48
+ self.config.d_model,
49
+ embedding_init=jax.nn.initializers.normal(self.config.init_std, self.dtype),
50
+ dtype=self.dtype,
51
+ )
52
+ self.encoder = FlaxBartEncoder(self.config, dtype=self.dtype, embed_tokens=self.shared)
53
+
54
+ # the decoder has a different config
55
+ decoder_config = BartConfig(self.config.to_dict())
56
+ decoder_config.max_position_embeddings = OUTPUT_LENGTH
57
+ decoder_config.vocab_size = OUTPUT_VOCAB_SIZE
58
+ self.decoder = FlaxBartDecoder(decoder_config, dtype=self.dtype, embed_tokens=self.decoder_embed)
59
+
60
+ class CustomFlaxBartForConditionalGenerationModule(FlaxBartForConditionalGenerationModule):
61
+ def setup(self):
62
+ self.model = CustomFlaxBartModule(config=self.config, dtype=self.dtype)
63
+ self.lm_head = nn.Dense(
64
+ OUTPUT_VOCAB_SIZE,
65
+ use_bias=False,
66
+ dtype=self.dtype,
67
+ kernel_init=jax.nn.initializers.normal(self.config.init_std, self.dtype),
68
+ )
69
+ self.final_logits_bias = self.param("final_logits_bias", self.bias_init, (1, OUTPUT_VOCAB_SIZE))
70
+
71
+ class CustomFlaxBartForConditionalGeneration(FlaxBartForConditionalGeneration):
72
+ module_class = CustomFlaxBartForConditionalGenerationModule
73
+
74
+ # create our model
75
+ # FIXME: Save tokenizer to hub so we can load from there
76
+ tokenizer = BartTokenizer.from_pretrained("facebook/bart-large-cnn")
77
+ model = CustomFlaxBartForConditionalGeneration.from_pretrained(BASE_MODEL)
78
+ model.config.force_bos_token_to_be_generated = False
79
+ model.config.forced_bos_token_id = None
80
+ model.config.forced_eos_token_id = None
81
+
82
+ vqgan = VQModel.from_pretrained("flax-community/vqgan_f16_16384")
83
+
84
+ def custom_to_pil(x):
85
+ x = np.clip(x, 0., 1.)
86
+ x = (255*x).astype(np.uint8)
87
+ x = Image.fromarray(x)
88
+ if not x.mode == "RGB":
89
+ x = x.convert("RGB")
90
+ return x
91
+
92
+ def generate(input, rng, params):
93
+ return model.generate(
94
+ **input,
95
+ max_length=257,
96
+ num_beams=1,
97
+ do_sample=True,
98
+ prng_key=rng,
99
+ eos_token_id=50000,
100
+ pad_token_id=50000,
101
+ params=params,
102
+ )
103
+
104
+ def get_images(indices, params):
105
+ return vqgan.decode_code(indices, params=params)
106
+
107
+ def plot_images(images):
108
+ fig = plt.figure(figsize=(40, 20))
109
+ columns = 4
110
+ rows = 2
111
+ plt.subplots_adjust(hspace=0, wspace=0)
112
+
113
+ for i in range(1, columns*rows +1):
114
+ fig.add_subplot(rows, columns, i)
115
+ plt.imshow(images[i-1])
116
+ plt.gca().axes.get_yaxis().set_visible(False)
117
+ plt.show()
118
+
119
+ def stack_reconstructions(images):
120
+ w, h = images[0].size[0], images[0].size[1]
121
+ img = Image.new("RGB", (len(images)*w, h))
122
+ for i, img_ in enumerate(images):
123
+ img.paste(img_, (i*w,0))
124
+ return img
125
+
126
+ p_generate = jax.pmap(generate, "batch")
127
+ p_get_images = jax.pmap(get_images, "batch")
128
+
129
+ bart_params = replicate(model.params)
130
+ vqgan_params = replicate(vqgan.params)
131
+
132
+ # ## CLIP Scoring
133
+ from transformers import CLIPProcessor, FlaxCLIPModel
134
+
135
+ clip = FlaxCLIPModel.from_pretrained("openai/clip-vit-base-patch32")
136
+ print("Initialize FlaxCLIPModel")
137
+ processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
138
+ print("Initialize CLIPProcessor")
139
+
140
+ def hallucinate(prompt, num_images=64):
141
+ prompt = [prompt] * jax.device_count()
142
+ inputs = tokenizer(prompt, return_tensors='jax', padding="max_length", truncation=True, max_length=128).data
143
+ inputs = shard(inputs)
144
+
145
+ all_images = []
146
+ for i in range(num_images // jax.device_count()):
147
+ key = random.randint(0, 1e7)
148
+ rng = jax.random.PRNGKey(key)
149
+ rngs = jax.random.split(rng, jax.local_device_count())
150
+ indices = p_generate(inputs, rngs, bart_params).sequences
151
+ indices = indices[:, :, 1:]
152
+
153
+ images = p_get_images(indices, vqgan_params)
154
+ images = np.squeeze(np.asarray(images), 1)
155
+ for image in images:
156
+ all_images.append(custom_to_pil(image))
157
+ return all_images
158
+
159
+ def clip_top_k(prompt, images, k=8):
160
+ inputs = processor(text=prompt, images=images, return_tensors="np", padding=True)
161
+ outputs = clip(**inputs)
162
+ logits = outputs.logits_per_text
163
+ scores = np.array(logits[0]).argsort()[-k:][::-1]
164
+ return [images[score] for score in scores]
165
+
166
+ def captioned_strip(images, caption):
167
+ increased_h = 0 if caption is None else 48
168
+ w, h = images[0].size[0], images[0].size[1]
169
+ img = Image.new("RGB", (len(images)*w, h + increased_h))
170
+ for i, img_ in enumerate(images):
171
+ img.paste(img_, (i*w, increased_h))
172
+
173
+ if caption is not None:
174
+ draw = ImageDraw.Draw(img)
175
+ font = ImageFont.truetype("/usr/share/fonts/truetype/liberation2/LiberationMono-Bold.ttf", 40)
176
+ draw.text((20, 3), caption, (255,255,255), font=font)
177
+ return img
178
+
179
+ def run_inference(prompt, num_images=64, num_preds=8):
180
+ images = hallucinate(prompt, num_images=num_images)
181
+ images = clip_top_k(prompt, images, k=num_preds)
182
+ predictions_strip = captioned_strip(images, None)
183
+ return predictions_strip
184
+
185
+ gr.Interface(run_inference,
186
+ 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')],
187
+ outputs=gr.outputs.Image(label='Generated image'),
188
+ title='DALLE-mini - HuggingFace Community Week',
189
+ description='This is a demo of the DALLE-mini model trained with Jax/Flax on TPU v3-8s during the HuggingFace Community Week',
190
+ article="<p style='text-align: center'> DALLE-mini by Boris Dayma et al. | <a href='https://github.com/borisdayma/dalle-mini'>GitHub</a></p>",
191
+ layout='vertical',
192
+ theme='huggingface',
193
+ examples=[['an armchair in the shape of an avocado']],
194
+ server_port=8999).launch(share=True)
app/requirements.txt → requirements.txt RENAMED
@@ -9,3 +9,5 @@ flax
9
  jupyter
10
  wandb
11
  ftfy
 
 
 
9
  jupyter
10
  wandb
11
  ftfy
12
+ streamlit
13
+ gradio