[diffusion] perf: support zero-cost weight offload and overlap with compute for wan-series (#15511)

python/sglang/multimodal_gen/runtime/loader/component_loader.py

@@ -43,6 +43,9 @@ from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
     get_diffusers_component_config,
     get_hf_config,
 )
+from sglang.multimodal_gen.runtime.utils.layerwise_offload import (
+    LayerwiseOffloadManager,
+)
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
 
@@ -709,6 +712,23 @@ class TransformerLoader(ComponentLoader):
 
         model = model.eval()
 
+        if server_args.dit_layerwise_offload and hasattr(model, "blocks"):
+            try:
+                num_layers = len(getattr(model, "blocks"))
+            except Exception:
+                num_layers = None
+
+            if isinstance(num_layers, int) and num_layers > 0:
+                mgr = LayerwiseOffloadManager(
+                    model,
+                    module_list_attr="blocks",
+                    num_layers=num_layers,
+                    enabled=True,
+                    pin_cpu_memory=server_args.pin_cpu_memory,
+                    auto_initialize=True,
+                )
+                setattr(model, "_layerwise_offload_manager", mgr)
+
         return model