[Refactor] Migrate existing pipelines to use SupportsComponentDiscovery for offload discovery (#3076)

Signed-off-by: Nick Cao <ncao@redhat.com>
Co-authored-by: Claude <noreply@anthropic.com>

.buildkite/test-ready.yml

@@ -61,7 +61,7 @@ steps:
     steps:
       - label: "Diffusion · Offloader Test"
         commands:
-          - timeout 10m pytest -s -v tests/diffusion/offloader -m "core_model and cuda" --run-level "core_model"
+          - timeout 40m pytest -s -v tests/diffusion/offloader -m "core_model and cuda" --run-level "core_model"
         agents:
           queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
         plugins:

examples/offline_inference/sensenova_u1/README.md

@@ -145,7 +145,7 @@ python end2end.py \
 ## Reducing GPU Memory Usage
 
 For hardware with limited VRAM, enable **module-wise CPU offload** with
-`--enable-cpu-offload`. The pipeline implements `SupportsModuleOffload`, so the
+`--enable-cpu-offload`. The pipeline implements `SupportsComponentDiscovery`, so the
 vision encoder (`vision_model`) and the Qwen3 LLM (`language_model`) are
 swapped between CPU and GPU on demand:
 

tests/diffusion/offloader/test_diffusion_cpu_offload.py

@@ -18,6 +18,8 @@ AUDIO_MODEL = {
 
 IMAGE_VIDEO_MODELS = {
     "riverclouds/qwen_image_random": {"cuda": 2200, "rocm": 2100},
+    "Tongyi-MAI/Z-Image-Turbo": {"cuda": 2500, "rocm": 2100},
+    "OmniGen2/OmniGen2": {"cuda": 2500, "rocm": 2100},
 }
 
 MODELS = {**AUDIO_MODEL, **IMAGE_VIDEO_MODELS}

tests/e2e/offline_inference/test_magi_human_expansion.py

@@ -25,10 +25,24 @@ _OMNI_RUNNER_PARAM = (
     },
 )
 
+_OMNI_RUNNER_PARAM_OFFLOAD = (
+    MODEL,
+    None,
+    {
+        "init_timeout": 1200,
+        "tensor_parallel_size": 2,
+        "enable_cpu_offload": True,
+    },
+)
+
 pytestmark = [
     pytest.mark.full_model,
     pytest.mark.diffusion,
-    pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True),
+    pytest.mark.parametrize(
+        "omni_runner",
+        [_OMNI_RUNNER_PARAM, _OMNI_RUNNER_PARAM_OFFLOAD],
+        indirect=True,
+    ),
 ]
 
 _SKIP_ISSUE_3236 = pytest.mark.skip(

tests/e2e/offline_inference/test_stable_audio_expansion.py

@@ -103,3 +103,23 @@ def test_stable_audio_quantization_and_teacache() -> None:
         )
     finally:
         m.close()
+
+
+@hardware_test(res={"cuda": "L4", "xpu": "B60"})
+def test_stable_audio_cpu_offload() -> None:
+    """Stable Audio Open with FP8 + CPU offload."""
+    m = Omni(
+        model="stabilityai/stable-audio-open-1.0",
+        quantization="fp8",
+        enable_cpu_offload=True,
+    )
+    try:
+        audio = generate_stable_audio_short_clip(m)
+        assert_audio_valid(
+            audio,
+            sample_rate=_SAMPLE_RATE,
+            channels=2,
+            duration_s=_CLIP_DURATION_S,
+        )
+    finally:
+        m.close()

tests/e2e/online_serving/test_bagel_expansion.py

@@ -65,6 +65,17 @@ def _get_diffusion_feature_cases(model: str):
     """
 
     return [
+        # CPU offload (single-card)
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--enable-cpu-offload",
+                ],
+            ),
+            id="single_card_cpu_offload",
+            marks=SINGLE_CARD_FEATURE_MARKS,
+        ),
         # TeaCache (single-card)
         pytest.param(
             OmniServerParams(

tests/e2e/online_serving/test_flux2_expansion.py

@@ -31,6 +31,19 @@ NEGATIVE_PROMPT = "blurry, low quality"
 # cross-feature compatibility more generally.
 def _get_diffusion_feature_cases(model: str):
     return [
+        # CPU offload + TP=2 (model needs TP=2)
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--enable-cpu-offload",
+                    "--tensor-parallel-size",
+                    "2",
+                ],
+            ),
+            id="cpu_offload",
+            marks=FOUR_CARD_FEATURE_MARKS,
+        ),
         # FP8 / Hybrid sequence parallelism
         pytest.param(
             OmniServerParams(

tests/e2e/online_serving/test_flux2_klein_expansion.py

@@ -27,6 +27,13 @@ _NUM_INFERENCE_STEPS = 4
 
 def _get_diffusion_feature_cases(model: str):
     return [
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=["--tensor-parallel-size", "2", "--enable-cpu-offload"],
+            ),
+            id="tp2_cpu_offload",
+        ),
         pytest.param(
             OmniServerParams(
                 model=model,

tests/e2e/online_serving/test_hunyuan_video_15_expansion.py

@@ -31,6 +31,17 @@ def _get_diffusion_feature_cases(model: str):
     Designed for 2x H100 environment per issue #1832.
     """
     return [
+        # (1 GPU) CPU offload
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--enable-cpu-offload",
+                ],
+            ),
+            id="single_card_cpu_offload",
+            marks=SINGLE_CARD_MARKS,
+        ),
         # (1 GPU) CacheDiT + Layerwise CPU offloading
         pytest.param(
             OmniServerParams(

tests/e2e/online_serving/test_nextstep_expansion.py

@@ -31,6 +31,20 @@ _DEFAULT_MODEL = "stepfun-ai/NextStep-1.1"
 def _get_diffusion_feature_cases(model: str):
     """Single online config: TP=4, explicit pipeline class."""
     return [
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--tensor-parallel-size",
+                    "2",
+                    "--model-class-name",
+                    "NextStep11Pipeline",
+                    "--enable-cpu-offload",
+                ],
+            ),
+            id="nextstep_tp2_cpu_offload",
+            marks=FOUR_CARD_MARKS,
+        ),
         pytest.param(
             OmniServerParams(
                 model=model,

tests/e2e/online_serving/test_qwen_image_edit_expansion.py

@@ -29,6 +29,16 @@ PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
 # If a similar test only involves one model, one can just define a global list variable.
 def _get_diffusion_feature_cases(model: str):
     return [
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--enable-cpu-offload",
+                ],
+            ),
+            id="cpu_offload",
+            marks=SINGLE_CARD_FEATURE_MARKS,
+        ),
         pytest.param(
             OmniServerParams(
                 model=model,

tests/e2e/online_serving/test_qwen_image_expansion.py

@@ -25,6 +25,14 @@ PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
 
 def _get_diffusion_feature_cases(model: str):
     return [
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=["--enable-cpu-offload"],
+            ),
+            id="cpu_offload",
+            marks=SINGLE_CARD_FEATURE_MARKS,
+        ),
         pytest.param(
             OmniServerParams(model=model, server_args=["--step-execution"]),
             id="step_execution",

tests/e2e/online_serving/test_qwen_image_layered_expansion.py

@@ -32,6 +32,14 @@ SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
 # sp_001          : cache_dit + Ulysses-SP 2
 # cfg_parallel_002: cache_dit + CFG-Parallel 2
 FEATURE_CASES = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            server_args=["--enable-cpu-offload"],
+        ),
+        id="cpu_offload",
+        marks=SINGLE_CARD_FEATURE_MARKS,
+    ),
     pytest.param(
         OmniServerParams(
             model=MODEL,

tests/e2e/online_serving/test_sd3_expansion.py

@@ -24,6 +24,20 @@ NEGATIVE_PROMPT = "blurry, low quality, distorted"
 # that provide good performance improvements.
 def _get_diffusion_feature_cases(model: str):
     return [
+        # CFG Parallel + Tensor Parallel + CPU offload
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--cfg-parallel-size",
+                    "2",
+                    "--tensor-parallel-size",
+                    "2",
+                    "--enable-cpu-offload",
+                ],
+            ),
+            marks=FOUR_CARD_FEATURE_MARKS,
+        ),
         # Cache-DiT + CFG Parallel + Tensor Parallel
         pytest.param(
             OmniServerParams(

tests/e2e/online_serving/test_wan22_expansion.py

@@ -85,7 +85,22 @@ def _get_wan22_feature_cases():
     """
     cases = []
 
-    # ---- CUDA cases (unchanged) ----
+    # ---- CUDA cases ----
+
+    # Single-card: CPU offload (applies to all models)
+    for model_path, model_key in WAN22_MODELS:
+        cases.append(
+            pytest.param(
+                OmniServerParams(
+                    model=model_path,
+                    server_args=["--enable-cpu-offload"],
+                ),
+                id=f"{model_key}_cpu_offload",
+                marks=CUDA_SINGLE_CARD_MARKS,
+            )
+        )
+
+    # Single-card: Cache-DiT (applies to all models)
     for model_path, model_key in WAN22_MODELS:
         cases.append(
             pytest.param(

tests/e2e/online_serving/test_wan_2_1_vace_expansion.py

@@ -37,6 +37,18 @@ PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
 
 def _get_vace_feature_cases():
     return [
+        # Single GPU: CPU offload
+        pytest.param(
+            OmniServerParams(
+                model=MODEL,
+                server_args=[
+                    "--enable-cpu-offload",
+                    "--vae-use-tiling",
+                ],
+            ),
+            id="single_card_cpu_offload",
+            marks=SINGLE_CARD_FEATURE_MARKS,
+        ),
         # Single GPU: Cache-DiT + layerwise CPU offload
         pytest.param(
             OmniServerParams(

vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py

@@ -5,6 +5,7 @@ import logging
 import math
 import os
 from collections.abc import Iterable
+from typing import ClassVar
 
 import torch
 import torch.distributed
@@ -23,7 +24,7 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
-from vllm_omni.diffusion.models.interface import SupportAudioInput, SupportImageInput
+from vllm_omni.diffusion.models.interface import SupportAudioInput, SupportImageInput, SupportsComponentDiscovery
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 
 try:
@@ -95,9 +96,15 @@ VIDEO_CONFIG = {
 }
 
 
-class DreamIDOmniPipeline(nn.Module, CFGParallelMixin, SupportImageInput, SupportAudioInput):
+class DreamIDOmniPipeline(
+    nn.Module, CFGParallelMixin, SupportImageInput, SupportAudioInput, SupportsComponentDiscovery
+):
     """DreamID-Omni pipeline for vLLM-Omni."""
 
+    _dit_modules: ClassVar[list[str]] = ["model"]
+    _encoder_modules: ClassVar[list[str]] = ["text_model"]
+    _vae_modules: ClassVar[list[str]] = ["vae_model_video", "vae_model_audio"]
+
     def __init__(
         self,
         *,

vllm_omni/diffusion/models/flux/pipeline_flux.py

@@ -5,7 +5,7 @@ import json
 import logging
 import os
 from collections.abc import Iterable
-from typing import Any
+from typing import Any, ClassVar
 
 import numpy as np
 import torch
@@ -28,6 +28,7 @@ from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineL
 from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
 from vllm_omni.diffusion.models.flux import FluxTransformer2DModel
 from vllm_omni.diffusion.models.flux.flux_pipeline_mixin import FluxPipelineMixin
+from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
 from vllm_omni.diffusion.models.t5_encoder import T5EncoderModel
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -66,7 +67,13 @@ def get_flux_post_process_func(
     return post_process_func
 
 
-class FluxPipeline(nn.Module, FluxPipelineMixin, CFGParallelMixin, DiffusionPipelineProfilerMixin):
+class FluxPipeline(
+    nn.Module, FluxPipelineMixin, CFGParallelMixin, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery
+):
+    _dit_modules: ClassVar[list[str]] = ["transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["text_encoder", "text_encoder_2"]
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+
     def __init__(
         self,
         *,

vllm_omni/diffusion/models/flux/pipeline_flux_kontext.py

@@ -9,7 +9,7 @@ from __future__ import annotations
 import json
 import os
 from collections.abc import Callable, Iterable
-from typing import Any, cast
+from typing import Any, ClassVar, cast
 
 import numpy as np
 import PIL.Image
@@ -30,7 +30,7 @@ from vllm_omni.diffusion.models.flux import (
     FluxKontextTransformer2DModel,
 )
 from vllm_omni.diffusion.models.flux.flux_pipeline_mixin import FluxPipelineMixin
-from vllm_omni.diffusion.models.interface import SupportImageInput
+from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -70,10 +70,19 @@ def get_flux_kontext_post_process_func(od_config: OmniDiffusionConfig) -> Callab
 
 
 class FluxKontextPipeline(
-    nn.Module, FluxPipelineMixin, SupportImageInput, ProgressBarMixin, DiffusionPipelineProfilerMixin
+    nn.Module,
+    FluxPipelineMixin,
+    SupportImageInput,
+    ProgressBarMixin,
+    DiffusionPipelineProfilerMixin,
+    SupportsComponentDiscovery,
 ):
     """FLUX.1-Kontext pipeline for image editing with text guidance."""
 
+    _dit_modules: ClassVar[list[str]] = ["transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["text_encoder", "text_encoder_2"]
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+
     support_image_input = True
 
     def __init__(

vllm_omni/diffusion/models/flux2/pipeline_flux2.py

@@ -6,7 +6,7 @@ import logging
 import math
 import os
 from collections.abc import Callable, Iterable
-from typing import Any, cast
+from typing import Any, ClassVar, cast
 
 import numpy as np
 import PIL.Image
@@ -30,7 +30,7 @@ from vllm_omni.diffusion.distributed.parallel_state import get_classifier_free_g
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.models.flux2 import Flux2Transformer2DModel
-from vllm_omni.diffusion.models.interface import SupportImageInput
+from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery
 from vllm_omni.diffusion.models.mistral_encoder import MistralEncoderModel
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
@@ -336,9 +336,20 @@ def retrieve_latents(encoder_output: torch.Tensor, generator: torch.Generator =
         raise AttributeError("Could not access latents of provided encoder_output")
 
 
-class Flux2Pipeline(nn.Module, CFGParallelMixin, SupportImageInput, ProgressBarMixin, DiffusionPipelineProfilerMixin):
+class Flux2Pipeline(
+    nn.Module,
+    CFGParallelMixin,
+    SupportImageInput,
+    ProgressBarMixin,
+    DiffusionPipelineProfilerMixin,
+    SupportsComponentDiscovery,
+):
     """Flux2 pipeline for text-to-image generation."""
 
+    _dit_modules: ClassVar[list[str]] = ["transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+
     _callback_tensor_inputs = ["latents", "prompt_embeds"]
 
     support_image_input = True

vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py

@@ -19,7 +19,7 @@ import json
 import math
 import os
 from collections.abc import Callable, Iterable
-from typing import Any, cast
+from typing import Any, ClassVar, cast
 
 import numpy as np
 import PIL.Image
@@ -43,7 +43,7 @@ from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_p
 from vllm_omni.diffusion.models.flux2_klein.flux2_klein_transformer import (
     Flux2Transformer2DModel,
 )
-from vllm_omni.diffusion.models.interface import SupportImageInput
+from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
@@ -181,9 +181,15 @@ def compute_empirical_mu(image_seq_len: int, num_steps: int) -> float:
     return float(mu)
 
 
-class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, DiffusionPipelineProfilerMixin):
+class Flux2KleinPipeline(
+    nn.Module, CFGParallelMixin, SupportImageInput, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery
+):
     """Flux2 klein pipeline for text-to-image generation."""
 
+    _dit_modules: ClassVar[list[str]] = ["transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+
     support_image_input = True
 
     def __init__(

vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py

@@ -17,7 +17,7 @@ import logging
 import os
 import re
 from collections.abc import Iterable
-from typing import cast
+from typing import ClassVar, cast
 
 import numpy as np
 import PIL.Image
@@ -46,6 +46,7 @@ from vllm_omni.diffusion.models.glm_image.glm_image_transformer import (
     GlmImageKVCache,
     GlmImageTransformer2DModel,
 )
+from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.inputs.data import OmniTextPrompt
@@ -238,7 +239,7 @@ def retrieve_latents(
         raise AttributeError("Could not access latents of provided encoder_output")
 
 
-class GlmImagePipeline(nn.Module, DiffusionPipelineProfilerMixin):
+class GlmImagePipeline(nn.Module, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery):
     """
     GLM-Image Pipeline for text-to-image and image-to-image generation.
 
@@ -255,6 +256,10 @@ class GlmImagePipeline(nn.Module, DiffusionPipelineProfilerMixin):
     4. VAE decodes final latents to image
     """
 
+    _dit_modules: ClassVar[list[str]] = ["transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+
     def __init__(
         self,
         *,

vllm_omni/diffusion/models/helios/pipeline_helios.py

@@ -9,7 +9,7 @@ import logging
 import math
 import os
 from collections.abc import Iterable
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, ClassVar
 
 import numpy as np
 import torch
@@ -26,6 +26,7 @@ from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.models.helios.helios_transformer import HeliosTransformer3DModel
 from vllm_omni.diffusion.models.helios.scheduling_helios import HeliosScheduler
+from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -151,13 +152,19 @@ def get_helios_pre_process_func(
     return pre_process_func
 
 
-class HeliosPipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipelineProfilerMixin):
+class HeliosPipeline(
+    nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery
+):
     """Helios text-to-video / image-to-video / video-to-video pipeline for vllm-omni.
 
     Supports T2V, I2V (with image input), and V2V (with video input).
     Implements chunked video generation with multi-term memory history context.
     """
 
+    _dit_modules: ClassVar[list[str]] = ["transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+
     def __init__(
         self,
         *,

vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py

@@ -24,7 +24,7 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
-from vllm_omni.diffusion.models.interface import SupportImageInput
+from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import (
     DiffusionPipelineProfilerMixin,
 )
@@ -336,10 +336,15 @@ class HunyuanImage3Pipeline(
     HunyuanImage3PreTrainedModel,
     GenerationMixin,
     SupportImageInput,
+    SupportsComponentDiscovery,
     DiffusionPipelineProfilerMixin,
 ):
     supports_step_execution: ClassVar[bool] = True
     support_image_input = True
+    _dit_modules: ClassVar[list[str]] = ["model"]
+    _encoder_modules: ClassVar[list[str]] = ["vision_model"]
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+
     hf_to_vllm_mapper = WeightsMapper(
         orig_to_new_prefix={
             "model.": "",

vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5.py

@@ -7,7 +7,7 @@ import logging
 import os
 import re
 from collections.abc import Iterable
-from typing import Any
+from typing import Any, ClassVar
 
 import numpy as np
 import torch
@@ -27,6 +27,7 @@ from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
 from vllm_omni.diffusion.models.hunyuan_video.hunyuan_video_15_transformer import HunyuanVideo15Transformer3DModel
+from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
 from vllm_omni.diffusion.models.t5_encoder import T5EncoderModel
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
@@ -86,7 +87,13 @@ def get_hunyuan_video_15_post_process_func(od_config: OmniDiffusionConfig):
     return post_process_func
 
 
-class HunyuanVideo15Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipelineProfilerMixin):
+class HunyuanVideo15Pipeline(
+    nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery
+):
+    _dit_modules: ClassVar[list[str]] = ["transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["text_encoder", "text_encoder_2"]
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+
     def __init__(
         self,
         *,

vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5_i2v.py

@@ -6,7 +6,7 @@ from __future__ import annotations
 import logging
 import os
 from collections.abc import Iterable
-from typing import Any, cast
+from typing import Any, ClassVar, cast
 
 import numpy as np
 import PIL.Image
@@ -40,7 +40,7 @@ from vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5 import
     get_hunyuan_video_15_post_process_func,
     retrieve_latents,
 )
-from vllm_omni.diffusion.models.interface import SupportImageInput
+from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
 from vllm_omni.diffusion.models.t5_encoder import T5EncoderModel
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
@@ -104,8 +104,17 @@ def get_hunyuan_video_15_i2v_pre_process_func(od_config: OmniDiffusionConfig):
 
 
 class HunyuanVideo15I2VPipeline(
-    nn.Module, CFGParallelMixin, SupportImageInput, ProgressBarMixin, DiffusionPipelineProfilerMixin
+    nn.Module,
+    CFGParallelMixin,
+    SupportImageInput,
+    ProgressBarMixin,
+    DiffusionPipelineProfilerMixin,
+    SupportsComponentDiscovery,
 ):
+    _dit_modules: ClassVar[list[str]] = ["transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["text_encoder", "text_encoder_2", "image_encoder"]
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+
     support_image_input = True
     color_format = "RGB"
 

vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image.py

@@ -9,7 +9,7 @@ import os
 import re
 from collections.abc import Iterable
 from functools import partial
-from typing import Any
+from typing import Any, ClassVar
 
 import numpy as np
 import torch
@@ -28,6 +28,7 @@ from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
+from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
 from vllm_omni.diffusion.models.longcat_image.longcat_image_transformer import LongCatImageTransformer2DModel
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -201,7 +202,11 @@ def get_prompt_language(prompt):
     return "en"
 
 
-class LongCatImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMixin):
+class LongCatImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery):
+    _dit_modules: ClassVar[list[str]] = ["transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+
     def __init__(
         self,
         *,

vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image_edit.py

@@ -6,7 +6,7 @@ import math
 import os
 import re
 from collections.abc import Iterable
-from typing import Any, cast
+from typing import Any, ClassVar, cast
 
 import numpy as np
 import PIL.Image
@@ -30,7 +30,7 @@ from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
-from vllm_omni.diffusion.models.interface import SupportImageInput
+from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery
 from vllm_omni.diffusion.models.longcat_image.longcat_image_transformer import (
     LongCatImageTransformer2DModel,
 )
@@ -222,7 +222,13 @@ def split_quotation(prompt, quote_pairs=None):
     return result
 
 
-class LongCatImageEditPipeline(nn.Module, CFGParallelMixin, SupportImageInput, DiffusionPipelineProfilerMixin):
+class LongCatImageEditPipeline(
+    nn.Module, CFGParallelMixin, SupportImageInput, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery
+):
+    _dit_modules: ClassVar[list[str]] = ["transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+
     def __init__(
         self,
         *,

vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py

@@ -157,7 +157,11 @@ class _VideoAudioScheduler:
         return ((video_out, audio_out),)
 
 
-class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
+class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, SupportsComponentDiscovery):
+    _dit_modules: ClassVar[list[str]] = ["transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
+    _vae_modules: ClassVar[list[str]] = ["vae", "audio_vae"]
+
     # Audio is diffused jointly with video; warmup must size audio tokens.
     dummy_run_num_frames = 2
 

vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py

@@ -12,7 +12,7 @@ import os
 from collections.abc import Iterable
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Literal
+from typing import Any, ClassVar, Literal
 
 import numpy as np
 import torch
@@ -46,6 +46,7 @@ from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
 from vllm_omni.diffusion.model_loader.diffusers_loader import (
     DiffusersPipelineLoader,
 )
+from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
 from vllm_omni.diffusion.models.t5_encoder.t5_gemma_encoder import T5GemmaEncoderModelTP
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import (
@@ -1679,7 +1680,11 @@ def _resolve_subdir(
 # ===========================================================================
 # Main Pipeline
 # ===========================================================================
-class MagiHumanPipeline(nn.Module, ProgressBarMixin, DiffusionPipelineProfilerMixin):
+class MagiHumanPipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery, DiffusionPipelineProfilerMixin):
+    _dit_modules: ClassVar[list[str]] = ["dit", "sr_dit"]
+    _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
+    _vae_modules: ClassVar[list[str]] = ["vae", "audio_vae"]
+
     def __init__(self, od_config: OmniDiffusionConfig, **kwargs):
         super().__init__()
         model_path = od_config.model

vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from collections.abc import Iterable
-from typing import Any
+from typing import Any, ClassVar
 
 import torch
 from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
@@ -11,6 +11,7 @@ from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm
 from vllm.config import VllmConfig
 from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper
 
+from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
 from vllm_omni.model_executor.models.output_templates import OmniOutput
 from vllm_omni.transformers_utils.configs.mammoth_moda2 import Mammothmoda2Config
 
@@ -19,7 +20,7 @@ from .rope_real import RotaryPosEmbedReal
 from .schedulers import FlowMatchEulerDiscreteScheduler
 
 
-class MammothModa2DiTPipeline(nn.Module):
+class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
     """
     MammothModa2 DiT + VAE generation stage (non-autoregressive).
 
@@ -28,6 +29,10 @@ class MammothModa2DiTPipeline(nn.Module):
 
     """
 
+    _dit_modules: ClassVar[list[str]] = ["gen_transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["gen_image_condition_refiner"]
+    _vae_modules: ClassVar[list[str]] = ["gen_vae"]
+
     have_multimodal_outputs = True
 
     # Load only gen_* weights; ignore llm_model.* to prevent loading the entire LLM backbone in the DiT stage.

vllm_omni/diffusion/models/nextstep_1_1/pipeline_nextstep_1_1.py

@@ -5,7 +5,7 @@
 import os
 import re
 from collections.abc import Iterable
-from typing import Literal
+from typing import ClassVar, Literal
 
 import numpy as np
 import torch
@@ -28,6 +28,7 @@ from vllm_omni.diffusion.distributed.parallel_state import (
 )
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
 from vllm_omni.diffusion.models.nextstep_1_1.modeling_flux_vae import AutoencoderKL
 from vllm_omni.diffusion.models.nextstep_1_1.modeling_nextstep import (
     NextStepConfig,
@@ -131,7 +132,7 @@ def get_nextstep11_post_process_func(od_config: OmniDiffusionConfig):
     return post_process_func
 
 
-class NextStep11Pipeline(nn.Module, DiffusionPipelineProfilerMixin):
+class NextStep11Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProfilerMixin):
     """
     NextStep-1.1 Pipeline for text-to-image generation.
 
@@ -140,6 +141,10 @@ class NextStep11Pipeline(nn.Module, DiffusionPipelineProfilerMixin):
     to generate images autoregressively.
     """
 
+    _dit_modules: ClassVar[list[str]] = ["model"]
+    _encoder_modules: ClassVar[list[str]] = []
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+
     def __init__(
         self,
         *,

vllm_omni/diffusion/models/ovis_image/pipeline_ovis_image.py

@@ -19,7 +19,7 @@ import inspect
 import json
 import os
 from collections.abc import Callable, Iterable
-from typing import Any
+from typing import Any, ClassVar
 
 import numpy as np
 import torch
@@ -39,6 +39,7 @@ from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
+from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
 from vllm_omni.diffusion.models.ovis_image.ovis_image_transformer import OvisImageTransformer2DModel
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -142,7 +143,11 @@ def retrieve_timesteps(
     return timesteps, num_inference_steps
 
 
-class OvisImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMixin):
+class OvisImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery):
+    _dit_modules: ClassVar[list[str]] = ["transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+
     def __init__(
         self,
         *,

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py

@@ -28,6 +28,7 @@ from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
 from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
+from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
 from vllm_omni.diffusion.models.qwen_image.cfg_parallel import (
     QwenImageCFGParallelMixin,
 )
@@ -251,7 +252,13 @@ def apply_rotary_emb_qwen(
         return x_out.type_as(x)
 
 
-class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineProfilerMixin):
+class QwenImagePipeline(
+    nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery
+):
+    _dit_modules: ClassVar[list[str]] = ["transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+
     supports_step_execution: ClassVar[bool] = True
 
     def __init__(

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py

@@ -7,7 +7,7 @@ import logging
 import math
 import os
 from collections.abc import Iterable
-from typing import Any, cast
+from typing import Any, ClassVar, cast
 
 import numpy as np
 import PIL.Image
@@ -28,7 +28,7 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
-from vllm_omni.diffusion.models.interface import SupportImageInput
+from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery
 from vllm_omni.diffusion.models.qwen_image.cfg_parallel import (
     QwenImageCFGParallelMixin,
 )
@@ -222,7 +222,13 @@ def retrieve_latents(
         raise AttributeError("Could not access latents of provided encoder_output")
 
 
-class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMixin, DiffusionPipelineProfilerMixin):
+class QwenImageEditPipeline(
+    nn.Module, SupportImageInput, QwenImageCFGParallelMixin, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery
+):
+    _dit_modules: ClassVar[list[str]] = ["transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+
     def __init__(
         self,
         *,

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py

@@ -5,7 +5,7 @@ import json
 import logging
 import os
 from collections.abc import Iterable
-from typing import Any, cast
+from typing import Any, ClassVar, cast
 
 import numpy as np
 import PIL.Image
@@ -27,7 +27,7 @@ from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
 from vllm_omni.diffusion.model_metadata import QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES
-from vllm_omni.diffusion.models.interface import SupportImageInput
+from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery
 from vllm_omni.diffusion.models.qwen_image.cfg_parallel import (
     QwenImageCFGParallelMixin,
 )
@@ -186,8 +186,12 @@ def get_qwen_image_edit_plus_post_process_func(
 
 
 class QwenImageEditPlusPipeline(
-    nn.Module, SupportImageInput, QwenImageCFGParallelMixin, DiffusionPipelineProfilerMixin
+    nn.Module, SupportImageInput, QwenImageCFGParallelMixin, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery
 ):
+    _dit_modules: ClassVar[list[str]] = ["transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+
     def __init__(
         self,
         *,

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py

@@ -7,7 +7,7 @@ import logging
 import math
 import os
 from collections.abc import Iterable
-from typing import Any, cast
+from typing import Any, ClassVar, cast
 
 import numpy as np
 import PIL.Image
@@ -25,7 +25,7 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
-from vllm_omni.diffusion.models.interface import SupportImageInput
+from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery
 from vllm_omni.diffusion.models.qwen_image.autoencoder_kl_qwenimage import (
     AutoencoderKLQwenImage,
 )
@@ -203,7 +203,13 @@ def retrieve_latents(
         raise AttributeError("Could not access latents of provided encoder_output")
 
 
-class QwenImageLayeredPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMixin, DiffusionPipelineProfilerMixin):
+class QwenImageLayeredPipeline(
+    nn.Module, SupportImageInput, QwenImageCFGParallelMixin, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery
+):
+    _dit_modules: ClassVar[list[str]] = ["transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+
     color_format = "RGBA"
 
     def __init__(

vllm_omni/diffusion/models/sd3/pipeline_sd3.py

@@ -3,6 +3,7 @@ import json
 import logging
 import os
 from collections.abc import Iterable
+from typing import ClassVar
 
 import torch
 from diffusers.image_processor import VaeImageProcessor
@@ -20,6 +21,7 @@ from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
+from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
 from vllm_omni.diffusion.models.sd3.sd3_transformer import (
     SD3Transformer2DModel,
 )
@@ -129,7 +131,11 @@ def retrieve_timesteps(
     return timesteps, num_inference_steps
 
 
-class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMixin):
+class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery):
+    _dit_modules: ClassVar[list[str]] = ["transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["text_encoder", "text_encoder_2", "text_encoder_3"]
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+
     def __init__(
         self,
         *,

vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py

@@ -29,7 +29,7 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
-from vllm_omni.diffusion.models.interface import SupportAudioOutput
+from vllm_omni.diffusion.models.interface import SupportAudioOutput, SupportsComponentDiscovery
 from vllm_omni.diffusion.models.stable_audio.stable_audio_transformer import (
     StableAudioDiTModel,
     StableAudioSchedulerWrapper,
@@ -65,7 +65,7 @@ def get_stable_audio_post_process_func(
     return post_process_func
 
 
-class StableAudioPipeline(nn.Module, SupportAudioOutput, DiffusionPipelineProfilerMixin):
+class StableAudioPipeline(nn.Module, SupportAudioOutput, SupportsComponentDiscovery, DiffusionPipelineProfilerMixin):
     """
     Pipeline for text-to-audio generation using Stable Audio Open.
 
@@ -84,6 +84,11 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput, DiffusionPipelineProfil
     support_audio_output: ClassVar[bool] = True
     audio_sample_rate: ClassVar[int] = 44100
 
+    _dit_modules: ClassVar[list[str]] = ["transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+    _resident_modules: ClassVar[list[str]] = ["projection_model"]
+
     def __init__(
         self,
         *,

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py

@@ -8,7 +8,7 @@ import logging
 import os
 import time
 from collections.abc import Iterable
-from typing import Any, cast
+from typing import Any, ClassVar, cast
 
 import PIL.Image
 import torch
@@ -28,6 +28,7 @@ from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
 from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
+from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_zero
 from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler
 from vllm_omni.diffusion.models.wan2_2.scheduling_wan_euler import WanEulerScheduler
@@ -257,8 +258,17 @@ def get_wan22_pre_process_func(
 
 
 class Wan22Pipeline(
-    nn.Module, PipelineParallelMixin, CFGParallelMixin, ProgressBarMixin, DiffusionPipelineProfilerMixin
+    nn.Module,
+    PipelineParallelMixin,
+    CFGParallelMixin,
+    ProgressBarMixin,
+    DiffusionPipelineProfilerMixin,
+    SupportsComponentDiscovery,
 ):
+    _dit_modules: ClassVar[list[str]] = ["transformer", "transformer_2"]
+    _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+
     def __init__(
         self,
         *,

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py

@@ -7,7 +7,7 @@ import logging
 import os
 import time
 from collections.abc import Iterable
-from typing import Any, cast
+from typing import Any, ClassVar, cast
 
 import numpy as np
 import PIL.Image
@@ -28,7 +28,7 @@ from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
 from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
-from vllm_omni.diffusion.models.interface import SupportImageInput
+from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_zero
 from vllm_omni.diffusion.models.utils import _load_json
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
@@ -143,6 +143,7 @@ class Wan22I2VPipeline(
     CFGParallelMixin,
     ProgressBarMixin,
     DiffusionPipelineProfilerMixin,
+    SupportsComponentDiscovery,
 ):
     """
     Wan2.2 Image-to-Video Pipeline.
@@ -151,6 +152,10 @@ class Wan22I2VPipeline(
     Wan2.2-style I2V (with expand_timesteps for TI2V-5B).
     """
 
+    _dit_modules: ClassVar[list[str]] = ["transformer", "transformer_2"]
+    _encoder_modules: ClassVar[list[str]] = ["text_encoder", "image_encoder"]
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+
     def __init__(
         self,
         *,

vllm_omni/diffusion/models/z_image/pipeline_z_image.py

@@ -19,7 +19,7 @@ import inspect
 import json
 import os
 from collections.abc import Callable, Iterable
-from typing import Any
+from typing import Any, ClassVar
 
 import PIL.Image
 import torch
@@ -36,6 +36,7 @@ from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl import Distribu
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders
+from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
 from vllm_omni.diffusion.models.utils import create_transformers_model
 from vllm_omni.diffusion.models.z_image.z_image_transformer import (
     ZImageTransformer2DModel,
@@ -160,7 +161,11 @@ def retrieve_timesteps(
     return timesteps, num_inference_steps
 
 
-class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin):
+class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery):
+    _dit_modules: ClassVar[list[str]] = ["transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+
     def __init__(
         self,
         *,

vllm_omni/diffusion/offloader/sequential_backend.py

@@ -204,12 +204,6 @@ class ModelLevelOffloadBackend(OffloadBackend):
             return
 
         modules = ModuleDiscovery.discover(pipeline)
-        if not modules.dits:
-            logger.warning("No DiT/transformer modules found, skipping model-level offloading")
-            return
-        if not modules.encoders:
-            logger.warning("No encoder modules found, skipping model-level offloading")
-            return
 
         # Move encoders to GPU
         for enc in modules.encoders:
@@ -229,6 +223,17 @@ class ModelLevelOffloadBackend(OffloadBackend):
             except Exception as exc:
                 logger.warning("Failed to move resident module '%s' to GPU: %s", name, exc)
 
+        if not modules.dits:
+            logger.warning("No DiT/transformer modules found, skipping model-level offloading")
+            return
+
+        if not modules.encoders:
+            # Nothing to swap against — move DiTs to GPU and skip hooks.
+            for dit in modules.dits:
+                dit.to(self.device)
+            logger.warning("No encoder modules found, skipping model-level offloading")
+            return
+
         # Apply sequential offloading hooks
         apply_sequential_offload(
             dit_modules=modules.dits,