diff --git a/.buildkite/cuda/test-nightly.yml b/.buildkite/cuda/test-nightly.yml index 0c5d94d81..986047bf6 100644 --- a/.buildkite/cuda/test-nightly.yml +++ b/.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 diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 58bca4809..beb70949d 100644 --- a/docs/models/supported_models.md +++ b/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 diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 34b642b9e..4a46c6b25 100644 --- a/docs/user_guide/diffusion_features.md +++ b/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** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index 2d1f9823a..2ee09e178 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/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( diff --git a/recipes/HiDream-ai/HiDream-O1-Image.md b/recipes/HiDream-ai/HiDream-O1-Image.md new file mode 100644 index 000000000..1864324b2 --- /dev/null +++ b/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. diff --git a/recipes/README.md b/recipes/README.md index 8b93124d0..1a384d08e 100644 --- a/recipes/README.md +++ b/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) | diff --git a/tests/assets/hidream_o1/hidream_o1_baseline.png b/tests/assets/hidream_o1/hidream_o1_baseline.png new file mode 100644 index 000000000..38920f301 Binary files /dev/null and b/tests/assets/hidream_o1/hidream_o1_baseline.png differ diff --git a/tests/diffusion/models/hidream_o1_image/test_hidream_o1_image_cache.py b/tests/diffusion/models/hidream_o1_image/test_hidream_o1_image_cache.py new file mode 100644 index 000000000..e5581cd77 --- /dev/null +++ b/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, + ) diff --git a/tests/e2e/accuracy/test_hidream_o1_image.py b/tests/e2e/accuracy/test_hidream_o1_image.py new file mode 100644 index 000000000..923e53a47 --- /dev/null +++ b/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, + ) diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py index 737eba753..1b40b1e5f 100644 --- a/tests/model_tests/diffusion/test_alignment.py +++ b/tests/model_tests/diffusion/test_alignment.py @@ -74,6 +74,7 @@ EXCLUDED_MODELS = [ "Cosmos3OmniPipeline", "DiffusersAdapterPipeline", "HiDreamImagePipeline", + "HiDreamO1ImagePipeline", "DreamZeroPipeline", "StableDiffusionXLPipeline", "Gr00tN1d7Pipeline", diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 2abd478c7..786e79e9d 100644 --- a/vllm_omni/diffusion/data.py +++ b/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() diff --git a/vllm_omni/diffusion/models/hidream_o1_image/__init__.py b/vllm_omni/diffusion/models/hidream_o1_image/__init__.py new file mode 100644 index 000000000..fe2723c76 --- /dev/null +++ b/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", +] diff --git a/vllm_omni/diffusion/models/hidream_o1_image/hidream_o1_image_transformer.py b/vllm_omni/diffusion/models/hidream_o1_image/hidream_o1_image_transformer.py new file mode 100644 index 000000000..8a0b1955c --- /dev/null +++ b/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"] diff --git a/vllm_omni/diffusion/models/hidream_o1_image/pipeline_hidream_o1_image.py b/vllm_omni/diffusion/models/hidream_o1_image/pipeline_hidream_o1_image.py new file mode 100644 index 000000000..de16df9de --- /dev/null +++ b/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) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index b248e6690..527ad2f74 100644 --- a/vllm_omni/diffusion/registry.py +++ b/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", diff --git a/vllm_omni/diffusion/utils/hf_utils.py b/vllm_omni/diffusion/utils/hf_utils.py index 78d8c6536..5c41685f0 100644 --- a/vllm_omni/diffusion/utils/hf_utils.py +++ b/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)