From f0f121ddf2218ec53eefcfc86ea12c243c0488c7 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Sat, 16 May 2026 14:29:53 +0300 Subject: [PATCH] Fix FunPackAdvisorLLM: thinking tokens, semantic gate, missing import - Strip ... blocks in decode so Qwen3 chain-of-thought doesn't pollute the repaired prompt and cause validation to reject it - Pass enable_thinking kwarg to apply_chat_template (with TypeError fallback for models that don't support it); propagate thinking flag through tokenize -> generate so the token budget is expanded by 2048 when thinking mode is active, giving room for reasoning + response - Return 1.0 from _v2_text_semantic_similarity when the clip has no encode_from_tokens_scheduled - generation-only clips can't compute similarity so they shouldn't be penalized with a 0.0 semantic gate - Add FunPackAdvisorLLM to standalone (non-package) imports in __init__.py so tests can run without the full ComfyUI runtime - Update three stale advisor tests: 5-value _v2_prompt_advisor unpack, "Encoded as:" prefix change, 6-return-value unpacking after negative conditioning was removed Co-Authored-By: Claude Sonnet 4.6 --- __init__.py | 1 + conditioning.py | 29 +++++++++++++++++++++++------ tests/test_v2_category_alignment.py | 12 ++++++------ 3 files changed, 30 insertions(+), 12 deletions(-) diff --git a/__init__.py b/__init__.py index f57c26c8..b6fe3953 100644 --- a/__init__.py +++ b/__init__.py @@ -23,6 +23,7 @@ if __package__: else: # Standalone tests may not have the full ComfyUI/CUDA runtime loaded. from conditioning import ( + FunPackAdvisorLLM, FunPackLorebookEnhancer, FunPackPromptCombiner, FunPackPromptEnhancer, diff --git a/conditioning.py b/conditioning.py index 1e0a58e8..56fe6544 100644 --- a/conditioning.py +++ b/conditioning.py @@ -7010,6 +7010,9 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner): def _v2_text_semantic_similarity(self, clip, left, right, encode_cache=None): if clip is None: return 0.0 + if not hasattr(clip, "encode_from_tokens_scheduled"): + # Generation-only clip (e.g. FunPackAdvisorLLM) can't encode; skip semantic gate. + return 1.0 left = self._v2_clean_phrase_text(left) right = self._v2_clean_phrase_text(right) if not left or not right: @@ -11164,6 +11167,7 @@ class _FunPackAdvisorLLMWrapper: ) def tokenize(self, text, system_prompt=None, skip_template=False, **kwargs): + thinking = bool(kwargs.get("thinking", True)) if skip_template: # Text already has role tokens baked in (e.g. from TextGenerateLTX2Prompt). input_ids = self._tokenizer.encode( @@ -11177,21 +11181,31 @@ class _FunPackAdvisorLLMWrapper: ] else: messages = [{"role": "user", "content": str(text)}] - input_ids = self._tokenizer.apply_chat_template( - messages, tokenize=True, add_generation_prompt=True, return_tensors="pt" - ).to(self._device) - return {"input_ids": input_ids, "prompt_length": int(input_ids.shape[1])} + try: + input_ids = self._tokenizer.apply_chat_template( + messages, tokenize=True, add_generation_prompt=True, return_tensors="pt", + enable_thinking=thinking, + ).to(self._device) + except TypeError: + input_ids = self._tokenizer.apply_chat_template( + messages, tokenize=True, add_generation_prompt=True, return_tensors="pt" + ).to(self._device) + return {"input_ids": input_ids, "prompt_length": int(input_ids.shape[1]), "thinking": thinking} def generate(self, tokens, do_sample=True, max_length=800, temperature=0.7, top_k=50, top_p=0.92, min_p=0.0, repetition_penalty=1.3, no_repeat_ngram_size=5, presence_penalty=0.0, seed=None, **kwargs): input_ids = tokens["input_ids"] prompt_length = tokens.get("prompt_length", int(input_ids.shape[1])) + thinking = tokens.get("thinking", False) if seed: torch.manual_seed(int(seed)) + # When thinking mode is active, add extra budget for the reasoning chain so the + # actual response is not crowded out by thinking tokens. + effective_max = int(max_length) + 2048 if thinking else int(max_length) required = dict( do_sample=do_sample, - max_new_tokens=int(max_length), + max_new_tokens=effective_max, temperature=float(temperature), top_k=int(top_k), top_p=float(top_p), @@ -11225,7 +11239,10 @@ class _FunPackAdvisorLLMWrapper: output_ids = generated["output_ids"] prompt_length = generated.get("prompt_length", 0) new_tokens = output_ids[0][prompt_length:] - return self._tokenizer.decode(new_tokens, skip_special_tokens=skip_special_tokens) + text = self._tokenizer.decode(new_tokens, skip_special_tokens=skip_special_tokens) + # Strip chain-of-thought thinking blocks emitted by models like Qwen3. + text = re.sub(r"[\s\S]*?", "", text, flags=re.IGNORECASE).strip() + return text _FUNPACK_ADVISOR_LLM_CACHE = {} diff --git a/tests/test_v2_category_alignment.py b/tests/test_v2_category_alignment.py index 8513b898..6b99360b 100644 --- a/tests/test_v2_category_alignment.py +++ b/tests/test_v2_category_alignment.py @@ -1471,7 +1471,7 @@ def test_refiner_v2_advisor_uses_explicit_system_prompt_previous_prompt_thinking "REPAIRED_PROMPT: person smoking, smoke trails drifting upward" ) - prompt, status, diagnostic, applied = refiner._v2_prompt_advisor( + prompt, status, diagnostic, applied, _ = refiner._v2_prompt_advisor( clip, "Diagnostics", "person smoking", @@ -1489,8 +1489,8 @@ def test_refiner_v2_advisor_uses_explicit_system_prompt_previous_prompt_thinking assert applied is False assert "diagnostics only" in status assert diagnostic == "add clearer smoke motion." - assert "analyzing a video generation prompt" in advisor_prompt - assert "Previous prompt (what caused the feedback): old encoded prompt" in advisor_prompt + assert "Prompt to analyze: person smoking" in advisor_prompt + assert "Previous prompt (what caused this rating): old encoded prompt" in advisor_prompt assert kwargs["image"] is image assert kwargs["thinking"] is True assert clip.generate_kwargs["seed"] == 123 @@ -1542,7 +1542,7 @@ def test_refiner_v2_advisor_repair_applies_validated_generated_prompt(tmp_path): "REPAIRED_PROMPT: person smoking, smoke trails drifting upward" ) - _, status, training_info, _, _, encoded_prompts = refiner.refine_v2( + _, status, training_info, _, encoded_prompts, _ = refiner.refine_v2( "person smoking", clip, "Missing details", @@ -1554,8 +1554,8 @@ def test_refiner_v2_advisor_repair_applies_validated_generated_prompt(tmp_path): state = json.loads(state_path.read_text(encoding="utf-8")) assert "Advisor: applied repair" in status - assert "Encoded: advisor repaired prompt" in training_info - assert encoded_prompts == "Positive prompt: person smoking, smoke trails drifting upward\n\nNegative prompt: " + assert "Encoded as: advisor repaired prompt" in training_info + assert "Positive prompt: person smoking, smoke trails drifting upward" in encoded_prompts assert state["last_run"]["encoded_prompt"] == "person smoking, smoke trails drifting upward" assert state["last_run"]["advisor"]["applied"] is True