[Feature] Improve diffusion LoRA compatibility for SDXL (#4825)

Signed-off-by: Thanaji Rao Thakkalapelli <thanaji.rao.thakkalapelli@intel.com>

tests/diffusion/lora/test_lora_manager.py

@@ -575,3 +575,45 @@ def test_lora_manager_discovers_bagel_component(monkeypatch):
     assert "bagel.language_model.qkv_proj" in manager._lora_modules
     # Verify the module was actually replaced in the tree (not just recorded)
     assert isinstance(pipeline.bagel.language_model.qkv_proj, _DummyBaseLayerWithLoRA)
+
+
+def test_lora_manager_discovers_unet_component(monkeypatch):
+    """Verify that _replace_layers_with_lora finds layers under 'unet'."""
+    import vllm_omni.diffusion.lora.manager as manager_mod
+
+    monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DummyBaseLayerWithLoRA)
+
+    def _fake_from_layer_diffusion(*, layer: torch.nn.Module, **_kwargs):
+        if isinstance(layer, _FakeLinearBase):
+            return _DummyBaseLayerWithLoRA(layer)
+        return layer
+
+    replace_calls: list[str] = []
+
+    monkeypatch.setattr(manager_mod, "from_layer_diffusion", _fake_from_layer_diffusion)
+    monkeypatch.setattr(
+        manager_mod,
+        "replace_submodule",
+        lambda root, name, sub: fake_replace_submodule(root, name, sub, replace_calls),
+    )
+
+    # Pipeline with a 'unet' component (no 'transformer')
+    pipeline = torch.nn.Module()
+    pipeline.unet = torch.nn.Module()
+    pipeline.unet.down_block = torch.nn.Module()
+    pipeline.unet.down_block.proj = _FakeLinearBase()
+
+    manager = DiffusionLoRAManager(
+        pipeline=pipeline,
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+        max_cached_adapters=1,
+    )
+
+    peft_helper = type("_PH", (), {"r": 1})()
+    manager._replace_layers_with_lora(peft_helper)
+
+    assert "down_block.proj" in replace_calls
+    assert "unet.down_block.proj" in manager._lora_modules
+    # Verify the module was actually replaced in the tree (not just recorded)
+    assert isinstance(pipeline.unet.down_block.proj, _DummyBaseLayerWithLoRA)

vllm_omni/diffusion/lora/manager.py

@@ -368,7 +368,17 @@ class DiffusionLoRAManager:
 
         # Default denoising components plus any a pipeline opts into via
         # ``_lora_components`` (opt-in; other pipelines unchanged).
-        default_components = ("transformer", "transformer_2", "dit", "bagel")
+        #
+        # 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.
+        default_components = (
+            "transformer",
+            "transformer_2",
+            "dit",
+            "bagel",
+            "unet",
+        )
         extra_components = tuple(getattr(self.pipeline, "_lora_components", ()) or ())
         for component_name in (*default_components, *extra_components):
             if not hasattr(self.pipeline, component_name):