RomanNum3ral
/
ComfyUI
mirror of https://github.com/Comfy-Org/ComfyUI.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Fix thinking handling (#15611)

This commit is contained in:
Jukka Seppänen 2026-08-14 22:17:18 +03:00 committed by GitHub
parent 1c6d8d45b3
commit 55b6a9b11d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 10 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
comfy/text_encoders/gemma4.py
View File

@ -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 <think>/</think> tags
# Only a close that ends a thought channel becomes </think>: generation primed with
# another channel leaves its opener 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()
text = re.sub(r"<\|channel>\w*\n?|<channel\|>|<\|turn>\w*\n?|<turn\|>", "", text)
text = text.replace("<eos>", "").strip()
return text

10
comfy_extras/nodes_textgen.py
View File

@ -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"<think>.*?</think>", "", text, flags=re.DOTALL)
if "</think>" in text: # unclosed/truncated reasoning: keep what follows the last close
text = text.rsplit("</think>", 1)[-1]
text = re.sub(r"</?think>|<\|channel>\w*\n?|<channel\|>|<\|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"<think>.*?(?:</think>|$)", "", out.args[0], flags=re.DOTALL).strip()
return io.NodeOutput(text or prompt)
class TextgenExtension(ComfyExtension):