feat(diffusion): select Wan VAE decode mode automatically

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -32,6 +32,9 @@ from vllm_omni.diffusion.compile import regionally_compile
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
+    prepare_pipeline_wan_spatial_shard_decode,
+)
 from vllm_omni.diffusion.forward_context import set_forward_context
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.models.interface import (
@@ -266,6 +269,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                     custom_pipeline_name=custom_pipeline_name,
                     device=self.device,
                 )
+                prepare_pipeline_wan_spatial_shard_decode(self.pipeline)
         time_after_load = time.perf_counter()
 
         logger.info(