[Refactor][Diffusion] Complete generic backend plan consumption (#7313)

Signed-off-by: specture724 <specture724@gmail.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>

tests/diffusion/offloader/test_layerwise_backend.py

@@ -24,6 +24,11 @@ from tests.diffusion.offloader.helpers import (
     patch_offload_runtime,
 )
 from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
+from vllm_omni.diffusion.offloader.block_discovery import (
+    get_blocks_attr_names,
+    get_blocks_from_dit,
+    set_blocks_attr_names,
+)
 from vllm_omni.diffusion.offloader.layerwise_backend import (
     LayerWiseOffloadBackend,
     LayerwiseOffloadHook,
@@ -209,21 +214,21 @@ class _NoAttrsModel(nn.Module):
 class TestGetBlocksFromDit:
     def test_get_blocks_from_dit_single_block_attr(self):
         model = _SingleBlockModel(num_blocks=3)
-        attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model)
+        attr_names, blocks = get_blocks_from_dit(model)
         assert attr_names == ["blocks"]
         assert len(blocks) == 3
         assert all(isinstance(b, _DummyBlock) for b in blocks)
 
     def test_get_blocks_from_dit_multi_block_attrs(self):
         model = _MultiBlockModel(num_transformer=2, num_single=3)
-        attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model)
+        attr_names, blocks = get_blocks_from_dit(model)
         assert set(attr_names) == {"transformer_blocks", "single_transformer_blocks"}
         assert len(blocks) == 5
         assert all(isinstance(b, _DummyBlock) for b in blocks)
 
     def test_get_blocks_from_dit_empty_blocks(self):
         model = _EmptyBlocksModel()
-        attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model)
+        attr_names, blocks = get_blocks_from_dit(model)
         assert attr_names == []
         assert blocks == []
 
@@ -233,17 +238,17 @@ class TestGetBlocksFromDit:
             AttributeError,
             match="Attribute 'nonexistent_blocks' declared in _layerwise_offload_blocks_attrs does not exist",
         ):
-            LayerWiseOffloadBackend.get_blocks_from_dit(model)
+            get_blocks_from_dit(model)
 
     def test_get_blocks_from_dit_no_attrs_defined(self):
         model = _NoAttrsModel(num_blocks=3)
-        attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model)
+        attr_names, blocks = get_blocks_from_dit(model)
         assert attr_names == []
         assert blocks == []
 
     def test_get_blocks_from_dit_deprecated_single_attr(self):
         model = _DeprecatedSingleAttrModel(num_blocks=2)
-        attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model)
+        attr_names, blocks = get_blocks_from_dit(model)
         assert attr_names == ["blocks"]
         assert len(blocks) == 2
 
@@ -251,17 +256,17 @@ class TestGetBlocksFromDit:
 class TestGetBlocksAttrNames:
     def test_get_blocks_attr_names_new_format(self):
         model = _MultiBlockModel()
-        attrs = LayerWiseOffloadBackend.get_blocks_attr_names(model)
+        attrs = get_blocks_attr_names(model)
         assert attrs == ["transformer_blocks", "single_transformer_blocks"]
 
     def test_get_blocks_attr_names_no_attrs(self):
         model = _NoAttrsModel()
-        attrs = LayerWiseOffloadBackend.get_blocks_attr_names(model)
+        attrs = get_blocks_attr_names(model)
         assert attrs == []
 
     def test_set_blocks_attr_names(self):
         model = _NoAttrsModel()
-        LayerWiseOffloadBackend.set_blocks_attr_names(model, ["new_blocks"])
+        set_blocks_attr_names(model, ["new_blocks"])
         assert hasattr(model.__class__, "_layerwise_offload_blocks_attrs")
         assert model.__class__._layerwise_offload_blocks_attrs == ["new_blocks"]