[Enhancement] Offload transformer after switch to transformer-2 (#3224)

Signed-off-by: Chendi Xue <chendi.xue@intel.com>
Co-authored-by: Canlin Guo <canlinguosdu@gmail.com>

docs/user_guide/diffusion/cpu_offload_diffusion.md

@@ -17,7 +17,7 @@ Both strategies use pinned memory for faster CPU-GPU transfers. The strategies a
 Model-level offloading implements mutual exclusion between DiT transformer and encoder modules using pre forward hooks:
 
 - **When encoders run**: DiT transformer is offloaded to CPU
-- **When DiT runs**: Encoders are offloaded to CPU
+- **When DiT runs**: Encoders are offloaded to CPU, if more than one dit models, only one loaded on GPU, others get offloaded to CPU.
 - **VAE**: Stays resident on GPU
 
 Before each module's forward pass, the hook automatically moves it to GPU while offloading the other module group to CPU. Transfers use pinned memory for speed.

vllm_omni/diffusion/offloader/sequential_backend.py

@@ -141,11 +141,12 @@ def apply_sequential_offload(
         ... )
         >>> # Modules of pipeline now automatically swap between CPU and GPU
     """
-    # Register hooks on DiT modules (offload encoders when DiT runs)
-    for dit_mod in dit_modules:
+    # Register hooks on DiT modules (offload encoders AND other DiTs when a DiT runs)
+    for i, dit_mod in enumerate(dit_modules):
+        other_dits = [d for j, d in enumerate(dit_modules) if j != i]
         registry = HookRegistry.get_or_create(dit_mod)
         hook = SequentialOffloadHook(
-            offload_targets=encoder_modules,
+            offload_targets=encoder_modules + other_dits,
             device=device,
             pin_memory=pin_memory,
             use_hsdp=use_hsdp,