[Feature] Support multi-LoRA composition for diffusion models

[ultism]

Summary

Implements multi-LoRA composition for diffusion models, enabling multiple LoRA adapters to be active simultaneously on a single request with independent per-adapter scales. This is a core workflow for the diffusion community (CivitAI, ComfyUI, etc.) and was tracked in RFC #2149.

Key changes:

• Slot-based buffer management: DiffusionLoRAManager now manages N GPU buffer slots (max_loras) instead of a single slot. Each adapter occupies one slot; unused slots are reset.
• Additive composition in apply(): DiffusionBaseLinearLayerWithLoRA.apply() iterates all active adapter slots and accumulates LoRA deltas: y += Σ (x @ Aᵢᵀ) @ Bᵢᵀ
• API surface updated to plural: lora_request/lora_scale → lora_requests: list/lora_scales: list across sampling params, entrypoints, and OpenAI-compatible protocol
• OpenAI endpoint accepts single dict or list of dicts for the lora field in /v1/images/generations, /v1/images/edits, /v1/chat/completions, and /v1/videos/generations
• LRU cache manages adapter registration independently of activation — max_cpu_loras controls cache size, max_loras controls concurrent GPU slots

New config parameter:

max_loras: int = 1   # max adapters composed per request (GPU slot count)
max_cpu_loras: int    # max adapters cached on CPU (LRU), defaults to max_loras

API Usage

Python (offline inference)

from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.lora.request import LoRARequest

params = OmniDiffusionSamplingParams(
    lora_requests=[lora_a, lora_b],
    lora_scales=[0.8, 0.6],
    # ... other params
)

OpenAI-compatible endpoint

{
  "prompt": "1girl, portrait",
  "lora": [
    {"name": "style_a", "path": "/path/to/lora_a", "scale": 0.8},
    {"name": "style_b", "path": "/path/to/lora_b", "scale": 0.6}
  ]
}

E2E Validation

Tested on RTX 5090 with Z-Image Turbo base model and two CivitAI LoRAs:

• Jinx Arcane (character style)
• Character Design Sheet Helper (layout style)

XY grid (Jinx scale × Chardesign scale, values: 0, 0.25, 0.5, 1.0):

The grid demonstrates:

• Each axis independently controls its respective adapter's influence
• Composition is additive — both styles blend when both scales are non-zero
• Scale 0 correctly disables the adapter without affecting the other

Files Changed (15 files, +775/-381)

Area	Files	Description
Core	diffusion/lora/manager.py	Multi-slot activation, extracted _set_lora_for_layer(), _activate_adapters()
Core	diffusion/lora/layers/base_linear.py	Multi-adapter apply() loop, per-slot active-slice tracking
Config	diffusion/data.py	Added max_loras param, validation
Data	inputs/data.py	lora_request → lora_requests, lora_scale → lora_scales
Worker	diffusion/worker/diffusion_worker.py	Plumbed max_loras, updated set_active_adapter calls
Engine	engine/async_omni_engine.py	Pass max_loras through engine config
Entrypoints	entrypoints/async_omni_diffusion.py	Plural lora_requests parameter
Entrypoints	entrypoints/openai/utils.py	New parse_lora_requests() for dict | list[dict] | None
Entrypoints	entrypoints/openai/api_server.py	Updated image gen/edit endpoints
Entrypoints	entrypoints/openai/serving_chat.py	Updated chat-based image generation
Entrypoints	entrypoints/openai/serving_video.py	Updated video generation
Protocol	entrypoints/openai/protocol/images.py	lora field accepts list[dict]
Protocol	entrypoints/openai/protocol/videos.py	lora field accepts list[dict]
Tests	tests/diffusion/lora/test_lora_manager.py	7 new multi-LoRA unit tests + updated existing tests
Tests	tests/e2e/offline_inference/test_diffusion_lora.py	New test_diffusion_multi_lora_composition E2E test

Test Plan

• Updated existing LoRA unit tests to use the new plural API (set_active_adapters, slot-indexed set_calls/reset_calls)
• Added 7 new multi-LoRA unit tests: multi-slot activation, unused slot reset, max_loras exceeded, mismatched lengths, empty deactivation, zero-scale filtering, scale comparison — all passed
• Added E2E test test_diffusion_multi_lora_composition: creates two synthetic LoRA adapters perturbing different QKV slices, verifies composed output differs from each adapter alone and from baseline — passed
• Manual E2E: XY grid with real CivitAI LoRAs on RTX 5090 (Z-Image Turbo, 16 image combinations) — passed

Breaking Changes

• OmniDiffusionSamplingParams.lora_request → lora_requests (list)
• OmniDiffusionSamplingParams.lora_scale → lora_scales (list)
• DiffusionLoRAManager.set_active_adapter() → set_active_adapters()

These are internal APIs not yet part of a stable release.

Closes #2149

🤖 Generated with Claude Code

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: d04ada9f74

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Set default scales when lora_requests override sampling params

When generate_batch(..., lora_requests=...) (and similarly generate) is used, only sampling_params.lora_requests is populated while sampling_params.lora_scales is left at its default empty list. The worker then calls set_active_adapters() and hits the length check (lora_requests vs lora_scales), so requests fail with ValueError whenever callers rely on this new convenience argument without manually setting scales in sampling_params.

Useful? React with 👍 / 👎.

[ultism]

Thanks for catching this! Fixed in commit ca4d418 — when lora_requests is passed via generate() / generate_batch() without an explicit lora_scales, lora_scales now defaults to [1.0] * len(lora_requests) so the length check in set_active_adapters no longer trips. Callers using the convenience argument without specifying scales will get uniform 1.0 scaling, which matches the most common use case.

[chatgpt-codex-connector]

Filter zero-scale adapters before enforcing max_loras

This limit check runs before zero-scale adapters are removed, so requests like [adapter_a, adapter_b] with scales [1.0, 0.0] are rejected when max_loras=1 even though only one adapter would actually be active. Since zero scales are explicitly treated as disabled later in the method, enforcing max_loras on the unfiltered list causes avoidable request failures.

Useful? React with 👍 / 👎.

[ultism]

Thanks for flagging — this is actually intentional. In this design, scale=0.0 is semantically distinct from omitting the adapter from lora_requests:

• 0.0 → adapter stays registered and occupies a LoRA cache slot, but contributes nothing to this forward pass.
• Absent from the list → slot is released.

So max_loras is enforced against the full lora_requests list on purpose. If we filtered zero-scale adapters first, the same list [A, B] would pass under max_loras=1 with scales [1.0, 0.0] but fail with [1.0, 1.0] — a state-dependent capacity error that's harder to reason about than rejecting up front.

To release a slot, callers should remove the adapter from the request rather than pass 0.0. I've added an inline comment on the check to make this contract explicit.

[yenuo26]

please add @pytest.mark.diffusion @pytest.mark.advanced_model
For more information about tags, please refer to docs/contributing/ci/tests_markers.md

[ultism]

Both markers are in place — @pytest.mark.diffusion and @pytest.mark.advanced_model are applied to test_diffusion_model at test_diffusion_lora.py:93-94, matching docs/contributing/ci/tests_markers.md. Thanks for flagging.

[SamitHuang]

For the test plan, please give the whole test script or add an example

[SamitHuang]

please also update the lora feature docs.

[ultism]

Done — docs/user_guide/diffusion/lora.md has been reworked to cover the new flows:

• 355ecea5 — rewrote the guide around init-time vs per-request LoRA.
• 5be2773f — clarified that sampling_params_list is stage-indexed (not per-prompt) and noted CLI auto-sizes max_loras.
• 7f5f7fd8 — documented the --axis XYZ system for Cartesian-product grids across lora_scale[i]/prompt/guidance_scale/num_inference_steps/seed.

The guide now has separate "Init-time LoRA" and "Per-request LoRA" sections with zero-scale slot semantics, max_loras sizing, multi-LoRA composition examples, and CLI snippets. Let me know if anything's still missing.

[lishunyang12]

left a few comments. the zero-scale filtering ordering bug is the main one.

[lishunyang12]

+1 on the codex bot's comment — max_loras is enforced before zero-scale adapters are filtered out a few lines below, so [a, b] with scales [1.0, 0.0] blows up when max_loras=1 even though only one adapter would actually be active.

Move the len(lora_requests) > self.max_loras check to after the zero-scale filter loop, operating on active_requests instead.

[ultism]

Good catch, but this ordering is intentional and I'd like to keep it. The key semantic distinction is that scale=0.0 is not equivalent to omitting the adapter — a zero-scale adapter is still registered and occupies a LoRA cache slot, it just contributes nothing to this forward pass. Releasing the slot requires removing the adapter from the request entirely.

If we moved the max_loras check after the zero-scale filter, [A, B] with scales [1.0, 0.0] would pass under max_loras=1, but the same list with scales [1.0, 1.0] (or even [1.0, 0.5]) would fail. That state-dependent capacity behavior is more surprising than the current upfront rejection.

I've pushed a commit that adds an inline comment on the check, so the intent is obvious at the call site. Happy to revisit if you see a concrete use case where zero-scale "soft disable" without a slot is required — but that feels like a distinct feature (e.g. an explicit "reserved but inactive" flag) rather than overloading 0.0.

Will address the _are_active_at_scales ordering TODO and the dead _parse_lora_request cleanup in a follow-up commit.

[lishunyang12]

_are_active_at_scales is order-sensitive — requesting [A, B] then [B, A] with the same scales will trigger a full re-activation even though the composed result is identical (addition is commutative). Probably fine for now but worth a TODO.

[ultism]

Good call — added the TODO at manager.py:508-511 inline above the length check, noting the multiset comparison as the follow-up. Keeping the order-sensitive behavior for now since it's correct, just suboptimal.

[lishunyang12]

Nit: is _parse_lora_request (singular) still called anywhere? Both call sites now use _parse_lora_requests. If it's dead code, remove it.

[ultism]

Cleaned up in the multi-LoRA refactor — the singular _parse_lora_request no longer exists anywhere in the tree (rg _parse_lora_request returns nothing). What's at api_server.py:1672 now is the plural _parse_lora_requests, which is a thin wrapper that translates ValueError from parse_lora_requests (in entrypoints/openai/utils.py) into HTTPException(400) — kept because the util is also called from non-HTTP paths (serving_video.py, serving_chat.py) that shouldn't raise HTTP errors.

[ultism]

Re: test plan — added both:

• E2E test: tests/e2e/offline_inference/test_diffusion_lora.py covers init-time LoRA, per-request single/multi-LoRA, zero-scale slot semantics, and max_loras rejection, using a synthetic Z-Image PEFT adapter (tagged @pytest.mark.diffusion @pytest.mark.advanced_model).
• Runnable example: examples/offline_inference/text_to_image/text_to_image.py now exposes --lora-path (init-time), --lora-paths + --lora-scales (per-request composition), and --axis x=…/y=…/z=… for Cartesian-product XYZ grids over lora_scale[i] / prompt / guidance_scale / num_inference_steps / seed.

Full walkthroughs with CLI snippets are in docs/user_guide/diffusion/lora.md. cc @SamitHuang

[ultism]

@hsliuustc0106 Could you take a look when you have time? The branch has been rebased onto latest upstream/main and pre-commit is now passing. Ready for review.