[Feature] HunyuanImage-3.0 IT2I: multi-image input + prompt API cleanup (#3444)
Signed-off-by: TaffyOfficial <2324465096@qq.com> Signed-off-by: TaffyOfficial <wu15922848573@outlook.com> Signed-off-by: skf1999 <13234016272@163.com> Signed-off-by: zuiho <2324465096@qq.com> Signed-off-by: Claude Code <noreply@anthropic.com> Signed-off-by: zuiho <wu15922848573@outlook.com> Signed-off-by: TaffyOfficial <2587297563@qq.com> Co-authored-by: TaffyOfficial <2324465096@qq.com> Co-authored-by: TaffyOfficial <wu15922848573@outlook.com> Co-authored-by: skf1999 <13234016272@163.com>
This commit is contained in:
17 files changed
+1559
-534
No files matched your search
@@ -0,0 +1,255 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Multi-image input regression for HunyuanImage3 IT2I prompt construction.
|
||||
|
||||
The official HunyuanImage-3.0-Instruct supports up to 3 reference images
|
||||
per IT2I request ("Multi-Image Fusion"; see hunyuan3.0_ins/README.md
|
||||
section 200-216 + line 500). Each cond image becomes its own user-role
|
||||
message and `apply_general_template` concatenates successive user
|
||||
messages back-to-back inside ONE user_prefix/user_suffix wrap (see
|
||||
hunyuan3.0_ins/tokenization_hunyuan_image_3.py:1399-1400, 1499-1515).
|
||||
The lightweight `<img>` + `multi_modal_data` builder used by the example
|
||||
flow must match that contract: N consecutive `<img>` placeholders sit
|
||||
between `User: ` and the user prompt, with no separator between them.
|
||||
|
||||
This file pins:
|
||||
1. N consecutive `<img>` placeholders for N=1/2/3 across both the
|
||||
string builder (`build_prompt`) and the token builder
|
||||
(`build_prompt_tokens`).
|
||||
2. The N=1 path stays bit-identical to the legacy single-image builder
|
||||
(regression guard so default callers don't notice).
|
||||
3. N=2 / N=3 token sequences differ from N=1 by exactly (N-1) extra
|
||||
`<img>` ids inserted between `User: ` and `user_prompt`.
|
||||
4. Validation: N<1 and N>3 raise ValueError (hard cap N<=3 mirrors
|
||||
official upstream).
|
||||
5. Text-only tasks ignore `num_images` (no validation, no extra ids).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import (
|
||||
MAX_IMAGES_PER_REQUEST,
|
||||
build_prompt,
|
||||
build_prompt_tokens,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class FakeTokenizer:
|
||||
"""Recording fake tokenizer mirroring the one in test_prompt_utils.
|
||||
|
||||
Special token ids: `<|startoftext|>`=1, `<img>`=2, `<think>`=3,
|
||||
`<recaption>`=4. encode() returns one id per character starting at
|
||||
100, so substring-position assertions are stable.
|
||||
"""
|
||||
|
||||
SPECIAL = {
|
||||
"<|startoftext|>": 1,
|
||||
"<img>": 2,
|
||||
"<think>": 3,
|
||||
"<recaption>": 4,
|
||||
}
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.encode_calls: list[str] = []
|
||||
|
||||
def convert_tokens_to_ids(self, tok: str) -> int:
|
||||
return self.SPECIAL.get(tok, 0)
|
||||
|
||||
def encode(self, text: str, add_special_tokens: bool = False) -> list[int]:
|
||||
self.encode_calls.append(text)
|
||||
return list(range(100, 100 + len(text)))
|
||||
|
||||
|
||||
_IMAGE_TASK_COMBOS = (
|
||||
("i2t", None),
|
||||
("it2i", "think"),
|
||||
("it2i", "recaption"),
|
||||
)
|
||||
_TEXT_ONLY_TASK_COMBOS = (("t2t", None),)
|
||||
|
||||
|
||||
# -------------------- string builder --------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("task,bot_task", _IMAGE_TASK_COMBOS)
|
||||
@pytest.mark.parametrize("num_images", [1, 2, 3])
|
||||
def test_build_prompt_emits_N_consecutive_img_placeholders(task: str, bot_task: str | None, num_images: int):
|
||||
"""N=1/2/3 -> exactly N `<img>` substrings appear consecutively
|
||||
between `User: ` and the user prompt, with no separator between them."""
|
||||
s = build_prompt("HELLO", task=task, bot_task=bot_task, num_images=num_images)
|
||||
assert s.count("<img>") == num_images, (
|
||||
f"task={task} bot_task={bot_task} num_images={num_images}: expected {num_images} <img> "
|
||||
f"placeholders, found {s.count('<img>')} -- prompt was: {s!r}"
|
||||
)
|
||||
|
||||
# All `<img>` placeholders must form one contiguous run "<img><img>..."
|
||||
# immediately after `User: ` and before HELLO.
|
||||
user_idx = s.index("User: ") + len("User: ")
|
||||
hello_idx = s.index("HELLO")
|
||||
between = s[user_idx:hello_idx]
|
||||
assert between == "<img>" * num_images, (
|
||||
f"region between `User: ` and prompt must be exactly N <img> placeholders; got {between!r}"
|
||||
)
|
||||
|
||||
|
||||
def test_build_prompt_default_num_images_matches_legacy():
|
||||
"""num_images default = 1 must produce a string bit-identical to the
|
||||
pre-multi-image behavior (single `<img>` placeholder)."""
|
||||
legacy = build_prompt("HELLO", task="it2i", bot_task="think")
|
||||
explicit = build_prompt("HELLO", task="it2i", bot_task="think", num_images=1)
|
||||
assert legacy == explicit, "default num_images=1 must match legacy single-image output"
|
||||
|
||||
|
||||
# -------------------- token builder --------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("task,bot_task", _IMAGE_TASK_COMBOS)
|
||||
def test_build_prompt_tokens_inserts_N_img_ids(task: str, bot_task: str | None):
|
||||
"""N=1/2/3 -> the resulting id sequence contains exactly N copies of
|
||||
img_id (=2) sitting consecutively after the `User: ` segment."""
|
||||
tok = FakeTokenizer()
|
||||
ids_n1 = build_prompt_tokens("hi", tok, task=task, bot_task=bot_task, num_images=1).token_ids
|
||||
tok = FakeTokenizer()
|
||||
ids_n2 = build_prompt_tokens("hi", tok, task=task, bot_task=bot_task, num_images=2).token_ids
|
||||
tok = FakeTokenizer()
|
||||
ids_n3 = build_prompt_tokens("hi", tok, task=task, bot_task=bot_task, num_images=3).token_ids
|
||||
|
||||
assert ids_n1.count(2) == 1
|
||||
assert ids_n2.count(2) == 2
|
||||
assert ids_n3.count(2) == 3
|
||||
|
||||
# Each additional image must extend the sequence by exactly one img_id,
|
||||
# not shift other tokens around.
|
||||
assert len(ids_n2) == len(ids_n1) + 1
|
||||
assert len(ids_n3) == len(ids_n1) + 2
|
||||
|
||||
# The img_ids must be CONSECUTIVE (no other token between successive
|
||||
# `<img>` placeholders -- mirrors the official `process_successive_message`
|
||||
# wrapping where successive user messages share one user_prefix/suffix).
|
||||
for ids, n in [(ids_n2, 2), (ids_n3, 3)]:
|
||||
first = ids.index(2)
|
||||
for k in range(n):
|
||||
assert ids[first + k] == 2, (
|
||||
f"img_ids must be consecutive starting at position {first} for n={n}; got {ids[first : first + n]!r}"
|
||||
)
|
||||
|
||||
|
||||
def test_build_prompt_tokens_default_num_images_matches_legacy():
|
||||
"""num_images default = 1 must produce the same id sequence as
|
||||
omitting the parameter (regression guard for existing single-image
|
||||
callers)."""
|
||||
tok_a = FakeTokenizer()
|
||||
legacy = build_prompt_tokens("hi", tok_a, task="it2i", bot_task="think").token_ids
|
||||
tok_b = FakeTokenizer()
|
||||
explicit = build_prompt_tokens("hi", tok_b, task="it2i", bot_task="think", num_images=1).token_ids
|
||||
assert legacy == explicit
|
||||
# Also: encode() must have been called on the same set of segments,
|
||||
# so segment boundaries are preserved.
|
||||
assert tok_a.encode_calls == tok_b.encode_calls
|
||||
|
||||
|
||||
# -------------------- validation --------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("task,bot_task", _IMAGE_TASK_COMBOS)
|
||||
@pytest.mark.parametrize("bad", [0, -1, MAX_IMAGES_PER_REQUEST + 1, 99])
|
||||
def test_build_prompt_rejects_out_of_range_num_images(task: str, bot_task: str | None, bad: int):
|
||||
with pytest.raises(ValueError, match="num_images must be in"):
|
||||
build_prompt("hi", task=task, bot_task=bot_task, num_images=bad)
|
||||
with pytest.raises(ValueError, match="num_images must be in"):
|
||||
build_prompt_tokens("hi", FakeTokenizer(), task=task, bot_task=bot_task, num_images=bad)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("task,bot_task", _TEXT_ONLY_TASK_COMBOS)
|
||||
@pytest.mark.parametrize("num_images", [0, 1, 2, 99])
|
||||
def test_text_only_tasks_ignore_num_images(task: str, bot_task: str | None, num_images: int):
|
||||
"""Validation only kicks in for image-input tasks; t2t et al. accept
|
||||
any num_images and emit zero `<img>` placeholders."""
|
||||
s = build_prompt("hi", task=task, bot_task=bot_task, num_images=num_images)
|
||||
assert "<img>" not in s
|
||||
ids = build_prompt_tokens("hi", FakeTokenizer(), task=task, bot_task=bot_task, num_images=num_images).token_ids
|
||||
assert 2 not in ids
|
||||
|
||||
|
||||
# -------------------- real HF tokenizer regression --------------------
|
||||
|
||||
_HUNYUAN_MODEL_ID = "tencent/HunyuanImage-3.0-Instruct"
|
||||
|
||||
|
||||
def _hf_cached(model_id: str) -> bool:
|
||||
hf_home = os.environ.get("HF_HOME") or os.path.expanduser("~/.cache/huggingface")
|
||||
snap_dir = os.path.join(hf_home, "hub", f"models--{model_id.replace('/', '--')}", "snapshots")
|
||||
return os.path.isdir(snap_dir) and any(os.scandir(snap_dir))
|
||||
|
||||
|
||||
@pytest.mark.skipif(not _hf_cached(_HUNYUAN_MODEL_ID), reason=f"{_HUNYUAN_MODEL_ID} tokenizer not in HF cache")
|
||||
@pytest.mark.parametrize("num_images", [1, 2, 3])
|
||||
def test_real_tokenizer_emits_n_consecutive_img_ids(num_images: int):
|
||||
"""Real `AutoTokenizer.from_pretrained(...)` (the production path) must
|
||||
encode N=1/2/3 prompts to a sequence with exactly N consecutive `<img>`
|
||||
token-ids in the right place — proves the placeholder layout from
|
||||
`build_prompt_tokens` survives a real BPE tokenizer, not just FakeTokenizer.
|
||||
"""
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
tok = AutoTokenizer.from_pretrained(_HUNYUAN_MODEL_ID, trust_remote_code=True)
|
||||
img_id = tok.convert_tokens_to_ids("<img>")
|
||||
assert img_id is not None and img_id >= 0, f"<img> not in tokenizer vocab; got id={img_id}"
|
||||
|
||||
ids = build_prompt_tokens("hi", tok, task="it2i", bot_task="think", num_images=num_images).token_ids
|
||||
|
||||
# Exactly N copies of <img> id, all consecutive.
|
||||
img_positions = [i for i, x in enumerate(ids) if x == img_id]
|
||||
assert len(img_positions) == num_images, (
|
||||
f"expected {num_images} <img> ids, got {len(img_positions)} at positions {img_positions}"
|
||||
)
|
||||
assert img_positions == list(range(img_positions[0], img_positions[0] + num_images)), (
|
||||
f"<img> ids must be contiguous; got positions {img_positions}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not _hf_cached(_HUNYUAN_MODEL_ID), reason=f"{_HUNYUAN_MODEL_ID} tokenizer not in HF cache")
|
||||
def test_real_tokenizer_n_plus_one_extends_by_exactly_one_img_id():
|
||||
"""Going from N to N+1 images must extend the encoded id sequence by
|
||||
exactly one extra `<img>` token-id and shift nothing else. Catches
|
||||
accidental separator tokens between successive `<img>` placeholders
|
||||
that a FakeTokenizer (deterministic encode) can't surface."""
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
tok = AutoTokenizer.from_pretrained(_HUNYUAN_MODEL_ID, trust_remote_code=True)
|
||||
img_id = tok.convert_tokens_to_ids("<img>")
|
||||
|
||||
ids_n1 = build_prompt_tokens("hi", tok, task="it2i", bot_task="think", num_images=1).token_ids
|
||||
ids_n2 = build_prompt_tokens("hi", tok, task="it2i", bot_task="think", num_images=2).token_ids
|
||||
ids_n3 = build_prompt_tokens("hi", tok, task="it2i", bot_task="think", num_images=3).token_ids
|
||||
|
||||
assert len(ids_n2) == len(ids_n1) + 1, f"N=2 should be N=1 + 1 token; got {len(ids_n2)} vs {len(ids_n1)}"
|
||||
assert len(ids_n3) == len(ids_n1) + 2, f"N=3 should be N=1 + 2 tokens; got {len(ids_n3)} vs {len(ids_n1)}"
|
||||
|
||||
# Insert one img_id at the existing position; everything else unchanged.
|
||||
p1 = ids_n1.index(img_id)
|
||||
assert ids_n2[: p1 + 1] == ids_n1[: p1 + 1] + [], "prefix before extra <img> must match N=1"
|
||||
assert ids_n2[p1] == img_id and ids_n2[p1 + 1] == img_id, "two consecutive <img> ids at the insertion point"
|
||||
assert ids_n2[p1 + 2 :] == ids_n1[p1 + 1 :], "tail after the extra <img> must match N=1's tail"
|
||||
# N=3 same pattern, three in a row.
|
||||
assert ids_n3[p1 : p1 + 3] == [img_id, img_id, img_id]
|
||||
assert ids_n3[p1 + 3 :] == ids_n1[p1 + 1 :]
|
||||
|
||||
|
||||
@pytest.mark.skipif(not _hf_cached(_HUNYUAN_MODEL_ID), reason=f"{_HUNYUAN_MODEL_ID} tokenizer not in HF cache")
|
||||
def test_real_tokenizer_default_n1_byte_identical_to_legacy():
|
||||
"""Default `num_images=1` must produce the exact same id sequence as
|
||||
omitting the parameter — pins the legacy single-image regression
|
||||
against the real tokenizer (not just FakeTokenizer)."""
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
tok = AutoTokenizer.from_pretrained(_HUNYUAN_MODEL_ID, trust_remote_code=True)
|
||||
legacy = build_prompt_tokens("hi", tok, task="it2i", bot_task="think").token_ids
|
||||
explicit = build_prompt_tokens("hi", tok, task="it2i", bot_task="think", num_images=1).token_ids
|
||||
assert legacy == explicit, "real tokenizer: default num_images=1 must be byte-identical to legacy"
|
||||
@@ -1,20 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Regression tests for HunyuanImage3 prompt construction (PR #3243).
|
||||
|
||||
Two layers:
|
||||
1. Pure-logic tests with a recording fake tokenizer -- protect the
|
||||
prompt template structure (BOS, User:/Assistant: framing, trigger
|
||||
placement, image placeholder position) and protect the segment-
|
||||
by-segment tokenization contract (each segment must hit
|
||||
`tokenizer.encode` in isolation).
|
||||
2. Real-tokenizer regression -- run when the HunyuanImage3-Instruct
|
||||
tokenizer is in the local HF cache. Asserts the segment-tokenized
|
||||
output diverges from the naive full-string encode, which is the
|
||||
bug-tripping fixture for the cross-segment BPE merge fix
|
||||
(commit 7bd429ed).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import ast
|
||||
@@ -24,7 +9,9 @@ import pathlib
|
||||
import pytest
|
||||
|
||||
from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import (
|
||||
_TASK_PRESETS,
|
||||
HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS,
|
||||
available_bot_tasks,
|
||||
available_tasks,
|
||||
build_prompt,
|
||||
build_prompt_tokens,
|
||||
@@ -34,18 +21,7 @@ from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import (
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
# -------------------- Pure-logic structural tests --------------------
|
||||
|
||||
|
||||
class FakeTokenizer:
|
||||
"""Minimal tokenizer stub that records every encode() call.
|
||||
|
||||
Returns deterministic ids from convert_tokens_to_ids while
|
||||
encode() returns one id per character starting at 100. This lets
|
||||
tests both verify segmentation (by inspecting `encode_calls`) and
|
||||
locate substrings inside the returned id list.
|
||||
"""
|
||||
|
||||
SPECIAL = {
|
||||
"<|startoftext|>": 1,
|
||||
"<img>": 2,
|
||||
@@ -72,85 +48,131 @@ class FakeTokenizer:
|
||||
|
||||
|
||||
def test_available_tasks_covers_all_modalities():
|
||||
tasks = set(available_tasks())
|
||||
assert tasks >= {
|
||||
"t2t",
|
||||
"i2t",
|
||||
assert set(available_tasks()) == {"t2t", "i2t", "it2i", "t2i"}
|
||||
|
||||
|
||||
def test_available_bot_tasks_covers_all_modes():
|
||||
assert set(available_bot_tasks()) == {None, "think", "recaption", "think_recaption", "vanilla"}
|
||||
|
||||
|
||||
def test_legacy_task_presets_still_available():
|
||||
assert {
|
||||
"it2i_think",
|
||||
"it2i_recaption",
|
||||
"it2i_think_recaption",
|
||||
"t2i_think",
|
||||
"t2i_recaption",
|
||||
"t2i_vanilla",
|
||||
} <= set(_TASK_PRESETS)
|
||||
|
||||
|
||||
def test_legacy_base_task_omitted_bot_task_keeps_plain_mode():
|
||||
prompt = build_prompt("HELLO", task="i2t")
|
||||
assert prompt.endswith("Assistant: ")
|
||||
assert not prompt.endswith("<think>")
|
||||
|
||||
result = build_prompt_tokens("hi", FakeTokenizer(), task="i2t")
|
||||
assert result.system_prompt_type == "en_unified"
|
||||
assert result.token_ids[-1] not in {
|
||||
FakeTokenizer.SPECIAL["<think>"],
|
||||
FakeTokenizer.SPECIAL["<recaption>"],
|
||||
}
|
||||
|
||||
|
||||
def test_resolve_stop_token_ids_uses_answer_for_generation_tasks():
|
||||
def test_legacy_composite_task_with_none_bot_task_keeps_encoded_mode():
|
||||
prompt = build_prompt("HELLO", task="it2i_think", bot_task=None)
|
||||
assert prompt.endswith("Assistant: <think>")
|
||||
|
||||
result = build_prompt_tokens("hi", FakeTokenizer(), task="it2i_recaption", bot_task=None)
|
||||
assert result.token_ids[-1] == FakeTokenizer.SPECIAL["<recaption>"]
|
||||
|
||||
|
||||
def test_default_prompt_still_uses_it2i_think_mode():
|
||||
prompt = build_prompt("HELLO")
|
||||
assert prompt.endswith("Assistant: <think>")
|
||||
|
||||
result = build_prompt_tokens("hi", FakeTokenizer())
|
||||
assert result.system_prompt_type == "en_unified"
|
||||
assert result.token_ids[-1] == FakeTokenizer.SPECIAL["<think>"]
|
||||
|
||||
|
||||
def test_resolve_stop_token_ids_image_tasks_stop_on_ratio_range():
|
||||
"""Image-output tasks stop on any ``<img_ratio_*>`` token.
|
||||
|
||||
Mirrors upstream ``modeling_hunyuan_image_3.py::generate_image``
|
||||
(line 3289-3303): when ``need_ratio`` is true,
|
||||
``final_stop_tokens = list(range(start_ratio, end_ratio + 1)) +
|
||||
ratio_token_other_slices``. AR stops AT the ratio token sampled
|
||||
after ``<img_size_*>``; the bridge then strips the trailing ratio
|
||||
token before passing the cot to DiT.
|
||||
"""
|
||||
tok = FakeTokenizer()
|
||||
|
||||
start = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_0>"]
|
||||
end = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_32>"]
|
||||
other_start = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_33>"]
|
||||
other_end = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_36>"]
|
||||
expected = list(range(start, end + 1)) + list(range(other_start, other_end + 1))
|
||||
|
||||
# Image-output: t2i / it2i stop on the full ratio token range.
|
||||
for bot in ("think", "recaption", "think_recaption", "vanilla"):
|
||||
assert resolve_stop_token_ids(task="t2i", bot_task=bot, tokenizer=tok) == expected
|
||||
assert resolve_stop_token_ids(task="it2i", bot_task=bot, tokenizer=tok) == expected
|
||||
|
||||
# Text-output: i2t / t2t comprehension stops on <answer> (response sits inside).
|
||||
answer_id = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<answer>"]
|
||||
assert resolve_stop_token_ids(task="t2i_think", tokenizer=tok) == [answer_id]
|
||||
assert resolve_stop_token_ids(task="t2i_recaption", tokenizer=tok) == [answer_id]
|
||||
assert resolve_stop_token_ids(task="i2t", bot_task=None, tokenizer=tok) == [answer_id]
|
||||
assert resolve_stop_token_ids(task="t2t", bot_task=None, tokenizer=tok) == [answer_id]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"task",
|
||||
"task,bot_task",
|
||||
[
|
||||
"t2t",
|
||||
"i2t",
|
||||
"it2i_think",
|
||||
"it2i_recaption",
|
||||
"it2i_think_recaption",
|
||||
"t2i_think",
|
||||
"t2i_recaption",
|
||||
("t2t", None),
|
||||
("i2t", None),
|
||||
("it2i", "think"),
|
||||
("it2i", "recaption"),
|
||||
("it2i", "think_recaption"),
|
||||
("t2i", "think"),
|
||||
("t2i", "recaption"),
|
||||
("t2i", "think_recaption"),
|
||||
],
|
||||
)
|
||||
def test_build_prompt_string_structure_chat_template(task: str):
|
||||
"""Chat-template tasks must produce <|startoftext|>...User: ...Assistant: ...
|
||||
with image placeholder (when applicable) and trigger tag AFTER `Assistant: `."""
|
||||
s = build_prompt("HELLO", task=task)
|
||||
|
||||
def test_build_prompt_string_structure_chat_template(task: str, bot_task: str | None):
|
||||
s = build_prompt("HELLO", task=task, bot_task=bot_task)
|
||||
assert s.startswith("<|startoftext|>")
|
||||
assert "User: " in s
|
||||
assert "Assistant: " in s
|
||||
assert s.index("User: ") < s.index("HELLO") < s.index("Assistant: ")
|
||||
|
||||
if task.startswith(("i2t", "it2i")):
|
||||
assert s.index("User: ") < s.index("<img>") < s.index("HELLO"), (
|
||||
"<img> placeholder must sit between `User: ` and the user prompt"
|
||||
)
|
||||
if task in ("i2t", "it2i"):
|
||||
assert s.index("User: ") < s.index("<img>") < s.index("HELLO")
|
||||
else:
|
||||
assert "<img>" not in s
|
||||
|
||||
# Trigger tag must be the FINAL token of the prompt (after `Assistant: `).
|
||||
# Note: the system prompt itself mentions <think>/<recaption> as mode
|
||||
# documentation, so substring index() catches the wrong occurrence -- use
|
||||
# endswith() which directly captures "trigger is at the tail" (the Part A
|
||||
# fix: trigger goes AFTER `Assistant: `, not before user_prompt).
|
||||
if task in ("it2i_think", "t2i_think", "it2i_think_recaption"):
|
||||
assert s.endswith("Assistant: <think>"), (
|
||||
f"Trigger <think> must be appended right after `Assistant: ` (Part A fix). Got tail: ...{s[-40:]!r}"
|
||||
)
|
||||
if task in ("it2i_recaption", "t2i_recaption"):
|
||||
assert s.endswith("Assistant: <recaption>"), (
|
||||
f"Trigger <recaption> must be appended right after `Assistant: ` (Part A fix). Got tail: ...{s[-40:]!r}"
|
||||
)
|
||||
if task in ("t2t", "i2t"):
|
||||
assert s.endswith("Assistant: "), "Plain (no-trigger) task must end at `Assistant: ` with no trailing tag."
|
||||
if bot_task in ("think", "think_recaption"):
|
||||
assert s.endswith("Assistant: <think>")
|
||||
elif bot_task == "recaption":
|
||||
assert s.endswith("Assistant: <recaption>")
|
||||
elif bot_task is None:
|
||||
assert s.endswith("Assistant: ")
|
||||
|
||||
|
||||
def test_build_prompt_vanilla_uses_pretrain_template():
|
||||
"""t2i_vanilla is the only task that bypasses chat structure -- direct
|
||||
text->image generation driven by the vanilla system prompt."""
|
||||
s = build_prompt("HELLO", task="t2i_vanilla")
|
||||
s = build_prompt("HELLO", task="t2i", bot_task="vanilla")
|
||||
assert s.startswith("<|startoftext|>")
|
||||
assert "User: " not in s
|
||||
assert "Assistant: " not in s
|
||||
assert "<think>" not in s
|
||||
assert "<recaption>" not in s
|
||||
assert s.endswith("HELLO")
|
||||
|
||||
|
||||
def test_build_prompt_vanilla_rejects_non_t2i_task():
|
||||
with pytest.raises(ValueError, match="bot_task='vanilla'"):
|
||||
build_prompt("x", task="it2i", bot_task="vanilla")
|
||||
with pytest.raises(ValueError, match="bot_task='vanilla'"):
|
||||
build_prompt_tokens("x", FakeTokenizer(), task="i2t", bot_task="vanilla")
|
||||
|
||||
|
||||
def test_build_prompt_unknown_task_raises():
|
||||
with pytest.raises(ValueError, match="Unknown task"):
|
||||
build_prompt("x", task="bogus")
|
||||
@@ -158,127 +180,83 @@ def test_build_prompt_unknown_task_raises():
|
||||
build_prompt_tokens("x", FakeTokenizer(), task="bogus")
|
||||
|
||||
|
||||
def test_build_prompt_unknown_bot_task_raises():
|
||||
with pytest.raises(ValueError, match="Unknown bot_task"):
|
||||
build_prompt("x", task="t2i", bot_task="bogus")
|
||||
with pytest.raises(ValueError, match="Unknown bot_task"):
|
||||
build_prompt_tokens("x", FakeTokenizer(), task="t2i", bot_task="bogus")
|
||||
|
||||
|
||||
def test_build_prompt_tokens_segments_each_boundary():
|
||||
"""Regression for cross-segment BPE merge bug (commit 7bd429ed):
|
||||
each template segment must hit tokenizer.encode() independently;
|
||||
user_prompt MUST NOT be concatenated with the following separator
|
||||
in the same encode() call."""
|
||||
tok = FakeTokenizer()
|
||||
build_prompt_tokens("写诗。", tok, task="i2t")
|
||||
|
||||
# Each canonical segment is encoded in its own call.
|
||||
build_prompt_tokens("写诗。", tok, task="i2t", bot_task=None)
|
||||
assert "User: " in tok.encode_calls
|
||||
assert "写诗。" in tok.encode_calls, (
|
||||
"user_prompt must be encoded alone -- if it is concatenated with the "
|
||||
"trailing separator, BPE will merge across the boundary (the PR-#3243 bug)."
|
||||
)
|
||||
assert "写诗。" in tok.encode_calls
|
||||
assert "\n\nAssistant: " in tok.encode_calls
|
||||
|
||||
# No call must contain user_prompt glued to neighboring text.
|
||||
for call in tok.encode_calls:
|
||||
if call != "写诗。":
|
||||
assert "写诗。" not in call, f"user_prompt leaked into a multi-segment encode call: {call!r}"
|
||||
assert "写诗。" not in call
|
||||
|
||||
|
||||
def test_build_prompt_tokens_image_placeholder_present_for_image_tasks():
|
||||
tok = FakeTokenizer()
|
||||
result = build_prompt_tokens("hi", tok, task="i2t")
|
||||
result = build_prompt_tokens("hi", tok, task="i2t", bot_task=None)
|
||||
ids = result.token_ids
|
||||
assert ids[0] == FakeTokenizer.SPECIAL["<|startoftext|>"], "BOS (<|startoftext|>) must be the first token"
|
||||
assert FakeTokenizer.SPECIAL["<img>"] in ids, "<img> placeholder must be present for i2t/it2i tasks"
|
||||
assert ids[0] == FakeTokenizer.SPECIAL["<|startoftext|>"]
|
||||
assert FakeTokenizer.SPECIAL["<img>"] in ids
|
||||
|
||||
|
||||
def test_build_prompt_tokens_no_image_for_text_only_tasks():
|
||||
tok = FakeTokenizer()
|
||||
result = build_prompt_tokens("hi", tok, task="t2t")
|
||||
result = build_prompt_tokens("hi", tok, task="t2t", bot_task=None)
|
||||
ids = result.token_ids
|
||||
assert FakeTokenizer.SPECIAL["<img>"] not in ids, "<img> must NOT appear for text-only tasks"
|
||||
assert FakeTokenizer.SPECIAL["<img>"] not in ids
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"task,trigger_id",
|
||||
"task,bot_task,trigger_id",
|
||||
[
|
||||
("it2i_think", FakeTokenizer.SPECIAL["<think>"]),
|
||||
("t2i_think", FakeTokenizer.SPECIAL["<think>"]),
|
||||
("it2i_recaption", FakeTokenizer.SPECIAL["<recaption>"]),
|
||||
("t2i_recaption", FakeTokenizer.SPECIAL["<recaption>"]),
|
||||
("it2i", "think", FakeTokenizer.SPECIAL["<think>"]),
|
||||
("t2i", "think", FakeTokenizer.SPECIAL["<think>"]),
|
||||
("t2i", "think_recaption", FakeTokenizer.SPECIAL["<think>"]),
|
||||
("it2i", "recaption", FakeTokenizer.SPECIAL["<recaption>"]),
|
||||
("t2i", "recaption", FakeTokenizer.SPECIAL["<recaption>"]),
|
||||
("it2i_think", None, FakeTokenizer.SPECIAL["<think>"]),
|
||||
("it2i_recaption", None, FakeTokenizer.SPECIAL["<recaption>"]),
|
||||
],
|
||||
)
|
||||
def test_build_prompt_tokens_trigger_is_last_token(task: str, trigger_id: int):
|
||||
"""Trigger tag id must be the LAST token (after `Assistant: ` segment)."""
|
||||
def test_build_prompt_tokens_trigger_is_last_token(task: str, bot_task: str | None, trigger_id: int):
|
||||
tok = FakeTokenizer()
|
||||
result = build_prompt_tokens("hi", tok, task=task)
|
||||
ids = result.token_ids
|
||||
assert ids[-1] == trigger_id
|
||||
result = build_prompt_tokens("hi", tok, task=task, bot_task=bot_task)
|
||||
assert result.token_ids[-1] == trigger_id
|
||||
|
||||
|
||||
def test_build_prompt_tokens_no_trigger_for_plain_tasks():
|
||||
"""Tasks without trigger_tag (t2t / i2t) must NOT append a trigger id."""
|
||||
tok = FakeTokenizer()
|
||||
result = build_prompt_tokens("hi", tok, task="t2t")
|
||||
ids = result.token_ids
|
||||
assert ids[-1] not in {
|
||||
result = build_prompt_tokens("hi", tok, task="t2t", bot_task=None)
|
||||
assert result.token_ids[-1] not in {
|
||||
FakeTokenizer.SPECIAL["<think>"],
|
||||
FakeTokenizer.SPECIAL["<recaption>"],
|
||||
}
|
||||
|
||||
|
||||
# -------------------- end2end.py wiring guard --------------------
|
||||
|
||||
|
||||
def _repo_root() -> pathlib.Path:
|
||||
# tests/diffusion/models/hunyuan_image3/test_prompt_utils.py -> repo root
|
||||
return pathlib.Path(__file__).resolve().parents[4]
|
||||
|
||||
|
||||
def test_end2end_routes_through_shared_prompt_utils():
|
||||
"""Regression for the *delivery vector* of PR #3243.
|
||||
|
||||
Background: the wrong-template bug that PR #3243 fixes was introduced
|
||||
when end2end.py grew its own hand-rolled prompt builder that diverged
|
||||
from the canonical instruct chat template. To prevent that exact
|
||||
failure mode from recurring, end2end.py MUST:
|
||||
1. Import the prompt builders from the shared prompt_utils module.
|
||||
2. NOT redefine `build_prompt` or `build_prompt_tokens` locally.
|
||||
|
||||
A local redefinition is precisely how a future merge can silently
|
||||
re-introduce a pretrain-style template (trigger BEFORE user_prompt,
|
||||
no User:/Assistant: framing, etc.) without touching prompt_utils,
|
||||
bypassing every other test in this file.
|
||||
"""
|
||||
end2end_path = _repo_root() / "examples" / "offline_inference" / "hunyuan_image3" / "end2end.py"
|
||||
assert end2end_path.is_file(), f"end2end.py not found at {end2end_path}"
|
||||
|
||||
tree = ast.parse(end2end_path.read_text(encoding="utf-8"))
|
||||
|
||||
local_func_names = {n.name for n in ast.walk(tree) if isinstance(n, ast.FunctionDef)}
|
||||
forbidden = {"build_prompt", "build_prompt_tokens"}
|
||||
redefined = local_func_names & forbidden
|
||||
assert not redefined, (
|
||||
f"end2end.py defines {sorted(redefined)} locally. This is exactly how "
|
||||
"the wrong prompt template re-entered the example before PR #3243. "
|
||||
"Use the shared `vllm_omni.diffusion.models.hunyuan_image3.prompt_utils` "
|
||||
"helpers instead."
|
||||
)
|
||||
assert not (local_func_names & {"build_prompt", "build_prompt_tokens"})
|
||||
|
||||
imported_from_prompt_utils: set[str] = set()
|
||||
for node in ast.walk(tree):
|
||||
if isinstance(node, ast.ImportFrom) and node.module and node.module.endswith("hunyuan_image3.prompt_utils"):
|
||||
imported_from_prompt_utils.update(alias.name for alias in node.names)
|
||||
expected_imports = {
|
||||
"_TASK_PRESETS",
|
||||
"build_prompt_tokens",
|
||||
"resolve_stop_token_ids",
|
||||
}
|
||||
assert expected_imports <= imported_from_prompt_utils, (
|
||||
"end2end.py must import the HunyuanImage3 prompt and stop-token helpers from "
|
||||
"vllm_omni.diffusion.models.hunyuan_image3.prompt_utils -- the shared "
|
||||
"module is the single source of truth for the AR-prefill template and "
|
||||
"bot_task-derived AR stop token ids."
|
||||
)
|
||||
|
||||
|
||||
# -------------------- Real-tokenizer regression --------------------
|
||||
expected_imports = {"build_prompt_tokens", "resolve_stop_token_ids", "resolve_sys_type"}
|
||||
assert expected_imports <= imported_from_prompt_utils
|
||||
|
||||
|
||||
_HUNYUAN_MODEL_ID = "tencent/HunyuanImage-3.0-Instruct"
|
||||
@@ -290,41 +268,14 @@ def _hf_cached(model_id: str) -> bool:
|
||||
return os.path.isdir(snap_dir) and any(os.scandir(snap_dir))
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not _hf_cached(_HUNYUAN_MODEL_ID),
|
||||
reason=f"{_HUNYUAN_MODEL_ID} tokenizer not in HF cache",
|
||||
)
|
||||
@pytest.mark.skipif(not _hf_cached(_HUNYUAN_MODEL_ID), reason=f"{_HUNYUAN_MODEL_ID} tokenizer not in HF cache")
|
||||
def test_segment_tokenize_diverges_from_full_string_encode():
|
||||
"""Regression for PR #3243 segment-tokenization fix.
|
||||
|
||||
The naive `tokenizer.encode(build_prompt(...))` lets BPE merge tokens
|
||||
across segment boundaries (notably `。\\n\\n` -> a single id), which
|
||||
drifts the AR prefill away from HF's apply_chat_template output. The
|
||||
segment-by-segment build_prompt_tokens must produce a STRICTLY
|
||||
DIFFERENT id sequence on a prompt that triggers the merge.
|
||||
|
||||
If someone "simplifies" build_prompt_tokens to call encode() on the
|
||||
full string, this assertion fires.
|
||||
"""
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
tok = AutoTokenizer.from_pretrained(_HUNYUAN_MODEL_ID, trust_remote_code=True)
|
||||
|
||||
user_prompt = "写一首关于夜的诗。"
|
||||
result = build_prompt_tokens(user_prompt, tok, task="i2t")
|
||||
result = build_prompt_tokens(user_prompt, tok, task="i2t", bot_task=None)
|
||||
seg_ids = result.token_ids
|
||||
full_ids = tok.encode(build_prompt(user_prompt, task="i2t"), add_special_tokens=False)
|
||||
|
||||
assert seg_ids != full_ids, (
|
||||
"build_prompt_tokens output equals naive full-string encode -- "
|
||||
"the BPE-merge-bypass behavior is no longer exercised. This means "
|
||||
"the segment-by-segment fix from PR #3243 has been silently undone."
|
||||
)
|
||||
|
||||
# Segmenting prevents merges, so the segment id list should have AT LEAST
|
||||
# as many tokens as the merged version (a merge consumes 2+ ids -> 1).
|
||||
assert len(seg_ids) >= len(full_ids), (
|
||||
f"segment-encoded length ({len(seg_ids)}) shorter than full-string "
|
||||
f"merged length ({len(full_ids)}) -- impossible if segmentation is "
|
||||
f"genuinely bypassing merges."
|
||||
)
|
||||
full_ids = tok.encode(build_prompt(user_prompt, task="i2t", bot_task=None), add_special_tokens=False)
|
||||
assert seg_ids != full_ids
|
||||
assert len(seg_ids) >= len(full_ids)
|
||||
@@ -93,7 +93,13 @@ def _run(stage_config_path: str, output_path: Path) -> tuple[Image.Image, str, f
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
|
||||
result = build_prompt_tokens(PROMPT, tokenizer, task="it2i_recaption", sys_type="en_unified")
|
||||
result = build_prompt_tokens(
|
||||
PROMPT,
|
||||
tokenizer,
|
||||
task="it2i",
|
||||
bot_task="recaption",
|
||||
sys_type="en_unified",
|
||||
)
|
||||
token_ids = result.token_ids
|
||||
system_prompt_type = result.system_prompt_type
|
||||
|
||||
|
||||
@@ -1349,8 +1349,16 @@ def test_image_edit_parameter_default(async_omni_test_client):
|
||||
engine = async_omni_test_client.app.state.engine_client
|
||||
captured_sampling_params = engine.captured_sampling_params_list[-1]
|
||||
|
||||
assert captured_sampling_params.width == 24
|
||||
assert captured_sampling_params.height == 16
|
||||
# size="auto" on multi-stage pipelines deliberately leaves the diffusion
|
||||
# stages sampling_params width/height unset so AR-driven pipelines (e.g.
|
||||
# HunyuanImage-3.0) can let ar2diffusion override the final bucket from
|
||||
# the AR-predicted ratio token; see
|
||||
# test_image_edits_size_auto_preserves_bridge_size for the contract.
|
||||
# Single-stage diffusion (test_image_edit_parameter_default_single_stage)
|
||||
# still pins width/height to the input image size via api_servers
|
||||
# gen_params, which is unchanged.
|
||||
assert captured_sampling_params.width is None
|
||||
assert captured_sampling_params.height is None
|
||||
assert captured_sampling_params.num_outputs_per_prompt == 1
|
||||
assert captured_sampling_params.num_inference_steps == 4
|
||||
assert captured_sampling_params.guidance_scale == 7.5
|
||||
@@ -1649,3 +1657,59 @@ def test_extract_images_from_result():
|
||||
assert len(images) == 1
|
||||
assert isinstance(images[0], Image.Image)
|
||||
assert images[0].size == (32, 32)
|
||||
|
||||
|
||||
def test_image_edits_size_auto_preserves_bridge_size(async_omni_stage_configs_only_client):
|
||||
"""size=auto must NOT pin the diffusion stage sampling_params.height/width.
|
||||
|
||||
Regression: prior to the fix, edit_images resolved size=auto to the
|
||||
first input image dimensions and forwarded them through gen_params +
|
||||
extra_body to the diffusion stages sampling_params. AR-driven
|
||||
pipelines (e.g. HunyuanImage-3.0) rely on ar2diffusions
|
||||
bridge to override the final bucket via the AR-predicted ratio token,
|
||||
and the DiT pre_process_func only fills sampling_params from the
|
||||
bridge value when sampling_params.width is None (see
|
||||
pipeline_hunyuan_image3.py:290). Non-None width from the input image
|
||||
silently suppressed the AR decision, producing the wrong bucket
|
||||
(e.g. 1024x1024 square instead of the AR-decided 1280x720 landscape
|
||||
for multi-image fusion).
|
||||
|
||||
Cross-pins the multi-image fix at the API level: 2 reference images
|
||||
with bot_task=think must produce 2 <img> placeholders in the captured
|
||||
AR prompt (build_prompt called with num_images=2).
|
||||
"""
|
||||
img_a = make_test_image_bytes((32, 32))
|
||||
img_b = make_test_image_bytes((128, 64))
|
||||
response = async_omni_stage_configs_only_client.post(
|
||||
"/v1/images/edits",
|
||||
files=[("image", img_a), ("image", img_b)],
|
||||
data={
|
||||
"prompt": "fuse",
|
||||
"size": "auto",
|
||||
"bot_task": "think",
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
|
||||
engine = async_omni_stage_configs_only_client.app.state.engine_client
|
||||
captured = engine.captured_sampling_params_list
|
||||
assert captured is not None
|
||||
assert len(captured) == 2
|
||||
|
||||
diffusion_params = captured[1]
|
||||
assert diffusion_params.height is None, (
|
||||
f"size=auto leaked into diffusion sampling_params.height={diffusion_params.height}; "
|
||||
"must stay None so AR-driven pipelines can apply the bridges decision."
|
||||
)
|
||||
assert diffusion_params.width is None, (
|
||||
f"size=auto leaked into diffusion sampling_params.width={diffusion_params.width}; "
|
||||
"must stay None so AR-driven pipelines can apply the bridges decision."
|
||||
)
|
||||
|
||||
KEY = "prompt"
|
||||
IMG = "<img>"
|
||||
captured_prompt = engine.captured_prompt
|
||||
if isinstance(captured_prompt, dict) and isinstance(captured_prompt.get("prompt"), str):
|
||||
assert captured_prompt["prompt"].count("<img>") == 2, (
|
||||
f"N=2 reference images must emit 2 <img> placeholders in AR prompt; got {captured_prompt[KEY].count(IMG)} -- prompt: {captured_prompt[KEY]!r}"
|
||||
)
|
||||
@@ -91,3 +91,275 @@ def test_build_multistage_generation_inputs_applies_stage_specific_overrides(ser
|
||||
assert engine.default_sampling_params_list[1].lora_request is None
|
||||
assert engine.default_sampling_params_list[2].resolution == 640
|
||||
assert engine.default_sampling_params_list[2].lora_request is None
|
||||
|
||||
|
||||
def test_build_multistage_generation_inputs_multi_image_emits_n_img_placeholders(serving_chat):
|
||||
"""N reference images with bot_task set must emit N <img> placeholders.
|
||||
|
||||
Regression: prior to the multi-image online fix, build_prompt was
|
||||
called without num_images, defaulting to 1. A 2-image edit request
|
||||
would only get a single <img> placeholder in the AR prompt; vLLMs
|
||||
_process_multimodal then raised
|
||||
AssertionError(Failed to apply prompt replacement for mm_items[image][1])
|
||||
when trying to replace the second image (no placeholder left for it).
|
||||
|
||||
Pins the contract that build_prompt() is invoked with the actual image
|
||||
count so multi-image IT2I is wired correctly through the online
|
||||
/v1/images/edits path.
|
||||
"""
|
||||
from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
|
||||
|
||||
engine = SimpleNamespace(
|
||||
stage_configs=[
|
||||
SimpleNamespace(stage_type="llm", is_comprehension=True),
|
||||
SimpleNamespace(stage_type="diffusion", is_comprehension=False),
|
||||
],
|
||||
default_sampling_params_list=[
|
||||
SamplingParams(temperature=0.0),
|
||||
OmniDiffusionSamplingParams(),
|
||||
],
|
||||
)
|
||||
IMG = "<img>"
|
||||
images = [Image.new("RGB", (32, 32), color="red") for _ in range(3)]
|
||||
|
||||
for n in (1, 2, 3):
|
||||
engine_prompt, _ = OmniOpenAIServingChat._build_multistage_generation_inputs(
|
||||
serving_chat,
|
||||
engine=engine,
|
||||
prompt="edit me",
|
||||
extra_body={"bot_task": "think"},
|
||||
reference_images=images[:n],
|
||||
gen_params=OmniDiffusionSamplingParams(),
|
||||
)
|
||||
prompt_str = engine_prompt["prompt"]
|
||||
assert prompt_str.count("<img>") == n, (
|
||||
f"N={n}: expected {n} <img> placeholders, got {prompt_str.count(IMG)} -- prompt: {prompt_str!r}"
|
||||
)
|
||||
|
||||
|
||||
def test_build_multistage_generation_inputs_tokenizer_path_emits_prompt_token_ids(serving_chat):
|
||||
"""When a tokenizer is provided, the helper must emit HF byte-for-byte
|
||||
prompt_token_ids and forward use_system_prompt to the engine prompt.
|
||||
|
||||
Regression: prior to the HF-byte-equivalent fix, online IT2I always
|
||||
passed the prompt as a single string. The engine then BPE-merged across
|
||||
chat-template segment boundaries (e.g. user_prompt-ending punctuation
|
||||
plus the trailing \n\n before \"Assistant: \") producing a token
|
||||
sequence that differs from HF apply_chat_template / offline
|
||||
end2end.py. AR generated different cot_text (706 tokens / 1190 chars
|
||||
vs offline 661 / 1118 for the same inputs) and DiT produced a visually
|
||||
different image (yin-yang on brushed-metal vs three-blue swirl on
|
||||
canvas) under the same seed.
|
||||
|
||||
Pins:
|
||||
1. engine_prompt[\"prompt_token_ids\"] is set when tokenizer is passed.
|
||||
2. engine_prompt[\"prompt\"] stays as the raw user prompt -- the DiT
|
||||
side rebuilds its own system prefix via use_system_prompt.
|
||||
3. engine_prompt[\"use_system_prompt\"] == \"en_unified\" so
|
||||
ar2diffusion forwards the matching system prompt to DiT.
|
||||
4. N reference images emit N <img> token ids in the AR sequence.
|
||||
"""
|
||||
from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
|
||||
|
||||
# Minimal FakeTokenizer mirroring tests/diffusion/.../test_hunyuan_image3_it2i_multi_image.py
|
||||
class FakeTokenizer:
|
||||
SPECIAL = {
|
||||
"<|startoftext|>": 1,
|
||||
"<img>": 2,
|
||||
"<think>": 3,
|
||||
"<recaption>": 4,
|
||||
}
|
||||
|
||||
def convert_tokens_to_ids(self, tok: str) -> int:
|
||||
return self.SPECIAL.get(tok, 0)
|
||||
|
||||
def encode(self, text: str, add_special_tokens: bool = False) -> list[int]:
|
||||
return list(range(100, 100 + len(text)))
|
||||
|
||||
engine = SimpleNamespace(
|
||||
stage_configs=[
|
||||
SimpleNamespace(stage_type="llm", is_comprehension=True),
|
||||
SimpleNamespace(stage_type="diffusion", is_comprehension=False),
|
||||
],
|
||||
default_sampling_params_list=[
|
||||
SamplingParams(temperature=0.0),
|
||||
OmniDiffusionSamplingParams(),
|
||||
],
|
||||
)
|
||||
PROMPT_KEY = "prompt"
|
||||
USP_KEY = "use_system_prompt"
|
||||
images = [Image.new("RGB", (32, 32), color="red") for _ in range(3)]
|
||||
|
||||
for n in (1, 2, 3):
|
||||
tok = FakeTokenizer()
|
||||
engine_prompt, _ = OmniOpenAIServingChat._build_multistage_generation_inputs(
|
||||
serving_chat,
|
||||
engine=engine,
|
||||
prompt="edit me",
|
||||
extra_body={"bot_task": "think"},
|
||||
reference_images=images[:n],
|
||||
gen_params=OmniDiffusionSamplingParams(),
|
||||
tokenizer=tok,
|
||||
)
|
||||
# (1) prompt_token_ids must be set and non-empty
|
||||
assert "prompt_token_ids" in engine_prompt, f"N={n}: prompt_token_ids missing"
|
||||
token_ids = engine_prompt["prompt_token_ids"]
|
||||
assert isinstance(token_ids, list) and len(token_ids) > 0, f"N={n}: prompt_token_ids empty"
|
||||
# (2) raw prompt preserved (DiT bridge needs raw user text)
|
||||
assert engine_prompt["prompt"] == "edit me", (
|
||||
f"N={n}: prompt must stay raw user text, got {engine_prompt[PROMPT_KEY]!r}"
|
||||
)
|
||||
# (3) use_system_prompt forwarded for ar2diffusion bridge
|
||||
assert engine_prompt.get("use_system_prompt") == "en_unified", (
|
||||
f"N={n}: use_system_prompt must be en_unified, got {engine_prompt.get(USP_KEY)!r}"
|
||||
)
|
||||
# (4) N <img> token ids (id=2 in FakeTokenizer)
|
||||
img_count = token_ids.count(2)
|
||||
assert img_count == n, f"N={n}: expected {n} <img> token ids in prompt_token_ids, got {img_count}"
|
||||
|
||||
|
||||
def test_build_multistage_generation_inputs_bot_task_semantic_changes_trigger_and_sys(serving_chat):
|
||||
"""Passing bot_task=think_recaption (vs default "think") must flip the
|
||||
resolved sys_type to en_think_recaption (and trigger tag is still
|
||||
<think>). Pins that the API actually plumbs the bot_task semantic
|
||||
through to build_prompt rather than ignoring it.
|
||||
"""
|
||||
from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
|
||||
|
||||
engine = SimpleNamespace(
|
||||
stage_configs=[
|
||||
SimpleNamespace(stage_type="llm", is_comprehension=True),
|
||||
SimpleNamespace(stage_type="diffusion", is_comprehension=False),
|
||||
],
|
||||
default_sampling_params_list=[
|
||||
SamplingParams(temperature=0.0),
|
||||
OmniDiffusionSamplingParams(),
|
||||
],
|
||||
)
|
||||
images = [Image.new("RGB", (32, 32), color="red")]
|
||||
|
||||
# Default bot_task (think) -> en_unified system prompt baked into the
|
||||
# legacy string path. Use legacy build_prompt (tokenizer=None) so the
|
||||
# rendered prompt is a string we can grep.
|
||||
think_prompt, _ = OmniOpenAIServingChat._build_multistage_generation_inputs(
|
||||
serving_chat,
|
||||
engine=engine,
|
||||
prompt="edit me",
|
||||
extra_body={"task": "it2i", "bot_task": "think"},
|
||||
reference_images=images,
|
||||
gen_params=OmniDiffusionSamplingParams(),
|
||||
)
|
||||
# think_recaption -> en_think_recaption system prompt (different content).
|
||||
recap_prompt, _ = OmniOpenAIServingChat._build_multistage_generation_inputs(
|
||||
serving_chat,
|
||||
engine=engine,
|
||||
prompt="edit me",
|
||||
extra_body={"task": "it2i", "bot_task": "think_recaption"},
|
||||
reference_images=images,
|
||||
gen_params=OmniDiffusionSamplingParams(),
|
||||
)
|
||||
assert think_prompt["prompt"] != recap_prompt["prompt"], (
|
||||
"bot_task semantic must change the rendered system prompt: "
|
||||
f"think/think_recaption produced identical strings (len={len(think_prompt['prompt'])})"
|
||||
)
|
||||
|
||||
|
||||
def test_build_multistage_generation_inputs_sys_type_override(serving_chat):
|
||||
"""Caller-supplied sys_type must override the bot_task-derived default.
|
||||
Mirrors offline `--bot-task think_recaption --sys-type en_unified`
|
||||
where the user wants think_recaptions trigger but the unified system
|
||||
prompt body.
|
||||
"""
|
||||
from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
|
||||
|
||||
engine = SimpleNamespace(
|
||||
stage_configs=[
|
||||
SimpleNamespace(stage_type="llm", is_comprehension=True),
|
||||
SimpleNamespace(stage_type="diffusion", is_comprehension=False),
|
||||
],
|
||||
default_sampling_params_list=[
|
||||
SamplingParams(temperature=0.0),
|
||||
OmniDiffusionSamplingParams(),
|
||||
],
|
||||
)
|
||||
images = [Image.new("RGB", (32, 32), color="red")]
|
||||
|
||||
# think_recaption defaults sys_type -> en_think_recaption.
|
||||
default_sys, _ = OmniOpenAIServingChat._build_multistage_generation_inputs(
|
||||
serving_chat,
|
||||
engine=engine,
|
||||
prompt="edit me",
|
||||
extra_body={"task": "it2i", "bot_task": "think_recaption"},
|
||||
reference_images=images,
|
||||
gen_params=OmniDiffusionSamplingParams(),
|
||||
)
|
||||
# sys_type=en_unified overrides -> same system body as bot_task=think.
|
||||
overridden, _ = OmniOpenAIServingChat._build_multistage_generation_inputs(
|
||||
serving_chat,
|
||||
engine=engine,
|
||||
prompt="edit me",
|
||||
extra_body={"task": "it2i", "bot_task": "think_recaption", "sys_type": "en_unified"},
|
||||
reference_images=images,
|
||||
gen_params=OmniDiffusionSamplingParams(),
|
||||
)
|
||||
plain_think, _ = OmniOpenAIServingChat._build_multistage_generation_inputs(
|
||||
serving_chat,
|
||||
engine=engine,
|
||||
prompt="edit me",
|
||||
extra_body={"task": "it2i", "bot_task": "think"},
|
||||
reference_images=images,
|
||||
gen_params=OmniDiffusionSamplingParams(),
|
||||
)
|
||||
|
||||
# Override must (a) differ from the no-override default, and (b) equal
|
||||
# the prompt that bot_task=think produces (both end up with
|
||||
# en_unified system body + <think> trigger).
|
||||
assert overridden["prompt"] != default_sys["prompt"], (
|
||||
"sys_type override must change the rendered prompt body vs the bot_task default"
|
||||
)
|
||||
assert overridden["prompt"] == plain_think["prompt"], (
|
||||
"sys_type=en_unified + bot_task=think_recaption must produce the same prompt as "
|
||||
"bot_task=think (both = en_unified system body + <think> trigger)"
|
||||
)
|
||||
|
||||
|
||||
def test_build_multistage_generation_inputs_custom_system_prompt(serving_chat):
|
||||
"""`extra_body["system_prompt"]` must reach build_prompt as
|
||||
`custom_system_prompt`, enabling sys_type="custom" callers to inject
|
||||
a verbatim system body. Without this plumbing the sys_type="custom"
|
||||
branch in get_system_prompt() returns None and silently drops the
|
||||
user-supplied content.
|
||||
"""
|
||||
from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
|
||||
|
||||
engine = SimpleNamespace(
|
||||
stage_configs=[
|
||||
SimpleNamespace(stage_type="llm", is_comprehension=True),
|
||||
SimpleNamespace(stage_type="diffusion", is_comprehension=False),
|
||||
],
|
||||
default_sampling_params_list=[
|
||||
SamplingParams(temperature=0.0),
|
||||
OmniDiffusionSamplingParams(),
|
||||
],
|
||||
)
|
||||
images = [Image.new("RGB", (32, 32), color="red")]
|
||||
|
||||
marker = "ZZZ_CUSTOM_SYSTEM_PROMPT_MARKER_ZZZ"
|
||||
|
||||
out, _ = OmniOpenAIServingChat._build_multistage_generation_inputs(
|
||||
serving_chat,
|
||||
engine=engine,
|
||||
prompt="edit me",
|
||||
extra_body={
|
||||
"task": "it2i",
|
||||
"bot_task": "think",
|
||||
"sys_type": "custom",
|
||||
"system_prompt": marker,
|
||||
},
|
||||
reference_images=images,
|
||||
gen_params=OmniDiffusionSamplingParams(),
|
||||
)
|
||||
assert marker in out["prompt"], (
|
||||
f"custom system_prompt content must reach the rendered prompt; "
|
||||
f"marker {marker!r} not found in prompt of length {len(out['prompt'])}"
|
||||
)
|
||||
@@ -0,0 +1,92 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for HunyuanImage3 stage input processor."""
|
||||
|
||||
import builtins
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import (
|
||||
HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS,
|
||||
)
|
||||
from vllm_omni.model_executor.stage_input_processors.hunyuan_image3 import (
|
||||
_extract_ratio_index,
|
||||
_truncate_at_cot_end,
|
||||
ar2diffusion,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _source_output(token_ids: list[int], text: str = ""):
|
||||
return SimpleNamespace(
|
||||
outputs=[
|
||||
SimpleNamespace(
|
||||
token_ids=token_ids,
|
||||
cumulative_token_ids=token_ids,
|
||||
text=text,
|
||||
)
|
||||
],
|
||||
multimodal_output=None,
|
||||
)
|
||||
|
||||
|
||||
def test_extract_ratio_index_uses_fixed_special_token_ids():
|
||||
ratio_33 = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_33>"]
|
||||
ratio_36 = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_36>"]
|
||||
|
||||
assert _extract_ratio_index([1, ratio_33, 2]) == 33
|
||||
assert _extract_ratio_index([1, ratio_33, 2, ratio_36]) == 36
|
||||
|
||||
|
||||
def test_truncate_at_cot_end_strips_tail_after_recaption_marker():
|
||||
text = _truncate_at_cot_end("body text</recaption><answer><boi><img_size_1024><img_ratio_0>")
|
||||
assert text == "body text</recaption>"
|
||||
|
||||
|
||||
def test_ar2diffusion_applies_ratio_and_truncates_tail_without_tokenizer(monkeypatch: pytest.MonkeyPatch):
|
||||
real_import = builtins.__import__
|
||||
|
||||
def _block_transformers_import(name, *args, **kwargs):
|
||||
if name == "transformers" or name.startswith("transformers."):
|
||||
raise AssertionError("ar2diffusion must not import transformers on the bridge path")
|
||||
return real_import(name, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(builtins, "__import__", _block_transformers_import)
|
||||
|
||||
end_recaption = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["</recaption>"]
|
||||
answer = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<answer>"]
|
||||
boi = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<boi>"]
|
||||
size = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_size_1024>"]
|
||||
ratio_0 = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_0>"]
|
||||
token_ids = [100, 101, end_recaption, answer, boi, size, ratio_0]
|
||||
|
||||
result = ar2diffusion(
|
||||
[_source_output(token_ids, text="decoded without special tokens")],
|
||||
prompt=[{"prompt": "edit", "height": 64, "width": 64}],
|
||||
)
|
||||
|
||||
assert len(result) == 1
|
||||
assert (result[0]["height"], result[0]["width"]) == (512, 2048)
|
||||
assert result[0]["extra"]["ar_generated_text"] == "decoded without special tokens"
|
||||
assert "ar_token_ids" not in result[0]["extra"]
|
||||
|
||||
|
||||
def test_ar2diffusion_forwards_custom_system_prompt_body():
|
||||
end_think = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["</think>"]
|
||||
marker = "CUSTOM_SYSTEM_BODY"
|
||||
|
||||
result = ar2diffusion(
|
||||
[_source_output([100, end_think], text="thought</think>")],
|
||||
prompt=[
|
||||
{
|
||||
"prompt": "edit",
|
||||
"use_system_prompt": "custom",
|
||||
"system_prompt": marker,
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
assert result[0]["use_system_prompt"] == "custom"
|
||||
assert result[0]["system_prompt"] == marker
|
||||
Reference in new issue
Block a user