Merge branch 'feat/batched-te' into integration

# Conflicts:
#	web/conf_builder/conf-builder-main.js

ProjectStructure.md

@@ -300,6 +300,7 @@ All routes registered via `PromptServer.instance.routes`.
 * **`batch_encoding.py`** (378 lines)
   * `encode_prompt_with_combinators(clip, text, clip_skip)` — supports AND, CAT, AVG, BREAK combinators
   * `batch_encode_prompts()` — bulk prompt encoding with cache integration
+  * `encode_prompts_batched()` — `te_batch_size` > 1: several eligible prompts per text-encoder forward (right-padded, per-prompt replay of the unmodified `_encode_single`; allowlist `TE_BATCH_FAMILIES`, self-check, OOM halving, `DEFAULT_TE_BATCH_SIZE`)
   * **Combinators:** `AND` (multi-cond with optional weights like `:1.5`), `CAT` (tensor concatenation), `AVG(weight)` (weighted average), `BREAK` (77-token segment break)
 
 * **`model_loader.py`** (805 lines)

README.md

@@ -171,6 +171,7 @@ If you'd rather build the graph yourself:
 * **LoRA Compatibility Detection:** Automatically detects and skips incompatible LoRAs with detailed error reporting, preventing log spam.
 * **Graceful Interruption:** Cancel button now stops ALL remaining jobs (not just current generation) and saves all completed work including pending remote VAE decodes.
 * **VAE Batching:** Includes a `vae_batch_size` input to batch decode images, significantly speeding up large grid runs.
+* **Text-Encoder Batching:** The `te_batch_size` Run Setting encodes several prompts per text-encoder pass (MOSS-Image, Krea 2, Qwen-Image), cutting pre-encode time on prompt-heavy grids. Off by default.
 * **Live Dashboard Updates:** Configure `flush_batch_every` to update the dashboard incrementally (e.g., every 4 images) instead of waiting for the entire batch to complete.
 * **Remote VAE Support:** Offload VAE decoding to remote servers (HuggingFace endpoints or local) for 20-30% faster generation and lower VRAM usage.
 * **VAE Selection & Iteration:** Test multiple VAEs per config — add individual VAEs from a searchable dropdown or expand entire VAE folders. Each VAE becomes a dimension in the Cartesian grid.
@@ -370,6 +371,11 @@ This suite consists of three nodes found under the `sampling/testing` category:
   - `12-24`: For 24GB+ VRAM  
   - `-1`: Decode all images at once (fastest, but risky)
   
+* **`te_batch_size`** (Run Setting "Text Encoder Batch Size"): How many prompts to encode per text-encoder forward. `1` (default) = off.
+  - Only MOSS-Image, Krea 2 and Qwen-Image encoders batch; every other encoder ignores it.
+  - Weighted prompts and `AND` / `CAT` / `AVG()` / `BREAK` prompts are still encoded one at a time.
+  - The first batch is self-checked against a single-prompt encode; on any mismatch the run falls back to one-by-one automatically. A CUDA OOM halves the batch and retries.
+
 * **`flush_batch_every`**: Update dashboard every X images (0 = use VAE batch size).
   - `0`: Wait until VAE batch is full
   - `4`: Update dashboard every 4 images (recommended for live monitoring)
@@ -757,6 +763,7 @@ Set **Model Type = Diffusion Model**, pick the diffusion file (e.g. `moss_image_
 * The CLIP Type list is read from your installed ComfyUI's `CLIPLoader`, so new model families appear without a plugin update. If you leave it at `stable_diffusion` for a diffusion model, the generator picks the CLIP type that matches the detected model family and logs it (MOSS-Image needs `moss_image` to get its chat template; with `stable_diffusion` ComfyUI builds a generic Qwen3.5 encoder).
 * The empty latent follows the model's latent format: 16 px per latent cell for the Flux.2 family, 8 for SD / SDXL / Flux.1. Use sizes that are multiples of 16; a 832x1216 job renders at 832x1216. Earlier releases always used 8, which made 16x models render at double size and run ~10x slower than the stock workflows.
 * Pre-encoding loads the text encoder once per pass instead of once per prompt (ComfyUI's dynamic VRAM loader re-prepares a model on every encode, which is the repeated `Model ... prepared for dynamic VRAM loading` log line). What remains is the encoder itself: the Qwen3.5-9B encoder costs roughly 0.8 s per prompt on a 16 GB card, the same as the stock workflow, so a 60-prompt grid spends about a minute encoding before the first image. SDXL's CLIPs take about 0.15 s per prompt, Krea 2 and Qwen-Image about 0.3 s. The pre-encode summary line reports the measured ms/prompt.
+* Text-encoder batching (`te_batch_size`, default 1 = off) puts several prompts through one encoder forward, which amortises the encoder's fixed per-forward cost. Prompts that can't be batched (weighted, `AND`/`CAT`/`AVG()`/`BREAK`) and encoders other than MOSS-Image / Krea 2 / Qwen-Image keep the per-prompt path. The summary line shows how many prompts were batched.
 
 ### Model-Specific Prompts
 
@@ -1109,6 +1116,9 @@ Create a consistent quality stack using the Config Builder: add 3 LoRAs manually
 1. Use identical LoRA/prompt settings across models for fair comparison.
 2. Use `add_random_seeds_to_gens: 2-3` to evaluate model consistency.
 
+### For MOSS-Image / Krea 2 / Qwen-Image grids with many prompts
+1. Try `te_batch_size: 8` (Run Settings > Text Encoder Batch Size) to cut pre-encode time; lower it if you see an OOM retry message.
+
 ### For Memory-Constrained Systems
 1. Lower `vae_batch_size` to 1-2.
 

batch_encoding.py

@@ -5,8 +5,10 @@ Supports prompt combinators: AND, CAT, AVG(), BREAK (native).
 """
 
 import contextlib
+import numbers
 import re
 import time
+import weakref
 import torch
 import gc
 import comfy.model_management
@@ -92,11 +94,377 @@ def encode_pair_cached(clip_model, cache, positive, negative, clip_skip=0):
     return cache["positive"][positive], cache["negative"][negative]
 
 
+# ============================================================
+# === SEVERAL PROMPTS THROUGH ONE TEXT-ENCODER FORWARD ===
+# ============================================================
+# MOSS-Image's Qwen3.5-9B encoder spends ~0.55 s per forward in fixed overhead
+# (24 Gated DeltaNet layers, each a 63-iteration Python loop that doesn't depend on
+# the batch size). Encoding B prompts in ONE forward amortises it.
+#
+# Design "batched forward, per-prompt replay": the INNER model's ``encode`` (the call
+# ClipTokenWeightEncoder.encode_token_weights makes) is shadowed by a replay object.
+# The first call runs the real encode on the right-padded batch; every call returns
+# that prompt's slice, shaped exactly like a one-row forward. The existing
+# ``_encode_single`` then runs per prompt, so template stripping, mask popping,
+# clip_skip, device moves and dict building are the unmodified single-prompt code.
+
+# Default for the ``te_batch_size`` Run Setting. 1 = off (today's per-prompt path).
+DEFAULT_TE_BATCH_SIZE = 1
+
+# clip_name -> family label. An allowlist: a family is added only after the live check passes.
+TE_BATCH_FAMILIES = {"qwen35_9b": "MOSS-Image", "qwen3vl_4b": "Krea 2", "qwen25_7b": "Qwen-Image"}
+TE_BATCH_MAX_TOKENS = 4096        # cap on rows * longest-row tokens per forward (memory guard)
+TE_BATCH_SELFCHECK_RELL2 = 1e-3   # batched vs single cond tolerance
+
+_TE_BATCH_DISABLED = weakref.WeakSet()
+_TE_BATCH_DISABLED_IDS = set()    # fallback for objects that cannot be weak-referenced
+
+
+class _TEBatchUnsupported(Exception):
+    """The encoder's output didn't have the shape the batched path relies on."""
+
+
+def _te_disabled(clip_model):
+    try:
+        return clip_model in _TE_BATCH_DISABLED
+    except TypeError:
+        return id(clip_model) in _TE_BATCH_DISABLED_IDS
+
+
+def _te_disable(clip_model):
+    try:
+        _TE_BATCH_DISABLED.add(clip_model)
+    except TypeError:
+        _TE_BATCH_DISABLED_IDS.add(id(clip_model))
+
+
+def _te_batch_target(clip_model):
+    """-> ((inner, clip_name, pad_id), None) when this CLIP can batch, else (None, reason)."""
+    csm = getattr(clip_model, "cond_stage_model", None)
+    clip_attr = getattr(csm, "clip", None)
+    clip_name = getattr(csm, "clip_name", None)
+    if not isinstance(clip_attr, str) or not isinstance(clip_name, str):
+        return None, "multi-encoder or unsupported CLIP wrapper"
+    if clip_name not in TE_BATCH_FAMILIES:
+        return None, f"{clip_name} is not a verified batchable encoder"
+    inner = getattr(csm, clip_attr, None)
+    if inner is None or not callable(getattr(inner, "encode", None)) or not callable(getattr(inner, "process_tokens", None)):
+        return None, "encoder has no encode/process_tokens"
+    if getattr(inner, "enable_attention_masks", None) is not True or getattr(inner, "return_attention_masks", None) is not True:
+        return None, "encoder does not use attention masks"
+    special = getattr(inner, "special_tokens", None) or {}
+    if special.get("end") is not None:
+        return None, "encoder has an end token"
+    pad = special.get("pad")
+    if not isinstance(pad, numbers.Integral) or isinstance(pad, bool):
+        return None, "encoder has no pad token"
+    if getattr(clip_model, "use_clip_schedule", False):
+        return None, "scheduled CLIP"
+    if "encode" in vars(inner):
+        return None, "encoder already wrapped"
+    return (inner, clip_name, int(pad)), None
+
+
+def _batchable_ids(clip_model, text, clip_name, pad_id):
+    """-> ((tokens_dict, ids_list), None) or (None, reason). Tokenises ONCE; the dict is reused for the encode."""
+    if _has_combinators(text) or "BREAK" in text:
+        return None, "combinator"
+    tokens = clip_model.tokenize(text)
+    if list(tokens) != [clip_name] or len(tokens[clip_name]) != 1:
+        return None, "multi-section"
+    ids = []
+    for item in tokens[clip_name][0]:
+        if not isinstance(item, (tuple, list)) or len(item) < 2:
+            return None, "unexpected token entry"
+        t, w = item[0], item[1]
+        if not isinstance(t, numbers.Integral) or isinstance(t, bool):
+            return None, "embedding or image token"
+        if w != 1.0:
+            return None, "weighted"
+        ids.append(int(t))
+    if not ids:
+        return None, "empty"
+    if pad_id in ids:
+        return None, "pad id inside prompt"
+    return (tokens, ids), None
+
+
+def _own(t):
+    return t.clone() if hasattr(t, "clone") else t.copy()
+
+
+class _ReplayEncode:
+    """Stands in for ``inner.encode``: one real forward on the padded batch, then per-prompt slices."""
+
+    def __init__(self, real_encode, id_lists, pad_id):
+        self.real = real_encode
+        self.id_lists = id_lists
+        self.pad = pad_id
+        self.index = {tuple(ids): i for i, ids in enumerate(id_lists)}
+        self.T = max(len(ids) for ids in id_lists)
+        self.out = None
+        self.forwards = 0
+        self.misses = 0
+
+    def __call__(self, to_encode):
+        key = tuple(to_encode[0]) if len(to_encode) == 1 else None
+        if key not in self.index:
+            self.misses += 1  # safety net: unknown input runs the normal single path
+            return self.real(to_encode)
+        if self.out is None:
+            padded = [list(ids) + [self.pad] * (self.T - len(ids)) for ids in self.id_lists]
+            self.out = self.real(padded)
+            self.forwards += 1
+        return self._slice(self.index[key])
+
+    def _slice(self, i):
+        out = self.out
+        z, pooled = out[0], out[1]
+        B = len(self.id_lists)
+        L = len(self.id_lists[i])
+        if len(z.shape) < 2 or z.shape[0] != B or z.shape[-2] != self.T:
+            raise _TEBatchUnsupported(f"unexpected output shape {tuple(z.shape)}")
+        if pooled is not None:
+            raise _TEBatchUnsupported("pooled output")  # EOS-position pooled would move with padding
+        z_i = _own(z[i:i + 1, ..., :L, :])
+        if len(out) == 2:
+            return (z_i, None)
+        extra_i = {}
+        for k, v in out[2].items():
+            if k == "attention_mask":
+                if int(v[i].sum()) != L:
+                    raise _TEBatchUnsupported("attention mask length mismatch")
+                extra_i[k] = _own(v[i:i + 1, :L])
+            elif hasattr(v, "shape") and len(v.shape) and v.shape[0] == B:
+                extra_i[k] = _own(v[i:i + 1])
+            else:
+                raise _TEBatchUnsupported(f"extra output {k}")
+        return (z_i, None, extra_i)
+
+
+def _tensor_stats(a, b):
+    """-> (||a-b||, ||b||, max|a-b|) for torch tensors or numpy arrays."""
+    if hasattr(a, "detach"):
+        a32, b32 = a.detach().float(), b.detach().float()
+        d = a32 - b32
+        return float(d.norm()), float(b32.norm()), (float(d.abs().max()) if d.numel() else 0.0)
+    import numpy as np
+    a64, b64 = np.asarray(a, dtype=np.float64), np.asarray(b, dtype=np.float64)
+    d = a64 - b64
+    return float(np.linalg.norm(d)), float(np.linalg.norm(b64)), (float(np.abs(d).max()) if d.size else 0.0)
+
+
+def _compare_conds(batched, single):
+    """-> (ok, reason, worst_relL2, worst_maxabs, shapes_equal). Compares keys, shapes, dtypes and values."""
+    if len(batched) != len(single):
+        return False, "conditioning length differs", 0.0, 0.0, False
+    worst_rel = worst_abs = 0.0
+    for (ca, da), (cb, db) in zip(batched, single):
+        pairs = [("cond", ca, cb)]
+        if sorted(da) != sorted(db):
+            return False, f"keys differ {sorted(da)} vs {sorted(db)}", worst_rel, worst_abs, False
+        for k in da:
+            pairs.append((k, da[k], db[k]))
+        for name, a, b in pairs:
+            if hasattr(a, "shape") or hasattr(b, "shape"):
+                if not (hasattr(a, "shape") and hasattr(b, "shape")) or tuple(a.shape) != tuple(b.shape):
+                    return False, f"{name} shape differs", worst_rel, worst_abs, False
+                if str(a.dtype) != str(b.dtype):
+                    return False, f"{name} dtype differs", worst_rel, worst_abs, False
+                dn, bn, mx = _tensor_stats(a, b)
+                worst_rel = max(worst_rel, dn / max(bn, 1e-12))
+                worst_abs = max(worst_abs, mx)
+            elif a != b:
+                return False, f"{name} value differs", worst_rel, worst_abs, True
+    ok = worst_rel <= TE_BATCH_SELFCHECK_RELL2
+    return ok, ("" if ok else f"relL2 {worst_rel:.2e} > {TE_BATCH_SELFCHECK_RELL2:g}"), worst_rel, worst_abs, True
+
+
+def _is_oom(e):
+    oom = getattr(comfy.model_management, "OOM_EXCEPTION", None)
+    if isinstance(oom, type) and isinstance(e, oom):
+        return True
+    return "out of memory" in str(e).lower()
+
+
+def encode_prompts_batched(clip_model, prompts, batch_size, clip_skip=0, verify=False):
+    """
+    Encode as many ``prompts`` as possible several-at-a-time in one text-encoder forward.
+
+    Returns (results: dict prompt -> conditioning, leftovers: list[str], stats: dict).
+    ``leftovers`` are prompts the caller must still encode the normal way (ineligible
+    prompt/encoder, singleton batches, failures). ``batch_size <= 1`` returns everything as
+    leftovers and touches nothing. The first batch of every call is self-checked against a
+    single-prompt encode of its shortest prompt (most padding); on mismatch batching is
+    disabled for this CLIP and every prompt is returned as a leftover.
+    """
+    prompts = list(prompts)
+    stats = {"batched": 0, "forwards": 0, "max_rows": 0, "reason": None, "verify": None}
+    try:
+        batch_size = int(batch_size)
+    except (TypeError, ValueError):
+        batch_size = 1
+    if batch_size <= 1 or not prompts:
+        return {}, prompts, stats
+    if _te_disabled(clip_model):
+        stats["reason"] = "disabled earlier in this run (self-check or unsupported output)"
+        return {}, prompts, stats
+    target, reason = _te_batch_target(clip_model)
+    if target is None:
+        stats["reason"] = reason
+        return {}, prompts, stats
+    inner, clip_name, pad_id = target
+
+    # eligibility + dedupe by token ids
+    tok = {}
+    by_ids = {}
+    for p in prompts:
+        got, _why = _batchable_ids(clip_model, p, clip_name, pad_id)
+        if got is None:
+            continue
+        tok[p] = got[0]
+        by_ids.setdefault(tuple(got[1]), []).append(p)
+    order = sorted(by_ids, key=len, reverse=True)  # longest first, stable
+
+    results = {}
+    singles = set()   # prompts that go the legacy way (singleton batches / OOM at 1)
+    state = {"bs": batch_size, "forwards": 0, "max_rows": 0, "checked": False, "worst_rel": 0.0,
+             "worst_abs": 0.0, "shapes_equal": True, "verified": 0, "verify_fail": 0}
+
+    def _fail(why):
+        _te_disable(clip_model)
+        print(f"[GridTester] ⚠️ TE batching self-check failed ({why}) - encoding per prompt")
+        return {}, prompts, stats
+
+    def _interrupted():
+        if comfy.model_management.processing_interrupted():
+            print(f"\n[GridTester] 🛑 INTERRUPTED during text-encoder batching - Stopping all encoding")
+            raise comfy.model_management.InterruptProcessingException()
+
+    def _run(batch):
+        """Encode ``batch`` (list of ids tuples) in one forward. Returns {prompt: cond}."""
+        id_lists = [list(ids) for ids in batch]
+        replay = _ReplayEncode(inner.encode, id_lists, pad_id)
+        local = {}
+        inner.encode = replay
+        try:
+            for ids in batch:
+                _interrupted()
+                p = by_ids[ids][0]
+                local[p] = _encode_single(clip_model, p, clip_skip, tokens=tok[p])
+        finally:
+            # always un-shadow, whatever happened (incl. CUDA OOM / interrupt)
+            vars(inner).pop("encode", None)
+            replay.out = None
+        if replay.misses:
+            raise _TEBatchUnsupported("replay miss")
+        state["forwards"] += replay.forwards
+        state["max_rows"] = max(state["max_rows"], len(batch))
+        return local
+
+    def _single_ref(p):
+        return _encode_single(clip_model, p, clip_skip, tokens=tok[p])
+
+    def _accept(batch, local):
+        """Self-check (first batch) / verify (dev), then merge. Returns a failure reason or None."""
+        skip = None
+        if not state["checked"]:
+            skip = by_ids[batch[-1]][0]  # shortest prompt of the first batch = most padding
+            ref = _single_ref(skip)
+            ok, why, rel, mx, same = _compare_conds(local[skip], ref)
+            if not ok:
+                return why
+            local[skip] = ref
+            state["checked"] = True
+            state["worst_rel"] = max(state["worst_rel"], rel)
+            state["worst_abs"] = max(state["worst_abs"], mx)
+        if verify:
+            for ids in batch:
+                p = by_ids[ids][0]
+                state["verified"] += 1
+                if p == skip:
+                    continue
+                ok, why, rel, mx, same = _compare_conds(local[p], _single_ref(p))
+                state["worst_rel"] = max(state["worst_rel"], rel)
+                state["worst_abs"] = max(state["worst_abs"], mx)
+                state["shapes_equal"] = state["shapes_equal"] and same
+                if not ok:
+                    state["verify_fail"] += 1
+        results.update(local)
+        return None
+
+    def _process(batch):
+        try:
+            local = _run(batch)
+        except comfy.model_management.InterruptProcessingException:
+            raise
+        except _TEBatchUnsupported:
+            raise
+        except Exception as e:
+            if not _is_oom(e):
+                raise
+            try:
+                comfy.model_management.soft_empty_cache()
+            except Exception:
+                pass
+            if len(batch) <= 1:
+                singles.update(by_ids[batch[0]])
+                return None
+            half = len(batch) // 2
+            print(f"[GridTester] ⚠️ TE batch OOM at B={len(batch)} - retrying at B={half}")
+            state["bs"] = min(state["bs"], half)
+            for part in (batch[:half], batch[half:]):
+                if len(part) <= 1:
+                    singles.update(by_ids[part[0]])
+                else:
+                    bad = _process(part)
+                    if bad:
+                        return bad
+            return None
+        return _accept(batch, local)
+
+    pending = list(order)
+    try:
+        while pending:
+            _interrupted()
+            tmax = len(pending[0])
+            n = 1
+            while n < state["bs"] and n < len(pending) and (n + 1) * tmax <= TE_BATCH_MAX_TOKENS:
+                n += 1
+            batch, pending = pending[:n], pending[n:]
+            if len(batch) == 1:
+                singles.update(by_ids[batch[0]])
+                continue
+            bad = _process(batch)
+            if bad:
+                return _fail(bad)
+    except _TEBatchUnsupported as e:
+        return _fail(f"unsupported encoder output: {e}")
+
+    # prompts that tokenised identically to a batched one share its conditioning
+    for plist in by_ids.values():
+        if plist[0] in results:
+            for other in plist[1:]:
+                results[other] = results[plist[0]]
+    leftovers = [p for p in prompts if p not in results]
+    stats["batched"] = len(results)
+    stats["forwards"] = state["forwards"]
+    stats["max_rows"] = state["max_rows"]
+    if verify and state["verified"]:
+        print(f"[GridTester] 🧪 TE batch verify: {state['verified']} prompts, worst relL2={state['worst_rel']:.2e}, "
+              f"worst max|Δ|={state['worst_abs']:.2e}, shapes equal: {'yes' if state['shapes_equal'] else 'no'}"
+              f"{'' if not state['verify_fail'] else ', FAILED: ' + str(state['verify_fail'])}")
+        stats["verify"] = {"prompts": state["verified"], "worst_rel": state["worst_rel"],
+                           "worst_abs": state["worst_abs"], "shapes_equal": state["shapes_equal"],
+                           "failed": state["verify_fail"]}
+    return results, leftovers, stats
+
+
 # ============================================================
 # === PROMPT COMBINATOR ENCODING (AND / CAT / AVG / BREAK) ===
 # ============================================================
 
-def _encode_single(clip_model, text, clip_skip=0):
+def _encode_single(clip_model, text, clip_skip=0, tokens=None):
     """
     Encode a single text segment. BREAK is handled natively by clip.tokenize().
 
@@ -104,6 +472,8 @@ def _encode_single(clip_model, text, clip_skip=0):
         clip_model: CLIP model instance
         text: Prompt text (may contain BREAK keyword)
         clip_skip: CLIP layer skip value
+        tokens: Pre-tokenized ``clip_model.tokenize(text)`` result (the batched
+            text-encoder path tokenises once to check eligibility and reuses it)
 
     Returns:
         Conditioning list: [[cond_tensor, pooled_dict]]
@@ -113,17 +483,25 @@ def _encode_single(clip_model, text, clip_skip=0):
         original_layer = clip_model.cond_stage_model.clip_layer
         clip_model.cond_stage_model.set_clip_options({"layer": clip_skip})
 
-    tokens = clip_model.tokenize(text)  # BREAK handled natively here
-    pooled_dict = clip_model.encode_from_tokens(tokens, return_pooled=True, return_dict=True)
+    try:
+        if tokens is None:
+            tokens = clip_model.tokenize(text)  # BREAK handled natively here
+        pooled_dict = clip_model.encode_from_tokens(tokens, return_pooled=True, return_dict=True)
+    finally:
+        # restore even when the encode raises (e.g. CUDA OOM that the batched path retries)
+        if original_layer is not None:
+            clip_model.cond_stage_model.set_clip_options({"layer": original_layer})
     cond = pooled_dict.pop("cond")
     result = [[cond, pooled_dict]]
 
-    if original_layer is not None:
-        clip_model.cond_stage_model.set_clip_options({"layer": original_layer})
-
     return result
 
 
+def _has_combinators(prompt):
+    """True when the prompt uses AND / CAT / AVG() (the slow path of encode_prompt_with_combinators)."""
+    return ' AND ' in prompt or ' CAT ' in prompt or 'AVG(' in prompt
+
+
 def _parse_and_segments(prompt):
     """
     Split prompt on AND keyword. Returns list of (text, weight) tuples.
@@ -268,7 +646,7 @@ def encode_prompt_with_combinators(clip_model, prompt, clip_skip=0):
         Conditioning list: [[cond_tensor, pooled_dict], ...]
     """
     # Fast path: if no combinators present, use direct encoding
-    if ' AND ' not in prompt and ' CAT ' not in prompt and 'AVG(' not in prompt:
+    if not _has_combinators(prompt):
         return _encode_single(clip_model, prompt, clip_skip)
 
     # 1. Split on AND
@@ -317,7 +695,8 @@ def encode_prompt_with_combinators(clip_model, prompt, clip_skip=0):
     return final
 
 
-def batch_encode_with_cache(clip_model, prompts, cond_cache, prompt_type="positive", batch_size=64, clip_skip=0):
+def batch_encode_with_cache(clip_model, prompts, cond_cache, prompt_type="positive", batch_size=64, clip_skip=0,
+                            te_batch_size=DEFAULT_TE_BATCH_SIZE, te_batch_verify=False):
     """
     Batch encode prompts while checking persistent cache first.
     Only encodes prompts that aren't already cached.
@@ -329,6 +708,9 @@ def batch_encode_with_cache(clip_model, prompts, cond_cache, prompt_type="positi
         prompt_type: "positive" or "negative"
         batch_size: Number of prompts to encode at once
         clip_skip: Number of CLIP layers to skip from the end (0 = use last layer, -1 = skip 1 layer, -2 = skip 2 layers)
+        te_batch_size: Encode up to this many eligible prompts per text-encoder forward
+            (MOSS-Image / Krea 2 / Qwen-Image only; 1 = off, the per-prompt path)
+        te_batch_verify: Dev only. Re-encode every batched prompt singly and print the worst difference
         
     Returns:
         dict: Mapping of prompt text to conditioning tensors
@@ -360,10 +742,28 @@ def batch_encode_with_cache(clip_model, prompts, cond_cache, prompt_type="positi
         # the whole loop; without it ComfyUI re-prepares it on every single encode.
         _t_enc_start = time.time()
         
-        prompts_list = list(prompts_to_encode)
-        total_batches = (len(prompts_list) + batch_size - 1) // batch_size
+        prompts_all = list(prompts_to_encode)
+        prompts_list = prompts_all
+        te_stats = None
         
-        with torch.no_grad(), pinned_clip(clip_model, prompts_list):
+        with torch.no_grad(), pinned_clip(clip_model, prompts_all):
+            if te_batch_size and te_batch_size > 1:
+                # Several prompts per text-encoder forward; whatever can't be batched comes
+                # back as leftovers and takes the unchanged per-prompt loop below.
+                try:
+                    batched, prompts_list, te_stats = encode_prompts_batched(
+                        clip_model, prompts_all, te_batch_size, clip_skip, te_batch_verify)
+                except comfy.model_management.InterruptProcessingException:
+                    raise
+                except Exception as e:
+                    print(f"[GridTester] ⚠️ TE batching failed ({e}) - encoding per prompt")
+                    batched, prompts_list, te_stats = {}, prompts_all, None
+                for prompt, conditioning in batched.items():
+                    results[prompt] = conditioning
+                    cond_cache.set(prompt, conditioning, prompt_type)
+                if te_stats and te_stats.get("reason"):
+                    print(f"[GridTester] ℹ️ TE batching off for this encoder ({te_stats['reason']})")
+            total_batches = (len(prompts_list) + batch_size - 1) // batch_size
             for batch_idx in range(0, len(prompts_list), batch_size):
                 batch_prompts = prompts_list[batch_idx:batch_idx + batch_size]
                 current_batch = (batch_idx // batch_size) + 1
@@ -399,14 +799,22 @@ def batch_encode_with_cache(clip_model, prompts, cond_cache, prompt_type="positi
                     print(f"[GridTester]   Batch {current_batch}/{total_batches} | Encoded {encoded_count}/{len(prompts_list)}")
         
         _enc_elapsed = time.time() - _t_enc_start
-        _per_prompt_ms = (_enc_elapsed / len(prompts_list) * 1000.0) if prompts_list else 0.0
-        print(f"[GridTester] ✅ Batch encoding complete! {len(prompts_list)} prompts in {_enc_elapsed:.1f}s "
-              f"({_per_prompt_ms:.0f} ms/prompt, encoder prepared once)")
+        _per_prompt_ms = (_enc_elapsed / len(prompts_all) * 1000.0) if prompts_all else 0.0
+        _te_note = ""
+        if te_stats and te_stats.get("batched"):
+            _te_note = (f"; {te_stats['batched']} batched in {te_stats['forwards']} forwards of <={te_stats['max_rows']}, "
+                        f"{len(prompts_list)} per-prompt")
+            if te_stats.get("verify"):
+                v = te_stats["verify"]
+                _te_note += f", verify worst relL2={v['worst_rel']:.1e}"
+        print(f"[GridTester] ✅ Batch encoding complete! {len(prompts_all)} prompts in {_enc_elapsed:.1f}s "
+              f"({_per_prompt_ms:.0f} ms/prompt, encoder prepared once{_te_note})")
     
     return results
 
 
-def batch_encode_prompts(patched_clip, unique_positives, unique_negatives, cond_cache, clip_skip=0, enable_disk_cache=True):
+def batch_encode_prompts(patched_clip, unique_positives, unique_negatives, cond_cache, clip_skip=0, enable_disk_cache=True,
+                         te_batch_size=DEFAULT_TE_BATCH_SIZE, te_batch_verify=False):
     """
     Batch encode both positive and negative prompts.
     
@@ -417,6 +825,8 @@ def batch_encode_prompts(patched_clip, unique_positives, unique_negatives, cond_
         cond_cache: ConditioningCache instance
         clip_skip: Number of CLIP layers to skip from the end
         enable_disk_cache: Whether to save cache to disk (passed through to cache)
+        te_batch_size: Prompts per text-encoder forward (1 = off), see batch_encode_with_cache
+        te_batch_verify: Dev only, see batch_encode_with_cache
         
     Returns:
         dict: conditioning_cache with "positive" and "negative" keys
@@ -437,7 +847,9 @@ def batch_encode_prompts(patched_clip, unique_positives, unique_negatives, cond_
             cond_cache, 
             prompt_type="positive",
             batch_size=16,
-            clip_skip=clip_skip
+            clip_skip=clip_skip,
+            te_batch_size=te_batch_size,
+            te_batch_verify=te_batch_verify
         )
 
         conditioning_cache["negative"] = batch_encode_with_cache(
@@ -446,7 +858,9 @@ def batch_encode_prompts(patched_clip, unique_positives, unique_negatives, cond_
             cond_cache, 
             prompt_type="negative",
             batch_size=16,
-            clip_skip=clip_skip
+            clip_skip=clip_skip,
+            te_batch_size=te_batch_size,
+            te_batch_verify=te_batch_verify
         )
     
     # Save and print cache stats (only saves if disk cache is enabled)

config_builder_node.py

@@ -928,6 +928,12 @@ class UltimateConfigBuilder:
             session_settings["vae_batch_size"] = int(state.get("vae_batch_size", 1))
         except (TypeError, ValueError):
             session_settings["vae_batch_size"] = 1
+        # Text-encoder batching: prompts per TE forward (1 = off). Session-level like vae_batch_size.
+        from .batch_encoding import DEFAULT_TE_BATCH_SIZE
+        try:
+            session_settings["te_batch_size"] = max(1, int(state.get("te_batch_size", DEFAULT_TE_BATCH_SIZE)))
+        except (TypeError, ValueError):
+            session_settings["te_batch_size"] = DEFAULT_TE_BATCH_SIZE
 
         # session_name moves into session_settings too so the Generator's
         # session_name widget can be removed (Phase 2).

generation_orchestrator.py

@@ -16,7 +16,7 @@ import comfy.sd  # Required for async workers
 from comfy.model_management import InterruptProcessingException
 
 from .trigger_words import collect_unique_prompts_with_triggers, build_prompt_with_triggers, clear_trigger_caches
-from .batch_encoding import batch_encode_prompts, encode_prompt_with_combinators, pinned_clip
+from .batch_encoding import batch_encode_prompts, encode_prompt_with_combinators, pinned_clip, DEFAULT_TE_BATCH_SIZE
 from .manifest_utils import load_existing_manifest, save_manifest
 from .model_loader import (
     load_checkpoint, load_loras, cleanup_model_references,
@@ -1152,9 +1152,20 @@ def run_generation_loop(
         if clip_skip != 0:
             print(f"[GridTester] 🔧 Using clip_skip={clip_skip}")
 
+        # Text-encoder batching (several prompts per TE forward). Session-level Run Setting;
+        # te_batch_verify is a hidden dev key (never emitted by the Builder) used by the live check.
+        try:
+            te_batch_size = max(1, int((session_settings or {}).get("te_batch_size", DEFAULT_TE_BATCH_SIZE)))
+        except (TypeError, ValueError):
+            te_batch_size = DEFAULT_TE_BATCH_SIZE
+        te_batch_verify = bool((session_settings or {}).get("te_batch_verify", False))
+        if te_batch_size > 1:
+            print(f"[GridTester] ⚙️ te_batch_size={te_batch_size} (text-encoder batching)")
+
         try:
             conditioning_cache = batch_encode_prompts(
-                patched_clip, unique_positives, unique_negatives, cond_cache, clip_skip, enable_disk_cache=save_conditioning_cache_to_file
+                patched_clip, unique_positives, unique_negatives, cond_cache, clip_skip, enable_disk_cache=save_conditioning_cache_to_file,
+                te_batch_size=te_batch_size, te_batch_verify=te_batch_verify
             )
         except InterruptProcessingException:
             print(f"\n[GridTester] 🛑 INTERRUPTED during pre-encoding - Stopping all jobs")

tests/conftest.py

@@ -99,6 +99,7 @@ _GEN_ORCH_DEPS = {
     },
     "batch_encoding": {
         "batch_encode_prompts": lambda *a, **kw: None,
+        "DEFAULT_TE_BATCH_SIZE": 1,
         "encode_prompt_with_combinators": lambda *a, **kw: None,
         "pinned_clip": lambda clip=None, *a, **kw: __import__("contextlib").nullcontext(clip),
         "encode_pair_cached": lambda *a, **kw: (None, None),

tests/test_state_to_configs_json.py

@@ -200,3 +200,21 @@ def test_legacy_dual_array_preserves_cartesian_form():
     expected = "legacyLora:[0.5, 1.0]:[0.7, 1.0]"
     assert any(l == expected for l in loras), \
         f"Expected legacy dual-array form {expected!r}. Got: {loras}"
+
+
+def _te_batch_size_for(**overrides):
+    state = make_state()
+    state.update(overrides)
+    parsed = json.loads(UltimateConfigBuilder.state_to_configs_json(state))
+    return parsed["_session_settings"]["te_batch_size"]
+
+
+def test_te_batch_size_defaults_off_and_is_clamped():
+    """te_batch_size is a session-level Run Setting: default 1 (off), junk/<1 -> 1."""
+    assert _te_batch_size_for() == 1
+    assert _te_batch_size_for(te_batch_size=8) == 8
+    assert _te_batch_size_for(te_batch_size="8") == 8
+    assert _te_batch_size_for(te_batch_size="junk") == 1
+    assert _te_batch_size_for(te_batch_size=None) == 1
+    assert _te_batch_size_for(te_batch_size=0) == 1
+    assert _te_batch_size_for(te_batch_size=-3) == 1

tests/test_te_batching.py

@@ -0,0 +1,474 @@
+"""Several prompts through ONE text-encoder forward (batch_encoding.encode_prompts_batched).
+
+The batched path must hand back exactly what the per-prompt path returns (same dict
+keys, shapes, dtypes, values), fall back for anything it cannot batch safely, and
+never leave the inner encoder's ``encode`` shadowed. Fake numpy tensors, no torch.
+"""
+import contextlib
+import importlib.util
+import os
+import sys
+import types
+
+import pytest
+
+np = pytest.importorskip("numpy")
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+_PKG = os.path.basename(_ROOT)
+_mm = sys.modules.get("comfy.model_management")
+if _mm is not None and "comfy" in sys.modules and not hasattr(sys.modules["comfy"], "model_management"):
+    sys.modules["comfy"].model_management = _mm
+_spec = importlib.util.spec_from_file_location(f"{_PKG}.batch_encoding_real_te", os.path.join(_ROOT, "batch_encoding.py"))
+be = importlib.util.module_from_spec(_spec)
+be.__package__ = _PKG
+_spec.loader.exec_module(be)
+if not hasattr(be.torch, "no_grad"):
+    be.torch = types.SimpleNamespace(no_grad=contextlib.nullcontext,
+                                     cuda=types.SimpleNamespace(is_available=lambda: False))
+
+
+class _Interrupt(Exception):
+    pass
+
+
+class _MM:
+    """Replaces be.comfy.model_management for the duration of a test."""
+
+    def __init__(self):
+        self.interrupt_on_call = None
+        self.calls = 0
+        self.empty_cache_calls = 0
+        self.OOM_EXCEPTION = MemoryError
+        self.InterruptProcessingException = _Interrupt
+
+    def processing_interrupted(self):
+        self.calls += 1
+        return self.interrupt_on_call is not None and self.calls >= self.interrupt_on_call
+
+    def soft_empty_cache(self):
+        self.empty_cache_calls += 1
+
+
+@pytest.fixture(autouse=True)
+def mm(monkeypatch):
+    fake = _MM()
+    monkeypatch.setattr(be.comfy, "model_management", fake, raising=False)
+    be._TE_BATCH_DISABLED.clear()
+    be._TE_BATCH_DISABLED_IDS.clear()
+    return fake
+
+
+PAD = 0
+VEC = np.array([1.0, 2.0, 3.0, 4.0])
+
+
+class FakeInner:
+    special_tokens = {"pad": PAD}
+    enable_attention_masks = True
+    return_attention_masks = True
+
+    def __init__(self, pooled=None, bias_when_batched=0.0, oom_over=None, fail=None):
+        self.calls = []            # list of row lists passed to the REAL encode
+        self.layers = []
+        self.owner = None
+        self.pooled = pooled
+        self.bias = bias_when_batched
+        self.oom_over = oom_over
+        self.fail = fail
+
+    def process_tokens(self, *a, **kw):  # presence only
+        raise AssertionError("not used")
+
+    def encode(self, rows):
+        T = len(rows[0])
+        assert all(len(r) == T for r in rows), "rows must be equal length"
+        self.calls.append([list(r) for r in rows])
+        if self.owner is not None:
+            self.layers.append(getattr(self.owner, "clip_layer", None))
+        if self.fail:
+            raise self.fail
+        if self.oom_over is not None and len(rows) > self.oom_over:
+            raise MemoryError("CUDA out of memory")
+        ids = np.array(rows, dtype=np.float64)
+        z = np.cumsum(ids, axis=1)[:, :, None] * VEC[None, None, :]
+        if len(rows) > 1:
+            z = z + self.bias
+        mask = np.ones(ids.shape, dtype=np.int64)
+        for b, r in enumerate(rows):
+            if PAD in r:
+                mask[b, r.index(PAD):] = 0
+        pooled = None if self.pooled is None else np.zeros((len(rows), 2))
+        return z.astype(np.float32), pooled, {"attention_mask": mask}
+
+    def encode_token_weights(self, token_weight_pairs):
+        to_encode = [[t[0] for t in section] for section in token_weight_pairs]
+        o = self.encode(to_encode)
+        return o[0], o[1], {"attention_mask": o[2]["attention_mask"]}
+
+
+class FakeCSM:
+    def __init__(self, inner, clip_name="qwen35_9b", with_layer=False):
+        self.clip = clip_name
+        self.clip_name = clip_name
+        setattr(self, clip_name, inner)
+        inner.owner = self
+        if with_layer:
+            self.clip_layer = -1
+
+    def set_clip_options(self, opts):
+        if "layer" in opts:
+            self.clip_layer = opts["layer"]
+
+    def encode_token_weights(self, tokens):
+        return getattr(self, self.clip).encode_token_weights(tokens[self.clip_name])
+
+
+class FakeKreaCSM(FakeCSM):
+    """Outer post-processing: strip a 3-token template prefix and pop the all-ones mask."""
+
+    def encode_token_weights(self, tokens):
+        z, pooled, extra = super().encode_token_weights(tokens)
+        extra = dict(extra)
+        assert int(extra.pop("attention_mask").min()) == 1
+        return z[:, 3:], pooled, extra
+
+
+class FakeCLIP:
+    def __init__(self, csm, clip_name="qwen35_9b"):
+        self.cond_stage_model = csm
+        self.clip_name = clip_name
+        self.patcher = object()
+
+    def tokenize(self, text):
+        t = text.strip()
+        if " BREAK " in t:
+            parts = t.split(" BREAK ")
+            return {self.clip_name: [[(ord(c) % 50 + 1, 1.0) for c in p] for p in parts]}
+        if t.startswith("(") and t.endswith(")"):
+            return {self.clip_name: [[(ord(c) % 50 + 1, 1.2) for c in t[1:-1]]]}
+        return {self.clip_name: [[(PAD if c == "#" else ord(c) % 50 + 1, 1.0) for c in t]]}
+
+    def load_model(self, tokens={}):
+        return self.patcher
+
+    def encode_from_tokens(self, tokens, return_pooled=False, return_dict=False):
+        out = self.cond_stage_model.encode_token_weights(tokens)
+        d = {"cond": out[0], "pooled_output": out[1]}
+        if len(out) > 2:
+            d.update(out[2])
+        return d
+
+
+def make(krea=False, with_layer=False, **inner_kw):
+    inner = FakeInner(**inner_kw)
+    csm = (FakeKreaCSM if krea else FakeCSM)(inner, with_layer=with_layer)
+    return FakeCLIP(csm), inner, csm
+
+
+def single(clip, text, skip=0):
+    return be.encode_prompt_with_combinators(clip, text, skip)
+
+
+def assert_same(a, b):
+    assert len(a) == len(b)
+    (ca, da), (cb, db) = a[0], b[0]
+    assert ca.shape == cb.shape and ca.dtype == cb.dtype
+    assert np.array_equal(ca, cb)
+    assert sorted(da) == sorted(db)
+    for k in da:
+        if hasattr(da[k], "shape"):
+            assert da[k].shape == db[k].shape and np.array_equal(da[k], db[k])
+        else:
+            assert da[k] == db[k]
+
+
+def multi_row_calls(inner):
+    return [c for c in inner.calls if len(c) > 1]
+
+
+PROMPTS5 = ["abcdefghij", "abcdefgh", "abcdef", "abcd", "ab"]
+
+
+# ---- 1/2: grouping, slicing, padding -----------------------------------------
+
+def test_batches_of_four_with_singleton_leftover_and_equal_results():
+    clip, inner, _ = make()
+    results, leftovers, stats = be.encode_prompts_batched(clip, PROMPTS5, 4)
+    assert [len(c) for c in multi_row_calls(inner)] == [4]
+    assert leftovers == ["ab"]                       # singleton batch -> legacy path
+    assert stats["batched"] == 4 and stats["forwards"] == 1 and stats["max_rows"] == 4
+    for p in PROMPTS5[:4]:
+        ref = single(clip, p)
+        assert_same(results[p], ref)
+        assert results[p][0][0].shape == (1, len(p), 4)
+        assert results[p][0][1]["attention_mask"].shape == (1, len(p))
+        assert results[p][0][1]["attention_mask"].min() == 1
+
+
+def test_rows_are_right_padded_longest_first():
+    clip, inner, _ = make()
+    be.encode_prompts_batched(clip, PROMPTS5[::-1], 4)
+    rows = multi_row_calls(inner)[0]
+    assert all(len(r) == 10 for r in rows)
+    lens = [len([t for t in r if t != PAD]) for r in rows]
+    assert lens == [10, 8, 6, 4]
+    for r, n in zip(rows, lens):
+        assert r[n:] == [PAD] * (10 - n) and PAD not in r[:n]
+
+
+# ---- 3: outer post-processing survives per item -----------------------------------
+
+def test_krea_like_outer_strip_and_mask_pop_per_item():
+    clip, inner, _ = make(krea=True)
+    prompts = ["abcdefgh", "abcdef", "abcd"]
+    results, leftovers, _ = be.encode_prompts_batched(clip, prompts, 4)
+    assert leftovers == []
+    for p in prompts:
+        ref = single(clip, p)
+        assert_same(results[p], ref)
+        assert results[p][0][0].shape == (1, len(p) - 3, 4)
+        assert "attention_mask" not in results[p][0][1]
+
+
+# ---- 4: fallbacks ----------------------------------------------------------------
+
+@pytest.mark.parametrize("bad", ["(weighted)", "a AND b cc", "a CAT b cc", "a AVG(0.5) b", "abc BREAK def", "ab#cd"])
+def test_ineligible_prompts_never_enter_a_batched_forward(bad):
+    clip, inner, _ = make()
+    prompts = [bad, "abcdefgh", "abcdef", "abcd"]
+    results, leftovers, _ = be.encode_prompts_batched(clip, prompts, 4)
+    assert bad in leftovers and bad not in results
+    for p in prompts[1:]:
+        assert p in results
+    badids = [t[0] for t in clip.tokenize(bad)[clip.clip_name][0]]
+    for rows in multi_row_calls(inner):
+        assert badids not in rows
+
+
+def test_multi_key_tokens_dict_falls_back():
+    clip, inner, _ = make()
+    orig = clip.tokenize
+    clip.tokenize = lambda t: dict(orig(t), extra=[[(1, 1.0)]])
+    results, leftovers, stats = be.encode_prompts_batched(clip, PROMPTS5, 4)
+    assert results == {} and leftovers == PROMPTS5 and inner.calls == []
+
+
+def test_non_allowlisted_clip_name():
+    inner = FakeInner()
+    clip = FakeCLIP(FakeCSM(inner, clip_name="t5xxl"), clip_name="t5xxl")
+    results, leftovers, stats = be.encode_prompts_batched(clip, PROMPTS5, 4)
+    assert results == {} and leftovers == PROMPTS5 and "t5xxl" in stats["reason"] and inner.calls == []
+
+
+def test_encoder_without_masks_or_with_end_token_is_excluded():
+    clip, inner, _ = make()
+    inner.enable_attention_masks = False
+    assert be.encode_prompts_batched(clip, PROMPTS5, 4)[0] == {}
+    inner.enable_attention_masks = True
+    inner.special_tokens = {"pad": PAD, "end": 7}
+    assert be.encode_prompts_batched(clip, PROMPTS5, 4)[0] == {}
+    assert inner.calls == []
+
+
+def test_multi_encoder_wrapper_and_clip_schedule_are_excluded():
+    clip, inner, csm = make()
+    del csm.clip
+    assert be.encode_prompts_batched(clip, PROMPTS5, 4)[0] == {}
+    clip2, inner2, _ = make()
+    clip2.use_clip_schedule = True
+    assert be.encode_prompts_batched(clip2, PROMPTS5, 4)[0] == {}
+
+
+def test_pooled_output_disables_batching_for_the_run(capsys):
+    clip, inner, _ = make(pooled=True)
+    results, leftovers, _ = be.encode_prompts_batched(clip, PROMPTS5, 4)
+    assert results == {} and leftovers == PROMPTS5
+    assert "encode" not in vars(inner)
+    n = len(inner.calls)
+    results, leftovers, stats = be.encode_prompts_batched(clip, PROMPTS5, 4)
+    assert results == {} and leftovers == PROMPTS5 and len(inner.calls) == n and "disabled" in stats["reason"]
+
+
+# ---- 5: off means off ------------------------------------------------------------
+
+def test_batch_size_one_touches_nothing():
+    clip, inner, _ = make()
+    results, leftovers, stats = be.encode_prompts_batched(clip, PROMPTS5, 1)
+    assert results == {} and leftovers == PROMPTS5 and inner.calls == []
+    assert "encode" not in vars(inner)
+
+
+# ---- 6: self-check -----------------------------------------------------------------
+
+def test_self_check_failure_falls_back_and_stays_disabled(capsys):
+    clip, inner, _ = make(bias_when_batched=100.0)
+    results, leftovers, _ = be.encode_prompts_batched(clip, PROMPTS5, 4)
+    assert results == {} and leftovers == PROMPTS5
+    assert "self-check failed" in capsys.readouterr().out
+    assert "encode" not in vars(inner)
+    # the per-prompt path is unaffected and matches a never-batched clip
+    ref_clip, _, _ = make()
+    for p in PROMPTS5:
+        assert_same(single(clip, p), single(ref_clip, p))
+    n = len(inner.calls)
+    assert be.encode_prompts_batched(clip, PROMPTS5, 4)[0] == {}
+    assert len(inner.calls) == n
+
+
+# ---- 7: OOM ------------------------------------------------------------------------
+
+def test_oom_halves_the_batch_and_still_gives_correct_results(capsys, mm):
+    clip, inner, _ = make(oom_over=2)
+    prompts = ["abcdefghijkl"[:n] for n in range(12, 4, -1)]  # 8 distinct lengths
+    results, leftovers, _ = be.encode_prompts_batched(clip, prompts, 8)
+    out = capsys.readouterr().out
+    assert "OOM at B=8 - retrying at B=4" in out and "OOM at B=4 - retrying at B=2" in out
+    assert leftovers == [] and mm.empty_cache_calls >= 1
+    assert "encode" not in vars(inner)
+    for p in prompts:
+        assert_same(results[p], single(clip, p))
+
+
+# ---- 8: the shadow never leaks ----------------------------------------------------------
+
+def test_shadow_removed_after_success():
+    clip, inner, _ = make()
+    be.encode_prompts_batched(clip, PROMPTS5, 4)
+    assert "encode" not in vars(inner)
+
+
+def test_shadow_removed_after_an_exception():
+    clip, inner, _ = make(fail=RuntimeError("boom"))
+    with pytest.raises(RuntimeError):
+        be.encode_prompts_batched(clip, PROMPTS5, 4)
+    assert "encode" not in vars(inner)
+
+
+def test_shadow_removed_after_oom_that_never_recovers():
+    clip, inner, _ = make(oom_over=0)   # every forward raises
+    results, leftovers, _ = be.encode_prompts_batched(clip, PROMPTS5, 4)
+    assert results == {} and sorted(leftovers) == sorted(PROMPTS5)
+    assert "encode" not in vars(inner)
+
+
+@pytest.mark.parametrize("k", [1, 3])
+def test_shadow_removed_after_interrupt(mm, k):
+    clip, inner, _ = make()
+    mm.interrupt_on_call = k   # k=1: before install; k=3: inside the shadowed block
+    with pytest.raises(_Interrupt):
+        be.encode_prompts_batched(clip, PROMPTS5, 4)
+    assert "encode" not in vars(inner)
+
+
+# ---- 9: clip_skip -----------------------------------------------------------------------
+
+def test_clip_skip_applies_to_the_batched_forward_and_is_restored():
+    clip, inner, csm = make(with_layer=True)
+    results, _, _ = be.encode_prompts_batched(clip, PROMPTS5, 4, clip_skip=-2)
+    batched_layers = [lay for c, lay in zip(inner.calls, inner.layers) if len(c) > 1]
+    assert batched_layers == [-2]
+    assert csm.clip_layer == -1
+
+
+def test_clip_layer_is_restored_when_the_encode_raises():
+    clip, inner, csm = make(with_layer=True, fail=RuntimeError("boom"))
+    with pytest.raises(RuntimeError):
+        be._encode_single(clip, "abc", -2)
+    assert csm.clip_layer == -1
+
+
+# ---- 10: duplicates ---------------------------------------------------------------------
+
+def test_prompts_with_identical_tokens_share_one_row():
+    clip, inner, _ = make()
+    prompts = ["abcdef", " abcdef ", "abcd", "ab"]
+    results, leftovers, _ = be.encode_prompts_batched(clip, prompts, 4)
+    assert leftovers == []
+    rows = multi_row_calls(inner)[0]
+    assert len(rows) == 3
+    assert_same(results["abcdef"], results[" abcdef "])
+
+
+# ---- 11: token cap ----------------------------------------------------------------------
+
+def test_token_cap_prevents_oversized_batches():
+    clip, inner, _ = make()
+    long_prompts = ["a" * 3000, "b" * 3000, "c" * 3000]
+    results, leftovers, _ = be.encode_prompts_batched(clip, long_prompts, 8)
+    assert multi_row_calls(inner) == []          # 2 x 3000 > TE_BATCH_MAX_TOKENS: nothing batched
+    assert sorted(leftovers) == sorted(long_prompts)
+
+
+def test_token_cap_still_allows_batches_that_fit():
+    clip, inner, _ = make()
+    prompts = ["a" * 2000, "b" * 2000, "c" * 2000]
+    results, leftovers, _ = be.encode_prompts_batched(clip, prompts, 8)
+    sizes = [len(c) for c in multi_row_calls(inner)]
+    assert sizes == [2]
+    assert len(results) == 2 and len(leftovers) == 1
+
+
+def test_verify_mode_reports_every_prompt(capsys):
+    clip, inner, _ = make()
+    results, _, stats = be.encode_prompts_batched(clip, PROMPTS5[:4], 4, verify=True)
+    assert stats["verify"]["prompts"] == 4 and stats["verify"]["shapes_equal"] and stats["verify"]["failed"] == 0
+    assert "TE batch verify: 4 prompts" in capsys.readouterr().out
+
+
+# ---- 12: integration through batch_encode_with_cache ----------------------------------
+
+class MemoryCache:
+    def __init__(self):
+        self.store = {}
+        self.sets = 0
+
+    def get(self, prompt, prompt_type):
+        return self.store.get((prompt_type, prompt))
+
+    def set(self, prompt, cond, prompt_type):
+        self.sets += 1
+        self.store[(prompt_type, prompt)] = cond
+
+    def save(self):
+        pass
+
+
+def test_batch_encode_with_cache_batches_caches_and_reports(capsys):
+    clip, inner, _ = make()
+    cache = MemoryCache()
+    res = be.batch_encode_with_cache(clip, PROMPTS5 + ["a AND b cc"], cache, te_batch_size=4)
+    assert set(res) == set(PROMPTS5 + ["a AND b cc"]) and cache.sets == 6
+    assert all(v is not None for v in res.values())
+    out = capsys.readouterr().out
+    assert "Batch encoding complete! 6 prompts" in out
+    assert "4 batched in 1 forwards of <=4, 2 per-prompt" in out
+    for p in PROMPTS5[:4]:
+        assert_same(res[p], single(clip, p))
+
+
+def test_batch_encode_with_cache_default_is_the_legacy_path(capsys):
+    clip, inner, _ = make()
+    cache = MemoryCache()
+    res = be.batch_encode_with_cache(clip, PROMPTS5, cache)
+    assert len(res) == 5 and cache.sets == 5
+    assert [len(c) for c in inner.calls] == [1] * 5     # one forward per prompt, no shadow
+    assert "batched in" not in capsys.readouterr().out
+
+
+def test_ineligible_encoder_prints_off_note_and_still_encodes(capsys):
+    inner = FakeInner()
+    clip = FakeCLIP(FakeCSM(inner, clip_name="t5xxl"), clip_name="t5xxl")
+    res = be.batch_encode_with_cache(clip, PROMPTS5, MemoryCache(), te_batch_size=8)
+    assert len(res) == 5
+    assert "TE batching off for this encoder" in capsys.readouterr().out
+
+
+def test_batch_encode_prompts_passes_te_kwargs_through(capsys):
+    clip, inner, _ = make()
+    cache = MemoryCache()
+    out = be.batch_encode_prompts(clip, set(PROMPTS5), {"abcdefg", "abcde"}, cache, te_batch_size=4)
+    assert len(out["positive"]) == 5 and len(out["negative"]) == 2
+    assert multi_row_calls(inner)

web/conf_builder/conf-builder-config-management.js

@@ -6960,6 +6960,7 @@ export function renderRunSettingsSection(node, container) {
     if (node.state.save_conditioning_cache_to_file === undefined) node.state.save_conditioning_cache_to_file = false;
     if (node.state.enable_model_cache === undefined) node.state.enable_model_cache = false;
     if (node.state.vae_batch_size === undefined) node.state.vae_batch_size = 1;
+    if (node.state.te_batch_size === undefined) node.state.te_batch_size = 1;
     if (node.state.add_random_seeds_to_gens === undefined) node.state.add_random_seeds_to_gens = 0;
 
     const section = document.createElement("div");
@@ -7118,6 +7119,8 @@ export function renderRunSettingsSection(node, container) {
         tooltip: "Experimental: intelligent model/LoRA caching with async background preloading. Speeds up generation when switching LoRAs frequently. Loads cached LoRAs from RAM instead of disk. Disable to reduce RAM/VRAM usage." });
     _addRunSetting(content, { stateKey: "vae_batch_size", label: "VAE Batch Size", kind: "int", min: -1, max: 64, defaultValue: 1,
         tooltip: "How many images to encode/decode per VAE pass. Lower = less VRAM. -1 = process all at once. Default: 4." });
+    _addRunSetting(content, { stateKey: "te_batch_size", label: "Text Encoder Batch Size", kind: "int", min: 1, max: 32, defaultValue: 1,
+        tooltip: "Encode this many prompts per text-encoder pass (MOSS-Image, Krea 2, Qwen-Image only; other encoders ignore it). 1 = off. 8 is a good value once verified; weighted/AND/CAT/AVG/BREAK prompts are still encoded one by one." });
     _addRunSetting(content, { stateKey: "add_random_seeds_to_gens", label: "Additional Random Seeds", kind: "int", min: 0, max: 100, defaultValue: 0,
         tooltip: "Generate this many extra images per config using additional random seeds. 0 = disabled." });
 

web/conf_builder/conf-builder-main.js

@@ -211,7 +211,8 @@ app.registerExtension({
                     save_conditioning_cache_to_file: false,
                     enable_model_cache: false,
                     vae_batch_size: 1,
-                    add_random_seeds_to_gens: 0
+                    add_random_seeds_to_gens: 0,
+                    te_batch_size: 1
                 };
 
                 // Create HTML container immediately
@@ -323,6 +324,7 @@ app.registerExtension({
                             if (this.state.enable_model_cache === undefined) this.state.enable_model_cache = false;
                             if (this.state.vae_batch_size === undefined) this.state.vae_batch_size = 1;
                             if (this.state.add_random_seeds_to_gens === undefined) this.state.add_random_seeds_to_gens = 0;
+                            if (this.state.te_batch_size === undefined) this.state.te_batch_size = 1;
 
                             this.saveState();
                             this.renderUI();