[diffusion][feature] Add component-selective offload policies (#5929)

Signed-off-by: lishunyang12 <lishunyang12@163.com>

tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py

@@ -9,6 +9,7 @@ from PIL import Image
 from torch import nn
 
 from tests.diffusion.models.wan2_2.conftest import StubScheduler, StubTransformer, StubVAE, noop_progress_bar
+from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import _WAN_TEXT_ENCODER_OFFLOAD_PLAN
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v import (
     Wan22I2VPipeline,
     get_wan22_i2v_post_process_func,
@@ -30,6 +31,10 @@ def test_wan22_i2v_postprocess_honors_request_output_type() -> None:
     assert output is video
 
 
+def test_i2v_pipeline_declares_text_encoder_offload_blocks() -> None:
+    assert Wan22I2VPipeline._offload_plan is _WAN_TEXT_ENCODER_OFFLOAD_PLAN
+
+
 def _make_i2v_pipeline(*, expand_timesteps: bool) -> Wan22I2VPipeline:
     pipeline = object.__new__(Wan22I2VPipeline)
     nn.Module.__init__(pipeline)

tests/diffusion/models/wan2_2/test_wan22_pipeline_helpers.py

@@ -18,6 +18,15 @@ from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
 
+def test_t2v_pipeline_declares_text_encoder_offload_blocks() -> None:
+    plan = wan22_module.Wan22Pipeline._offload_plan
+
+    assert plan.encoder_component_types == {"text_encoder": "text_encoder"}
+    assert plan.encoder_block_attrs == {"text_encoder": ("encoder.block",)}
+    assert plan.encoder_dlo_weight_replication == frozenset({"text_encoder"})
+    assert plan.on_demand_component_paths == frozenset()
+
+
 class _LatentDist:
     def sample(self, generator):
         assert isinstance(generator, torch.Generator)

tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py

@@ -392,52 +392,31 @@ def test_encode_audio_skips_unshard_reshard_when_not_fsdp():
     assert "audio_emb" in result
 
 
-def test_s2v_pipeline_skips_cpu_offload_when_hsdp_enabled():
-    """Test that transformer.to('cpu') is NOT called when HSDP is active."""
-    from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import Wan22S2VPipeline
-
+@pytest.mark.parametrize(
+    ("components", "use_hsdp", "expected"),
+    [
+        (None, False, True),
+        (["dit"], False, True),
+        (["text_encoder"], False, False),
+        (["dit"], True, False),
+    ],
+)
+def test_s2v_dit_release_honors_component_selection(components, use_hsdp, expected):
     pipeline = object.__new__(Wan22S2VPipeline)
     nn.Module.__init__(pipeline)
+    compact = None if components is None else {"mode": "module", "components": components}
+    pipeline.od_config = SimpleNamespace(
+        diffusion_offload_config=compact,
+        enable_cpu_offload=True,
+        enable_layerwise_offload=False,
+        enable_distributed_layerwise_offload=False,
+        dlo_use_allgather=True,
+        dlo_resident_layers=0,
+        pin_cpu_memory=True,
+        parallel_config=SimpleNamespace(use_hsdp=use_hsdp),
+    )
 
-    od_config = MagicMock()
-    od_config.enable_cpu_offload = True
-    parallel_config = MagicMock()
-    parallel_config.use_hsdp = True
-    od_config.parallel_config = parallel_config
-    pipeline.od_config = od_config
-
-    mock_transformer = MagicMock()
-    pipeline.transformer = mock_transformer
-
-    # Simulate the offload decision from the forward loop
-    if pipeline.od_config.enable_cpu_offload and not getattr(pipeline.od_config.parallel_config, "use_hsdp", False):
-        pipeline.transformer.to("cpu")
-
-    mock_transformer.to.assert_not_called()
-
-
-def test_s2v_pipeline_allows_cpu_offload_when_hsdp_disabled():
-    """Test that transformer.to('cpu') IS called when HSDP is not active."""
-    from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import Wan22S2VPipeline
-
-    pipeline = object.__new__(Wan22S2VPipeline)
-    nn.Module.__init__(pipeline)
-
-    od_config = MagicMock()
-    od_config.enable_cpu_offload = True
-    parallel_config = MagicMock()
-    parallel_config.use_hsdp = False
-    od_config.parallel_config = parallel_config
-    pipeline.od_config = od_config
-
-    mock_transformer = MagicMock()
-    pipeline.transformer = mock_transformer
-
-    # Simulate the offload decision from the forward loop
-    if pipeline.od_config.enable_cpu_offload and not getattr(pipeline.od_config.parallel_config, "use_hsdp", False):
-        pipeline.transformer.to("cpu")
-
-    mock_transformer.to.assert_called_once_with("cpu")
+    assert pipeline._should_release_dit_before_decode() is expected
 
 
 def test_s2v_pipeline_hsdp_forward_complete_process():

tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from dataclasses import dataclass, field
 from types import SimpleNamespace
 
 import pytest
@@ -20,6 +21,46 @@ from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
 
+def test_vace_default_flow_shift_preserves_compact_offload_runtime_state(monkeypatch) -> None:
+    @dataclass
+    class RuntimeConfig:
+        flow_shift: float | None = None
+        diffusion_offload_config: dict[str, object] = field(
+            default_factory=lambda: {
+                "mode": "layer",
+                "components": ["text_encoder"],
+            }
+        )
+        enable_layerwise_offload: bool = True
+        post_init_calls: int = 0
+
+        def __post_init__(self) -> None:
+            self.post_init_calls += 1
+
+    config = RuntimeConfig()
+    resolved_offload = object()
+    config._resolved_diffusion_offload = resolved_offload
+    config._diffusion_offload_flags_materialized = True
+    captured = {}
+
+    def capture_base_init(self, *, od_config, prefix="") -> None:
+        captured.update(od_config=od_config, prefix=prefix)
+
+    monkeypatch.setattr(Wan22VACEPipeline.__mro__[1], "__init__", capture_base_init)
+
+    Wan22VACEPipeline(od_config=config, prefix="stage")
+
+    forwarded = captured["od_config"]
+    assert forwarded is not config
+    assert config.flow_shift is None
+    assert forwarded.flow_shift == 3.0
+    assert forwarded.diffusion_offload_config is config.diffusion_offload_config
+    assert forwarded._resolved_diffusion_offload is resolved_offload
+    assert forwarded._diffusion_offload_flags_materialized is True
+    assert forwarded.post_init_calls == 1
+    assert captured["prefix"] == "stage"
+
+
 def _make_vace_pipeline() -> Wan22VACEPipeline:
     pipeline = object.__new__(Wan22VACEPipeline)
     nn.Module.__init__(pipeline)