[Refactor][Diffusion] Remove model-specific names from LoRA and ModelOpt loader defaults (#5907)

Signed-off-by: Alicia <115451386+congw729@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/diffusion/lora/test_lora_manager.py

@@ -729,8 +729,10 @@ def test_lora_manager_discovers_bagel_component(monkeypatch):
         lambda root, name, sub: fake_replace_submodule(root, name, sub, replace_calls),
     )
 
-    # Pipeline with a 'bagel' component (no 'transformer')
+    # Pipeline with a 'bagel' component (no 'transformer') declared via
+    # _lora_components, matching how BagelPipeline opts in.
     pipeline = torch.nn.Module()
+    pipeline._lora_components = ["bagel"]
     pipeline.bagel = torch.nn.Module()
     pipeline.bagel.language_model = torch.nn.Module()
     pipeline.bagel.language_model.qkv_proj = _FakeLinearBase()

tests/diffusion/models/bagel/test_bagel_lora.py

@@ -122,8 +122,10 @@ class TestStage1DiTLoRA:
             lambda root, name, sub: fake_replace_submodule(root, name, sub, replace_calls),
         )
 
-        # Build pipeline with bagel component
+        # Build pipeline with bagel component declared via _lora_components,
+        # matching how BagelPipeline opts in.
         pipeline = torch.nn.Module()
+        pipeline._lora_components = ["bagel"]
         pipeline.bagel = torch.nn.Module()
         lm = torch.nn.Module()
         lm.stacked_params_mapping = [
@@ -194,8 +196,11 @@ class TestBagelLoRARoundTrip:
 
         monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", DummyBaseLayerWithLoRA)
 
-        # Build pipeline with bagel.language_model.foo (simple non-packed layer)
+        # Build pipeline with bagel.language_model.foo (simple non-packed layer);
+        # the bagel component is declared via _lora_components, matching how
+        # BagelPipeline opts in.
         pipeline = torch.nn.Module()
+        pipeline._lora_components = ["bagel"]
         pipeline.bagel = torch.nn.Module()
         lm = torch.nn.Module()
         lm.foo = _FakeLinearBase()

vllm_omni/diffusion/lora/manager.py

@@ -394,17 +394,20 @@ class DiffusionLoRAManager:
             fully_sharded_loras=False,
         )
 
-        # Default denoising components, declared DiT components, and any a
-        # pipeline opts into via ``_lora_components``.
+        # Components scanned for LoRA-capable layers: framework defaults,
+        # declared DiT components, and any a pipeline opts into via
+        # ``_lora_components``. The defaults only cover the generic diffusers
+        # naming convention: ``transformer`` (plus ``transformer_2`` for
+        # dual-DiT pipelines such as Wan2.2) and ``unet`` for SDXL-style
+        # pipelines. Model-specific attribute names must be declared by the
+        # pipeline itself via ``_dit_modules`` or ``_lora_components``.
         #
-        # NOTE: SDXL-style pipelines expose the denoiser as ``unet``.
-        # Without scanning this component, adapters can load/activate while
-        # effectively applying to zero layers, producing base-identical output.
+        # NOTE: if the denoiser component is not scanned here, adapters can
+        # load/activate while effectively applying to zero layers, producing
+        # base-identical output.
         default_components = (
             "transformer",
             "transformer_2",
-            "dit",
-            "bagel",
             "unet",
         )
         declared_components = tuple(getattr(self.pipeline, "_dit_modules", ()) or ())

vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py

@@ -17,10 +17,17 @@ MODEL_OPT_SCALE_SUFFIXES = (
     ".weight_scale_2",
     ".weight_scale_inv",
 )
+# Fused projections following the generic diffusers Attention naming
+# convention. ``to_q``/``to_k``/``to_v`` and the joint-attention
+# ``add_q_proj``/``add_k_proj``/``add_v_proj`` shard names are fused into
+# ``to_qkv``/``add_kv_proj`` by most vllm-omni DiT implementations, many of
+# which do not declare their own ``packed_modules_mapping``; keeping these
+# two entries here avoids repeating them in every such model. Model-specific
+# fusions (e.g. Z-Image's ``w13``) must be declared on the model class via
+# ``packed_modules_mapping``, which is merged on top of these defaults.
 DEFAULT_PACKED_MODULES_MAPPING = {
     "to_qkv": ("to_q", "to_k", "to_v"),
     "add_kv_proj": ("add_q_proj", "add_k_proj", "add_v_proj"),
-    "w13": ("w1", "w3"),
 }
 FP8_DTYPES = tuple(
     dtype

vllm_omni/diffusion/models/bagel/pipeline_bagel.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """
 BagelPipeline implementation for vLLM-Omni.
 """
@@ -316,6 +316,11 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf
         "bagel.connector",
         "bagel.vit_pos_embed",
     ]
+    # LoRA scanning target. ``_dit_modules`` holds the dotted path
+    # "language_model.model", which DiffusionLoRAManager cannot resolve as a
+    # pipeline attribute, so the top-level ``bagel`` module (LLM plus ViT,
+    # connector, and VAE projections) is declared explicitly here.
+    _lora_components: ClassVar[list[str]] = ["bagel"]
     supports_step_execution: ClassVar[bool] = True
 
     def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):

vllm_omni/diffusion/models/z_image/z_image_transformer.py

@@ -612,6 +612,15 @@ class ZImageTransformer2DModel(CachedTransformer):
         Note: Our "Sequence Parallelism" (SP) corresponds to "Context Parallelism" (CP) in diffusers.
     """
 
+    # Fused projections and their checkpoint shard names (mirrors
+    # stacked_params_mapping in load_weights). ``w13`` is the Z-Image FFN
+    # fusion of ``w1``/``w3``; declared here so quantized-checkpoint adapters
+    # can resolve shard keys without a framework-level default entry.
+    packed_modules_mapping = {
+        "to_qkv": ["to_q", "to_k", "to_v"],
+        "w13": ["w1", "w3"],
+    }
+
     _repeated_blocks = ["ZImageTransformerBlock"]
     _layerwise_offload_blocks_attrs = ["layers"]