[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>
2026-05-14 16:58:40 +08:00
1 parent c4f859bf56
commit 3f63aaf982
@@ -1349,8 +1349,16 @@ def test_image_edit_parameter_default(async_omni_test_client):
    engine = async_omni_test_client.app.state.engine_client
    captured_sampling_params = engine.captured_sampling_params_list[-1]

-    assert captured_sampling_params.width == 24
-    assert captured_sampling_params.height == 16
+    # size="auto" on multi-stage pipelines deliberately leaves the diffusion
+    # stages sampling_params width/height unset so AR-driven pipelines (e.g.
+    # HunyuanImage-3.0) can let ar2diffusion override the final bucket from
+    # the AR-predicted ratio token; see
+    # test_image_edits_size_auto_preserves_bridge_size for the contract.
+    # Single-stage diffusion (test_image_edit_parameter_default_single_stage)
+    # still pins width/height to the input image size via api_servers
+    # gen_params, which is unchanged.
+    assert captured_sampling_params.width is None
+    assert captured_sampling_params.height is None
    assert captured_sampling_params.num_outputs_per_prompt == 1
    assert captured_sampling_params.num_inference_steps == 4
    assert captured_sampling_params.guidance_scale == 7.5
@@ -1649,3 +1657,59 @@ def test_extract_images_from_result():
    assert len(images) == 1
    assert isinstance(images[0], Image.Image)
    assert images[0].size == (32, 32)
+
+
+def test_image_edits_size_auto_preserves_bridge_size(async_omni_stage_configs_only_client):
+    """size=auto must NOT pin the diffusion stage sampling_params.height/width.
+
+    Regression: prior to the fix, edit_images resolved size=auto to the
+    first input image dimensions and forwarded them through gen_params +
+    extra_body to the diffusion stages sampling_params. AR-driven
+    pipelines (e.g. HunyuanImage-3.0) rely on ar2diffusions
+    bridge to override the final bucket via the AR-predicted ratio token,
+    and the DiT pre_process_func only fills sampling_params from the
+    bridge value when sampling_params.width is None (see
+    pipeline_hunyuan_image3.py:290). Non-None width from the input image
+    silently suppressed the AR decision, producing the wrong bucket
+    (e.g. 1024x1024 square instead of the AR-decided 1280x720 landscape
+    for multi-image fusion).
+
+    Cross-pins the multi-image fix at the API level: 2 reference images
+    with bot_task=think must produce 2 <img> placeholders in the captured
+    AR prompt (build_prompt called with num_images=2).
+    """
+    img_a = make_test_image_bytes((32, 32))
+    img_b = make_test_image_bytes((128, 64))
+    response = async_omni_stage_configs_only_client.post(
+        "/v1/images/edits",
+        files=[("image", img_a), ("image", img_b)],
+        data={
+            "prompt": "fuse",
+            "size": "auto",
+            "bot_task": "think",
+        },
+    )
+    assert response.status_code == 200, response.text
+
+    engine = async_omni_stage_configs_only_client.app.state.engine_client
+    captured = engine.captured_sampling_params_list
+    assert captured is not None
+    assert len(captured) == 2
+
+    diffusion_params = captured[1]
+    assert diffusion_params.height is None, (
+        f"size=auto leaked into diffusion sampling_params.height={diffusion_params.height}; "
+        "must stay None so AR-driven pipelines can apply the bridges decision."
+    )
+    assert diffusion_params.width is None, (
+        f"size=auto leaked into diffusion sampling_params.width={diffusion_params.width}; "
+        "must stay None so AR-driven pipelines can apply the bridges decision."
+    )
+
+    KEY = "prompt"
+    IMG = "<img>"
+    captured_prompt = engine.captured_prompt
+    if isinstance(captured_prompt, dict) and isinstance(captured_prompt.get("prompt"), str):
+        assert captured_prompt["prompt"].count("<img>") == 2, (
+            f"N=2 reference images must emit 2 <img> placeholders in AR prompt; got {captured_prompt[KEY].count(IMG)} -- prompt: {captured_prompt[KEY]!r}"
+        )
@@ -91,3 +91,275 @@ def test_build_multistage_generation_inputs_applies_stage_specific_overrides(ser
    assert engine.default_sampling_params_list[1].lora_request is None
    assert engine.default_sampling_params_list[2].resolution == 640
    assert engine.default_sampling_params_list[2].lora_request is None
+
+
+def test_build_multistage_generation_inputs_multi_image_emits_n_img_placeholders(serving_chat):
+    """N reference images with bot_task set must emit N <img> placeholders.
+
+    Regression: prior to the multi-image online fix, build_prompt was
+    called without num_images, defaulting to 1. A 2-image edit request
+    would only get a single <img> placeholder in the AR prompt; vLLMs
+    _process_multimodal then raised
+    AssertionError(Failed to apply prompt replacement for mm_items[image][1])
+    when trying to replace the second image (no placeholder left for it).
+
+    Pins the contract that build_prompt() is invoked with the actual image
+    count so multi-image IT2I is wired correctly through the online
+    /v1/images/edits path.
+    """
+    from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
+
+    engine = SimpleNamespace(
+        stage_configs=[
+            SimpleNamespace(stage_type="llm", is_comprehension=True),
+            SimpleNamespace(stage_type="diffusion", is_comprehension=False),
+        ],
+        default_sampling_params_list=[
+            SamplingParams(temperature=0.0),
+            OmniDiffusionSamplingParams(),
+        ],
+    )
+    IMG = "<img>"
+    images = [Image.new("RGB", (32, 32), color="red") for _ in range(3)]
+
+    for n in (1, 2, 3):
+        engine_prompt, _ = OmniOpenAIServingChat._build_multistage_generation_inputs(
+            serving_chat,
+            engine=engine,
+            prompt="edit me",
+            extra_body={"bot_task": "think"},
+            reference_images=images[:n],
+            gen_params=OmniDiffusionSamplingParams(),
+        )
+        prompt_str = engine_prompt["prompt"]
+        assert prompt_str.count("<img>") == n, (
+            f"N={n}: expected {n} <img> placeholders, got {prompt_str.count(IMG)} -- prompt: {prompt_str!r}"
+        )
+
+
+def test_build_multistage_generation_inputs_tokenizer_path_emits_prompt_token_ids(serving_chat):
+    """When a tokenizer is provided, the helper must emit HF byte-for-byte
+    prompt_token_ids and forward use_system_prompt to the engine prompt.
+
+    Regression: prior to the HF-byte-equivalent fix, online IT2I always
+    passed the prompt as a single string. The engine then BPE-merged across
+    chat-template segment boundaries (e.g. user_prompt-ending punctuation
+    plus the trailing \n\n before \"Assistant: \") producing a token
+    sequence that differs from HF apply_chat_template / offline
+    end2end.py. AR generated different cot_text (706 tokens / 1190 chars
+    vs offline 661 / 1118 for the same inputs) and DiT produced a visually
+    different image (yin-yang on brushed-metal vs three-blue swirl on
+    canvas) under the same seed.
+
+    Pins:
+      1. engine_prompt[\"prompt_token_ids\"] is set when tokenizer is passed.
+      2. engine_prompt[\"prompt\"] stays as the raw user prompt -- the DiT
+         side rebuilds its own system prefix via use_system_prompt.
+      3. engine_prompt[\"use_system_prompt\"] == \"en_unified\" so
+         ar2diffusion forwards the matching system prompt to DiT.
+      4. N reference images emit N <img> token ids in the AR sequence.
+    """
+    from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
+
+    # Minimal FakeTokenizer mirroring tests/diffusion/.../test_hunyuan_image3_it2i_multi_image.py
+    class FakeTokenizer:
+        SPECIAL = {
+            "<|startoftext|>": 1,
+            "<img>": 2,
+            "<think>": 3,
+            "<recaption>": 4,
+        }
+
+        def convert_tokens_to_ids(self, tok: str) -> int:
+            return self.SPECIAL.get(tok, 0)
+
+        def encode(self, text: str, add_special_tokens: bool = False) -> list[int]:
+            return list(range(100, 100 + len(text)))
+
+    engine = SimpleNamespace(
+        stage_configs=[
+            SimpleNamespace(stage_type="llm", is_comprehension=True),
+            SimpleNamespace(stage_type="diffusion", is_comprehension=False),
+        ],
+        default_sampling_params_list=[
+            SamplingParams(temperature=0.0),
+            OmniDiffusionSamplingParams(),
+        ],
+    )
+    PROMPT_KEY = "prompt"
+    USP_KEY = "use_system_prompt"
+    images = [Image.new("RGB", (32, 32), color="red") for _ in range(3)]
+
+    for n in (1, 2, 3):
+        tok = FakeTokenizer()
+        engine_prompt, _ = OmniOpenAIServingChat._build_multistage_generation_inputs(
+            serving_chat,
+            engine=engine,
+            prompt="edit me",
+            extra_body={"bot_task": "think"},
+            reference_images=images[:n],
+            gen_params=OmniDiffusionSamplingParams(),
+            tokenizer=tok,
+        )
+        # (1) prompt_token_ids must be set and non-empty
+        assert "prompt_token_ids" in engine_prompt, f"N={n}: prompt_token_ids missing"
+        token_ids = engine_prompt["prompt_token_ids"]
+        assert isinstance(token_ids, list) and len(token_ids) > 0, f"N={n}: prompt_token_ids empty"
+        # (2) raw prompt preserved (DiT bridge needs raw user text)
+        assert engine_prompt["prompt"] == "edit me", (
+            f"N={n}: prompt must stay raw user text, got {engine_prompt[PROMPT_KEY]!r}"
+        )
+        # (3) use_system_prompt forwarded for ar2diffusion bridge
+        assert engine_prompt.get("use_system_prompt") == "en_unified", (
+            f"N={n}: use_system_prompt must be en_unified, got {engine_prompt.get(USP_KEY)!r}"
+        )
+        # (4) N <img> token ids (id=2 in FakeTokenizer)
+        img_count = token_ids.count(2)
+        assert img_count == n, f"N={n}: expected {n} <img> token ids in prompt_token_ids, got {img_count}"
+
+
+def test_build_multistage_generation_inputs_bot_task_semantic_changes_trigger_and_sys(serving_chat):
+    """Passing bot_task=think_recaption (vs default "think") must flip the
+    resolved sys_type to en_think_recaption (and trigger tag is still
+    <think>). Pins that the API actually plumbs the bot_task semantic
+    through to build_prompt rather than ignoring it.
+    """
+    from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
+
+    engine = SimpleNamespace(
+        stage_configs=[
+            SimpleNamespace(stage_type="llm", is_comprehension=True),
+            SimpleNamespace(stage_type="diffusion", is_comprehension=False),
+        ],
+        default_sampling_params_list=[
+            SamplingParams(temperature=0.0),
+            OmniDiffusionSamplingParams(),
+        ],
+    )
+    images = [Image.new("RGB", (32, 32), color="red")]
+
+    # Default bot_task (think) -> en_unified system prompt baked into the
+    # legacy string path. Use legacy build_prompt (tokenizer=None) so the
+    # rendered prompt is a string we can grep.
+    think_prompt, _ = OmniOpenAIServingChat._build_multistage_generation_inputs(
+        serving_chat,
+        engine=engine,
+        prompt="edit me",
+        extra_body={"task": "it2i", "bot_task": "think"},
+        reference_images=images,
+        gen_params=OmniDiffusionSamplingParams(),
+    )
+    # think_recaption -> en_think_recaption system prompt (different content).
+    recap_prompt, _ = OmniOpenAIServingChat._build_multistage_generation_inputs(
+        serving_chat,
+        engine=engine,
+        prompt="edit me",
+        extra_body={"task": "it2i", "bot_task": "think_recaption"},
+        reference_images=images,
+        gen_params=OmniDiffusionSamplingParams(),
+    )
+    assert think_prompt["prompt"] != recap_prompt["prompt"], (
+        "bot_task semantic must change the rendered system prompt: "
+        f"think/think_recaption produced identical strings (len={len(think_prompt['prompt'])})"
+    )
+
+
+def test_build_multistage_generation_inputs_sys_type_override(serving_chat):
+    """Caller-supplied sys_type must override the bot_task-derived default.
+    Mirrors offline `--bot-task think_recaption --sys-type en_unified`
+    where the user wants think_recaptions trigger but the unified system
+    prompt body.
+    """
+    from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
+
+    engine = SimpleNamespace(
+        stage_configs=[
+            SimpleNamespace(stage_type="llm", is_comprehension=True),
+            SimpleNamespace(stage_type="diffusion", is_comprehension=False),
+        ],
+        default_sampling_params_list=[
+            SamplingParams(temperature=0.0),
+            OmniDiffusionSamplingParams(),
+        ],
+    )
+    images = [Image.new("RGB", (32, 32), color="red")]
+
+    # think_recaption defaults sys_type -> en_think_recaption.
+    default_sys, _ = OmniOpenAIServingChat._build_multistage_generation_inputs(
+        serving_chat,
+        engine=engine,
+        prompt="edit me",
+        extra_body={"task": "it2i", "bot_task": "think_recaption"},
+        reference_images=images,
+        gen_params=OmniDiffusionSamplingParams(),
+    )
+    # sys_type=en_unified overrides -> same system body as bot_task=think.
+    overridden, _ = OmniOpenAIServingChat._build_multistage_generation_inputs(
+        serving_chat,
+        engine=engine,
+        prompt="edit me",
+        extra_body={"task": "it2i", "bot_task": "think_recaption", "sys_type": "en_unified"},
+        reference_images=images,
+        gen_params=OmniDiffusionSamplingParams(),
+    )
+    plain_think, _ = OmniOpenAIServingChat._build_multistage_generation_inputs(
+        serving_chat,
+        engine=engine,
+        prompt="edit me",
+        extra_body={"task": "it2i", "bot_task": "think"},
+        reference_images=images,
+        gen_params=OmniDiffusionSamplingParams(),
+    )
+
+    # Override must (a) differ from the no-override default, and (b) equal
+    # the prompt that bot_task=think produces (both end up with
+    # en_unified system body + <think> trigger).
+    assert overridden["prompt"] != default_sys["prompt"], (
+        "sys_type override must change the rendered prompt body vs the bot_task default"
+    )
+    assert overridden["prompt"] == plain_think["prompt"], (
+        "sys_type=en_unified + bot_task=think_recaption must produce the same prompt as "
+        "bot_task=think (both = en_unified system body + <think> trigger)"
+    )
+
+
+def test_build_multistage_generation_inputs_custom_system_prompt(serving_chat):
+    """`extra_body["system_prompt"]` must reach build_prompt as
+    `custom_system_prompt`, enabling sys_type="custom" callers to inject
+    a verbatim system body. Without this plumbing the sys_type="custom"
+    branch in get_system_prompt() returns None and silently drops the
+    user-supplied content.
+    """
+    from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
+
+    engine = SimpleNamespace(
+        stage_configs=[
+            SimpleNamespace(stage_type="llm", is_comprehension=True),
+            SimpleNamespace(stage_type="diffusion", is_comprehension=False),
+        ],
+        default_sampling_params_list=[
+            SamplingParams(temperature=0.0),
+            OmniDiffusionSamplingParams(),
+        ],
+    )
+    images = [Image.new("RGB", (32, 32), color="red")]
+
+    marker = "ZZZ_CUSTOM_SYSTEM_PROMPT_MARKER_ZZZ"
+
+    out, _ = OmniOpenAIServingChat._build_multistage_generation_inputs(
+        serving_chat,
+        engine=engine,
+        prompt="edit me",
+        extra_body={
+            "task": "it2i",
+            "bot_task": "think",
+            "sys_type": "custom",
+            "system_prompt": marker,
+        },
+        reference_images=images,
+        gen_params=OmniDiffusionSamplingParams(),
+    )
+    assert marker in out["prompt"], (
+        f"custom system_prompt content must reach the rendered prompt; "
+        f"marker {marker!r} not found in prompt of length {len(out['prompt'])}"
+    )