From c857b6c65736e45fd9a8d6618c78e712fe0bc96a Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Sun, 12 Apr 2026 21:25:19 +0300
Subject: [PATCH] Translate think tokens

---
 comfy_extras/nodes_textgen.py | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/comfy_extras/nodes_textgen.py b/comfy_extras/nodes_textgen.py
index b4f793f9a..0d4cf3a2b 100644
--- a/comfy_extras/nodes_textgen.py
+++ b/comfy_extras/nodes_textgen.py
@@ -71,7 +71,15 @@ class TextGenerate(io.ComfyNode):
             seed=seed
         )
 
-        generated_text = clip.decode(generated_ids, skip_special_tokens=True)
+        generated_text = clip.decode(generated_ids, skip_special_tokens=not thinking)
+
+        if thinking:
+            # Translate Gemma4 thinking channel markers to standard <think>/</think> tags
+            generated_text = generated_text.replace("<|channel>thought\n", "<think>\n")
+            generated_text = generated_text.replace("<channel|>", "</think>")
+            # Strip remaining special tokens
+            generated_text = generated_text.replace("<turn|>", "").replace("<eos>", "").strip()
+
         return io.NodeOutput(generated_text)