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<channel|> 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

@@ -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<turn|>\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 + "<audio|>"
-                llama_text = f"{system}<|turn>user\n{text}{media}<turn|>\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<channel|>" if self.prime_empty_thought and not thinking else ""
+                llama_text = f"{system}<|turn>user\n{text}{media}<turn|>\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).

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<channel|>"
+
+# 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<turn|>\n<|turn>user\n{PROMPT}<turn|>\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}<turn|>\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<turn|>\n<|turn>user\n{PROMPT}<turn|>\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}<turn|>\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