From a196f67682f609ee99e4f66432aecb8971d7ccf7 Mon Sep 17 00:00:00 2001 From: Glary-Bot Date: Tue, 4 Aug 2026 06:32:35 +0000 Subject: [PATCH] Only prime the empty thought block on Gemma4 variants that want it Google ships different canonical chat templates per variant: 12B/31B append <|channel>thought\n to the generation prompt when thinking is off, E2B/E4B do not. #14304 added it for every variant, which cues E2B/E4B into reasoning that then leaks into the answer untagged. Gate it on a tokenizer-owned prime_empty_thought flag so each variant renders its own template. --- comfy/text_encoders/gemma4.py | 18 ++++---- tests-unit/comfy_test/gemma4_template_test.py | 43 ++++++++++++------- 2 files changed, 38 insertions(+), 23 deletions(-) diff --git a/comfy/text_encoders/gemma4.py b/comfy/text_encoders/gemma4.py index 93fac702f..a4af2ea84 100644 --- a/comfy/text_encoders/gemma4.py +++ b/comfy/text_encoders/gemma4.py @@ -1183,6 +1183,7 @@ def _get_aspect_ratio_preserving_size(height, width, patch_size, max_patches, po class Gemma4_Tokenizer(): tokenizer_json_data = None + prime_empty_thought = False def state_dict(self): if self.tokenizer_json_data is not None: @@ -1310,9 +1311,7 @@ class Gemma4_Tokenizer(): if llama_template is not None: llama_text = llama_template.format(text) else: - # 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. + # Build template from modalities present system = "<|turn>system\n<|think|>\n\n" if thinking else "" media = "" if len(images) > 0: @@ -1335,7 +1334,9 @@ 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 + "" - llama_text = f"{system}<|turn>user\n{text}{media}\n<|turn>model\n" + # 12B/31B prime a closed thought block for non-thinking mode, E2B/E4B must not: it cues them into reasoning inline. + model_open = "<|channel>thought\n" if self.prime_empty_thought and not thinking else "" + llama_text = f"{system}<|turn>user\n{text}{media}\n<|turn>model\n{model_open}" text_tokens = super().tokenize_with_weights(llama_text, return_word_ids) @@ -1418,6 +1419,7 @@ class Gemma4Tokenizer(sd1_clip.SD1Tokenizer): class Gemma4UnifiedSDTokenizer(Gemma4SDTokenizer): """Encoder-free (gemma4_unified) audio: raw 16kHz waveform frames instead of mel spectrogram.""" embedding_size = 3840 + prime_empty_thought = True def _extract_audio_features(self, waveform, sample_rate): audio = self._resample_16k(waveform, sample_rate) @@ -1489,7 +1491,7 @@ def gemma4_te(dtype_llama=None, llama_quantization_metadata=None, model_class=No # Variants -def _make_variant(config_cls): +def _make_variant(config_cls, prime_empty_thought=False): audio = config_cls.audio_config is not None bases = (Gemma4AudioMixin, Gemma4Base) if audio else (Gemma4Base,) class Variant(*bases): @@ -1499,8 +1501,8 @@ def _make_variant(config_cls): if audio: self._init_audio(self.model.config, dtype, device, operations) embedding_size = config_cls.hidden_size - if embedding_size != Gemma4SDTokenizer.embedding_size: - tok_cls = type('T', (Gemma4SDTokenizer,), {'embedding_size': embedding_size}) + if embedding_size != Gemma4SDTokenizer.embedding_size or prime_empty_thought: + tok_cls = type('T', (Gemma4SDTokenizer,), {'embedding_size': embedding_size, 'prime_empty_thought': prime_empty_thought}) class Tokenizer(Gemma4Tokenizer): tokenizer_class = tok_cls Variant.tokenizer = Tokenizer @@ -1510,7 +1512,7 @@ def _make_variant(config_cls): Gemma4_E4B = _make_variant(Gemma4Config) Gemma4_E2B = _make_variant(Gemma4_E2B_Config) -Gemma4_31B = _make_variant(Gemma4_31B_Config) +Gemma4_31B = _make_variant(Gemma4_31B_Config, prime_empty_thought=True) # Gemma4 12B Unified: encoder-free multimodal, distinct base/tokenizer (not via _make_variant). diff --git a/tests-unit/comfy_test/gemma4_template_test.py b/tests-unit/comfy_test/gemma4_template_test.py index 104e2f0ae..69a2d9231 100644 --- a/tests-unit/comfy_test/gemma4_template_test.py +++ b/tests-unit/comfy_test/gemma4_template_test.py @@ -1,5 +1,6 @@ """Gemma4 chat template regression tests.""" +import pytest import torch from comfy.cli_args import args @@ -7,9 +8,15 @@ from comfy.cli_args import args if not torch.cuda.is_available(): args.cpu = True -from comfy.text_encoders.gemma4 import Gemma4_Tokenizer # noqa: E402 +import comfy.text_encoders.gemma4 as gemma4 # noqa: E402 PROMPT = "describe a cute anime girl with fennec ears" +THOUGHT_BLOCK = "<|channel>thought\n" + +# E2B/E4B and 12B/31B ship different canonical chat templates: only the latter prime a +# closed thought block when thinking is off. +NO_PRIMING = [gemma4.Gemma4_E2B, gemma4.Gemma4_E4B] +PRIMING = [gemma4.Gemma4_31B, gemma4.Gemma4_12B] class _CaptureTemplate: @@ -19,28 +26,34 @@ class _CaptureTemplate: return {} -class _Gemma4TemplateProbe(Gemma4_Tokenizer, _CaptureTemplate): - pass - - -def build_template(**kwargs): - probe = _Gemma4TemplateProbe() +def build_template(variant, **kwargs): + prime = variant.tokenizer.tokenizer_class.prime_empty_thought + probe = type("Probe", (gemma4.Gemma4_Tokenizer, _CaptureTemplate), {"prime_empty_thought": prime})() 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) +@pytest.mark.parametrize("variant", NO_PRIMING + PRIMING) +def test_thinking_enabled_only_asks_via_the_system_turn(variant): + template = build_template(variant, skip_template=False, thinking=True) + assert template == f"<|turn>system\n<|think|>\n\n<|turn>user\n{PROMPT}\n<|turn>model\n" + + +@pytest.mark.parametrize("variant", NO_PRIMING) +def test_thinking_disabled_does_not_prime_a_thought_channel(variant): + template = build_template(variant, 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" +@pytest.mark.parametrize("variant", PRIMING) +def test_thinking_disabled_primes_a_thought_channel(variant): + template = build_template(variant, skip_template=False, thinking=False) + assert template == f"<|turn>user\n{PROMPT}\n<|turn>model\n{THOUGHT_BLOCK}" -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 +@pytest.mark.parametrize("variant", NO_PRIMING + PRIMING) +@pytest.mark.parametrize("thinking", [False, True]) +def test_skip_template_passes_text_through_unchanged(variant, thinking): + assert build_template(variant, skip_template=True, thinking=thinking) == PROMPT