From e957d16f8dda9a424382b78fc02119a36aa9f436 Mon Sep 17 00:00:00 2001 From: Nick Cao Date: Wed, 17 Jun 2026 09:33:06 -0400 Subject: [PATCH] [Refactor] Migrate existing pipelines to use SupportsComponentDiscovery for offload discovery (#3076) Signed-off-by: Nick Cao Co-authored-by: Claude --- .buildkite/test-ready.yml | 2 +- .../offline_inference/sensenova_u1/README.md | 2 +- .../offloader/test_diffusion_cpu_offload.py | 2 ++ .../test_magi_human_expansion.py | 16 ++++++++++++++- .../test_stable_audio_expansion.py | 20 +++++++++++++++++++ .../online_serving/test_bagel_expansion.py | 11 ++++++++++ .../online_serving/test_flux2_expansion.py | 13 ++++++++++++ .../test_flux2_klein_expansion.py | 7 +++++++ .../test_hunyuan_video_15_expansion.py | 11 ++++++++++ .../online_serving/test_nextstep_expansion.py | 14 +++++++++++++ .../test_qwen_image_edit_expansion.py | 10 ++++++++++ .../test_qwen_image_expansion.py | 8 ++++++++ .../test_qwen_image_layered_expansion.py | 8 ++++++++ .../e2e/online_serving/test_sd3_expansion.py | 14 +++++++++++++ .../online_serving/test_wan22_expansion.py | 17 +++++++++++++++- .../test_wan_2_1_vace_expansion.py | 12 +++++++++++ .../dreamid_omni/pipeline_dreamid_omni.py | 11 ++++++++-- .../diffusion/models/flux/pipeline_flux.py | 11 ++++++++-- .../models/flux/pipeline_flux_kontext.py | 15 +++++++++++--- .../diffusion/models/flux2/pipeline_flux2.py | 17 +++++++++++++--- .../flux2_klein/pipeline_flux2_klein.py | 12 ++++++++--- .../models/glm_image/pipeline_glm_image.py | 9 +++++++-- .../models/helios/pipeline_helios.py | 11 ++++++++-- .../hunyuan_image3/pipeline_hunyuan_image3.py | 7 ++++++- .../pipeline_hunyuan_video_1_5.py | 11 ++++++++-- .../pipeline_hunyuan_video_1_5_i2v.py | 15 +++++++++++--- .../longcat_image/pipeline_longcat_image.py | 9 +++++++-- .../pipeline_longcat_image_edit.py | 12 ++++++++--- .../diffusion/models/ltx2/pipeline_ltx2.py | 6 +++++- .../models/magi_human/pipeline_magi_human.py | 9 +++++++-- .../pipeline_mammothmoda2_dit.py | 9 +++++++-- .../nextstep_1_1/pipeline_nextstep_1_1.py | 9 +++++++-- .../models/ovis_image/pipeline_ovis_image.py | 9 +++++++-- .../models/qwen_image/pipeline_qwen_image.py | 9 ++++++++- .../qwen_image/pipeline_qwen_image_edit.py | 12 ++++++++--- .../pipeline_qwen_image_edit_plus.py | 10 +++++++--- .../qwen_image/pipeline_qwen_image_layered.py | 12 ++++++++--- .../diffusion/models/sd3/pipeline_sd3.py | 8 +++++++- .../stable_audio/pipeline_stable_audio.py | 9 +++++++-- .../models/wan2_2/pipeline_wan2_2.py | 14 +++++++++++-- .../models/wan2_2/pipeline_wan2_2_i2v.py | 9 +++++++-- .../models/z_image/pipeline_z_image.py | 9 +++++++-- .../diffusion/offloader/sequential_backend.py | 17 ++++++++++------ 43 files changed, 392 insertions(+), 66 deletions(-) diff --git a/.buildkite/test-ready.yml b/.buildkite/test-ready.yml index 99d874475..851ef6934 100644 --- a/.buildkite/test-ready.yml +++ b/.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: diff --git a/examples/offline_inference/sensenova_u1/README.md b/examples/offline_inference/sensenova_u1/README.md index 3943b8d4e..a2e245526 100644 --- a/examples/offline_inference/sensenova_u1/README.md +++ b/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: diff --git a/tests/diffusion/offloader/test_diffusion_cpu_offload.py b/tests/diffusion/offloader/test_diffusion_cpu_offload.py index 92e2abc7e..7d155fef9 100644 --- a/tests/diffusion/offloader/test_diffusion_cpu_offload.py +++ b/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} diff --git a/tests/e2e/offline_inference/test_magi_human_expansion.py b/tests/e2e/offline_inference/test_magi_human_expansion.py index 9dd4a1d19..1b3550913 100644 --- a/tests/e2e/offline_inference/test_magi_human_expansion.py +++ b/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( diff --git a/tests/e2e/offline_inference/test_stable_audio_expansion.py b/tests/e2e/offline_inference/test_stable_audio_expansion.py index 96365c7d3..2785dab24 100644 --- a/tests/e2e/offline_inference/test_stable_audio_expansion.py +++ b/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() diff --git a/tests/e2e/online_serving/test_bagel_expansion.py b/tests/e2e/online_serving/test_bagel_expansion.py index a8f52c311..fa17791e5 100644 --- a/tests/e2e/online_serving/test_bagel_expansion.py +++ b/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( diff --git a/tests/e2e/online_serving/test_flux2_expansion.py b/tests/e2e/online_serving/test_flux2_expansion.py index 3aab4889f..c900b91e8 100644 --- a/tests/e2e/online_serving/test_flux2_expansion.py +++ b/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( diff --git a/tests/e2e/online_serving/test_flux2_klein_expansion.py b/tests/e2e/online_serving/test_flux2_klein_expansion.py index a3d4f004e..98ccf0014 100644 --- a/tests/e2e/online_serving/test_flux2_klein_expansion.py +++ b/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, diff --git a/tests/e2e/online_serving/test_hunyuan_video_15_expansion.py b/tests/e2e/online_serving/test_hunyuan_video_15_expansion.py index 681d8ff23..0999ea8b5 100644 --- a/tests/e2e/online_serving/test_hunyuan_video_15_expansion.py +++ b/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( diff --git a/tests/e2e/online_serving/test_nextstep_expansion.py b/tests/e2e/online_serving/test_nextstep_expansion.py index 75290db5d..24384d29c 100644 --- a/tests/e2e/online_serving/test_nextstep_expansion.py +++ b/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, diff --git a/tests/e2e/online_serving/test_qwen_image_edit_expansion.py b/tests/e2e/online_serving/test_qwen_image_edit_expansion.py index 1457ad636..b3d6f8b19 100644 --- a/tests/e2e/online_serving/test_qwen_image_edit_expansion.py +++ b/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, diff --git a/tests/e2e/online_serving/test_qwen_image_expansion.py b/tests/e2e/online_serving/test_qwen_image_expansion.py index f1cf449f7..c900ca4ce 100644 --- a/tests/e2e/online_serving/test_qwen_image_expansion.py +++ b/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", diff --git a/tests/e2e/online_serving/test_qwen_image_layered_expansion.py b/tests/e2e/online_serving/test_qwen_image_layered_expansion.py index 4e79beab7..c26e21039 100644 --- a/tests/e2e/online_serving/test_qwen_image_layered_expansion.py +++ b/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, diff --git a/tests/e2e/online_serving/test_sd3_expansion.py b/tests/e2e/online_serving/test_sd3_expansion.py index 6be0fe4e9..c299ae8f1 100644 --- a/tests/e2e/online_serving/test_sd3_expansion.py +++ b/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( diff --git a/tests/e2e/online_serving/test_wan22_expansion.py b/tests/e2e/online_serving/test_wan22_expansion.py index fd0221ddc..85eee128a 100644 --- a/tests/e2e/online_serving/test_wan22_expansion.py +++ b/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( diff --git a/tests/e2e/online_serving/test_wan_2_1_vace_expansion.py b/tests/e2e/online_serving/test_wan_2_1_vace_expansion.py index 4f0e96447..d58bb9e32 100644 --- a/tests/e2e/online_serving/test_wan_2_1_vace_expansion.py +++ b/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( diff --git a/vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py b/vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py index 47b905de7..1b764e4a5 100644 --- a/vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py +++ b/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, *, diff --git a/vllm_omni/diffusion/models/flux/pipeline_flux.py b/vllm_omni/diffusion/models/flux/pipeline_flux.py index 03bef7bd6..ecbc0536d 100644 --- a/vllm_omni/diffusion/models/flux/pipeline_flux.py +++ b/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, *, diff --git a/vllm_omni/diffusion/models/flux/pipeline_flux_kontext.py b/vllm_omni/diffusion/models/flux/pipeline_flux_kontext.py index c3bea7dd1..d57909837 100644 --- a/vllm_omni/diffusion/models/flux/pipeline_flux_kontext.py +++ b/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__( diff --git a/vllm_omni/diffusion/models/flux2/pipeline_flux2.py b/vllm_omni/diffusion/models/flux2/pipeline_flux2.py index 948f922d4..a92d7e55c 100644 --- a/vllm_omni/diffusion/models/flux2/pipeline_flux2.py +++ b/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 diff --git a/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py b/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py index aea6d7e03..74598319f 100644 --- a/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py +++ b/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__( diff --git a/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py b/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py index 102456d9f..e6075ce10 100644 --- a/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py +++ b/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, *, diff --git a/vllm_omni/diffusion/models/helios/pipeline_helios.py b/vllm_omni/diffusion/models/helios/pipeline_helios.py index c3a47ebaa..9b625b6b9 100644 --- a/vllm_omni/diffusion/models/helios/pipeline_helios.py +++ b/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, *, diff --git a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py index 2aef06f20..483f345fe 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py +++ b/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.": "", diff --git a/vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5.py b/vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5.py index 0684f6749..aa5b9b83c 100644 --- a/vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5.py +++ b/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, *, diff --git a/vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5_i2v.py b/vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5_i2v.py index 583ef87ed..3a3cc46dd 100644 --- a/vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5_i2v.py +++ b/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" diff --git a/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image.py b/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image.py index ec942c25f..e8c23c1f0 100644 --- a/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image.py +++ b/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, *, diff --git a/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image_edit.py b/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image_edit.py index 4bfa00a97..54de7e9f9 100644 --- a/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image_edit.py +++ b/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, *, diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py index a409776d2..85927354f 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py +++ b/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 diff --git a/vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py b/vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py index 65da889f1..2794ddddd 100644 --- a/vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py +++ b/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 diff --git a/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py b/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py index e55338963..e230fab33 100644 --- a/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py +++ b/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. diff --git a/vllm_omni/diffusion/models/nextstep_1_1/pipeline_nextstep_1_1.py b/vllm_omni/diffusion/models/nextstep_1_1/pipeline_nextstep_1_1.py index 4fa56ea93..a0d27c9d0 100644 --- a/vllm_omni/diffusion/models/nextstep_1_1/pipeline_nextstep_1_1.py +++ b/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, *, diff --git a/vllm_omni/diffusion/models/ovis_image/pipeline_ovis_image.py b/vllm_omni/diffusion/models/ovis_image/pipeline_ovis_image.py index f77c5d419..cfe0159c1 100644 --- a/vllm_omni/diffusion/models/ovis_image/pipeline_ovis_image.py +++ b/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, *, diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py index 655641fb4..498d18503 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py +++ b/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__( diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py index a83aa83ad..482dbac7f 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py +++ b/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, *, diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py index 7d2170e09..ffae53ff5 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py +++ b/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, *, diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py index 19711cd8c..055edac35 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py +++ b/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__( diff --git a/vllm_omni/diffusion/models/sd3/pipeline_sd3.py b/vllm_omni/diffusion/models/sd3/pipeline_sd3.py index 7eeffdde4..7dce56aac 100644 --- a/vllm_omni/diffusion/models/sd3/pipeline_sd3.py +++ b/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, *, diff --git a/vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py b/vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py index f59215e0c..3f49450e5 100644 --- a/vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py +++ b/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, *, diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py index 85a585960..1b4cee123 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py +++ b/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, *, diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py index 68eb3b11b..ffb21b245 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py +++ b/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, *, diff --git a/vllm_omni/diffusion/models/z_image/pipeline_z_image.py b/vllm_omni/diffusion/models/z_image/pipeline_z_image.py index 388b6316c..d9f383249 100644 --- a/vllm_omni/diffusion/models/z_image/pipeline_z_image.py +++ b/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, *, diff --git a/vllm_omni/diffusion/offloader/sequential_backend.py b/vllm_omni/diffusion/offloader/sequential_backend.py index 8fb675b69..0edf77651 100644 --- a/vllm_omni/diffusion/offloader/sequential_backend.py +++ b/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,