[feature]: support Hidream-I1-Full model (#2572)

Signed-off-by: wangmengdi <160612767@qq.com>
Co-authored-by: wangmengdi <wangmengdi@cmss.chinamobile.com>

docs/models/supported_models.md

@@ -71,4 +71,6 @@ th {
 | `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | |
 |`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | |
 | `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+|`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | |
+
 ✅︎ indicates the model is supported on that backend. Empty cells mean not listed as supported on that backend.

docs/user_guide/diffusion_features.md

@@ -115,6 +115,7 @@ The following tables show which models support each feature:
 | **FLUX.1-Kontext-dev**   |     ❌     |     ❌      |           ❌           |       ❌        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **FLUX.2-dev**           |     ✅     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **GLM-Image**            |     ❌     |     ❌      |           ❌           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ❌        |        ❌         |
+| **Hidream-I1-Full**        |     ❌     |     ❌      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **HunyuanImage3**        |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ✅        |        ❌         |
 | **LongCat-Image**        |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **LongCat-Image-Edit**   |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |

examples/offline_inference/text_to_image/README.md

@@ -35,6 +35,7 @@ This folder provides several entrypoints for experimenting with text-to-image di
 | `black-forest-labs/FLUX.2-klein-9B` | 1024 x 1024 | 37.1 | 32.3 |
 | `black-forest-labs/FLUX.2-dev` | 1024 x 1024 | 65.7 | >80 (CPU offload required) |
 | `HunyuanImage-3.0` | 1024 x 1024 | 80.0 (TP≥3)  | 160 |
+| `HiDream-I1-Full` | 1024 x 1024 | 63.7  | 57.7 |
 
 !!! info
 *Peak VRAM:  based on basic single-card usage, batch size =1, without any acceleration/optimization features. FLUX.2-dev requires `--enable-cpu-offload` on a single 80 GiB GPU.
@@ -94,6 +95,7 @@ python text_to_image.py \
 | `--lora-scale` | float | `1.0` | Scale factor for LoRA weights |
 | `--use-system-prompt` | str | `None` | System prompt preset: `en_unified`, `en_vanilla`, `en_recaption`, `en_think_recaption`, `dynamic`, `None`, or custom text. Recommended: `en_unified`. Only for HunyuanImage-3.0.|
 | `--system-prompt` | str | `None` | Custom system prompt text. Only used when `--use-system-prompt` is set to `custom`. Only for HunyuanImage-3.0.|
+| `--auxiliary-text-encoder` | str | `None` | Supplementary auxiliary text encoder parameters model name or path (especially for Hidream-l1-full). |
 
 **NextStep-1.1 specific arguments:**
 
@@ -177,6 +179,24 @@ python examples/offline_inference/text_to_image/text_to_image.py \
   --output flux2-dev.png
 ```
 
+
+### HiDream-I1-Full Models
+
+The `--auxiliary-text-encoder` parameter is required when running HiDream‑I1‑Full:
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model HiDream-ai/HiDream-I1-Full \
+  --prompt "The setting sun of late autumn dyes the riverside with a warm orange hue" \
+  --seed 42 \
+  --guidance-scale 5.0 \
+  --tensor-parallel-size 1 \
+  --num-images-per-prompt 1 \
+  --num-inference-steps 50 \
+  --auxiliary-text-encoder meta-llama/Meta-Llama-3.1-8B-Instruct \
+  --output /output.png
+```
+
 ### Batch Requests (Multiple Prompts)
 
 You can pass multiple prompts in a single `generate` call.

examples/offline_inference/text_to_image/text_to_image.py

@@ -300,6 +300,12 @@ def parse_args() -> argparse.Namespace:
         default=None,
         help=("Custom system prompt. Used when --use-system-prompt is custom. "),
     )
+    parser.add_argument(
+        "--auxiliary-text-encoder",
+        type=str,
+        default=None,
+        help="Supplementary auxiliary text encoder parameters model name or path (especially for Hidream-l1-full).",
+    )
     current_omni_platform.pre_register_and_update(parser)
     from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 
@@ -392,6 +398,7 @@ def main():
         "profiler_config": args.profiler_config,
         "init_timeout": args.init_timeout,
         "stage_init_timeout": args.stage_init_timeout,
+        "auxiliary_text_encoder": args.auxiliary_text_encoder,
         **lora_args,
         **quant_kwargs,
     }

pyproject.toml

@@ -241,6 +241,8 @@ extend-ignore-identifiers-re = [
     ".*MoBA",
     ".*temperal_downsample",
     ".*_nd",
+    ".*nd$",
+    ".*EmbedND$",
     ".*nothink.*",
     ".*NOTHINK.*",
     ".*nin.*",

recipes/HiDream/HiDream-I1-Image.md

@@ -0,0 +1,106 @@
+# HiDream-I1-Full
+
+> Text-to-image serving and benchmark
+
+## Summary
+
+- Vendor: HiDream.ai
+- Model: `HiDream-ai/HiDream-I1-Full`
+- Task: Text-to-image generation
+- Mode: Online serving with the OpenAI-compatible API
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe when you want a known-good starting point for serving
+`HiDream-ai/HiDream-I1-Full` on a single 80 GB A800, validate the normal online-serving path.
+
+## References
+
+- Model: <https://huggingface.co/HiDream-ai/HiDream-I1-Full>
+- Upstream or canonical docs:
+  [`docs/user_guide/examples/online_serving/text_to_image.md`](../../docs/user_guide/examples/online_serving/text_to_image.md)
+- Related example under `examples/`:
+  [`examples/online_serving/text_to_image/README.md`](../../examples/online_serving/text_to_image/README.md)
+
+
+## Hardware Support
+
+This recipe currently documents one CUDA GPU serving configuration.
+
+## GPU
+
+### 1x A800 80GB
+
+#### Environment
+
+- OS: Linux
+- Python: 3.10+
+- Driver / runtime: NVIDIA CUDA environment with an A100 80 GB GPU
+- vLLM version: Match the repository requirements for your checkout
+- vLLM-Omni version or commit: Use the commit you are deploying from
+
+#### Command
+
+Start the baseline server:
+
+```bash
+export MODEL_NAME_OR_PATH=HiDream-ai/HiDream-I1-Full
+vllm serve ${MODEL_NAME_OR_PATH} \
+   --omni \
+   --port 8092 \
+   --auxiliary-text-encoder meta-llama/Llama-3.1-8B-Instruct \
+   --tensor-parallel-size 1 \
+   --vae_use_slicing \
+   --vae_use_tiling
+```
+
+
+You can also use the example launcher and pass the extra flags through:
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model HiDream-ai/HiDream-I1-Full \
+  --prompt "The setting sun of late autumn dyes the riverside with a warm orange hue" \
+  --seed 42 \
+  --guidance-scale 5.0 \
+  --tensor-parallel-size 1 \
+  --num-images-per-prompt 1 \
+  --num-inference-steps 50 \
+  --auxiliary-text-encoder meta-llama/Llama-3.1-8B-Instruct \
+  --output /workspace/mnt/cmss-wmd/result/output.png
+
+```
+
+#### Verification
+
+Run the existing client example after the server is ready:
+
+```bash
+python examples/online_serving/text_to_image/openai_chat_client.py \
+  --server http://localhost:8092 \
+  --prompt "The setting sun of late autumn dyes the riverside with a warm orange hue" \
+  --output /tmp/hidream_i1_recipe.png
+```
+
+For a direct API smoke test:
+
+```bash
+curl -s http://localhost:8092/v1/chat/completions   -H "Content-Type: application/json"   -d '{
+    "messages": [
+      {"role": "user", "content": "The setting sun of late autumn dyes the riverside with a warm orange hue"}
+    ],
+    "extra_body": {
+      "num_inference_steps": 50,
+      "seed": 42,
+      "height": 1024,
+      "width": 1024
+    }
+  }' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output_online.png
+```
+
+
+
+#### Notes
+
+- Key flags: `--auxiliary-text-encoder` designates the path for the auxiliary text model meta-llama/Llama-3.1-8B-Instruct. For HiDream-I1-Full, unspecified use defaults meta-llama/Llama-3.1-8B-Instruct (downloaded from official Hugging Face), and custom paths are supported.

tests/e2e/online_serving/test_hidream_i1_full_expansion.py

@@ -0,0 +1,69 @@
+"""
+End-to-end diffusion coverage for HiDream-I1-Full in online serving mode.
+
+This test verifies that HiDream-I1-Full can be launched,
+accepts text-to-image requests through the OpenAI-compatible API, and returns
+valid generated images with the requested resolution.
+
+assert_diffusion_response validates successful generation and the expected
+image resolution.
+"""
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
+
+MODEL = "HiDream-ai/HiDream-I1-Full"
+PROMPT = "A cinematic mountain landscape at sunrise, dramatic clouds, ultra-detailed, realistic photography."
+NEGATIVE_PROMPT = "low quality, blurry, distorted, deformed, watermark"
+
+SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
+
+
+def _get_hidream_i1_image_feature_cases(model: str):
+    """Return HiDream-I1-Full diffusion feature cases."""
+
+    return [
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--auxiliary-text-encoder",
+                    "meta-llama/Meta-Llama-3.1-8B-Instruct",
+                ],
+            ),
+            id="default",
+            marks=SINGLE_CARD_FEATURE_MARKS,
+        ),
+    ]
+
+
+@pytest.mark.parametrize(
+    "omni_server",
+    _get_hidream_i1_image_feature_cases(MODEL),
+    indirect=True,
+)
+def test_hidream_i1_image(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+):
+    """Validate HiDream-I1-Full online serving with CPU offload."""
+
+    messages = dummy_messages_from_mix_data(content_text=PROMPT)
+
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "extra_body": {
+            "height": 512,
+            "width": 512,
+            "num_inference_steps": 2,
+            "negative_prompt": NEGATIVE_PROMPT,
+            "seed": 42,
+        },
+    }
+
+    openai_client.send_diffusion_request(request_config)

vllm_omni/diffusion/data.py

@@ -11,7 +11,7 @@ from typing import TYPE_CHECKING, Any
 import diffusers
 import torch
 from PIL import Image
-from pydantic import model_validator
+from pydantic import Field, model_validator
 from typing_extensions import Self
 from vllm.config.utils import config
 from vllm.logger import init_logger
@@ -593,6 +593,9 @@ class OmniDiffusionConfig:
     # Maximum number of sequences to generate in a batch
     max_num_seqs: int = 1
 
+    # Supplementary model specific parameters
+    extras: dict[str, Any] = Field(default_factory=dict)
+
     @property
     def is_moe(self) -> bool:
         num_experts = self.tf_model_config.get("num_experts", None)

vllm_omni/diffusion/models/hidream_image/__init__.py

@@ -0,0 +1,17 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""HiDream Image diffusion model components."""
+
+from vllm_omni.diffusion.models.hidream_image.hidream_image_transformer import (
+    HiDreamImageTransformer2DModel,
+)
+from vllm_omni.diffusion.models.hidream_image.pipeline_hidream_image import (
+    HiDreamImagePipeline,
+    get_hidream_image_post_process_func,
+)
+
+__all__ = [
+    "HiDreamImageTransformer2DModel",
+    "HiDreamImagePipeline",
+    "get_hidream_image_post_process_func",
+]

vllm_omni/diffusion/registry.py

@@ -261,6 +261,11 @@ _DIFFUSION_MODELS = {
         "pipeline_diffusers_adapter",
         "DiffusersAdapterPipeline",
     ),
+    "HiDreamImagePipeline": (
+        "hidream_image",
+        "pipeline_hidream_image",
+        "HiDreamImagePipeline",
+    ),
 }
 
 
@@ -482,6 +487,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "OmniVoicePipeline": "get_omnivoice_post_process_func",
     "DreamIDOmniPipeline": "get_dreamid_omni_post_process_func",
     "SenseNovaU1Pipeline": "get_sensenova_u1_post_process_func",
+    "HiDreamImagePipeline": "get_hidream_image_post_process_func",
 }
 
 _DIFFUSION_PRE_PROCESS_FUNCS = {

vllm_omni/engine/async_omni_engine.py

@@ -1950,6 +1950,10 @@ class AsyncOmniEngine:
             "force_cutlass_fp8": bool(kwargs.get("force_cutlass_fp8", False)),
             "enable_diffusion_pipeline_profiler": kwargs.get("enable_diffusion_pipeline_profiler", False),
             "enable_ar_profiler": kwargs.get("enable_ar_profiler", False),
+            "extras": {
+                "auxiliary_text_encoder": kwargs.get("auxiliary_text_encoder", None),
+                "default_llama_model_id": kwargs.get("default_llama_model_id", "meta-llama/Meta-Llama-3.1-8B-Instruct"),
+            },
             **(
                 {
                     "profiler_config": asdict(kwargs["profiler_config"])

vllm_omni/entrypoints/cli/serve.py

@@ -650,6 +650,16 @@ class OmniServeCommand(CLISubcommand):
             action="store_true",
             help="Enable AR stage profiler to include AR stage timing in stage_durations.",
         )
+
+        # Supplementary auxiliary text encoder parameters
+        # (e.g., the meta llama/meta llama-3.1-8b-instrument used by hidream)
+        omni_config_group.add_argument(
+            "--auxiliary-text-encoder",
+            type=str,
+            default=None,
+            help="Auxiliary text encoder parameters model name or path (especially for Hidream-l1-full).",
+        )
+
         # Stash via type(self) so the docs hook (which execs this function in a
         # sandboxed globals dict via ``DummySelf``) doesn't fail on a NameError.
         type(self)._parser = serve_parser