feat: update wandb inference

dev/inference/samples.txt

@@ -1,6 +1,6 @@
 white snow covered mountain under blue sky during daytime
-aerial view of the beach at night
 aerial view of the beach during daytime
+aerial view of the beach at night
 a beautiful sunset at a beach with a shell on the shore
 a farmhouse surrounded by beautiful flowers
 a photo of a fantasy version of New York City
@@ -29,6 +29,7 @@ an avocado armchair flying into space
 a cute avocado armchair singing karaoke on stage in front of a crowd of strawberry shaped lamps
 an illustration of an avocado in a christmas sweater staring at its reflection in a mirror
 illustration of an avocado armchair
+illustration of an armchair in the shape of an avocado
 illustration of an avocado armchair getting married to a pineapple
 a muscular banana sitting upright on a bench smoking watching a banana on television, high definition photography
 Mohammed Ali and Mike Tyson in a hypothetical match
@@ -48,6 +49,8 @@ a dog eating worthlessness
 an elephant made of carrots
 an elephant on a unicycle during a circus
 photography of a penguin watching television
+a penguin is walking on the Moon, Earth is in the background
+a penguin standing on a tower of books holds onto a rope from a helicopter
 rat wearing a crown
 a portrait of a nightmare creature watching at you
 a white room full of a black substance
@@ -66,6 +69,8 @@ Cartoon of a carrot with big eyes
 logo of a robot wearing glasses and reading a book
 a bottle of coca-cola on a table
 a cactus lifting weights
+illustration of a cactus lifting weigths
+logo of a cactus lifting weights
 a living room with two white armchairs and a painting of the collosseum. The painting is mounted above a modern fireplace.
 a long line of alternating green and red blocks
 a long line of green blocks on a beach at subset

dev/inference/wandb-backend.ipynb

@@ -47,7 +47,10 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "run_ids = ['k76r0v39']\n",
+    "run_ids = ['he9rrc3q']\n",
+    "# poorly shuffled 1nj161cl\n",
+    "# well shuffled he9rrc3q\n",
+    "# non normalized 1fwxpyfh ! requires changing normalize_text\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",
@@ -115,16 +118,6 @@
     "        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,
@@ -286,7 +279,8 @@
     "\n",
     "            # generate images\n",
     "            images = []\n",
-    "            for i in tqdm(range(num_images // jax.device_count()), desc='Generating Images'):\n",
+    "            pbar = tqdm(range(num_images // jax.device_count()), desc='Generating Images', leave=None)\n",
+    "            for i in pbar:\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",
@@ -296,7 +290,7 @@
     "                    images.append(Image.fromarray(np.asarray(img * 255, dtype=np.uint8)))\n",
     "\n",
     "            # get clip scores\n",
-    "            print('Calculating 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",