Fix FunPackAdvisorLLM: thinking tokens, semantic gate, missing import

- Strip <think>...</think> 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 <noreply@anthropic.com>
This commit is contained in:
digital-garbage
2026-05-16 14:36:49 +03:00
co-authored by Claude Sonnet 4.6
parent 8f7c61c6ce
commit f0f121ddf2
3 changed files with 30 additions and 12 deletions
+1
View File
@@ -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,
+23 -6
View File
@@ -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"<think>[\s\S]*?</think>", "", text, flags=re.IGNORECASE).strip()
return text
_FUNPACK_ADVISOR_LLM_CACHE = {}
+6 -6
View File
@@ -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