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

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

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)