From 198138e51cd2b70f63a6797b020e836351686d55 Mon Sep 17 00:00:00 2001 From: Jason Hoku Date: Tue, 3 Nov 2026 22:16:43 -0800 Subject: [PATCH] perf(encode): prepare the text encoder once per pre-encode loop, not once per prompt ComfyUI's CLIP.encode_from_tokens() calls load_model() -> model_management.load_models_gpu() on every encode. With the dynamic VRAM loader (comfy_aimdo, 0.4x+) that re-runs ModelPatcherDynamic.load() for a model that is already resident: a walk over every module plus free_memory bookkeeping, visible as one Model MOSSImageTEModel_ prepared for dynamic VRAM loading. 7002MB Staged. ... log line per prompt (N prompts -> N+1 lines). USCG's pre-encode paid it once per prompt, and again on the single-encode paths (model switch, distributed, upscale, Florence2 hires). Measured on the 5060 Ti (ComfyUI 0.42.5), MOSS-Image TE = Qwen3.5-9B w4a8, same process state: template path (CLIPTextEncode), per prompt ....... 0.92 s (re-prepare every encode) USCG batch encode with pinned_clip, per prompt ... 0.84 s (one prepare per pass) USCG job, 21 prompts + 21 one-step images ........ 36.8 s -> 33.0 s for scale: SDXL 0.15 s, Krea 2 0.29 s, Qwen-Image 0.32 s per prompt on the template path The rest is the encoder forward itself: ~0.55 s fixed + ~2 ms/token for Qwen3.5-9B (24 of its 32 layers are Gated DeltaNet run through the pure-torch chunked recurrence, fp32 compute), so a 60-prompt MOSS grid still spends ~50 s encoding before the first image. That part is identical to the official workflow. batch_encoding.pinned_clip(clip, texts): load once through load_model() (sized by the longest text so any memory estimate still runs), shadow clip.load_model with a no-op for the block, restore on exit or exception, nesting safe. encode_pair_cached(): positive+negative with one load, nothing when both are cached. The pre-encode summary line now prints total time and ms/prompt. Applied to batch_encode_with_cache, the model-switch pre-encode in the main loop, the distributed group pre-encode and per-job encode, the deferred-upscale re-encode fallback, the dashboard upscale runner and the Florence2 hires step. Tests: tests/test_pinned_clip.py (10) against the real batch_encoding module; the conftest stub gained pinned_clip/encode_pair_cached. 137 passed, same 15 pre-existing torch-stub failures. Co-Authored-By: Claude Fable 5.1 --- README.md | 1 + batch_encoding.py | 94 ++++++++++++++++- florence2_hires.py | 30 ++---- generation_orchestrator.py | 96 +++++++++--------- tests/conftest.py | 2 + tests/test_pinned_clip.py | 200 +++++++++++++++++++++++++++++++++++++ upscale_runner.py | 13 +-- 7 files changed, 352 insertions(+), 84 deletions(-) create mode 100644 tests/test_pinned_clip.py diff --git a/README.md b/README.md index 163619a..eb9193c 100644 --- a/README.md +++ b/README.md @@ -748,6 +748,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. ### Model-Specific Prompts diff --git a/batch_encoding.py b/batch_encoding.py index 75e82cf..56b088a 100644 --- a/batch_encoding.py +++ b/batch_encoding.py @@ -4,12 +4,94 @@ Handles efficient batch encoding of prompts with caching. Supports prompt combinators: AND, CAT, AVG(), BREAK (native). """ +import contextlib import re +import time import torch import gc import comfy.model_management +# ============================================================ +# === KEEP THE TEXT ENCODER RESIDENT ACROSS A RUN OF ENCODES === +# ============================================================ + +@contextlib.contextmanager +def pinned_clip(clip_model, sample_texts=None): + """ + Load a CLIP/text encoder ONCE and keep ComfyUI from re-preparing it on every + encode inside the block. + + comfy.sd.CLIP.encode_from_tokens() calls self.load_model() -> + comfy.model_management.load_models_gpu([patcher]) on EVERY encode. Since the + dynamic VRAM loader (comfy_aimdo, ComfyUI 0.4x+) load_models_gpu() re-runs + ModelPatcherDynamic.load() even for a model that is already resident: a Python + walk over every module of the encoder plus free_memory()/get_free_memory() + bookkeeping. That is the repeated + + Model MOSSImageTEModel_ prepared for dynamic VRAM loading. 7002MB Staged. ... + + log line, once per prompt (N prompts -> N+1 lines). For a ~500-module LLM text + encoder (Qwen3.5-9B for MOSS-Image, Qwen2.5-VL for Qwen-Image, Qwen3-VL for + Krea 2) the re-prepare costs about as much as the encode itself; SDXL's small + CLIPs hide it, which is why SDXL pre-encoding always looked fast. + + Inside this context the encoder is loaded once through the normal load_model() + path (sized by the longest of ``sample_texts`` so any memory estimate still + runs) and ``clip_model.load_model`` is shadowed by a no-op returning the + already-loaded patcher. Nothing in a plain encode loop touches model + management, so the encoder stays resident. The shadow is removed on exit, even + on exceptions, so the CLIP behaves normally afterwards (sampling may then evict + it as usual). Nesting is safe: the inner block restores the outer shadow. + """ + if clip_model is None or not hasattr(clip_model, "load_model"): + yield clip_model + return + + tokens = {} + if sample_texts: + try: + longest = max(sample_texts, key=len) + tokens = clip_model.tokenize(longest) + except Exception: + tokens = {} + try: + patcher = clip_model.load_model(tokens) + except TypeError: + patcher = clip_model.load_model() # very old ComfyUI: load_model() took no tokens + if patcher is None: + patcher = getattr(clip_model, "patcher", None) + + inst = vars(clip_model) + had_shadow = "load_model" in inst + previous = inst.get("load_model") + clip_model.load_model = lambda tokens={}: patcher + try: + yield clip_model + finally: + if had_shadow: + clip_model.load_model = previous + else: + inst.pop("load_model", None) + + +def encode_pair_cached(clip_model, cache, positive, negative, clip_skip=0): + """ + Fill cache["positive"][positive] / cache["negative"][negative] if missing, + loading the encoder at most once for the pair, and return (pos_cond, neg_cond). + Does not touch the encoder at all when both are already cached. + """ + need_pos = positive not in cache["positive"] + need_neg = negative not in cache["negative"] + if need_pos or need_neg: + with pinned_clip(clip_model): + if need_pos: + cache["positive"][positive] = encode_prompt_with_combinators(clip_model, positive, clip_skip) + if need_neg: + cache["negative"][negative] = encode_prompt_with_combinators(clip_model, negative, clip_skip) + return cache["positive"][positive], cache["negative"][negative] + + # ============================================================ # === PROMPT COMBINATOR ENCODING (AND / CAT / AVG / BREAK) === # ============================================================ @@ -274,13 +356,14 @@ def batch_encode_with_cache(clip_model, prompts, cond_cache, prompt_type="positi clip_skip_msg = f" (clip_skip={clip_skip})" if clip_skip != 0 else "" print(f"[GridTester] 🚀 Batch encoding {len(prompts_to_encode)} {prompt_type} prompts{clip_skip_msg} (batch_size={batch_size})") - # Force model to stay in VRAM - comfy.model_management.load_models_gpu([clip_model.patcher]) + # The encoder is loaded ONCE by pinned_clip (below) and stays resident for + # 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 - with torch.no_grad(): + with torch.no_grad(), pinned_clip(clip_model, prompts_list): 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 @@ -315,7 +398,10 @@ def batch_encode_with_cache(clip_model, prompts, cond_cache, prompt_type="positi if current_batch % 5 == 0 or current_batch == total_batches: print(f"[GridTester] Batch {current_batch}/{total_batches} | Encoded {encoded_count}/{len(prompts_list)}") - print(f"[GridTester] ✅ Batch encoding complete!") + _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)") return results diff --git a/florence2_hires.py b/florence2_hires.py index 0d22889..c984eee 100644 --- a/florence2_hires.py +++ b/florence2_hires.py @@ -725,39 +725,25 @@ def run_florence2_step( # and crash with "mat1 and mat2 shapes cannot be multiplied". So we maintain # a per-clip-identity conditioning cache inside ckpt_cache. try: - from .batch_encoding import encode_prompt_with_combinators + from .batch_encoding import encode_pair_cached except ImportError: - from batch_encoding import encode_prompt_with_combinators + from batch_encoding import encode_pair_cached session_clip = fallback_handles[1] if clip_handle is session_clip: # Same clip as session — reuse the upscale_runner's shared cache. - if positive_prompt not in conditioning_cache["positive"]: - conditioning_cache["positive"][positive_prompt] = encode_prompt_with_combinators( - clip_handle, positive_prompt, clip_skip - ) - if negative_prompt not in conditioning_cache["negative"]: - conditioning_cache["negative"][negative_prompt] = encode_prompt_with_combinators( - clip_handle, negative_prompt, clip_skip - ) - pos_cond = conditioning_cache["positive"][positive_prompt] - neg_cond = conditioning_cache["negative"][negative_prompt] + pos_cond, neg_cond = encode_pair_cached( + clip_handle, conditioning_cache, positive_prompt, negative_prompt, clip_skip + ) else: # Different clip (per-manifest model) — use a per-clip-id cache inside # ckpt_cache so multiple images with the same per-item model still share. per_clip = ckpt_cache.setdefault("_cond_cache", {}) clip_id = id(clip_handle) bucket = per_clip.setdefault(clip_id, {"positive": {}, "negative": {}}) - if positive_prompt not in bucket["positive"]: - bucket["positive"][positive_prompt] = encode_prompt_with_combinators( - clip_handle, positive_prompt, clip_skip - ) - if negative_prompt not in bucket["negative"]: - bucket["negative"][negative_prompt] = encode_prompt_with_combinators( - clip_handle, negative_prompt, clip_skip - ) - pos_cond = bucket["positive"][positive_prompt] - neg_cond = bucket["negative"][negative_prompt] + pos_cond, neg_cond = encode_pair_cached( + clip_handle, bucket, positive_prompt, negative_prompt, clip_skip + ) # 10. VAE encode -> generate_image (project's known-working sampling path) try: diff --git a/generation_orchestrator.py b/generation_orchestrator.py index 6dafb08..12f7a82 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 +from .batch_encoding import batch_encode_prompts, encode_prompt_with_combinators, pinned_clip from .manifest_utils import load_existing_manifest, save_manifest from .model_loader import ( load_checkpoint, load_loras, cleanup_model_references, @@ -470,12 +470,13 @@ def run_deferred_upscales( print(f"[GridTester] ⚠️ Deferred upscale #{job_idx+1}: conditioning not in cache, re-encoding") try: clip_skip = conf.get("clip_skip", 0) - if final_positive is None: - final_positive = encode_prompt_with_combinators(patched_clip, actual_positive_prompt, clip_skip) - conditioning_cache["positive"][actual_positive_prompt] = final_positive - if final_negative is None: - final_negative = encode_prompt_with_combinators(patched_clip, actual_negative_prompt, clip_skip) - conditioning_cache["negative"][actual_negative_prompt] = final_negative + with pinned_clip(patched_clip): # one encoder load for both prompts + if final_positive is None: + final_positive = encode_prompt_with_combinators(patched_clip, actual_positive_prompt, clip_skip) + conditioning_cache["positive"][actual_positive_prompt] = final_positive + if final_negative is None: + final_negative = encode_prompt_with_combinators(patched_clip, actual_negative_prompt, clip_skip) + conditioning_cache["negative"][actual_negative_prompt] = final_negative except Exception as e: print(f"[GridTester] ⚠️ Deferred upscale #{job_idx+1}: re-encode failed: {e}, skipping") continue @@ -1586,29 +1587,28 @@ def run_generation_loop( if model_unique_positives: print(f"[GridTester] 🧠 Batch encoding {len(model_unique_positives)} prompts for {target_model_name}") - import comfy.model_management as mm_batch - mm_batch.load_models_gpu([patched_clip.patcher]) clip_skip = conf.get("clip_skip", 0) try: - for prompt in model_unique_positives: - # Check for interrupt before each prompt encoding - if mm.processing_interrupted(): - print(f"\n[GridTester] 🛑 INTERRUPTED during positive encoding - Stopping all encoding") - raise InterruptProcessingException() + with pinned_clip(patched_clip, model_unique_positives + model_unique_negatives): + for prompt in model_unique_positives: + # Check for interrupt before each prompt encoding + if mm.processing_interrupted(): + print(f"\n[GridTester] 🛑 INTERRUPTED during positive encoding - Stopping all encoding") + raise InterruptProcessingException() - if prompt not in conditioning_cache["positive"]: - conditioning_cache["positive"][prompt] = encode_prompt_with_combinators(patched_clip, prompt, clip_skip) + if prompt not in conditioning_cache["positive"]: + conditioning_cache["positive"][prompt] = encode_prompt_with_combinators(patched_clip, prompt, clip_skip) - for prompt in model_unique_negatives: - # Check for interrupt before each prompt encoding - if mm.processing_interrupted(): - print(f"\n[GridTester] 🛑 INTERRUPTED during negative encoding - Stopping all encoding") - raise InterruptProcessingException() + for prompt in model_unique_negatives: + # Check for interrupt before each prompt encoding + if mm.processing_interrupted(): + print(f"\n[GridTester] 🛑 INTERRUPTED during negative encoding - Stopping all encoding") + raise InterruptProcessingException() - if prompt not in conditioning_cache["negative"]: - conditioning_cache["negative"][prompt] = encode_prompt_with_combinators(patched_clip, prompt, clip_skip) + if prompt not in conditioning_cache["negative"]: + conditioning_cache["negative"][prompt] = encode_prompt_with_combinators(patched_clip, prompt, clip_skip) except InterruptProcessingException: print(f"\n[GridTester] 🛑 INTERRUPTED during encoding - Stopping all jobs") @@ -2773,22 +2773,21 @@ def _preencode_all_conditionings( else: patched_model, patched_clip = loaded_model, loaded_clip - # Force CLIP onto GPU for encoding - import comfy.model_management as mm_enc - mm_enc.load_models_gpu([patched_clip.patcher]) - clip_skip = first_conf.get("clip_skip", 0) serialized_positive = {} serialized_negative = {} - for prompt in unique_positives: - cond = encode_prompt_with_combinators(patched_clip, prompt, clip_skip) - serialized_positive[prompt] = _conditioning_to_serializable(cond) + # One encoder load for the whole group (pinned_clip); ComfyUI would otherwise + # re-prepare the text encoder on every encode. + with pinned_clip(patched_clip, list(unique_positives) + list(unique_negatives)): + for prompt in unique_positives: + cond = encode_prompt_with_combinators(patched_clip, prompt, clip_skip) + serialized_positive[prompt] = _conditioning_to_serializable(cond) - for prompt in unique_negatives: - cond = encode_prompt_with_combinators(patched_clip, prompt, clip_skip) - serialized_negative[prompt] = _conditioning_to_serializable(cond) + for prompt in unique_negatives: + cond = encode_prompt_with_combinators(patched_clip, prompt, clip_skip) + serialized_negative[prompt] = _conditioning_to_serializable(cond) encoded_all[group_key] = { "positive": serialized_positive, @@ -3132,23 +3131,20 @@ def _run_distributed_generation( cached_lora_key = current_lora_string # --- Prompt Encoding --- - if optional_positive: - final_positive = optional_positive - else: - if actual_positive_prompt not in conditioning_cache["positive"]: - clip_skip = conf.get("clip_skip", 0) - conditioning_cache["positive"][actual_positive_prompt] = \ - encode_prompt_with_combinators(patched_clip, actual_positive_prompt, clip_skip) - final_positive = conditioning_cache["positive"][actual_positive_prompt] - - if optional_negative: - final_negative = optional_negative - else: - if actual_negative_prompt not in conditioning_cache["negative"]: - clip_skip = conf.get("clip_skip", 0) - conditioning_cache["negative"][actual_negative_prompt] = \ - encode_prompt_with_combinators(patched_clip, actual_negative_prompt, clip_skip) - final_negative = conditioning_cache["negative"][actual_negative_prompt] + # One encoder load covers both prompts (pinned_clip); nothing is loaded when both are cached. + need_pos = not optional_positive and actual_positive_prompt not in conditioning_cache["positive"] + need_neg = not optional_negative and actual_negative_prompt not in conditioning_cache["negative"] + if need_pos or need_neg: + clip_skip = conf.get("clip_skip", 0) + with pinned_clip(patched_clip): + if need_pos: + conditioning_cache["positive"][actual_positive_prompt] = \ + encode_prompt_with_combinators(patched_clip, actual_positive_prompt, clip_skip) + if need_neg: + conditioning_cache["negative"][actual_negative_prompt] = \ + encode_prompt_with_combinators(patched_clip, actual_negative_prompt, clip_skip) + final_positive = optional_positive if optional_positive else conditioning_cache["positive"][actual_positive_prompt] + final_negative = optional_negative if optional_negative else conditioning_cache["negative"][actual_negative_prompt] # --- Create Latent --- latent_channels = get_latent_channels(loaded_model, optional_latent) diff --git a/tests/conftest.py b/tests/conftest.py index 20bb524..62950d5 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -100,6 +100,8 @@ _GEN_ORCH_DEPS = { "batch_encoding": { "batch_encode_prompts": lambda *a, **kw: None, "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), }, "manifest_utils": { "load_existing_manifest": lambda *a, **kw: {}, diff --git a/tests/test_pinned_clip.py b/tests/test_pinned_clip.py new file mode 100644 index 0000000..1adfcfa --- /dev/null +++ b/tests/test_pinned_clip.py @@ -0,0 +1,200 @@ +"""The text encoder is prepared once per encode loop, not once per prompt. + +REGRESSION (2026-11): comfy.sd.CLIP.encode_from_tokens() calls load_model() -> +model_management.load_models_gpu() on EVERY encode, and ComfyUI's dynamic VRAM +loader re-runs the patcher's load() walk even for a resident model (the repeated +"Model MOSSImageTEModel_ prepared for dynamic VRAM loading" log line, N prompts -> +N+1 lines). For the 7 GB Qwen3.5-9B encoder that was ~0.9 s per prompt, which made +USCG's pre-encode of 60 prompts take about a minute. pinned_clip() loads the +encoder once and shadows load_model for the duration of the loop. +""" +import contextlib +import importlib.util +import os +import sys +import types + +import pytest + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +# tests/conftest.py registers a STUB `batch_encoding` for the orchestrator's relative +# imports, so load the real module file under a private name (same trick as +# test_latent_geometry.py does for model_loader). +_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 not hasattr(_mm, "processing_interrupted"): + _mm.processing_interrupted = lambda: False +if _mm is not None and "comfy" in sys.modules and not hasattr(sys.modules["comfy"], "model_management"): + sys.modules["comfy"].model_management = _mm # `import comfy.model_management` on stubs does not wire the attribute +_spec = importlib.util.spec_from_file_location(f"{_PKG}.batch_encoding_real", 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"): + # torch stub (no real torch in this interpreter): the fast path only needs these + be.torch = types.SimpleNamespace(no_grad=contextlib.nullcontext, + cuda=types.SimpleNamespace(is_available=lambda: False)) + + +# ---- fakes ----------------------------------------------------------------- + +class FakeCLIP: + """Mimics comfy.sd.CLIP: encode_from_tokens() calls self.load_model() every time.""" + + def __init__(self): + self.patcher = object() + self.load_calls = 0 + self.load_tokens = [] + self.encode_calls = 0 + self.cond_stage_model = types.SimpleNamespace() # no clip_layer -> clip_skip branch is skipped + + def tokenize(self, text): + return {"qwen": [[(len(text), 1.0)]]} + + def load_model(self, tokens={}): + self.load_calls += 1 + self.load_tokens.append(tokens) + return self.patcher + + def encode_from_tokens(self, tokens, return_pooled=False, return_dict=False): + self.load_model(tokens) + self.encode_calls += 1 + return {"cond": f"cond{self.encode_calls}", "pooled_output": "pooled"} + + +class OldCLIP(FakeCLIP): + """Ancient ComfyUI: load_model() took no arguments.""" + + def load_model(self): + self.load_calls += 1 + return self.patcher + + def encode_from_tokens(self, tokens, return_pooled=False, return_dict=False): + self.load_model() + self.encode_calls += 1 + return {"cond": "c", "pooled_output": "p"} + + +class MemoryCache: + def __init__(self): + self.store = {} + + def get(self, prompt, prompt_type): + return self.store.get((prompt_type, prompt)) + + def set(self, prompt, cond, prompt_type): + self.store[(prompt_type, prompt)] = cond + + +TEXTS = ["a cat", "a dog on a very long sunny beach with umbrellas", "sky", "tea", "a bird"] + + +# ---- pinned_clip ----------------------------------------------------------- + +def test_many_encodes_load_the_encoder_once(): + clip = FakeCLIP() + with be.pinned_clip(clip, TEXTS): + for t in TEXTS: + be._encode_single(clip, t) + assert clip.encode_calls == len(TEXTS) + assert clip.load_calls == 1, "encoder must be prepared once for the whole loop" + + +def test_shadow_is_removed_after_the_block(): + clip = FakeCLIP() + with be.pinned_clip(clip, TEXTS): + be._encode_single(clip, "x") + assert "load_model" not in vars(clip) + be._encode_single(clip, "y") + assert clip.load_calls == 2, "normal per-encode load_model behaviour is back" + + +def test_shadow_is_removed_on_exception(): + clip = FakeCLIP() + with pytest.raises(RuntimeError): + with be.pinned_clip(clip, TEXTS): + raise RuntimeError("interrupted") + assert "load_model" not in vars(clip) + be._encode_single(clip, "y") + assert clip.load_calls == 2 + + +def test_nested_blocks_restore_the_outer_shadow(): + clip = FakeCLIP() + with be.pinned_clip(clip, TEXTS): + outer = vars(clip)["load_model"] + with be.pinned_clip(clip): + be._encode_single(clip, "inner") + assert vars(clip)["load_model"] is outer + be._encode_single(clip, "outer") + assert "load_model" not in vars(clip) + # the inner block's load_model call hits the outer shadow, so the encoder is + # loaded exactly once for the whole nest and never per encode + assert clip.load_calls == 1 + assert clip.encode_calls == 2 + + +def test_single_load_is_sized_by_the_longest_prompt(): + clip = FakeCLIP() + with be.pinned_clip(clip, TEXTS): + pass + longest = max(TEXTS, key=len) + assert clip.load_tokens == [clip.tokenize(longest)] + + +def test_none_and_non_clip_pass_through(): + with be.pinned_clip(None) as c: + assert c is None + plain = object() + with be.pinned_clip(plain) as c: + assert c is plain + + +def test_old_load_model_signature_still_works(): + clip = OldCLIP() + with be.pinned_clip(clip, TEXTS): + for t in TEXTS: + be._encode_single(clip, t) + assert clip.load_calls == 1 + assert "load_model" not in vars(clip) + + +# ---- batch_encode_with_cache / encode_pair_cached -------------------------- + +def test_batch_encode_with_cache_prepares_the_encoder_once(capsys): + clip = FakeCLIP() + prompts = [f"prompt number {i} with some words" for i in range(10)] + results = be.batch_encode_with_cache(clip, prompts, MemoryCache(), prompt_type="positive", batch_size=4) + assert clip.encode_calls == 10 + assert clip.load_calls == 1 + assert set(results) == set(prompts) + cond = results[prompts[0]] + assert isinstance(cond, list) and cond[0][0] == "cond1" and cond[0][1] == {"pooled_output": "pooled"} + out = capsys.readouterr().out + assert "10 prompts in" in out and "ms/prompt" in out + assert "load_model" not in vars(clip) + + +def test_batch_encode_with_cache_skips_the_encoder_on_full_cache_hit(): + clip = FakeCLIP() + cache = MemoryCache() + for p in TEXTS: + cache.set(p, [["cached", {}]], "positive") + results = be.batch_encode_with_cache(clip, TEXTS, cache, prompt_type="positive") + assert clip.load_calls == 0 and clip.encode_calls == 0 + assert all(results[p] == [["cached", {}]] for p in TEXTS) + + +def test_encode_pair_cached_loads_once_and_never_on_hits(): + clip = FakeCLIP() + cache = {"positive": {}, "negative": {}} + pos, neg = be.encode_pair_cached(clip, cache, "a cat", "blurry") + assert clip.load_calls == 1 and clip.encode_calls == 2 + assert pos is cache["positive"]["a cat"] and neg is cache["negative"]["blurry"] + pos2, neg2 = be.encode_pair_cached(clip, cache, "a cat", "blurry") + assert (pos2, neg2) == (pos, neg) + assert clip.load_calls == 1, "both cached -> encoder not touched" + be.encode_pair_cached(clip, cache, "a dog", "blurry") + assert clip.load_calls == 2 and clip.encode_calls == 3, "only the missing half is encoded" diff --git a/upscale_runner.py b/upscale_runner.py index e45d1db..744e38f 100644 --- a/upscale_runner.py +++ b/upscale_runner.py @@ -16,7 +16,7 @@ import numpy as np from PIL import Image as PILImage from .model_loader import load_checkpoint, load_vae_by_name -from .batch_encoding import encode_prompt_with_combinators +from .batch_encoding import encode_prompt_with_combinators, encode_pair_cached from .image_generation import upscale_image, decode_latent_with_vae, create_image_metadata from .manifest_utils import save_manifest @@ -248,13 +248,10 @@ def _run_upscale_thread(job, target_items, upscale_config, meta, manifest_data, neg_prompt = item.get("negative", meta.get("negative", "")) clip_skip = item.get("clip_skip", 0) - if pos_prompt not in conditioning_cache["positive"]: - conditioning_cache["positive"][pos_prompt] = encode_prompt_with_combinators(patched_clip, pos_prompt, clip_skip) - if neg_prompt not in conditioning_cache["negative"]: - conditioning_cache["negative"][neg_prompt] = encode_prompt_with_combinators(patched_clip, neg_prompt, clip_skip) - - final_positive = conditioning_cache["positive"][pos_prompt] - final_negative = conditioning_cache["negative"][neg_prompt] + # One encoder load for the pair (ComfyUI re-prepares the TE on every encode otherwise) + final_positive, final_negative = encode_pair_cached( + patched_clip, conditioning_cache, pos_prompt, neg_prompt, clip_skip + ) # HiRes prompt adjustment hires_positive_cond = final_positive