File size: 15,411 Bytes
378a628
 
 
 
91d8a29
378a628
 
 
 
 
b8bbe68
353365f
 
 
c85fbb6
b8bbe68
353365f
 
 
378a628
 
 
 
 
 
 
 
 
91d8a29
 
378a628
 
 
91d8a29
 
378a628
91d8a29
 
c85fbb6
91d8a29
ff051c9
 
 
 
91d8a29
 
ff051c9
 
 
91d8a29
 
ff051c9
91d8a29
 
c85fbb6
91d8a29
378a628
 
353365f
 
91d8a29
353365f
 
91d8a29
353365f
 
 
 
38705a9
bf3640d
353365f
bf3640d
 
353365f
 
378a628
 
91d8a29
378a628
 
91d8a29
378a628
 
2d169e3
 
bf3640d
9a553a4
 
 
 
 
 
378a628
 
 
 
91d8a29
bf3640d
353365f
 
 
bf3640d
 
 
 
 
c85fbb6
 
9a553a4
 
 
 
c85fbb6
 
9a553a4
353365f
 
 
 
91d8a29
c85fbb6
 
91d8a29
c85fbb6
 
 
 
 
 
91d8a29
378a628
 
 
 
2d169e3
 
38705a9
2d169e3
353365f
 
 
b8bbe68
 
378a628
 
91d8a29
378a628
 
 
 
ff051c9
bf3640d
ff051c9
 
 
 
bf3640d
ff051c9
378a628
 
 
 
91d8a29
378a628
 
 
 
bf3640d
ff051c9
bf3640d
 
378a628
 
 
 
91d8a29
378a628
 
 
 
 
bf3640d
378a628
9a553a4
ff051c9
378a628
 
 
 
 
 
91d8a29
bf3640d
378a628
 
 
bf3640d
 
 
 
 
 
 
 
 
 
 
 
 
 
378a628
 
 
 
91d8a29
ff051c9
378a628
 
 
ff051c9
 
378a628
 
 
 
91d8a29
ff051c9
353365f
91d8a29
ff051c9
 
 
 
 
 
 
 
 
 
 
9a553a4
 
2d169e3
ff051c9
 
91d8a29
ff051c9
 
 
 
 
 
 
 
 
 
 
 
9a553a4
ff051c9
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2d169e3
ff051c9
 
 
 
 
 
 
 
91d8a29
ff051c9
 
 
 
 
 
 
 
 
 
2d169e3
ff051c9
 
 
 
c85fbb6
ff051c9
 
 
 
 
2d169e3
ff051c9
 
 
2d169e3
9a553a4
 
 
 
 
 
 
 
 
c85fbb6
9a553a4
 
 
 
 
 
 
 
 
 
ff051c9
 
 
 
9a553a4
 
 
 
 
 
 
 
 
 
353365f
 
 
bf3640d
378a628
 
b8bbe68
ff051c9
bf3640d
 
 
b8bbe68
378a628
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ff2a984-b8b2-4a69-89cf-0d16da2393c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tempfile\n",
    "from functools import partial\n",
    "import random\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "from tqdm.notebook import tqdm\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "from flax.training.common_utils import shard, shard_prng_key\n",
    "from flax.jax_utils import replicate\n",
    "import wandb\n",
    "from dalle_mini.model import CustomFlaxBartForConditionalGeneration\n",
    "from vqgan_jax.modeling_flax_vqgan import VQModel\n",
    "from transformers import BartTokenizer, CLIPProcessor, FlaxCLIPModel\n",
    "from dalle_mini.text import TextNormalizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23e00271-941c-4e1b-b6a9-107a1b77324d",
   "metadata": {},
   "outputs": [],
   "source": [
    "run_ids = ['3kaut6e8']\n",
    "ENTITY, PROJECT = 'wandb', 'hf-flax-dalle-mini'\n",
    "VQGAN_REPO, VQGAN_COMMIT_ID = 'dalle-mini/vqgan_imagenet_f16_16384', None\n",
    "normalize_text = False\n",
    "latest_only = True   # log only latest or all versions\n",
    "suffix = ''           # mainly for duplicate inference runs with a deleted version\n",
    "add_clip_32 = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92f4557c-fd7f-4edc-81c2-de0b0a10c270",
   "metadata": {},
   "outputs": [],
   "source": [
    "run_ids = ['k76r0v39']\n",
    "ENTITY, PROJECT = 'dalle-mini', 'dalle-mini'  # used only for training run\n",
    "VQGAN_REPO, VQGAN_COMMIT_ID = 'dalle-mini/vqgan_imagenet_f16_16384', None\n",
    "normalize_text = True\n",
    "latest_only = True    # log only latest or all versions\n",
    "suffix = ''           # mainly for duplicate inference runs with a deleted version\n",
    "add_clip_32 = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93b2e24b-f0e5-4abe-a3ec-0aa834cc3bf3",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 8\n",
    "num_images = 128\n",
    "top_k = 8\n",
    "text_normalizer = TextNormalizer() if normalize_text else None\n",
    "padding_item = 'NONE'\n",
    "seed = random.randint(0, 2**32-1)\n",
    "key = jax.random.PRNGKey(seed)\n",
    "api = wandb.Api()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6a878fa-4bf5-4978-abb5-e235841d765b",
   "metadata": {},
   "outputs": [],
   "source": [
    "vqgan = VQModel.from_pretrained(VQGAN_REPO, revision=VQGAN_COMMIT_ID)\n",
    "clip = FlaxCLIPModel.from_pretrained(\"openai/clip-vit-base-patch16\")\n",
    "processor = CLIPProcessor.from_pretrained(\"openai/clip-vit-base-patch16\")\n",
    "clip_params = replicate(clip.params)\n",
    "vqgan_params = replicate(vqgan.params)\n",
    "\n",
    "if add_clip_32:\n",
    "    clip32 = FlaxCLIPModel.from_pretrained(\"openai/clip-vit-base-patch32\")\n",
    "    processor32 = CLIPProcessor.from_pretrained(\"openai/clip-vit-base-patch32\")\n",
    "    clip32_params = replicate(clip32.params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a500dd07-dbc3-477d-80d4-2b73a3b83ef3",
   "metadata": {},
   "outputs": [],
   "source": [
    "@partial(jax.pmap, axis_name=\"batch\")\n",
    "def p_decode(indices, params):\n",
    "    return vqgan.decode_code(indices, params=params)\n",
    "\n",
    "@partial(jax.pmap, axis_name=\"batch\")\n",
    "def p_clip(inputs, params):\n",
    "    logits = clip(params=params, **inputs).logits_per_image\n",
    "    return logits\n",
    "\n",
    "if add_clip_32:\n",
    "    @partial(jax.pmap, axis_name=\"batch\")\n",
    "    def p_clip32(inputs, params):\n",
    "        logits = clip32(params=params, **inputs).logits_per_image\n",
    "        return logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebf4f7bf-2efa-46cc-b3f4-2d7a54f7b2cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "clip_params['logit_scale']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e57797ab-0b3a-4490-be58-03d8d1c23fe9",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('samples.txt', encoding='utf8') as f:\n",
    "    samples = [l.strip() for l in f.readlines()]\n",
    "    # make list multiple of batch_size by adding elements\n",
    "    samples_to_add = [padding_item] * (-len(samples) % batch_size)\n",
    "    samples.extend(samples_to_add)\n",
    "    # reshape\n",
    "    samples = [samples[i:i+batch_size] for i in range(0, len(samples), batch_size)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3e02d9d-4ee1-49e7-a7bc-4d8b139e9614",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_artifact_versions(run_id, latest_only=False):\n",
    "    try:\n",
    "        if latest_only:\n",
    "            return [api.artifact(type='bart_model', name=f'{ENTITY}/{PROJECT}/model-{run_id}:latest')]\n",
    "        else:\n",
    "            return api.artifact_versions(type_name='bart_model', name=f'{ENTITY}/{PROJECT}/model-{run_id}', per_page=10000)\n",
    "    except:\n",
    "        return []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0d7ed17-7abb-4a31-ab3c-a12b9039a570",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_training_config(run_id):\n",
    "    training_run = api.run(f'{ENTITY}/{PROJECT}/{run_id}')\n",
    "    config = training_run.config\n",
    "    return config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e784a43-626d-4e8d-9e47-a23775b2f35f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# retrieve inference run details\n",
    "def get_last_inference_version(run_id):\n",
    "    try:\n",
    "        inference_run = api.run(f'dalle-mini/dalle-mini/{run_id}-clip16{suffix}')\n",
    "        return inference_run.summary.get('version', None)\n",
    "    except:\n",
    "        return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1cc9993-1bfc-4ec6-a004-c056189c42ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "# compile functions - needed only once per run\n",
    "def pmap_model_function(model):\n",
    "    \n",
    "    @partial(jax.pmap, axis_name=\"batch\")\n",
    "    def _generate(tokenized_prompt, key, params):\n",
    "        return model.generate(\n",
    "            **tokenized_prompt,\n",
    "            do_sample=True,\n",
    "            num_beams=1,\n",
    "            prng_key=key,\n",
    "            params=params\n",
    "        )\n",
    "    \n",
    "    return _generate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23b2444c-67a9-44d7-abd1-187ed83a9431",
   "metadata": {},
   "outputs": [],
   "source": [
    "run_id = run_ids[0]\n",
    "# TODO: turn everything into a class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bba70f33-af8b-4eb3-9973-7be672301a0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "artifact_versions = get_artifact_versions(run_id, latest_only)\n",
    "last_inference_version = get_last_inference_version(run_id)\n",
    "training_config = get_training_config(run_id)\n",
    "run = None\n",
    "p_generate = None\n",
    "model_files = ['config.json', 'flax_model.msgpack', 'merges.txt', 'special_tokens_map.json', 'tokenizer.json', 'tokenizer_config.json', 'vocab.json']\n",
    "for artifact in artifact_versions:\n",
    "    print(f'Processing artifact: {artifact.name}')\n",
    "    version = int(artifact.version[1:])\n",
    "    results = []\n",
    "    if add_clip_32:\n",
    "        results32 = []\n",
    "    columns = ['Caption'] + [f'Image {i+1}' for i in range(top_k)] + [f'Score {i+1}' for i in range(top_k)]\n",
    "    \n",
    "    if latest_only:\n",
    "        assert last_inference_version is None or version > last_inference_version\n",
    "    else:\n",
    "        if last_inference_version is None:\n",
    "            # we should start from v0\n",
    "            assert version == 0\n",
    "        elif version <= last_inference_version:\n",
    "            print(f'v{version} has already been logged (versions logged up to v{last_inference_version}')\n",
    "        else:\n",
    "            # check we are logging the correct version\n",
    "            assert version == last_inference_version + 1\n",
    "\n",
    "    # start/resume corresponding run\n",
    "    if run is None:\n",
    "        run = wandb.init(job_type='inference', entity='dalle-mini', project='dalle-mini', config=training_config, id=f'{run_id}-clip16{suffix}', resume='allow')\n",
    "\n",
    "    # work in temporary directory\n",
    "    with tempfile.TemporaryDirectory() as tmp:\n",
    "\n",
    "        # download model files\n",
    "        artifact = run.use_artifact(artifact)\n",
    "        for f in model_files:\n",
    "            artifact.get_path(f).download(tmp)\n",
    "\n",
    "        # load tokenizer and model\n",
    "        tokenizer = BartTokenizer.from_pretrained(tmp)\n",
    "        model = CustomFlaxBartForConditionalGeneration.from_pretrained(tmp)\n",
    "        model_params = replicate(model.params)\n",
    "\n",
    "        # pmap model function needs to happen only once per model config\n",
    "        if p_generate is None:\n",
    "            p_generate = pmap_model_function(model)\n",
    "\n",
    "        # process one batch of captions\n",
    "        for batch in tqdm(samples):\n",
    "            processed_prompts = [text_normalizer(x) for x in batch] if normalize_text else list(batch)\n",
    "\n",
    "            # repeat the prompts to distribute over each device and tokenize\n",
    "            processed_prompts = processed_prompts * jax.device_count()\n",
    "            tokenized_prompt = tokenizer(processed_prompts, return_tensors='jax', padding='max_length', truncation=True, max_length=128).data\n",
    "            tokenized_prompt = shard(tokenized_prompt)\n",
    "\n",
    "            # generate images\n",
    "            images = []\n",
    "            for i in tqdm(range(num_images // jax.device_count()), desc='Generating Images'):\n",
    "                key, subkey = jax.random.split(key)\n",
    "                encoded_images = p_generate(tokenized_prompt, shard_prng_key(subkey), model_params)\n",
    "                encoded_images = encoded_images.sequences[..., 1:]\n",
    "                decoded_images = p_decode(encoded_images, vqgan_params)\n",
    "                decoded_images = decoded_images.clip(0., 1.).reshape((-1, 256, 256, 3))\n",
    "                for img in decoded_images:\n",
    "                    images.append(Image.fromarray(np.asarray(img * 255, dtype=np.uint8)))\n",
    "\n",
    "            # get clip scores\n",
    "            print('Calculating CLIP scores')\n",
    "            clip_inputs = processor(text=batch, images=images, return_tensors='np', padding='max_length', max_length=77, truncation=True).data\n",
    "            # each shard will have one prompt, images need to be reorganized to be associated to the correct shard\n",
    "            images_per_prompt_indices = np.asarray(range(0, len(images), batch_size))\n",
    "            clip_inputs['pixel_values'] = jnp.concatenate(list(clip_inputs['pixel_values'][images_per_prompt_indices + i] for i in range(batch_size)))\n",
    "            clip_inputs = shard(clip_inputs)\n",
    "            logits = p_clip(clip_inputs, clip_params)\n",
    "            logits = logits.reshape(-1, num_images)\n",
    "            top_scores = logits.argsort()[:, -top_k:][..., ::-1]\n",
    "            logits = jax.device_get(logits)\n",
    "            # add to results table\n",
    "            for i, (idx, scores, sample) in enumerate(zip(top_scores, logits, batch)):\n",
    "                if sample == padding_item: continue\n",
    "                cur_images = [images[x] for x in images_per_prompt_indices + i]\n",
    "                top_images = [wandb.Image(cur_images[x]) for x in idx]\n",
    "                top_scores = [scores[x] for x in idx]\n",
    "                results.append([sample] + top_images + top_scores)\n",
    "                \n",
    "            # get clip 32 scores - TODO: this should be refactored as it is same code as above\n",
    "            if add_clip_32:\n",
    "                print('Calculating CLIP 32 scores')\n",
    "                clip_inputs = processor32(text=batch, images=images, return_tensors='np', padding='max_length', max_length=77, truncation=True).data\n",
    "                # each shard will have one prompt, images need to be reorganized to be associated to the correct shard\n",
    "                images_per_prompt_indices = np.asarray(range(0, len(images), batch_size))\n",
    "                clip_inputs['pixel_values'] = jnp.concatenate(list(clip_inputs['pixel_values'][images_per_prompt_indices + i] for i in range(batch_size)))\n",
    "                clip_inputs = shard(clip_inputs)\n",
    "                logits = p_clip32(clip_inputs, clip32_params)\n",
    "                logits = logits.reshape(-1, num_images)\n",
    "                top_scores = logits.argsort()[:, -top_k:][..., ::-1]\n",
    "                logits = jax.device_get(logits)\n",
    "                # add to results table\n",
    "                for i, (idx, scores, sample) in enumerate(zip(top_scores, logits, batch)):\n",
    "                    if sample == padding_item: continue\n",
    "                    cur_images = [images[x] for x in images_per_prompt_indices + i]\n",
    "                    top_images = [wandb.Image(cur_images[x]) for x in idx]\n",
    "                    top_scores = [scores[x] for x in idx]\n",
    "                    results32.append([sample] + top_images + top_scores)\n",
    "\n",
    "    # log results\n",
    "    table = wandb.Table(columns=columns, data=results)\n",
    "    run.log({'Samples': table, 'version': version})\n",
    "    wandb.finish()\n",
    "    \n",
    "    if add_clip_32:        \n",
    "        run = wandb.init(job_type='inference', entity='dalle-mini', project='dalle-mini', config=training_config, id=f'{run_id}-clip32{suffix}', resume='allow')\n",
    "        table = wandb.Table(columns=columns, data=results32)\n",
    "        run.log({'Samples': table, 'version': version})\n",
    "        wandb.finish()\n",
    "        run = None  # ensure we don't log on this run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e4c7d0c-2848-4f88-b967-82fd571534f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: not implemented\n",
    "def log_runs(runs):\n",
    "    for run in tqdm(runs):\n",
    "        log_run(run)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}