[Misc] Support WorkerWrapperBase and CustomPipeline for Diffusion Worker (#764)

Signed-off-by: knlnguyen1802 <knlnguyen1802@gmail.com>

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -70,6 +70,8 @@ class DiffusionModelRunner:
     def load_model(
         self,
         memory_pool_context_fn: callable | None = None,
+        load_format: str | None = None,
+        custom_pipeline_name: str | None = None,
     ) -> None:
         """
         Load the diffusion model, apply compilation and offloading.
@@ -77,7 +79,17 @@ class DiffusionModelRunner:
         Args:
             memory_pool_context_fn: Optional function that returns a context manager
                 for memory pool allocation (used for sleep mode).
+            load_format: Format for loading model weights. Supported formats:
+                - "default" (default): Automatically detect and use the default format based on configuration
+                - "custom_pipeline": Init model from a custom pipeline class specified by `custom_pipeline_name`
+                - "dummy": Skip actual weight loading, useful for testing and custom pipelines that
+                    don't require default weights.
+            custom_pipeline_name: Optional custom pipeline class name to use.
         """
+
+        if load_format == "dummy":
+            return
+
         load_device = (
             "cpu" if self.od_config.enable_cpu_offload or self.od_config.enable_layerwise_offload else str(self.device)
         )
@@ -97,6 +109,8 @@ class DiffusionModelRunner:
                 self.pipeline = model_loader.load_model(
                     od_config=self.od_config,
                     load_device=load_device,
+                    load_format=load_format,
+                    custom_pipeline_name=custom_pipeline_name,
                 )
         time_after_load = time.perf_counter()