[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)
|
||||
Reference in new issue
Block a user