feat: cleanup

dev/inference/wandb-backend.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "id": "4ff2a984-b8b2-4a69-89cf-0d16da2393c8",
    "metadata": {},
    "outputs": [],
@@ -26,51 +26,42 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
-   "id": "92f4557c-fd7f-4edc-81c2-de0b0a10c270",
+   "execution_count": null,
+   "id": "23e00271-941c-4e1b-b6a9-107a1b77324d",
    "metadata": {},
    "outputs": [],
    "source": [
-    "run_ids = ['rjf3rycy']\n",
-    "ENTITY, PROJECT = 'dalle-mini', 'dalle-mini'  # used only for training run\n",
+    "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 = True\n",
-    "latest_only = False   # log only latest or all versions\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"
+    "add_clip_32 = True"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
-   "id": "23e00271-941c-4e1b-b6a9-107a1b77324d",
+   "execution_count": null,
+   "id": "92f4557c-fd7f-4edc-81c2-de0b0a10c270",
    "metadata": {},
    "outputs": [],
    "source": [
-    "run_ids = ['3kaut6e8']\n",
-    "ENTITY, PROJECT = 'wandb', 'hf-flax-dalle-mini'\n",
+    "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 = False\n",
-    "latest_only = True   # log only latest or all versions\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 = True"
+    "add_clip_32 = False"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "id": "93b2e24b-f0e5-4abe-a3ec-0aa834cc3bf3",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "INFO:absl:Unable to initialize backend 'tpu_driver': NOT_FOUND: Unable to find driver in registry given worker: \n",
-      "INFO:absl:Unable to initialize backend 'gpu': NOT_FOUND: Could not find registered platform with name: \"cuda\". Available platform names are: TPU Interpreter Host\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "batch_size = 8\n",
     "num_images = 128\n",
@@ -84,18 +75,10 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "id": "c6a878fa-4bf5-4978-abb5-e235841d765b",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Working with z of shape (1, 256, 16, 16) = 65536 dimensions.\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "vqgan = VQModel.from_pretrained(VQGAN_REPO, revision=VQGAN_COMMIT_ID)\n",
     "clip = FlaxCLIPModel.from_pretrained(\"openai/clip-vit-base-patch16\")\n",
@@ -111,7 +94,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "id": "a500dd07-dbc3-477d-80d4-2b73a3b83ef3",
    "metadata": {},
    "outputs": [],
@@ -134,29 +117,17 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": null,
    "id": "ebf4f7bf-2efa-46cc-b3f4-2d7a54f7b2cb",
    "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "ShardedDeviceArray([4.6051702, 4.6051702, 4.6051702, 4.6051702, 4.6051702,\n",
-       "                    4.6051702, 4.6051702, 4.6051702], dtype=float32)"
-      ]
-     },
-     "execution_count": 7,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
+   "outputs": [],
    "source": [
     "clip_params['logit_scale']"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": null,
    "id": "e57797ab-0b3a-4490-be58-03d8d1c23fe9",
    "metadata": {},
    "outputs": [],
@@ -172,7 +143,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": null,
    "id": "f3e02d9d-4ee1-49e7-a7bc-4d8b139e9614",
    "metadata": {},
    "outputs": [],
@@ -189,7 +160,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": null,
    "id": "f0d7ed17-7abb-4a31-ab3c-a12b9039a570",
    "metadata": {},
    "outputs": [],
@@ -202,7 +173,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": null,
    "id": "7e784a43-626d-4e8d-9e47-a23775b2f35f",
    "metadata": {},
    "outputs": [],
@@ -218,7 +189,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": null,
    "id": "d1cc9993-1bfc-4ec6-a004-c056189c42ac",
    "metadata": {},
    "outputs": [],
@@ -241,7 +212,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": null,
    "id": "23b2444c-67a9-44d7-abd1-187ed83a9431",
    "metadata": {},
    "outputs": [],
@@ -252,19 +223,10 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": null,
    "id": "bba70f33-af8b-4eb3-9973-7be672301a0b",
    "metadata": {},
-   "outputs": [
-    {
-     "ename": "SyntaxError",
-     "evalue": "EOL while scanning string literal (1745443972.py, line 60)",
-     "output_type": "error",
-     "traceback": [
-      "\u001b[0;36m  File \u001b[0;32m\"/tmp/ipykernel_402605/1745443972.py\"\u001b[0;36m, line \u001b[0;32m60\u001b[0m\n\u001b[0;31m    for i in tqdm(range(num_images // jax.device_count()), desc='Generating Images):\u001b[0m\n\u001b[0m                                                                                    ^\u001b[0m\n\u001b[0;31mSyntaxError\u001b[0m\u001b[0;31m:\u001b[0m EOL while scanning string literal\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "artifact_versions = get_artifact_versions(run_id, latest_only)\n",
     "last_inference_version = get_last_inference_version(run_id)\n",
@@ -281,8 +243,7 @@
     "    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",
-    "        pass\n",
-    "        #assert last_inference_version is None or version > last_inference_version\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",
@@ -325,7 +286,7 @@
     "\n",
     "            # generate images\n",
     "            images = []\n",
-    "            for i in tqdm(range(num_images // jax.device_count()), desc='Generating 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",
@@ -386,36 +347,6 @@
     "        run = None  # ensure we don't log on this run"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "fdcd09d6-079c-461a-a81a-d9e650d3b099",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "p_clip32"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "7d86ceee-c9ac-4860-abad-410cadd16c3c",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "clip_inputs['attention_mask'].shape, clip_inputs['pixel_values'].shape"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "fbba4858-da2d-4dd5-97b7-ce3ab4746f96",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "clip_inputs['input_ids'].shape"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -428,14 +359,6 @@
     "    for run in tqdm(runs):\n",
     "        log_run(run)"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "a7a5fdf5-3c6e-421b-96a8-5115f730328c",
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {