[Model] Add HiDream-O1-Image support (#5194)

Signed-off-by: yixiaoer <shin@yixiaoer.sg>

.buildkite/cuda/test-nightly.yml

@@ -255,6 +255,12 @@ steps:
           - pytest -s -v tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py -m full_model --run-level full_model
         mirror_hardwares: h100_4
 
+      - label: ":full_moon: HiDream-O1-Image · Accuracy Test"
+        timeout_in_minutes: 180
+        commands:
+          - pytest -s -v tests/e2e/accuracy/test_hidream_o1_image.py -m full_model --run-level full_model
+        mirror_hardwares: h100_2
+
       - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image"
         key: nightly-diffusion-x2iat-performance-qwen-image
         timeout_in_minutes: 180

docs/models/supported_models.md

@@ -109,6 +109,7 @@ th {
 | `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `GepardTalkerForConditionalGeneration` | Gepard-1.0 | `nineninesix/gepard-1.0` | ✅︎ | | | | — |
 |`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | | — |
+| `HiDreamO1ImagePipeline` | HiDream-O1-Image | `HiDream-ai/HiDream-O1-Image` | ✅︎ | | | | — |
 
 The `Recipe` column links the corresponding deployment recipe. `Published` is
 preferred when a corresponding `recipes.vllm.ai` page exists; otherwise the

docs/user_guide/diffusion_features.md

@@ -124,6 +124,7 @@ The following tables show which models support each feature:
 | **FLUX.2-dev**           |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |          ✅           |       ❌        |        ❌         |
 | **GLM-Image**            |     ❌     |     ❌      |           ❌           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **Hidream-I1-Full**        |     ❌     |     ❌      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
+| **HiDream-O1-Image**     |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **HunyuanImage3**        |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ✅        |        ✅*        |
 | **Krea 2**               |     ❌     |     ✅      |           ❌           |       ❌        |         ❌         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **LongCat-Image**        |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |

examples/offline_inference/text_to_image/text_to_image.py

@@ -89,7 +89,7 @@ def parse_args() -> argparse.Namespace:
         "Qwen/Qwen-Image, Tongyi-MAI/Z-Image-Turbo, Qwen/Qwen-Image-2512, stepfun-ai/NextStep-1.1, "
         "black-forest-labs/FLUX.1-dev, black-forest-labs/FLUX.2-klein-9B, "
         "black-forest-labs/FLUX.2-dev, tencent/HunyuanImage-3.0-Instruct, "
-        "meituan-longcat/LongCat-Image, OvisAI/Ovis-Image, "
+        "HiDream-ai/HiDream-O1-Image, meituan-longcat/LongCat-Image, OvisAI/Ovis-Image, "
         "stabilityai/stable-diffusion-3.5-medium, Tongyi-MAI/Z-Image-Turbo and etc.",
     )
     parser.add_argument(

recipes/HiDream-ai/HiDream-O1-Image.md

@@ -0,0 +1,113 @@
+# HiDream-O1-Image
+
+> HiDream-O1 text-to-image generation through the shared offline example
+
+## Summary
+
+- Vendor: HiDream.ai
+- Model: `HiDream-ai/HiDream-O1-Image`
+- Task: Text-to-image
+- Mode: Offline inference
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe to generate native-resolution images with HiDream-O1 on one
+or two CUDA GPUs. The two-GPU command shards the Qwen3-VL attention and MLP
+weights, as well as the timestep and pixel projections.
+
+## References
+
+- Upstream model:
+  [`HiDream-ai/HiDream-O1-Image`](https://huggingface.co/HiDream-ai/HiDream-O1-Image)
+- Related offline example:
+  [`examples/offline_inference/text_to_image/text_to_image.py`](../../examples/offline_inference/text_to_image/text_to_image.py)
+- Related issue:
+  [`vllm-project/vllm-omni#3479`](https://github.com/vllm-project/vllm-omni/issues/3479)
+
+## Hardware Support
+
+## GPU
+
+### 1x H100 80GB
+
+#### Environment
+
+- OS: Linux
+- Python: Match the repository requirements for your checkout
+- Driver / runtime: NVIDIA CUDA
+- vLLM version: Match the repository requirements for your checkout
+- vLLM-Omni version or commit: Use the commit you are deploying from
+
+#### Command
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model HiDream-ai/HiDream-O1-Image \
+  --prompt "A cat is sitting next to a sign that says 'HiDream-O1 vLLM-Omni'" \
+  --height 2048 \
+  --width 2048 \
+  --num-inference-steps 50 \
+  --guidance-scale 5.0 \
+  --seed 42 \
+  --output hidream_o1_output.png
+```
+
+#### Verification
+
+```bash
+python -c "from PIL import Image; im = Image.open('hidream_o1_output.png'); print(im.mode, im.size)"
+```
+
+Expected output:
+
+```text
+RGB (2048, 2048)
+```
+
+#### Notes
+
+- The pipeline supports BF16 text-to-image generation.
+- Requested sizes are matched to the closest resolution supported by the checkpoint.
+
+### 2x H100 80GB
+
+#### Environment
+
+- OS: Linux
+- Python: Match the repository requirements for your checkout
+- Driver / runtime: NVIDIA CUDA
+- vLLM version: Match the repository requirements for your checkout
+- vLLM-Omni version or commit: Use the commit you are deploying from
+
+#### Command
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model HiDream-ai/HiDream-O1-Image \
+  --tensor-parallel-size 2 \
+  --prompt "A cat is sitting next to a sign that says 'HiDream-O1 vLLM-Omni'" \
+  --height 2048 \
+  --width 2048 \
+  --num-inference-steps 50 \
+  --guidance-scale 5.0 \
+  --seed 42 \
+  --output hidream_o1_tp2_output.png
+```
+
+#### Verification
+
+```bash
+python -c "from PIL import Image; im = Image.open('hidream_o1_tp2_output.png'); print(im.mode, im.size)"
+```
+
+Expected output:
+
+```text
+RGB (2048, 2048)
+```
+
+#### Notes
+
+- Tensor parallelism shards the attention, MLP, timestep, and pixel projection weights across both GPUs.
+- Image editing and image-conditioned generation are not supported.

recipes/README.md

@@ -47,6 +47,7 @@ recipes/
 | [`ByteDance/Lance.md`](./ByteDance/Lance.md) | Unified AR+diffusion: text/img/video gen + understanding (Lance 3B) | 1x B300 / A100 80GB |
 | [`fishaudio/Fish-Speech-S2-Pro.md`](./fishaudio/Fish-Speech-S2-Pro.md) | Online serving for TTS | 1x A800 80GB |
 | [`Helios/Helios.md`](./Helios/Helios.md) | Text-to-video, image-to-video, and video-to-video generation | 1x NVIDIA H20 |
+| [`HiDream-ai/HiDream-O1-Image.md`](./HiDream-ai/HiDream-O1-Image.md) | Text-to-image with shared offline example | 1x or 2x H100 80GB |
 | [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md) | Online serving for multimodal chat + standalone TTS | 4x H100 / 1x H100 80GB |
 | [`inclusionAI/Ming-omni-tts.md`](./inclusionAI/Ming-omni-tts.md) | Offline + online dense Ming TTS/audio generation | 1x H100 80GB / 1x AMD MI300X (ROCm 7.2) |
 | [`k2-fsa/OmniVoice.md`](./k2-fsa/OmniVoice.md) | Offline and online multilingual text-to-speech | 1x AMD MI300X (ROCm 7.2) |

tests/diffusion/models/hidream_o1_image/test_hidream_o1_image_cache.py

@@ -0,0 +1,67 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+from cache_dit import ForwardPattern
+from torch import nn
+
+from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig, CacheDiTBackend
+from vllm_omni.diffusion.models.hidream_o1_image.hidream_o1_image_transformer import (
+    HiDreamO1TextModel,
+)
+from vllm_omni.diffusion.models.hidream_o1_image.pipeline_hidream_o1_image import (
+    HiDreamO1ImagePipeline,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def _empty_text_model() -> HiDreamO1TextModel:
+    model = HiDreamO1TextModel.__new__(HiDreamO1TextModel)
+    nn.Module.__init__(model)
+    model.layers = nn.ModuleList([nn.Identity()])
+    return model
+
+
+def test_cache_dit_targets_the_decoder_module_that_owns_layers():
+    text_model = _empty_text_model()
+    pipeline = SimpleNamespace(
+        _dit_modules=HiDreamO1ImagePipeline._dit_modules,
+        model=SimpleNamespace(
+            model=SimpleNamespace(language_model=text_model),
+        ),
+    )
+
+    with (
+        patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter") as block_adapter_cls,
+        patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit") as cache_dit,
+    ):
+        backend = CacheDiTBackend()
+        backend.enable(pipeline)
+        backend.refresh(pipeline, num_inference_steps=50)
+
+    adapter_config = HiDreamO1TextModel._cache_dit_adapter_config
+    assert isinstance(adapter_config, CacheDiTAdapterConfig)
+    assert adapter_config.block_forward_patterns == {"layers": ForwardPattern.Pattern_3}
+    assert adapter_config.has_separate_cfg is True
+
+    block_adapter_cls.assert_called_once_with(
+        transformer=text_model,
+        blocks=[text_model.layers],
+        forward_pattern=[ForwardPattern.Pattern_3],
+        has_separate_cfg=True,
+        check_forward_pattern=True,
+    )
+    cache_dit.enable_cache.assert_called_once()
+    enable_call = cache_dit.enable_cache.call_args
+    assert enable_call.args == (block_adapter_cls.return_value,)
+    assert enable_call.kwargs["cache_config"].Fn_compute_blocks == 1
+    assert enable_call.kwargs["calibrator_config"] is None
+    cache_dit.refresh_context.assert_called_once_with(
+        text_model,
+        num_inference_steps=50,
+        verbose=True,
+    )

tests/e2e/accuracy/test_hidream_o1_image.py

@@ -0,0 +1,65 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import pytest
+from PIL import Image
+
+from tests.e2e.accuracy.helpers import assert_similarity, model_output_dir
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.diffusion.utils.image_output import extract_images_from_outputs
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
+
+MODEL_ID = "HiDream-ai/HiDream-O1-Image"
+PROMPT = "A cat is sitting next to a sign that says 'HiDream-O1 vLLM-Omni'"
+BASELINE_PATH = Path(__file__).resolve().parents[2] / "assets" / "hidream_o1" / "hidream_o1_baseline.png"
+SSIM_THRESHOLD = 0.8
+PSNR_THRESHOLD = 19.0
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=2)
+def test_hidream_o1_image_tp2_matches_baseline(
+    accuracy_artifact_root: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setenv("DIFFUSION_ATTENTION_BACKEND", "TORCH_SDPA")
+    model = os.environ.get("HIDREAM_O1_IMAGE_MODEL", MODEL_ID)
+    with OmniRunner(
+        model,
+        parallel_config=DiffusionParallelConfig(tensor_parallel_size=2),
+        enforce_eager=True,
+    ) as runner:
+        outputs = runner.omni.generate(
+            {"prompt": PROMPT},
+            OmniDiffusionSamplingParams(
+                height=2048,
+                width=2048,
+                num_inference_steps=50,
+                guidance_scale=5.0,
+                seed=42,
+            ),
+        )
+        images = extract_images_from_outputs(outputs)
+        assert images
+        image = images[0].convert("RGB")
+
+    output_dir = model_output_dir(accuracy_artifact_root, MODEL_ID)
+    image.save(output_dir / "vllm_omni_tp2.png")
+    assert BASELINE_PATH.exists()
+    baseline = Image.open(BASELINE_PATH).convert("RGB")
+    resized = image.resize(baseline.size, Image.Resampling.LANCZOS)
+    assert_similarity(
+        model_name=MODEL_ID,
+        vllm_image=resized,
+        diffusers_image=baseline,
+        ssim_threshold=SSIM_THRESHOLD,
+        psnr_threshold=PSNR_THRESHOLD,
+    )

tests/model_tests/diffusion/test_alignment.py

@@ -74,6 +74,7 @@ EXCLUDED_MODELS = [
     "Cosmos3OmniPipeline",
     "DiffusersAdapterPipeline",
     "HiDreamImagePipeline",
+    "HiDreamO1ImagePipeline",
     "DreamZeroPipeline",
     "StableDiffusionXLPipeline",
     "Gr00tN1d7Pipeline",

vllm_omni/diffusion/data.py

@@ -613,6 +613,10 @@ def resolve_model_class_name(
     model_type = cfg.get("model_type")
     architectures = cfg.get("architectures") or []
 
+    from vllm_omni.diffusion.utils.hf_utils import _looks_like_hidream_o1
+
+    if _looks_like_hidream_o1(model, cfg):
+        return "HiDreamO1ImagePipeline"
     if model_type == "bagel" or "BagelForConditionalGeneration" in architectures:
         return "BagelPipeline"
     if (
@@ -1318,7 +1322,17 @@ class OmniDiffusionConfig:
                 model_type = cfg.get("model_type")
                 architectures = cfg.get("architectures") or []
 
-                if model_type == "bagel" or "BagelForConditionalGeneration" in architectures:
+                from vllm_omni.diffusion.utils.hf_utils import _looks_like_hidream_o1
+
+                is_hidream_o1 = _looks_like_hidream_o1(self.model, cfg)
+                if self.model_class_name == "HiDreamO1ImagePipeline" and not is_hidream_o1:
+                    raise ValueError(f"Checkpoint {self.model} does not have the HiDream-O1 signature")
+
+                if is_hidream_o1:
+                    self.model_class_name = "HiDreamO1ImagePipeline"
+                    self.set_tf_model_config(TransformerConfig())
+                    self.update_multimodal_support()
+                elif model_type == "bagel" or "BagelForConditionalGeneration" in architectures:
                     self.model_class_name = "BagelPipeline"
                     self.set_tf_model_config(TransformerConfig())
                     self.update_multimodal_support()

vllm_omni/diffusion/models/hidream_o1_image/__init__.py

@@ -0,0 +1,15 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from vllm_omni.diffusion.models.hidream_o1_image.hidream_o1_image_transformer import (
+    HiDreamO1ImageTransformer,
+)
+from vllm_omni.diffusion.models.hidream_o1_image.pipeline_hidream_o1_image import (
+    HiDreamO1ImagePipeline,
+    get_hidream_o1_image_post_process_func,
+)
+
+__all__ = [
+    "HiDreamO1ImageTransformer",
+    "HiDreamO1ImagePipeline",
+    "get_hidream_o1_image_post_process_func",
+]

vllm_omni/diffusion/models/hidream_o1_image/hidream_o1_image_transformer.py

@@ -0,0 +1,572 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Copyright 2025 The Qwen Team and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import math
+from collections.abc import Iterable
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import torch
+from cache_dit import ForwardPattern
+from torch import nn
+from transformers import Qwen3VLConfig
+from vllm.distributed import get_tensor_model_parallel_world_size
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.model_executor.models.qwen3 import Qwen3MLP
+
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig
+
+if TYPE_CHECKING:
+    from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+
+
+def _prefix(parent: str, child: str) -> str:
+    return f"{parent}.{child}" if parent else child
+
+
+def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def _apply_rotary_pos_emb(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    cos = cos.unsqueeze(1)
+    sin = sin.unsqueeze(1)
+    return (
+        query * cos + _rotate_half(query) * sin,
+        key * cos + _rotate_half(key) * sin,
+    )
+
+
+class HiDreamO1RotaryEmbedding(nn.Module):
+    def __init__(self, config) -> None:
+        super().__init__()
+        rope_parameters = getattr(config, "rope_parameters", None) or getattr(config, "rope_scaling", None) or {}
+        head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+        rope_theta = rope_parameters.get("rope_theta", getattr(config, "rope_theta", 10000.0))
+        inv_freq = 1.0 / (float(rope_theta) ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.mrope_section = rope_parameters.get("mrope_section", [24, 20, 20])
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_ids: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if position_ids.ndim == 2:
+            position_ids = position_ids.unsqueeze(0).expand(3, -1, -1)
+
+        inv_freq = self.inv_freq[None, None, :, None].float().to(hidden_states.device)
+        inv_freq = inv_freq.expand(3, position_ids.shape[1], -1, 1)
+        positions = position_ids[:, :, None, :].float()
+        device_type = hidden_states.device.type if hidden_states.device.type != "mps" else "cpu"
+        with torch.autocast(device_type=device_type, enabled=False):
+            freqs = (inv_freq @ positions).transpose(2, 3)
+            freqs_t = freqs[0]
+            for dim, offset in enumerate((1, 2), start=1):
+                length = self.mrope_section[dim] * 3
+                freqs_t[..., offset:length:3] = freqs[dim, ..., offset:length:3]
+            emb = torch.cat((freqs_t, freqs_t), dim=-1)
+            cos = emb.cos()
+            sin = emb.sin()
+        return cos.to(hidden_states.dtype), sin.to(hidden_states.dtype)
+
+
+class HiDreamO1Attention(nn.Module):
+    def __init__(
+        self,
+        config,
+        *,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+        tp_size = get_tensor_model_parallel_world_size()
+        if config.num_attention_heads % tp_size != 0:
+            raise ValueError(
+                f"num_attention_heads={config.num_attention_heads} must be divisible by tensor_parallel_size={tp_size}"
+            )
+        if config.num_key_value_heads >= tp_size:
+            if config.num_key_value_heads % tp_size != 0:
+                raise ValueError(
+                    f"num_key_value_heads={config.num_key_value_heads} must be divisible by "
+                    f"tensor_parallel_size={tp_size}"
+                )
+        elif tp_size % config.num_key_value_heads != 0:
+            raise ValueError(
+                f"tensor_parallel_size={tp_size} must be divisible by num_key_value_heads={config.num_key_value_heads}"
+            )
+
+        self.qkv_proj = QKVParallelLinear(
+            config.hidden_size,
+            self.head_dim,
+            config.num_attention_heads,
+            config.num_key_value_heads,
+            bias=config.attention_bias,
+            quant_config=quant_config,
+            prefix=_prefix(prefix, "qkv_proj"),
+        )
+        self.o_proj = RowParallelLinear(
+            config.num_attention_heads * self.head_dim,
+            config.hidden_size,
+            bias=config.attention_bias,
+            quant_config=quant_config,
+            prefix=_prefix(prefix, "o_proj"),
+        )
+        self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.scaling = self.head_dim**-0.5
+        self.attn = Attention(
+            num_heads=self.qkv_proj.num_heads,
+            head_size=self.head_dim,
+            causal=False,
+            softmax_scale=self.scaling,
+            num_kv_heads=self.qkv_proj.num_kv_heads,
+            prefix=_prefix(prefix, "attn"),
+            role="self",
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None,
+        full_attn_spans: list[list[tuple[int, int]]] | None,
+    ) -> torch.Tensor:
+        batch_size, seq_len, _ = hidden_states.shape
+        qkv, _ = self.qkv_proj(hidden_states)
+        q_size = self.qkv_proj.num_heads * self.head_dim
+        kv_size = self.qkv_proj.num_kv_heads * self.head_dim
+        query, key, value = qkv.split([q_size, kv_size, kv_size], dim=-1)
+
+        query = query.view(batch_size, seq_len, self.qkv_proj.num_heads, self.head_dim).transpose(1, 2)
+        key = key.view(batch_size, seq_len, self.qkv_proj.num_kv_heads, self.head_dim).transpose(1, 2)
+        value = value.view(batch_size, seq_len, self.qkv_proj.num_kv_heads, self.head_dim).transpose(1, 2)
+        query = self.q_norm(query)
+        key = self.k_norm(key)
+        query, key = _apply_rotary_pos_emb(query, key, *position_embeddings)
+        query = query.transpose(1, 2)
+        key = key.transpose(1, 2)
+        value = value.transpose(1, 2)
+
+        attn_output = self.attn(
+            query,
+            key,
+            value,
+            AttentionMetadata(
+                attn_mask=attention_mask,
+                full_attn_spans=full_attn_spans,
+            ),
+        )
+        attn_output = attn_output.reshape(batch_size, seq_len, q_size)
+        attn_output, _ = self.o_proj(attn_output)
+        return attn_output
+
+
+class HiDreamO1DecoderLayer(nn.Module):
+    def __init__(
+        self,
+        config,
+        *,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.self_attn = HiDreamO1Attention(
+            config,
+            quant_config=quant_config,
+            prefix=_prefix(prefix, "self_attn"),
+        )
+        self.mlp = Qwen3MLP(
+            hidden_size=config.hidden_size,
+            intermediate_size=config.intermediate_size,
+            hidden_act=config.hidden_act,
+            quant_config=quant_config,
+            prefix=_prefix(prefix, "mlp"),
+        )
+        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None,
+        full_attn_spans: list[list[tuple[int, int]]] | None,
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states = self.self_attn(
+            hidden_states,
+            position_embeddings,
+            attention_mask,
+            full_attn_spans,
+        )
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        return residual + hidden_states
+
+
+class HiDreamO1TextModel(nn.Module):
+    _cache_dit_adapter_config = CacheDiTAdapterConfig(
+        block_forward_patterns={"layers": ForwardPattern.Pattern_3},
+        has_separate_cfg=True,
+    )
+
+    def __init__(
+        self,
+        config,
+        *,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.embed_tokens = VocabParallelEmbedding(
+            config.vocab_size,
+            config.hidden_size,
+            quant_config=quant_config,
+            prefix=_prefix(prefix, "embed_tokens"),
+        )
+        self.layers = nn.ModuleList(
+            [
+                HiDreamO1DecoderLayer(
+                    config,
+                    quant_config=quant_config,
+                    prefix=_prefix(prefix, f"layers.{layer_idx}"),
+                )
+                for layer_idx in range(config.num_hidden_layers)
+            ]
+        )
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = HiDreamO1RotaryEmbedding(config)
+
+    def forward(
+        self,
+        inputs_embeds: torch.Tensor,
+        position_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+        full_attn_spans: list[list[tuple[int, int]]] | None,
+    ) -> torch.Tensor:
+        position_embeddings = self.rotary_emb(inputs_embeds, position_ids)
+        hidden_states = inputs_embeds
+        for layer in self.layers:
+            hidden_states = layer(
+                hidden_states,
+                position_embeddings=position_embeddings,
+                attention_mask=attention_mask,
+                full_attn_spans=full_attn_spans,
+            )
+        return self.norm(hidden_states)
+
+
+class TimestepEmbedder(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        *,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+        frequency_embedding_size: int = 256,
+    ) -> None:
+        super().__init__()
+        self.mlp = nn.Sequential(
+            ColumnParallelLinear(
+                frequency_embedding_size,
+                hidden_size,
+                bias=True,
+                gather_output=False,
+                quant_config=quant_config,
+                prefix=_prefix(prefix, "mlp.0"),
+                return_bias=False,
+            ),
+            nn.SiLU(),
+            RowParallelLinear(
+                hidden_size,
+                hidden_size,
+                bias=True,
+                input_is_parallel=True,
+                quant_config=quant_config,
+                prefix=_prefix(prefix, "mlp.2"),
+                return_bias=False,
+            ),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+    @staticmethod
+    def timestep_embedding(timestep: torch.Tensor, dim: int, max_period: int = 10000) -> torch.Tensor:
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period) * torch.arange(0, half, dtype=torch.float32, device=timestep.device) / half
+        )
+        args = timestep[:, None].float() * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        return embedding
+
+    def forward(self, timestep: torch.Tensor) -> torch.Tensor:
+        embedding = self.timestep_embedding(timestep * 1000, self.frequency_embedding_size)
+        return self.mlp(embedding.to(self.mlp[0].weight.dtype))
+
+
+class BottleneckPatchEmbed(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        *,
+        patch_size: int,
+        in_channels: int,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        bottleneck_dim = hidden_size // 4
+        self.proj1 = ColumnParallelLinear(
+            patch_size * patch_size * in_channels,
+            bottleneck_dim,
+            bias=False,
+            gather_output=False,
+            quant_config=quant_config,
+            prefix=_prefix(prefix, "proj1"),
+            return_bias=False,
+        )
+        self.proj2 = RowParallelLinear(
+            bottleneck_dim,
+            hidden_size,
+            bias=True,
+            input_is_parallel=True,
+            quant_config=quant_config,
+            prefix=_prefix(prefix, "proj2"),
+            return_bias=False,
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.proj2(self.proj1(hidden_states))
+
+
+class FinalLayer(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        *,
+        patch_size: int,
+        out_channels: int,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.linear = ColumnParallelLinear(
+            hidden_size,
+            patch_size * patch_size * out_channels,
+            bias=True,
+            gather_output=True,
+            quant_config=quant_config,
+            prefix=_prefix(prefix, "linear"),
+            return_bias=False,
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.linear(hidden_states)
+
+
+@dataclass
+class HiDreamO1ImageOutput:
+    x_pred: torch.Tensor
+
+
+class HiDreamO1ImageModel(nn.Module):
+    def __init__(
+        self,
+        config: Qwen3VLConfig,
+        *,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.config = config
+        self.language_model = HiDreamO1TextModel(
+            config.text_config,
+            quant_config=quant_config,
+            prefix=_prefix(prefix, "language_model"),
+        )
+        self.patch_size = 32
+        self.in_channels = 3
+        hidden_size = config.text_config.hidden_size
+        self.t_embedder1 = TimestepEmbedder(
+            hidden_size,
+            quant_config=quant_config,
+            prefix=_prefix(prefix, "t_embedder1"),
+        )
+        self.x_embedder = BottleneckPatchEmbed(
+            hidden_size,
+            patch_size=self.patch_size,
+            in_channels=self.in_channels,
+            quant_config=quant_config,
+            prefix=_prefix(prefix, "x_embedder"),
+        )
+        self.final_layer2 = FinalLayer(
+            hidden_size,
+            patch_size=self.patch_size,
+            out_channels=self.in_channels,
+            quant_config=quant_config,
+            prefix=_prefix(prefix, "final_layer2"),
+        )
+        self.tms_token_id = 151673
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        position_ids: torch.Tensor,
+        vinputs: torch.Tensor,
+        timestep: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+        full_attn_spans: list[list[tuple[int, int]]] | None,
+    ) -> HiDreamO1ImageOutput:
+        inputs_embeds = self.language_model.embed_tokens(input_ids)
+        timestep_embeds = self.t_embedder1(timestep.to(inputs_embeds.device))
+        timestep_mask = input_ids == self.tms_token_id
+        inputs_embeds = torch.where(
+            timestep_mask.unsqueeze(-1),
+            timestep_embeds.unsqueeze(1).expand_as(inputs_embeds),
+            inputs_embeds,
+        )
+
+        image_embeds = self.x_embedder(vinputs.to(inputs_embeds.device)).to(inputs_embeds.dtype)
+        inputs_embeds = torch.cat([inputs_embeds, image_embeds], dim=1)
+        batch_size, total_seq_len, _ = inputs_embeds.shape
+
+        if attention_mask is not None:
+            if attention_mask.shape[0] == 1 and batch_size > 1:
+                attention_mask = attention_mask.expand(batch_size, -1, -1, -1)
+            expected_mask_shape = (batch_size, 1, total_seq_len, total_seq_len)
+            if attention_mask.shape != expected_mask_shape:
+                raise ValueError(
+                    f"attention_mask must have shape {expected_mask_shape}, got {tuple(attention_mask.shape)}"
+                )
+            attention_mask = attention_mask.to(device=inputs_embeds.device, dtype=inputs_embeds.dtype)
+
+        if full_attn_spans is None or len(full_attn_spans) != batch_size:
+            raise ValueError(
+                f"full_attn_spans must contain one entry per batch item; got {full_attn_spans!r} "
+                f"for batch_size={batch_size}"
+            )
+
+        hidden_states = self.language_model(
+            inputs_embeds,
+            position_ids,
+            attention_mask,
+            full_attn_spans,
+        )
+        return HiDreamO1ImageOutput(x_pred=self.final_layer2(hidden_states))
+
+
+class HiDreamO1ImageTransformer(nn.Module):
+    def __init__(
+        self,
+        config: Qwen3VLConfig,
+        *,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.config = config
+        self.quant_config = quant_config
+        self.model = HiDreamO1ImageModel(
+            config,
+            quant_config=quant_config,
+            prefix=_prefix(prefix, "model"),
+        )
+
+    @property
+    def supports_piecewise_attention(self) -> bool:
+        return all(
+            layer.self_attn.attn.attn_backend.supports_piecewise_spans for layer in self.model.language_model.layers
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        position_ids: torch.Tensor,
+        vinputs: torch.Tensor,
+        timestep: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+        full_attn_spans: list[list[tuple[int, int]]] | None = None,
+        **_: object,
+    ) -> HiDreamO1ImageOutput:
+        return self.model(
+            input_ids=input_ids,
+            position_ids=position_ids,
+            vinputs=vinputs,
+            timestep=timestep,
+            attention_mask=attention_mask,
+            full_attn_spans=full_attn_spans,
+        )
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        stacked_params_mapping = [
+            (".qkv_proj", ".q_proj", "q"),
+            (".qkv_proj", ".k_proj", "k"),
+            (".qkv_proj", ".v_proj", "v"),
+            (".gate_up_proj", ".gate_proj", 0),
+            (".gate_up_proj", ".up_proj", 1),
+        ]
+        params = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+
+        for name, loaded_weight in weights:
+            if name.startswith("model.visual.") or name == "lm_head.weight":
+                continue
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+                param = params[name]
+                param.weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                if name not in params:
+                    continue
+                param = params[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+
+        return loaded_params
+
+
+__all__ = ["HiDreamO1ImageTransformer"]

vllm_omni/diffusion/models/hidream_o1_image/pipeline_hidream_o1_image.py

@@ -0,0 +1,732 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from typing import ClassVar, TypedDict
+
+import numpy as np
+import torch
+from einops import rearrange
+from PIL import Image
+from torch import nn
+from transformers import AutoProcessor, PreTrainedTokenizerBase, Qwen3VLConfig
+from transformers.processing_utils import ProcessorMixin
+from vllm.logger import init_logger
+from vllm.model_executor.models.utils import AutoWeightsLoader
+
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.utils import get_local_device
+from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.models.hidream_o1_image.hidream_o1_image_transformer import (
+    HiDreamO1ImageTransformer,
+)
+from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
+from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler
+from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import (
+    DiffusionPipelineProfilerMixin,
+)
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+
+logger = init_logger(__name__)
+
+
+# upstream: HiDream-O1-Image models/pipeline.py L14-18 @21bcd30471ac
+PATCH_SIZE = 32
+TIMESTEP_TOKEN_NUM = 1
+NOISE_SCALE = 8.0
+T_EPS = 0.001
+
+# upstream: HiDream-O1-Image models/utils.py L7-19 @21bcd30471ac
+PREDEFINED_RESOLUTIONS: tuple[tuple[int, int], ...] = (
+    (2048, 2048),
+    (2304, 1728),
+    (1728, 2304),
+    (2560, 1440),
+    (1440, 2560),
+    (2496, 1664),
+    (1664, 2496),
+    (3104, 1312),
+    (1312, 3104),
+    (2304, 1792),
+    (1792, 2304),
+)
+
+
+class HiDreamO1TextSample(TypedDict):
+    input_ids: torch.Tensor
+    position_ids: torch.Tensor
+    attention_mask: torch.Tensor | None
+    full_attn_spans: list[list[tuple[int, int]]]
+    vinput_mask: torch.Tensor
+
+
+def find_closest_resolution(width: int, height: int) -> tuple[int, int]:
+    img_ratio = width / height
+    best_res: tuple[int, int] | None = None
+    min_diff = float("inf")
+    for w, h in PREDEFINED_RESOLUTIONS:
+        diff = abs(w / h - img_ratio)
+        if diff < min_diff:
+            min_diff = diff
+            best_res = (w, h)
+    assert best_res is not None
+    return best_res
+
+
+def get_rope_index_fix_point(
+    *,
+    spatial_merge_size: int,
+    image_token_id: int,
+    video_token_id: int,
+    vision_start_token_id: int,
+    input_ids: torch.Tensor | None = None,
+    image_grid_thw: torch.Tensor | None = None,
+    video_grid_thw: torch.Tensor | None = None,
+    attention_mask: torch.Tensor | None = None,
+    skip_vision_start_token: list[int],
+    fix_point: int = 4096,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if video_grid_thw is not None:
+        video_grid_thw = torch.repeat_interleave(video_grid_thw, video_grid_thw[:, 0], dim=0)
+        video_grid_thw[:, 0] = 1
+
+    mrope_position_deltas = []
+    if input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None):
+        total_input_ids = input_ids
+        if attention_mask is None:
+            attention_mask = torch.ones_like(total_input_ids)
+        position_ids = torch.ones(
+            3,
+            input_ids.shape[0],
+            input_ids.shape[1],
+            dtype=input_ids.dtype,
+            device=input_ids.device,
+        )
+        image_index, video_index = 0, 0
+        attention_mask = attention_mask.to(total_input_ids.device)
+        for i, input_ids in enumerate(total_input_ids):
+            input_ids = input_ids[attention_mask[i] == 1]
+            image_nums, video_nums = 0, 0
+            vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1)
+            vision_tokens = input_ids[vision_start_indices + 1]
+            image_nums = (vision_tokens == image_token_id).sum()
+            video_nums = (vision_tokens == video_token_id).sum()
+            input_tokens = input_ids.tolist()
+            llm_pos_ids_list: list[torch.Tensor] = []
+            st = 0
+            remain_images, remain_videos = image_nums, video_nums
+            for _ in range(image_nums + video_nums):
+                if image_token_id in input_tokens and remain_images > 0:
+                    ed_image = input_tokens.index(image_token_id, st)
+                else:
+                    ed_image = len(input_tokens) + 1
+                if video_token_id in input_tokens and remain_videos > 0:
+                    ed_video = input_tokens.index(video_token_id, st)
+                else:
+                    ed_video = len(input_tokens) + 1
+                if ed_image < ed_video:
+                    t, h, w = (
+                        image_grid_thw[image_index][0],
+                        image_grid_thw[image_index][1],
+                        image_grid_thw[image_index][2],
+                    )
+                    image_index += 1
+                    remain_images -= 1
+                    ed = ed_image
+                else:
+                    t, h, w = (
+                        video_grid_thw[video_index][0],
+                        video_grid_thw[video_index][1],
+                        video_grid_thw[video_index][2],
+                    )
+                    video_index += 1
+                    remain_videos -= 1
+                    ed = ed_video
+                llm_grid_t, llm_grid_h, llm_grid_w = (
+                    t.item(),
+                    h.item() // spatial_merge_size,
+                    w.item() // spatial_merge_size,
+                )
+                text_len = ed - st
+
+                text_len -= skip_vision_start_token[image_index - 1]
+                text_len = max(0, text_len)
+
+                st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+                llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+                t_index = torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten()
+                h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten()
+                w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten()
+
+                if skip_vision_start_token[image_index - 1]:
+                    if fix_point > 0:
+                        fix_point = fix_point - st_idx
+                        llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + fix_point + st_idx)
+                        fix_point = 0
+                else:
+                    llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx)
+                st = ed + llm_grid_t * llm_grid_h * llm_grid_w
+
+            if st < len(input_tokens):
+                st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+                text_len = len(input_tokens) - st
+                llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+            llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
+            position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device)
+            mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i]))
+        mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1)
+        return position_ids, mrope_position_deltas
+    else:
+        if attention_mask is not None:
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            position_ids = position_ids.unsqueeze(0).expand(3, -1, -1).to(attention_mask.device)
+            max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[0]
+            mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1]
+        else:
+            position_ids = (
+                torch.arange(input_ids.shape[1], device=input_ids.device)
+                .view(1, 1, -1)
+                .expand(3, input_ids.shape[0], -1)
+            )
+            mrope_position_deltas = torch.zeros(
+                [input_ids.shape[0], 1],
+                device=input_ids.device,
+                dtype=input_ids.dtype,
+            )
+        return position_ids, mrope_position_deltas
+
+
+def build_t2i_text_sample(
+    *,
+    prompt: str,
+    height: int,
+    width: int,
+    tokenizer: PreTrainedTokenizerBase,
+    processor: ProcessorMixin,
+    model_config: Qwen3VLConfig,
+) -> dict[str, torch.Tensor]:
+    image_token_id = model_config.image_token_id
+    video_token_id = model_config.video_token_id
+    vision_start_token_id = model_config.vision_start_token_id
+    image_len = (height // PATCH_SIZE) * (width // PATCH_SIZE)
+
+    boi_token = getattr(tokenizer, "boi_token", "<|boi_token|>")
+    tms_token = getattr(tokenizer, "tms_token", "<|tms_token|>")
+
+    messages = [{"role": "user", "content": prompt}]
+    template_caption = (
+        processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        + boi_token
+        + tms_token * TIMESTEP_TOKEN_NUM
+    )
+    input_ids = tokenizer.encode(template_caption, return_tensors="pt", add_special_tokens=False)
+
+    image_grid_thw = torch.tensor([1, height // PATCH_SIZE, width // PATCH_SIZE], dtype=torch.int64).unsqueeze(0)
+
+    vision_tokens = torch.zeros((1, image_len), dtype=input_ids.dtype) + image_token_id
+    vision_tokens[0, 0] = vision_start_token_id
+    input_ids_pad = torch.cat([input_ids, vision_tokens], dim=-1)
+
+    position_ids, _ = get_rope_index_fix_point(
+        spatial_merge_size=1,
+        image_token_id=image_token_id,
+        video_token_id=video_token_id,
+        vision_start_token_id=vision_start_token_id,
+        input_ids=input_ids_pad,
+        image_grid_thw=image_grid_thw,
+        video_grid_thw=None,
+        attention_mask=None,
+        skip_vision_start_token=[1],
+    )
+
+    txt_seq_len = input_ids.shape[-1]
+    all_seq_len = position_ids.shape[-1]
+
+    token_types = torch.zeros((1, all_seq_len), dtype=input_ids.dtype)
+    bgn = txt_seq_len - TIMESTEP_TOKEN_NUM
+    token_types[0, bgn : bgn + image_len + TIMESTEP_TOKEN_NUM] = 1
+    token_types[0, txt_seq_len - TIMESTEP_TOKEN_NUM : txt_seq_len] = 3
+
+    vinput_mask = token_types == 1
+    token_types_bin = (token_types > 0).to(token_types.dtype)
+
+    return {
+        "input_ids": input_ids,
+        "position_ids": position_ids,
+        "token_types": token_types_bin,
+        "vinput_mask": vinput_mask,
+    }
+
+
+def build_hidream_o1_attention_mask(
+    token_types: torch.Tensor,
+    *,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    """Build the mixed causal/full mask once while preparing a request.
+
+    Text rows stay causal; image and timestep rows get full attention.
+    """
+    if token_types.ndim == 1:
+        token_types = token_types.unsqueeze(0)
+    elif token_types.ndim == 2 and token_types.shape[1] == 1 and token_types.shape[0] != 1:
+        token_types = token_types.transpose(0, 1)
+    if token_types.ndim != 2:
+        raise ValueError(f"token_types must have one or two dimensions, got shape {tuple(token_types.shape)}")
+
+    batch_size, seq_len = token_types.shape
+    min_value = torch.finfo(dtype).min
+    attention_mask = torch.triu(
+        torch.full(
+            (batch_size, seq_len, seq_len),
+            min_value,
+            dtype=dtype,
+            device=token_types.device,
+        ),
+        diagonal=1,
+    )
+    attention_mask[token_types.bool()] = 0
+    return attention_mask.unsqueeze(1)
+
+
+def build_hidream_o1_full_attn_spans(
+    token_types: torch.Tensor,
+) -> list[list[tuple[int, int]]]:
+    """Encode the mixed causal/full rows without allocating an ``S x S`` mask."""
+    if token_types.ndim == 1:
+        token_types = token_types.unsqueeze(0)
+    elif token_types.ndim == 2 and token_types.shape[1] == 1 and token_types.shape[0] != 1:
+        token_types = token_types.transpose(0, 1)
+    if token_types.ndim != 2:
+        raise ValueError(f"token_types must have one or two dimensions, got shape {tuple(token_types.shape)}")
+
+    spans: list[list[tuple[int, int]]] = []
+    for sample_types in token_types:
+        full_rows = sample_types.bool().nonzero(as_tuple=False).flatten()
+        if full_rows.numel() == 0:
+            spans.append([])
+            continue
+
+        start = int(full_rows[0].item())
+        end = int(full_rows[-1].item()) + 1
+        expected = torch.arange(start, end, device=full_rows.device)
+        if not torch.equal(full_rows, expected):
+            raise ValueError("HiDream-O1 full-attention rows must form one contiguous span")
+        spans.append([(start, end)])
+    return spans
+
+
+def build_hidream_o1_scheduler(
+    *,
+    scheduler_name: str = "default",
+    num_inference_steps: int = 50,
+    shift: float = 3.0,
+    device: torch.device | str | None = None,
+    timesteps_list: list[int] | None = None,
+) -> FlowUniPCMultistepScheduler:
+    if scheduler_name != "default":
+        raise NotImplementedError(f"unsupported scheduler_name={scheduler_name!r}")
+
+    scheduler = FlowUniPCMultistepScheduler(
+        num_train_timesteps=1000,
+        shift=shift,
+        prediction_type="flow_prediction",
+        use_dynamic_shifting=False,
+    )
+    scheduler.set_timesteps(num_inference_steps, device=device)
+
+    if timesteps_list is not None:
+        scheduler.timesteps = torch.tensor(timesteps_list, device=device, dtype=torch.long)
+        sigmas = [t.item() / 1000.0 for t in scheduler.timesteps]
+        sigmas.append(0.0)
+        scheduler.sigmas = torch.tensor(sigmas, device=device)
+
+    return scheduler
+
+
+def get_hidream_o1_image_post_process_func(od_config: OmniDiffusionConfig):
+    del od_config
+
+    def post_process_func(
+        output_data: tuple[torch.Tensor, int, int],
+    ) -> Image.Image:
+        """Convert HiDream-O1 patch-space output to a PIL RGB image."""
+        if not isinstance(output_data, tuple) or len(output_data) != 3:
+            raise TypeError(
+                f"HiDreamO1 postprocess expects (patch_tensor, height, width), got {type(output_data).__name__}"
+            )
+
+        z, height, width = output_data
+
+        if not torch.is_tensor(z):
+            raise TypeError(
+                f"HiDreamO1 postprocess expects the first tuple element to be a torch.Tensor, got {type(z).__name__}"
+            )
+
+        height = int(height)
+        width = int(width)
+
+        if height <= 0 or width <= 0:
+            raise ValueError(
+                f"HiDreamO1 postprocess requires positive height/width, got height={height}, width={width}"
+            )
+
+        if height % PATCH_SIZE != 0 or width % PATCH_SIZE != 0:
+            raise ValueError(
+                f"HiDreamO1 output resolution must be divisible by PATCH_SIZE={PATCH_SIZE}, got {width}x{height}"
+            )
+
+        h_patches = height // PATCH_SIZE
+        w_patches = width // PATCH_SIZE
+        expected_image_len = h_patches * w_patches
+        expected_patch_dim = 3 * PATCH_SIZE * PATCH_SIZE
+
+        if z.ndim != 3:
+            raise ValueError(
+                f"HiDreamO1 patch tensor must have shape (B, image_len, patch_dim), got shape={tuple(z.shape)}"
+            )
+
+        if z.shape[0] != 1:
+            raise ValueError(f"HiDreamO1 postprocess currently expects batch size 1, got batch_size={z.shape[0]}")
+
+        if z.shape[1] != expected_image_len:
+            raise ValueError(
+                "HiDreamO1 patch count mismatch: "
+                f"got {z.shape[1]}, expected {expected_image_len} "
+                f"for resolution {width}x{height}"
+            )
+
+        if z.shape[2] != expected_patch_dim:
+            raise ValueError(f"HiDreamO1 patch dimension mismatch: got {z.shape[2]}, expected {expected_patch_dim}")
+
+        # Preserve upstream decode order exactly:
+        #   1. map patch values from [-1, 1] to [0, 1]
+        #   2. restore BCHW image layout
+        #   3. convert to float32 on CPU
+        img = (z + 1) / 2
+
+        img = rearrange(
+            img.cpu().float(),
+            "B (H W) (C p1 p2) -> B C (H p1) (W p2)",
+            H=h_patches,
+            W=w_patches,
+            p1=PATCH_SIZE,
+            p2=PATCH_SIZE,
+        )
+
+        arr = np.round(
+            np.clip(
+                img[0].numpy().transpose(1, 2, 0) * 255,
+                0,
+                255,
+            )
+        ).astype(np.uint8)
+
+        return Image.fromarray(arr).convert("RGB")
+
+    return post_process_func
+
+
+class HiDreamO1ImagePipeline(nn.Module, DiffusionPipelineProfilerMixin, ProgressBarMixin):
+    supports_request_batch = False
+    dummy_run_num_frames: ClassVar[int] = 0
+    _dit_modules: ClassVar[list[str]] = ["model.model.language_model"]
+
+    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
+        super().__init__()
+        self.od_config = od_config
+        self.prefix = prefix
+        self.model_dir = od_config.model
+        self.dtype = od_config.dtype if od_config.dtype is not None else torch.bfloat16
+        self.device = get_local_device()
+        self.weights_sources = [
+            DiffusersPipelineLoader.ComponentSource(
+                model_or_path=od_config.model,
+                subfolder=None,
+                revision=od_config.revision,
+                prefix="model.",
+                fall_back_to_pt=True,
+            )
+        ]
+
+        self._validate_static_config()
+
+        self.setup_diffusion_pipeline_profiler(
+            profiler_targets=["forward"],
+            enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler,
+        )
+
+        self._init_processor_and_model()
+        self._validate_tms_token_id()
+
+    def _validate_static_config(self) -> None:
+        if self.dtype != torch.bfloat16:
+            raise NotImplementedError(f"only bfloat16 is supported; got {self.dtype!r}")
+        cfg_parallel_size = int(self.od_config.parallel_config.cfg_parallel_size)
+        if cfg_parallel_size > 1:
+            raise NotImplementedError(f"cfg_parallel_size={cfg_parallel_size} not supported")
+
+    def _init_processor_and_model(self) -> None:
+        logger.info(
+            "HiDreamO1ImagePipeline: loading processor and model config from %s",
+            self.model_dir,
+        )
+        self.processor = AutoProcessor.from_pretrained(
+            self.model_dir,
+            revision=self.od_config.revision,
+        )
+        model_config = Qwen3VLConfig.from_pretrained(
+            self.model_dir,
+            revision=self.od_config.revision,
+        )
+        self.model = HiDreamO1ImageTransformer(
+            model_config,
+            quant_config=self.od_config.quantization_config,
+            prefix=self.prefix,
+        )
+
+        self._add_special_tokens(self.processor)
+        self.tokenizer = (
+            self.processor if isinstance(self.processor, PreTrainedTokenizerBase) else self.processor.tokenizer
+        )
+
+    def _validate_tms_token_id(self) -> None:
+        actual_tms_ids = self.tokenizer.encode(self.tokenizer.tms_token, add_special_tokens=False)
+        expected_tms_id = int(self.model.model.tms_token_id)
+        if actual_tms_ids != [expected_tms_id]:
+            raise RuntimeError(f"tms_token maps to {actual_tms_ids}, model expects [{expected_tms_id}]")
+
+    @staticmethod
+    def _add_special_tokens(processor: ProcessorMixin | PreTrainedTokenizerBase) -> None:
+        tok = processor if isinstance(processor, PreTrainedTokenizerBase) else processor.tokenizer
+        tok.boi_token = "<|boi_token|>"
+        tok.bor_token = "<|bor_token|>"
+        tok.eor_token = "<|eor_token|>"
+        tok.bot_token = "<|bot_token|>"
+        tok.tms_token = "<|tms_token|>"
+
+    def _resolve_generation_params(
+        self,
+        req: DiffusionRequestBatch,
+    ) -> tuple[str, int, int, int, int, float]:
+        sp = req.sampling_params
+
+        if sp.timesteps is not None:
+            raise NotImplementedError("sp.timesteps not supported")
+        if sp.latents is not None:
+            raise NotImplementedError("sp.latents not supported")
+        if sp.num_outputs_per_prompt != 1:
+            raise NotImplementedError(f"num_outputs_per_prompt={sp.num_outputs_per_prompt} not supported")
+
+        if not req.prompts:
+            raise ValueError("req.prompts is empty")
+        if len(req.prompts) != 1:
+            raise NotImplementedError(f"only one prompt per request; got {len(req.prompts)}")
+        first_prompt = req.prompts[0]
+        if isinstance(first_prompt, str):
+            prompt = first_prompt
+        elif isinstance(first_prompt, dict):
+            raw_prompt = first_prompt.get("prompt")
+            if not isinstance(raw_prompt, str):
+                raise TypeError(f"dict prompt must contain a str 'prompt' field; got {type(raw_prompt).__name__}")
+            prompt = raw_prompt
+        else:
+            raise TypeError(f"prompt must be str or dict; got {type(first_prompt).__name__}")
+
+        raw_h = 2048 if sp.height is None else int(sp.height)
+        raw_w = 2048 if sp.width is None else int(sp.width)
+        if raw_h <= 0 or raw_w <= 0:
+            raise ValueError(f"h/w must be positive, got h={raw_h} w={raw_w}")
+        snapped_w, snapped_h = find_closest_resolution(raw_w, raw_h)
+        if (snapped_w, snapped_h) != (raw_w, raw_h):
+            logger.info(
+                "HiDreamO1: resolution snapped from %dx%d to %dx%d",
+                raw_w,
+                raw_h,
+                snapped_w,
+                snapped_h,
+            )
+
+        steps = 50 if sp.num_inference_steps is None else int(sp.num_inference_steps)
+        if steps <= 0:
+            raise ValueError(f"num_inference_steps must be positive, got {steps}")
+
+        generator = sp.generator
+        if isinstance(generator, list):
+            if len(generator) != 1:
+                raise NotImplementedError("generator list only supported for single-output requests with one generator")
+            generator = generator[0]
+
+        generator_seed: int | None = None
+        if generator is not None:
+            if not hasattr(generator, "initial_seed"):
+                raise TypeError(f"sp.generator must expose initial_seed(); got {type(generator).__name__}")
+            generator_seed = int(generator.initial_seed())
+
+        if sp.seed is not None and generator_seed is not None and int(sp.seed) != generator_seed:
+            raise ValueError(f"seed/generator mismatch: seed={int(sp.seed)} generator.initial_seed()={generator_seed}")
+
+        if sp.seed is None and generator_seed is None:
+            raise RuntimeError("expected request initialization to resolve a concrete seed")
+        seed = int(sp.seed) if sp.seed is not None else generator_seed
+        assert seed is not None
+        guidance_scale = float(sp.guidance_scale)
+        return prompt, snapped_h, snapped_w, steps, seed, guidance_scale
+
+    def _prepare_noise_and_patchify(
+        self,
+        height: int,
+        width: int,
+        seed: int,
+        dtype: torch.dtype,
+        device: torch.device,
+    ) -> torch.Tensor:
+        gen = torch.Generator("cpu").manual_seed(seed + 1)
+        noise = NOISE_SCALE * torch.randn((1, 3, height, width), generator=gen)
+        noise = noise.to(device, dtype)
+        return rearrange(
+            noise,
+            "B C (H p1) (W p2) -> B (H W) (C p1 p2)",
+            p1=PATCH_SIZE,
+            p2=PATCH_SIZE,
+        )
+
+    def _forward_once(
+        self,
+        sample: HiDreamO1TextSample,
+        z_in: torch.Tensor,
+        t_pixeldit: torch.Tensor,
+    ) -> torch.Tensor:
+        with torch.autocast(self.device.type, dtype=self.dtype, cache_enabled=False):
+            outputs = self.model(
+                input_ids=sample["input_ids"],
+                position_ids=sample["position_ids"],
+                vinputs=z_in,
+                timestep=t_pixeldit.reshape(-1).to(self.device),
+                attention_mask=sample["attention_mask"],
+                full_attn_spans=sample["full_attn_spans"],
+            )
+
+        if outputs.x_pred is None:
+            raise RuntimeError("model returned x_pred=None")
+        if outputs.x_pred.shape[-1] != z_in.shape[-1]:
+            raise RuntimeError(f"x_pred patch_dim {outputs.x_pred.shape[-1]} != z_in {z_in.shape[-1]}")
+
+        mask = sample["vinput_mask"][0]
+        return outputs.x_pred[0, mask].unsqueeze(0)
+
+    def _build_and_validate_sample(
+        self,
+        prompt: str,
+        height: int,
+        width: int,
+        expected_image_len: int,
+    ) -> HiDreamO1TextSample:
+        sample = build_t2i_text_sample(
+            prompt=prompt,
+            height=height,
+            width=width,
+            tokenizer=self.tokenizer,
+            processor=self.processor,
+            model_config=self.model.config,
+        )
+        actual_image_len = int(sample["vinput_mask"].sum().item())
+        if actual_image_len != expected_image_len:
+            raise RuntimeError(
+                f"vinput_mask has {actual_image_len} slots, expected {expected_image_len} for {height}x{width}"
+            )
+        full_attn_spans = build_hidream_o1_full_attn_spans(sample["token_types"])
+        inputs = {key: value.to(self.device) for key, value in sample.items()}
+        attention_mask = None
+        if not self.model.supports_piecewise_attention:
+            attention_mask = build_hidream_o1_attention_mask(
+                inputs["token_types"],
+                dtype=self.dtype,
+            )
+        return {
+            "input_ids": inputs["input_ids"],
+            "position_ids": inputs["position_ids"],
+            "attention_mask": attention_mask,
+            "full_attn_spans": full_attn_spans,
+            "vinput_mask": inputs["vinput_mask"],
+        }
+
+    @torch.inference_mode()
+    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
+        prompt, height, width, num_inference_steps, seed, guidance_scale = self._resolve_generation_params(req)
+
+        expected_image_len = (height // PATCH_SIZE) * (width // PATCH_SIZE)
+        cond_sample = self._build_and_validate_sample(
+            prompt,
+            height,
+            width,
+            expected_image_len,
+        )
+        samples: list[HiDreamO1TextSample] = [cond_sample]
+        if guidance_scale > 1.0:
+            uncond_sample = self._build_and_validate_sample(
+                " ",
+                height,
+                width,
+                expected_image_len,
+            )
+            samples.append(uncond_sample)
+
+        z = self._prepare_noise_and_patchify(
+            height,
+            width,
+            seed,
+            self.dtype,
+            self.device,
+        )
+
+        sched = build_hidream_o1_scheduler(
+            scheduler_name="default",
+            num_inference_steps=num_inference_steps,
+            shift=3.0,
+            device=self.device,
+        )
+
+        with self.progress_bar(total=len(sched.timesteps)) as pbar:
+            for step_t in sched.timesteps:
+                step_t_f32 = step_t.float()
+                t_pixeldit = 1.0 - step_t_f32 / 1000.0
+                sigma = (step_t_f32 / 1000.0).to(dtype=torch.float32).clamp_min(T_EPS)
+
+                x_pred_cond = self._forward_once(samples[0], z, t_pixeldit)
+                v_cond = (x_pred_cond.to(torch.float32) - z.to(torch.float32)) / sigma
+
+                if len(samples) > 1:
+                    x_pred_uncond = self._forward_once(
+                        samples[1],
+                        z,
+                        t_pixeldit,
+                    )
+                    v_uncond = (x_pred_uncond.to(torch.float32) - z.to(torch.float32)) / sigma
+                    v_guided = v_uncond + guidance_scale * (v_cond - v_uncond)
+                else:
+                    v_guided = v_cond
+
+                model_output = -v_guided
+
+                z = sched.step(
+                    model_output,
+                    step_t_f32,
+                    z.float(),
+                    return_dict=False,
+                )[0].to(self.dtype)
+                pbar.update()
+
+        return DiffusionOutput(
+            output=(z, height, width),
+            stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
+        )
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights)

vllm_omni/diffusion/registry.py

@@ -336,6 +336,11 @@ _DIFFUSION_MODELS = {
         "pipeline_hidream_image",
         "HiDreamImagePipeline",
     ),
+    "HiDreamO1ImagePipeline": (
+        "hidream_o1_image",
+        "pipeline_hidream_o1_image",
+        "HiDreamO1ImagePipeline",
+    ),
     "DreamZeroPipeline": (
         "dreamzero",
         "pipeline_dreamzero",
@@ -600,6 +605,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "Cosmos3OmniDiffusersPipeline": "get_cosmos3_post_process_func",
     "Cosmos3OmniPipeline": "get_cosmos3_post_process_func",
     "HiDreamImagePipeline": "get_hidream_image_post_process_func",
+    "HiDreamO1ImagePipeline": "get_hidream_o1_image_post_process_func",
     "StableDiffusionXLPipeline": "get_sdxl_image_post_process_func",
     "Krea2Pipeline": "get_krea2_post_process_func",
     "HunyuanImage3ForCausalMM": "get_hunyuan_image3_post_process_func",

vllm_omni/diffusion/utils/hf_utils.py

@@ -71,6 +71,30 @@ def _looks_like_dreamzero(model_name: str) -> bool:
         return False
 
 
+HIDREAM_O1_SIGNATURE_WEIGHTS = (
+    "model.final_layer2.linear.weight",
+    "final_layer2.linear.weight",
+)
+
+
+def _looks_like_hidream_o1(model_name: str, config: Mapping | None = None) -> bool:
+    """Detect HiDream-O1 without matching regular Qwen3-VL checkpoints."""
+    try:
+        cfg = config if config is not None else get_hf_file_to_dict("config.json", model_name)
+        if not isinstance(cfg, Mapping) or cfg.get("model_type") != "qwen3_vl":
+            return False
+
+        index = get_hf_file_to_dict("model.safetensors.index.json", model_name)
+        if not isinstance(index, Mapping):
+            return False
+        weight_map = index.get("weight_map")
+        if not isinstance(weight_map, Mapping):
+            return False
+        return any(key in weight_map for key in HIDREAM_O1_SIGNATURE_WEIGHTS)
+    except (AttributeError, OSError, TypeError, ValueError):
+        return False
+
+
 @lru_cache
 def is_diffusion_model(model_name: str) -> bool:
     """Check if a model is a diffusion model.
@@ -115,7 +139,4 @@ def is_diffusion_model(model_name: str) -> bool:
     except Exception as e:
         logger.debug("Failed to load diffusers config via DiffusionPipeline: %s", e)
 
-        # Bagel and DreamZero are not diffusers pipelines (no model_index.json),
-        # but are still diffusion-style models in vllm-omni. Detect them via
-        # config.json.
-    return _looks_like_bagel(model_name) or _looks_like_dreamzero(model_name)
+    return _looks_like_bagel(model_name) or _looks_like_dreamzero(model_name) or _looks_like_hidream_o1(model_name)