anima_closeout v17: the bed called transformer.forward() directly, which is NOT the trainer's path — concat_padding_mask is on so forward()'s padding_mask=None default dies in prepare_embedded_sequence, and it also skips the llm_adapter and the per-block multiband window dispatch. The bed now runs pipe.to_layers() (InitialLayer -> LLMAdapterLayer -> 28x TransformerLayer -> FinalLayer) with the exact input tuple prepare_inputs builds, so it inherits the zero padding mask, the adapter call and the band dispatch instead of re-implementing any of them. The image_spec join from v16 is confirmed working on the live cache: 512 unique specs, 184+189+139 rows joined.

colab/anima_closeout.ipynb

@@ -935,7 +935,7 @@
     "# arm. Reuses the fork's eval cache (latents + text embeds) and the fork's\n",
     "# own transformer + llm_adapter call path (mirrored from\n",
     "# cosmos_predict2.py InitialLayer/LLMAdapterLayer).\n",
-    "import json, os, sys\n",
+    "import json, os, sys, time\n",
     "sys.path.insert(0, \".\")                     # fork root (cwd)\n",
     "import utils.common                         # bind fork utils BEFORE ComfyUI\n",
     "sys.path.append(\"submodules/ComfyUI\")\n",
@@ -1059,7 +1059,22 @@
     "pipe = cosmos_predict2.CosmosPredict2Pipeline(mc)\n",
     "pipe.load_diffusion_model()\n",
     "tf = pipe.transformer.to(DEV).eval()\n",
-    "from models.aleph_relay import attach_aleph_relays, band_weights\n",
+    "from models.aleph_relay import attach_aleph_relays\n",
+    "\n",
+    "# THE fork's own execution path (cosmos_predict2.py:460-470):\n",
+    "#   InitialLayer -> LLMAdapterLayer -> 28x TransformerLayer -> FinalLayer\n",
+    "# Calling transformer.forward() directly is NOT what the trainer does and\n",
+    "# does not work: concat_padding_mask is on, so forward()'s padding_mask=None\n",
+    "# default crashes in prepare_embedded_sequence \u2014 InitialLayer builds a zero\n",
+    "# padding mask itself (:612). The layer path also applies the llm_adapter\n",
+    "# (LLMAdapterLayer) and sets each block's multiband window from the actual\n",
+    "# timesteps (TransformerLayer:663-668), so the bed inherits all of it\n",
+    "# instead of re-implementing it. The layers hold references to tf.blocks,\n",
+    "# so adapters attached per-arm below are picked up without rebuilding.\n",
+    "LAYERS = pipe.to_layers()\n",
+    "print(f\"model layers: {len(LAYERS)} \"\n",
+    "      f\"({type(LAYERS[0]).__name__} .. {type(LAYERS[-1]).__name__})\",\n",
+    "      flush=True)\n",
     "\n",
     "def clear_adapters():\n",
     "    for b in tf.blocks:\n",
@@ -1067,25 +1082,17 @@
     "            if hasattr(b, attr):\n",
     "                delattr(b, attr)\n",
     "\n",
-    "# text embeds through llm_adapter ONCE per row (t-independent) \u2014 the exact\n",
-    "# call LLMAdapterLayer makes (cosmos_predict2.py:643-649)\n",
-    "la = getattr(tf, \"llm_adapter\", None)\n",
-    "embeds = []\n",
-    "with torch.no_grad():\n",
-    "    for r in rows:\n",
-    "        e = r[\"prompt_embeds\"].to(DEV, torch.bfloat16)\n",
-    "        if e.ndim == 2:\n",
-    "            e = e.unsqueeze(0)\n",
-    "        if la is not None:\n",
-    "            am = r[\"attn_mask\"].to(DEV).reshape(1, -1)\n",
-    "            ids = r[\"t5_input_ids\"].to(DEV).reshape(1, -1)\n",
-    "            t5m = r[\"t5_attn_mask\"].to(DEV).reshape(1, -1)\n",
-    "            e = la(source_hidden_states=e, target_input_ids=ids,\n",
-    "                   target_attention_mask=t5m, source_attention_mask=am)\n",
-    "            e = e.clone()\n",
-    "            e[~t5m.bool()] = 0\n",
-    "        embeds.append(e)\n",
-    "print(\"embeds prepared (post-llm_adapter)\", flush=True)\n",
+    "# conditioning exactly as prepare_inputs hands it over (:412): the CACHED\n",
+    "# embeds, un-adapted \u2014 LLMAdapterLayer does the adapter call itself.\n",
+    "conds = []\n",
+    "for r in rows:\n",
+    "    e = r[\"prompt_embeds\"]\n",
+    "    if e.ndim == 2:\n",
+    "        e = e.unsqueeze(0)\n",
+    "    conds.append((e, r[\"attn_mask\"].reshape(1, -1),\n",
+    "                  r[\"t5_input_ids\"].reshape(1, -1),\n",
+    "                  r[\"t5_attn_mask\"].reshape(1, -1)))\n",
+    "print(f\"conditioning prepared for {len(conds)} rows\", flush=True)\n",
     "\n",
     "# fixed noise bank + the fork's exact eval-t transform:\n",
     "# t = sigmoid(Normal.icdf(q))  (logit_normal, scale 1, no shift \u2014\n",
@@ -1106,33 +1113,32 @@
     "results = {\"quantiles\": QUANTILES, \"t_of_q\": T_OF_Q, \"n_rows\": len(rows),\n",
     "           \"arms\": {}}\n",
     "\n",
-    "def eval_arm(name, lesion=None, needs_bands=False):\n",
+    "def eval_arm(name, lesion=None):\n",
+    "    t0 = time.time()\n",
     "    per_q = {}\n",
     "    with torch.no_grad():\n",
     "        for q in QUANTILES:\n",
     "            t = T_OF_Q[q]\n",
-    "            if needs_bands:\n",
-    "                w = band_weights(torch.tensor([t], device=DEV))\n",
-    "                for b in tf.blocks:\n",
-    "                    b.aleph_w_bands = w\n",
     "            losses = []\n",
-    "            for l, nz, e in zip(lats, noises, embeds):\n",
-    "                ld = l.to(DEV, torch.bfloat16)\n",
-    "                nd = nz.to(DEV, torch.bfloat16)\n",
+    "            for l, nz, cd in zip(lats, noises, conds):\n",
+    "                ld = l.to(DEV)                       # float32, as :409\n",
+    "                nd = nz.to(DEV)\n",
     "                x_t = (1 - t) * ld + t * nd\n",
-    "                ts = torch.full((ld.shape[0],), t, device=DEV,\n",
-    "                                dtype=torch.bfloat16)\n",
-    "                pred = tf(x_t, ts, e)\n",
-    "                if isinstance(pred, (list, tuple)):\n",
-    "                    pred = pred[0]\n",
+    "                ts = torch.full((ld.shape[0], 1), t, device=DEV,\n",
+    "                                dtype=torch.float32)   # (B,1), as :456\n",
+    "                out = (x_t, ts, cd[0].to(DEV), cd[1].to(DEV),\n",
+    "                       cd[2].to(DEV), cd[3].to(DEV))\n",
+    "                for layer in LAYERS:\n",
+    "                    out = layer(out)\n",
+    "                pred = out[0] if isinstance(out, (list, tuple)) else out\n",
     "                losses.append(((pred.float().cpu() - (nz - l)) ** 2)\n",
     "                              .mean().item())\n",
     "            per_q[str(q)] = sum(losses) / len(losses)\n",
     "    results[\"arms\"][name] = {\"loss_by_quantile\": per_q,\n",
     "                             \"loss_mean\": sum(per_q.values()) / len(per_q),\n",
     "                             \"lesion\": lesion}\n",
-    "    print(f\"[{name}] mean {results['arms'][name]['loss_mean']:.6f}\",\n",
-    "          flush=True)\n",
+    "    print(f\"[{name}] mean {results['arms'][name]['loss_mean']:.6f} \"\n",
+    "          f\"({(time.time() - t0) / 60:.1f} min)\", flush=True)\n",
     "\n",
     "clear_adapters()\n",
     "eval_arm(\"frozen\")\n",
@@ -1154,15 +1160,14 @@
     "        gates.extend(torch.as_tensor(g).float().flatten().tolist())\n",
     "    for b in tf.blocks:                        # toggle assert per arm\n",
     "        b.aleph_relay.enabled = False\n",
-    "    eval_arm(name + \"_toggled_off\", needs_bands=(mode == \"multiband3\"))\n",
+    "    eval_arm(name + \"_toggled_off\")\n",
     "    assert results[\"arms\"][name + \"_toggled_off\"][\"loss_by_quantile\"] \\\n",
     "        == f_ref, f\"{name}: TOGGLE LAW VIOLATED\"\n",
     "    for b in tf.blocks:\n",
     "        b.aleph_relay.enabled = True\n",
     "        if arm.get(\"lesion\") is not None:\n",
     "            b.aleph_relay.band_enabled[arm[\"lesion\"]] = False\n",
-    "    eval_arm(name, lesion=arm.get(\"lesion\"),\n",
-    "             needs_bands=(mode == \"multiband3\"))\n",
+    "    eval_arm(name, lesion=arm.get(\"lesion\"))\n",
     "    results[\"arms\"][name][\"gates_fp32\"] = {\n",
     "        \"min\": min(gates), \"max\": max(gates),\n",
     "        \"moved\": sum(1 for g in gates if abs(g + 3.0) > 1e-3)}\n",