From f075294467b62edb7607fe845da4d856a93c4281 Mon Sep 17 00:00:00 2001 From: Glary-Bot Date: Tue, 4 Aug 2026 06:17:46 +0000 Subject: [PATCH] Fix Gemma4 ignoring thinking=false PR #14304 started priming an empty, already-closed thought block (<|channel>thought\n) on the model turn when thinking is off. Google's canonical Gemma 4 chat template never does this: <|think|> in the system turn is the only thought channel switch, and the generation prompt ends at <|turn>model\n. Priming a closed block instead cues the model into reasoning, which then leaks into the answer untagged. The rendered template now matches google/gemma-4-e4b-it chat_template.jinja byte for byte for thinking on and off. --- comfy/text_encoders/gemma4.py | 8 ++-- tests-unit/comfy_test/gemma4_template_test.py | 41 +++++++++++++++++++ 2 files changed, 45 insertions(+), 4 deletions(-) create mode 100644 tests-unit/comfy_test/gemma4_template_test.py diff --git a/comfy/text_encoders/gemma4.py b/comfy/text_encoders/gemma4.py index 5163c1676..93fac702f 100644 --- a/comfy/text_encoders/gemma4.py +++ b/comfy/text_encoders/gemma4.py @@ -1310,7 +1310,9 @@ class Gemma4_Tokenizer(): if llama_template is not None: llama_text = llama_template.format(text) else: - # Build template from modalities present + # Build template from modalities present. + # <|think|> in the system turn is the only thought channel switch. Priming a closed + # thought block on the model turn does not disable it, the model reasons inline instead. system = "<|turn>system\n<|think|>\n\n" if thinking else "" media = "" if len(images) > 0: @@ -1333,9 +1335,7 @@ class Gemma4_Tokenizer(): num_samples = int(waveform.shape[-1] * 16000 / sample_rate) if sample_rate != 16000 else waveform.shape[-1] n_audio_tokens = self._audio_token_count(num_samples) media += "<|audio>" + "<|audio|>" * n_audio_tokens + "" - # Non-thinking mode primes an empty thought channel so the model answers directly. - model_open = "" if thinking else "<|channel>thought\n" - llama_text = f"{system}<|turn>user\n{text}{media}\n<|turn>model\n{model_open}" + llama_text = f"{system}<|turn>user\n{text}{media}\n<|turn>model\n" text_tokens = super().tokenize_with_weights(llama_text, return_word_ids) diff --git a/tests-unit/comfy_test/gemma4_template_test.py b/tests-unit/comfy_test/gemma4_template_test.py new file mode 100644 index 000000000..cd82ea09c --- /dev/null +++ b/tests-unit/comfy_test/gemma4_template_test.py @@ -0,0 +1,41 @@ +from comfy.cli_args import args + +args.cpu = True # importing comfy.model_management probes the torch device at import time + +from comfy.text_encoders.gemma4 import Gemma4_Tokenizer + +PROMPT = "describe a cute anime girl with fennec ears" + + +class _CaptureTemplate: + """Stands in for SDTokenizer.tokenize_with_weights so the built template is checked without model files.""" + def tokenize_with_weights(self, text, return_word_ids=False, **kwargs): + self.llama_text = text + return {} + + +class _Gemma4TemplateProbe(Gemma4_Tokenizer, _CaptureTemplate): + pass + + +def build_template(**kwargs): + probe = _Gemma4TemplateProbe() + probe.tokenize_with_weights(PROMPT, **kwargs) + return probe.llama_text + + +def test_thinking_disabled_does_not_prime_a_thought_channel(): + template = build_template(skip_template=False, thinking=False) + assert template == f"<|turn>user\n{PROMPT}\n<|turn>model\n" + assert "channel" not in template + assert "<|think|>" not in template + + +def test_thinking_enabled_asks_for_the_thought_channel(): + template = build_template(skip_template=False, thinking=True) + assert template == f"<|turn>system\n<|think|>\n\n<|turn>user\n{PROMPT}\n<|turn>model\n" + + +def test_skip_template_passes_text_through_unchanged(): + assert build_template(skip_template=True, thinking=False) == PROMPT + assert build_template(skip_template=True, thinking=True) == PROMPT