From dc072a12bd08024c3b90e4c16299bfcda42f5c5d Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Wed, 12 Aug 2026 23:02:41 +0800 Subject: [PATCH] Fix LTX2 channels-last review scope Signed-off-by: Hongsheng Liu --- tests/config/test_omni_config.py | 14 ++++------- tests/diffusion/models/ltx2/test_ltx2_vae.py | 15 ++++++++++++ tests/diffusion/test_diffusion_engine.py | 12 ---------- .../test_async_omni_diffusion_config.py | 23 ++----------------- vllm_omni/config/omni_config.py | 1 - vllm_omni/config/stage_config.py | 1 - vllm_omni/diffusion/data.py | 4 ---- .../diffusion/models/ltx2/ltx2_components.py | 13 ++++++++++- vllm_omni/engine/arg_utils.py | 1 - vllm_omni/engine/async_omni_engine.py | 2 +- vllm_omni/entrypoints/cli/serve.py | 9 -------- 11 files changed, 35 insertions(+), 60 deletions(-) diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index 363aa7071..2dab0224d 100644 --- a/tests/config/test_omni_config.py +++ b/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(): diff --git a/tests/diffusion/models/ltx2/test_ltx2_vae.py b/tests/diffusion/models/ltx2/test_ltx2_vae.py index 4d69712db..ac612d315 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_vae.py +++ b/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"), [ diff --git a/tests/diffusion/test_diffusion_engine.py b/tests/diffusion/test_diffusion_engine.py index 5edb3c311..a0c6acb02 100644 --- a/tests/diffusion/test_diffusion_engine.py +++ b/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", [ diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 47045b1ef..8a4888057 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/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() diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index f0a88496f..4b5c73a39 100644 --- a/vllm_omni/config/omni_config.py +++ b/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 diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index b0f24f5d6..496aea208 100644 --- a/vllm_omni/config/stage_config.py +++ b/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 diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index db8ce6f4d..137fb9860 100644 --- a/vllm_omni/diffusion/data.py +++ b/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: diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_components.py b/vllm_omni/diffusion/models/ltx2/ltx2_components.py index 24ddef958..a01dc597d 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_components.py +++ b/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, diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 3bee3e277..a3aa96f60 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/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 diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index f7588ba4a..7bbaee67e 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/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"), }, diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index d65bfc8ff..8aa34a32b 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/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",