fix: forward parsed phrase, not raw text, in SAM3 tokenizer fast path - #15839
fix: forward parsed phrase, not raw text, in SAM3 tokenizer fast path#15839a-yeyang wants to merge 1 commit into
Conversation
SAM3TokenizerWrapper.tokenize_with_weights() takes a fast path when there's
a single prompt with max_detections == 1 (i.e. a bare "person" or a
"person:1" with the default max_detections). _parse_prompts() correctly
strips the ":N" suffix in both cases and returns [("person", 1)], but the
fast path forwarded the raw, unparsed `text` to the inner tokenizer instead
of the parsed phrase -- so "person:1" was literally encoded as "person:1"
rather than "person", producing a different (and sometimes badly wrong)
embedding even though _parse_prompts already determined the phrase should
be identical to the bare "person" case.
The issue's own measurements show this can be catastrophic for some
wordings: on a test image, "person" masks 31.6% of the frame (correct),
while "person:1" -- which should be equivalent -- masks only 0.6% (wrong).
"person:2" and "person:1,person:1" already took the (correct) multi-prompt
path and were unaffected.
Fix: pass the already-parsed phrase to the fast path instead of the raw
text, falling back to the original text only when parsing produced nothing
(e.g. an empty string) to avoid indexing into an empty list. This is the
exact fix suggested in the issue report.
Added tests-unit/comfy_test/sam3_tokenizer_test.py, following the pattern
in tests-unit/comfy_test/gemma4_template_test.py: a small capture stand-in
subs for the inner SDTokenizer so the fix is verified against the real
SAM3TokenizerWrapper.tokenize_with_weights() code path without needing
real vocab/model files. Covers: the exact "person:1" regression, whitespace
variants, the unaffected bare-prompt and multi-prompt paths, and an empty
prompt edge case that would otherwise IndexError.
Fixes Comfy-Org#15811
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: ASSERTIVE Plan: Pro Plus Run ID: 📒 Files selected for processing (2)
Included review availability: Your plan provides up to 10 included reviews per hour; 8 remain after this review. 📜 Recent review details🧰 Additional context used📓 Path-based instructions (6)**/*📄 CodeRabbit inference engine (AGENTS.md)
Files:
**/*.py📄 CodeRabbit inference engine (AGENTS.md)
Files:
**/*.{py,json}📄 CodeRabbit inference engine (AGENTS.md)
Files:
**/*.{py,md,txt,json}📄 CodeRabbit inference engine (AGENTS.md)
Files:
**⚙️ CodeRabbit configuration file
Files:
comfy/**⚙️ CodeRabbit configuration file
Files:
🔇 Additional comments (3)
📝 WalkthroughWalkthroughSAM3 single-prompt tokenization now encodes parsed prompt text, so Merge Risk: ⚪ Minimal · up to The change forwards the parsed SAM3 phrase for the single-prompt fast path and adds focused regression coverage; no actionable merge-blocking risk remains beyond normal checks. 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment Warning |
Problem
In the SAM3 detection prompt grammar,
:Nsetsmax_detectionsfor a category, and omitting it defaults to1(_parse_prompts():result.append((part, 1))). Soperson:1is documented to be equivalent toperson.It isn't.
SAM3TokenizerWrapper.tokenize_with_weights()takes a fast path when there's a single prompt withmax_detections == 1— i.e. exactly thefooandfoo:1cases — but forwards the raw, unparsedtextto the inner tokenizer instead of the phrase_parse_prompts()already extracted:So
person:1is literally encoded as the string"person:1", not"person". Per the issue's own measurements, this can be catastrophic depending on wording: on a test image,personmasked 31.6% of the frame (correct), whileperson:1— which should produce the same embedding — masked only 0.6% (a wrong, unrelated region).person:2andperson:1,person:1already took the (correct) multi-prompt path and were unaffected.Fix
Forward the already-parsed phrase to the fast path instead of the raw text, falling back to the original text only when parsing produced nothing (an empty string) to avoid indexing into an empty list:
This is the exact fix suggested in the issue report.
Testing
Added
tests-unit/comfy_test/sam3_tokenizer_test.py, following the pattern already used intests-unit/comfy_test/gemma4_template_test.py: a small capture stand-in subs for the innerSDTokenizerso the fix is verified against the realSAM3TokenizerWrapper.tokenize_with_weights()code path without needing real vocab/model files. Covers: the exactperson:1regression, a whitespace-padded variant, the unaffected bare-prompt and multi-prompt paths (person:2,person:1,person:1), and an empty-prompt edge case.I verified the fix's correctness against the actual, unmodified
sam3_clip.pysource (loading it directly with lightweight stand-ins for itstorch/transformersdependency chain, since I could not get those installed in my sandbox) — confirmed the same assertions fail against the pre-fix code and pass against the post-fix code.ruff check comfy/text_encoders/sam3_clip.py tests-unit/comfy_test/sam3_tokenizer_test.py→ All checks passed.python -m py_compileon both files succeeds. I was not able to run the test file itself end-to-end with the real torch/transformers dependency chain locally (network constraints); it will run under CI'stest-unit.yml, which installs the full dependency set.Fixes #15811