From ccdb2b51aa5d87ff80022b67790fd9beb36f46af Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=B1=AA=E5=BF=97=E9=B9=8F?= Date: Sun, 26 Jul 2026 16:33:17 +0800 Subject: [PATCH] [Refactor]: Intrduce x_to_text.py in examples (#5384) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: princepride Signed-off-by: 汪志鹏 --- docs/.nav.yml | 1 - .../offline_inference/x_to_text/README.md | 93 ++++++++++++++++ .../offline_inference/x_to_text/x_to_text.py | 102 ++++++++++++++++++ recipes/Bagel/BAGEL-7B-MoT.md | 23 +++- recipes/MammothModa2/MammothModa2-Preview.md | 29 ++++- recipes/Tencent/HunyuanImage-3.0-Instruct.md | 33 +++++- tests/model_extras/test_model_extras.py | 54 ++++++++++ vllm_omni/model_extras/__init__.py | 4 + vllm_omni/model_extras/bagel.py | 17 +++ vllm_omni/model_extras/hunyuan_image3.py | 37 +++++++ .../model_extras/mammothmodal2_preview.py | 24 +++++ vllm_omni/model_extras/registry.py | 49 +++++++++ 12 files changed, 457 insertions(+), 9 deletions(-) create mode 100644 examples/offline_inference/x_to_text/README.md create mode 100644 examples/offline_inference/x_to_text/x_to_text.py create mode 100644 vllm_omni/model_extras/hunyuan_image3.py diff --git a/docs/.nav.yml b/docs/.nav.yml index e7157b573..f3244cfe2 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -22,7 +22,6 @@ nav: - Image-To-Image: user_guide/examples/offline_inference/image_to_image.md - Image-To-Video: user_guide/examples/offline_inference/image_to_video.md - InternVLA-A1: user_guide/examples/offline_inference/internvla_a1.md - - MammothModa2-Preview: user_guide/examples/offline_inference/mammothmodal2_preview.md - MiMo-Audio: user_guide/examples/offline_inference/mimo_audio.md - Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md - Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md diff --git a/examples/offline_inference/x_to_text/README.md b/examples/offline_inference/x_to_text/README.md new file mode 100644 index 000000000..fb3f7ae14 --- /dev/null +++ b/examples/offline_inference/x_to_text/README.md @@ -0,0 +1,93 @@ +# X-To-Text + +Generate text from text or image inputs with vLLM-Omni's shared offline entrypoint. + +- `x_to_text.py`: command-line script for text-to-text (T2T) and image-to-text (I2T) inference. + +## Supported Models + +| Model | T2T | I2T | Default deploy config | +| --- | --- | --- | --- | +| `ByteDance-Seed/BAGEL-7B-MoT` | Yes | Yes | `vllm_omni/deploy/bagel.yaml` | +| `tencent/HunyuanImage-3.0-Instruct` | Yes | Yes | `vllm_omni/deploy/hunyuan_image3_ar.yaml` | +| `bytedance-research/MammothModa2-Preview` | Yes | Yes | `vllm_omni/deploy/mammoth_moda2_ar.yaml` | + +The script recognizes these model families from `config.json`, applies the +model-specific prompt format, and selects an AR-only deploy for HunyuanImage-3 +and MammothModa2. BAGEL uses its registered default deploy. + +## Text-To-Text + +### BAGEL + +```bash +python examples/offline_inference/x_to_text/x_to_text.py \ + --model ByteDance-Seed/BAGEL-7B-MoT \ + --prompt "Where is the capital of France?" +``` + +### HunyuanImage-3 + +```bash +python examples/offline_inference/x_to_text/x_to_text.py \ + --model tencent/HunyuanImage-3.0-Instruct \ + --prompt "Explain multimodal inference in three concise sentences." +``` + +### MammothModa2 + +```bash +python examples/offline_inference/x_to_text/x_to_text.py \ + --model bytedance-research/MammothModa2-Preview \ + --prompt "Explain multimodal inference in three concise sentences." +``` + +## Image-To-Text + +Pass one image with `--image`. The default question can be replaced with any +instruction supported by the checkpoint. + +```bash +python examples/offline_inference/x_to_text/x_to_text.py \ + --model ByteDance-Seed/BAGEL-7B-MoT \ + --image image.png \ + --prompt "Describe this image in detail." +``` + +Use the same command with either of the other supported checkpoints: + +```bash +python examples/offline_inference/x_to_text/x_to_text.py \ + --model tencent/HunyuanImage-3.0-Instruct \ + --image image.png \ + --prompt "Describe this image in detail." +``` + +```bash +python examples/offline_inference/x_to_text/x_to_text.py \ + --model bytedance-research/MammothModa2-Preview \ + --image image.png \ + --prompt "Describe this image in detail." +``` + +## Key Arguments + +| Argument | Default | Description | +| --- | --- | --- | +| `--model` | Required | Hugging Face model ID or local checkpoint path | +| `--prompt` | Required | User question or instruction | +| `--image` | None | Input image; enables I2T when supplied | +| `--output` | None | Optional text output file | +| `--deploy-config` | Model default | Override the deploy YAML | +| `--max-tokens` | `512` | Maximum generated tokens | +| `--temperature` | `0.0` | Sampling temperature | +| `--top-p` | `1.0` | Nucleus sampling probability | +| `--seed` | `42` | Sampling seed | +| `--trust-remote-code` | Off | Allow checkpoint-provided Python code | +| `--enforce-eager` | Off | Disable graph compilation | + +## Hardware Notes + +The committed HunyuanImage-3 AR-only default uses four GPUs with TP=4. To use a +different topology, pass a custom AR-only YAML through `--deploy-config`. +BAGEL and MammothModa2 use their repository deploy defaults unless overridden. diff --git a/examples/offline_inference/x_to_text/x_to_text.py b/examples/offline_inference/x_to_text/x_to_text.py new file mode 100644 index 000000000..9937a57ff --- /dev/null +++ b/examples/offline_inference/x_to_text/x_to_text.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared offline example for any supported input modality to text.""" + +from __future__ import annotations + +import argparse +from pathlib import Path +from typing import Any + +from PIL import Image + +from vllm_omni import Omni +from vllm_omni.model_extras import build_x_to_text_prompt, get_x_to_text_model_family + +_REPO_ROOT = Path(__file__).resolve().parents[3] +_HUNYUAN_AR_DEPLOY_CONFIG = _REPO_ROOT / "vllm_omni" / "deploy" / "hunyuan_image3_ar.yaml" +_MAMMOTH_AR_DEPLOY_CONFIG = _REPO_ROOT / "vllm_omni" / "deploy" / "mammoth_moda2_ar.yaml" + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Run shared offline x-to-text inference (currently T2T and I2T).") + parser.add_argument("--model", required=True, help="Model name or local path.") + parser.add_argument("--prompt", required=True, help="Question or text prompt.") + parser.add_argument("--image", help="Optional input image. Supplying it selects image-to-text.") + parser.add_argument("--output", help="Optional path to write the generated text.") + parser.add_argument("--deploy-config", help="Optional deploy YAML override.") + parser.add_argument("--max-tokens", type=int, default=512) + parser.add_argument("--temperature", type=float, default=0.0) + parser.add_argument("--top-p", type=float, default=1.0) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--trust-remote-code", action="store_true") + parser.add_argument("--enforce-eager", action="store_true", default=None) + parser.add_argument("--log-stats", action="store_true") + return parser.parse_args() + + +def _extract_text(outputs: list[Any]) -> str: + chunks: list[str] = [] + for output in outputs: + request_output = getattr(output, "request_output", output) + for completion in getattr(request_output, "outputs", None) or []: + chunks.append(getattr(completion, "text", "") or "") + return "".join(chunks).strip() + + +def main() -> None: + args = parse_args() + family = get_x_to_text_model_family(args.model) + image = Image.open(args.image).convert("RGB") if args.image else None + + prompt_dict, stop_token_ids = build_x_to_text_prompt( + model_family=family, + model=args.model, + prompt=args.prompt, + has_image=image is not None, + ) + + if image is not None: + prompt_dict["multi_modal_data"] = {"image": image} + + omni_kwargs: dict[str, Any] = { + "model": args.model, + "mode": "image-to-text" if image is not None else "text-to-text", + "trust_remote_code": args.trust_remote_code or family in {"hunyuan_image3", "mammoth_moda2"}, + "log_stats": args.log_stats, + } + if args.enforce_eager is not None: + omni_kwargs["enforce_eager"] = args.enforce_eager + if args.deploy_config: + omni_kwargs["deploy_config"] = args.deploy_config + elif family == "hunyuan_image3": + omni_kwargs["deploy_config"] = str(_HUNYUAN_AR_DEPLOY_CONFIG) + elif family == "mammoth_moda2": + omni_kwargs["deploy_config"] = str(_MAMMOTH_AR_DEPLOY_CONFIG) + + omni = Omni(**omni_kwargs) + try: + sampling_params = list(omni.default_sampling_params_list or []) + for params in sampling_params: + if hasattr(params, "max_tokens"): + params.max_tokens = args.max_tokens + if hasattr(params, "temperature"): + params.temperature = args.temperature + if hasattr(params, "top_p"): + params.top_p = args.top_p + if hasattr(params, "seed"): + params.seed = args.seed + if stop_token_ids is not None and hasattr(params, "stop_token_ids"): + params.stop_token_ids = stop_token_ids + outputs = list(omni.generate([prompt_dict], sampling_params_list=sampling_params or None)) + finally: + omni.close() + + text = _extract_text(outputs) + print(text) + if args.output: + Path(args.output).write_text(text + "\n", encoding="utf-8") + + +if __name__ == "__main__": + main() diff --git a/recipes/Bagel/BAGEL-7B-MoT.md b/recipes/Bagel/BAGEL-7B-MoT.md index a7b0524e4..4c2d75aca 100644 --- a/recipes/Bagel/BAGEL-7B-MoT.md +++ b/recipes/Bagel/BAGEL-7B-MoT.md @@ -6,7 +6,7 @@ - Vendor: ByteDance Seed - Model: `ByteDance-Seed/BAGEL-7B-MoT` -- Task: Text-to-image and image-to-image generation +- Task: Text-to-image, image-to-image, text-to-text, and image-to-text - Mode: Offline inference and OpenAI-compatible online serving - Maintainer: Community @@ -26,6 +26,8 @@ parameters through the pipeline-declared `extra_args` contract. [`examples/offline_inference/text_to_image/text_to_image.py`](../../examples/offline_inference/text_to_image/text_to_image.py) - Related offline image-to-image example: [`examples/offline_inference/image_to_image/image_edit.py`](../../examples/offline_inference/image_to_image/image_edit.py) +- Related offline understanding example: + [`examples/offline_inference/x_to_text/x_to_text.py`](../../examples/offline_inference/x_to_text/x_to_text.py) - Related online example: [`examples/online_serving/text_to_image/openai_chat_client.py`](../../examples/online_serving/text_to_image/openai_chat_client.py) - Related online image-to-image example: @@ -118,6 +120,25 @@ The `--extra-args` JSON forwards BAGEL-specific parameters (e.g. `cfg_text_scale `cfg_img_scale`, `cfg_interval`, `cfg_renorm_type`) into `OmniDiffusionSamplingParams.extra_args` via the model-extras registry. +Run text-to-text with the shared understanding example. BAGEL's default +`bagel.yaml` deploy config is discovered from the checkpoint, so no model-specific +example or explicit deploy config is required: + +```bash +python examples/offline_inference/x_to_text/x_to_text.py \ + --model ByteDance-Seed/BAGEL-7B-MoT \ + --prompt "Where is the capital of France?" +``` + +Add `--image` to run image-to-text with the same example: + +```bash +python examples/offline_inference/x_to_text/x_to_text.py \ + --model ByteDance-Seed/BAGEL-7B-MoT \ + --image /path/to/input.jpg \ + --prompt "Please describe this image in detail." +``` + #### Online Commands Start the OpenAI-compatible server: diff --git a/recipes/MammothModa2/MammothModa2-Preview.md b/recipes/MammothModa2/MammothModa2-Preview.md index 4e4846fba..c94560254 100644 --- a/recipes/MammothModa2/MammothModa2-Preview.md +++ b/recipes/MammothModa2/MammothModa2-Preview.md @@ -1,6 +1,6 @@ # MammothModa2-Preview -> MammothModa2-Preview text-to-image generation through the shared offline image example +> MammothModa2-Preview generation and understanding through shared offline examples ## Summary @@ -13,8 +13,8 @@ ## When to use this recipe -Use this recipe to run MammothModa2-Preview text-to-image through the shared -offline image example (`text_to_image.py`) instead of a model-specific script. +Use this recipe to run MammothModa2-Preview through shared task-oriented +offline examples. Text-to-image uses the shared image example (`text_to_image.py`) instead of a model-specific script. The generic example formats the AR prompt, drives the AR → DiT stage pipeline, and forwards MammothModa2-specific generation parameters through the pipeline-declared `extra_body` contract. @@ -29,8 +29,10 @@ flags. Image size uses the standard `--height` / `--width` flags. - Upstream model: [`bytedance-research/MammothModa2-Preview`](https://huggingface.co/bytedance-research/MammothModa2-Preview) -- Related offline example: +- Related text-to-image example: [`examples/offline_inference/text_to_image/text_to_image.py`](../../examples/offline_inference/text_to_image/text_to_image.py) +- Related T2T/I2T example: + [`examples/offline_inference/x_to_text/x_to_text.py`](../../examples/offline_inference/x_to_text/x_to_text.py) - Declared parameters: [`vllm_omni/model_extras/mammothmodal2_preview.py`](../../vllm_omni/model_extras/mammothmodal2_preview.py) - Deploy config: @@ -100,6 +102,25 @@ so unknown keys for MammothModa2 are silently dropped: `--height` and `--width` must be multiples of 16. +Run text-to-text through the shared understanding example. It recognizes the +MammothModa2 checkpoint and automatically selects `mammoth_moda2_ar.yaml`: + +```bash +python examples/offline_inference/x_to_text/x_to_text.py \ + --model ./MammothModa2-Preview \ + --prompt "Explain multimodal generation in three sentences." +``` + +Add an image for image-to-text or image summarization. The shared example +uses MammothModa2's chat and vision-token template: + +```bash +python examples/offline_inference/x_to_text/x_to_text.py \ + --model ./MammothModa2-Preview \ + --image /path/to/input.jpg \ + --prompt "Please summarize the content of this image." +``` + #### Verification The example writes the generated image to the `--output` path. Confirm the file diff --git a/recipes/Tencent/HunyuanImage-3.0-Instruct.md b/recipes/Tencent/HunyuanImage-3.0-Instruct.md index 3e6adf130..f169e6a6b 100644 --- a/recipes/Tencent/HunyuanImage-3.0-Instruct.md +++ b/recipes/Tencent/HunyuanImage-3.0-Instruct.md @@ -8,8 +8,8 @@ - Vendor: Tencent Hunyuan - Model: `tencent/HunyuanImage-3.0-Instruct` -- Task: Text-to-image generation -- Mode: Online serving and performance benchmarking, DiT stage only +- Task: Text-to-image generation; text-to-text and image-to-text understanding +- Mode: Offline understanding, online serving, and performance benchmarking - Maintainer: Community ## When to use this recipe @@ -31,8 +31,10 @@ FP8/NVFP4 configuration: ## References - Model: -- Offline example: +- Existing model-specific offline example: [`examples/offline_inference/hunyuan_image3`](../../examples/offline_inference/hunyuan_image3) +- 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: [#2495](https://github.com/vllm-project/vllm-omni/pull/2495) for DiT performance CI, [#3055](https://github.com/vllm-project/vllm-omni/pull/3055) for GEBench accuracy CI, @@ -42,6 +44,31 @@ FP8/NVFP4 configuration: ## GPU +### Shared T2T/I2T offline example + +For text-to-text, the shared x-to-text example automatically selects the +AR-only Hunyuan deploy config and applies the checkpoint's prompt tokens and +stop-token rules: + +```bash +python examples/offline_inference/x_to_text/x_to_text.py \ + --model tencent/HunyuanImage-3.0-Instruct \ + --prompt "Explain why the sky appears blue." +``` + +Add `--image` for image-to-text: + +```bash +python examples/offline_inference/x_to_text/x_to_text.py \ + --model tencent/HunyuanImage-3.0-Instruct \ + --image /path/to/input.jpg \ + --prompt "Describe the content of this image." +``` + +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. + ### 4x H100/H800 80GB #### Environment diff --git a/tests/model_extras/test_model_extras.py b/tests/model_extras/test_model_extras.py index 76ef84e7e..7c03616bc 100644 --- a/tests/model_extras/test_model_extras.py +++ b/tests/model_extras/test_model_extras.py @@ -12,12 +12,66 @@ from vllm_omni.model_extras import ( build_image_to_image_prompt, build_image_to_video_prompt, build_text_to_image_prompt, + build_x_to_text_prompt, get_extra_body_params, get_extra_output_params, + get_x_to_text_model_family, should_init_extra_args_for_non_diffusion_stages, ) +@pytest.mark.core_model +@pytest.mark.cpu +@pytest.mark.parametrize( + ("config", "expected"), + [ + ({"model_type": "bagel"}, "bagel"), + ({"architectures": ["HunyuanImage3ForConditionalGeneration"]}, "hunyuan_image3"), + ({"model_type": "mammothmoda2"}, "mammoth_moda2"), + ({"model_type": "qwen2_vl"}, "generic"), + ], +) +def test_x_to_text_model_family(monkeypatch: pytest.MonkeyPatch, config: dict[str, object], expected: str) -> None: + monkeypatch.setattr("vllm.transformers_utils.config.get_hf_file_to_dict", lambda *args: config) + assert get_x_to_text_model_family("unused") == expected + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_bagel_x_to_text_prompt_builder() -> None: + prompt, stop_token_ids = build_x_to_text_prompt("bagel", "unused", "Describe it.", has_image=True) + assert prompt == { + "prompt": "<|im_start|>user\n<|image_pad|>\nDescribe it.<|im_end|>\n<|im_start|>assistant\n", + "modalities": ["text"], + } + assert stop_token_ids is None + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_mammoth_x_to_text_prompt_builder() -> None: + prompt, stop_token_ids = build_x_to_text_prompt("mammoth_moda2", "unused", "Hello.", has_image=False) + assert prompt == { + "prompt": ( + "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n" + "<|im_start|>user\nHello.<|im_end|>\n" + "<|im_start|>assistant\n" + ), + "modalities": ["text"], + "additional_information": {"omni_task": ["chat"]}, + } + assert stop_token_ids is None + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_generic_x_to_text_prompt_builder() -> None: + assert build_x_to_text_prompt("generic", "unused", "Hello.", has_image=False) == ( + {"prompt": "Hello.", "modalities": ["text"]}, + None, + ) + + @pytest.mark.core_model @pytest.mark.cpu def test_bagel_extra_registry_declares_request_and_response_params() -> None: diff --git a/vllm_omni/model_extras/__init__.py b/vllm_omni/model_extras/__init__.py index 5311f495b..1a3d71870 100644 --- a/vllm_omni/model_extras/__init__.py +++ b/vllm_omni/model_extras/__init__.py @@ -5,9 +5,11 @@ from vllm_omni.model_extras.registry import ( build_image_to_image_prompt, build_image_to_video_prompt, build_text_to_image_prompt, + build_x_to_text_prompt, get_extra_body_params, get_extra_output_params, get_model_class_name, + get_x_to_text_model_family, should_init_extra_args_for_non_diffusion_stages, ) @@ -15,8 +17,10 @@ __all__ = [ "build_image_to_image_prompt", "build_image_to_video_prompt", "build_text_to_image_prompt", + "build_x_to_text_prompt", "get_extra_body_params", "get_extra_output_params", "get_model_class_name", + "get_x_to_text_model_family", "should_init_extra_args_for_non_diffusion_stages", ] diff --git a/vllm_omni/model_extras/bagel.py b/vllm_omni/model_extras/bagel.py index e8cbcbf11..93646aad1 100644 --- a/vllm_omni/model_extras/bagel.py +++ b/vllm_omni/model_extras/bagel.py @@ -50,6 +50,23 @@ def build_image_to_image_prompt( return img_prompt +def build_x_to_text_prompt( + model: str, + prompt: str, + has_image: bool, +) -> tuple[dict[str, Any], list[int] | None]: + """Build BAGEL's chat prompt for text-output tasks.""" + del model + image_token = "<|image_pad|>\n" if has_image else "" + return ( + { + "prompt": f"<|im_start|>user\n{image_token}{prompt}<|im_end|>\n<|im_start|>assistant\n", + "modalities": ["text"], + }, + None, + ) + + BAGEL_EXTRA_BODY_PARAMS = frozenset( { "cfg_text_scale", diff --git a/vllm_omni/model_extras/hunyuan_image3.py b/vllm_omni/model_extras/hunyuan_image3.py new file mode 100644 index 000000000..8a53dfd54 --- /dev/null +++ b/vllm_omni/model_extras/hunyuan_image3.py @@ -0,0 +1,37 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from typing import Any + + +def build_x_to_text_prompt( + model: str, + prompt: str, + has_image: bool, +) -> tuple[dict[str, Any], list[int] | None]: + """Build HunyuanImage-3 prompt tokens and stopping rules for T2T/I2T.""" + 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) + return ( + { + "prompt": prompt, + "prompt_token_ids": built.token_ids, + "modalities": ["text"], + "use_system_prompt": resolve_sys_type(None), + }, + resolve_stop_token_ids(task=task, bot_task=None, tokenizer=tokenizer, image_size="auto"), + ) diff --git a/vllm_omni/model_extras/mammothmodal2_preview.py b/vllm_omni/model_extras/mammothmodal2_preview.py index c17e428a8..ddce493f6 100644 --- a/vllm_omni/model_extras/mammothmodal2_preview.py +++ b/vllm_omni/model_extras/mammothmodal2_preview.py @@ -66,5 +66,29 @@ MAMMOTHMODA2_PREVIEW_EXTRA_BODY_PARAMS = frozenset( "num_inference_steps", } ) + + +def build_x_to_text_prompt( + model: str, + prompt: str, + has_image: bool, +) -> tuple[dict[str, Any], list[int] | None]: + """Build MammothModa2's chat prompt for text-output tasks.""" + del model + vision = "<|vision_start|><|image_pad|><|vision_end|>" if has_image else "" + return ( + { + "prompt": ( + "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n" + f"<|im_start|>user\n{vision}{prompt}<|im_end|>\n" + "<|im_start|>assistant\n" + ), + "modalities": ["text"], + "additional_information": {"omni_task": ["chat"]}, + }, + None, + ) + + MAMMOTHMODA2_PREVIEW_EXTRA_OUTPUT_PARAMS = frozenset() MAMMOTHMODA2_PREVIEW_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES = True diff --git a/vllm_omni/model_extras/registry.py b/vllm_omni/model_extras/registry.py index ca13a35c5..44051c081 100644 --- a/vllm_omni/model_extras/registry.py +++ b/vllm_omni/model_extras/registry.py @@ -23,6 +23,7 @@ from vllm_omni.model_extras.bagel import ( from vllm_omni.model_extras.bagel import ( build_text_to_image_prompt as build_bagel_text_to_image_prompt, ) +from vllm_omni.model_extras.bagel import build_x_to_text_prompt as build_bagel_x_to_text_prompt from vllm_omni.model_extras.cosmos3 import ( COSMOS3_EXTRA_BODY_PARAMS, COSMOS3_EXTRA_OUTPUT_PARAMS, @@ -34,6 +35,7 @@ from vllm_omni.model_extras.helios import ( HELIOS_EXTRA_BODY_PARAMS, HELIOS_EXTRA_OUTPUT_PARAMS, ) +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.lingbot_video import LINGBOT_VIDEO_EXTRA_BODY_PARAMS from vllm_omni.model_extras.magi_human import ( MAGI_HUMAN_EXTRA_BODY_PARAMS, @@ -47,6 +49,9 @@ from vllm_omni.model_extras.mammothmodal2_preview import ( from vllm_omni.model_extras.mammothmodal2_preview import ( build_text_to_image_prompt as build_mammothmoda2_text_to_image_prompt, ) +from vllm_omni.model_extras.mammothmodal2_preview import ( + build_x_to_text_prompt as build_mammothmoda2_x_to_text_prompt, +) from vllm_omni.model_extras.ming_flash_omni import ( MING_FLASH_OMNI_EXTRA_BODY_PARAMS, MING_FLASH_OMNI_EXTRA_OUTPUT_PARAMS, @@ -89,6 +94,50 @@ ImageToVideoPromptBuilder = Callable[ ], dict[str, Any], ] +XToTextPromptBuilder = Callable[[str, str, bool], tuple[dict[str, Any], list[int] | None]] + + +def default_x_to_text_prompt( + model: str, + prompt: str, + has_image: bool, +) -> tuple[dict[str, Any], list[int] | None]: + del model, has_image + return {"prompt": prompt, "modalities": ["text"]}, None + + +_X_TO_TEXT_SPECS: dict[str, XToTextPromptBuilder] = { + "bagel": build_bagel_x_to_text_prompt, + "hunyuan_image3": build_hunyuan_x_to_text_prompt, + "mammoth_moda2": build_mammothmoda2_x_to_text_prompt, +} + + +def get_x_to_text_model_family(model: str) -> str: + """Resolve a text-output prompt family from a checkpoint's config.json.""" + from vllm.transformers_utils.config import get_hf_file_to_dict + + config = get_hf_file_to_dict("config.json", model) or {} + model_type = str(config.get("model_type", "")).lower() + architectures = {str(value).lower() for value in (config.get("architectures") or [])} + if model_type == "bagel" or "bagelforconditionalgeneration" in architectures: + return "bagel" + if model_type == "hunyuan_image_3_moe" or any("hunyuanimage3" in value for value in architectures): + return "hunyuan_image3" + if "mammoth" in model_type or any("mammothmoda2" in value for value in architectures): + return "mammoth_moda2" + return "generic" + + +def build_x_to_text_prompt( + model_family: str, + model: str, + prompt: str, + has_image: bool, +) -> tuple[dict[str, Any], list[int] | None]: + """Build a model-aware T2T/I2T prompt and optional stop-token ids.""" + builder = _X_TO_TEXT_SPECS.get(model_family, default_x_to_text_prompt) + return builder(model, prompt, has_image) def default_text_to_image_prompt(