[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 <liuhongsheng4@huawei.com>

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["<img>"] in ids, "<img> 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["<img>"] not in ids, "<img> 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["<think>"],
         FakeTokenizer.SPECIAL["<recaption>"],
@@ -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, (

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,
             }