diff --git a/ProjectStructure.md b/ProjectStructure.md index ba96c67..9395766 100644 --- a/ProjectStructure.md +++ b/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) diff --git a/README.md b/README.md index a6896bc..a7f0642 100644 --- a/README.md +++ b/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. diff --git a/batch_encoding.py b/batch_encoding.py index 56b088a..6e5091d 100644 --- a/batch_encoding.py +++ b/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) diff --git a/config_builder_node.py b/config_builder_node.py index 02df6b4..fe4f7b5 100644 --- a/config_builder_node.py +++ b/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). diff --git a/generation_orchestrator.py b/generation_orchestrator.py index 27c163d..9b4aaf0 100644 --- a/generation_orchestrator.py +++ b/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") diff --git a/tests/conftest.py b/tests/conftest.py index 84602a3..5db2609 100644 --- a/tests/conftest.py +++ b/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), diff --git a/tests/test_state_to_configs_json.py b/tests/test_state_to_configs_json.py index eda9faf..fab7824 100644 --- a/tests/test_state_to_configs_json.py +++ b/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 diff --git a/tests/test_te_batching.py b/tests/test_te_batching.py new file mode 100644 index 0000000..68afd85 --- /dev/null +++ b/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) diff --git a/web/conf_builder/conf-builder-config-management.js b/web/conf_builder/conf-builder-config-management.js index cc893c5..6b4a820 100644 --- a/web/conf_builder/conf-builder-config-management.js +++ b/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." }); diff --git a/web/conf_builder/conf-builder-main.js b/web/conf_builder/conf-builder-main.js index ac93969..3f934b3 100644 --- a/web/conf_builder/conf-builder-main.js +++ b/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();