[SOKOUDJOU-LEOPOLD]

Purpose

This PR adds support for HiDream-O1-Image (HiDream-ai/HiDream-O1-Image and HiDream-ai/HiDream-O1-Image-Dev), a ~8B-parameter image generation model from HiDream.ai. It supports text-to-image, instruction-based image editing, multi-reference subject personalization, and layout-bbox spatial conditioning — all in a single model with no VAE and no separate text encoder.

Why

HiDream-O1-Image has an architecture that nothing in vLLM-Omni currently handles:

• No VAE. The model diffuses directly on raw 32×32 pixel patches. Output is assembled from final hidden states with no latent encoder/decoder.
• No separate text encoder. A single Qwen3-VL backbone handles everything — causal attention over text for conditioning, bidirectional attention over image tokens for the diffusion pass. This is called a Pixel-level Unified Transformer (UiT).
• Timestep conditioning via token substitution, not AdaLN. A special <|tms_token|> placeholder has its embedding swapped for the sinusoidal timestep embedding at every step.
• Mixed causal/bidirectional attention in one forward call. Text tokens are causal; image tokens attend bidirectionally over the full sequence. Expressed through AttentionMetadata + full_attn_spans.
• 3D MRoPE with fix_point anchoring. Custom get_rope_index_fix_point pins reference image positions at a fixed offset to avoid RoPE extrapolation.

What

New files — vllm_omni/diffusion/models/hidream_o1_image/ (4 files):

• qwen3_vl_uit_transformer.py — HiDreamO1UiTModel backbone. Lazy-imports HF Qwen3-VL (gated on transformers >= 4.57.1). Overrides attention/MLP projections with TP-capable parallel linears via _build_attention_subclasses(quant_config) closure. Carries _cache_dit_adapter_config, _sp_plan, _hsdp_shard_conditions, _layerwise_offload_blocks_attrs.
• pipeline_hidream_o1_image.py — HiDreamO1ImagePipeline denoise loop. Handles T2I, editing, personalization, layout-bbox. Implements SupportsComponentDiscovery (empty _vae_modules, nested _encoder_modules = ["transformer.visual"]), DiffusionPipelineProfilerMixin, CFGParallelMixin.
• utils_hidream_o1.py — pure utilities: build_packed_attention_metadata (builds both full_attn_spans and dense attn_mask), get_rope_index_fix_point, patchify/depatchify, layout-bbox rendering.
• __init__.py — module exports.

Parallelism — all wired in the model files above:

Strategy	Implementation
Tensor Parallel	QKVParallelLinear (fused Q/K/V), MergedColumnParallelLinear (gate+up), RowParallelLinear (o_proj, down_proj). stacked_params_mapping in load_weights() remaps split checkpoint weights into fused params.
Sequence Parallel (Ulysses + Ring)	_sp_plan splits inputs_embeds/position_ids before decoder stack, gathers final_layer2 after. attn_metadata excluded — Ulysses AllToAll restores full sequence before attention. T2I only; editing path documented as known SP limitation.
CFG-Parallel	Zero code — inherited from CFGParallelMixin.
HSDP	_hsdp_shard_conditions = [_is_transformer_block] — each Qwen3VLTextDecoderLayer is a shard unit.
Cache-DiT	_cache_dit_adapter_config with ForwardPattern.Pattern_0, has_separate_cfg=True, check_forward_pattern=False. Added layers property alias on HiDreamO1UiTModel pointing to language_model.layers for cache-dit auto-detection.
CPU Offload	SupportsComponentDiscovery. _vae_modules = [] (no VAE). _encoder_modules = ["transformer.visual"] — first nested encoder path in the codebase, resolved via operator.attrgetter.
Layerwise Offload	_layerwise_offload_blocks_attrs = ["language_model.layers"]

Integration:

• vllm_omni/diffusion/registry.py — registered with get_hidream_o1_image_post_process_func ([-1, 1] → uint8, no VAE decode).
• docs/models/supported_models.md, docs/user_guide/diffusion_features.md, docs/user_guide/diffusion/cpu_offload_diffusion.md — rows added.
• .buildkite/cuda/test-ready.yml, test-merge.yml, test-nightly.yml — CI entries added.
• examples/offline_inference/hidream_o1_image/, examples/online_serving/hidream_o1_image/ — runnable examples with READMEs.
• recipes/HiDream/HiDream-O1-Image.md — full serving recipe with two-bucket performance methodology.

Test Plan

Unit tests (47 tests, CPU-only, pytest -m core_model):

File	Coverage
tests/diffusion/models/hidream_o1_image/test_packed_sequence.py	build_packed_attention_metadata for T2I, editing, multi-ref, layout-bbox. Verifies full_attn_spans coordinates, attn_mask, token types, causal vs bidirectional pattern.
tests/diffusion/models/hidream_o1_image/test_embed_cache.py	Vision-tower called exactly once per request across all denoising steps.
tests/diffusion/models/hidream_o1_image/test_layout_bbox_utils.py	Bbox parsing, layout image rendering, coordinate clamping, batch handling.

E2E offline (tests/e2e/offline_inference/test_hidream_o1_image.py): full pipeline forward pass, output shape and pixel range verified.

E2E online (tests/e2e/online_serving/test_hidream_o1_image.py): T2I (Dev + Full variants), image editing, multi-reference personalization, layout-bbox.

Nightly expansion (tests/e2e/online_serving/test_hidream_o1_image_expansion.py): 8 parametrized server configs:

id	flags	GPUs
default	—	1× H100
cache_dit	--cache-backend cache_dit	1× H100
tp2	--tensor-parallel-size 2	2× H100
ulysses2	--ulysses-degree 2	2× H100
cfg_parallel2	--cfg-parallel-size 2	2× H100
hsdp	--use-hsdp	1× H100
cpu_offload	--enable-cpu-offload	1× H100
layerwise_offload	--enable-layerwise-offload	1× H100

vLLM Version: 0.23.0

vLLM-Omni Commit: 1ff2f22

Test Result

Unit tests (47): all passed on CPU.

E2E and expansion tests require H100 hardware, tracked through CI (test-merge and test-nightly pipelines).

Pixel-level output comparison against the HiDream reference repo is explicitly deferred and tracked in #5304 .
H100 benchmark numbers are explicitly deferred and tracked in #5305 .

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[Copilot]

Pull request overview

Adds first-class vLLM-Omni diffusion support for HiDream-O1-Image (Qwen3-VL–based Pixel-level Unified Transformer) including text-to-image, reference-image editing/personalization, and layout-bbox conditioning, plus integration into registry/docs/CI and a test+example suite.

Changes:

• Introduces a new HiDream-O1-Image model implementation (UiT backbone + diffusion pipeline) and registers it in the diffusion registry/post-processing map.
• Adds utilities for packed attention metadata, MRoPE position id construction, patchify/depatchify, and layout-bbox rendering.
• Adds unit + E2E tests, Buildkite coverage, docs updates, and offline/online runnable examples + a serving recipe.

Reviewed changes

Copilot reviewed 26 out of 27 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
vllm_omni/diffusion/registry.py	Registers HiDreamO1ImagePipeline and its post-process function.
vllm_omni/diffusion/models/hidream_o1_image/utils_hidream_o1.py	Adds core tensor/PIL utilities (masking, MRoPE indices, patch ops, layout-bbox helpers).
vllm_omni/diffusion/models/hidream_o1_image/qwen3_vl_uit_transformer.py	Implements the Qwen3-VL–based UiT backbone with vLLM-Omni attention + TP-capable linears.
vllm_omni/diffusion/models/hidream_o1_image/pipeline_hidream_o1_image.py	Implements denoising pipeline for T2I/editing/personalization/layout-bbox + weight loading and caching.
vllm_omni/diffusion/models/hidream_o1_image/init.py	Exposes lazy imports for pipeline and post-processing function.
tests/e2e/online_serving/test_hidream_o1_image.py	Adds L2/L3 online-serving smoke tests for T2I, editing, and multi-ref.
tests/e2e/online_serving/test_hidream_o1_image_expansion.py	Adds L4 nightly expansion matrix (cache_dit/tp/sp/cfg-parallel/hsdp/offload).
tests/e2e/offline_inference/test_hidream_o1_image.py	Adds offline inference smoke test.
tests/diffusion/models/hidream_o1_image/test_packed_sequence.py	Unit tests for packed sequence masking, MRoPE, patchify/depatchify, ref preprocessing.
tests/diffusion/models/hidream_o1_image/test_layout_bbox_utils.py	Unit tests for bbox parsing/rendering and ref/layout image generation.
tests/diffusion/models/hidream_o1_image/test_embed_cache.py	Unit tests for vision-embed caching contract.
tests/diffusion/models/hidream_o1_image/init.py	Package marker for hidream_o1_image diffusion tests.
recipes/HiDream/HiDream-O1-Image.md	Adds a full serving recipe and profiling methodology placeholders.
examples/online_serving/hidream_o1_image/run_server.sh	Adds a convenience script to start the server.
examples/online_serving/hidream_o1_image/run_curl_hidream_o1.sh	Adds a curl example for /v1/images/generations.
examples/online_serving/hidream_o1_image/README.md	Documents online-serving usage for HiDream-O1-Image.
examples/online_serving/hidream_o1_image/openai_client_hidream_o1.py	Adds a simple OpenAI-compatible Python client example.
examples/offline_inference/hidream_o1_image/text_to_image_hidream_o1.py	Adds an offline T2I driver with full/dev presets.
examples/offline_inference/hidream_o1_image/README.md	Documents offline scripts and arguments for T2I/edit/layout control.
examples/offline_inference/hidream_o1_image/layout_control_hidream_o1.py	Adds an offline layout-bbox controlled generation script.
examples/offline_inference/hidream_o1_image/image_edit_hidream_o1.py	Adds an offline editing/multi-ref personalization script.
docs/user_guide/diffusion/cpu_offload_diffusion.md	Documents CPU/layerwise offload support for the new pipeline.
docs/user_guide/diffusion_features.md	Updates diffusion feature support matrix to include HiDream-O1-Image.
docs/models/supported_models.md	Adds HiDream-O1-Image to supported models list.
.buildkite/cuda/test-ready.yml	Adds L2 test-ready step for HiDream-O1-Image serving.
.buildkite/cuda/test-nightly.yml	Adds nightly expansion test entry for HiDream-O1-Image.
.buildkite/cuda/test-merge.yml	Adds L3 test-merge step for HiDream-O1-Image serving.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[SOKOUDJOU-LEOPOLD]

Invalid.

It would add a redundant operation (dense_mask[b, ~gen] &= (~gen).unsqueeze(0)) that zeroes bits already guaranteed to be zero by construction. The docstring phrase "causal-only-over-text" is accurate as is, text tokens already only see text columns because the layout puts all gen tokens at the tail of every sequence.

[SOKOUDJOU-LEOPOLD]

Valid.
The docstring was written at the end of Phase 3 and still says "single GPU, no Cache-DiT / TP / SP yet", but Phases 5–8 added all of those. It's now stale and misleading.

[SOKOUDJOU-LEOPOLD]

Valid.

• Lines 388–394 build ref_pixel_grid_thw but with a subtle bug: it calls resize_pilimage(pil, ref_max) inside the comprehension expression (for _ in [resize_pilimage(...)]) but then reads pil.size (the original unresized dimensions), not the resized ones. This produces wrong grid values AND does a wasted resize call.
• Lines 396–401 immediately overwrite ref_pixel_grid_thw correctly: it resizes, captures pil_r.size, and uses rw, rh.

[SOKOUDJOU-LEOPOLD]

valid

[SOKOUDJOU-LEOPOLD]

Valid.
load_layout_bboxes calls os.path.exists() on any user-supplied string in online serving a remote client controls multi_modal_data["layout_bboxes"], so they can probe server-local file existence and read arbitrary JSON files from disk.
To fix:

• use json.loads() directly, only accept inline JSON, never read files from user data.
• preload the file with load_layout_bboxes() before putting data in multi_modal_data, since that's operator-controlled code anyway.

[SOKOUDJOU-LEOPOLD]

Hi @hsliuustc0106 , Please can you review my PR?

[hsliuustc0106]

Joint review with #5194. These are comment-only findings; no approval or change-request verdict.

[hsliuustc0106]

[P1] Add automatic model-class resolution before relying on this command. The official checkpoint config declares Qwen3VLForConditionalGeneration, but the resolver returns that unregistered architecture while the registry contains HiDreamO1ImagePipeline, so vllm serve "$MODEL" --omni will fail unless the caller supplies --model-class-name. Please add the HiDream signature mapping to both resolver/enrichment paths; carrying #5194's diffusion detector alone is not sufficient because detection does not select the pipeline class.

[hsliuustc0106]

[P1] Guard this combination before advertising sequence parallelism for the full pipeline. With reference images, visual_pos_masks remains full-length while _sp_plan splits only inputs_embeds and position_ids; _deepstack_process() then indexes local hidden states with the full mask. Edit/personalization/layout with --ulysses-degree > 1 can therefore shape-error or produce incorrect deepstack embeddings. Please either implement mask-aware split/gather or reject this mode explicitly and add a reference-image SP test.

[hsliuustc0106]

[P1] Apply the documented Dev sampling defaults to this online path. The server defaults to the Dev variant and this request sends 28 steps, but omits guidance and shift; the pipeline consequently substitutes guidance 5.0 and shift 3.0, enabling CFG and using the full-model shift instead of the Dev preset (1.0/1.0). Please make defaults model-variant-aware or include and enforce the Dev parameters in the server example.

[hsliuustc0106]

[P2] Refresh the vision-embedding cache before constructing the negative CFG branch. cached_img_emb and cached_ds_emb are captured before the positive branch populates embed_storage; the negative branch still receives those stale None values and _embed_storage=None, so reference-image CFG recomputes the vision tower on the first denoise step. Please refresh the cached values after the positive call and add an actual pipeline call-count test.

[vllm-omni-review-bot]

Omni ReviewBot: no human activity for 29 days

@SOKOUDJOU-LEOPOLD this pull request has had no human commit, comment or review since 2026-08-10. Per repository policy it may be closed if it stays inactive.

To keep it moving, any one of these is enough: push an update, reply to the open blocker, or post the current plan and timeline.

[vllm-omni-review-bot]

Omni ReviewBot: no human activity for 44 days

@SOKOUDJOU-LEOPOLD this pull request has had no human commit, comment or review since 2026-08-10. Please consider closing this PR if the work is no longer planned. The author or a maintainer decides whether to change the PR state.

To keep it moving, any one of these is enough: push an update, reply to the open blocker, or post the current plan and timeline.

[vllm-omni-review-bot]

Omni ReviewBot routing record

Assigned Strict on zcode (GLM-5.3-Flash) under experiment fleet-strict-cursor-grok46-zcode-glm53flash-5050-c5-z10-20261002.

[vllm-omni-review-bot]

Omni ReviewBot attempt record

Review attempt ended as timeout (Strict attempt outlived its budget; falling back to direct/cursor/auto).

[vllm-omni-review-bot]

Omni ReviewBot review

Changes since the previous review

• 4 new inline finding(s); 0 finding(s) below.

CI at 43db9cd47b29 (2026-10-09T16:29:35.281736+00:00): required check(s) blocking: buildkite/vllm-omni (missing).

Note: The assigned review arm strict/zcode/GLM-5.3-Flash could not complete this review, so it was produced by the fallback arm direct/cursor/auto. It is excluded from the routing experiment.

Full review analysis

PR description

This PR adds HiDream-O1-Image (HiDreamO1ImagePipeline / HiDreamO1UiTModel): a Qwen3-VL pixel-space diffusion pipeline with no VAE, covering text-to-image, reference editing, multi-reference personalization, and layout-bbox conditioning. Denoising patches 32×32 pixels, swaps a <|tms_token|> embedding for the timestep, and registers a [-1, 1] to RGB post-process. Docs, examples, and Buildkite steps advertise Cache-DiT, TP, Ulysses, CFG-parallel, HSDP, and CPU/layerwise offload. The online serve script still does not select HiDreamO1ImagePipeline, and edit-plus-sequence-parallel plus the Dev sampling defaults are not actually enforced.

Change flow

flowchart LR
  req["[EXISTING] OmniDiffusionRequest"]:::existing
  reg["[CHANGED] diffusion registry"]:::changed
  pipe["[NEW] HiDreamO1ImagePipeline"]:::new
  uit["[NEW] HiDreamO1UiTModel"]:::new
  rgb["[NEW] patch RGB post-process"]:::new
  req --> reg --> pipe --> uit --> rgb
classDef existing fill:#e5e7eb,stroke:#6b7280,color:#111827
classDef changed fill:#fef3c7,stroke:#d97706,color:#451a03,stroke-width:2px
classDef new fill:#dcfce7,stroke:#16a34a,color:#052e16,stroke-width:2px
classDef removed fill:#fee2e2,stroke:#dc2626,color:#450a0a,stroke-width:2px

Loading

See inline comments below.

---

_{🤖 This review was generated by InferMatrix Copilot, an open-source repo-maintenance agent for PR review, CI debugging and issue triage. Try it on your own repo, and ⭐ star it if it helped!}

[vllm-omni-review-bot]

[P1] Edit vision tower receives target and ref grids

Evidence and suggested fix

_build_edit_sample returns pixel_values from the VLM processor (K reference images) but sets image_grid_thw to image_grid_thw_all, which concatenates those VLM grids with the target patch grid and the pixel-ref grids. forward_generation passes that tensor straight into get_image_features. Qwen3-VL splits pixel_values by the product of grid_thw, so the extra rows make the vision-tower length disagree with the image-token mask and masked_scatter. Every edit, personalization, and layout request can throw or scatter the wrong embeddings. Pass only image_grid_thw_vlm into the vision tower; keep the concatenated grid for get_rope_index_fix_point.

[vllm-omni-review-bot]

[P1] Two-GPU HiDream rows are only on the 1-GPU nightly shard

Evidence and suggested fix

test_hidream_o1_image_expansion.py is added only to the H100 single-GPU step (.buildkite/cuda/test-nightly.yml around line 213), whose command is -m "full_model and diffusion and H100 and not distributed_cuda" and mirror_hardwares: h100_1. The tp2, ulysses2, and cfg_parallel2 params request 2 GPUs. They are either deselected by not distributed_cuda or executed on one GPU. No h100_2 step lists this file (the multi-GPU steps cover Qwen-Image, layered, and Bagel only). The PR test plan's 2×H100 rows are not what nightly runs. Put those params on a distributed_cuda / h100_2 step.

[vllm-omni-review-bot]

[P2] Merge CI adds unrelated Longcat and SenseNova jobs

Evidence and suggested fix

The HiDream test-merge hunk also inserts "Diffusion · Longcat Image Edit Test" and "Diffusion · Longcat Image Test", whose sources are tests/e2e/online_serving/test_longcat_image*.py and vllm_omni/diffusion/models/longcat_image/. test-nightly.yml line 228 likewise adds tests/e2e/offline_inference/test_sensenova_u1_expansion.py. Neither model is part of HiDream-O1. Those jobs add H100 time and can fail this PR for unrelated tests. Drop them from this diff.

[vllm-omni-review-bot]

[P2] Non-test diff exceeds the size budget

Evidence and suggested fix

The supplied diff size is 2592 non-test added lines against a budget of 1000 (14 counted files, 13 exempt). The largest contributors are qwen3_vl_uit_transformer.py (694), pipeline_hidream_o1_image.py (679), and utils_hidream_o1.py (592). The PR body describes the new model and deferred pixel/benchmark follow-ups (#5304, #5305) but does not give a split plan or a concrete size exemption. Split the backbone, pipeline, and examples/CI into reviewable PRs, or state an explicit exemption for this new-model drop. Size is not itself a correctness defect.

[vllm-omni-review-bot]

Omni ReviewBot: finding feedback

[p1] Reference-image sequence parallel is not rejected — vllm_omni/diffusion/models/hidream_o1_image/qwen3_vl_uit_transformer.py:500

See the review for details. If you are the PR author and disagree, react 👎 here; the maintainer will see your disagreement.

[vllm-omni-review-bot]

Omni ReviewBot: finding feedback

[p1] Two-GPU HiDream rows are only on the 1-GPU nightly shard — .buildkite/cuda/test-nightly.yml:229

See the review for details. If you are the PR author and disagree, react 👎 here; the maintainer will see your disagreement.

[vllm-omni-review-bot]

Omni ReviewBot: finding feedback

[p1] Dev examples and nightly tests use full CFG defaults — examples/online_serving/hidream_o1_image/run_curl_hidream_o1.sh:14

See the review for details. If you are the PR author and disagree, react 👎 here; the maintainer will see your disagreement.

[vllm-omni-review-bot]

Omni ReviewBot: finding feedback

[p1] Online serve does not select the pipeline class — examples/online_serving/hidream_o1_image/run_server.sh:13

See the review for details. If you are the PR author and disagree, react 👎 here; the maintainer will see your disagreement.

[vllm-omni-review-bot]

Omni ReviewBot: finding feedback

[p1] CFG negative branch sees a stale vision cache — vllm_omni/diffusion/models/hidream_o1_image/pipeline_hidream_o1_image.py:595

See the review for details. If you are the PR author and disagree, react 👎 here; the maintainer will see your disagreement.

[vllm-omni-review-bot]

Omni ReviewBot: finding feedback

[p1] Edit vision tower receives target and ref grids — vllm_omni/diffusion/models/hidream_o1_image/pipeline_hidream_o1_image.py:445

See the review for details. If you are the PR author and disagree, react 👎 here; the maintainer will see your disagreement.

[hsliuustc0106]

@SOKOUDJOU-LEOPOLD Are you planning to continue work on this PR? Please reply by 2026-10-19 10:37:13 UTC.

A yes or a brief plan is enough; code changes do not need to be finished within these three days. If you are waiting on a review or other maintainer action, please let us know.

If there is no reply, new commit, or maintainer takeover by then, we will recheck and close this PR as inactive. It can be reopened when work resumes.

[hsliuustc0106]

Closing as inactive after the 3-day deadline with no author reply. Feel free to reopen anytime.

3 天期限已过且作者未回复，先关闭；随时可以 reopen。