diff --git a/examples/offline_inference/hunyuan_image3/end2end.py b/examples/offline_inference/hunyuan_image3/end2end.py
index 0f41150b6..8d29b0f58 100644
--- a/examples/offline_inference/hunyuan_image3/end2end.py
+++ b/examples/offline_inference/hunyuan_image3/end2end.py
@@ -192,9 +192,9 @@ def main():
# Format prompts
formatted_prompts: list[OmniPromptType] = []
for p in prompts:
- token_ids = build_prompt_tokens(p, tokenizer, task=task, sys_type=args.sys_type)
- preset_sys_type, _, _ = _TASK_PRESETS[task]
- effective_sys_type = args.sys_type or preset_sys_type
+ result = build_prompt_tokens(p, tokenizer, task=task, sys_type=args.sys_type)
+ token_ids = result.token_ids
+ effective_sys_type = result.system_prompt_type
# `prompt_token_ids` drives the AR stage (matches HF byte-for-byte).
# `prompt` and `use_system_prompt` are forwarded by ar2diffusion to
diff --git a/tests/diffusion/models/hunyuan_image3/test_prompt_utils.py b/tests/diffusion/models/hunyuan_image3/test_prompt_utils.py
index f899e5a65..9972d328d 100644
--- a/tests/diffusion/models/hunyuan_image3/test_prompt_utils.py
+++ b/tests/diffusion/models/hunyuan_image3/test_prompt_utils.py
@@ -180,14 +180,16 @@ def test_build_prompt_tokens_segments_each_boundary():
def test_build_prompt_tokens_image_placeholder_present_for_image_tasks():
tok = FakeTokenizer()
- ids = build_prompt_tokens("hi", tok, task="i2t")
+ result = build_prompt_tokens("hi", tok, task="i2t")
+ ids = result.token_ids
assert ids[0] == FakeTokenizer.SPECIAL["<|startoftext|>"], "BOS (<|startoftext|>) must be the first token"
assert FakeTokenizer.SPECIAL["
"] in ids, "
placeholder must be present for i2t/it2i tasks"
def test_build_prompt_tokens_no_image_for_text_only_tasks():
tok = FakeTokenizer()
- ids = build_prompt_tokens("hi", tok, task="t2t")
+ result = build_prompt_tokens("hi", tok, task="t2t")
+ ids = result.token_ids
assert FakeTokenizer.SPECIAL["
"] not in ids, "
must NOT appear for text-only tasks"
@@ -203,14 +205,16 @@ def test_build_prompt_tokens_no_image_for_text_only_tasks():
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)."""
tok = FakeTokenizer()
- ids = build_prompt_tokens("hi", tok, task=task)
+ result = build_prompt_tokens("hi", tok, task=task)
+ ids = result.token_ids
assert 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()
- ids = build_prompt_tokens("hi", tok, task="t2t")
+ result = build_prompt_tokens("hi", tok, task="t2t")
+ ids = result.token_ids
assert ids[-1] not in {
FakeTokenizer.SPECIAL[""],
FakeTokenizer.SPECIAL[""],
@@ -305,7 +309,8 @@ def test_segment_tokenize_diverges_from_full_string_encode():
tok = AutoTokenizer.from_pretrained(_HUNYUAN_MODEL_ID, trust_remote_code=True)
user_prompt = "写一首关于夜的诗。"
- seg_ids = build_prompt_tokens(user_prompt, tok, task="i2t")
+ result = build_prompt_tokens(user_prompt, tok, task="i2t")
+ seg_ids = result.token_ids
full_ids = tok.encode(build_prompt(user_prompt, task="i2t"), add_special_tokens=False)
assert seg_ids != full_ids, (
diff --git a/tests/e2e/accuracy/test_hunyuan_image3.py b/tests/e2e/accuracy/test_hunyuan_image3.py
index 2c96ee479..93671e7bb 100644
--- a/tests/e2e/accuracy/test_hunyuan_image3.py
+++ b/tests/e2e/accuracy/test_hunyuan_image3.py
@@ -93,7 +93,9 @@ 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)
- token_ids = build_prompt_tokens(PROMPT, tokenizer, task="it2i_recaption", sys_type="en_unified")
+ result = build_prompt_tokens(PROMPT, tokenizer, task="it2i_recaption", sys_type="en_unified")
+ token_ids = result.token_ids
+ system_prompt_type = result.system_prompt_type
with OmniRunner(MODEL_NAME, stage_configs_path=stage_config_path) as runner:
params_list = list(runner.omni.default_sampling_params_list)
@@ -108,7 +110,7 @@ def _run(stage_config_path: str, output_path: Path) -> tuple[Image.Image, str, f
{
"prompt_token_ids": token_ids,
"prompt": PROMPT,
- "use_system_prompt": "en_unified",
+ "use_system_prompt": system_prompt_type,
"height": HEIGHT,
"width": WIDTH,
}
diff --git a/vllm_omni/diffusion/models/hunyuan_image3/prompt_utils.py b/vllm_omni/diffusion/models/hunyuan_image3/prompt_utils.py
index bfc0146d8..3cd679b68 100644
--- a/vllm_omni/diffusion/models/hunyuan_image3/prompt_utils.py
+++ b/vllm_omni/diffusion/models/hunyuan_image3/prompt_utils.py
@@ -17,6 +17,7 @@ canonical mapping for both flows.
from __future__ import annotations
+from dataclasses import dataclass
from typing import Any
from .system_prompt import get_system_prompt
@@ -128,13 +129,19 @@ def build_prompt(
return "".join(parts)
+@dataclass
+class PromptTokensResult:
+ token_ids: list[int] # The tokenized prompt
+ system_prompt_type: str # The effective system prompt type used
+
+
def build_prompt_tokens(
user_prompt: str,
tokenizer,
task: str = "it2i_think",
sys_type: str | None = None,
custom_system_prompt: str | None = None,
-) -> list[int]:
+) -> PromptTokensResult:
"""Segment-by-segment tokenization that matches HF apply_chat_template.
Calling tokenizer.encode(build_prompt(...)) on the full string lets BPE
@@ -144,6 +151,9 @@ def build_prompt_tokens(
each segment independently and concatenates token_ids, so no cross-
boundary merge happens. We replicate that here and feed the result to
Omni via OmniTokensPrompt (prompt_token_ids).
+
+ Returns:
+ PromptTokensResult
"""
if task not in _TASK_PRESETS:
raise ValueError(f"Unknown task {task!r}. Choose from: {available_tasks()}")
@@ -162,7 +172,11 @@ def build_prompt_tokens(
# protect, fall back to whole-string encode.
if task == "t2i_vanilla":
s = build_prompt(user_prompt, task, sys_type, custom_system_prompt)
- return tokenizer.encode(s, add_special_tokens=False)
+ token_ids = tokenizer.encode(s, add_special_tokens=False)
+ return PromptTokensResult(
+ token_ids=token_ids,
+ system_prompt_type=effective_sys_type,
+ )
system_prompt = get_system_prompt(effective_sys_type, preset_bot_task, custom_system_prompt)
# Do NOT strip -- HF apply_chat_template keeps the system prompt's
@@ -180,7 +194,11 @@ def build_prompt_tokens(
ids += tokenizer.encode("\n\nAssistant: ", add_special_tokens=False)
if trig_id is not None:
ids += [trig_id]
- return ids
+
+ return PromptTokensResult(
+ token_ids=ids,
+ system_prompt_type=effective_sys_type,
+ )
__all__ = ["build_prompt", "build_prompt_tokens", "resolve_stop_token_ids", _TASK_PRESETS]
diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py
index 24d2e6592..99827454e 100644
--- a/vllm_omni/entrypoints/openai/serving_chat.py
+++ b/vllm_omni/entrypoints/openai/serving_chat.py
@@ -2255,26 +2255,32 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
else:
engine_prompt_data = {"image": reference_images}
- prompt_token_ids: list[int] | None = None
+ engine_prompt: OmniTextPrompt = {"prompt": prompt}
if bot_task:
from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import (
build_prompt,
build_prompt_tokens,
)
+ prompt_token_ids: list[int] | None = None
+ system_prompt_type: str | None = None
if tokenizer is not None:
- prompt_token_ids = build_prompt_tokens(prompt, tokenizer, task=bot_task)
+ result = build_prompt_tokens(prompt, tokenizer, task=bot_task)
+ prompt_token_ids = result.token_ids
+ system_prompt_type = result.system_prompt_type
else:
prompt = build_prompt(prompt, task=bot_task)
+ engine_prompt["prompt"] = prompt
if reference_images and len(reference_images) == 1:
engine_prompt_data = {"image": reference_images[0]}
modalities = ["image"]
+ if prompt_token_ids is not None:
+ engine_prompt["prompt_token_ids"] = prompt_token_ids
+ if system_prompt_type is not None:
+ engine_prompt["use_system_prompt"] = system_prompt_type
- engine_prompt: OmniTextPrompt = {"prompt": prompt}
engine_prompt["modalities"] = modalities
- if prompt_token_ids is not None:
- engine_prompt["prompt_token_ids"] = prompt_token_ids
if negative_prompt is not None:
engine_prompt["negative_prompt"] = negative_prompt