[Bugfix][HunyuanImage3] Fix offline AR garbage output by switching to Instruct chat template (#3243)

Signed-off-by: zuiho <wu15922848573@outlook.com>
Signed-off-by: TaffyOfficial <2324465096@qq.com>
Signed-off-by: zuiho-kai <31877877+zuiho-kai@users.noreply.github.com>
Signed-off-by: zuiho <2324465096@qq.com>
Co-authored-by: TaffyOfficial <2324465096@qq.com>
Co-authored-by: zuiho-kai <31877877+zuiho-kai@users.noreply.github.com>

vllm_omni/diffusion/models/hunyuan_image3/prompt_utils.py

@@ -0,0 +1,152 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Shared prompt-template construction for HunyuanImage-3.0-Instruct.
+
+Single source of truth for the AR-prefill prompt format used by the
+example scripts and any downstream caller that needs to build
+HunyuanImage3 chat-template token sequences without invoking the full
+diffusion pipeline tokenizer wrapper.
+
+The DiT pipeline (`pipeline_hunyuan_image3.py`) builds prompts through
+`TokenizerWrapper.apply_chat_template`, which eagerly consumes
+`JointImageInfo` objects produced by image preprocessing. The example
+flow uses an `<img>` placeholder + `multi_modal_data` instead, so it
+needs a lighter-weight builder that only requires a HF tokenizer. This
+module provides that builder; the task -> template mapping below is the
+canonical mapping for both flows.
+"""
+
+from __future__ import annotations
+
+from .system_prompt import get_system_prompt
+
+# task -> (sys_type, bot_task, trigger_tag)
+_TASK_PRESETS: dict[str, tuple[str, str | None, str | None]] = {
+    "t2t": ("en_unified", None, None),
+    "i2t": ("en_unified", None, None),
+    "it2i_think": ("en_unified", "think", "<think>"),
+    "it2i_recaption": ("en_unified", "recaption", "<recaption>"),
+    "t2i_think": ("en_unified", "think", "<think>"),
+    "t2i_recaption": ("en_unified", "recaption", "<recaption>"),
+    "t2i_vanilla": ("en_vanilla", "image", None),
+}
+
+
+def available_tasks() -> list[str]:
+    """Sorted list of task keys accepted by `build_prompt` / `build_prompt_tokens`."""
+    return sorted(_TASK_PRESETS)
+
+
+def build_prompt(
+    user_prompt: str,
+    task: str = "it2i_think",
+    sys_type: str | None = None,
+    custom_system_prompt: str | None = None,
+) -> str:
+    """Build a HunyuanImage-3.0 prompt as a string (legacy/compat path).
+
+    NOTE: when this string is passed to the engine, the engine's tokenizer
+    will run a single BPE pass over the whole string, which can merge
+    tokens across segment boundaries (e.g. `。\\n\\n` -> id 3490). For
+    inputs that need to match HF baseline byte-for-byte, use
+    `build_prompt_tokens` instead and feed the result via prompt_token_ids.
+    """
+    if task not in _TASK_PRESETS:
+        raise ValueError(f"Unknown task {task!r}. Choose from: {available_tasks()}")
+
+    preset_sys_type, preset_bot_task, trigger_tag = _TASK_PRESETS[task]
+    effective_sys_type = sys_type or preset_sys_type
+
+    system_prompt = get_system_prompt(effective_sys_type, preset_bot_task, custom_system_prompt)
+    sys_text = system_prompt.strip() if system_prompt else ""
+
+    has_image_input = task.startswith("i2t") or task.startswith("it2i")
+
+    # t2i_vanilla: pretrain mode for direct text->image generation. The
+    # vanilla system prompt drives the model with no chat structure.
+    if task == "t2i_vanilla":
+        parts = ["<|startoftext|>"]
+        if sys_text:
+            parts.append(sys_text)
+        parts.append(user_prompt)
+        return "".join(parts)
+
+    # All other tasks (t2t / i2t / t2i_think / t2i_recaption /
+    # it2i_think / it2i_recaption) use HunyuanImage3 Instruct chat template:
+    #   <|startoftext|>{system?}\n\nUser: {<img>?}{user_prompt}\n\nAssistant: {trigger?}
+    # generation_config.json declares sequence_template="instruct", so the
+    # AR prefill MUST use this template -- verified to match HF's
+    # apply_chat_template output token-for-token (modulo BPE boundary merges).
+    # The trigger_tag (e.g. <think>) MUST come AFTER the `Assistant: ` prefix:
+    # if it goes BEFORE user_prompt (the old pretrain layout) the model puts
+    # the user's instructions inside the "thinking section" and collapses
+    # into repetition garbage under greedy decoding.
+    parts = ["<|startoftext|>"]
+    if sys_text:
+        parts.append(f"{sys_text}\n\n")
+    parts.append("User: ")
+    if has_image_input:
+        parts.append("<img>")
+    parts.append(user_prompt)
+    parts.append("\n\nAssistant: ")
+    if trigger_tag:
+        parts.append(trigger_tag)
+
+    return "".join(parts)
+
+
+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]:
+    """Segment-by-segment tokenization that matches HF apply_chat_template.
+
+    Calling tokenizer.encode(build_prompt(...)) on the full string lets BPE
+    merge tokens across segment boundaries (e.g. user_prompt ends with `。`
+    and the next segment is `\\n\\n` -> they merge into a single token id
+    3490 instead of HF's [1811, 271]). HF's apply_chat_template tokenizes
+    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).
+    """
+    if task not in _TASK_PRESETS:
+        raise ValueError(f"Unknown task {task!r}. Choose from: {available_tasks()}")
+
+    preset_sys_type, preset_bot_task, trigger_tag = _TASK_PRESETS[task]
+    effective_sys_type = sys_type or preset_sys_type
+
+    bos_id = tokenizer.convert_tokens_to_ids("<|startoftext|>")
+    img_id = tokenizer.convert_tokens_to_ids("<img>")
+    trig_id = tokenizer.convert_tokens_to_ids(trigger_tag) if trigger_tag else None
+
+    has_image_input = task.startswith("i2t") or task.startswith("it2i")
+
+    # t2i_vanilla uses pretrain template with no chat structure; the vanilla
+    # system prompt drives the model directly. No segment boundaries to
+    # 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)
+
+    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
+    # natural trailing newline; stripping it would shift one token id.
+    sys_text = system_prompt or ""
+
+    ids: list[int] = [bos_id]
+    if sys_text:
+        ids += tokenizer.encode(sys_text, add_special_tokens=False)
+        ids += tokenizer.encode("\n\n", add_special_tokens=False)
+    ids += tokenizer.encode("User: ", add_special_tokens=False)
+    if has_image_input:
+        ids += [img_id]
+    ids += tokenizer.encode(user_prompt, add_special_tokens=False)
+    ids += tokenizer.encode("\n\nAssistant: ", add_special_tokens=False)
+    if trig_id is not None:
+        ids += [trig_id]
+    return ids
+
+
+__all__ = ["build_prompt", "build_prompt_tokens", "available_tasks"]