From 55b6a9b11dffecdd65a3ccd5eb6a1b3a178c96dc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jukka=20Sepp=C3=A4nen?= <40791699+kijai@users.noreply.github.com> Date: Fri, 14 Aug 2026 22:17:18 +0300 Subject: [PATCH] Fix thinking handling (#15611) --- comfy/text_encoders/gemma4.py | 9 ++++++--- comfy_extras/nodes_textgen.py | 10 ++++------ 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/comfy/text_encoders/gemma4.py b/comfy/text_encoders/gemma4.py index 61bc3a3f0..da116ab27 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 @@ -1624,11 +1625,13 @@ 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 / tags + # Only a close that ends a thought channel becomes : generation primed with + # another channel leaves its opener in the prompt, so its close is not reasoning. + text = re.sub(r"<\|channel>thought\n(.*?)", r"\n\1", text, flags=re.DOTALL) text = text.replace("<|channel>thought\n", "\n") - text = text.replace("", "") # Strip remaining special tokens - text = text.replace("", "").replace("", "").strip() + text = re.sub(r"<\|channel>\w*\n?||<\|turn>\w*\n?|", "", text) + text = text.replace("", "").strip() return text diff --git a/comfy_extras/nodes_textgen.py b/comfy_extras/nodes_textgen.py index 40004652c..ef5858a7d 100644 --- a/comfy_extras/nodes_textgen.py +++ b/comfy_extras/nodes_textgen.py @@ -256,12 +256,10 @@ class TextGenerateLTX2Prompt(TextGenerate): out = super().execute(clip, formatted_prompt, max_length, sampling_mode, image=image, thinking=thinking, use_default_template=use_default_template, video=video, audio=audio) - text = out.args[0] - text = re.sub(r".*?", "", text, flags=re.DOTALL) - if "" in text: # unclosed/truncated reasoning: keep what follows the last close - text = text.rsplit("", 1)[-1] - text = re.sub(r"|<\|channel>\w*\n?||<\|turn>\w*\n?", "", text).strip() - return io.NodeOutput(text) + # Drop reasoning, including a block left unclosed by max_length. Both system prompts ask + # for the original prompt back when there is nothing to give; empty conditions on nothing. + text = re.sub(r".*?(?:|$)", "", out.args[0], flags=re.DOTALL).strip() + return io.NodeOutput(text or prompt) class TextgenExtension(ComfyExtension):