[Example] Migrate HunyuanImage-3.0 to model_extras + shared task examples (#5559)

Signed-off-by: suyanli220 <suyanli220@gmail.com>
Signed-off-by: suyan.li <suyan.li@bytedance.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: suyan.li <suyan.li@bytedance.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/model_extras/test_shared_script_ar_integration.py

@@ -0,0 +1,340 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Integration coverage for the shared task examples' ``_apply_ar_stage_inputs``.
+
+The unit tests in ``test_model_extras.py`` cover ``get_ar_input_builder`` and
+``build_ar_stage_inputs`` in isolation, but nothing previously exercised the
+*glue* in ``text_to_image.py`` / ``image_edit.py`` that wires them together:
+resolving the model's declared ``ar_input_builder`` via the registry, loading
+the AR tokenizer, and writing the result onto ``prompt_dict`` /
+``sampling_params_list``. A regression in that glue -- e.g. the registry key
+mismatch fixed in this PR (``HunyuanImage3ForCausalMM`` vs
+``HunyuanImage3Pipeline``), or the ``trust_remote_code`` opt-in being
+silently bypassed -- would not have been caught by the isolated unit tests
+above, since they call ``build_ar_stage_inputs`` directly rather than going
+through the shared script's own code path.
+
+These tests dynamically load the example scripts by file path (``examples/``
+is not an installed package) and call ``_apply_ar_stage_inputs`` the same way
+``main()`` does, with a fake tokenizer standing in for a real HF download.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import pytest
+
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_REPO_ROOT = Path(__file__).resolve().parents[2]
+
+
+def _load_module(name: str, relpath: str):
+    path = _REPO_ROOT / relpath
+    spec = importlib.util.spec_from_file_location(name, path)
+    assert spec is not None and spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+class _FakeHunyuanTokenizer:
+    """Mirrors test_model_extras.py's fixture: distinct ids, no network I/O."""
+
+    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 [500 + len(text)]
+
+
+@dataclass
+class _FakeARSamplingParams:
+    """Stands in for an LLM (non-diffusion) stage's sampling params."""
+
+    stop_token_ids: list[int] | None = None
+
+
+def _patch_tokenizer_loader(monkeypatch: pytest.MonkeyPatch, module: Any) -> list[dict[str, Any]]:
+    """Patch the ``AutoTokenizer`` the module imports locally, recording calls."""
+    calls: list[dict[str, Any]] = []
+
+    class _FakeAutoTokenizer:
+        @staticmethod
+        def from_pretrained(model: str, trust_remote_code: bool = False) -> _FakeHunyuanTokenizer:
+            calls.append({"model": model, "trust_remote_code": trust_remote_code})
+            return _FakeHunyuanTokenizer()
+
+    fake_transformers = type(sys)("transformers")
+    fake_transformers.AutoTokenizer = _FakeAutoTokenizer
+    monkeypatch.setitem(sys.modules, "transformers", fake_transformers)
+    return calls
+
+
+def _patch_failing_tokenizer_loader(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Patch AutoTokenizer.from_pretrained to always raise, simulating a
+    missing --trust-remote-code / network / cache failure."""
+
+    class _FailingAutoTokenizer:
+        @staticmethod
+        def from_pretrained(model: str, trust_remote_code: bool = False) -> Any:
+            raise OSError("simulated tokenizer load failure")
+
+    fake_transformers = type(sys)("transformers")
+    fake_transformers.AutoTokenizer = _FailingAutoTokenizer
+    monkeypatch.setitem(sys.modules, "transformers", fake_transformers)
+
+
+@pytest.mark.parametrize(
+    "module_name,relpath",
+    [
+        ("_shared_text_to_image", "examples/offline_inference/text_to_image/text_to_image.py"),
+        ("_shared_image_edit", "examples/offline_inference/image_to_image/image_edit.py"),
+    ],
+)
+def test_apply_ar_stage_inputs_routes_through_real_registry(
+    monkeypatch: pytest.MonkeyPatch, module_name: str, relpath: str
+) -> None:
+    """The shared script's glue must resolve HunyuanImage3's *real* registry
+    entry (keyed on the runtime ``model_class_name``) and write its output
+    onto the request -- not a stub, not a hardcoded id."""
+    from vllm_omni.model_extras import get_ar_input_builder
+
+    module = _load_module(module_name, relpath)
+    tokenizer_calls = _patch_tokenizer_loader(monkeypatch, module)
+
+    ar_input_builder = get_ar_input_builder("HunyuanImage3ForCausalMM")
+    assert ar_input_builder is not None, (
+        "get_ar_input_builder('HunyuanImage3ForCausalMM') returned None -- "
+        "if this regresses, the shared script silently skips AR-stage "
+        "input building entirely (this is exactly how the registry-key bug "
+        "this PR fixes went undetected)."
+    )
+
+    prompt_dict: dict[str, Any] = {}
+    ar_params = _FakeARSamplingParams()
+    diffusion_params = OmniDiffusionSamplingParams()
+    sampling_params_list: list[Any] = [ar_params, diffusion_params]
+
+    module._apply_ar_stage_inputs(
+        ar_input_builder,
+        model="tencent/HunyuanImage-3.0-Instruct",
+        prompt_text="a red panda",
+        extra_body={"bot_task": "think", "use_system_prompt": "en_recaption"},
+        num_images=0,
+        height=1024,
+        width=1024,
+        prompt_dict=prompt_dict,
+        sampling_params_list=sampling_params_list,
+        trust_remote_code=False,
+    )
+
+    # The tokenizer was loaded exactly once, for the model the user asked for.
+    assert len(tokenizer_calls) == 1
+    assert tokenizer_calls[0]["model"] == "tencent/HunyuanImage-3.0-Instruct"
+
+    # Byte-for-byte tokenizer path: prompt_token_ids populated, not the string form.
+    assert prompt_dict.get("prompt_token_ids"), "AR prefill token-ids were not written onto the request"
+    assert "prompt" not in prompt_dict
+    assert prompt_dict.get("use_system_prompt") == "en_recaption"
+    assert prompt_dict.get("modalities") == ["image"]
+
+    # stop_token_ids reached the AR-stage params, and only those. The diffusion
+    # stage must be left untouched: OmniDiffusionSamplingParams has no
+    # stop_token_ids field, so the helper skipping it means the attribute is
+    # never created (reading it would raise AttributeError, not return None).
+    assert ar_params.stop_token_ids, "AR stage did not receive stop_token_ids"
+    assert not hasattr(diffusion_params, "stop_token_ids"), "diffusion stage params must not receive stop_token_ids"
+
+
+@pytest.mark.parametrize(
+    "module_name,relpath,cli_flag",
+    [
+        ("_shared_text_to_image_trc", "examples/offline_inference/text_to_image/text_to_image.py", False),
+        ("_shared_text_to_image_trc", "examples/offline_inference/text_to_image/text_to_image.py", True),
+    ],
+)
+def test_apply_ar_stage_inputs_respects_trust_remote_code_option(
+    monkeypatch: pytest.MonkeyPatch, module_name: str, relpath: str, cli_flag: bool
+) -> None:
+    """The tokenizer load must use the caller's resolved --trust-remote-code
+    value, not silently default to True regardless of user intent."""
+    from vllm_omni.model_extras import get_ar_input_builder
+
+    module = _load_module(module_name, relpath)
+    tokenizer_calls = _patch_tokenizer_loader(monkeypatch, module)
+    ar_input_builder = get_ar_input_builder("HunyuanImage3ForCausalMM")
+
+    module._apply_ar_stage_inputs(
+        ar_input_builder,
+        model="tencent/HunyuanImage-3.0-Instruct",
+        prompt_text="a red panda",
+        extra_body={},
+        num_images=0,
+        height=1024,
+        width=1024,
+        prompt_dict={},
+        sampling_params_list=[_FakeARSamplingParams(), OmniDiffusionSamplingParams()],
+        trust_remote_code=cli_flag,
+    )
+
+    assert tokenizer_calls[0]["trust_remote_code"] is cli_flag
+
+
+@pytest.mark.parametrize(
+    "module_name,relpath",
+    [
+        ("_shared_text_to_image_validate", "examples/offline_inference/text_to_image/text_to_image.py"),
+        ("_shared_image_edit_validate", "examples/offline_inference/image_to_image/image_edit.py"),
+    ],
+)
+def test_apply_ar_stage_inputs_runs_declared_tokenizer_validator(
+    monkeypatch: pytest.MonkeyPatch, module_name: str, relpath: str
+) -> None:
+    """get_ar_tokenizer_validator's hook must run on a real tokenizer load,
+    and a validation failure must propagate -- not get swallowed by the
+    tokenizer-load try/except (which would silently fall back to the string
+    prompt instead of surfacing the drift)."""
+    from vllm_omni.model_extras import get_ar_input_builder, get_ar_tokenizer_validator
+
+    module = _load_module(module_name, relpath)
+    _patch_tokenizer_loader(monkeypatch, module)
+    ar_input_builder = get_ar_input_builder("HunyuanImage3ForCausalMM")
+    validator = get_ar_tokenizer_validator("HunyuanImage3ForCausalMM")
+    assert validator is not None
+
+    # The fake tokenizer's ids intentionally don't match
+    # HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS, so the real validator must reject it.
+    with pytest.raises(ValueError, match="no longer match"):
+        module._apply_ar_stage_inputs(
+            ar_input_builder,
+            model="tencent/HunyuanImage-3.0-Instruct",
+            prompt_text="a red panda",
+            extra_body={},
+            num_images=0,
+            height=1024,
+            width=1024,
+            prompt_dict={},
+            sampling_params_list=[_FakeARSamplingParams(), OmniDiffusionSamplingParams()],
+            trust_remote_code=False,
+            validate_tokenizer=validator,
+        )
+
+
+@pytest.mark.parametrize(
+    "module_name,relpath,required_argv",
+    [
+        ("_shared_text_to_image_cli", "examples/offline_inference/text_to_image/text_to_image.py", []),
+        # image_edit.py's parser declares --image/--prompt as required (no defaults);
+        # omitting them makes parse_args() raise SystemExit(2) before the assertion
+        # below is ever reached.
+        (
+            "_shared_image_edit_cli",
+            "examples/offline_inference/image_to_image/image_edit.py",
+            ["--image", "dummy.png", "--prompt", "x"],
+        ),
+    ],
+)
+@pytest.mark.parametrize("argv,expected", [([], False), (["--trust-remote-code"], True)])
+def test_trust_remote_code_flag_parses_from_real_cli(
+    monkeypatch: pytest.MonkeyPatch,
+    module_name: str,
+    relpath: str,
+    required_argv: list[str],
+    argv: list[str],
+    expected: bool,
+) -> None:
+    """The other tests in this file call _apply_ar_stage_inputs directly with
+    an already-resolved trust_remote_code bool, which never exercises
+    argparse itself. Closing that specific gap doesn't need a GPU or a model
+    load -- parse_args() only touches sys.argv -- so cover it here instead of
+    the (expensive, one-model-load-per-invocation) full CLI subprocess a
+    stricter GPU e2e test would need."""
+    module = _load_module(module_name, relpath)
+    monkeypatch.setattr(sys, "argv", ["prog", *required_argv, *argv])
+    args = module.parse_args()
+    assert args.trust_remote_code is expected
+
+
+@pytest.mark.parametrize(
+    "module_name,relpath",
+    [
+        ("_shared_text_to_image_failfast", "examples/offline_inference/text_to_image/text_to_image.py"),
+        ("_shared_image_edit_failfast", "examples/offline_inference/image_to_image/image_edit.py"),
+    ],
+)
+def test_apply_ar_stage_inputs_tokenizer_load_failure_raises_by_default(
+    monkeypatch: pytest.MonkeyPatch, module_name: str, relpath: str
+) -> None:
+    """A tokenizer load failure must surface as an error by default, not
+    silently degrade to the string-prompt form -- that degraded form can
+    complete "successfully" with subtly wrong stop tokens, masking the real
+    problem (e.g. a missing --trust-remote-code)."""
+    from vllm_omni.model_extras import get_ar_input_builder
+
+    module = _load_module(module_name, relpath)
+    _patch_failing_tokenizer_loader(monkeypatch)
+    ar_input_builder = get_ar_input_builder("HunyuanImage3ForCausalMM")
+
+    with pytest.raises(OSError, match="simulated tokenizer load failure"):
+        module._apply_ar_stage_inputs(
+            ar_input_builder,
+            model="tencent/HunyuanImage-3.0-Instruct",
+            prompt_text="a red panda",
+            extra_body={},
+            num_images=0,
+            height=1024,
+            width=1024,
+            prompt_dict={},
+            sampling_params_list=[_FakeARSamplingParams(), OmniDiffusionSamplingParams()],
+            trust_remote_code=False,
+        )
+
+
+@pytest.mark.parametrize(
+    "module_name,relpath",
+    [
+        ("_shared_text_to_image_fallback", "examples/offline_inference/text_to_image/text_to_image.py"),
+        ("_shared_image_edit_fallback", "examples/offline_inference/image_to_image/image_edit.py"),
+    ],
+)
+def test_apply_ar_stage_inputs_allow_tokenizer_fallback_option(
+    monkeypatch: pytest.MonkeyPatch, module_name: str, relpath: str
+) -> None:
+    """With allow_tokenizer_fallback=True (the CLI's --allow-tokenizer-fallback),
+    a tokenizer load failure degrades to the string-prompt form instead of
+    raising -- for a deliberate offline-compat run, not the default."""
+    from vllm_omni.model_extras import get_ar_input_builder
+
+    module = _load_module(module_name, relpath)
+    _patch_failing_tokenizer_loader(monkeypatch)
+    ar_input_builder = get_ar_input_builder("HunyuanImage3ForCausalMM")
+
+    prompt_dict: dict[str, Any] = {}
+    module._apply_ar_stage_inputs(
+        ar_input_builder,
+        model="tencent/HunyuanImage-3.0-Instruct",
+        prompt_text="a red panda",
+        extra_body={"bot_task": "recaption"},
+        num_images=0,
+        height=1024,
+        width=1024,
+        prompt_dict=prompt_dict,
+        sampling_params_list=[_FakeARSamplingParams(), OmniDiffusionSamplingParams()],
+        trust_remote_code=False,
+        allow_tokenizer_fallback=True,
+    )
+
+    # No tokenizer -> string-prompt form, not the byte-for-byte token-ids form.
+    assert "prompt_token_ids" not in prompt_dict
+    assert prompt_dict.get("prompt") and "a red panda" in prompt_dict["prompt"]