From 0d91fbbbb7fe8b0e3b59e35403d2d2123969ae3f Mon Sep 17 00:00:00 2001 From: dengyunyang <584797741@qq.com> Date: Tue, 12 May 2026 12:13:30 +0800 Subject: [PATCH] [Bugfix] Align the AR and DiT prompt formatting across both online and offline modes. (#3516) Signed-off-by: dengyunyang <584797741@qq.com> Co-authored-by: Hongsheng Liu --- .../hunyuan_image3/end2end.py | 6 ++--- .../hunyuan_image3/test_prompt_utils.py | 15 ++++++++---- tests/e2e/accuracy/test_hunyuan_image3.py | 6 +++-- .../models/hunyuan_image3/prompt_utils.py | 24 ++++++++++++++++--- vllm_omni/entrypoints/openai/serving_chat.py | 16 +++++++++---- 5 files changed, 49 insertions(+), 18 deletions(-) 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