From 4c06bb09ed9985da1c83c72ad704eb297c989986 Mon Sep 17 00:00:00 2001 From: Lancer Date: Sun, 3 May 2026 15:53:28 +0800 Subject: [PATCH] [Feat] ERNIE image model (T2I) (#2861) Signed-off-by: Lancer --- docs/models/supported_models.md | 2 +- docs/user_guide/diffusion_features.md | 1 + pyproject.toml | 4 + recipes/Baidu/ERNIE-Image.md | 149 ++++ recipes/README.md | 2 + .../diffusion/models/ernie_image/__init__.py | 0 .../models/ernie_image/test_ernie_image_pe.py | 135 ++++ .../ernie_image/test_ernie_image_rope.py | 187 +++++ .../models/ernie_image/test_ernie_image_tp.py | 66 ++ .../e2e/offline_inference/test_ernie_image.py | 58 ++ .../test_ernie_image_expansion.py | 63 ++ .../diffusion/cache/cache_dit_backend.py | 64 ++ .../diffusion/models/ernie_image/__init__.py | 14 + .../ernie_image/ernie_image_transformer.py | 646 ++++++++++++++++++ .../ernie_image/pipeline_ernie_image.py | 560 +++++++++++++++ vllm_omni/diffusion/registry.py | 6 + 16 files changed, 1956 insertions(+), 1 deletion(-) create mode 100644 recipes/Baidu/ERNIE-Image.md create mode 100644 tests/diffusion/models/ernie_image/__init__.py create mode 100644 tests/diffusion/models/ernie_image/test_ernie_image_pe.py create mode 100644 tests/diffusion/models/ernie_image/test_ernie_image_rope.py create mode 100644 tests/diffusion/models/ernie_image/test_ernie_image_tp.py create mode 100644 tests/e2e/offline_inference/test_ernie_image.py create mode 100644 tests/e2e/online_serving/test_ernie_image_expansion.py create mode 100644 vllm_omni/diffusion/models/ernie_image/__init__.py create mode 100644 vllm_omni/diffusion/models/ernie_image/ernie_image_transformer.py create mode 100644 vllm_omni/diffusion/models/ernie_image/pipeline_ernie_image.py diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index b81ff5b99..6b305f16c 100644 --- a/docs/models/supported_models.md +++ b/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. diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 3e9ddb942..e249d2e23 100644 --- a/docs/user_guide/diffusion_features.md +++ b/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. diff --git a/pyproject.toml b/pyproject.toml index d21f3c475..3c3f7a026 100644 --- a/pyproject.toml +++ b/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" diff --git a/recipes/Baidu/ERNIE-Image.md b/recipes/Baidu/ERNIE-Image.md new file mode 100644 index 000000000..b09b61c9c --- /dev/null +++ b/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: +- Upstream model card (Turbo): +- GitHub: +- Blog: + +## 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 ` — distribute + model weights across N GPUs. + - **SP (Sequence Parallelism):** `--usp ` (Ulysses SP) and `--ring ` + (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` diff --git a/recipes/README.md b/recipes/README.md index 2dcad5f98..01f76d1a8 100644 --- a/recipes/README.md +++ b/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 diff --git a/tests/diffusion/models/ernie_image/__init__.py b/tests/diffusion/models/ernie_image/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/diffusion/models/ernie_image/test_ernie_image_pe.py b/tests/diffusion/models/ernie_image/test_ernie_image_pe.py new file mode 100644 index 000000000..d1dc90715 --- /dev/null +++ b/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)) diff --git a/tests/diffusion/models/ernie_image/test_ernie_image_rope.py b/tests/diffusion/models/ernie_image/test_ernie_image_rope.py new file mode 100644 index 000000000..22d8efaad --- /dev/null +++ b/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" diff --git a/tests/diffusion/models/ernie_image/test_ernie_image_tp.py b/tests/diffusion/models/ernie_image/test_ernie_image_tp.py new file mode 100644 index 000000000..3688c7c61 --- /dev/null +++ b/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) diff --git a/tests/e2e/offline_inference/test_ernie_image.py b/tests/e2e/offline_inference/test_ernie_image.py new file mode 100644 index 000000000..c790bbf7c --- /dev/null +++ b/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) diff --git a/tests/e2e/online_serving/test_ernie_image_expansion.py b/tests/e2e/online_serving/test_ernie_image_expansion.py new file mode 100644 index 000000000..1f2447677 --- /dev/null +++ b/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) diff --git a/vllm_omni/diffusion/cache/cache_dit_backend.py b/vllm_omni/diffusion/cache/cache_dit_backend.py index 1b1800dac..5bf31108b 100644 --- a/vllm_omni/diffusion/cache/cache_dit_backend.py +++ b/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, } ) diff --git a/vllm_omni/diffusion/models/ernie_image/__init__.py b/vllm_omni/diffusion/models/ernie_image/__init__.py new file mode 100644 index 000000000..3a7374fd3 --- /dev/null +++ b/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"] diff --git a/vllm_omni/diffusion/models/ernie_image/ernie_image_transformer.py b/vllm_omni/diffusion/models/ernie_image/ernie_image_transformer.py new file mode 100644 index 000000000..aa9b0d8b7 --- /dev/null +++ b/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 diff --git a/vllm_omni/diffusion/models/ernie_image/pipeline_ernie_image.py b/vllm_omni/diffusion/models/ernie_image/pipeline_ernie_image.py new file mode 100644 index 000000000..7b5d80571 --- /dev/null +++ b/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 diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 37f519944..d7b058ef2 100644 --- a/vllm_omni/diffusion/registry.py +++ b/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",