anima_closeout v16: stop deriving cache directory names. v15 built the AR dir as ar_frames_{ar}x{f} but bucket_suffix (dataset.py:60-70) joins 2-tuples with an UNDERSCORE (ar_frames_0.630_1) and 4-tuples with 'x' (cache_0.630x400x640x1). The bed now WALKS the cache root, discovering every 'latents' and 'text_embeddings_*' dir (uncond_* excluded by prefix), and joins on image_spec, which is globally unique — so no directory-naming convention can break it again. Verified on a synthetic cache built by the fork's own ParquetCache with names from the fork's own bucket_suffix: 3 AR/size bucket pairs, shuffled latents, uncond + iteration_order decoys -> 24/24 rows correctly paired, each image exactly once.

colab/anima_closeout.ipynb

@@ -966,10 +966,26 @@
     "# (dataset.py:312), so we join on image_spec and iterate in sorted spec\n",
     "# order, which also makes the bed reproducible run-to-run.\n",
     "cache_root = CFG[\"eval_cache_dir\"]\n",
-    "size_dirs = sorted(d for d in os.listdir(cache_root)\n",
-    "                   if d.startswith(\"cache_\")\n",
-    "                   and os.path.isdir(os.path.join(cache_root, d)))\n",
-    "assert size_dirs, f\"no cache_* size-bucket dirs under {cache_root}\"\n",
+    "# DISCOVER both kinds of cache by walking the tree. Do NOT derive one\n",
+    "# directory name from the other: size buckets are\n",
+    "# cache_{ar:.3f}x{w}x{h}x{f} but AR buckets are ar_frames_{ar:.3f}_{f} \u2014\n",
+    "# 'x' vs '_' (bucket_suffix, dataset.py:60-70). image_spec is globally\n",
+    "# unique per image, so one map over every text-embedding cache joins any\n",
+    "# latents cache unambiguously, whatever the directories are called.\n",
+    "lat_dirs, te_dirs = [], []\n",
+    "for _dp, _dns, _ in os.walk(cache_root):\n",
+    "    for _d in _dns:\n",
+    "        if _d == \"latents\":\n",
+    "            lat_dirs.append(os.path.join(_dp, _d))\n",
+    "        elif _d.startswith(\"text_embeddings_\"):    # excludes uncond_*\n",
+    "            te_dirs.append(os.path.join(_dp, _d))\n",
+    "lat_dirs.sort(); te_dirs.sort()\n",
+    "print(\"latents caches:   \",\n",
+    "      [os.path.relpath(d, cache_root) for d in lat_dirs], flush=True)\n",
+    "print(\"text-embed caches:\",\n",
+    "      [os.path.relpath(d, cache_root) for d in te_dirs], flush=True)\n",
+    "assert lat_dirs, f\"no latents cache under {cache_root}\"\n",
+    "assert te_dirs, f\"no text_embeddings_* cache under {cache_root}\"\n",
     "\n",
     "def read_cache(d):\n",
     "    shards = sorted(p for p in os.listdir(d)\n",
@@ -1000,40 +1016,38 @@
     "\n",
     "LAT = [\"latents\"]\n",
     "TE = [\"prompt_embeds\", \"attn_mask\", \"t5_input_ids\", \"t5_attn_mask\"]\n",
+    "\n",
+    "te_map = {}\n",
+    "for d in te_dirs:\n",
+    "    tbs = read_cache(d)\n",
+    "    for c in TE:\n",
+    "        assert c + \"__shape\" in tbs[0].schema.names, \\\n",
+    "            (d, tbs[0].schema.names[:8])\n",
+    "    for s, loc in by_spec(tbs, \"text embedding\").items():\n",
+    "        te_map.setdefault(s, loc)\n",
+    "print(f\"text embeddings: {len(te_map)} unique image_spec\", flush=True)\n",
+    "\n",
     "rows = []\n",
-    "for sd in size_dirs:\n",
-    "    parts = sd[len(\"cache_\"):].split(\"x\")        # {ar}x{w}x{h}x{frames}\n",
-    "    ar_dir = os.path.join(cache_root, f\"ar_frames_{parts[0]}x{parts[-1]}\")\n",
-    "    assert os.path.isdir(ar_dir), (\n",
-    "        f\"{sd}: expected AR-bucket dir {ar_dir} \u2014 text embeddings are \"\n",
-    "        \"cached per AR bucket, not per size bucket\")\n",
-    "    te_names = sorted(d for d in os.listdir(ar_dir)\n",
-    "                      if d.startswith(\"text_embeddings_\"))\n",
-    "    assert len(te_names) == 1, \\\n",
-    "        f\"{ar_dir}: expected one text_embeddings_* cache, found {te_names}\"\n",
-    "    lat = by_spec(read_cache(os.path.join(cache_root, sd, \"latents\")),\n",
-    "                  \"latents\")\n",
-    "    te = by_spec(read_cache(os.path.join(ar_dir, te_names[0])),\n",
-    "                 \"text embedding\")\n",
+    "for d in lat_dirs:\n",
+    "    tbs = read_cache(d)\n",
     "    for c in LAT:\n",
-    "        _s = next(iter(lat.values()))[0].schema.names\n",
-    "        assert c + \"__shape\" in _s, (sd, _s[:8])\n",
-    "    for c in TE:\n",
-    "        _s = next(iter(te.values()))[0].schema.names\n",
-    "        assert c + \"__shape\" in _s, (ar_dir, _s[:8])\n",
-    "    _missing = [s for s in lat if s not in te]\n",
+    "        assert c + \"__shape\" in tbs[0].schema.names, \\\n",
+    "            (d, tbs[0].schema.names[:8])\n",
+    "    lat = by_spec(tbs, \"latents\")\n",
+    "    _missing = [s for s in lat if s not in te_map]\n",
     "    assert not _missing, (\n",
-    "        f\"{sd}: {len(_missing)} latents have no text embedding for their \"\n",
-    "        f\"image_spec (e.g. {_missing[:1]}) \u2014 caches are out of sync\")\n",
+    "        f\"{os.path.relpath(d, cache_root)}: {len(_missing)} latents have \"\n",
+    "        f\"no text embedding for their image_spec (e.g. {_missing[:1]}) \u2014 \"\n",
+    "        \"caches are out of sync\")\n",
     "    for s in sorted(lat):\n",
     "        lt, lj = lat[s]\n",
-    "        tt, tj = te[s]\n",
+    "        tt, tj = te_map[s]\n",
     "        r = {c: decol(lt, c, lj) for c in LAT}\n",
     "        r.update({c: decol(tt, c, tj) for c in TE})\n",
     "        rows.append(r)\n",
-    "    print(f\"  {sd}: {len(lat)} rows joined on image_spec \"\n",
-    "          f\"(AR pool {len(te)})\", flush=True)\n",
-    "print(f\"eval cache rows: {len(rows)} from {len(size_dirs)} size bucket(s)\",\n",
+    "    print(f\"  {os.path.relpath(d, cache_root)}: {len(lat)} rows joined \"\n",
+    "          \"on image_spec\", flush=True)\n",
+    "print(f\"eval cache rows: {len(rows)} from {len(lat_dirs)} latents cache(s)\",\n",
     "      flush=True)\n",
     "rows = rows[:int(os.environ.get(\"VB_MAX_ROWS\", \"512\"))]\n",
     "assert rows\n",