diff --git a/README.md b/README.md
index 57937b4..7658780 100644
--- a/README.md
+++ b/README.md
@@ -142,8 +142,8 @@ Every input includes an in-node tooltip.
| Output | Description |
| --- | --- |
| `RESPONSE` | Final model response with Qwen thinking blocks removed. |
-| `THINKING` | Extracted `...` reasoning when present in model output. |
-| `PERF` | llama.cpp `llama_perf_context_print` lines. |
+| `THINKING` | Extracted reasoning when present in model output. |
+| `PERF` | Parsed llama.cpp prompt and generation speed, for example `[ Prompt: 80.0 t/s | Generation: 76.6 t/s ]`. |
Both outputs include ComfyUI output tooltips.
diff --git a/llama_cli.py b/llama_cli.py
index db6480f..2da296d 100644
--- a/llama_cli.py
+++ b/llama_cli.py
@@ -13,11 +13,11 @@ import comfy.model_management
from .llama_binary import ensure_llama_cli_paths
-THINK_BLOCK_RE = re.compile(r"]*>.*?", flags=re.IGNORECASE | re.DOTALL)
-LLAMA_PERF_PREFIXES = (
- "llama_perf_context_print:",
- "common_perf_print:",
-)
+PROMPT_ECHO_END = "... (truncated)"
+PROMPT_PADDING = " " * 501
+PERF_RE = re.compile(r"\[\s*Prompt:\s*[^|\]]+\|\s*Generation:\s*[^\]]+\]")
+START_THINKING = "[Start thinking]"
+END_THINKING = "[End thinking]"
def tensor_to_temp_png(image) -> Path:
@@ -48,7 +48,7 @@ def _write_prompt_file(system_prompt: str, prompt: str) -> Path:
chunks = []
if system_prompt:
chunks.append(system_prompt.strip())
- chunks.append(prompt.strip())
+ chunks.append(prompt.strip() + PROMPT_PADDING)
return _write_temp_text_file("qwen-gguf-prompt-", "\n\n".join(chunks))
@@ -160,9 +160,7 @@ def run_llama_cli(
raise RuntimeError(
f"llama.cpp inference failed with exit code {result.returncode}:\n{stderr}"
)
- response, thinking = _extract_thinking(result.stdout)
- perf = _extract_perf(result.stderr)
- return response, thinking, perf
+ return _parse_response(result.stdout + "\n" + result.stderr)
def _stop_process(process: subprocess.Popen) -> None:
@@ -194,35 +192,21 @@ def _communicate_with_interrupt(process: subprocess.Popen, timeout_seconds: int)
continue
-def _extract_perf(stderr: str) -> str:
- lines = [
- line.strip() for line in str(stderr or "").splitlines()
- if any(prefix in line for prefix in LLAMA_PERF_PREFIXES)
- ]
- return "\n".join(lines)
-
-
-def _extract_thinking(text: str) -> tuple[str, str]:
+def _parse_response(text: str) -> tuple[str, str, str]:
text = str(text or "")
- thinking = ""
+ if PROMPT_ECHO_END in text:
+ text = text.split(PROMPT_ECHO_END, 1)[1]
- # llama.cpp returns raw stdout; Qwen reasoning may appear as a full block,
- # a closing-only block, or an unfinished block if generation is truncated.
- match = THINK_BLOCK_RE.search(text)
- if match:
- thinking = re.sub(r"?think[^>]*>", "", match.group(0), flags=re.IGNORECASE).strip()
- text = THINK_BLOCK_RE.sub("", text).strip()
- elif "" in text:
- before, after = text.split("", 1)
- thinking = before.strip()
- text = after.strip()
- elif "" in text:
- before, after = text.split("", 1)
- thinking = after.strip()
- text = before.strip()
+ perf_match = PERF_RE.search(text)
+ perf = perf_match.group(0).strip() if perf_match else ""
+ content = text[:perf_match.start()] if perf_match else text
+ content = content.strip()
+ if not content.startswith(START_THINKING):
+ return content, "", perf
- for token in ("<|im_end|>", "<|im_start|>", "<|endoftext|>", " [end of text]"):
- text = text.replace(token, "")
+ thinking_text = content[len(START_THINKING):]
+ if END_THINKING not in thinking_text:
+ return "", thinking_text.strip(), perf
- response = text.strip()
- return response, thinking
+ thinking, response = thinking_text.split(END_THINKING, 1)
+ return response.strip(), thinking.strip(), perf
diff --git a/nodes.py b/nodes.py
index 5405b7b..50a62ba 100644
--- a/nodes.py
+++ b/nodes.py
@@ -134,8 +134,8 @@ class QwenGGUF:
RETURN_NAMES = ("RESPONSE", "THINKING", "PERF")
OUTPUT_TOOLTIPS = (
"Final model response with Qwen thinking blocks removed.",
- "Extracted reasoning when present in model output.",
- "llama.cpp llama_perf_context_print lines.",
+ "Extracted reasoning when present in model output.",
+ "llama.cpp prompt and generation speed.",
)
FUNCTION = "generate"
CATEGORY = "Qwen GGUF"