diff --git a/tests-unit/comfy_test/gemma4_decode_test.py b/tests-unit/comfy_test/gemma4_decode_test.py new file mode 100644 index 000000000..bd932ec61 --- /dev/null +++ b/tests-unit/comfy_test/gemma4_decode_test.py @@ -0,0 +1,52 @@ +"""Gemma4 channel-marker decode regression tests.""" + +import torch + +from comfy.cli_args import args + +if not torch.cuda.is_available(): + args.cpu = True + +from comfy.text_encoders.gemma4 import Gemma4SDTokenizer # noqa: E402 + + +class _CannedTokenizer: + """Stands in for the wrapped tokenizer so decode is checked without model files.""" + def __init__(self, text): + self.text = text + + def decode(self, token_ids, skip_special_tokens=False): + return self.text + + +def decode(text): + tokenizer = Gemma4SDTokenizer.__new__(Gemma4SDTokenizer) + tokenizer.tokenizer = _CannedTokenizer(text) + return tokenizer.decode([]) + + +def test_thought_channel_becomes_think_tags(): + assert decode("<|channel>thought\nreasoningthe answer") == "\nreasoningthe answer" + + +def test_primed_empty_thought_channel_closes_immediately(): + assert decode("<|channel>thought\nthe answer") == "\nthe answer" + + +def test_thought_channel_left_unclosed_still_opens(): + assert decode("<|channel>thought\nreasoning") == "\nreasoning" + + +def test_close_of_a_non_thought_channel_is_not_reasoning(): + # Non-thinking LTX2 prompt enhancement primes a "final" channel in the prompt, so only its + # close is generated. Reading that close as made the whole answer look like + # reasoning and the node returned an empty string. + assert decode("the answer") == "the answer" + + +def test_thought_close_does_not_consume_a_later_channel(): + assert decode("<|channel>thought\nreasoning<|channel>final\nthe answer") == "\nreasoningthe answer" + + +def test_turn_and_eos_markers_are_stripped(): + assert decode("<|turn>model\nthe answer") == "the answer"