fix: correct clip params

tools/inference/log_inference_samples.ipynb

@@ -24,25 +24,6 @@
     "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",
-    "# Alamy - 3kaut6e8\n",
-    "# YFCC - to do\n",
-    "# HF spaces - 4oh3u7ca\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 = False"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -50,13 +31,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "run_ids = ['2u5lk3uw']\n",
-    "# poorly shuffled 1nj161cl\n",
-    "# well shuffled he9rrc3q\n",
-    "# non normalized 1fwxpyfh ! requires changing normalize_text\n",
+    "run_ids = ['63otg87g']\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",
+    "VQGAN_REPO, VQGAN_COMMIT_ID = 'dalle-mini/vqgan_imagenet_f16_16384', 'e93a26e7707683d349bf5d5c41c5b0ef69b677a9'\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"
@@ -85,7 +62,7 @@
     "batch_size = 8\n",
     "num_images = 128\n",
     "top_k = 8\n",
-    "text_normalizer = TextNormalizer() if normalize_text else None\n",
+    "text_normalizer = TextNormalizer()\n",
     "padding_item = 'NONE'\n",
     "seed = random.randint(0, 2**32-1)\n",
     "key = jax.random.PRNGKey(seed)\n",
@@ -230,7 +207,7 @@
    "outputs": [],
    "source": [
     "run_id = run_ids[0]\n",
-    "# TODO: turn everything into a class or loop over runs"
+    "# TODO: loop over runs"
    ]
   },
   {
@@ -287,7 +264,7 @@
     "\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",
+    "            processed_prompts = [text_normalizer(x) for x in batch] if model.config.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",
@@ -296,7 +273,7 @@
     "\n",
     "            # generate images\n",
     "            images = []\n",
-    "            pbar = tqdm(range(num_images // jax.device_count()), desc='Generating Images', leave=None)\n",
+    "            pbar = tqdm(range(num_images // jax.device_count()), desc='Generating Images', leave=True)\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",
@@ -312,7 +289,7 @@
     "                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, clip32_params)\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",
@@ -348,6 +325,14 @@
     "        wandb.finish()\n",
     "        run = None  # ensure we don't log on this run"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "415d3f54-7226-43de-9eea-4283a948dc93",
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {