anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: suyanli220 <suyanli220@gmail.com>
Signed-off-by: suyan.li <suyan.li@bytedance.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: suyan.li <suyan.li@bytedance.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
SuyanLiClaude Opus 4.8suyan.liHongsheng Liu
 authored and GitHub committed 2026-09-11 16:47:25 +08:00
1 parent 96360ccb00
commit a7ab91fb85
17 files changed
+1627 -544
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

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