[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>

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py

@@ -471,8 +471,21 @@ class Resolution:
         return f"{self.h}x{self.w}"
 
 
+# Baked-in extras matching the official model's
+# `HunyuanImage3ImageProcessor.vae_reso_group` (image_processor.py:147-152).
+# These four aspect buckets sit at ratio_token indices 33-36 in the trained
+# model and the AR was trained to address them, so any deviation breaks the
+# ratio-token vocab → output-shape lookup.
+HUNYUAN_IMAGE3_EXTRA_RESOLUTIONS: tuple[str, ...] = (
+    "1024x768",
+    "1280x720",
+    "768x1024",
+    "720x1280",
+)
+
+
 class ResolutionGroup:
-    def __init__(self, base_size=None, step=None, align=1):
+    def __init__(self, base_size=None, step=None, align=1, extra_resolutions=None):
         self.align = align
         self.base_size = base_size
         assert base_size % align == 0, f"base_size {base_size} is not divisible by align {align}"
@@ -486,6 +499,11 @@ class ResolutionGroup:
         self.step = step
         self.data = self._calc_by_step()
 
+        if extra_resolutions is not None:
+            for er in extra_resolutions:
+                if not any(r.ratio == er.ratio for r in self.data):
+                    self.data.append(er)
+
         self.ratio = np.array([x.ratio for x in self.data])
         self.attr = ["" for _ in range(len(self.data))]
         self.prefix_space = 0
@@ -1351,7 +1369,10 @@ class HunyuanImage3ImageProcessor:
     def __init__(self, config):
         self.config = config
 
-        self.reso_group = ResolutionGroup(base_size=config.image_base_size)
+        self.reso_group = ResolutionGroup(
+            base_size=config.image_base_size,
+            extra_resolutions=[Resolution(s) for s in HUNYUAN_IMAGE3_EXTRA_RESOLUTIONS],
+        )
         self.vae_processor = transforms.Compose(
             [
                 transforms.ToTensor(),

vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py

@@ -283,11 +283,13 @@ def get_hunyuan_image_3_pre_process_func(od_config: OmniDiffusionConfig):
                 cond_image_infos = [_build_cond_joint_image(image) for image in image_list]
                 prompt["additional_information"]["batch_cond_image_info"] = cond_image_infos
 
+                bridge_h = prompt.get("height") if isinstance(prompt, dict) else None
+                bridge_w = prompt.get("width") if isinstance(prompt, dict) else None
                 first_image_w, first_image_h = _to_pil_image(image_list[0]).size
                 if request.sampling_params.width is None:
-                    request.sampling_params.width = int(first_image_w)
+                    request.sampling_params.width = int(bridge_w or first_image_w)
                 if request.sampling_params.height is None:
-                    request.sampling_params.height = int(first_image_h)
+                    request.sampling_params.height = int(bridge_h or first_image_h)
 
             request.prompts[i] = prompt
 
@@ -539,7 +541,12 @@ class HunyuanImage3Pipeline(
         timestep_scatter_index: BatchRaggedTensor,
     ):
         batch_size, seq_len, n_embd = x.shape
-        # batch_size x n x n_embd
+        # `_encode_cond_image` returns `t` as list[Tensor] for the
+        # multi-image branch (outer length = batch_size, currently fixed
+        # at 1 by the stage runtime `max_batch_size`); flatten to a Tensor
+        # before reshape.
+        if isinstance(t, list):
+            t = torch.cat([ti.reshape(-1) for ti in t], dim=0)
         timestep_scatter_src = self.timestep_emb(t.reshape(-1)).reshape(batch_size, -1, n_embd)
         x.scatter_(
             dim=1,
@@ -627,7 +634,11 @@ class HunyuanImage3Pipeline(
             if isinstance(vae_encode_result, torch.Tensor):
                 latents = vae_encode_result
             else:
-                latents = vae_encode_result.latent_dist.sample()
+                # Match HunyuanImage-3's cond encode path: sample the
+                # posterior, but use a fixed generator so repeated online
+                # requests are deterministic.
+                _cond_vae_gen = torch.Generator(device=image.device).manual_seed(0)
+                latents = vae_encode_result.latent_dist.sample(_cond_vae_gen)
             if hasattr(config, "shift_factor") and config.shift_factor:
                 latents.sub_(config.shift_factor)
             if hasattr(config, "scaling_factor") and config.scaling_factor:
@@ -1353,25 +1364,21 @@ class HunyuanImage3Pipeline(
         use_system_prompt = extra_args.get("use_system_prompt")
         system_prompt = extra_args.get("system_prompt")
         # Fall back to per-prompt use_system_prompt forwarded by ar2diffusion
-        if use_system_prompt is None and req.prompts:
+        if req.prompts:
             first_prompt = req.prompts[0]
             if isinstance(first_prompt, dict):
-                use_system_prompt = first_prompt.get("use_system_prompt")
+                if use_system_prompt is None:
+                    use_system_prompt = first_prompt.get("use_system_prompt")
+                if system_prompt is None:
+                    system_prompt = first_prompt.get("system_prompt")
         if use_system_prompt is not None:
             system_prompt = get_system_prompt(use_system_prompt, "image", system_prompt)
             system_prompt = system_prompt.strip() if system_prompt is not None else ""
         prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
 
-        # Extract AR-generated CoT/recaption text from each prompt's extra dict.
-        # The AR-side stage input processor (``ar2diffusion``) already prepends
-        # the trigger tag (e.g. ``<think>``) when the AR used the KV-reuse
-        # pretrain format, so ``ar_generated_text`` is a self-contained string
-        # and ``get_cot_sections()`` can parse the think/recaption structure
-        # directly.
-        cot_text_list = []
-        for p in req.prompts:
-            extra = p.get("extra", {}) if isinstance(p, dict) else {}
-            cot_text_list.append(extra.get("ar_generated_text") or None)
+        cot_text_list = [
+            (p.get("extra", {}).get("ar_generated_text") if isinstance(p, dict) else None) or None for p in req.prompts
+        ]
         cot_text = (
             [self._normalize_cot_text(t) for t in cot_text_list] if any(t is not None for t in cot_text_list) else None
         )

vllm_omni/diffusion/models/hunyuan_image3/prompt_utils.py

@@ -11,8 +11,23 @@ The DiT pipeline (`pipeline_hunyuan_image3.py`) builds prompts through
 `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.
+module provides that builder; the (task, bot_task) -> template mapping
+below is the canonical mapping for both flows.
+
+Two orthogonal axes:
+
+  * `task` selects the I/O modality combination, which only controls
+    whether `<img>` placeholders are emitted between `User: ` and the
+    user prompt: ``i2t`` / ``it2i`` produce them, ``t2t`` / ``t2i`` do
+    not.
+
+  * `bot_task` selects the prompting mode and drives both the system
+    prompt and the trigger tag appended after ``Assistant: ``. ``None``
+    (default) gives a plain Assistant turn under the unified prompt;
+    ``think`` / ``recaption`` switch the trigger tag to ``<think>`` /
+    ``<recaption>``; ``think_recaption`` swaps the system prompt for
+    the dedicated combined-mode template; ``vanilla`` drops the chat
+    structure entirely (pretrain template, ``t2i`` only).
 """
 
 from __future__ import annotations
@@ -45,143 +60,234 @@ HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS: dict[str, int] = {
     "<img_ratio_36>": 130106,
 }
 
-# task -> (sys_type, bot_task, trigger_tag)
+# bot_task -> (sys_type, trigger_tag).
+# ``vanilla`` is special-cased downstream: it bypasses the chat template
+# (no ``User:`` / ``Assistant:`` framing) and is only valid with
+# ``task='t2i'``.
+_BOT_TASK_PRESETS: dict[str | None, tuple[str, str | None]] = {
+    None: ("en_unified", None),
+    "think": ("en_unified", "<think>"),
+    "recaption": ("en_unified", "<recaption>"),
+    "think_recaption": ("en_think_recaption", "<think>"),
+    "vanilla": ("en_vanilla", None),
+}
+
+_TASKS: frozenset[str] = frozenset({"t2t", "i2t", "it2i", "t2i"})
+
+
+class _DefaultBotTask:
+    pass
+
+
+_DEFAULT_BOT_TASK = _DefaultBotTask()
+
+# Legacy composite task alias -> (task, bot_task). Keep this during rebase so
+# older callers and intermediate commits still resolve cleanly.
 _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>"),
     "it2i_think_recaption": ("en_unified", "think_recaption", "<think>"),
-    "t2i": ("en_unified", "image", None),
-    "t2i_vanilla": ("en_vanilla", "image", None),
+    "t2i": ("en_unified", None, None),
+    "t2i_vanilla": ("en_vanilla", "vanilla", None),
     "t2i_think": ("en_unified", "think", "<think>"),
     "t2i_recaption": ("en_unified", "recaption", "<recaption>"),
 }
 
+_LEGACY_COMPOSITE_TASKS: frozenset[str] = frozenset(_TASK_PRESETS) - {"t2t", "i2t", "t2i"}
+
+
+def _normalize_task_and_bot_task(
+    task: str,
+    bot_task: str | None | _DefaultBotTask,
+) -> tuple[str, str | None]:
+    bot_task_was_omitted = bot_task is _DEFAULT_BOT_TASK
+    if task in _TASK_PRESETS:
+        _, legacy_bot_task, _ = _TASK_PRESETS[task]
+        base_task = task.split("_", 1)[0]
+        if base_task == "t2i" and task == "t2i":
+            base_task = "t2i"
+        if task in ("t2t", "i2t", "t2i"):
+            base_task = task
+        if bot_task_was_omitted:
+            bot_task = legacy_bot_task
+        elif task in _LEGACY_COMPOSITE_TASKS and bot_task is None:
+            # Composite task names already encode the legacy bot_task. Keep
+            # calls like build_prompt_tokens(task="it2i_think", bot_task=None)
+            # on their historical meaning; explicit None is the plain-mode
+            # escape hatch only for the new two-axis base tasks.
+            bot_task = legacy_bot_task
+        task = base_task
+    elif bot_task_was_omitted:
+        bot_task = "think"
+    return task, bot_task
+
 
 def available_tasks() -> list[str]:
-    """Sorted list of task keys accepted by `build_prompt` / `build_prompt_tokens`."""
-    return sorted(_TASK_PRESETS)
+    """Sorted list of `task` values accepted by the prompt builders."""
+    return sorted(_TASKS)
+
+
+def available_bot_tasks() -> list[str | None]:
+    """Sorted list of `bot_task` values (with ``None`` first)."""
+    rest = sorted(k for k in _BOT_TASK_PRESETS if k is not None)
+    return [None, *rest]
+
+
+def resolve_sys_type(bot_task: str | None) -> str:
+    """Default system-prompt type for a given ``bot_task``."""
+    if bot_task not in _BOT_TASK_PRESETS:
+        raise ValueError(f"Unknown bot_task {bot_task!r}. Choose from: {available_bot_tasks()}")
+    return _BOT_TASK_PRESETS[bot_task][0]
 
 
 def resolve_stop_token_ids(
-    task: str = "it2i_think",
-    bot_task: str = "think",
+    task: str = "it2i",
+    bot_task: str | None | _DefaultBotTask = _DEFAULT_BOT_TASK,
     tokenizer: Any | None = None,
-):
+) -> list[int]:
+    """AR stop-token ids for a given (task, bot_task) generation request.
+
+    Image-output tasks (``it2i`` / ``t2i``) stop on any ``<img_ratio_*>``
+    token. Upstream ``modeling_hunyuan_image_3.py::generate_image``
+    (line 3289-3303) sets ``final_stop_tokens`` to the full ratio token
+    range when ``need_ratio`` is true, then strips the trailing ratio
+    token before passing the cot to the image stage. AR's natural
+    trajectory under ``_stage_transitions`` is
+    ``</recaption><answer><boi><img_size_base><img_ratio_X>``; stopping
+    AT the ratio token means KV ends exactly at the prefix DiT reuses,
+    and ``ar2diffusion`` can read the ratio off the last sampled token
+    without AR wasting decode steps on ``<|endoftext|>``.
+
+    Text-output tasks (``i2t`` / ``t2t``) stop on ``<answer>`` -- the AR
+    is the final stage, and the comprehension response sits inside the
+    ``<answer>`` body so the answer-open is the natural cot/recaption
+    terminator.
+    """
+    task, bot_task = _normalize_task_and_bot_task(task, bot_task)
+    if task not in _TASKS:
+        raise ValueError(f"Unknown task {task!r}. Choose from: {available_tasks()}")
+    if bot_task not in _BOT_TASK_PRESETS:
+        raise ValueError(f"Unknown bot_task {bot_task!r}. Choose from: {available_bot_tasks()}")
+    if task in ("it2i", "t2i"):
+        # Main ratio range: <img_ratio_0> .. <img_ratio_32>.
+        start = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_0>"]
+        end = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_32>"]
+        stops = list(range(start, end + 1))
+        # Other slices (upstream tokenizer ``ratio_token_other_slices``):
+        # <img_ratio_33> .. <img_ratio_36>.
+        other_start = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_33>"]
+        other_end = HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<img_ratio_36>"]
+        stops.extend(range(other_start, other_end + 1))
+        return stops
     return [HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS["<answer>"]]
 
 
+# Upstream "Multi-Image Fusion" caps reference images at 3 per request.
+MAX_IMAGES_PER_REQUEST = 3
+
+
+def _validate_num_images(num_images: int) -> None:
+    if not (1 <= num_images <= MAX_IMAGES_PER_REQUEST):
+        raise ValueError(f"num_images must be in [1, {MAX_IMAGES_PER_REQUEST}], got {num_images}")
+
+
+def _resolve_preset(task: str, bot_task: str | None) -> tuple[str, str | None]:
+    """Validate (task, bot_task) and return ``(sys_type, trigger_tag)``."""
+    task, bot_task = _normalize_task_and_bot_task(task, bot_task)
+    if task not in _TASKS:
+        raise ValueError(f"Unknown task {task!r}. Choose from: {available_tasks()}")
+    if bot_task not in _BOT_TASK_PRESETS:
+        raise ValueError(f"Unknown bot_task {bot_task!r}. Choose from: {available_bot_tasks()}")
+    if bot_task == "vanilla" and task != "t2i":
+        raise ValueError(f"bot_task='vanilla' is only valid with task='t2i' (pretrain template); got task={task!r}")
+    return _BOT_TASK_PRESETS[bot_task]
+
+
 def build_prompt(
     user_prompt: str,
-    task: str = "it2i_think",
+    task: str = "it2i",
+    bot_task: str | None | _DefaultBotTask = _DEFAULT_BOT_TASK,
     sys_type: str | None = None,
     custom_system_prompt: str | None = None,
+    num_images: int = 1,
 ) -> 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]
+    """Build a HunyuanImage-3.0 prompt as a string (legacy/compat path)."""
+    task, bot_task = _normalize_task_and_bot_task(task, bot_task)
+    preset_sys_type, trigger_tag = _resolve_preset(task, bot_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 ""
+    system_prompt = get_system_prompt(effective_sys_type, bot_task, custom_system_prompt)
+    sys_text = system_prompt or ""
 
-    has_image_input = task.startswith("i2t") or task.startswith("it2i")
+    has_image_input = task in ("i2t", "it2i")
+    if has_image_input:
+        _validate_num_images(num_images)
 
-    # t2i_vanilla: pretrain mode for direct text->image generation. The
-    # vanilla system prompt drives the model with no chat structure.
-    if task == "t2i_vanilla":
+    if bot_task == "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.extend(["<img>"] * num_images)
     parts.append(user_prompt)
     parts.append("\n\nAssistant: ")
     if trigger_tag:
         parts.append(trigger_tag)
-
     return "".join(parts)
 
 
 @dataclass
 class PromptTokensResult:
-    token_ids: list[int]  # The tokenized prompt
-    system_prompt_type: str  # The effective system prompt type used
+    token_ids: list[int]
+    system_prompt_type: str
 
 
 def build_prompt_tokens(
     user_prompt: str,
     tokenizer,
-    task: str = "it2i_think",
+    task: str = "it2i",
+    bot_task: str | None | _DefaultBotTask = _DEFAULT_BOT_TASK,
     sys_type: str | None = None,
     custom_system_prompt: str | None = None,
+    num_images: int = 1,
 ) -> PromptTokensResult:
-    """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).
-
-    Returns:
-        PromptTokensResult
-    """
-    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]
+    """Segment-by-segment tokenization that matches HF apply_chat_template."""
+    task, bot_task = _normalize_task_and_bot_task(task, bot_task)
+    preset_sys_type, trigger_tag = _resolve_preset(task, bot_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")
+    has_image_input = task in ("i2t", "it2i")
+    if has_image_input:
+        _validate_num_images(num_images)
 
-    # 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)
+    if bot_task == "vanilla":
+        s = build_prompt(
+            user_prompt,
+            task=task,
+            bot_task=bot_task,
+            sys_type=sys_type,
+            custom_system_prompt=custom_system_prompt,
+        )
         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
-    # natural trailing newline; stripping it would shift one token id.
+    system_prompt = get_system_prompt(effective_sys_type, bot_task, custom_system_prompt)
     sys_text = system_prompt or ""
 
     ids: list[int] = [bos_id]
@@ -190,7 +296,7 @@ def build_prompt_tokens(
         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 += [img_id] * num_images
     ids += tokenizer.encode(user_prompt, add_special_tokens=False)
     ids += tokenizer.encode("\n\nAssistant: ", add_special_tokens=False)
     if trig_id is not None:
@@ -202,4 +308,14 @@ def build_prompt_tokens(
     )
 
 
-__all__ = ["build_prompt", "build_prompt_tokens", "resolve_stop_token_ids", _TASK_PRESETS]
+__all__ = [
+    "HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS",
+    "MAX_IMAGES_PER_REQUEST",
+    "_TASK_PRESETS",
+    "available_bot_tasks",
+    "available_tasks",
+    "build_prompt",
+    "build_prompt_tokens",
+    "resolve_stop_token_ids",
+    "resolve_sys_type",
+]