Fix LTX2 channels-last review scope

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/config/test_omni_config.py

@@ -779,13 +779,8 @@ def test_structured_diffusion_config_rejects_non_boolean_compile_dynamic():
         omni_config_module._DiffusionConfigProjection(diffusion_compile_dynamic="false")
 
 
-def test_structured_diffusion_config_rejects_non_boolean_vae_channels_last_3d():
-    with pytest.raises(ValidationError, match="vae_use_channels_last_3d"):
-        omni_config_module._DiffusionConfigProjection(vae_use_channels_last_3d="true")
-
-
-def test_from_pipeline_config_routes_vae_channels_last_3d(tmp_path):
-    deploy_path = tmp_path / "dreamzero_vae_layout.yaml"
+def test_from_pipeline_config_routes_ltx2_model_extra(tmp_path):
+    deploy_path = tmp_path / "dreamzero_ltx2_extras.yaml"
     deploy_path.write_text(
         "\n".join(
             [
@@ -793,14 +788,15 @@ def test_from_pipeline_config_routes_vae_channels_last_3d(tmp_path):
                 "async_chunk: false",
                 "stages:",
                 "  - stage_id: 0",
-                "    vae_use_channels_last_3d: true",
+                "    extras:",
+                "      ltx2_use_channels_last_3d: true",
             ]
         )
     )
 
     stage = _from_pipeline_key("dreamzero", deploy_config_path=str(deploy_path)).stage_by_id(0)
 
-    assert stage.diffusion_config.vae_use_channels_last_3d is True
+    assert stage.diffusion_config.extras["ltx2_use_channels_last_3d"] is True
 
 
 def test_structured_diffusion_config_rejects_invalid_compile_granularity():

tests/diffusion/models/ltx2/test_ltx2_vae.py

@@ -12,6 +12,21 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 class TestLTXChannelsLast3D:
+    @pytest.mark.parametrize(
+        ("extras", "expected"),
+        [({}, False), ({"ltx2_use_channels_last_3d": True}, True)],
+    )
+    def test_channels_last_opt_in_is_ltx2_model_extra(self, extras, expected):
+        from vllm_omni.diffusion.models.ltx2.ltx2_components import _ltx2_use_channels_last_3d
+
+        assert _ltx2_use_channels_last_3d(SimpleNamespace(extras=extras)) is expected
+
+    def test_channels_last_opt_in_rejects_non_boolean_model_extra(self):
+        from vllm_omni.diffusion.models.ltx2.ltx2_components import _ltx2_use_channels_last_3d
+
+        with pytest.raises(TypeError, match="ltx2_use_channels_last_3d"):
+            _ltx2_use_channels_last_3d(SimpleNamespace(extras={"ltx2_use_channels_last_3d": "true"}))
+
     @pytest.mark.parametrize(
         ("causal", "kernel_size"),
         [

tests/diffusion/test_diffusion_engine.py

@@ -534,18 +534,6 @@ class TestDiffusionCompileConfig:
         with pytest.raises(TypeError, match="diffusion_compile_dynamic"):
             OmniDiffusionConfig(model="test", diffusion_compile_dynamic="false")
 
-    def test_channels_last_3d_is_opt_in_and_strict(self) -> None:
-        assert OmniDiffusionConfig(model="test").vae_use_channels_last_3d is False
-        assert (
-            OmniDiffusionConfig.from_kwargs(
-                model="test",
-                vae_use_channels_last_3d=True,
-            ).vae_use_channels_last_3d
-            is True
-        )
-        with pytest.raises(TypeError, match="vae_use_channels_last_3d"):
-            OmniDiffusionConfig(model="test", vae_use_channels_last_3d="true")
-
     @pytest.mark.parametrize(
         "kwargs, feature",
         [

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -20,7 +20,7 @@ def test_default_stage_config_includes_cache_backend():
             "cache_backend": "cache_dit",
             "cache_config": '{"Fn_compute_blocks": 2}',
             "vae_use_slicing": True,
-            "vae_use_channels_last_3d": True,
+            "extras": {"ltx2_use_channels_last_3d": True},
             "ulysses_degree": 2,
         }
     )[0]
@@ -30,7 +30,7 @@ def test_default_stage_config_includes_cache_backend():
     assert engine_args["cache_backend"] == "cache_dit"
     assert engine_args["cache_config"]["Fn_compute_blocks"] == 2
     assert engine_args["vae_use_slicing"] is True
-    assert engine_args["vae_use_channels_last_3d"] is True
+    assert engine_args["extras"]["ltx2_use_channels_last_3d"] is True
     assert engine_args["parallel_config"].ulysses_degree == 2
     assert engine_args["model_stage"] == "diffusion"
 
@@ -373,25 +373,6 @@ def test_serve_cli_accepts_diffusion_compile_controls():
     assert stage_cfg["engine_args"]["diffusion_compile_dynamic"] is False
 
 
-def test_serve_cli_accepts_vae_channels_last_3d_opt_in():
-    parser = TrackingArgumentParser()
-    subparsers = parser.add_subparsers(dest="command")
-    OmniServeCommand().subparser_init(subparsers)
-    args = parser.parse_args(
-        [
-            "serve",
-            "Lightricks/LTX-2",
-            "--omni",
-            "--vae-use-channels-last-3d",
-        ]
-    )
-
-    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(args.get_explicit_kwargs_dict())[0]
-
-    assert args.vae_use_channels_last_3d is True
-    assert stage_cfg["engine_args"]["vae_use_channels_last_3d"] is True
-
-
 def test_serve_cli_accepts_diffusion_attention_backend():
     """Ensure diffusion serve CLI exposes the shorthand backend flag."""
     parser = TrackingArgumentParser()

vllm_omni/config/omni_config.py

@@ -582,7 +582,6 @@ class _DiffusionConfigProjection:
     fa_deterministic: bool = False
     vae_use_slicing: bool = False
     vae_use_tiling: bool = False
-    vae_use_channels_last_3d: bool = Field(default=False, strict=True)
     mask_strategy_file_path: str | None = None
     skip_time_steps: int = 15
     VSA_sparsity: float = 0.0

vllm_omni/config/stage_config.py

@@ -408,7 +408,6 @@ class StageDeployConfig:
     step_execution: bool | None = None
     vae_use_slicing: bool | None = None
     vae_use_tiling: bool | None = None
-    vae_use_channels_last_3d: bool | None = None
     boundary_ratio: float | None = None
     flow_shift: float | None = None
     diffusion_kv_cache_dtype: str | None = None

vllm_omni/diffusion/data.py

@@ -713,8 +713,6 @@ class OmniDiffusionConfig:
     # VAE memory optimization parameters
     vae_use_slicing: bool = False
     vae_use_tiling: bool = False
-    # Opt in because CL3D performance varies materially by model shape and GPU.
-    vae_use_channels_last_3d: bool = False
 
     # STA (Sliding Tile Attention) parameters
     mask_strategy_file_path: str | None = None
@@ -934,8 +932,6 @@ class OmniDiffusionConfig:
             )
         if not isinstance(self.diffusion_compile_dynamic, bool):
             raise TypeError(f"diffusion_compile_dynamic must be a bool, got {type(self.diffusion_compile_dynamic)!r}")
-        if not isinstance(self.vae_use_channels_last_3d, bool):
-            raise TypeError(f"vae_use_channels_last_3d must be a bool, got {type(self.vae_use_channels_last_3d)!r}")
         self.diffusion_kv_mode = parse_diffusion_kv_cache_mode(self.diffusion_kv_mode)
 
         if self.omni_kv_config is None:

vllm_omni/diffusion/models/ltx2/ltx2_components.py

@@ -57,6 +57,17 @@ _LTX_COMPONENT_SUBFOLDERS = (
 )
 logger = logging.getLogger(__name__)
 
+_LTX2_CHANNELS_LAST_3D_EXTRA = "ltx2_use_channels_last_3d"
+
+
+def _ltx2_use_channels_last_3d(od_config: Any) -> bool:
+    """Read the LTX2-only channels-last opt-in from model extras."""
+    extras = getattr(od_config, "extras", {}) or {}
+    enabled = extras.get(_LTX2_CHANNELS_LAST_3D_EXTRA, False)
+    if not isinstance(enabled, bool):
+        raise TypeError(f"{_LTX2_CHANNELS_LAST_3D_EXTRA} must be a bool, got {type(enabled)!r}")
+    return enabled
+
 
 @dataclass(frozen=True)
 class LTXComponentProfile:
@@ -365,7 +376,7 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None:
         local_files_only=local_files_only,
         dtype=dtype,
     )
-    if getattr(od_config, "vae_use_channels_last_3d", False):
+    if _ltx2_use_channels_last_3d(od_config):
         pipeline.vae.enable_channels_last_3d()
     pipeline.audio_vae = _load_component(
         AutoencoderKLLTX2Audio,

vllm_omni/engine/arg_utils.py

@@ -542,7 +542,6 @@ class OrchestratorArgs:
     step_execution: bool = False
     vae_use_slicing: bool = False
     vae_use_tiling: bool = False
-    vae_use_channels_last_3d: bool = False
     enable_multithread_weight_load: bool = True
     num_weight_load_threads: int = 4
     enable_cpu_offload: bool = False

vllm_omni/engine/async_omni_engine.py

@@ -1029,7 +1029,6 @@ class AsyncOmniEngine:
             "request_batch_max_wait_ms": kwargs.get("request_batch_max_wait_ms", 0.0),
             "vae_use_slicing": kwargs.get("vae_use_slicing", False),
             "vae_use_tiling": kwargs.get("vae_use_tiling", False),
-            "vae_use_channels_last_3d": kwargs.get("vae_use_channels_last_3d", False),
             "cache_backend": cache_backend,
             "cache_config": cache_config,
             "enable_cache_dit_summary": kwargs.get("enable_cache_dit_summary", False),
@@ -1076,6 +1075,7 @@ class AsyncOmniEngine:
             "streaming_output": kwargs.get("diffusion_streaming_output", False),
             "enable_ar_profiler": kwargs.get("enable_ar_profiler", False),
             "extras": {
+                **(kwargs.get("extras") or {}),
                 "auxiliary_text_encoder": kwargs.get("auxiliary_text_encoder", None),
                 "default_llama_model_id": kwargs.get("default_llama_model_id", "meta-llama/Meta-Llama-3.1-8B-Instruct"),
             },

vllm_omni/entrypoints/cli/serve.py

@@ -668,15 +668,6 @@ class OmniServeCommand(CLISubcommand):
             action="store_true",
             help="Enable VAE tiling for memory optimization (useful for mitigating OOM issues).",
         )
-        omni_config_group.add_argument(
-            "--vae-use-channels-last-3d",
-            action="store_true",
-            help=(
-                "Use channels_last_3d Conv3d weights for supported VAEs. This is opt-in because performance "
-                "depends on model shape and GPU architecture."
-            ),
-        )
-
         # Parallel weight loading (faster diffusion startup)
         omni_config_group.add_argument(
             "--disable-multithread-weight-load",