From d8ff27d2610aaab0867cbd1eeb433a7b4aca7b8c Mon Sep 17 00:00:00 2001 From: Nguyen Thanh Dat Date: Mon, 24 Aug 2026 10:40:32 +0700 Subject: [PATCH] Strip the :N suffix on SAM3's single-prompt path (#15811) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `:N` sets max_detections per category and defaults to 1, so `person:1` is by definition the same detection prompt as `person`. `_parse_prompts()` strips the suffix, but the fast path in `SAM3TokenizerWrapper.tokenize_with_weights()` forwarded the raw text — and that branch is taken for exactly the `foo:1` shape, so the encoder grounded on the literal "person:1". Forward the parsed phrase instead, falling back to the raw text when nothing parses (empty prompt). This also makes the fast path agree with the multi-prompt path on parentheses, which `_parse_prompts()` already strips. --- comfy/text_encoders/sam3_clip.py | 7 +- .../comfy_test/sam3_prompt_parse_test.py | 138 ++++++++++++++++++ 2 files changed, 144 insertions(+), 1 deletion(-) create mode 100644 tests-unit/comfy_test/sam3_prompt_parse_test.py diff --git a/comfy/text_encoders/sam3_clip.py b/comfy/text_encoders/sam3_clip.py index 11cb7d9dbff..b7367f90f76 100644 --- a/comfy/text_encoders/sam3_clip.py +++ b/comfy/text_encoders/sam3_clip.py @@ -51,7 +51,12 @@ def __init__(self, embedding_directory=None, tokenizer_data={}): def tokenize_with_weights(self, text: str, return_word_ids=False, **kwargs): parsed = _parse_prompts(text) if len(parsed) <= 1 and (not parsed or parsed[0][1] == 1): - return super().tokenize_with_weights(text, return_word_ids, **kwargs) + # Forward the PARSED phrase, not the raw text: this branch is taken + # for `foo` and for `foo:1`, and the raw form still carries the + # `:1` suffix, so the encoder grounded on "person:1" instead of + # "person". Falls back to `text` when nothing parsed (empty prompt). + single = parsed[0][0] if parsed else text + return super().tokenize_with_weights(single, return_word_ids, **kwargs) # Tokenize each prompt part separately, store per-part batches and metadata inner = getattr(self, self.clip) per_prompt = [] diff --git a/tests-unit/comfy_test/sam3_prompt_parse_test.py b/tests-unit/comfy_test/sam3_prompt_parse_test.py new file mode 100644 index 00000000000..cd95614316e --- /dev/null +++ b/tests-unit/comfy_test/sam3_prompt_parse_test.py @@ -0,0 +1,138 @@ +"""#15811 — SAM3's single-prompt fast path forwarded the raw prompt text. + +`:N` sets max_detections per category and defaults to 1, so `person:1` is by +definition the same detection prompt as `person`. `_parse_prompts()` strips the +suffix correctly, but `SAM3TokenizerWrapper.tokenize_with_weights()` took the +fast path for exactly the `foo:1` shape (one prompt, max_detections == 1) and +handed `super()` the unparsed string — so the encoder grounded on the literal +`"person:1"`. + +The module is loaded here with `comfy.sd1_clip` stubbed, so the forwarding +contract is asserted without torch, transformers or any model weights. The +stub is scoped to the load and never enters `sys.modules` under the real +module name. +""" + +import importlib.util +import sys +import types +from contextlib import contextmanager +from pathlib import Path + +import pytest + +MODULE_PATH = Path(__file__).resolve().parents[2] / "comfy" / "text_encoders" / "sam3_clip.py" + + +class _RecordingTokenizer: + """Stands in for `sd1_clip.SD1Tokenizer`, recording what it is asked to tokenize.""" + + def __init__(self, *args, **kwargs): + self.seen = [] + self.clip = "clip_l" + self.clip_name = "l" + setattr(self, self.clip, self) + + def tokenize_with_weights(self, text, return_word_ids=False, **kwargs): + self.seen.append(text) + return [[("token", text)]] + + +@contextmanager +def _stubbed_sd1_clip(): + """Load sam3_clip with a stub base module, restoring sys.modules afterwards.""" + stub = types.ModuleType("comfy.sd1_clip") + stub.SDClipModel = type("SDClipModel", (), {"__init__": lambda self, **kwargs: None}) + stub.SDTokenizer = type("SDTokenizer", (), {"__init__": lambda self, **kwargs: None}) + stub.SD1ClipModel = type("SD1ClipModel", (), {"__init__": lambda self, **kwargs: None}) + stub.SD1Tokenizer = _RecordingTokenizer + + comfy_pkg = sys.modules.get("comfy") or types.ModuleType("comfy") + saved = { + "comfy": sys.modules.get("comfy"), + "comfy.sd1_clip": sys.modules.get("comfy.sd1_clip"), + } + saved_attr = getattr(comfy_pkg, "sd1_clip", None) + try: + sys.modules["comfy"] = comfy_pkg + sys.modules["comfy.sd1_clip"] = stub + comfy_pkg.sd1_clip = stub + + spec = importlib.util.spec_from_file_location("sam3_clip_under_test", MODULE_PATH) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + yield module + finally: + for name, previous in saved.items(): + if previous is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = previous + if saved_attr is None: + if hasattr(comfy_pkg, "sd1_clip"): + delattr(comfy_pkg, "sd1_clip") + else: + comfy_pkg.sd1_clip = saved_attr + + +@pytest.fixture +def sam3(): + with _stubbed_sd1_clip() as module: + yield module + + +def _tokenized(sam3_module, prompt): + """The strings the inner tokenizer was actually handed for `prompt`.""" + wrapper = sam3_module.SAM3TokenizerWrapper() + wrapper.tokenize_with_weights(prompt) + return getattr(wrapper, wrapper.clip).seen + + +def test_max_detections_suffix_does_not_reach_the_encoder(sam3): + assert _tokenized(sam3, "person:1") == ["person"] + + +def test_explicit_one_matches_the_bare_prompt(sam3): + """The equivalence the prompt syntax promises: `:1` is the default.""" + assert _tokenized(sam3, "person:1") == _tokenized(sam3, "person") + + +@pytest.mark.parametrize( + "prompt,expected", + [ + ("person", ["person"]), + ("person:1", ["person"]), + (" person : 1 ", ["person"]), + ("a person on a bike:1", ["a person on a bike"]), + # `_parse_prompts` strips parentheses (SAM3 has weights disabled), and + # the multi-prompt path already tokenized the stripped phrase; the fast + # path now agrees with it instead of encoding the literal brackets. + ("(person)", ["person"]), + ], +) +def test_single_prompt_shapes(sam3, prompt, expected): + assert _tokenized(sam3, prompt) == expected + + +def test_empty_prompt_still_reaches_the_tokenizer(sam3): + # Nothing parses out of these, so the original text is forwarded unchanged + # rather than being turned into None. + for prompt in ["", " ", ",", " , "]: + assert _tokenized(sam3, prompt) == [prompt] + + +def test_multi_prompt_path_is_untouched(sam3): + wrapper = sam3.SAM3TokenizerWrapper() + out = wrapper.tokenize_with_weights("person:2, car") + + assert getattr(wrapper, wrapper.clip).seen == ["person", "car"] + assert [max_det for _batches, max_det in out["sam3_per_prompt"]] == [2, 1] + + +def test_parse_prompts_itself_is_unchanged(sam3): + assert sam3._parse_prompts("person") == [("person", 1)] + assert sam3._parse_prompts("person:1") == [("person", 1)] + assert sam3._parse_prompts("person:3") == [("person", 3)] + assert sam3._parse_prompts("person:1.4") == [("person", 1)] + assert sam3._parse_prompts("person, car:2") == [("person", 1), ("car", 2)] + assert sam3._parse_prompts("") == []