From 4a2259b7cda356a5fe94dbbac506614948c54753 Mon Sep 17 00:00:00 2001 From: da03 Date: Sat, 17 Oct 2026 11:11:23 -0400 Subject: [PATCH] Keep adjacent prompt images independent during tokenization --- programasweights/_runtime_vision.py | 154 +++++++++++++++------------- tests/test_vision_prompt.py | 21 +++- 2 files changed, 102 insertions(+), 73 deletions(-) diff --git a/programasweights/_runtime_vision.py b/programasweights/_runtime_vision.py index 2f0a4a6..25b8f69 100644 --- a/programasweights/_runtime_vision.py +++ b/programasweights/_runtime_vision.py @@ -169,14 +169,6 @@ def _eval_prompt(self, parts): marker = mtmd.mtmd_default_marker().decode("utf-8") prompt, urls = _to_prompt(*parts, image_marker=marker) - encoded = prompt.encode("utf-8") - - input_text = mtmd.mtmd_input_text() - input_text.text = encoded - input_text.text_len = len(encoded) - input_text.add_special = False - input_text.parse_special = True - bitmaps = [] chunks = None try: @@ -188,70 +180,90 @@ def _eval_prompt(self, parts): raise RuntimeError("Failed to create image bitmap.") bitmaps.append(bitmap) - chunks = mtmd.mtmd_input_chunks_init() - if not chunks: - raise RuntimeError("Failed to allocate prompt chunks.") - - bitmap_array = ( - mtmd.mtmd_bitmap_p_ctypes * len(bitmaps) - )(*bitmaps) - result = mtmd.mtmd_tokenize( - handler.mtmd_ctx, - chunks, - ctypes.byref(input_text), - bitmap_array, - len(bitmaps), - ) - if result != 0: - raise ValueError(f"Prompt tokenization failed: {result}.") - - for index in range(mtmd.mtmd_input_chunks_size(chunks)): - chunk = mtmd.mtmd_input_chunks_get(chunks, index) - if not chunk: - raise RuntimeError("Missing prompt chunk.") - - size = mtmd.mtmd_input_chunk_get_n_tokens(chunk) - if model.n_tokens + size >= model.n_ctx(): - raise ValueError("Prompt leaves no room for output.") - - kind = mtmd.mtmd_input_chunk_get_type(chunk) - if kind == mtmd.MTMD_INPUT_CHUNK_TYPE_TEXT: - count = ctypes.c_size_t() - tokens = mtmd.mtmd_input_chunk_get_tokens_text( - chunk, ctypes.byref(count) - ) - if count.value: - if not tokens: - raise RuntimeError("Missing text tokens.") - model.eval([tokens[i] for i in range(count.value)]) - - elif kind == mtmd.MTMD_INPUT_CHUNK_TYPE_IMAGE: - start = model.n_tokens - end = self.lib.llama_pos(0) - result = mtmd.mtmd_helper_eval_chunk_single( - handler.mtmd_ctx, - model._ctx.ctx, - chunk, - self.lib.llama_pos(start), - self.lib.llama_seq_id(0), - model.n_batch, - True, - ctypes.byref(end), - ) - if result != 0: - raise RuntimeError( - f"Image evaluation failed: {result}." + # PAW images are independent stills. Some native Qwen backends + # merge adjacent same-size bitmaps as video frames, changing the + # trained representation. Tokenize one image per call without + # inserting separators or changing any literal prompt text. + text_segments = prompt.split(marker) + groups = [(prompt, [])] if not bitmaps else [ + (text_segments[i] + marker + + (text_segments[-1] if i == len(bitmaps) - 1 else ""), [bitmap]) + for i, bitmap in enumerate(bitmaps) + ] + for text, group in groups: + encoded = text.encode("utf-8") + input_text = mtmd.mtmd_input_text() + input_text.text = encoded + input_text.text_len = len(encoded) + input_text.add_special = False + input_text.parse_special = True + chunks = mtmd.mtmd_input_chunks_init() + if not chunks: + raise RuntimeError("Failed to allocate prompt chunks.") + + bitmap_array = ( + mtmd.mtmd_bitmap_p_ctypes * len(group) + )(*group) + result = mtmd.mtmd_tokenize( + handler.mtmd_ctx, + chunks, + ctypes.byref(input_text), + bitmap_array, + len(group), + ) + if result != 0: + raise ValueError(f"Prompt tokenization failed: {result}.") + + for index in range(mtmd.mtmd_input_chunks_size(chunks)): + chunk = mtmd.mtmd_input_chunks_get(chunks, index) + if not chunk: + raise RuntimeError("Missing prompt chunk.") + + size = mtmd.mtmd_input_chunk_get_n_tokens(chunk) + if model.n_tokens + size >= model.n_ctx(): + raise ValueError("Prompt leaves no room for output.") + + kind = mtmd.mtmd_input_chunk_get_type(chunk) + if kind == mtmd.MTMD_INPUT_CHUNK_TYPE_TEXT: + count = ctypes.c_size_t() + tokens = mtmd.mtmd_input_chunk_get_tokens_text( + chunk, ctypes.byref(count) + ) + if count.value: + if not tokens: + raise RuntimeError("Missing text tokens.") + model.eval([tokens[i] for i in range(count.value)]) + + elif kind == mtmd.MTMD_INPUT_CHUNK_TYPE_IMAGE: + start = model.n_tokens + end = self.lib.llama_pos(0) + result = mtmd.mtmd_helper_eval_chunk_single( + handler.mtmd_ctx, + model._ctx.ctx, + chunk, + self.lib.llama_pos(start), + self.lib.llama_seq_id(0), + model.n_batch, + True, + ctypes.byref(end), ) - if not start < end.value < model.n_ctx(): - raise RuntimeError("Invalid image end position.") - - # Image embeddings occupy these positions. These zeros - # are Python bookkeeping, not tokens sent to the model. - model.input_ids[start:end.value] = 0 - model.n_tokens = end.value - model._requires_eval = False - else: - raise ValueError("Unsupported prompt chunk type.") + if result != 0: + raise RuntimeError( + f"Image evaluation failed: {result}." + ) + if not start < end.value < model.n_ctx(): + raise RuntimeError("Invalid image end position.") + + # Image embeddings occupy these positions. These zeros + # are Python bookkeeping, not tokens sent to the model. + model.input_ids[start:end.value] = 0 + model.n_tokens = end.value + model._requires_eval = False + else: + raise ValueError("Unsupported prompt chunk type.") + + mtmd.mtmd_input_chunks_free(chunks) + chunks = None if model.n_tokens == 0: raise ValueError("The prompt produced no tokens.") diff --git a/tests/test_vision_prompt.py b/tests/test_vision_prompt.py index e6c2e9e..c93e184 100644 --- a/tests/test_vision_prompt.py +++ b/tests/test_vision_prompt.py @@ -43,7 +43,7 @@ def tolist(self): @pytest.fixture def native(): state = NS(chunks=[], flags=None, loaded=[], freed=[], chunk_freed=False, - images=[], fail_tokenize=False, fail_image=False) + images=[], fail_tokenize=False, fail_image=False, prompts=[]) model = NS(n_tokens=0, input_ids=Vector([0] * 256), n_ctx=lambda: 256, _ctx=NS(ctx=None), n_batch=128) def evaluate(tokens): @@ -54,6 +54,7 @@ def tokenize(ctx, chunks, ptr, bitmaps, count): prompt = ptr._obj state.flags = (prompt.add_special, prompt.parse_special) state.text = prompt.text[:prompt.text_len].decode() + state.prompts.append(state.text) state.chunks = [] for i, piece in enumerate(state.text.split("<__media__>")): if i: @@ -103,7 +104,8 @@ def test_exact_text_and_image_order_no_wrapping_or_special_token_insertion(nativ state, handler, model, lib = native prompt = evaluate_prompt(handler, model, [text(" pre"), text("fix "), image("B"), text("mid"), image("A"), image("B"), text("end\n ")], lib) - assert state.text == " prefix <__media__>mid<__media__><__media__>end\n " + assert "".join(state.prompts) == " prefix <__media__>mid<__media__><__media__>end\n " + assert all(p.count("<__media__>") <= 1 for p in state.prompts) assert state.flags == (False, True) assert state.loaded == ["B", "A", "B"] assert state.images == [(1, True), (2, True), (3, True)] @@ -266,3 +268,18 @@ def test_evaluation_error_prevents_generation(): with pytest.raises(RuntimeError, match='native failure'): _Runtime._complete_prompt(owner, ('input',), max_tokens=8, temperature=0) owner.llm.create_completion.assert_not_called() + + +def test_adjacent_images_are_not_exposed_to_native_video_merging(native): + state, handler, model, lib = native + tokenize = handler._mtmd_cpp.mtmd_tokenize + def merges_video(*args): + result = tokenize(*args) + if "<__media__><__media__>" in state.text: + state.chunks = [c for i, c in enumerate(state.chunks) + if not (i and c.kind == state.chunks[i - 1].kind == 2)] + return result + handler._mtmd_cpp.mtmd_tokenize = merges_video + evaluate_prompt(handler, model, [image("A"), image("B")], lib) + assert state.images == [(1, True), (2, True)] + assert "".join(state.prompts) == "<__media__><__media__>"