diff --git a/comfy_extras/nodes_textgen.py b/comfy_extras/nodes_textgen.py index 40004652c..9664caece 100644 --- a/comfy_extras/nodes_textgen.py +++ b/comfy_extras/nodes_textgen.py @@ -211,6 +211,24 @@ AESTHETIC QUALITY (in addition to the above, without breaking the objective capt """ +LTX2_MARKERS = re.compile(r"|<\|channel>\w*\n?||<\|turn>\w*\n?") + + +def parse_ltx2_prompt(generated_text, prompt): + """Strip reasoning blocks and channel markers from the generated LTX2 prompt. + + A close marker with no opening tag leaves nothing after the split, so fall back to the + text before it and then to the user prompt. Both system prompts ask for the original + prompt back when the model has nothing to give, and an empty string here silently + generates the video with no conditioning at all. + """ + text = re.sub(r".*?", "", generated_text, flags=re.DOTALL) + if "" in text: # unclosed/truncated reasoning: keep what follows the last close + head, _, tail = text.rpartition("") + text = tail if LTX2_MARKERS.sub("", tail).strip() else head + return LTX2_MARKERS.sub("", text).strip() or prompt + + class TextGenerateLTX2Prompt(TextGenerate): @classmethod def define_schema(cls): @@ -256,12 +274,7 @@ 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) + return io.NodeOutput(parse_ltx2_prompt(out.args[0], prompt)) class TextgenExtension(ComfyExtension): diff --git a/tests-unit/comfy_extras_test/nodes_textgen_test.py b/tests-unit/comfy_extras_test/nodes_textgen_test.py new file mode 100644 index 000000000..62ef63210 --- /dev/null +++ b/tests-unit/comfy_extras_test/nodes_textgen_test.py @@ -0,0 +1,42 @@ +from comfy_extras.nodes_textgen import parse_ltx2_prompt + + +USER_PROMPT = "a fennec girl waving at the camera" + + +class TestParseLTX2Prompt: + @staticmethod + def _parse(generated_text: str) -> str: + return parse_ltx2_prompt(generated_text, USER_PROMPT) + + def test_plain_text(self): + assert self._parse("an enhanced prompt") == "an enhanced prompt" + + def test_dangling_close_keeps_text_before_it(self): + # Non-thinking mode primes a "final" channel, whose close decodes to a lone . + assert self._parse("an enhanced prompt") == "an enhanced prompt" + + def test_text_after_reasoning_block(self): + assert self._parse("reasoningan enhanced prompt") == "an enhanced prompt" + + def test_reasoning_only_falls_back_to_prompt(self): + assert self._parse("reasoning") == USER_PROMPT + + def test_empty_generation_falls_back_to_prompt(self): + assert self._parse("") == USER_PROMPT + + def test_multiline_reasoning_block(self): + assert self._parse("\nline one\nline two\n\nan enhanced prompt") == "an enhanced prompt" + + def test_reasoning_block_before_dangling_close(self): + assert self._parse("reasoningan enhanced prompt") == "an enhanced prompt" + + def test_channel_markers_are_stripped(self): + assert self._parse("<|channel>final\nan enhanced prompt") == "an enhanced prompt" + + def test_turn_markers_are_stripped(self): + assert self._parse("<|turn>model\nan enhanced prompt") == "an enhanced prompt" + + def test_never_returns_empty(self): + for generated_text in ["an enhanced prompt", "x", "xan enhanced prompt", "an enhanced prompt", "", "", " "]: + assert self._parse(generated_text) != ""