[Bugfix] T5 text encoder to render correct text in FLUX.1-dev (#2760)

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

tests/diffusion/models/t5_encoder/test_t5_encoder_prefix.py

@@ -0,0 +1,164 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for T5EncoderModel prefix handling and weight loading fix."""
+
+import pytest
+import torch
+from transformers import T5Config
+from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config
+
+from vllm_omni.diffusion.models.t5_encoder.t5_encoder import T5EncoderModel
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+_SMALL_T5_CONFIG = dict(
+    d_model=64,
+    d_kv=8,
+    d_ff=128,
+    num_heads=8,
+    num_layers=2,
+    vocab_size=256,
+    relative_attention_num_buckets=32,
+    relative_attention_max_distance=128,
+    is_gated_act=True,
+    dense_act_fn="gelu_new",
+    layer_norm_epsilon=1e-6,
+    feed_forward_proj="gated-gelu",
+)
+
+_T5_MODULE = "vllm_omni.diffusion.models.t5_encoder.t5_encoder"
+
+
+@pytest.fixture
+def t5_config() -> T5Config:
+    return T5Config(**_SMALL_T5_CONFIG)
+
+
+@pytest.fixture(scope="function", autouse=True)
+def setup_vllm_config(monkeypatch, mocker):
+    """Set up VllmConfig and TP=2 mocks for tests."""
+    device_config = DeviceConfig(device="cpu")
+
+    monkeypatch.setattr("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", lambda: 2)
+    monkeypatch.setattr(f"{_T5_MODULE}.get_tensor_model_parallel_world_size", lambda: 2)
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.vocab_parallel_embedding.get_tensor_model_parallel_world_size",
+        lambda: 2,
+    )
+
+    monkeypatch.setattr(f"{_T5_MODULE}.get_tensor_model_parallel_rank", lambda: 0)
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.vocab_parallel_embedding.get_tensor_model_parallel_rank",
+        lambda: 0,
+    )
+
+    mock_tp_group = mocker.MagicMock()
+    mock_tp_group.world_size = 2
+    mocker.patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_tp_group)
+
+    monkeypatch.setattr(f"{_T5_MODULE}.get_act_fn", lambda _: torch.nn.GELU())
+
+    with set_current_vllm_config(VllmConfig(device_config=device_config)):
+        yield
+
+
+class TestT5EncoderModelPrefixHandling:
+    """Test that T5EncoderModel correctly handles prefix attribute."""
+
+    def test_prefix_stored_in_model(self, t5_config):
+        """Test that prefix is stored in the model when provided."""
+        prefix = "text_encoder"
+        model = T5EncoderModel(t5_config, prefix=prefix)
+        assert hasattr(model, "prefix")
+        assert model.prefix == prefix
+
+    def test_prefix_empty_by_default(self, t5_config):
+        """Test that prefix defaults to empty string when not provided."""
+        model = T5EncoderModel(t5_config)
+        assert hasattr(model, "prefix")
+        assert model.prefix == ""
+
+
+class TestT5EncoderModelWeightLoadingWithPrefix:
+    """Test weight loading with prefix handling."""
+
+    def test_load_weights_with_prefix(self, t5_config):
+        """Test that weights without prefix are loaded when model has prefix."""
+        config = T5Config(**{**_SMALL_T5_CONFIG, "num_layers": 1})
+        model = T5EncoderModel(config, prefix="text_encoder")
+
+        inner_dim = config.num_heads * config.d_kv
+
+        weights = [
+            ("encoder.block.0.layer.0.SelfAttention.q.weight", torch.randn(inner_dim, config.d_model)),
+            ("encoder.block.0.layer.0.SelfAttention.k.weight", torch.randn(inner_dim, config.d_model)),
+            ("encoder.block.0.layer.0.SelfAttention.v.weight", torch.randn(inner_dim, config.d_model)),
+        ]
+
+        loaded = model.load_weights(weights)
+        assert len(loaded) > 0
+
+    def test_load_weights_embed_tokens_shared_sync(self, t5_config):
+        """Test that embed_tokens and shared weights are synced."""
+        model = T5EncoderModel(t5_config, prefix="text_encoder")
+
+        d_model = t5_config.d_model
+        vocab_size = t5_config.vocab_size
+
+        embed_weight = torch.randn(vocab_size, d_model)
+        weights = [
+            ("encoder.embed_tokens.weight", embed_weight.clone()),
+        ]
+
+        model.load_weights(weights)
+
+        shared_param = model.shared.weight
+        embed_param = model.encoder.embed_tokens.weight
+
+        assert torch.allclose(shared_param, embed_param), (
+            "shared and embed_tokens should have the same weights after loading"
+        )
+
+    def test_load_weights_shared_without_prefix(self, t5_config):
+        """Test shared.weight is recognized without relying on dot context."""
+        model = T5EncoderModel(t5_config, prefix="text_encoder")
+
+        shared_weight = torch.randn(t5_config.vocab_size, t5_config.d_model)
+        loaded = model.load_weights([("shared.weight", shared_weight)])
+
+        assert "shared.weight" in loaded
+        assert torch.allclose(model.shared.weight, model.encoder.embed_tokens.weight)
+
+    def test_unmatched_weights_are_not_reported_loaded(self, t5_config):
+        """Test that skipped checkpoint weights are not added to loaded_params."""
+        model = T5EncoderModel(t5_config, prefix="text_encoder")
+
+        loaded = model.load_weights(
+            [
+                (
+                    "text_encoder.encoder.block.0.layer.0.SelfAttention.missing.weight",
+                    torch.randn(t5_config.d_model, t5_config.d_model),
+                ),
+            ]
+        )
+
+        assert loaded == set()
+
+
+class TestT5EncoderModelWeightLoadingWithoutPrefix:
+    """Test weight loading without prefix."""
+
+    def test_load_weights_without_prefix(self, t5_config):
+        """Test that weights without prefix are loaded correctly."""
+        config = T5Config(**{**_SMALL_T5_CONFIG, "num_layers": 1})
+        model = T5EncoderModel(config)
+
+        inner_dim = config.num_heads * config.d_kv
+
+        weights = [
+            ("encoder.block.0.layer.0.SelfAttention.q.weight", torch.randn(inner_dim, config.d_model)),
+        ]
+
+        loaded = model.load_weights(weights)
+        assert len(loaded) > 0