From fb2e3535af3745d27f6932ba81439e784f39c459 Mon Sep 17 00:00:00 2001
From: shiyichuan <93317314+CarrotSwordsman@users.noreply.github.com>
Date: Mon, 21 Sep 2026 17:01:03 +0800
Subject: [PATCH] [Bugfix][Engine] Preserve text_encoder_tp_size in
 OmniEngineArgs (#7652)

Signed-off-by: mershi <mershi@tencent.com>
Co-authored-by: mershi <mershi@tencent.com>
---
 tests/engine/test_arg_utils.py | 37 ++++++++++++++++++++++++++++++++++
 vllm_omni/engine/arg_utils.py  |  6 ++++++
 2 files changed, 43 insertions(+)

diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py
index 7132d7512..d40a9c226 100644
--- a/tests/engine/test_arg_utils.py
+++ b/tests/engine/test_arg_utils.py
@@ -543,3 +543,40 @@ def test_tensor_parallel_size_none_is_handled():
     )
     assert isinstance(args, dict)
     assert "tensor_parallel_size" not in args
+
+
+# For https://github.com/vllm-project/vllm-omni/issues/7564
+def test_from_cli_args_preserves_text_encoder_tp_size():
+    """`--text-encoder-tp-size` must survive from_cli_args field filtering.
+
+    Library callers build engine args via ``OmniEngineArgs.from_cli_args``;
+    the dataclass field filter drops any namespace attribute the dataclass
+    does not declare, silently resetting the diffusion text-encoder TP to 1.
+    """
+    engine_args = OmniEngineArgs.from_cli_args(
+        SimpleNamespace(text_encoder_tp_size=2),
+    )
+    assert engine_args.text_encoder_tp_size == 2
+
+
+# For https://github.com/vllm-project/vllm-omni/issues/7564
+def test_text_encoder_tp_size_reaches_default_diffusion_parallel_config():
+    """The preserved CLI value must land in DiffusionParallelConfig.
+
+    Mirrors the library flow: from_cli_args -> kwargs -> the generic
+    diffusion fallback, which resolves the field through
+    ``DiffusionParallelConfig.from_stage_overrides``.
+    """
+    from dataclasses import asdict
+
+    from vllm_omni.config.config_factory import StageConfigFactory
+
+    engine_args = OmniEngineArgs.from_cli_args(
+        SimpleNamespace(text_encoder_tp_size=2),
+    )
+    stage_cfg = StageConfigFactory.create_default_diffusion(
+        {k: v for k, v in asdict(engine_args).items() if v is not None},
+    )[0]
+
+    parallel_config = stage_cfg["engine_args"]["parallel_config"]
+    assert parallel_config["text_encoder_tp_size"] == 2
diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py
index 939b19d50..8fb6c80e9 100644
--- a/vllm_omni/engine/arg_utils.py
+++ b/vllm_omni/engine/arg_utils.py
@@ -223,6 +223,12 @@ class OmniEngineArgs(EngineArgs):
     # Diffusion request-mode batch admission (forwarded to OmniDiffusionConfig).
     request_batch_max_wait_ms: float = 0.0
     fa_deterministic: bool = False
+    # Tensor-parallel degree for the diffusion text encoder (forwarded to
+    # DiffusionParallelConfig via the generic diffusion fallback). Declared
+    # here so ``from_cli_args`` field filtering keeps ``--text-encoder-tp-size``
+    # for library callers (#7564); registered pipelines may consume it through
+    # their own stage_cli_aliases or deploy YAML.
+    text_encoder_tp_size: int | None = None
 
     @classmethod
     def _add_omni_specific_args(cls, parser: argparse.ArgumentParser) -> argparse.ArgumentParser: