From a7ab91fb85b1bb513d34f0a19c3289c55a38b013 Mon Sep 17 00:00:00 2001 From: SuyanLi <126558907+suyanli220@users.noreply.github.com> Date: Fri, 11 Sep 2026 16:47:25 +0800 Subject: [PATCH] [Example] Migrate HunyuanImage-3.0 to model_extras + shared task examples (#5559) Signed-off-by: suyanli220 Signed-off-by: suyan.li Co-authored-by: Claude Opus 4.8 Co-authored-by: suyan.li Co-authored-by: Hongsheng Liu --- .../offline_inference/hunyuan_image3.md | 35 -- .../hunyuan_image3/README.md | 154 -------- .../hunyuan_image3/end2end.py | 321 ----------------- .../image_to_image/image_edit.py | 126 +++++++ .../text_to_image/text_to_image.py | 126 +++++++ recipes/Tencent/HunyuanImage-3.0-Instruct.md | 52 ++- .../hunyuan_image3/test_prompt_utils.py | 129 ++++++- .../offline_inference/test_hunyuanimage3.py | 122 ++++++- .../openai_api/test_image_server.py | 69 ++++ ...test_serving_chat_multistage_generation.py | 81 +++++ tests/model_extras/test_model_extras.py | 241 +++++++++++++ .../test_shared_script_ar_integration.py | 340 ++++++++++++++++++ .../models/hunyuan_image3/prompt_utils.py | 136 +++++++ vllm_omni/entrypoints/openai/serving_chat.py | 11 +- vllm_omni/model_extras/__init__.py | 4 + vllm_omni/model_extras/hunyuan_image3.py | 171 ++++++++- vllm_omni/model_extras/registry.py | 53 +++ 17 files changed, 1627 insertions(+), 544 deletions(-) delete mode 100644 docs/user_guide/examples/offline_inference/hunyuan_image3.md delete mode 100644 examples/offline_inference/hunyuan_image3/README.md delete mode 100644 examples/offline_inference/hunyuan_image3/end2end.py create mode 100644 tests/model_extras/test_shared_script_ar_integration.py diff --git a/docs/user_guide/examples/offline_inference/hunyuan_image3.md b/docs/user_guide/examples/offline_inference/hunyuan_image3.md deleted file mode 100644 index 636aa0c11..000000000 --- a/docs/user_guide/examples/offline_inference/hunyuan_image3.md +++ /dev/null @@ -1,35 +0,0 @@ -# HunyuanImage-3.0 Image-to-Text Inference - -Source . - - -This example demonstrates how to run HunyuanImage-3.0 Image-to-Text with the vLLM-Omni. - -## Local CLI Usage - -Download the example image: - -```bash -wget https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg -``` - -Run example: - -```bash -python image_to_text.py \ - --image cherry_blossom.jpg \ - --prompt "<|startoftext|>You are an assistant that understands images and outputs text.Describe the content of the picture." -``` - -Key arguments: - -- `--model`: Model used. Default is: tencent/HunyuanImage-3.0-Instruct (Optional). -- `--image`: Path to input image (required). -- `--prompt`: Text description used to guide image understanding (required). - -## Example materials - -??? abstract "image_to_text.py" - ``````py - --8<-- "examples/offline_inference/hunyuan_image3/image_to_text.py" - `````` diff --git a/examples/offline_inference/hunyuan_image3/README.md b/examples/offline_inference/hunyuan_image3/README.md deleted file mode 100644 index b4d5af2c7..000000000 --- a/examples/offline_inference/hunyuan_image3/README.md +++ /dev/null @@ -1,154 +0,0 @@ -# HunyuanImage-3.0-Instruct - -This example runs HunyuanImage-3.0-Instruct offline with the unified deploy -YAMLs under `vllm_omni/deploy/`. - -## Deploy Configs - -| File | Topology | Default use | -| :--- | :--- | :--- | -| `vllm_omni/deploy/hunyuan_image_3_moe.yaml` | AR + DiT | Default for `text2img` and `img2img`. | -| `vllm_omni/deploy/hunyuan_image3_ar.yaml` | AR only | Default for `img2text` and `text2text`. | -| `vllm_omni/deploy/hunyuan_image3_dit.yaml` | DiT only | Standalone diffusion stage. Pass it explicitly with `--deploy-config`. | - -The example chooses a deploy config automatically when `--deploy-config` is -omitted: - -| `--modality` | `mode` passed to Omni | Default deploy | -| :--- | :--- | :--- | -| `text2img` | `text-to-image` | `hunyuan_image_3_moe.yaml` | -| `img2img` | `image-editing` | `hunyuan_image_3_moe.yaml` | -| `img2text` | `image-to-text` | `hunyuan_image3_ar.yaml` | -| `text2text` | `text-to-text` | `hunyuan_image3_ar.yaml` | - -`--modality` is an offline example convenience flag. It maps to the internal -`mode` argument passed to `Omni(...)` by this script. HunyuanImage3 uses -separate deploy YAMLs for AR + DiT, AR-only, and DiT-only topologies, so the -stage topology is selected by the deploy file rather than by YAML mode -overrides. - -Online serving does not expose a `--modality` flag or accept `mode` as an API -request field. Choose the deploy topology when starting the server with -`--deploy-config`, then use the OpenAI-compatible endpoint and request shape for -the scenario. The `modalities` request field is used by the chat completions -path; the image endpoints infer the image task from the endpoint and payload. - -| Online scenario | Server deploy | Request | -| :--- | :--- | :--- | -| Text to image | `--deploy-config vllm_omni/deploy/hunyuan_image_3_moe.yaml` | `POST /v1/images/generations`, or `POST /v1/chat/completions` with `"modalities": ["image"]`. | -| Image editing | `--deploy-config vllm_omni/deploy/hunyuan_image_3_moe.yaml` | `POST /v1/images/edits`. | -| Image/text to text | `--deploy-config vllm_omni/deploy/hunyuan_image3_ar.yaml` | `POST /v1/chat/completions` for text output, for example with `"modalities": ["text"]`. | -| DiT-only image generation | `--deploy-config vllm_omni/deploy/hunyuan_image3_dit.yaml` | `POST /v1/images/generations`. | - -## Run Examples - -Text to image, using the default AR + DiT deploy: - -```bash -python examples/offline_inference/hunyuan_image3/end2end.py \ - --model tencent/HunyuanImage-3.0-Instruct \ - --modality text2img \ - --prompts "A cute cat sitting on a windowsill watching the sunset" -``` - -Image editing, using the default AR + DiT deploy: - -```bash -python examples/offline_inference/hunyuan_image3/end2end.py \ - --model tencent/HunyuanImage-3.0-Instruct \ - --modality img2img \ - --image-path /path/to/image.png \ - --prompts "Make the petals neon pink" -``` - -Image to text, using the AR-only deploy: - -```bash -python examples/offline_inference/hunyuan_image3/end2end.py \ - --model tencent/HunyuanImage-3.0-Instruct \ - --modality img2text \ - --image-path /path/to/image.jpg \ - --prompts "Describe the content of the picture." -``` - -Text to text, using the AR-only deploy: - -```bash -python examples/offline_inference/hunyuan_image3/end2end.py \ - --model tencent/HunyuanImage-3.0-Instruct \ - --modality text2text \ - --prompts "What is the capital of France?" -``` - -Standalone DiT, using the DiT-only deploy explicitly: - -```bash -python examples/offline_inference/hunyuan_image3/end2end.py \ - --model tencent/HunyuanImage-3.0-Instruct \ - --modality text2img \ - --deploy-config vllm_omni/deploy/hunyuan_image3_dit.yaml \ - --prompts "A cinematic portrait of an astronaut in a greenhouse" -``` - -Override the default full AR + DiT deploy explicitly: - -```bash -python examples/offline_inference/hunyuan_image3/end2end.py \ - --model tencent/HunyuanImage-3.0-Instruct \ - --modality text2img \ - --deploy-config vllm_omni/deploy/hunyuan_image_3_moe.yaml \ - --prompts "A cute cat" -``` - -## Additional Config - -You can pass diffusion worker `additional_config` from the offline example as a JSON object. -This maps to the upstream vLLM `VllmConfig.additional_config` platform extension field: -https://docs.vllm.ai/en/stable/api/vllm/config/#vllm.config.VllmConfig.additional_config - -```bash -python end2end.py --modality text2img \ - --prompts "A cute cat" \ - --additional-config '{"torchair_graph_config":{"enabled":true}}' -``` - - -## Key Arguments - -| Argument | Description | -| :--- | :--- | -| `--deploy-config` | Optional path for unified deploy YAML overrides. | -| `--additional-config` | JSON object forwarded to diffusion worker `additional_config`. | -| `--modality` | Offline-only convenience flag. One of `text2img`, `img2img`, `img2text`, `text2text`. It selects prompt formatting, internal `mode`, and default deploy config for this script. Online serving uses `--deploy-config` plus the endpoint and, for chat completions, request `modalities` instead. | -| `--steps` | Number of diffusion inference steps for image generation. | -| `--guidance-scale` | Classifier-free guidance scale for image generation. | -| `--height`, `--width` | Output image size for `text2img`. | -| `--bot-task` | Override prompt mode. `none`, `think`, `recaption`, `think_recaption`, or `vanilla`. | -| `--sys-type` | Override the system prompt type, for example `en_unified` or `en_vanilla`. | -| `--vae-use-tiling` | Enable VAE tiling for memory reduction. | - -## Notes - -- `hunyuan_image3_ar.yaml` is a 4-card AR-only text/comprehension deploy. -- `hunyuan_image3_dit.yaml` is a single-stage DiT deploy with `stage_id: 0`. -- Platform-specific overrides are included in the unified deploy YAMLs. - -## Prompt Format - -HunyuanImage-3.0-Instruct uses an instruct chat template: - -```text -<|startoftext|>{system_prompt} - -User: {?}{user_prompt} - -Assistant: {trigger_tag?} -``` - -- ``: Placeholder for each input image (single token; expanded by the multimodal pipeline). -- Trigger tags: `` for CoT and `` for recaptioning, placed after `Assistant: `. -- System prompt: Auto-selected from `task` and `bot_task`. -- `bot_task='vanilla'` with `task='t2i'` uses the bare pretrain template. - -The shared `vllm_omni.diffusion.models.hunyuan_image3.prompt_utils.build_prompt_tokens()` -helper handles segment-by-segment tokenization and matches HF `apply_chat_template`. diff --git a/examples/offline_inference/hunyuan_image3/end2end.py b/examples/offline_inference/hunyuan_image3/end2end.py deleted file mode 100644 index fbe4289be..000000000 --- a/examples/offline_inference/hunyuan_image3/end2end.py +++ /dev/null @@ -1,321 +0,0 @@ -""" -HunyuanImage-3.0-Instruct unified end-to-end inference script. -""" - -import argparse -import json -import os -from pathlib import Path - -from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import ( - MAX_IMAGES_PER_REQUEST, - build_prompt_tokens, - resolve_stop_token_ids, - resolve_sys_type, -) -from vllm_omni.entrypoints.omni import Omni -from vllm_omni.inputs.data import OmniPromptType - -_REPO_ROOT = Path(__file__).resolve().parents[3] -_DEFAULT_DEPLOY_CONFIG = str(_REPO_ROOT / "vllm_omni" / "deploy" / "hunyuan_image_3_moe.yaml") -_DEFAULT_AR_DEPLOY_CONFIG = str(_REPO_ROOT / "vllm_omni" / "deploy" / "hunyuan_image3_ar.yaml") - -_MODALITY_TASK_MAP: dict[str, tuple[str, str | None]] = { - "text2img": ("t2i", "think"), - "img2img": ("it2i", "think"), - "img2text": ("i2t", None), - "text2text": ("t2t", None), -} - -_MODALITY_DEFAULT_DEPLOY_CONFIG = { - "text2img": _DEFAULT_DEPLOY_CONFIG, - "img2img": _DEFAULT_DEPLOY_CONFIG, - "img2text": _DEFAULT_AR_DEPLOY_CONFIG, - "text2text": _DEFAULT_AR_DEPLOY_CONFIG, -} - -_MODALITY_MODE = { - "text2img": "text-to-image", - "img2img": "image-editing", - "img2text": "image-to-text", - "text2text": "text-to-text", -} - - -def parse_args(): - parser = argparse.ArgumentParser(description="HunyuanImage-3.0-Instruct end-to-end inference.") - parser.add_argument("--model", default="tencent/HunyuanImage-3.0-Instruct", help="Model name or local path.") - parser.add_argument( - "--modality", - default="text2img", - choices=list(_MODALITY_TASK_MAP), - ) - parser.add_argument("--prompts", nargs="+", default=None, help="Input text prompts.") - parser.add_argument( - "--image-path", - type=str, - default=None, - help="Input image path(s) for img2img/img2text. Comma-separated for multi-image (up to 3).", - ) - parser.add_argument("--output", type=str, default=".", help="Output directory to save results.") - parser.add_argument("--steps", type=int, default=50, help="Number of inference steps.") - parser.add_argument("--guidance-scale", type=float, default=5.0, help="Classifier-free guidance scale.") - parser.add_argument("--seed", type=int, default=42, help="Random seed.") - parser.add_argument("--height", type=int, default=None, help="Output image height.") - parser.add_argument("--width", type=int, default=None, help="Output image width.") - parser.add_argument("--vae-use-tiling", action="store_true", help="Enable VAE tiling.") - parser.add_argument( - "--bot-task", - type=str, - default=None, - choices=["none", "think", "recaption", "think_recaption", "vanilla"], - help="Override prompt mode. Default: auto from --modality.", - ) - parser.add_argument("--sys-type", type=str, default=None, help="Override system prompt type.") - parser.add_argument("--deploy-config", type=str, default=None, help="Custom deploy YAML path.") - parser.add_argument( - "--stream", - action="store_true", - default=False, - help="Enable streaming CoT display; print AR text token-by-token in real time.", - ) - parser.add_argument("--log-stats", action="store_true", default=False) - parser.add_argument("--init-timeout", type=int, default=300, help="Initialization timeout in seconds.") - parser.add_argument("--enforce-eager", action="store_true", help="Disable torch.compile.") - parser.add_argument( - "--diffusion-kv-cache-dtype", - type=str, - default=None, - help="Diffusion attention KV cache dtype, for example 'fp8'. Separate from vLLM --kv-cache-dtype.", - ) - parser.add_argument( - "--diffusion-kv-cache-skip-steps", - type=str, - default=None, - help="Denoising step selector to keep diffusion KV cache in native dtype, for example '0,1,4-6'.", - ) - parser.add_argument( - "--diffusion-kv-cache-skip-layers", - type=str, - default=None, - help="Transformer layer selector to keep diffusion KV cache in native dtype, for example '0-2,10'.", - ) - parser.add_argument( - "--additional-config", - type=str, - default=None, - help=( - "JSON object forwarded to Omni/additional_config, for example " - '\'{"torchair_graph_config":{"enabled":true}}\'. Different platforms may support different ' - "configs. Make sure the configs are valid for the platform you are using. " - "Contents must be hashable." - ), - ) - - return parser.parse_args() - - -def parse_additional_config(raw_value: str | None) -> dict | None: - """Parse a JSON string into an additional_config mapping.""" - if raw_value is None: - return None - - try: - additional_config = json.loads(raw_value) - except json.JSONDecodeError as exc: - raise ValueError(f"Invalid --additional-config JSON: {exc}") from exc - - if additional_config is None: - return None - if not isinstance(additional_config, dict): - raise ValueError(f"--additional-config must decode to a JSON object, got {type(additional_config).__name__}") - return additional_config - - -def main(): - args = parse_args() - os.makedirs(args.output, exist_ok=True) - additional_config = parse_additional_config(args.additional_config) - - task, default_bot_task = _MODALITY_TASK_MAP[args.modality] - if args.bot_task is None: - bot_task: str | None = default_bot_task - elif args.bot_task == "none": - bot_task = None - else: - bot_task = args.bot_task - - deploy_config = args.deploy_config or _MODALITY_DEFAULT_DEPLOY_CONFIG[args.modality] - - omni_kwargs = { - "model": args.model, - "vae_use_tiling": args.vae_use_tiling, - "log_stats": args.log_stats, - "init_timeout": args.init_timeout, - "enforce_eager": args.enforce_eager, - "mode": _MODALITY_MODE[args.modality], - "diffusion_kv_cache_dtype": args.diffusion_kv_cache_dtype, - "diffusion_kv_cache_skip_steps": args.diffusion_kv_cache_skip_steps, - "diffusion_kv_cache_skip_layers": args.diffusion_kv_cache_skip_layers, - } - - if additional_config is not None: - omni_kwargs["additional_config"] = additional_config - omni_kwargs["deploy_config"] = deploy_config - - omni = Omni(**omni_kwargs) - - prompts = args.prompts or ["A cute cat"] - input_images: list = [] - if args.modality in ("img2img", "img2text"): - if not args.image_path: - raise ValueError(f"--image-path required for {args.modality}, got: {args.image_path}") - from PIL import Image - - image_paths = [p.strip() for p in args.image_path.split(",") if p.strip()] - if len(image_paths) > MAX_IMAGES_PER_REQUEST: - raise ValueError( - f"--image-path accepts at most {MAX_IMAGES_PER_REQUEST} images for " - f"HunyuanImage-3.0 IT2I, got {len(image_paths)}: {args.image_path}" - ) - for image_path in image_paths: - if not os.path.exists(image_path): - raise ValueError(f"Image path does not exist: {image_path}") - input_images.append(Image.open(image_path).convert("RGB")) - if not input_images: - raise ValueError(f"--image-path produced no usable paths: {args.image_path!r}") - - from transformers import AutoTokenizer - - tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True) - mm_image_payload = (input_images[0] if len(input_images) == 1 else input_images) if input_images else None - - formatted_prompts: list[OmniPromptType] = [] - for prompt in prompts: - build_kwargs: dict = {"task": task, "bot_task": bot_task, "sys_type": args.sys_type} - if input_images: - build_kwargs["num_images"] = len(input_images) - result = build_prompt_tokens(prompt, tokenizer, **build_kwargs) - token_ids = result.token_ids - effective_sys_type = args.sys_type or resolve_sys_type(bot_task) - - prompt_dict: dict = { - "prompt_token_ids": token_ids, - "prompt": prompt, - "use_system_prompt": effective_sys_type, - } - if args.modality == "text2img": - prompt_dict["modalities"] = ["image"] - elif args.modality == "img2img": - prompt_dict["modalities"] = ["image"] - prompt_dict["multi_modal_data"] = {"image": mm_image_payload} - prompt_dict["height"] = input_images[0].height - prompt_dict["width"] = input_images[0].width - elif args.modality == "img2text": - prompt_dict["modalities"] = ["text"] - prompt_dict["multi_modal_data"] = {"image": mm_image_payload} - else: - prompt_dict["modalities"] = ["text"] - formatted_prompts.append(prompt_dict) - - from vllm.sampling_params import RequestOutputKind - - params_list = list(omni.default_sampling_params_list) - - from vllm_omni.inputs.data import OmniDiffusionSamplingParams - - if (args.height is None) != (args.width is None): - raise ValueError("--height and --width must both be specified or both omitted.") - user_specified_size = args.height is not None and args.width is not None - if args.modality in ("img2text", "text2text"): - ar_image_size = "auto" - elif user_specified_size: - ar_image_size = f"{args.width}x{args.height}" - else: - ar_image_size = None - ar_stop_token_ids = resolve_stop_token_ids( - task=task, bot_task=bot_task, tokenizer=tokenizer, image_size=ar_image_size - ) - print( - f"[AR Config] task={task}, bot_task={bot_task}, image_size={ar_image_size}, stop_token_ids={ar_stop_token_ids}" - ) - for sp in params_list: - if isinstance(sp, OmniDiffusionSamplingParams): - sp.num_inference_steps = args.steps - sp.guidance_scale = args.guidance_scale - sp.guidance_scale_provided = True - if args.seed is not None: - sp.seed = args.seed - if args.modality == "text2img": - sp.height = args.height - sp.width = args.width - elif hasattr(sp, "stop_token_ids"): - sp.stop_token_ids = ar_stop_token_ids - # When --stream is set, request DELTA output from the AR stage - # so we can display CoT text token-by-token in real time. - if args.stream and hasattr(sp, "output_kind"): - sp.output_kind = RequestOutputKind.DELTA - - print(f"\n{'=' * 60}") - print("HunyuanImage-3.0 Generation Configuration:") - print(f" Model: {args.model}") - print(f" Modality: {args.modality}") - print(f" Prompt task: {task}") - print(f" Bot task: {bot_task}") - print(f" Deploy config: {deploy_config}") - print(f" Num stages: {omni.num_stages}") - if args.modality in ("text2img", "img2img"): - print(f" Inference steps: {args.steps}") - print(f" Guidance scale: {args.guidance_scale}") - print(f" Seed: {args.seed}") - print(f" diffusion_kv_cache_dtype: {args.diffusion_kv_cache_dtype}") - print(f" diffusion_kv_cache_skip_steps: {args.diffusion_kv_cache_skip_steps}") - print(f" diffusion_kv_cache_skip_layers: {args.diffusion_kv_cache_skip_layers}") - if args.modality == "text2img": - print(f" Output size: {args.width}x{args.height}") - if args.image_path: - print(f" Input image: {args.image_path}") - if additional_config is not None: - print(f" Additional config: {additional_config}") - print(f" Prompts: {prompts}") - print(f"{'=' * 60}\n") - - # When --stream is set, print AR CoT text token-by-token in real time. - # Otherwise, collect and print the full AR text once when stage 0 finishes. - omni_outputs = omni.generate( - prompts=formatted_prompts, - sampling_params_list=params_list, - py_generator=True, - use_tqdm=False, - ) - img_idx = 0 - for req_output in omni_outputs: - ro = req_output - stage_id = getattr(req_output, "stage_id", None) - - # AR stage text — each CompletionOutput.text is already a delta when - # output_kind=DELTA, so we can print it directly (matching the pattern - # in serving_chat.py). - if stage_id == 0 and ro and getattr(ro, "outputs", None): - for o in ro.outputs: - text = getattr(o, "text", "") or "" - if text: - print(text, end="", flush=True) - # Non-streaming: one shot with full text — emit a trailing newline. - if not args.stream: - print(flush=True) - - # Collect images from diffusion stage - images = getattr(req_output, "images", None) - if not images and ro and hasattr(ro, "images"): - images = ro.images - if images: - for j, img in enumerate(images): - save_path = os.path.join(args.output, f"output_{img_idx}_{j}.png") - img.save(save_path) - print(f"\n[Output] Saved image to {save_path}") - img_idx += 1 - - -if __name__ == "__main__": - main() diff --git a/examples/offline_inference/image_to_image/image_edit.py b/examples/offline_inference/image_to_image/image_edit.py index 136d99857..cec7014de 100644 --- a/examples/offline_inference/image_to_image/image_edit.py +++ b/examples/offline_inference/image_to_image/image_edit.py @@ -96,6 +96,7 @@ from typing import Any import torch from PIL import Image +from vllm_omni.diffusion.data import logger from vllm_omni.diffusion.utils.image_output import extract_images_from_outputs from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args from vllm_omni.entrypoints.omni import Omni @@ -103,6 +104,8 @@ from vllm_omni.entrypoints.openai.stage_params import clone_sampling_params from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.model_extras import ( build_image_to_image_prompt, + get_ar_input_builder, + get_ar_tokenizer_validator, get_extra_body_params, get_model_class_name, should_init_extra_args_for_non_diffusion_stages, @@ -110,6 +113,91 @@ from vllm_omni.model_extras import ( from vllm_omni.platforms import current_omni_platform +def _apply_ar_stage_inputs( + ar_input_builder: Any, + *, + model: str, + prompt_text: str, + extra_body: dict[str, Any], + num_images: int, + height: int | None, + width: int | None, + prompt_dict: dict[str, Any], + sampling_params_list: list[Any], + text_output: bool = False, + trust_remote_code: bool = False, + validate_tokenizer: Any | None = None, + allow_tokenizer_fallback: bool = False, +) -> None: + """Apply a model's declared AR-stage inputs to the request in place. + + Mirrors the helper in ``text_to_image.py``: loads the tokenizer for + byte-for-byte HF-parity AR tokenization, asks the model's + ``ar_input_builder`` for the AR prefill + stop tokens, and writes them + onto ``prompt_dict`` and the non-diffusion (AR) stage params. + Model-agnostic -- only the declared-hook contract is assumed. + + ``trust_remote_code`` must be threaded in from the caller's own resolved + ``--trust-remote-code`` flag -- this is a separate tokenizer load outside + the main engine, so it needs the same explicit user opt-in rather than + defaulting to trusting whatever ``--model`` was passed. + + ``validate_tokenizer``, when the model declares one via + ``get_ar_tokenizer_validator``, is called on a successfully-loaded real + tokenizer -- outside the load's try/except, so a validation failure (a + model/tokenizer revision drifting from hardcoded special-token ids) + raises instead of being swallowed into the string-prompt fallback. + + ``allow_tokenizer_fallback`` (default ``False``, i.e. fail fast): a + tokenizer load failure normally raises, since silently degrading to the + string-prompt form can produce a request that completes with subtly + wrong stop tokens instead of surfacing the real problem (missing + ``--trust-remote-code``, network/cache issue, etc.). Set ``True`` only + for explicit unit tests or a deliberate offline-compat run where the + caller has already decided a degraded prompt is acceptable. + """ + try: + from transformers import AutoTokenizer + + ar_tokenizer: Any | None = AutoTokenizer.from_pretrained(model, trust_remote_code=trust_remote_code) + except Exception as exc: # noqa: BLE001 - re-raised unless the caller opts into the fallback + if not allow_tokenizer_fallback: + raise + logger.warning(f"AR tokenizer load failed ({exc}); falling back to string prompt (no BPE parity).") + ar_tokenizer = None + + if ar_tokenizer is not None and validate_tokenizer is not None: + validate_tokenizer(ar_tokenizer) + + ar_inputs = ar_input_builder( + prompt=prompt_text, + tokenizer=ar_tokenizer, + extra_body=extra_body, + num_images=num_images, + height=height, + width=width, + text_output=text_output, + ) + + if ar_inputs.prompt_token_ids is not None: + prompt_dict["prompt_token_ids"] = ar_inputs.prompt_token_ids + elif ar_inputs.prompt is not None: + prompt_dict["prompt"] = ar_inputs.prompt + if ar_inputs.use_system_prompt: + prompt_dict["use_system_prompt"] = ar_inputs.use_system_prompt + prompt_dict["modalities"] = ar_inputs.modalities + + # Apply stop_token_ids to exactly the stage(s) the builder names -- not + # "whichever stage isn't a diffusion stage," which would misfire on any + # future topology with more than one non-diffusion stage. + for stage_index in ar_inputs.stage_indices: + if stage_index >= len(sampling_params_list): + continue + stage_params = sampling_params_list[stage_index] + if not isinstance(stage_params, OmniDiffusionSamplingParams) and hasattr(stage_params, "stop_token_ids"): + stage_params.stop_token_ids = ar_inputs.stop_token_ids + + def parse_profiler_config(value: str) -> dict[str, Any]: try: config = json.loads(value) @@ -419,6 +507,22 @@ def parse_args() -> argparse.Namespace: default=None, help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.', ) + parser.add_argument( + "--trust-remote-code", + action="store_true", + help="Trust and execute custom modeling code from the model repo (required by e.g. HunyuanImage-3.0).", + ) + parser.add_argument( + "--allow-tokenizer-fallback", + action="store_true", + help=( + "For models with a declared ar_input_builder (e.g. HunyuanImage-3.0): if loading " + "the AR tokenizer fails, degrade to the string-prompt form (no BPE parity) instead " + "of failing the run. Off by default -- a tokenizer load failure usually means a " + "missing --trust-remote-code or a network/cache issue that's worth surfacing, not " + "silently masking with a possibly-wrong prompt." + ), + ) return parser.parse_args() @@ -492,6 +596,8 @@ def main(): ) if args.enforce_eager is not None: omni_kwargs["enforce_eager"] = args.enforce_eager + if args.trust_remote_code: + omni_kwargs["trust_remote_code"] = True if args.deploy_config: omni_kwargs["deploy_config"] = args.deploy_config omni = Omni(**omni_kwargs) @@ -584,6 +690,26 @@ def main(): if not diffusion_replaced and len(sampling_params_list) == 1: sampling_params_list = [diffusion_params] + # Models with an AR text stage (e.g. HunyuanImage3 image-editing) declare an + # ar_input_builder; build the AR prefill + stop tokens declaratively from the + # plain prompt + reference image count + extra_args. Others are untouched. + ar_input_builder = get_ar_input_builder(model_class_name) + if ar_input_builder is not None: + _apply_ar_stage_inputs( + ar_input_builder, + model=args.model, + prompt_text=args.prompt, + extra_body=extra_args_from_cli, + num_images=len(input_images), + height=args.height, + width=args.width, + prompt_dict=prompt_dict, + sampling_params_list=sampling_params_list, + trust_remote_code=args.trust_remote_code, + validate_tokenizer=get_ar_tokenizer_validator(model_class_name), + allow_tokenizer_fallback=args.allow_tokenizer_fallback, + ) + outputs = omni.generate(prompt_dict, sampling_params_list=sampling_params_list) generation_end = time.perf_counter() generation_time = generation_end - generation_start diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index 2ee09e178..adcc10669 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -24,6 +24,8 @@ from vllm_omni.model_extras import ( build_text_to_image_prompt as build_model_text_to_image_prompt, ) from vllm_omni.model_extras import ( + get_ar_input_builder, + get_ar_tokenizer_validator, get_extra_body_params, get_model_class_name, should_init_extra_args_for_non_diffusion_stages, @@ -377,10 +379,111 @@ def parse_args() -> argparse.Namespace: default=None, help="Supplementary auxiliary text encoder parameters model name or path (especially for Hidream-l1-full).", ) + parser.add_argument( + "--trust-remote-code", + action="store_true", + help="Trust and execute custom modeling code from the model repo (required by e.g. HunyuanImage-3.0).", + ) + parser.add_argument( + "--allow-tokenizer-fallback", + action="store_true", + help=( + "For models with a declared ar_input_builder (e.g. HunyuanImage-3.0): if loading " + "the AR tokenizer fails, degrade to the string-prompt form (no BPE parity) instead " + "of failing the run. Off by default -- a tokenizer load failure usually means a " + "missing --trust-remote-code or a network/cache issue that's worth surfacing, not " + "silently masking with a possibly-wrong prompt." + ), + ) current_omni_platform.pre_register_and_update(parser) return parser.parse_args() +def _apply_ar_stage_inputs( + ar_input_builder: Any, + *, + model: str, + prompt_text: str, + extra_body: dict[str, Any], + num_images: int, + height: int | None, + width: int | None, + prompt_dict: dict[str, Any], + sampling_params_list: list[Any], + text_output: bool = False, + trust_remote_code: bool = False, + validate_tokenizer: Any | None = None, + allow_tokenizer_fallback: bool = False, +) -> None: + """Apply a model's declared AR-stage inputs to the request in place. + + Loads the model tokenizer (for byte-for-byte HF-parity segment + tokenization) and asks the model's ``ar_input_builder`` for the AR + prefill + stop tokens, then writes them onto ``prompt_dict`` and the + non-diffusion (AR) stage sampling params. Kept model-agnostic: the + example only knows the declared-hook contract. + + ``trust_remote_code`` must be threaded in from the caller's own resolved + ``--trust-remote-code`` flag -- this is a separate tokenizer load outside + the main engine, so it needs the same explicit user opt-in rather than + defaulting to trusting whatever ``--model`` was passed. + + ``validate_tokenizer``, when the model declares one via + ``get_ar_tokenizer_validator``, is called on a successfully-loaded real + tokenizer -- outside the load's try/except, so a validation failure (a + model/tokenizer revision drifting from hardcoded special-token ids) + raises instead of being swallowed into the string-prompt fallback. + + ``allow_tokenizer_fallback`` (default ``False``, i.e. fail fast): a + tokenizer load failure normally raises, since silently degrading to the + string-prompt form can produce a request that completes with subtly + wrong stop tokens instead of surfacing the real problem (missing + ``--trust-remote-code``, network/cache issue, etc.). Set ``True`` only + for explicit unit tests or a deliberate offline-compat run where the + caller has already decided a degraded prompt is acceptable. + """ + try: + from transformers import AutoTokenizer + + ar_tokenizer: Any | None = AutoTokenizer.from_pretrained(model, trust_remote_code=trust_remote_code) + except Exception as exc: # noqa: BLE001 - re-raised unless the caller opts into the fallback + if not allow_tokenizer_fallback: + raise + logger.warning(f"AR tokenizer load failed ({exc}); falling back to string prompt (no BPE parity).") + ar_tokenizer = None + + if ar_tokenizer is not None and validate_tokenizer is not None: + validate_tokenizer(ar_tokenizer) + + ar_inputs = ar_input_builder( + prompt=prompt_text, + tokenizer=ar_tokenizer, + extra_body=extra_body, + num_images=num_images, + height=height, + width=width, + text_output=text_output, + ) + + if ar_inputs.prompt_token_ids is not None: + prompt_dict["prompt_token_ids"] = ar_inputs.prompt_token_ids + elif ar_inputs.prompt is not None: + prompt_dict["prompt"] = ar_inputs.prompt + if ar_inputs.use_system_prompt: + prompt_dict["use_system_prompt"] = ar_inputs.use_system_prompt + prompt_dict["modalities"] = ar_inputs.modalities + + # Apply stop_token_ids to exactly the stage(s) the builder names -- not + # "whichever stage isn't a diffusion stage," which would misfire on any + # future topology with more than one non-diffusion stage. + for stage_index in ar_inputs.stage_indices: + if stage_index >= len(sampling_params_list): + continue + stage_params = sampling_params_list[stage_index] + if not isinstance(stage_params, OmniDiffusionSamplingParams) and hasattr(stage_params, "stop_token_ids"): + stage_params.stop_token_ids = ar_inputs.stop_token_ids + + def main(): args = parse_args() generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed) @@ -469,6 +572,8 @@ def main(): omni_kwargs["tensor_parallel_size"] = args.tensor_parallel_size if args.enforce_eager is not None: omni_kwargs["enforce_eager"] = args.enforce_eager + if args.trust_remote_code: + omni_kwargs["trust_remote_code"] = True if args.deploy_config: omni_kwargs["deploy_config"] = args.deploy_config if use_nextstep: @@ -613,6 +718,27 @@ def main(): if not diffusion_replaced and len(sampling_params_list) == 1: sampling_params_list = [diffusion_params] + # Models with an AR text stage (e.g. HunyuanImage3) declare an + # ar_input_builder. When present, build the AR prefill token-ids and AR + # stop-token-ids declaratively from the plain prompt + extra_body, so this + # example stays model-agnostic. Models without one are untouched. + ar_input_builder = get_ar_input_builder(model_class_name) + if ar_input_builder is not None: + _apply_ar_stage_inputs( + ar_input_builder, + model=args.model, + prompt_text=args.prompt, + extra_body=user_extra, + num_images=0, + height=args.height, + width=args.width, + prompt_dict=prompt_dict, + sampling_params_list=sampling_params_list, + trust_remote_code=args.trust_remote_code, + validate_tokenizer=get_ar_tokenizer_validator(model_class_name), + allow_tokenizer_fallback=args.allow_tokenizer_fallback, + ) + outputs = omni.generate(prompt_dict, sampling_params_list=sampling_params_list) generation_end = time.perf_counter() diff --git a/recipes/Tencent/HunyuanImage-3.0-Instruct.md b/recipes/Tencent/HunyuanImage-3.0-Instruct.md index 4614df382..eb5ccedf6 100644 --- a/recipes/Tencent/HunyuanImage-3.0-Instruct.md +++ b/recipes/Tencent/HunyuanImage-3.0-Instruct.md @@ -8,7 +8,7 @@ - Vendor: Tencent Hunyuan - Model: `tencent/HunyuanImage-3.0-Instruct` -- Task: Text-to-image generation; text-to-text and image-to-text understanding +- Task: Text-to-image generation, image editing; text-to-text and image-to-text understanding - Mode: Offline understanding, online serving, and performance benchmarking - Maintainer: Community @@ -31,8 +31,10 @@ FP8/NVFP4 configuration: ## References - Model: -- Existing model-specific offline example: - [`examples/offline_inference/hunyuan_image3`](../../examples/offline_inference/hunyuan_image3) +- Shared T2I example: + [`examples/offline_inference/text_to_image/text_to_image.py`](../../examples/offline_inference/text_to_image/text_to_image.py) +- Shared IT2I example: + [`examples/offline_inference/image_to_image/image_edit.py`](../../examples/offline_inference/image_to_image/image_edit.py) - Shared T2T/I2T example: [`examples/offline_inference/x_to_text/x_to_text.py`](../../examples/offline_inference/x_to_text/x_to_text.py) - Related PRs: @@ -66,8 +68,48 @@ python examples/offline_inference/x_to_text/x_to_text.py \ ``` The AR-only default uses four GPUs. Pass `--deploy-config` to override the -layout. Image-output tasks remain documented by the existing model-specific -example and the DiT sections below. +layout. + +### Shared T2I/IT2I offline examples + +Text-to-image and image-editing route through the shared task examples, with +all HunyuanImage-3.0-specific knobs (`bot_task`, `use_system_prompt`, +`system_prompt`, `negative_prompt`) declared in +`vllm_omni/model_extras/hunyuan_image3.py` and passed via `--extra-body` / +`--extra-args`: + +```bash +python examples/offline_inference/text_to_image/text_to_image.py \ + --model tencent/HunyuanImage-3.0-Instruct \ + --deploy-config vllm_omni/deploy/hunyuan_image_3_moe.yaml \ + --trust-remote-code \ + --prompt "A cute cat sitting on a windowsill watching the sunset" \ + --height 1024 --width 1024 \ + --guidance-scale 5.0 --num-inference-steps 50 --seed 42 \ + --extra-body '{"bot_task": "think", "use_system_prompt": "en_recaption"}' \ + --output hunyuan_t2i.png +``` + +```bash +python examples/offline_inference/image_to_image/image_edit.py \ + --model tencent/HunyuanImage-3.0-Instruct \ + --deploy-config vllm_omni/deploy/hunyuan_image_3_moe.yaml \ + --trust-remote-code \ + --image /path/to/image.png \ + --prompt "Make the petals neon pink" \ + --guidance-scale 5.0 --num-inference-steps 50 --seed 42 \ + --extra-args '{"bot_task": "think"}' \ + --output hunyuan_edit.png +``` + +`image_edit.py` accepts up to 3 reference images (repeat `--image`) for +HunyuanImage-3.0's multi-image fusion. All four shared-script paths (t2i, +it2i, i2t, t2t) build the AR prefill and stop tokens through the same +declarative seam (`build_ar_stage_inputs` / `build_x_to_text_prompt`, both +wrapping `prompt_utils.build_ar_prompt_inputs`), so prompt formatting is +identical across them. The OpenAI server's `serving_chat.py` does not go +through this seam yet -- see that function's docstring for the resulting +`bot_task`-resolution divergence risk when omitted from `extra_body`. ### 4x H100/H800 80GB diff --git a/tests/diffusion/models/hunyuan_image3/test_prompt_utils.py b/tests/diffusion/models/hunyuan_image3/test_prompt_utils.py index 7cb7ea83a..7a467e359 100644 --- a/tests/diffusion/models/hunyuan_image3/test_prompt_utils.py +++ b/tests/diffusion/models/hunyuan_image3/test_prompt_utils.py @@ -13,9 +13,11 @@ from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import ( HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS, available_bot_tasks, available_tasks, + build_ar_prompt_inputs, build_prompt, build_prompt_tokens, resolve_stop_token_ids, + validate_special_token_ids, ) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -296,19 +298,42 @@ def _repo_root() -> pathlib.Path: return pathlib.Path(__file__).resolve().parents[4] -def test_end2end_routes_through_shared_prompt_utils(): - end2end_path = _repo_root() / "examples" / "offline_inference" / "hunyuan_image3" / "end2end.py" - tree = ast.parse(end2end_path.read_text(encoding="utf-8")) +def test_example_routes_through_shared_seam(): + """Neither the shared task examples nor the ``model_extras`` seam itself + should duplicate AR prompt tokenization. - local_func_names = {n.name for n in ast.walk(tree) if isinstance(n, ast.FunctionDef)} - assert not (local_func_names & {"build_prompt", "build_prompt_tokens"}) + After the model_extras migration, t2i / image-editing run through the + shared ``text_to_image.py`` / ``image_edit.py`` examples (both call + ``build_ar_stage_inputs``), and t2t / i2t run through the shared + ``x_to_text.py`` example (which calls ``build_x_to_text_prompt``). This + test pins that both ``model_extras.hunyuan_image3`` seam functions + delegate to ``prompt_utils.build_ar_prompt_inputs`` / ``build_ar_stage_inputs`` + rather than re-implementing tokenization inline. + """ + seam_path = _repo_root() / "vllm_omni" / "model_extras" / "hunyuan_image3.py" + seam_tree = ast.parse(seam_path.read_text(encoding="utf-8")) - imported_from_prompt_utils: set[str] = set() - for node in ast.walk(tree): + local_func_names = {n.name for n in ast.walk(seam_tree) if isinstance(n, ast.FunctionDef)} + assert not (local_func_names & {"build_prompt", "build_prompt_tokens", "resolve_stop_token_ids"}) + + # build_x_to_text_prompt must route through build_ar_stage_inputs, not + # re-derive the AR prefill / stop-token logic itself. + build_x_to_text_prompt_node = next( + n for n in ast.walk(seam_tree) if isinstance(n, ast.FunctionDef) and n.name == "build_x_to_text_prompt" + ) + called_names = { + node.func.id + for node in ast.walk(build_x_to_text_prompt_node) + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) + } + assert "build_ar_stage_inputs" in called_names + + # The model_extras seam itself must delegate to the shared prompt_utils. + routed_from_prompt_utils: set[str] = set() + for node in ast.walk(seam_tree): if isinstance(node, ast.ImportFrom) and node.module and node.module.endswith("hunyuan_image3.prompt_utils"): - imported_from_prompt_utils.update(alias.name for alias in node.names) - expected_imports = {"build_prompt_tokens", "resolve_stop_token_ids", "resolve_sys_type"} - assert expected_imports <= imported_from_prompt_utils + routed_from_prompt_utils.update(alias.name for alias in node.names) + assert "build_ar_prompt_inputs" in routed_from_prompt_utils _HUNYUAN_MODEL_ID = "tencent/HunyuanImage-3.0-Instruct" @@ -331,3 +356,87 @@ def test_segment_tokenize_diverges_from_full_string_encode(): full_ids = tok.encode(build_prompt(user_prompt, task="i2t", bot_task=None), add_special_tokens=False) assert seg_ids != full_ids assert len(seg_ids) >= len(full_ids) + + +def test_build_ar_prompt_inputs_tokenizer_path_matches_build_prompt_tokens(): + """With a tokenizer, the helper delegates to build_prompt_tokens (parity).""" + tok = FakeTokenizer() + ar = build_ar_prompt_inputs("a red panda", task="t2i", bot_task="think", tokenizer=tok) + expected = build_prompt_tokens("a red panda", FakeTokenizer(), task="t2i", bot_task="think") + assert ar.prompt is None + assert ar.prompt_token_ids == expected.token_ids + assert ar.system_prompt_type == expected.system_prompt_type + # think + no explicit size -> AR predicts ratio, so stop range is non-trivial. + assert len(ar.stop_token_ids) > 1 + + +def test_build_ar_prompt_inputs_string_fallback_without_tokenizer(): + """Without a tokenizer, the helper returns the build_prompt string form.""" + ar = build_ar_prompt_inputs("a red panda", task="t2i", bot_task="recaption") + assert ar.prompt_token_ids is None + assert ar.prompt == build_prompt("a red panda", task="t2i", bot_task="recaption") + assert ar.system_prompt_type + assert ar.stop_token_ids + + +def test_build_ar_prompt_inputs_explicit_size_stops_at_terminator(): + """Explicit image_size makes 'recaption' stop at only.""" + ar = build_ar_prompt_inputs("x", task="t2i", bot_task="recaption", image_size="1024x1024") + assert ar.stop_token_ids == [HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""]] + + +def test_build_ar_prompt_inputs_comprehension_stops_on_answer(): + """Text-output tasks (t2t / i2t) stop on .""" + ar = build_ar_prompt_inputs("hi", task="t2t") + assert ar.stop_token_ids == [HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""]] + + +class _MatchingTokenizer: + """Reports exactly HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS -- a "real" tokenizer + that hasn't drifted from the hardcoded table.""" + + def convert_tokens_to_ids(self, tok: str) -> int: + return HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[tok] + + +class _DriftedTokenizer: + """Reports a different id for one token -- simulates a model/tokenizer + revision bump shifting special-token ids out from under the hardcoded + table.""" + + def convert_tokens_to_ids(self, tok: str) -> int: + if tok == "": + return 999999 + return HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[tok] + + +class _BaseCheckpointTokenizer: + """Simulates tencent/HunyuanImage-3.0 (base): reports the same + model_class_name as -Instruct, but its vocab doesn't include the + Instruct-only - tokens -- convert_tokens_to_ids + falls back to unk_token_id for those, like a real HF tokenizer would.""" + + unk_token_id = 0 + _INSTRUCT_ONLY = {"", ""} + + def convert_tokens_to_ids(self, tok: str) -> int: + if tok in self._INSTRUCT_ONLY: + return self.unk_token_id + return HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[tok] + + +def test_validate_special_token_ids_passes_for_matching_tokenizer(): + validate_special_token_ids(_MatchingTokenizer()) # must not raise + + +def test_validate_special_token_ids_raises_on_drift(): + with pytest.raises(ValueError, match=".*expected 128023.*got 999999"): + validate_special_token_ids(_DriftedTokenizer()) + + +def test_validate_special_token_ids_tolerates_base_checkpoint_missing_tokens(): + """The base checkpoint's tokenizer not recognizing Instruct-only ratio + tokens (unk_token_id, not a real drifted id) must not be flagged as a + mismatch -- both checkpoints share model_class_name, so this hook fires + for either one.""" + validate_special_token_ids(_BaseCheckpointTokenizer()) # must not raise diff --git a/tests/e2e/offline_inference/test_hunyuanimage3.py b/tests/e2e/offline_inference/test_hunyuanimage3.py index 97179d963..8655bd63f 100644 --- a/tests/e2e/offline_inference/test_hunyuanimage3.py +++ b/tests/e2e/offline_inference/test_hunyuanimage3.py @@ -1,5 +1,8 @@ # ruff: noqa: E501 +import importlib.util +import sys from collections.abc import Generator +from pathlib import Path import pytest import torch @@ -10,13 +13,15 @@ from transformers import CLIPModel, CLIPProcessor from tests.helpers.runtime import OmniRunner from tests.helpers.stage_config import get_deploy_config_path from vllm_omni import Omni +from vllm_omni.entrypoints.openai.stage_params import clone_sampling_params from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.model_extras import get_ar_input_builder, get_ar_tokenizer_validator, get_model_class_name from vllm_omni.platforms import current_omni_platform PROMPT = "A brown and white dog is running on the grass" MODEL_NAME = "tencent/HunyuanImage-3.0" LOCAL_CLIP_PATH = "openai/clip-vit-base-patch32" -DEPLOY_CONFIG_PATH = get_deploy_config_path("hunyuan_image3.yaml") +DEPLOY_CONFIG_PATH = get_deploy_config_path("hunyuan_image_3_moe.yaml") pytestmark = [pytest.mark.advanced_model, pytest.mark.diffusion] @@ -328,6 +333,15 @@ def _generate_image(omni: Omni, use_system_prompt: str | None) -> Image.Image: return _extract_generated_image(outputs) +# A floor for "this is recognizably the same generation" against each +# case's precomputed reference embedding, not a quality bar -- catches a +# gross migration regression (wrong/blank/garbage image, or the wrong +# system prompt applied) without being flaky on run-to-run sampling noise. +# Matches the threshold test_shared_script_ar_path_reaches_generation uses +# for its own (seed=1234, en_recaption) comparison below. +MIN_SEMANTIC_SIMILARITY = 0.5 + + @pytest.mark.skipif(torch.accelerator.device_count() < 8, reason="Need at least 8 CUDA GPUs for this test.") @pytest.mark.parametrize("system_prompt_name,use_system_prompt,expected_embedding", SYSTEM_PROMPT_CASES) def test_system_prompt_scores( @@ -342,3 +356,109 @@ def test_system_prompt_scores( score = compare_semantic(expected_embedding, generated_image, clip_model, clip_processor) print(f"{system_prompt_name}: CLIP cosine similarity = {score:.6f}") + assert score >= MIN_SEMANTIC_SIMILARITY, ( + f"{system_prompt_name}: CLIP similarity {score:.4f} is below " + f"{MIN_SEMANTIC_SIMILARITY} against the reference (seed=1234) embedding -- " + "the migrated path may have diverged from the reference HunyuanImage-3.0 " + f"output for use_system_prompt={use_system_prompt!r}." + ) + + +def _load_text_to_image_module(): + """Dynamically load the shared example script by file path (``examples/`` + is not an installed package) so this test can call its real + ``_apply_ar_stage_inputs`` instead of re-implementing its logic.""" + repo_root = Path(__file__).resolve().parents[3] + path = repo_root / "examples/offline_inference/text_to_image/text_to_image.py" + spec = importlib.util.spec_from_file_location("_e2e_shared_text_to_image", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +@pytest.mark.skipif(torch.accelerator.device_count() < 8, reason="Need at least 8 CUDA GPUs for this test.") +def test_shared_script_ar_path_reaches_generation( + omni: Omni, + clip_bundle: tuple[CLIPModel, CLIPProcessor], +) -> None: + """Exercise the path this PR actually migrates HunyuanImage-3.0 onto: + + get_model_class_name -> get_ar_input_builder -> _apply_ar_stage_inputs + -> Omni.generate() + + using the *real* model_extras registry lookup (keyed on the engine's + live ``model_class_name``) and the shared ``text_to_image.py``'s own + helper -- not the hand-constructed ``sampling_params.extra_args`` that + ``_generate_image``/``test_system_prompt_scores`` above use, which never + touch the registry or the AR-tokenizer-loading path at all. This is the + one gap those tests leave: a regression in the registry key (the bug + this PR fixes), in ``_apply_ar_stage_inputs``, or in the + ``ar_tokenizer_validator`` wiring would not fail any of them. + + Reuses the ``en_recaption`` reference embedding from + ``test_system_prompt_scores`` so this also asserts the migrated path + produces the *same* output as the hand-constructed one, not just "didn't + crash." + """ + text_to_image = _load_text_to_image_module() + clip_model, clip_processor = clip_bundle + + model_class_name = get_model_class_name(omni) + ar_input_builder = get_ar_input_builder(model_class_name) + assert ar_input_builder is not None, ( + f"get_ar_input_builder({model_class_name!r}) returned None for the running " + "engine -- this is exactly the registry-key regression this PR fixes " + "(spec registered under the wrong model_class_name), and it would " + "silently skip AR-stage input building entirely." + ) + + generator_device = current_omni_platform.device_type or "cuda" + diffusion_params = OmniDiffusionSamplingParams( + seed=1234, + generator=torch.Generator(device=generator_device).manual_seed(1234), + num_outputs_per_prompt=1, + ) + defaults = list(omni.default_sampling_params_list or []) + sampling_params_list = [clone_sampling_params(p) for p in defaults] + for idx, params in enumerate(sampling_params_list): + if isinstance(params, OmniDiffusionSamplingParams): + sampling_params_list[idx] = diffusion_params + + prompt_dict: dict[str, object] = {"prompt": PROMPT, "modalities": ["image"]} + text_to_image._apply_ar_stage_inputs( + ar_input_builder, + model=MODEL_NAME, + prompt_text=PROMPT, + extra_body={"use_system_prompt": "en_recaption"}, + num_images=0, + height=None, + width=None, + prompt_dict=prompt_dict, + sampling_params_list=sampling_params_list, + trust_remote_code=True, + # validate_special_token_ids now tolerates tokens the loaded tokenizer + # doesn't recognize at all (unk_token_id) instead of requiring every + # entry in HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS to be present, so this is + # safe against MODEL_NAME's base checkpoint missing the Instruct-only + # / tokens. Wiring it here exercises the + # real production path (get_ar_tokenizer_validator via the registry), + # not just the isolated fake-tokenizer coverage in test_prompt_utils.py + # / test_shared_script_ar_integration.py. + validate_tokenizer=get_ar_tokenizer_validator(model_class_name), + ) + + outputs = omni.generate(prompt_dict, sampling_params_list=sampling_params_list) + generated_image = _extract_generated_image(outputs) + score = compare_semantic(SYSTEM_EN_RECAPTION, generated_image, clip_model, clip_processor) + print(f"shared-script AR path (en_recaption): CLIP cosine similarity = {score:.6f}") + # test_system_prompt_scores now asserts this same (seed=1234, en_recaption) + # case through the hand-constructed extra_args path against the same + # module-level MIN_SEMANTIC_SIMILARITY floor, so this reuses it too rather + # than picking an independent threshold for the same comparison. + assert score >= MIN_SEMANTIC_SIMILARITY, ( + f"shared-script AR path (en_recaption) CLIP similarity {score:.4f} is below " + f"{MIN_SEMANTIC_SIMILARITY} -- the migrated get_ar_input_builder path may have " + "diverged from the reference HunyuanImage-3.0 output for this prompt/system-prompt." + ) diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py index c45563617..60850248e 100644 --- a/tests/entrypoints/openai_api/test_image_server.py +++ b/tests/entrypoints/openai_api/test_image_server.py @@ -2314,3 +2314,72 @@ def test_image_edits_size_auto_preserves_bridge_size(async_omni_stage_configs_on assert captured_prompt["prompt"].count("") == 2, ( f"N=2 reference images must emit 2 placeholders in AR prompt; got {captured_prompt[KEY].count(IMG)} -- prompt: {captured_prompt[KEY]!r}" ) + + +@pytest.mark.xfail( + strict=True, + reason=( + "Pre-existing gap, not introduced by this change: the /v1/images/edits route in " + "api_server.py never calls resolve_stop_token_ids at all (grep: api_server.py has " + "no stop_token_ids reference), so the AR stage keeps SamplingParams' default empty " + "list. The omitted-bot_task fix landed in serving_chat.py, which is the only " + "production caller; that path is covered by " + "test_serving_chat_multistage_generation.py::" + "test_build_multistage_generation_inputs_omitted_bot_task_matches_prompt_default. " + "Wiring the images/edits route through the same seam needs a tokenizer in that " + "scope and a decision on whether HunyuanImage3 it2i should resolve AR stop tokens " + "unconditionally (today the whole AR block is gated on an explicit bot_task / " + "use_system_prompt / system_prompt), so it is tracked separately." + ), +) +def test_image_edits_omitted_bot_task_stop_tokens_match_prompt_default( + async_omni_stage_configs_only_client, +): + """Regression: an omitted bot_task must resolve identically for the AR + prompt and its stop_token_ids. + + build_prompt/build_prompt_tokens default an omitted bot_task per-task + (e.g. "think" for the base "it2i" task, since it isn't itself a key in + _TASK_PRESETS). resolve_stop_token_ids must land on that same default + to compute the matching stop set. Passing the raw (still-None) outer + bot_task variable to resolve_stop_token_ids -- instead of mirroring + build_kwargs's own omitted-or-not "bot_task" entry -- made it normalize + bot_task=None instead of "think", so it fell through to the full + stop range instead of the think/recaption-only pair. + + An explicit (non-"auto") size is required to reach this: with + need_ratio=True (size="auto"), resolve_stop_token_ids returns the full + ratio range regardless of bot_task, so the two code paths only visibly + disagree once a concrete size selects the narrower think/recaption stop + set. + """ + from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import ( + HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS, + ) + + img = make_test_image_bytes((64, 64)) + response = async_omni_stage_configs_only_client.post( + "/v1/images/edits", + files=[("image", img)], + data={ + "prompt": "make it neon", + "size": "512x512", + }, + ) + 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 + + ar_params = captured[0] + expected_stop_token_ids = [ + HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""], + HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""], + ] + assert ar_params.stop_token_ids == expected_stop_token_ids, ( + f"omitted bot_task with an explicit size must resolve stop_token_ids " + f"for the default 'think' bot_task ({expected_stop_token_ids}); got " + f"{ar_params.stop_token_ids} -- this is the full ratio range, meaning " + "resolve_stop_token_ids disagreed with build_prompt_tokens's default." + ) diff --git a/tests/entrypoints/openai_api/test_serving_chat_multistage_generation.py b/tests/entrypoints/openai_api/test_serving_chat_multistage_generation.py index 0dea97179..639ddf6c6 100644 --- a/tests/entrypoints/openai_api/test_serving_chat_multistage_generation.py +++ b/tests/entrypoints/openai_api/test_serving_chat_multistage_generation.py @@ -538,3 +538,84 @@ def test_build_multistage_generation_inputs_no_stop_token_ids_without_size(servi assert sampling_params_list[0].stop_token_ids == [], ( "Without bot_task, AR stage stop_token_ids must be the default empty list" ) + + +def test_build_multistage_generation_inputs_omitted_bot_task_matches_prompt_default(serving_chat): + """Regression: an omitted bot_task must resolve identically for the AR + prompt and for its stop_token_ids. + + ``build_prompt_tokens``/``build_prompt`` default an omitted bot_task + per-task via the ``_DEFAULT_BOT_TASK`` sentinel (``"think"`` for the base + ``it2i`` task, which isn't itself a key in ``_TASK_PRESETS``). + ``resolve_stop_token_ids`` must normalize from that same omitted-or-not + starting point. Passing the raw -- still ``None`` -- outer ``bot_task`` + variable instead of mirroring ``build_kwargs``'s own omitted-or-not + ``"bot_task"`` entry made it normalize ``bot_task=None`` (plain mode), so + it fell through to the full ```` range instead of the + think/recaption terminator pair. + + An explicit (non-"auto") size is required to observe this: with + ``need_ratio=True`` the ratio range is returned regardless of bot_task, so + the two paths only visibly disagree once a concrete size selects the + narrower stop set. + """ + from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import ( + HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS, + ) + from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat + + class FakeTokenizer: + SPECIAL = { + "<|startoftext|>": 1, + "": 2, + "": 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))) + + def _build(extra_body): + 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(), + ], + ) + _, sampling_params_list = OmniOpenAIServingChat._build_multistage_generation_inputs( + serving_chat, + engine=engine, + prompt="draw a cat", + extra_body=extra_body, + reference_images=[Image.new("RGB", (32, 32), color="red")], + gen_params=OmniDiffusionSamplingParams(height=768, width=1024), + tokenizer=FakeTokenizer(), + ) + return sampling_params_list[0].stop_token_ids + + # bot_task omitted entirely; use_system_prompt is what carries the request + # into the AR-prompt path. + omitted = _build({"use_system_prompt": "en_recaption"}) + # Same request with the default spelled out explicitly. + explicit = _build({"bot_task": "think", "use_system_prompt": "en_recaption"}) + + expected = [ + HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""], + HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""], + ] + assert omitted == expected, ( + f"omitted bot_task with an explicit size must resolve stop_token_ids for the " + f"default 'think' bot_task ({expected}); got {omitted}. A long list here is the " + f"full range, meaning resolve_stop_token_ids normalized bot_task=None " + f"(plain mode) instead of agreeing with build_prompt_tokens's default." + ) + assert omitted == explicit, ( + "omitting bot_task must resolve the same stop_token_ids as passing the default " + f"explicitly; omitted={omitted} explicit={explicit}" + ) diff --git a/tests/model_extras/test_model_extras.py b/tests/model_extras/test_model_extras.py index e55691e10..dcdf13678 100644 --- a/tests/model_extras/test_model_extras.py +++ b/tests/model_extras/test_model_extras.py @@ -3,6 +3,7 @@ from __future__ import annotations +import sys from typing import Any import pytest @@ -592,6 +593,193 @@ def test_declared_extra_args_apply_to_existing_sampling_params() -> None: } +@pytest.mark.core_model +@pytest.mark.cpu +def test_hunyuan_image3_extra_registry_declares_request_and_response_params() -> None: + assert get_extra_body_params("HunyuanImage3Pipeline") == frozenset( + { + "bot_task", + "use_system_prompt", + "system_prompt", + "negative_prompt", + } + ) + assert get_extra_output_params("HunyuanImage3Pipeline") == frozenset() + assert should_init_extra_args_for_non_diffusion_stages("HunyuanImage3Pipeline") is True + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_hunyuan_image3_extra_registry_matches_runtime_model_class_name() -> None: + # od_config.model_class_name resolves to "HunyuanImage3ForCausalMM" at + # runtime; "HunyuanImage3Pipeline" is kept only as a compatibility alias. + assert get_extra_body_params("HunyuanImage3ForCausalMM") == get_extra_body_params("HunyuanImage3Pipeline") + assert get_extra_output_params("HunyuanImage3ForCausalMM") == get_extra_output_params("HunyuanImage3Pipeline") + assert should_init_extra_args_for_non_diffusion_stages( + "HunyuanImage3ForCausalMM" + ) == should_init_extra_args_for_non_diffusion_stages("HunyuanImage3Pipeline") + + +class _FakeHunyuanTokenizer: + """Minimal tokenizer with distinct ids so AR token assertions are unambiguous.""" + + SPECIAL = { + "<|startoftext|>": 1, + "": 2, + "": 3, + "": 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]: + # Offset well above the special-id range to avoid collisions. + return [500 + len(text)] + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_hunyuan_image3_ar_input_builder_registered() -> None: + from vllm_omni.model_extras import get_ar_input_builder + + builder = get_ar_input_builder("HunyuanImage3Pipeline") + assert builder is not None + assert get_ar_input_builder("HunyuanImage3ForCausalMM") is builder + # Models without the hook return None. + assert get_ar_input_builder("WanVACEPipeline") is None + assert get_ar_input_builder(None) is None + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_hunyuan_image3_ar_tokenizer_validator_registered() -> None: + from vllm_omni.model_extras import get_ar_tokenizer_validator + + validator = get_ar_tokenizer_validator("HunyuanImage3Pipeline") + assert validator is not None + assert get_ar_tokenizer_validator("HunyuanImage3ForCausalMM") is validator + # Models without the hook return None. + assert get_ar_tokenizer_validator("WanVACEPipeline") is None + assert get_ar_tokenizer_validator(None) is None + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_hunyuan_image3_build_ar_stage_inputs_t2i_string_fallback() -> None: + from vllm_omni.model_extras.hunyuan_image3 import build_ar_stage_inputs + + ar = build_ar_stage_inputs( + prompt="a red panda", + tokenizer=None, + extra_body={"bot_task": "recaption"}, + num_images=0, + height=1024, + width=1024, + ) + # No tokenizer -> string prompt form; image modality; recaption + explicit + # size stops at only. + assert ar.prompt_token_ids is None + assert ar.prompt and "a red panda" in ar.prompt + assert ar.modalities == ["image"] + assert ar.use_system_prompt + assert len(ar.stop_token_ids) == 1 + # stop_token_ids belong to stage 0 (HunyuanImage3's single AR stage) -- + # declared explicitly, not inferred by the caller from stage type. + assert ar.stage_indices == [0] + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_hunyuan_image3_build_ar_stage_inputs_t2i_tokenizer_parity() -> None: + from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import build_prompt_tokens + from vllm_omni.model_extras.hunyuan_image3 import build_ar_stage_inputs + + ar = build_ar_stage_inputs( + prompt="a red panda", + tokenizer=_FakeHunyuanTokenizer(), + extra_body={"bot_task": "think"}, + num_images=0, + ) + expected = build_prompt_tokens("a red panda", _FakeHunyuanTokenizer(), task="t2i", bot_task="think") + assert ar.prompt is None + assert ar.prompt_token_ids == expected.token_ids + assert ar.modalities == ["image"] + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_hunyuan_image3_build_ar_stage_inputs_it2i_inserts_image_placeholders() -> None: + from vllm_omni.model_extras.hunyuan_image3 import build_ar_stage_inputs + + ar = build_ar_stage_inputs( + prompt="edit it", + tokenizer=_FakeHunyuanTokenizer(), + extra_body={}, + num_images=2, + ) + # Two reference images -> two placeholder ids (2) in the AR prefill. + assert ar.prompt_token_ids is not None + assert ar.prompt_token_ids.count(_FakeHunyuanTokenizer.SPECIAL[""]) == 2 + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_hunyuan_image3_build_ar_stage_inputs_comprehension_text_output() -> None: + from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import ( + HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS, + ) + from vllm_omni.model_extras.hunyuan_image3 import build_ar_stage_inputs + + # t2t (no image, text output) -> text modality, stop on . + ar_t2t = build_ar_stage_inputs("hi", tokenizer=None, extra_body={}, num_images=0, text_output=True) + assert ar_t2t.modalities == ["text"] + assert ar_t2t.stop_token_ids == [HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""]] + + # i2t (image, text output) -> single placeholder in the prefill. + ar_i2t = build_ar_stage_inputs( + "describe", tokenizer=_FakeHunyuanTokenizer(), extra_body={}, num_images=1, text_output=True + ) + assert ar_i2t.modalities == ["text"] + assert ar_i2t.prompt_token_ids.count(_FakeHunyuanTokenizer.SPECIAL[""]) == 1 + + +class _RealIdHunyuanTokenizer: + """Mirrors HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS exactly, so validate_ar_tokenizer + accepts it -- unlike _FakeHunyuanTokenizer above, whose ids are deliberately + wrong to exercise the drift-detection failure path.""" + + unk_token_id = -1 + + def convert_tokens_to_ids(self, tok: str) -> int: + from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import ( + HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS, + ) + + return HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS.get(tok, self.unk_token_id) + + def encode(self, text: str, add_special_tokens: bool = False) -> list[int]: + return [900 + len(text)] + + +def _patch_hunyuan_autotokenizer(monkeypatch: pytest.MonkeyPatch, tokenizer: object) -> list[str]: + """Patch the ``AutoTokenizer`` that ``hunyuan_image3.build_x_to_text_prompt`` + imports locally, recording the models it was asked to load.""" + calls: list[str] = [] + + class _FakeAutoTokenizer: + @staticmethod + def from_pretrained(model: str, trust_remote_code: bool = False) -> object: + calls.append(model) + assert trust_remote_code is True + return tokenizer + + fake_transformers = type(sys)("transformers") + fake_transformers.AutoTokenizer = _FakeAutoTokenizer + monkeypatch.setitem(sys.modules, "transformers", fake_transformers) + return calls + + @pytest.mark.core_model @pytest.mark.cpu def test_mammothmoda2_extra_registry_declares_request_and_response_params() -> None: @@ -619,6 +807,59 @@ def test_mammothmoda2_extra_registry_declares_request_and_response_params() -> N assert wrapper_prompt["additional_information"]["omni_task"] == ["t2i"] +@pytest.mark.core_model +@pytest.mark.cpu +def test_hunyuan_image3_x_to_text_prompt_builder_t2t(monkeypatch: pytest.MonkeyPatch) -> None: + """t2t (no image) routes through build_ar_stage_inputs and stops on .""" + from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import ( + HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS, + ) + from vllm_omni.model_extras.hunyuan_image3 import build_x_to_text_prompt as hunyuan_build_x_to_text_prompt + + calls = _patch_hunyuan_autotokenizer(monkeypatch, _RealIdHunyuanTokenizer()) + prompt_dict, stop_token_ids = hunyuan_build_x_to_text_prompt( + "tencent/HunyuanImage-3.0-Instruct", "What is the capital of France?", has_image=False + ) + + assert calls == ["tencent/HunyuanImage-3.0-Instruct"] + assert prompt_dict["modalities"] == ["text"] + assert prompt_dict["prompt_token_ids"] is not None + assert prompt_dict["prompt_token_ids"].count(HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""]) == 0 + assert stop_token_ids == [HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""]] + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_hunyuan_image3_x_to_text_prompt_builder_i2t(monkeypatch: pytest.MonkeyPatch) -> None: + """i2t (image) inserts exactly one placeholder and still stops on .""" + from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import ( + HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS, + ) + from vllm_omni.model_extras.hunyuan_image3 import build_x_to_text_prompt as hunyuan_build_x_to_text_prompt + + _patch_hunyuan_autotokenizer(monkeypatch, _RealIdHunyuanTokenizer()) + prompt_dict, stop_token_ids = hunyuan_build_x_to_text_prompt( + "tencent/HunyuanImage-3.0-Instruct", "Describe this image.", has_image=True + ) + + assert prompt_dict["modalities"] == ["text"] + assert prompt_dict["prompt_token_ids"].count(HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""]) == 1 + assert stop_token_ids == [HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""]] + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_hunyuan_image3_x_to_text_prompt_builder_validates_tokenizer(monkeypatch: pytest.MonkeyPatch) -> None: + """A tokenizer whose special-token ids have drifted must be rejected here too -- + x_to_text.py (unlike text_to_image.py/image_edit.py) has no registry-driven + ar_tokenizer_validator wiring of its own, so this seam must self-validate.""" + from vllm_omni.model_extras.hunyuan_image3 import build_x_to_text_prompt as hunyuan_build_x_to_text_prompt + + _patch_hunyuan_autotokenizer(monkeypatch, _FakeHunyuanTokenizer()) + with pytest.raises(ValueError, match="no longer match"): + hunyuan_build_x_to_text_prompt("tencent/HunyuanImage-3.0-Instruct", "Hello.", has_image=False) + + @pytest.mark.core_model @pytest.mark.cpu def test_mammothmoda2_text_to_image_prompt_builder() -> None: diff --git a/tests/model_extras/test_shared_script_ar_integration.py b/tests/model_extras/test_shared_script_ar_integration.py new file mode 100644 index 000000000..f781ca5fb --- /dev/null +++ b/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, "": 2, "": 3, "": 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"] diff --git a/vllm_omni/diffusion/models/hunyuan_image3/prompt_utils.py b/vllm_omni/diffusion/models/hunyuan_image3/prompt_utils.py index b64fe2e97..b4ecd7dca 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/prompt_utils.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/prompt_utils.py @@ -60,6 +60,48 @@ HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS: dict[str, int] = { "": 130106, } + +def validate_special_token_ids(tokenizer: Any) -> None: + """Fail fast if a real tokenizer's ids drift from the hardcoded table. + + ``HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS`` exists so stop-token / trigger-tag + resolution doesn't need a tokenizer round-trip, but that only holds as + long as the ids match ``tokenizer.json`` for whatever model repo/revision + is actually loaded. Call this once a real tokenizer is available (the + example scripts do, right after loading it) to catch drift immediately + instead of silently producing a request that runs to completion with the + wrong stop tokens. Not wired into ``build_prompt_tokens`` itself: callers + that intentionally pass a stand-in tokenizer with different ids (e.g. + this module's own unit tests) must not be forced through this check. + + Both ``tencent/HunyuanImage-3.0`` (base) and ``-Instruct`` report the + same ``model_class_name`` (``HunyuanImage3ForCausalMM``), so this hook + fires for either checkpoint -- but the base tokenizer's vocab doesn't + include the Instruct-only ````-```` tokens + (added for it2i's second aspect-ratio range). A token the tokenizer + doesn't recognize at all (``convert_tokens_to_ids`` returning + ``unk_token_id``, or ``None`` for tokenizers without an unk concept) is + "not part of this checkpoint's vocab," not "drifted"; only an id that + resolves to something *else* indicates the hardcoded table is stale. + """ + unk_id = getattr(tokenizer, "unk_token_id", None) + mismatches: dict[str, tuple[int, int]] = {} + for tok, expected_id in HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS.items(): + actual_id = tokenizer.convert_tokens_to_ids(tok) + if actual_id is None or (unk_id is not None and actual_id == unk_id): + continue # not in this checkpoint's vocab -- nothing to compare + if actual_id != expected_id: + mismatches[tok] = (expected_id, actual_id) + if mismatches: + details = ", ".join(f"{tok!r}: expected {exp}, got {act}" for tok, (exp, act) in mismatches.items()) + raise ValueError( + "HunyuanImage-3.0 tokenizer special-token ids no longer match " + f"HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS ({details}). The loaded " + "tokenizer/model revision has likely changed; update the hardcoded " + "table in prompt_utils.py to match." + ) + + # 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 @@ -317,12 +359,106 @@ def build_prompt_tokens( ) +@dataclass +class ARPromptInputs: + """Fully-resolved AR-stage inputs for one HunyuanImage3 request. + + Consolidates the three things every caller (offline example, OpenAI + server) must compute before submitting an AR request: the AR prefill + (as segment-tokenized ``prompt_token_ids`` when a tokenizer is + available, else a formatted ``prompt`` string), the resolved DiT + system-prompt type, and the AR ``stop_token_ids``. + """ + + prompt: str | None + prompt_token_ids: list[int] | None + system_prompt_type: str + stop_token_ids: list[int] + + +def build_ar_prompt_inputs( + user_prompt: str, + task: str = "t2i", + bot_task: str | None | _DefaultBotTask = _DEFAULT_BOT_TASK, + sys_type: str | None = None, + custom_system_prompt: str | None = None, + num_images: int = 1, + tokenizer: Any | None = None, + image_size: str | None = None, +) -> ARPromptInputs: + """Build the complete AR-stage inputs for a HunyuanImage3 request. + + Wraps :func:`build_prompt_tokens` (byte-for-byte HF parity when + ``tokenizer`` is given), the :func:`build_prompt` string fallback (when + no tokenizer is available), and :func:`resolve_stop_token_ids` behind + one call, so a caller doesn't have to re-derive the AR prefill/stop-token + logic -- and, critically, doesn't have to independently resolve + ``bot_task`` for the prompt-building call vs. the stop-token call, which + can silently diverge if a caller passes the same nominal ``bot_task`` + value to both but resolves "omitted" vs. "explicit None" differently for + each (this seam always forwards the identical resolved value to both). + + Used by all four offline shared task examples: ``text_to_image.py`` / + ``image_edit.py`` (via ``model_extras.hunyuan_image3.build_ar_stage_inputs``) + and ``x_to_text.py`` (via ``build_x_to_text_prompt``, which also calls + ``build_ar_stage_inputs``). The OpenAI server's ``serving_chat.py`` does + not currently call this seam -- it still independently calls + ``build_prompt``/``build_prompt_tokens``/``resolve_stop_token_ids``, and + can hit exactly the divergence described above when ``bot_task`` is + omitted from ``extra_body`` (pre-existing, unrelated to this module). + ``image_size`` follows the same convention as :func:`resolve_stop_token_ids` + (``None``/``"auto"`` lets the AR predict the aspect ratio; an explicit + ``"{w}x{h}"`` makes it stop at the terminator). + """ + if tokenizer is not None: + result = build_prompt_tokens( + user_prompt, + tokenizer, + task=task, + bot_task=bot_task, + sys_type=sys_type, + custom_system_prompt=custom_system_prompt, + num_images=num_images, + ) + prompt_str: str | None = None + prompt_token_ids: list[int] | None = result.token_ids + system_prompt_type = result.system_prompt_type + else: + prompt_str = build_prompt( + user_prompt, + task=task, + bot_task=bot_task, + sys_type=sys_type, + custom_system_prompt=custom_system_prompt, + num_images=num_images, + ) + prompt_token_ids = None + _, resolved_bot_task = _normalize_task_and_bot_task(task, bot_task) + system_prompt_type = sys_type or resolve_sys_type(resolved_bot_task) + + stop_token_ids = resolve_stop_token_ids( + task=task, + bot_task=bot_task, + tokenizer=tokenizer, + image_size=image_size, + ) + return ARPromptInputs( + prompt=prompt_str, + prompt_token_ids=prompt_token_ids, + system_prompt_type=system_prompt_type, + stop_token_ids=stop_token_ids, + ) + + __all__ = [ "HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS", + "validate_special_token_ids", "MAX_IMAGES_PER_REQUEST", "_TASK_PRESETS", + "ARPromptInputs", "available_bot_tasks", "available_tasks", + "build_ar_prompt_inputs", "build_prompt", "build_prompt_tokens", "resolve_stop_token_ids", diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index f5cb24efd..20bddaba6 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -3080,11 +3080,20 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): ar_image_size: str | None = None if height is not None and width is not None: ar_image_size = f"{width}x{height}" + # Reuse build_kwargs's own (possibly-omitted) "bot_task" entry rather + # than the raw outer `bot_task` variable: build_prompt_tokens/build_prompt + # above default an omitted bot_task per-task (e.g. "think" for t2i), and + # resolve_stop_token_ids must normalize from that same omitted-or-not + # starting point to agree on which stop tokens apply -- passing the raw + # `bot_task` (still None when omitted) made the two calls disagree. + ar_stop_kwargs: dict[str, Any] = {} + if "bot_task" in build_kwargs: + ar_stop_kwargs["bot_task"] = build_kwargs["bot_task"] ar_stop_token_ids = resolve_stop_token_ids( task=ar_task, - bot_task=bot_task, tokenizer=tokenizer, image_size=ar_image_size, + **ar_stop_kwargs, ) engine_prompt: OmniTextPrompt = {"prompt": prompt} diff --git a/vllm_omni/model_extras/__init__.py b/vllm_omni/model_extras/__init__.py index 143a9b224..2e1146a82 100644 --- a/vllm_omni/model_extras/__init__.py +++ b/vllm_omni/model_extras/__init__.py @@ -6,6 +6,8 @@ from vllm_omni.model_extras.registry import ( build_image_to_video_prompt, build_text_to_image_prompt, build_x_to_text_prompt, + get_ar_input_builder, + get_ar_tokenizer_validator, get_extra_body_params, get_extra_output_params, get_model_class_name, @@ -22,6 +24,8 @@ __all__ = [ "build_image_to_video_prompt", "build_text_to_image_prompt", "build_x_to_text_prompt", + "get_ar_input_builder", + "get_ar_tokenizer_validator", "get_extra_body_params", "get_extra_output_params", "get_model_class_name", diff --git a/vllm_omni/model_extras/hunyuan_image3.py b/vllm_omni/model_extras/hunyuan_image3.py index 8a53dfd54..46d48b513 100644 --- a/vllm_omni/model_extras/hunyuan_image3.py +++ b/vllm_omni/model_extras/hunyuan_image3.py @@ -3,6 +3,8 @@ from __future__ import annotations +from collections.abc import Mapping +from dataclasses import dataclass, field from typing import Any @@ -11,27 +13,162 @@ def build_x_to_text_prompt( prompt: str, has_image: bool, ) -> tuple[dict[str, Any], list[int] | None]: - """Build HunyuanImage-3 prompt tokens and stopping rules for T2T/I2T.""" + """Build HunyuanImage-3 prompt tokens and stopping rules for T2T/I2T. + + Routes through :func:`build_ar_stage_inputs`, the same declarative seam + the shared text_to_image / image_edit examples use, instead of + re-deriving the AR prefill and stop-token logic here. Unlike those + examples -- generic scripts that reach ``validate_ar_tokenizer`` only + indirectly via ``get_ar_tokenizer_validator`` -- this function is already + HunyuanImage3-specific, so it calls the validator directly rather than + requiring the (model-agnostic) ``x_to_text.py`` caller to know about it. + """ from transformers import AutoTokenizer - from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import ( - build_prompt_tokens, - resolve_stop_token_ids, - resolve_sys_type, - ) - - task = "i2t" if has_image else "t2t" tokenizer = AutoTokenizer.from_pretrained(model, trust_remote_code=True) - build_kwargs: dict[str, Any] = {"task": task, "bot_task": None} - if has_image: - build_kwargs["num_images"] = 1 - built = build_prompt_tokens(prompt, tokenizer, **build_kwargs) + validate_ar_tokenizer(tokenizer) + ar_inputs = build_ar_stage_inputs( + prompt=prompt, + tokenizer=tokenizer, + extra_body=None, + num_images=1 if has_image else 0, + text_output=True, + ) return ( { - "prompt": prompt, - "prompt_token_ids": built.token_ids, - "modalities": ["text"], - "use_system_prompt": resolve_sys_type(None), + "prompt": ar_inputs.prompt, + "prompt_token_ids": ar_inputs.prompt_token_ids, + "modalities": ar_inputs.modalities, + "use_system_prompt": ar_inputs.use_system_prompt, }, - resolve_stop_token_ids(task=task, bot_task=None, tokenizer=tokenizer, image_size="auto"), + ar_inputs.stop_token_ids, ) + + +@dataclass +class ARStageInputs: + """AR-stage inputs the shared example applies to a HunyuanImage3 request. + + ``prompt_token_ids`` (byte-for-byte HF parity) is preferred; ``prompt`` + is the string fallback used when no tokenizer is available. ``modalities`` + picks the AR output type (``["image"]`` for generation, ``["text"]`` for + comprehension). ``use_system_prompt`` feeds the DiT system-prompt prefix + and ``stop_token_ids`` terminate the AR decode. + + ``stage_indices`` names which stage(s) in the request's + ``sampling_params_list`` these ``stop_token_ids`` belong to -- the shared + scripts apply them by explicit index, not by scanning for "whichever + stage isn't a diffusion stage." HunyuanImage-3.0's topology always puts + its single AR stage at index 0 (see + ``vllm_omni.model_executor.models.hunyuan_image3.pipeline``); a future + model with multiple AR/understanding stages would declare its own + indices here instead of relying on type-based inference. + """ + + prompt: str | None + prompt_token_ids: list[int] | None + modalities: list[str] + use_system_prompt: str + stop_token_ids: list[int] + stage_indices: list[int] = field(default_factory=lambda: [0]) + + +def build_ar_stage_inputs( + prompt: str, + tokenizer: Any | None, + extra_body: Mapping[str, Any] | None, + *, + num_images: int = 0, + height: int | None = None, + width: int | None = None, + text_output: bool = False, +) -> ARStageInputs: + """Resolve HunyuanImage-3.0 AR-stage inputs declaratively. + + Invoked generically by the shared task examples (``text_to_image`` / + ``image_to_image`` / ``x_to_text``) whenever the model declares an + ``ar_input_builder``. All model-specific knobs (``bot_task`` / + ``use_system_prompt`` / ``system_prompt``) arrive via ``extra_body``, so + the example itself stays model-agnostic. The heavy lifting is delegated to + :func:`build_ar_prompt_inputs`. The OpenAI server's ``serving_chat.py`` + does not go through this seam -- it independently calls the same + underlying ``build_prompt``/``build_prompt_tokens``/``resolve_stop_token_ids`` + primitives (see that function's docstring for the resulting divergence risk). + """ + from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import ( + build_ar_prompt_inputs, + ) + + extra_body = extra_body or {} + has_image = num_images > 0 + if text_output: + task = "i2t" if has_image else "t2t" + else: + task = "it2i" if has_image else "t2i" + + # ar_image_size: None -> AR predicts aspect ratio; explicit "{w}x{h}" -> + # AR stops at the terminator. Comprehension (text) tasks always auto-size. + ar_image_size: str | None = None + if not text_output and height is not None and width is not None: + ar_image_size = f"{width}x{height}" + + kwargs: dict[str, Any] = { + "task": task, + "sys_type": extra_body.get("use_system_prompt"), + "custom_system_prompt": extra_body.get("system_prompt"), + "num_images": max(num_images, 1) if has_image else 1, + "tokenizer": tokenizer, + "image_size": ar_image_size, + } + # Only forward bot_task when the caller set it: present (incl. explicit + # ``None`` -> plain mode) overrides; absent lets each task pick its default + # trigger via build_ar_prompt_inputs' own sentinel default. + if "bot_task" in extra_body: + kwargs["bot_task"] = extra_body["bot_task"] + + resolved = build_ar_prompt_inputs(prompt, **kwargs) + return ARStageInputs( + prompt=resolved.prompt, + prompt_token_ids=resolved.prompt_token_ids, + modalities=["text"] if text_output else ["image"], + use_system_prompt=resolved.system_prompt_type, + stop_token_ids=resolved.stop_token_ids, + # HunyuanImage-3.0's AR stage is always stage 0 (see pipeline.py's + # HUNYUAN_IMAGE3_PIPELINE / _AR_PIPELINE topologies). + stage_indices=[0], + ) + + +def validate_ar_tokenizer(tokenizer: Any) -> None: + """Registry-declared ``ar_tokenizer_validator`` hook for HunyuanImage3. + + Called by the shared task examples right after they load a *real* + tokenizer for the AR stage, so a model/tokenizer revision bump that + silently shifts special-token ids fails loudly instead of producing a + request that completes with the wrong stop tokens. + """ + from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import ( + validate_special_token_ids, + ) + + validate_special_token_ids(tokenizer) + + +# Per-request, model-specific knobs (non-standard sampling-param fields). +# Standard knobs (guidance_scale / num_inference_steps / seed / height / width) +# stay on OmniDiffusionSamplingParams; engine-level knobs +# (diffusion_kv_cache_* / vae_use_tiling) stay as Omni() init args. +HUNYUAN_IMAGE3_EXTRA_BODY_PARAMS = frozenset( + { + "bot_task", + "use_system_prompt", + "system_prompt", + "negative_prompt", + } +) +# Text outputs for i2t/t2t are surfaced through the standard AR text path; +# no diffusion custom-output params are declared. +HUNYUAN_IMAGE3_EXTRA_OUTPUT_PARAMS = frozenset() +# HunyuanImage3 runs an AR stage before DiT; its extra_args must be initialised +# so bot_task / use_system_prompt reach the AR input path. +HUNYUAN_IMAGE3_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES = True diff --git a/vllm_omni/model_extras/registry.py b/vllm_omni/model_extras/registry.py index 4d2ad010a..f21691eab 100644 --- a/vllm_omni/model_extras/registry.py +++ b/vllm_omni/model_extras/registry.py @@ -28,7 +28,18 @@ from vllm_omni.model_extras.helios import ( HELIOS_EXTRA_BODY_PARAMS, HELIOS_EXTRA_OUTPUT_PARAMS, ) +from vllm_omni.model_extras.hunyuan_image3 import ( + HUNYUAN_IMAGE3_EXTRA_BODY_PARAMS, + HUNYUAN_IMAGE3_EXTRA_OUTPUT_PARAMS, + HUNYUAN_IMAGE3_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES, +) +from vllm_omni.model_extras.hunyuan_image3 import ( + build_ar_stage_inputs as build_hunyuan_image3_ar_stage_inputs, +) from vllm_omni.model_extras.hunyuan_image3 import build_x_to_text_prompt as build_hunyuan_x_to_text_prompt +from vllm_omni.model_extras.hunyuan_image3 import ( + validate_ar_tokenizer as validate_hunyuan_image3_ar_tokenizer, +) from vllm_omni.model_extras.lingbot_video import LINGBOT_VIDEO_EXTRA_BODY_PARAMS from vllm_omni.model_extras.ltx2 import ( LTX_EXTRA_BODY_PARAMS, @@ -233,6 +244,18 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = { "LTX2DistilledTwoStagePipeline", ) }, + # HunyuanImage3 keys on the model architecture name that ``get_model_class_name`` + # reports at runtime for its AR+DiT / AR-only / DiT-only deploys + # (``od_config.model_class_name`` == the diffusion registry arch key). The + # "HunyuanImage3Pipeline" alias is kept for any path that surfaces the + # pipeline class name instead. + "HunyuanImage3ForCausalMM": { + "extra_body_params": HUNYUAN_IMAGE3_EXTRA_BODY_PARAMS, + "extra_output_params": HUNYUAN_IMAGE3_EXTRA_OUTPUT_PARAMS, + "init_extra_args_for_non_diffusion_stages": HUNYUAN_IMAGE3_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES, + "ar_input_builder": build_hunyuan_image3_ar_stage_inputs, + "ar_tokenizer_validator": validate_hunyuan_image3_ar_tokenizer, + }, "SanaVideoPipeline": { "extra_body_params": SANA_VIDEO_EXTRA_BODY_PARAMS, }, @@ -263,6 +286,10 @@ for model_class_name in ("LTX2Pipeline", "LTX2TwoStagePipeline"): _EXTRA_SPECS[model_class_name]["transformer_config_subfolder_resolver"] = ltx_transformer_config_subfolder +# Alias: some code paths surface HunyuanImage3's pipeline class name rather than +# the architecture name; point both at the same spec. +_EXTRA_SPECS["HunyuanImage3Pipeline"] = _EXTRA_SPECS["HunyuanImage3ForCausalMM"] + # Multi-stage discovery reports the top-level wrapper rather than its DiT # submodule, so both names must resolve to the same request builders. _EXTRA_SPECS["MammothModa2ForConditionalGeneration"] = _EXTRA_SPECS["MammothModa2DiTPipeline"] @@ -358,6 +385,32 @@ def should_init_extra_args_for_non_diffusion_stages(model_class_name: str | None return bool(spec and spec.get("init_extra_args_for_non_diffusion_stages", False)) +def get_ar_input_builder(model_class_name: str | None) -> Callable[..., Any] | None: + """Return a model's AR-stage input builder, or ``None`` if undeclared. + + Models with a text/AR stage that needs template-formatted prompt tokens + and AR stop tokens (e.g. HunyuanImage3) declare an ``ar_input_builder``. + The shared task examples call it generically when present, so the example + scripts stay model-agnostic; models without one are unaffected. + """ + spec = _get_spec(model_class_name) + return spec.get("ar_input_builder") if spec is not None else None + + +def get_ar_tokenizer_validator(model_class_name: str | None) -> Callable[[Any], None] | None: + """Return a model's AR-tokenizer validator, or ``None`` if undeclared. + + Models whose AR prompt/stop-token logic depends on hardcoded special + token ids (e.g. HunyuanImage3) declare an ``ar_tokenizer_validator`` to + check those ids against whatever tokenizer actually loads at runtime. + The shared task examples call it generically, right after loading a real + tokenizer, so model/tokenizer revision drift fails loudly instead of + silently producing a request with the wrong stop tokens. + """ + spec = _get_spec(model_class_name) + return spec.get("ar_tokenizer_validator") if spec is not None else None + + def build_text_to_image_prompt( model_class_name: str | None, prompt: dict[str, Any],