[Diffusion] Add FLUX.2-dev CPU Offload (Layerwise) (#5256)

Signed-off-by: dpeng123 <diyapeng123@gmail.com>

tests/diffusion/models/flux2/test_flux2_transformer_tp.py

@@ -253,3 +253,28 @@ class TestFlux2PackedModuleMapping:
         assert len(loaded_params) == 2, "Should load two weights"
         assert "x_embedder.weight" in loaded_params
         assert any("to_qkv" in p for p in loaded_params), "to_q should be mapped to to_qkv"
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+class TestFlux2TransformerLayerwiseOffload:
+    def test_layerwise_offload_blocks_attrs(self):
+        assert Flux2Transformer2DModel._layerwise_offload_blocks_attrs == [
+            "transformer_blocks",
+            "single_transformer_blocks",
+        ]
+
+    def test_get_blocks_from_dit(self):
+        from vllm_omni.diffusion.offloader.layerwise_backend import LayerWiseOffloadBackend
+
+        # Block discovery only needs the ModuleList structure, not production dims.
+        model = Flux2Transformer2DModel(
+            num_layers=2,
+            num_single_layers=2,
+            num_attention_heads=2,
+            attention_head_dim=4,
+            joint_attention_dim=16,
+        )
+        attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model)
+        assert attr_names == ["transformer_blocks", "single_transformer_blocks"]
+        assert len(blocks) == 4