From b48a6a2ead5533ee1ebbf0a0e6e289def26e886f Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Wed, 23 Sep 2026 23:05:21 +0300 Subject: [PATCH] fix: enhancer instructions and image never reached ComfyUI text encoders ComfyUI tokenizers drop an unknown system_prompt kwarg, and Gemma skips its chat template by default. Instructions now go in the user turn of the model's own template, which also places the image. Co-Authored-By: Claude Opus 5.5 --- conditioning.py | 77 +++++++---------------------------- tests/test_prompt_enhancer.py | 25 ++++++++++++ 2 files changed, 40 insertions(+), 62 deletions(-) diff --git a/conditioning.py b/conditioning.py index 036bb4ad..8bbc7fc1 100644 --- a/conditioning.py +++ b/conditioning.py @@ -8793,35 +8793,6 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): lines.append(f" Axis {axis}: {role} | + {pos_str} | - {neg_str}") return "\n".join(lines) - def _v2_find_chat_tokenizer(self, clip): - """BFS through CLIP wrapper hierarchy to find a HuggingFace tokenizer - that exposes apply_chat_template. Returns None if not found.""" - seen = set() - queue = [clip] - attrs = ("tokenizer", "processor", "cond_stage_model", "patcher", "model", "transformer") - depth = 0 - while queue and depth < 5: - next_q = [] - for obj in queue: - if obj is None or id(obj) in seen: - continue - seen.add(id(obj)) - try: - if hasattr(obj, "apply_chat_template"): - return obj - for attr in attrs: - child = getattr(obj, attr, None) - if child is not None: - next_q.append(child) - tkns = getattr(obj, "tokenizers", None) - if isinstance(tkns, dict): - next_q.extend(tkns.values()) - except Exception: - continue - queue = next_q - depth += 1 - return None - def _v2_advisor_analysis_prompt( self, prompt, @@ -8950,41 +8921,23 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): max_length = max(int(min_floor), int(max_length or 800)) _started = time.time() try: - # Tokenization: three layers in order of preference. - # Layer 1: wrapper natively accepts system_prompt kwarg. - tokens = None - try: - tokens = clip.tokenize(str(user_prompt), image=image, min_length=1, thinking=bool(thinking), system_prompt=str(system_prompt)) - except TypeError: + # No ComfyUI tokenizer takes a system prompt: they swallow unknown kwargs, so + # passing one drops it silently. Only FunPack's own Advisor LLM honours it. + # Everything else gets the instructions as the first paragraph of the user turn, + # inside the model's own chat template — which is also what places the image + # marker, and how Gemma's official template carries a system prompt anyway. + if isinstance(clip, _FunPackAdvisorLLMWrapper): + tokens = clip.tokenize(str(user_prompt), system_prompt=str(system_prompt), + thinking=bool(thinking)) + else: + merged = f"{system_prompt}\n\n{user_prompt}" if str(system_prompt).strip() \ + else str(user_prompt) try: - tokens = clip.tokenize(str(user_prompt), image=image, system_prompt=str(system_prompt)) + # min_length=1: Gemma pads to 1024 on the left otherwise + tokens = clip.tokenize(merged, image=image, skip_template=False, + min_length=1, thinking=bool(thinking)) except TypeError: - pass - if tokens is None: - # Layer 2: apply chat template via underlying HuggingFace tokenizer. - chat_tokenizer = self._v2_find_chat_tokenizer(clip) - advisor_prompt = None - if chat_tokenizer is not None: - try: - messages = [ - {"role": "system", "content": str(system_prompt)}, - {"role": "user", "content": str(user_prompt)}, - ] - advisor_prompt = chat_tokenizer.apply_chat_template( - messages, tokenize=False, add_generation_prompt=True - ) - except Exception: - pass - if advisor_prompt is None: - # Layer 3: flat string with completion anchor. - advisor_prompt = str(system_prompt) + "\n\n" + str(user_prompt) + "\n\nOutput:" - try: - tokens = clip.tokenize(advisor_prompt, image=image, min_length=1, thinking=bool(thinking)) - except TypeError: - try: - tokens = clip.tokenize(advisor_prompt, image=image) - except TypeError: - tokens = clip.tokenize(advisor_prompt) + tokens = clip.tokenize(merged, image=image) generate_kwargs = dict( do_sample=bool(do_sample), max_length=max_length, diff --git a/tests/test_prompt_enhancer.py b/tests/test_prompt_enhancer.py index 1607707b..3ac8667b 100644 --- a/tests/test_prompt_enhancer.py +++ b/tests/test_prompt_enhancer.py @@ -438,3 +438,28 @@ def test_a_fixed_enhancer_seed_overrides_the_run_seed(): src = _refine_source() assert '_enhance_seed = int(_enhance_sampling.pop("seed", 0) or 0) or seed' in src assert src.count("seed=_enhance_seed") == 2 + + +# ── how the instructions reach the model ───────────────────────────────────── + +def test_instructions_ride_in_the_models_own_template(studio): + # ComfyUI's Gemma tokenizer swallows unknown kwargs and defaults to skip_template=True, + # so a system_prompt kwarg vanished and the raw prompt went in with no chat turn and no + # image marker. The model's own template must be asked for, with the instructions in it. + clip = FakeClip() + studio._v2_enhance_prompt(clip, "a cat", "BE VIVID", image="img") + text, kwargs = clip.calls[0] + assert text.startswith("BE VIVID") and text.endswith("a cat") + assert kwargs["skip_template"] is False and kwargs["min_length"] == 1 + assert kwargs["image"] == "img" + assert "system_prompt" not in kwargs + + +def test_the_funpack_advisor_llm_still_gets_a_real_system_turn(studio): + wrapper = C._FunPackAdvisorLLMWrapper.__new__(C._FunPackAdvisorLLMWrapper) + seen = {} + wrapper.tokenize = lambda text, **kw: seen.update(text=text, **kw) or {} + wrapper.generate = lambda tokens, **kw: [1] + wrapper.decode = lambda ids, skip_special_tokens=True: "a vivid cat" + studio._v2_enhance_prompt(wrapper, "a cat", "BE VIVID") + assert seen["text"] == "a cat" and seen["system_prompt"] == "BE VIVID"