feat(log_inference_samples): cleanup

tools/inference/log_inference_samples.ipynb

@@ -100,11 +100,12 @@
    "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",
+    "clip16 = FlaxCLIPModel.from_pretrained(\"openai/clip-vit-base-patch16\")\n",
+    "processor16 = CLIPProcessor.from_pretrained(\"openai/clip-vit-base-patch16\")\n",
+    "clip16_params = replicate(clip16.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",
@@ -123,8 +124,8 @@
     "    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",
+    "def p_clip16(inputs, params):\n",
+    "    logits = clip16(params=params, **inputs).logits_per_image\n",
     "    return logits\n",
     "\n",
     "if add_clip_32:\n",
@@ -229,7 +230,7 @@
    "outputs": [],
    "source": [
     "run_id = run_ids[0]\n",
-    "# TODO: turn everything into a class"
+    "# TODO: turn everything into a class or loop over runs"
    ]
   },
   {
@@ -248,10 +249,8 @@
     "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",
+    "    results16, results32 = [], []\n",
+    "    columns = ['Caption'] + [f'Image {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",
@@ -307,34 +306,13 @@
     "                for img in decoded_images:\n",
     "                    images.append(Image.fromarray(np.asarray(img * 255, dtype=np.uint8)))\n",
     "\n",
-    "            # get clip scores\n",
-    "            pbar.set_description('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",
+    "            def add_clip_results(results, processor, p_clip, clip_params):  \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_clip32(clip_inputs, clip32_params)\n",
+    "                logits = p_clip(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",
@@ -342,13 +320,24 @@
     "                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",
+    "                    top_images = [wandb.Image(cur_images[x], caption=f'Score: {scores[x]:.2f}') for x in idx]\n",
+    "                    results.append([sample] + top_images)\n",
+    "                    \n",
+    "            # get clip scores\n",
+    "            pbar.set_description('Calculating CLIP 16 scores')\n",
+    "            add_clip_results(results16, processor16, p_clip16, clip16_params)\n",
+    "                \n",
+    "            # get clip 32 scores\n",
+    "            if add_clip_32:\n",
+    "                pbar.set_description('Calculating CLIP 32 scores')\n",
+    "                add_clip_results(results32, processor32, p_clip32, clip32_params)\n",
+    "\n",
     "            pbar.close()\n",
     "\n",
+    "                \n",
+    "\n",
     "    # log results\n",
-    "    table = wandb.Table(columns=columns, data=results)\n",
+    "    table = wandb.Table(columns=columns, data=results16)\n",
     "    run.log({'Samples': table, 'version': version})\n",
     "    wandb.finish()\n",
     "    \n",
@@ -359,19 +348,6 @@
     "        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": {