[Feat] ERNIE image model (T2I) (#2861)

Signed-off-by: Lancer <maruixiang6688@gmail.com>

docs/models/supported_models.md

@@ -66,5 +66,5 @@ th {
 | `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | |
 | `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | |
 |`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | |
-
+| `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 ✅︎ indicates the model is supported on that backend. Empty cells mean not listed as supported on that backend.

docs/user_guide/diffusion_features.md

@@ -129,6 +129,7 @@ The following tables show which models support each feature:
 | **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
 | **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ |
 | **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ✅ | ❌ | ✅ (decode) | ✅ | ❌ |
+| **ERNIE-Image** | ❌ | ✅ | ✅ | ❓ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
 
 > Notes:
 > 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT.

pyproject.toml

@@ -228,6 +228,10 @@ extend-ignore-identifiers-re = [
     ".*nin.*",
     ".*[Oo]no_[Aa]nna.*",
     ".*cann.*",
+    ".*[Nn][Dd]3.*",
+]
+extend-ignore-re = [
+    "\\.\\.\\.n,d->\\.\\.\\.nd",
 ]
 [tool.typos.default.extend-words]
 ue = "ue"

recipes/Baidu/ERNIE-Image.md

@@ -0,0 +1,149 @@
+# ERNIE-Image for text-to-image generation
+
+## Summary
+
+- Vendor: Baidu
+- Model: `baidu/ERNIE-Image` / `baidu/ERNIE-Image-Turbo`
+- Task: Text-to-image generation
+- Mode: Online serving with the OpenAI-compatible API
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe when you want a known-good starting point for serving
+`baidu/ERNIE-Image` or `baidu/ERNIE-Image-Turbo` with vLLM-Omni for
+high-quality text-to-image generation.
+
+ERNIE-Image is an 8B-parameter Diffusion Transformer (DiT) model that achieves
+state-of-the-art performance among open-weight text-to-image models. Key
+strengths include:
+
+- **Text rendering:** Excellent for dense, long-form, and layout-sensitive
+  text — ideal for posters, infographics, and UI-like images.
+- **Instruction following:** Reliably handles complex prompts with multiple
+  objects, detailed relationships, and knowledge-intensive descriptions.
+- **Structured generation:** Effective for posters, comics, storyboards, and
+  multi-panel compositions where layout matters.
+- **Style coverage:** Supports realistic photography, design-oriented imagery,
+  and stylized aesthetic outputs.
+- **Practical deployment:** Can run on consumer GPUs with 24GB VRAM.
+
+Two model variants are provided:
+
+1. **ERNIE-Image** — The SFT model with stronger general-purpose capability
+   and instruction fidelity, typically using 50 inference steps.
+2. **ERNIE-Image-Turbo** — Optimized by DMD and RL for faster speed and
+   higher aesthetics, requiring only 8 inference steps.
+
+## References
+
+- Upstream model card: <https://huggingface.co/baidu/ERNIE-Image>
+- Upstream model card (Turbo): <https://huggingface.co/baidu/ERNIE-Image-Turbo>
+- GitHub: <https://github.com/baidu/ERNIE-Image>
+- Blog: <https://ernie-image.github.io/>
+
+## Hardware Support
+
+This recipe currently documents tested configurations for CUDA GPU serving.
+Add more sections for other hardware as community validation lands.
+
+## GPU
+
+### 1 x RTX 4090 (Single GPU, 24GB VRAM, Minimum Recommended)
+
+#### Command
+
+**ERNIE-Image (full model):**
+
+```bash
+vllm serve baidu/ERNIE-Image --omni \
+  --enable-layerwise-offload \
+  --port 8091
+```
+
+**ERNIE-Image-Turbo (distilled, faster):**
+
+```bash
+vllm serve baidu/ERNIE-Image-Turbo --omni \
+  --enable-layerwise-offload \
+  --port 8091
+```
+
+### 2 x RTX 4090 (Multi-GPU, 24GB VRAM)
+
+#### Command
+
+**ERNIE-Image (full model):**
+
+```bash
+vllm serve baidu/ERNIE-Image --omni \
+  --tensor-parallel-size 2 \
+  --enable-cpu-offload \
+  --port 8091
+```
+
+**ERNIE-Image-Turbo (distilled, faster):**
+
+```bash
+vllm serve baidu/ERNIE-Image-Turbo --omni \
+  --tensor-parallel-size 2 \
+  --enable-cpu-offload \
+  --port 8091
+```
+
+#### Verification
+
+After the server is ready, test with a simple request:
+
+```bash
+curl -X POST http://localhost:8091/v1/images/generations \
+  -H "Content-Type: application/json" \
+  -d '{
+    "prompt": "A photo of a cat sitting on a laptop keyboard, digital art style.",
+    "size": "1024x1024",
+    "num_inference_steps": 50,
+    "guidance_scale": 4.0,
+    "seed": 42
+  }' | jq -r '.data[0].b64_json' | base64 -d > output.png
+```
+
+For ERNIE-Image-Turbo, reduce the inference steps:
+
+```bash
+curl -X POST http://localhost:8091/v1/images/generations \
+  -H "Content-Type: application/json" \
+  -d '{
+    "prompt": "A photo of a cat sitting on a laptop keyboard, digital art style.",
+    "size": "1024x1024",
+    "num_inference_steps": 8,
+    "guidance_scale": 1.0,
+    "seed": 42
+  }' | jq -r '.data[0].b64_json' | base64 -d > output.png
+```
+
+#### Notes
+
+- **Memory usage:** The 8B model requires significant VRAM. Use
+  `--enable-cpu-offload` to reduce GPU memory footprint by offloading
+  components to CPU when not in use.
+- **Key flags:**
+  - `--omni` — enables vLLM-Omni diffusion serving.
+- **Advanced features:**
+  - **TP (Tensor Parallelism):** `--tensor-parallel-size <N>` — distribute
+    model weights across N GPUs.
+  - **SP (Sequence Parallelism):** `--usp <N>` (Ulysses SP) and `--ring <N>`
+    (Ring SP) for long-sequence workloads.
+  - **HSDP:** `--use-hsdp` enables Hybrid Sharded Data Parallelism; use
+    `--hsdp-shard-size` and `--hsdp-replicate-size` for fine-grained control.
+  - **Cache-DiT:** `--cache-backend cache_dit` caches DiT intermediate outputs
+    for faster generation; configure via `--cache-config`.
+  - **Layer offload:** `--enable-layerwise-offload` offloads DiT layers to CPU
+    for memory-constrained scenarios.
+- **Prompt Enhancer (PE):** ERNIE-Image includes an optional 3B-parameter
+  Prompt Enhancer model that expands brief user inputs into richer structured
+  descriptions, improving output quality. Set `use_pe=False` in the request
+  body to disable it if you prefer direct prompt processing or want to use
+  larger LLMs (e.g., Gemini, ChatGPT) for prompt enhancement instead.
+- **Recommended settings:**
+  - ERNIE-Image: `num_inference_steps=50`, `guidance_scale=4.0`
+  - ERNIE-Image-Turbo: `num_inference_steps=8`, `guidance_scale=1.0`

recipes/README.md

@@ -31,6 +31,8 @@ recipes/
   recipe for Wan2.2 14B on `8x Ascend NPU (A2/A3)`
 - [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md):
   online serving recipe for multimodal chat (`4x H100 80GB`) and standalone TTS (`1x H100 80GB`)
+- [`Baidu/ERNIE-Image.md`](./Baidu/ERNIE-Image.md): text-to-image serving
+  online serving recipe for ERNIE-Image 8B on `1x RTX 4090 24GB` or `2x RTX 4090 24GB`
 - [`fishaudio/Fish-Speech-S2-Pro.md`](./fishaudio/Fish-Speech-S2-Pro.md): online serving recipe for TTS on `1x A800 80GB`
 
 Within a single recipe file, include different hardware support sections such

tests/diffusion/models/ernie_image/test_ernie_image_pe.py

@@ -0,0 +1,135 @@
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.ernie_image.ernie_image_transformer import ErnieImageTransformer2DModel
+from vllm_omni.diffusion.models.ernie_image.pipeline_ernie_image import ErnieImagePipeline
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _TokenInputs(dict):
+    @property
+    def input_ids(self):
+        return self["input_ids"]
+
+    def to(self, _device):
+        return self
+
+
+class _FakeTokenizer:
+    pad_token_id = 0
+    eos_token_id = 1
+
+    def apply_chat_template(self, *_args, **_kwargs):
+        return "chat prompt"
+
+    def __call__(self, *_args, **_kwargs):
+        return _TokenInputs(input_ids=torch.tensor([[1, 2]]))
+
+    def decode(self, *_args, **_kwargs):
+        return "rank1-enhanced"
+
+
+class _FakePEModel:
+    def __init__(self):
+        self.calls = 0
+
+    def generate(self, **_kwargs):
+        self.calls += 1
+        return torch.tensor([[1, 2, 3, 4]])
+
+
+def test_enhance_prompt_uses_rank0_result_in_distributed(monkeypatch):
+    pipe = ErnieImagePipeline.__new__(ErnieImagePipeline)
+    pipe.use_pe = True
+    pipe.pe_tokenizer = _FakeTokenizer()
+    pipe.pe_model = _FakePEModel()
+
+    broadcasts = []
+
+    def fake_broadcast_object_list(values, src=0):
+        broadcasts.append((list(values), src))
+        values[0] = "rank0-enhanced"
+
+    monkeypatch.setattr(torch.distributed, "is_available", lambda: True)
+    monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True)
+    monkeypatch.setattr(torch.distributed, "get_world_size", lambda: 2)
+    monkeypatch.setattr(torch.distributed, "get_rank", lambda: 1)
+    monkeypatch.setattr(torch.distributed, "broadcast_object_list", fake_broadcast_object_list)
+
+    enhanced = pipe._enhance_prompt("original", torch.device("cpu"))
+
+    assert enhanced == "rank0-enhanced"
+    assert pipe.pe_model.calls == 0
+    assert broadcasts == [([None], 0)]
+
+
+def test_should_apply_pe_respects_sampling_extra_args():
+    req = SimpleNamespace(sampling_params=SimpleNamespace(extra_args={"apply_pe": False}))
+
+    assert ErnieImagePipeline._should_apply_pe(req) is False
+
+
+def test_should_apply_pe_disables_dummy_warmup_request():
+    req = SimpleNamespace(request_ids=["dummy_req_id"], sampling_params=SimpleNamespace(extra_args={}))
+
+    assert ErnieImagePipeline._should_apply_pe(req) is False
+
+
+def test_hybrid_ring_slices_full_attention_mask(monkeypatch):
+    import vllm_omni.diffusion.models.ernie_image.ernie_image_transformer as ernie_transformer
+
+    mask = torch.arange(8).view(1, 8)
+
+    monkeypatch.setattr(ernie_transformer, "_get_ring_parallel_info", lambda: (2, 1))
+
+    sliced = ErnieImageTransformer2DModel._slice_attention_mask_for_ring(mask)
+
+    assert sliced.tolist() == [[4, 5, 6, 7]]
+    assert sliced.is_contiguous()
+
+
+def test_resolve_model_path_downloads_repo_id_for_pe_discovery(monkeypatch):
+    calls = []
+
+    def fake_exists(path):
+        return path == "/cache/baidu/ERNIE-Image" or path == "/cache/baidu/ERNIE-Image/pe"
+
+    def fake_download(model, revision, allow_patterns):
+        calls.append((model, revision, allow_patterns))
+        return "/cache/baidu/ERNIE-Image"
+
+    monkeypatch.setattr("os.path.exists", fake_exists)
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.models.ernie_image.pipeline_ernie_image.download_weights_from_hf_specific",
+        fake_download,
+    )
+
+    from vllm_omni.diffusion.models.ernie_image.pipeline_ernie_image import _resolve_model_path_for_optional_pe
+
+    assert _resolve_model_path_for_optional_pe("baidu/ERNIE-Image", None) == "/cache/baidu/ERNIE-Image"
+    assert calls == [("baidu/ERNIE-Image", None, ["pe/*", "pe_tokenizer/*"])]
+
+
+def test_resize_dimensions_rounds_down_to_vae_scale_factor():
+    pipe = ErnieImagePipeline.__new__(ErnieImagePipeline)
+    pipe.vae_scale_factor = 16
+
+    assert pipe._resize_dimensions(1025, 1024) == (1024, 1024)
+    assert pipe._resize_dimensions(1024, 1025) == (1024, 1024)
+    assert pipe._resize_dimensions(1025, 1031) == (1024, 1024)
+
+
+def test_load_weights_accepts_transformer_prefix():
+    model = ErnieImageTransformer2DModel.__new__(ErnieImageTransformer2DModel)
+    param = torch.nn.Parameter(torch.zeros(1))
+
+    model.named_parameters = lambda: [("proj.weight", param)]
+    model.named_buffers = lambda: []
+
+    loaded = model.load_weights([("transformer.proj.weight", torch.ones(1))])
+
+    assert loaded == {"proj.weight"}
+    assert torch.equal(param.data, torch.ones(1))

tests/diffusion/models/ernie_image/test_ernie_image_rope.py

@@ -0,0 +1,187 @@
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.ernie_image.ernie_image_transformer import (
+    ErnieImageEmbedND3,
+    _apply_rotary_emb,
+    rope,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestErnieImageRopePositionEmbedding:
+    """Test ErnieImage RoPE position embedding functionality"""
+
+    def test_rope_output_shape(self):
+        """Verify rope produces correct output shape"""
+        dim = 64
+        theta = 256
+        pos = torch.tensor([1.0, 2.0, 3.0])
+
+        out = rope(pos, dim, theta)
+
+        assert out.shape == (3, dim // 2), f"Expected shape (3, {dim // 2}), got {out.shape}"
+        assert out.dtype == torch.float32
+
+    def test_rope_trigonometric_relationship(self):
+        """Verify rope output satisfies cos^2 + sin^2 = 1"""
+        dim = 64
+        theta = 256
+        pos = torch.tensor([1.0, 2.0, 3.0])
+
+        out = rope(pos, dim, theta)
+        cos_out = out.cos()
+        sin_out = out.sin()
+
+        cos_sq_sin_sq = cos_out**2 + sin_out**2
+        assert torch.allclose(cos_sq_sin_sq, torch.ones_like(cos_sq_sin_sq), atol=1e-6)
+
+    def test_embed_nd3_output_shape(self):
+        """Verify ErnieImageEmbedND3 produces correct output shape"""
+        head_dim = 128
+        axes_dim = (32, 48, 48)
+        theta = 256
+        B, S = 2, 100
+
+        pos_embed = ErnieImageEmbedND3(dim=head_dim, theta=theta, axes_dim=axes_dim)
+        ids = torch.randn(B, S, 3)
+
+        freqs_cos, freqs_sin = pos_embed(ids)
+
+        expected_dim = sum(axes_dim) // 2
+        assert freqs_cos.shape == (B, S, expected_dim), (
+            f"Expected shape ({B}, {S}, {expected_dim}), got {freqs_cos.shape}"
+        )
+        assert freqs_sin.shape == (B, S, expected_dim)
+
+    def test_embed_nd3_different_batch_elements(self):
+        """Verify different batch elements produce different embeddings"""
+        head_dim = 128
+        axes_dim = (32, 48, 48)
+        theta = 256
+        B, S = 2, 100
+
+        pos_embed = ErnieImageEmbedND3(dim=head_dim, theta=theta, axes_dim=axes_dim)
+
+        ids = torch.zeros(B, S, 3)
+        ids[0, :, 0] = 1.0
+        ids[1, :, 0] = 2.0
+
+        freqs_cos, freqs_sin = pos_embed(ids)
+
+        assert not torch.allclose(freqs_cos[0], freqs_cos[1]), (
+            "Different batch elements should produce different cos embeddings"
+        )
+        assert not torch.allclose(freqs_sin[0], freqs_sin[1]), (
+            "Different batch elements should produce different sin embeddings"
+        )
+
+    def test_apply_rotary_emb_output_shape(self):
+        """Verify _apply_rotary_emb preserves input shape"""
+        B, S, H, D = 2, 100, 32, 128
+
+        x = torch.randn(B, S, H, D)
+        freqs_cos = torch.randn(B, S, D // 2)
+        freqs_sin = torch.randn(B, S, D // 2)
+
+        out = _apply_rotary_emb(x, freqs_cos, freqs_sin)
+
+        assert out.shape == x.shape, f"Expected shape {x.shape}, got {out.shape}"
+
+    def test_apply_rotary_emb_no_nan_inf(self):
+        """Verify _apply_rotary_emb produces valid values"""
+        B, S, H, D = 2, 100, 32, 128
+
+        x = torch.randn(B, S, H, D)
+        freqs_cos = torch.randn(B, S, D // 2)
+        freqs_sin = torch.randn(B, S, D // 2)
+
+        out = _apply_rotary_emb(x, freqs_cos, freqs_sin)
+
+        assert not torch.isnan(out).any(), "Output should not contain NaN"
+        assert not torch.isinf(out).any(), "Output should not contain Inf"
+
+    def test_apply_rotary_emb_preserves_batch_dimension(self):
+        """Verify _apply_rotary_emb preserves batch dimension for CFG mode"""
+        B, S, H, D = 2, 100, 32, 128
+
+        x = torch.randn(B, S, H, D)
+
+        freqs_cos = torch.zeros(B, S, D // 2)
+        freqs_sin = torch.zeros(B, S, D // 2)
+        freqs_cos[0] = 1.0
+        freqs_cos[1] = 0.5
+
+        out = _apply_rotary_emb(x, freqs_cos, freqs_sin)
+
+        assert not torch.allclose(out[0], out[1]), "Different batch elements should produce different outputs"
+
+    def test_cfg_mode_text_lens_difference(self):
+        """Verify CFG mode with different text_lens produces different positional embeddings"""
+        head_dim = 128
+        axes_dim = (32, 48, 48)
+        theta = 256
+        B = 2
+        N_img = 64
+        Tmax = 50
+
+        text_lens = torch.tensor([2, 50])
+
+        text_ids = torch.cat(
+            [
+                torch.arange(Tmax, dtype=torch.float32).view(1, Tmax, 1).expand(B, -1, -1),
+                torch.zeros((B, Tmax, 2)),
+            ],
+            dim=-1,
+        )
+
+        grid_yx = torch.stack(
+            torch.meshgrid(
+                torch.arange(8, dtype=torch.float32),
+                torch.arange(8, dtype=torch.float32),
+                indexing="ij",
+            ),
+            dim=-1,
+        ).reshape(-1, 2)
+
+        image_ids = torch.cat(
+            [
+                text_lens.float().view(B, 1, 1).expand(-1, N_img, -1),
+                grid_yx.view(1, N_img, 2).expand(B, -1, -1),
+            ],
+            dim=-1,
+        )
+
+        ids = torch.cat([image_ids, text_ids], dim=1)
+
+        pos_embed = ErnieImageEmbedND3(dim=head_dim, theta=theta, axes_dim=axes_dim)
+        freqs_cos, freqs_sin = pos_embed(ids)
+
+        assert not torch.allclose(freqs_cos[0], freqs_cos[1]), (
+            "CFG mode: uncond and cond should have different positional embeddings"
+        )
+        assert not torch.allclose(freqs_sin[0], freqs_sin[1]), (
+            "CFG mode: uncond and cond should have different positional embeddings"
+        )
+
+    def test_apply_rotary_emb_matches_diffusers_style(self):
+        """Verify _apply_rotary_emb matches diffusers implementation style"""
+        B, S, H, D = 2, 100, 32, 128
+
+        x = torch.randn(B, S, H, D)
+        freqs = torch.randn(B, S, D // 2)
+        freqs_cos = freqs.cos()
+        freqs_sin = freqs.sin()
+
+        out_vllm = _apply_rotary_emb(x, freqs_cos, freqs_sin)
+
+        cos_ = freqs_cos.unsqueeze(2)
+        sin_ = freqs_sin.unsqueeze(2)
+        cos_ = torch.stack([cos_, cos_], dim=-1).reshape(*cos_.shape[:-1], -1)
+        sin_ = torch.stack([sin_, sin_], dim=-1).reshape(*sin_.shape[:-1], -1)
+        x1, x2 = x.chunk(2, dim=-1)
+        rotated = torch.cat((-x2, x1), dim=-1)
+        out_diffusers = x * cos_ + rotated * sin_
+
+        assert torch.allclose(out_vllm, out_diffusers, atol=1e-6), "vllm-omni and diffusers outputs should match"

tests/diffusion/models/ernie_image/test_ernie_image_tp.py

@@ -0,0 +1,66 @@
+import pytest
+
+from vllm_omni.diffusion.models.ernie_image.ernie_image_transformer import (
+    validate_ernie_image_tp_constraints,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestErnieImageTPConstraints:
+    """Test ErnieImage tensor parallelism constraints"""
+
+    def test_tp1_heads_divisible(self):
+        """Verify TP=1 works with default heads=32"""
+        heads_per_gpu = validate_ernie_image_tp_constraints(heads=32, tensor_parallel_size=1)
+        assert heads_per_gpu == 32
+
+    def test_tp2_heads_divisible(self):
+        """Verify TP=2 works with heads=32"""
+        heads_per_gpu = validate_ernie_image_tp_constraints(heads=32, tensor_parallel_size=2)
+        assert heads_per_gpu == 16
+
+    def test_tp4_heads_divisible(self):
+        """Verify TP=4 works with heads=32"""
+        heads_per_gpu = validate_ernie_image_tp_constraints(heads=32, tensor_parallel_size=4)
+        assert heads_per_gpu == 8
+
+    def test_tp8_heads_divisible(self):
+        """Verify TP=8 works with heads=32"""
+        heads_per_gpu = validate_ernie_image_tp_constraints(heads=32, tensor_parallel_size=8)
+        assert heads_per_gpu == 4
+
+    def test_tp3_heads_not_divisible(self):
+        """Verify TP=3 fails when heads=32 not divisible by 3"""
+        with pytest.raises(ValueError, match=r"num_attention_heads.*must be divisible by tensor_parallel_size"):
+            validate_ernie_image_tp_constraints(heads=32, tensor_parallel_size=3)
+
+    def test_tp5_heads_not_divisible(self):
+        """Verify TP=5 fails when heads=32 not divisible by 5"""
+        with pytest.raises(ValueError, match=r"num_attention_heads.*must be divisible by tensor_parallel_size"):
+            validate_ernie_image_tp_constraints(heads=32, tensor_parallel_size=5)
+
+    def test_tp6_heads_not_divisible(self):
+        """Verify TP=6 fails when heads=32 not divisible by 6"""
+        with pytest.raises(ValueError, match=r"num_attention_heads.*must be divisible by tensor_parallel_size"):
+            validate_ernie_image_tp_constraints(heads=32, tensor_parallel_size=6)
+
+    def test_custom_heads_tp2_divisible(self):
+        """Verify TP=2 works with custom heads=24"""
+        heads_per_gpu = validate_ernie_image_tp_constraints(heads=24, tensor_parallel_size=2)
+        assert heads_per_gpu == 12
+
+    def test_custom_heads_tp3_divisible(self):
+        """Verify TP=3 works with custom heads=24"""
+        heads_per_gpu = validate_ernie_image_tp_constraints(heads=24, tensor_parallel_size=3)
+        assert heads_per_gpu == 8
+
+    def test_tp_invalid_zero(self):
+        """Verify TP=0 fails"""
+        with pytest.raises(ValueError, match=r"tensor_parallel_size must be > 0"):
+            validate_ernie_image_tp_constraints(heads=32, tensor_parallel_size=0)
+
+    def test_tp_invalid_negative(self):
+        """Verify negative TP fails"""
+        with pytest.raises(ValueError, match=r"tensor_parallel_size must be > 0"):
+            validate_ernie_image_tp_constraints(heads=32, tensor_parallel_size=-1)

tests/e2e/offline_inference/test_ernie_image.py

@@ -0,0 +1,58 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+End-to-end test for ErnieImage text-to-image generation.
+
+Equivalent to running:
+    vllm serve baidu/ERNIE-Image --omni
+"""
+
+import pytest
+from PIL import Image
+
+from tests.helpers.runtime import OmniRunner
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+MODEL = "baidu/ERNIE-Image"
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+def test_ernie_image_text_to_image():
+    with OmniRunner(
+        MODEL,
+        parallel_config=DiffusionParallelConfig(
+            tensor_parallel_size=2,
+        ),
+        enable_cpu_offload=True,
+    ) as runner:
+        omni_outputs = list(
+            runner.omni.generate(
+                prompts=["A photo of a cat sitting on a laptop"],
+                sampling_params_list=OmniDiffusionSamplingParams(
+                    height=512,
+                    width=512,
+                    num_inference_steps=2,
+                    guidance_scale=4.0,
+                    seed=42,
+                ),
+            )
+        )
+
+    assert len(omni_outputs) > 0
+    output = omni_outputs[0]
+    images = None
+    if output.images:
+        images = output.images
+    elif hasattr(output, "request_output") and output.request_output:
+        for stage_out in output.request_output:
+            if hasattr(stage_out, "images") and stage_out.images:
+                images = stage_out.images
+                break
+
+    assert images is not None
+    assert len(images) > 0
+    assert isinstance(images[0], Image.Image)
+    assert images[0].size == (512, 512)

tests/e2e/online_serving/test_ernie_image_expansion.py

@@ -0,0 +1,63 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+End-to-end test for ErnieImage text-to-image generation via online serving.
+"""
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import (
+    OmniServer,
+    OmniServerParams,
+    OpenAIClientHandler,
+    dummy_messages_from_mix_data,
+)
+
+T2I_PROMPT = "A photo of a cat sitting on a laptop keyboard, digital art style."
+NEGATIVE_PROMPT = "blurry, low quality"
+MODEL = "baidu/ERNIE-Image"
+PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"}, num_cards=2)
+
+
+def _get_diffusion_feature_cases(model: str):
+    return [
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--tensor-parallel-size",
+                    "2",
+                    "--enable-cpu-offload",
+                ],
+            ),
+            id="tp_2_cpu_offload",
+            marks=PARALLEL_FEATURE_MARKS,
+        ),
+    ]
+
+
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@pytest.mark.parametrize(
+    "omni_server",
+    _get_diffusion_feature_cases(MODEL),
+    indirect=True,
+)
+def test_ernie_image(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+    """Test each diffusion feature with ERNIE-Image (text-to-image), one feature per case."""
+    messages = dummy_messages_from_mix_data(content_text=T2I_PROMPT)
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "extra_body": {
+            "height": 512,
+            "width": 512,
+            "num_inference_steps": 2,
+            "negative_prompt": NEGATIVE_PROMPT,
+            "guidance_scale": 4.0,
+            "seed": 42,
+        },
+    }
+    openai_client.send_diffusion_request(request_config)

vllm_omni/diffusion/cache/cache_dit_backend.py

@@ -1320,6 +1320,69 @@ def enable_cache_for_glm_image(pipeline: Any, cache_config: Any) -> Callable[[in
     return refresh_cache_context
 
 
+def enable_cache_for_ernie_image(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
+    """Enable cache-dit for ERNIE-Image pipeline.
+
+    ERNIE-Image blocks have signature:
+        forward(x, rotary_pos_emb, temb, attention_mask) -> x
+
+    Where x is hidden_states (concatenated image + text tokens).
+    This matches Pattern_3 which expects:
+        - Input: hidden_states only
+        - Output: hidden_states only
+
+    Args:
+        pipeline: The ERNIE-Image pipeline instance.
+        cache_config: DiffusionCacheConfig instance with cache configuration.
+    Returns:
+        A refresh function that can be called to update cache context with new num_inference_steps.
+    """
+    db_cache_config = _build_db_cache_config(cache_config)
+
+    calibrator_config = None
+    if cache_config.enable_taylorseer:
+        taylorseer_order = cache_config.taylorseer_order
+        calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
+        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
+
+    logger.info(
+        f"Enabling cache-dit on ERNIE-Image transformer: "
+        f"Fn={db_cache_config.Fn_compute_blocks}, "
+        f"Bn={db_cache_config.Bn_compute_blocks}, "
+        f"W={db_cache_config.max_warmup_steps}, "
+    )
+
+    cache_dit.enable_cache(
+        BlockAdapter(
+            transformer=pipeline.transformer,
+            blocks=pipeline.transformer.layers,
+            forward_pattern=ForwardPattern.Pattern_3,
+            check_forward_pattern=False,
+        ),
+        cache_config=db_cache_config,
+        calibrator_config=calibrator_config,
+    )
+
+    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
+        if cache_config.scm_steps_mask_policy is None:
+            cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose)
+        else:
+            cache_dit.refresh_context(
+                pipeline.transformer,
+                cache_config=DBCacheConfig().reset(
+                    num_inference_steps=num_inference_steps,
+                    steps_computation_mask=cache_dit.steps_mask(
+                        mask_policy=cache_config.scm_steps_mask_policy,
+                        total_steps=num_inference_steps,
+                    ),
+                    steps_computation_policy=cache_config.scm_steps_policy,
+                ),
+                verbose=verbose,
+            )
+
+    return refresh_cache_context
+
+
 # Register custom cache-dit enablers after function definitions
 CUSTOM_DIT_ENABLERS.update(
     {
@@ -1342,6 +1405,7 @@ CUSTOM_DIT_ENABLERS.update(
         "BagelPipeline": enable_cache_for_bagel,
         "GlmImagePipeline": enable_cache_for_glm_image,
         "Flux2Pipeline": enable_cache_for_flux2,
+        "ErnieImagePipeline": enable_cache_for_ernie_image,
     }
 )
 

vllm_omni/diffusion/models/ernie_image/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+ErnieImage diffusion model for vLLM-Omni.
+
+This module implements ERNIE-Image text-to-image generation with:
+- ErnieImageTransformer2DModel: Custom DiT transformer
+- ErnieImagePipeline: Full generation pipeline
+"""
+
+from vllm_omni.diffusion.models.ernie_image.ernie_image_transformer import ErnieImageTransformer2DModel
+from vllm_omni.diffusion.models.ernie_image.pipeline_ernie_image import ErnieImagePipeline
+
+__all__ = ["ErnieImageTransformer2DModel", "ErnieImagePipeline"]

vllm_omni/diffusion/models/ernie_image/ernie_image_transformer.py

@@ -0,0 +1,646 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_ernie_image.py
+
+from collections.abc import Iterable
+from types import SimpleNamespace
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+from diffusers.models.embeddings import TimestepEmbedding, Timesteps
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from vllm.distributed import get_tensor_model_parallel_world_size
+from vllm.logger import init_logger
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.sp_plan import (
+    SequenceParallelInput,
+    SequenceParallelOutput,
+)
+
+logger = init_logger(__name__)
+
+if TYPE_CHECKING:
+    from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+
+
+def _get_sequence_parallel_world_size_or_one() -> int:
+    try:
+        from vllm_omni.diffusion.distributed.parallel_state import get_sequence_parallel_world_size
+
+        return max(1, int(get_sequence_parallel_world_size()))
+    except Exception:
+        return 1
+
+
+def _get_ring_parallel_info() -> tuple[int, int]:
+    try:
+        from vllm_omni.diffusion.distributed.parallel_state import (
+            get_ring_parallel_rank,
+            get_ring_parallel_world_size,
+        )
+
+        return max(1, int(get_ring_parallel_world_size())), int(get_ring_parallel_rank())
+    except Exception:
+        return 1, 0
+
+
+def validate_ernie_image_tp_constraints(*, heads: int, tensor_parallel_size: int) -> int:
+    """Validate ErnieImage TP constraints without requiring a distributed context.
+
+    Args:
+        heads: Number of attention heads
+        tensor_parallel_size: Tensor parallel size
+
+    Returns:
+        Number of heads per GPU after TP sharding
+
+    Raises:
+        ValueError: If heads is not divisible by tensor_parallel_size
+    """
+    tp_size = int(tensor_parallel_size)
+    if tp_size <= 0:
+        raise ValueError(f"tensor_parallel_size must be > 0, got {tp_size}")
+    if heads % tp_size != 0:
+        raise ValueError(f"num_attention_heads ({heads}) must be divisible by tensor_parallel_size ({tp_size})")
+    return heads // tp_size
+
+
+def rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor:
+    assert dim % 2 == 0
+    scale = torch.arange(0, dim, 2, dtype=torch.float32, device=pos.device) / dim
+    omega = 1.0 / (theta**scale)
+    out = torch.einsum("...n,d->...nd", pos, omega)
+    return out.float()
+
+
+def _apply_rotary_emb(
+    x: torch.Tensor,
+    freqs_cos: torch.Tensor,
+    freqs_sin: torch.Tensor,
+) -> torch.Tensor:
+    cos_ = freqs_cos.unsqueeze(2).to(x.dtype)
+    sin_ = freqs_sin.unsqueeze(2).to(x.dtype)
+    cos_ = cos_.repeat_interleave(2, dim=-1)
+    sin_ = sin_.repeat_interleave(2, dim=-1)
+    rot_dim = cos_.shape[-1]
+    x_rot = x[..., :rot_dim]
+    x1, x2 = x_rot.chunk(2, dim=-1)
+    x_rotated = torch.cat((-x2, x1), dim=-1)
+    return torch.cat((x_rot * cos_ + x_rotated * sin_, x[..., rot_dim:]), dim=-1)
+
+
+class ErnieImageEmbedND3(nn.Module):
+    def __init__(self, dim: int, theta: int, axes_dim: tuple[int, int, int]):
+        super().__init__()
+        self.dim = dim
+        self.theta = theta
+        self.axes_dim = list(axes_dim)
+
+    def forward(self, ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        cos_out = []
+        sin_out = []
+        pos = ids.float()
+        for i in range(3):
+            emb = rope(pos[..., i], self.axes_dim[i], self.theta)
+            cos_out.append(emb.cos())
+            sin_out.append(emb.sin())
+        freqs_cos = torch.cat(cos_out, dim=-1)
+        freqs_sin = torch.cat(sin_out, dim=-1)
+        return freqs_cos, freqs_sin
+
+
+class ErnieImagePatchEmbedDynamic(nn.Module):
+    def __init__(self, in_channels: int, embed_dim: int, patch_size: int):
+        super().__init__()
+        self.patch_size = patch_size
+        self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size, bias=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.proj(x)
+        batch_size, dim, height, width = x.shape
+        return x.reshape(batch_size, dim, height * width).transpose(1, 2).contiguous()
+
+
+class UnifiedPrepare(nn.Module):
+    """Prepares hidden_states, RoPE embeddings, and attention mask for sequence parallel.
+
+    This module encapsulates the input projection, RoPE, and attention_mask computation.
+    This creates a module boundary where _sp_plan can shard outputs via split_output=True.
+
+    The key insight is that hidden_states, freqs_cos, and freqs_sin
+    must be sharded together to maintain dimension alignment for attention layers.
+
+    Note: Our _sp_plan corresponds to diffusers' _cp_plan (Context Parallelism).
+    """
+
+    def __init__(
+        self,
+        x_embedder: nn.Module,
+        text_proj: nn.Module | None,
+        pos_embed: ErnieImageEmbedND3,
+    ):
+        super().__init__()
+        self.x_embedder = x_embedder
+        self.text_proj = text_proj
+        self.pos_embed = pos_embed
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        text_bth: torch.Tensor,
+        text_lens: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Prepare hidden_states, RoPE embeddings, and attention mask.
+
+        Args:
+            hidden_states: Image latent [B, C, H, W]
+            text_bth: Text embeddings [B, Tmax, text_in_dim]
+            text_lens: Text sequence lengths [B]
+
+        Returns:
+            hidden_states: [S, B, dim] where S = N_img + Tmax
+            freqs_cos: [B, seq, head_dim] cosine component of RoPE
+            freqs_sin: [B, seq, head_dim] sine component of RoPE
+            attention_mask: [B, seq] boolean mask
+        """
+        device = hidden_states.device
+        B, C, H, W = hidden_states.shape
+        Hp, Wp = H, W
+        N_img = Hp * Wp
+
+        img_sbh = self.x_embedder(hidden_states).transpose(0, 1).contiguous()
+        if self.text_proj is not None and text_bth.numel() > 0:
+            text_bth = self.text_proj(text_bth)
+        Tmax = text_bth.shape[1]
+        text_sbh = text_bth.transpose(0, 1).contiguous()
+
+        hidden_states = torch.cat([img_sbh, text_sbh], dim=0)
+
+        text_ids = (
+            torch.cat(
+                [
+                    torch.arange(Tmax, device=device, dtype=torch.float32).view(1, Tmax, 1).expand(B, -1, -1),
+                    torch.zeros((B, Tmax, 2), device=device),
+                ],
+                dim=-1,
+            )
+            if Tmax > 0
+            else torch.zeros((B, 0, 3), device=device)
+        )
+        grid_yx = torch.stack(
+            torch.meshgrid(
+                torch.arange(Hp, device=device, dtype=torch.float32),
+                torch.arange(Wp, device=device, dtype=torch.float32),
+                indexing="ij",
+            ),
+            dim=-1,
+        ).reshape(-1, 2)
+        image_ids = torch.cat(
+            [text_lens.float().view(B, 1, 1).expand(-1, N_img, -1), grid_yx.view(1, N_img, 2).expand(B, -1, -1)],
+            dim=-1,
+        )
+        freqs_cos, freqs_sin = self.pos_embed(torch.cat([image_ids, text_ids], dim=1))
+
+        valid_text = (
+            torch.arange(Tmax, device=device).view(1, Tmax) < text_lens.view(B, 1)
+            if Tmax > 0
+            else torch.zeros((B, 0), device=device, dtype=torch.bool)
+        )
+        attention_mask = torch.cat([torch.ones((B, N_img), device=device, dtype=torch.bool), valid_text], dim=1)
+        sp_size = _get_sequence_parallel_world_size_or_one()
+        pad_size = (-hidden_states.shape[0]) % sp_size
+        if pad_size:
+            pad_hidden = torch.zeros((pad_size, B, hidden_states.shape[-1]), device=device, dtype=hidden_states.dtype)
+            pad_freq = torch.zeros((B, pad_size, freqs_cos.shape[-1]), device=device, dtype=freqs_cos.dtype)
+            pad_mask = torch.zeros((B, pad_size), device=device, dtype=attention_mask.dtype)
+            hidden_states = torch.cat([hidden_states, pad_hidden], dim=0)
+            freqs_cos = torch.cat([freqs_cos, pad_freq], dim=1)
+            freqs_sin = torch.cat([freqs_sin, pad_freq], dim=1)
+            attention_mask = torch.cat([attention_mask, pad_mask], dim=1)
+
+        return hidden_states, freqs_cos, freqs_sin, attention_mask
+
+
+class ErnieImageAttention(nn.Module):
+    def __init__(
+        self,
+        parallel_config: DiffusionParallelConfig,
+        query_dim: int,
+        heads: int = 8,
+        dim_head: int = 64,
+        dropout: float = 0.0,
+        bias: bool = False,
+        qk_norm: str = "rms_norm",
+        added_proj_bias: bool | None = True,
+        out_bias: bool = True,
+        eps: float = 1e-5,
+        out_dim: int = None,
+        elementwise_affine: bool = True,
+        quant_config: "QuantizationConfig | None" = None,
+    ):
+        super().__init__()
+        self.parallel_config = parallel_config
+        self.head_dim = dim_head
+        self.inner_dim = out_dim if out_dim is not None else dim_head * heads
+        self.query_dim = query_dim
+        self.out_dim = out_dim if out_dim is not None else query_dim
+        tp_size = get_tensor_model_parallel_world_size()
+        if heads % tp_size != 0:
+            raise ValueError(f"num_attention_heads ({heads}) must be divisible by tensor_parallel_size ({tp_size})")
+        self.heads = heads // tp_size
+        self.dropout = dropout
+
+        self.to_q = ColumnParallelLinear(
+            input_size=query_dim,
+            output_size=self.inner_dim,
+            bias=bias,
+            gather_output=False,
+            quant_config=quant_config,
+        )
+        self.to_k = ColumnParallelLinear(
+            input_size=query_dim,
+            output_size=self.inner_dim,
+            bias=bias,
+            gather_output=False,
+            quant_config=quant_config,
+        )
+        self.to_v = ColumnParallelLinear(
+            input_size=query_dim,
+            output_size=self.inner_dim,
+            bias=bias,
+            gather_output=False,
+            quant_config=quant_config,
+        )
+
+        self.kv_num_heads = self.heads
+
+        if qk_norm is not None:
+            if qk_norm == "layer_norm":
+                self.norm_q = nn.LayerNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+                self.norm_k = nn.LayerNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+            elif qk_norm == "rms_norm":
+                self.norm_q = RMSNorm(dim_head, eps=eps)
+                self.norm_k = RMSNorm(dim_head, eps=eps)
+            else:
+                raise ValueError(f"unknown qk_norm: {qk_norm}")
+
+        self.to_out = nn.ModuleList(
+            [
+                RowParallelLinear(
+                    self.inner_dim,
+                    self.out_dim,
+                    bias=out_bias,
+                    input_is_parallel=True,
+                    return_bias=False,
+                    quant_config=quant_config,
+                ),
+                nn.Dropout(dropout),
+            ]
+        )
+
+        self.attn = Attention(
+            num_heads=self.heads,
+            head_size=self.head_dim,
+            softmax_scale=1.0 / (self.head_dim**0.5),
+            causal=False,
+            num_kv_heads=self.kv_num_heads,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        image_rotary_emb: torch.Tensor | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        query, _ = self.to_q(hidden_states)
+        key, _ = self.to_k(hidden_states)
+        value, _ = self.to_v(hidden_states)
+
+        query = query.unflatten(-1, (self.heads, -1))
+        key = key.unflatten(-1, (self.heads, -1))
+        value = value.unflatten(-1, (self.heads, -1))
+
+        if hasattr(self, "norm_q"):
+            query = self.norm_q(query)
+            key = self.norm_k(key)
+
+        if image_rotary_emb is not None:
+            freqs_cos, freqs_sin = image_rotary_emb
+            query = _apply_rotary_emb(query, freqs_cos, freqs_sin)
+            key = _apply_rotary_emb(key, freqs_cos, freqs_sin)
+
+        attn_metadata = None
+        if attention_mask is not None:
+            if attention_mask.dim() == 3:
+                attention_mask = attention_mask.unsqueeze(1)
+            attn_metadata = AttentionMetadata(attn_mask=attention_mask)
+
+        hidden_states = self.attn(query, key, value, attn_metadata)
+        hidden_states = hidden_states.flatten(2, 3).to(query.dtype)
+
+        hidden_states = self.to_out[0](hidden_states.contiguous())
+        hidden_states = self.to_out[1](hidden_states)
+
+        return hidden_states
+
+
+class ErnieImageFeedForward(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        ffn_hidden_size: int,
+        bias: bool = False,
+        quant_config: "QuantizationConfig | None" = None,
+    ):
+        super().__init__()
+        self.gate_proj = ColumnParallelLinear(
+            hidden_size,
+            ffn_hidden_size,
+            bias=bias,
+            gather_output=False,
+            quant_config=quant_config,
+        )
+        self.up_proj = ColumnParallelLinear(
+            hidden_size,
+            ffn_hidden_size,
+            bias=bias,
+            gather_output=False,
+            quant_config=quant_config,
+        )
+        self.linear_fc2 = RowParallelLinear(
+            ffn_hidden_size,
+            hidden_size,
+            bias=bias,
+            input_is_parallel=True,
+            return_bias=False,
+            quant_config=quant_config,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gate_out, _ = self.gate_proj(x)
+        up_out, _ = self.up_proj(x)
+        return self.linear_fc2(up_out * torch.nn.functional.gelu(gate_out))
+
+
+class ErnieImageSharedAdaLNBlock(nn.Module):
+    def __init__(
+        self,
+        parallel_config: DiffusionParallelConfig,
+        hidden_size: int,
+        num_heads: int,
+        ffn_hidden_size: int,
+        eps: float = 1e-6,
+        qk_layernorm: bool = True,
+        quant_config: "QuantizationConfig | None" = None,
+    ):
+        super().__init__()
+        self.adaLN_sa_ln = RMSNorm(hidden_size, eps=eps)
+        self.self_attention = ErnieImageAttention(
+            parallel_config=parallel_config,
+            query_dim=hidden_size,
+            dim_head=hidden_size // num_heads,
+            heads=num_heads,
+            qk_norm="rms_norm" if qk_layernorm else None,
+            eps=eps,
+            bias=False,
+            out_bias=False,
+            quant_config=quant_config,
+        )
+        self.adaLN_mlp_ln = RMSNorm(hidden_size, eps=eps)
+        self.mlp = ErnieImageFeedForward(hidden_size, ffn_hidden_size, quant_config=quant_config)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        rotary_pos_emb: torch.Tensor,
+        temb: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = temb
+        residual = hidden_states
+        hidden_states = self.adaLN_sa_ln(hidden_states)
+        hidden_states = (hidden_states.float() * (1 + scale_msa.float()) + shift_msa.float()).to(hidden_states.dtype)
+        hidden_states_bsh = hidden_states.permute(1, 0, 2)
+
+        attn_out = self.self_attention(
+            hidden_states_bsh, attention_mask=attention_mask, image_rotary_emb=rotary_pos_emb
+        )
+        attn_out = attn_out.permute(1, 0, 2)
+        hidden_states = residual + (gate_msa.float() * attn_out.float()).to(hidden_states.dtype)
+        residual = hidden_states
+        hidden_states = self.adaLN_mlp_ln(hidden_states)
+        hidden_states = (hidden_states.float() * (1 + scale_mlp.float()) + shift_mlp.float()).to(hidden_states.dtype)
+        return residual + (gate_mlp.float() * self.mlp(hidden_states).float()).to(hidden_states.dtype)
+
+
+class ErnieImageAdaLNContinuous(nn.Module):
+    def __init__(self, hidden_size: int, eps: float = 1e-6):
+        super().__init__()
+        self.norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=eps)
+        self.linear = nn.Linear(hidden_size, hidden_size * 2)
+
+    def forward(self, x: torch.Tensor, conditioning: torch.Tensor) -> torch.Tensor:
+        scale, shift = self.linear(conditioning).chunk(2, dim=-1)
+        x = self.norm(x)
+        x = x * (1 + scale.unsqueeze(0)) + shift.unsqueeze(0)
+        return x
+
+
+class ErnieImageTransformer2DModel(nn.Module):
+    _repeated_blocks = ["ErnieImageSharedAdaLNBlock"]
+    _layerwise_offload_blocks_attrs = ["layers"]
+
+    @staticmethod
+    def _is_layer_block(name: str, module) -> bool:
+        return "layers" in name and name.split(".")[-1].isdigit()
+
+    _hsdp_shard_conditions = [_is_layer_block]
+
+    _sp_plan = {
+        "unified_prepare": {
+            0: SequenceParallelInput(split_dim=0, expected_dims=3, split_output=True, auto_pad=True),
+            1: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True, auto_pad=True),
+            2: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True, auto_pad=True),
+        },
+        "final_linear": SequenceParallelOutput(gather_dim=0, expected_dims=3),
+    }
+
+    @staticmethod
+    def _slice_attention_mask_for_ring(attention_mask: torch.Tensor | None) -> torch.Tensor | None:
+        if attention_mask is None:
+            return None
+        ring_size, ring_rank = _get_ring_parallel_info()
+        if ring_size <= 1:
+            return attention_mask
+        seq_len = attention_mask.shape[1]
+        if seq_len % ring_size != 0:
+            raise ValueError(
+                "ERNIE-Image hybrid SP requires attention_mask length to be divisible by ring_degree, "
+                f"but got seq_len={seq_len}, ring_degree={ring_size}."
+            )
+        return attention_mask.chunk(ring_size, dim=1)[ring_rank].contiguous()
+
+    def __init__(
+        self,
+        patch_size: int = 1,
+        in_channels: int = 128,
+        out_channels: int | None = None,
+        num_layers: int = 24,
+        num_attention_heads: int = 24,
+        ffn_hidden_size: int = 8192,
+        hidden_size: int = 3072,
+        text_in_dim: int = 2560,
+        rope_theta: int = 256,
+        rope_axes_dim: tuple[int, int, int] = (32, 48, 48),
+        eps: float = 1e-6,
+        qk_layernorm: bool = True,
+        od_config: OmniDiffusionConfig = None,
+        quant_config: "QuantizationConfig | None" = None,
+    ):
+        super().__init__()
+        self.out_channels = out_channels or in_channels
+        self.patch_size = patch_size
+        self.in_channels = in_channels
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers
+        self.num_attention_heads = num_attention_heads
+        self.head_dim = hidden_size // num_attention_heads
+        self.text_in_dim = text_in_dim
+
+        if od_config is not None:
+            self.parallel_config = od_config.parallel_config
+        else:
+            from vllm_omni.diffusion.data import DiffusionParallelConfig
+
+            self.parallel_config = DiffusionParallelConfig()
+
+        self.config = SimpleNamespace(
+            patch_size=patch_size,
+            in_channels=in_channels,
+            out_channels=self.out_channels,
+            num_layers=num_layers,
+            num_attention_heads=num_attention_heads,
+            ffn_hidden_size=ffn_hidden_size,
+            hidden_size=hidden_size,
+            text_in_dim=text_in_dim,
+            rope_theta=rope_theta,
+            rope_axes_dim=rope_axes_dim,
+            eps=eps,
+            qk_layernorm=qk_layernorm,
+        )
+
+        self.x_embedder = ErnieImagePatchEmbedDynamic(in_channels, hidden_size, patch_size)
+        self.text_proj = nn.Linear(text_in_dim, hidden_size, bias=False) if text_in_dim != hidden_size else None
+        self.time_proj = Timesteps(hidden_size, flip_sin_to_cos=False, downscale_freq_shift=0)
+        self.time_embedding = TimestepEmbedding(hidden_size, hidden_size)
+        self.pos_embed = ErnieImageEmbedND3(dim=self.head_dim, theta=rope_theta, axes_dim=rope_axes_dim)
+        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size))
+        nn.init.zeros_(self.adaLN_modulation[-1].weight)
+        nn.init.zeros_(self.adaLN_modulation[-1].bias)
+        self.unified_prepare = UnifiedPrepare(self.x_embedder, self.text_proj, self.pos_embed)
+        self.layers = nn.ModuleList(
+            [
+                ErnieImageSharedAdaLNBlock(
+                    parallel_config=self.parallel_config,
+                    hidden_size=hidden_size,
+                    num_heads=num_attention_heads,
+                    ffn_hidden_size=ffn_hidden_size,
+                    eps=eps,
+                    qk_layernorm=qk_layernorm,
+                    quant_config=quant_config,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+        self.final_norm = ErnieImageAdaLNContinuous(hidden_size, eps)
+        self.final_linear = nn.Linear(hidden_size, patch_size * patch_size * self.out_channels)
+        nn.init.zeros_(self.final_linear.weight)
+        nn.init.zeros_(self.final_linear.bias)
+        self.gradient_checkpointing = False
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return next(self.parameters()).dtype
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        timestep: torch.Tensor,
+        text_bth: torch.Tensor,
+        text_lens: torch.Tensor,
+        return_dict: bool = True,
+    ) -> torch.Tensor | Transformer2DModelOutput:
+        dtype = hidden_states.dtype
+        B, C, H, W = hidden_states.shape
+        p = self.patch_size
+        Hp, Wp = H // p, W // p
+
+        N_img = Hp * Wp
+        hidden_states, freqs_cos, freqs_sin, attention_mask = self.unified_prepare(hidden_states, text_bth, text_lens)
+        attention_mask = self._slice_attention_mask_for_ring(attention_mask)
+        rotary_pos_emb = (freqs_cos, freqs_sin)
+        S = hidden_states.shape[0]
+
+        sample = self.time_proj(timestep)
+        sample = sample.to(dtype=dtype)
+        c = self.time_embedding(sample)
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = [
+            t.unsqueeze(0).expand(S, -1, -1).contiguous() for t in self.adaLN_modulation(c).chunk(6, dim=-1)
+        ]
+
+        for layer in self.layers:
+            temb = [shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp]
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                hidden_states = self._gradient_checkpointing_func(
+                    layer,
+                    hidden_states,
+                    rotary_pos_emb,
+                    temb,
+                    attention_mask,
+                )
+            else:
+                hidden_states = layer(hidden_states, rotary_pos_emb, temb, attention_mask)
+        hidden_states = self.final_norm(hidden_states, c).type_as(hidden_states)
+        patches = self.final_linear(hidden_states)[:N_img].transpose(0, 1).contiguous()
+        output = (
+            patches.view(B, Hp, Wp, p, p, self.out_channels)
+            .permute(0, 5, 1, 3, 2, 4)
+            .contiguous()
+            .view(B, self.out_channels, H, W)
+        )
+
+        if not return_dict:
+            return (output,)
+        return Transformer2DModelOutput(sample=output)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        params_dict = dict(self.named_parameters())
+
+        for name, buffer in self.named_buffers():
+            if name.endswith(".beta") or name.endswith(".eps"):
+                params_dict[name] = buffer
+
+        loaded_params: set[str] = set()
+        for name, loaded_weight in weights:
+            name = name.replace("transformer.", "", 1)
+            if name not in params_dict and ".to_out.0." in name:
+                name = name.replace(".to_out.0.", ".to_out.")
+            param = params_dict.get(name)
+            if param is None:
+                continue
+            weight_loader = getattr(param, "weight_loader", default_weight_loader)
+            weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+        return loaded_params

vllm_omni/diffusion/models/ernie_image/pipeline_ernie_image.py

@@ -0,0 +1,560 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/ernie_image/pipeline_ernie_image.py
+
+import json
+import os
+from collections.abc import Callable, Iterable
+
+import torch
+import torch.nn as nn
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.models.autoencoders.autoencoder_kl_flux2 import AutoencoderKLFlux2
+from diffusers.schedulers.scheduling_flow_match_euler_discrete import (
+    FlowMatchEulerDiscreteScheduler,
+)
+from diffusers.utils.torch_utils import randn_tensor
+from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+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.ernie_image.ernie_image_transformer import ErnieImageTransformer2DModel
+from vllm_omni.diffusion.models.interface import SupportImageInput
+from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
+from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
+from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
+
+logger = init_logger(__name__)
+
+
+def _resolve_model_path_for_optional_pe(model: str, revision: str | None) -> str:
+    if os.path.exists(model):
+        return model
+    try:
+        return download_weights_from_hf_specific(model, revision, ["pe/*", "pe_tokenizer/*"])
+    except Exception as e:
+        logger.debug("Failed to resolve ERNIE-Image PE files for %s: %s", model, e)
+        return model
+
+
+def get_ernie_image_post_process_func(od_config: OmniDiffusionConfig):
+    if od_config.output_type == "latent":
+        return lambda x: x
+    model_name = od_config.model
+    if os.path.exists(model_name):
+        model_path = model_name
+    else:
+        model_path = download_weights_from_hf_specific(model_name, None, ["*"])
+
+    vae_config_path = os.path.join(model_path, "vae/config.json")
+    with open(vae_config_path) as f:
+        vae_config = json.load(f)
+        vae_scale_factor = 2 ** (len(vae_config["block_out_channels"]) - 1) if "block_out_channels" in vae_config else 8
+
+    image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2, do_normalize=False)
+
+    def post_process_func(images: torch.Tensor):
+        images = (images.clamp(-1, 1) + 1) / 2
+        return image_processor.postprocess(images)
+
+    return post_process_func
+
+
+class ErnieImagePipeline(
+    nn.Module, CFGParallelMixin, SupportImageInput, ProgressBarMixin, DiffusionPipelineProfilerMixin
+):
+    support_image_input = False
+
+    def __init__(
+        self,
+        *,
+        od_config: OmniDiffusionConfig,
+        prefix: str = "",
+        is_distilled: bool = False,
+    ):
+        super().__init__()
+        self.od_config = od_config
+        self.is_distilled = is_distilled
+        self.weights_sources = [
+            DiffusersPipelineLoader.ComponentSource(
+                model_or_path=od_config.model,
+                subfolder="transformer",
+                revision=None,
+                prefix="transformer.",
+                fall_back_to_pt=True,
+            )
+        ]
+
+        self._execution_device = get_local_device()
+        model = od_config.model
+        logger.info("Model path for initialization: %s", model)
+        local_files_only = os.path.exists(model)
+        logger.info("Local files only: %s", local_files_only)
+
+        self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+            model,
+            subfolder="scheduler",
+            local_files_only=local_files_only,
+        )
+
+        self.text_encoder = AutoModel.from_pretrained(
+            model,
+            subfolder="text_encoder",
+            torch_dtype=od_config.dtype,
+            local_files_only=local_files_only,
+        ).to(self._execution_device)
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model,
+            subfolder="tokenizer",
+            local_files_only=local_files_only,
+        )
+
+        self.vae = AutoencoderKLFlux2.from_pretrained(
+            model,
+            subfolder="vae",
+            torch_dtype=od_config.dtype,
+            local_files_only=local_files_only,
+        ).to(self._execution_device)
+
+        # Load PE (Prompt Enhancement) model if available. For repo IDs,
+        # resolve/download only PE files first so the existence check works.
+        pe_base_path = _resolve_model_path_for_optional_pe(model, getattr(od_config, "revision", None))
+        pe_model_path = os.path.join(pe_base_path, "pe")
+        if os.path.exists(pe_model_path):
+            try:
+                self.pe_model = AutoModelForCausalLM.from_pretrained(
+                    pe_model_path,
+                    torch_dtype=od_config.dtype,
+                    local_files_only=True,
+                    trust_remote_code=True,
+                ).to(self._execution_device)
+                self.pe_tokenizer = AutoTokenizer.from_pretrained(
+                    pe_base_path,
+                    subfolder="pe_tokenizer",
+                    local_files_only=True,
+                    trust_remote_code=True,
+                    use_fast=False,
+                )
+                self.use_pe = True
+                logger.info("Loaded PE model from %s", pe_model_path)
+            except Exception as e:
+                logger.warning("Failed to load PE model: %s", e)
+                self.pe_model = None
+                self.pe_tokenizer = None
+                self.use_pe = False
+        else:
+            self.pe_model = None
+            self.pe_tokenizer = None
+            self.use_pe = False
+
+        transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, ErnieImageTransformer2DModel)
+        self.transformer = ErnieImageTransformer2DModel(
+            quant_config=od_config.quantization_config, **transformer_kwargs
+        )
+
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels)) if getattr(self, "vae", None) else 16
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
+        self.tokenizer_max_length = 512
+        self.default_sample_size = 128
+
+        self._guidance_scale = None
+        self._attention_kwargs = None
+        self._num_timesteps = None
+        self._current_timestep = None
+        self._interrupt = False
+        self.setup_diffusion_pipeline_profiler(
+            enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
+        )
+
+    @staticmethod
+    def _distributed_prompt_sync_state() -> tuple[bool, int]:
+        try:
+            if not torch.distributed.is_available() or not torch.distributed.is_initialized():
+                return False, 0
+            if torch.distributed.get_world_size() <= 1:
+                return False, 0
+            return True, torch.distributed.get_rank()
+        except Exception:
+            return False, 0
+
+    @staticmethod
+    def _broadcast_enhanced_prompt(prompt: str | None) -> str | None:
+        values = [prompt]
+        torch.distributed.broadcast_object_list(values, src=0)
+        return values[0]
+
+    def _enhance_prompt(
+        self,
+        prompt: str,
+        device: torch.device,
+        width: int = 1024,
+        height: int = 1024,
+        temperature: float = 0.6,
+        top_p: float = 0.95,
+    ) -> str:
+        if not self.use_pe or self.pe_model is None:
+            return prompt
+
+        sync_prompt, rank = self._distributed_prompt_sync_state()
+        if sync_prompt and rank != 0:
+            try:
+                enhanced = self._broadcast_enhanced_prompt(None)
+                return enhanced if enhanced else prompt
+            except Exception as e:
+                logger.warning("PE enhancement broadcast failed: %s", e)
+                return prompt
+
+        try:
+            user_content = json.dumps(
+                {"prompt": prompt, "width": width, "height": height},
+                ensure_ascii=False,
+            )
+            messages = [{"role": "user", "content": user_content}]
+            input_text = self.pe_tokenizer.apply_chat_template(
+                messages,
+                tokenize=False,
+                add_generation_prompt=False,
+            )
+            inputs = self.pe_tokenizer(input_text, return_tensors="pt").to(device)
+            output_ids = self.pe_model.generate(
+                **inputs,
+                max_new_tokens=512,
+                do_sample=temperature != 1.0 or top_p != 1.0,
+                temperature=temperature,
+                top_p=top_p,
+                pad_token_id=self.pe_tokenizer.pad_token_id,
+                eos_token_id=self.pe_tokenizer.eos_token_id,
+            )
+            output_ids = output_ids[0][inputs.input_ids.shape[1] :]
+            result = self.pe_tokenizer.decode(output_ids, skip_special_tokens=True).strip()
+            enhanced = result if result else prompt
+        except Exception as e:
+            logger.warning("PE enhancement failed: %s", e)
+            enhanced = prompt
+
+        if sync_prompt:
+            try:
+                synced = self._broadcast_enhanced_prompt(enhanced)
+                return synced if synced else prompt
+            except Exception as e:
+                logger.warning("PE enhancement broadcast failed: %s", e)
+        return enhanced
+
+    @staticmethod
+    def _is_warmup_request(req: OmniDiffusionRequest) -> bool:
+        request_ids = getattr(req, "request_ids", None) or ()
+        return len(request_ids) == 1 and request_ids[0] == "dummy_req_id"
+
+    @staticmethod
+    def _should_apply_pe(req: OmniDiffusionRequest) -> bool:
+        if ErnieImagePipeline._is_warmup_request(req):
+            return False
+        extra_args = getattr(req.sampling_params, "extra_args", {}) or {}
+        return bool(extra_args.get("apply_pe", True))
+
+    def encode_prompt(
+        self,
+        prompt: str | list[str],
+        device: torch.device,
+        num_images_per_prompt: int = 1,
+        width: int = 1024,
+        height: int = 1024,
+        apply_pe: bool = True,
+    ) -> list[torch.Tensor]:
+        if isinstance(prompt, str):
+            prompt = [prompt]
+
+        text_hiddens = []
+
+        for p in prompt:
+            if apply_pe and self.use_pe and self.pe_model is not None:
+                enhanced = self._enhance_prompt(p, device, width=width, height=height)
+                logger.info("PE: original='%s...' enhanced='%s...'", p[:50], enhanced[:50])
+                p = enhanced
+            ids = self.tokenizer(
+                p,
+                add_special_tokens=True,
+                truncation=True,
+                padding=False,
+            )["input_ids"]
+
+            if len(ids) == 0:
+                if self.tokenizer.bos_token_id is not None:
+                    ids = [self.tokenizer.bos_token_id]
+                else:
+                    ids = [0]
+
+            input_ids = torch.tensor([ids], device=device)
+            with torch.no_grad():
+                outputs = self.text_encoder(input_ids=input_ids, output_hidden_states=True)
+            hidden = outputs.hidden_states[-2][0]
+
+            for _ in range(num_images_per_prompt):
+                text_hiddens.append(hidden)
+
+        return text_hiddens
+
+    @staticmethod
+    def _patchify_latents(latents: torch.Tensor) -> torch.Tensor:
+        b, c, h, w = latents.shape
+        latents = latents.view(b, c, h // 2, 2, w // 2, 2)
+        latents = latents.permute(0, 1, 3, 5, 2, 4)
+        return latents.reshape(b, c * 4, h // 2, w // 2)
+
+    @staticmethod
+    def _unpatchify_latents(latents: torch.Tensor) -> torch.Tensor:
+        b, c, h, w = latents.shape
+        latents = latents.reshape(b, c // 4, 2, 2, h, w)
+        latents = latents.permute(0, 1, 4, 2, 5, 3)
+        return latents.reshape(b, c // 4, h * 2, w * 2)
+
+    @staticmethod
+    def _pad_text(text_hiddens: list[torch.Tensor], device: torch.device, dtype: torch.dtype, text_in_dim: int):
+        B = len(text_hiddens)
+        if B == 0:
+            return torch.zeros((0, 0, text_in_dim), device=device, dtype=dtype), torch.zeros(
+                (0,), device=device, dtype=torch.long
+            )
+        normalized = [
+            th.squeeze(1).to(device).to(dtype) if th.dim() == 3 else th.to(device).to(dtype) for th in text_hiddens
+        ]
+        lens = torch.tensor([t.shape[0] for t in normalized], device=device, dtype=torch.long)
+        Tmax = int(lens.max().item())
+        text_bth = torch.zeros((B, Tmax, text_in_dim), device=device, dtype=dtype)
+        for i, t in enumerate(normalized):
+            text_bth[i, : t.shape[0], :] = t
+        return text_bth, lens
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+
+    @property
+    def do_classifier_free_guidance(self):
+        return self._guidance_scale is not None and self._guidance_scale > 1 and not self.is_distilled
+
+    @property
+    def num_timesteps(self):
+        return self._num_timesteps
+
+    @property
+    def current_timestep(self):
+        return self._current_timestep
+
+    @property
+    def interrupt(self):
+        return self._interrupt
+
+    def _resize_dimensions(self, height: int, width: int) -> tuple[int, int]:
+        resized_height = height - height % self.vae_scale_factor
+        resized_width = width - width % self.vae_scale_factor
+        return resized_height, resized_width
+
+    def check_inputs(
+        self,
+        prompt,
+        height,
+        width,
+        prompt_embeds=None,
+        callback_on_step_end_tensor_inputs=None,
+        guidance_scale=None,
+    ):
+        if (height is not None and height % self.vae_scale_factor != 0) or (
+            width is not None and width % self.vae_scale_factor != 0
+        ):
+            resized_height, resized_width = self._resize_dimensions(height, width)
+            logger.warning(
+                "`height` and `width` have to be divisible by %s "
+                "but are %s and %s. Dimensions will be resized to %s and %s accordingly",
+                self.vae_scale_factor,
+                height,
+                width,
+                resized_height,
+                resized_width,
+            )
+
+        if callback_on_step_end_tensor_inputs is not None and not all(
+            k in ["latents", "prompt_embeds"] for k in callback_on_step_end_tensor_inputs
+        ):
+            raise ValueError("`callback_on_step_end_tensor_inputs` must be a subset of ['latents', 'prompt_embeds'].")
+
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError("Cannot forward both `prompt` and `prompt_embeds`. Please forward only one.")
+        if prompt is None and prompt_embeds is None:
+            raise ValueError("Must provide either `prompt` or `prompt_embeds`.")
+
+    def forward(
+        self,
+        req: OmniDiffusionRequest,
+        prompt: str | list[str] | None = None,
+        negative_prompt: str | list[str] | None = "",
+        height: int = 1024,
+        width: int = 1024,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 4.0,
+        num_images_per_prompt: int = 1,
+        generator: torch.Generator | None = None,
+        latents: torch.Tensor | None = None,
+        prompt_embeds: list[torch.FloatTensor] | None = None,
+        negative_prompt_embeds: list[torch.FloatTensor] | None = None,
+        output_type: str = "pil",
+        return_dict: bool = True,
+        callback_on_step_end: Callable[[int, int, dict], None] | None = None,
+        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
+    ) -> DiffusionOutput:
+        if len(req.prompts) > 1:
+            logger.warning("This model only supports a single prompt. Taking only the first.")
+        first_prompt = req.prompts[0]
+        prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "")
+
+        height = req.sampling_params.height or height
+        width = req.sampling_params.width or width
+        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
+        guidance_scale = (
+            req.sampling_params.guidance_scale if req.sampling_params.guidance_scale is not None else guidance_scale
+        )
+        generator = req.sampling_params.generator or generator
+        num_images_per_prompt = (
+            req.sampling_params.num_outputs_per_prompt
+            if req.sampling_params.num_outputs_per_prompt > 0
+            else num_images_per_prompt
+        )
+
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+        dtype = self.transformer.dtype
+
+        self._guidance_scale = guidance_scale
+        self._current_timestep = None
+        self._interrupt = False
+
+        self.check_inputs(prompt=prompt, height=height, width=width, prompt_embeds=prompt_embeds)
+        height, width = self._resize_dimensions(height, width)
+
+        if prompt is not None:
+            if isinstance(prompt, str):
+                prompt = [prompt]
+
+        if prompt_embeds is not None:
+            text_hiddens = prompt_embeds
+        else:
+            text_hiddens = self.encode_prompt(
+                prompt,
+                device,
+                num_images_per_prompt,
+                width=width,
+                height=height,
+                apply_pe=self._should_apply_pe(req),
+            )
+
+        if self.do_classifier_free_guidance:
+            if negative_prompt_embeds is not None:
+                uncond_text_hiddens = negative_prompt_embeds
+            else:
+                if negative_prompt is None:
+                    negative_prompt = ""
+                if isinstance(negative_prompt, str):
+                    negative_prompt = [negative_prompt] * batch_size
+                uncond_text_hiddens = self.encode_prompt(
+                    negative_prompt, device, num_images_per_prompt, width=width, height=height, apply_pe=False
+                )
+
+        latent_h = height // self.vae_scale_factor
+        latent_w = width // self.vae_scale_factor
+        latent_channels = self.transformer.config.in_channels
+
+        if latents is None:
+            latents = randn_tensor(
+                (batch_size * num_images_per_prompt, latent_channels, latent_h, latent_w),
+                generator=generator,
+                device=device,
+                dtype=dtype,
+            )
+
+        sigmas = torch.linspace(1.0, 0.0, num_inference_steps + 1)
+        self.scheduler.set_timesteps(sigmas=sigmas[:-1], device=device)
+
+        if self.do_classifier_free_guidance:
+            cfg_text_hiddens = list(uncond_text_hiddens) + list(text_hiddens)
+        else:
+            cfg_text_hiddens = text_hiddens
+
+        text_bth, text_lens = self._pad_text(
+            text_hiddens=cfg_text_hiddens, device=device, dtype=dtype, text_in_dim=self.transformer.config.text_in_dim
+        )
+
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(self.scheduler.timesteps):
+                if self.interrupt:
+                    continue
+
+                self._current_timestep = t
+
+                if self.do_classifier_free_guidance:
+                    latent_model_input = torch.cat([latents, latents], dim=0)
+                    t_batch = torch.full(
+                        (batch_size * num_images_per_prompt * 2,), t.item(), device=device, dtype=dtype
+                    )
+                else:
+                    latent_model_input = latents
+                    t_batch = torch.full((batch_size * num_images_per_prompt,), t.item(), device=device, dtype=dtype)
+
+                pred = self.transformer(
+                    hidden_states=latent_model_input,
+                    timestep=t_batch,
+                    text_bth=text_bth,
+                    text_lens=text_lens,
+                    return_dict=False,
+                )[0]
+
+                if self.do_classifier_free_guidance:
+                    pred_uncond, pred_cond = pred.chunk(2, dim=0)
+                    pred = pred_uncond + guidance_scale * (pred_cond - pred_uncond)
+
+                latents = self.scheduler.step(pred, t, latents, return_dict=False)[0]
+
+                if callback_on_step_end is not None:
+                    callback_kwargs = {k: locals()[k] for k in callback_on_step_end_tensor_inputs}
+                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+                    latents = callback_outputs.pop("latents", latents)
+
+                progress_bar.update()
+
+        self._current_timestep = None
+
+        if output_type == "latent":
+            return DiffusionOutput(output=latents)
+
+        bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(device)
+        bn_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + self.vae.config.batch_norm_eps).to(device)
+
+        latents = latents * bn_std + bn_mean
+
+        latents = self._unpatchify_latents(latents)
+
+        images = self.vae.decode(latents, return_dict=False)[0]
+
+        return DiffusionOutput(
+            output=images, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None
+        )
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loaded_weights = set()
+        transformer_weights = (
+            (name.replace("transformer.", "", 1), weight) for name, weight in weights if name.startswith("transformer.")
+        )
+        loaded_weights |= {f"transformer.{name}" for name in self.transformer.load_weights(transformer_weights)}
+        loaded_weights |= {f"vae.{name}" for name, _ in self.vae.named_parameters()}
+        loaded_weights |= {f"text_encoder.{name}" for name, _ in self.text_encoder.named_parameters()}
+        return loaded_weights

vllm_omni/diffusion/registry.py

@@ -165,6 +165,11 @@ _DIFFUSION_MODELS = {
         "pipeline_flux2_klein",
         "Flux2KleinPipeline",
     ),
+    "ErnieImagePipeline": (
+        "ernie_image",
+        "pipeline_ernie_image",
+        "ErnieImagePipeline",
+    ),
     "NextStep11Pipeline": (
         "nextstep_1_1",
         "pipeline_nextstep_1_1",
@@ -431,6 +436,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "StableDiffusion3Pipeline": "get_sd3_image_post_process_func",
     "FluxKontextPipeline": "get_flux_kontext_post_process_func",
     "Flux2KleinPipeline": "get_flux2_klein_post_process_func",
+    "ErnieImagePipeline": "get_ernie_image_post_process_func",
     "NextStep11Pipeline": "get_nextstep11_post_process_func",
     "FluxPipeline": "get_flux_post_process_func",
     "OmniGen2Pipeline": "get_omnigen2_post_process_func",