From f3cfb0866dcd320266562a5491ea3d07b45f2819 Mon Sep 17 00:00:00 2001
From: bymyself <cbyrne@comfy.org>
Date: Tue, 11 Aug 2026 17:45:13 -0700
Subject: [PATCH] Make Gemma 4 channel close markers channel aware

decode() rewrote every <channel|> into </think>, including the close of a
channel that is not the thought channel. Only close a thought channel that was
actually opened in the decoded text; other channel markers are left alone for
the caller to strip.
---
 comfy/text_encoders/gemma4.py                 |  7 ++-
 .../comfy_test/test_gemma4_tokenizer.py       | 45 +++++++++++++++++++
 2 files changed, 50 insertions(+), 2 deletions(-)
 create mode 100644 tests-unit/comfy_test/test_gemma4_tokenizer.py

diff --git a/comfy/text_encoders/gemma4.py b/comfy/text_encoders/gemma4.py
index 5b0b968c9..2d1e45a75 100644
--- a/comfy/text_encoders/gemma4.py
+++ b/comfy/text_encoders/gemma4.py
@@ -6,6 +6,7 @@ import numpy as np
 from tokenizers import Tokenizer
 from dataclasses import dataclass
 import math
+import re
 
 from comfy import sd1_clip
 import comfy.model_management
@@ -1401,9 +1402,11 @@ class Gemma4SDTokenizer(Gemma4_Tokenizer, sd1_clip.SDTokenizer):
 
     def decode(self, token_ids, **kwargs):
         text = super().decode(token_ids, skip_special_tokens=False)
-        # Translate thinking channel markers to standard <think>/</think> tags
+        # Translate thinking channel markers to standard <think>/</think> tags. Only a close
+        # that ends a thought channel becomes </think>; generation primed with another channel
+        # leaves its opening marker in the prompt, so its close is not reasoning.
+        text = re.sub(r"<\|channel>thought\n(.*?)<channel\|>", r"<think>\n\1</think>", text, flags=re.DOTALL)
         text = text.replace("<|channel>thought\n", "<think>\n")
-        text = text.replace("<channel|>", "</think>")
         # Strip remaining special tokens
         text = text.replace("<turn|>", "").replace("<eos>", "").strip()
         return text
diff --git a/tests-unit/comfy_test/test_gemma4_tokenizer.py b/tests-unit/comfy_test/test_gemma4_tokenizer.py
new file mode 100644
index 000000000..c4d86f867
--- /dev/null
+++ b/tests-unit/comfy_test/test_gemma4_tokenizer.py
@@ -0,0 +1,45 @@
+import torch
+
+from comfy.cli_args import args as cli_args
+
+if not torch.cuda.is_available():
+    cli_args.cpu = True
+
+from comfy.text_encoders.gemma4 import Gemma4SDTokenizer  # noqa: E402
+
+
+class _StubTokenizer:
+    """Returns a canned decode so the marker translation can be tested without model files."""
+    def __init__(self, text):
+        self.text = text
+
+    def decode(self, token_ids, skip_special_tokens=False):
+        return self.text
+
+
+def decode(text):
+    tokenizer = Gemma4SDTokenizer.__new__(Gemma4SDTokenizer)
+    tokenizer.tokenizer = _StubTokenizer(text)
+    return tokenizer.decode([])
+
+
+class TestGemma4Decode:
+    def test_thought_channel_becomes_think_tags(self):
+        assert decode("<|channel>thought\nreasoning<channel|>the answer") == "<think>\nreasoning</think>the answer"
+
+    def test_primed_empty_thought_channel(self):
+        assert decode("<|channel>thought\n<channel|>the answer") == "<think>\n</think>the answer"
+
+    def test_unclosed_thought_channel(self):
+        assert decode("<|channel>thought\nreasoning") == "<think>\nreasoning"
+
+    def test_other_channel_close_is_not_reasoning(self):
+        # Non-thinking LTX2 prompt enhancement primes a "final" channel, so only its close is
+        # generated. Turning that into </think> made the whole answer look like reasoning.
+        assert decode("the answer<channel|>") == "the answer<channel|>"
+
+    def test_other_channel_kept_after_a_thought_channel(self):
+        assert decode("<|channel>thought\nreasoning<channel|><|channel>final\nthe answer<channel|>") == "<think>\nreasoning</think><|channel>final\nthe answer<channel|>"
+
+    def test_turn_and_eos_are_stripped(self):
+        assert decode("the answer<turn|><eos>") == "the answer"