[diffusion] feat: support multiple LoRA adapters loading and application (#16667)

Co-authored-by: niehen6174 <niehen.6174@gmail.com>
Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
WenhaoZhang
2026-01-10 23:32:45 +08:00
committed by GitHub
parent 76d4881794
commit 5c72be1e51
10 changed files with 578 additions and 180 deletions

View File

@@ -241,17 +241,21 @@ The server supports dynamic loading, merging, and unmerging of LoRA adapters.
#### Set LoRA Adapter
Loads a LoRA adapter and merges its weights into the model.
Loads one or more LoRA adapters and merges their weights into the model. Supports both single LoRA (backward compatible) and multiple LoRA adapters.
**Endpoint:** `POST /v1/set_lora`
**Parameters:**
- `lora_nickname` (string, required): A unique identifier for this LoRA
- `lora_path` (string, optional): Path to the `.safetensors` file or Hugging Face repo ID. Required for the first load; optional if re-activating a cached nickname
- `target` (string, optional): Which transformer(s) to apply the LoRA to. One of "all" (default), "transformer", "transformer_2", "critic"
- `strength` (float, optional): LoRA strength for merge, default 1.0. Values < 1.0 reduce the effect, values > 1.0 amplify the effect
- `lora_nickname` (string or list of strings, required): A unique identifier for the LoRA adapter(s). Can be a single string or a list of strings for multiple LoRAs
- `lora_path` (string or list of strings/None, optional): Path to the `.safetensors` file(s) or Hugging Face repo ID(s). Required for the first load; optional if re-activating a cached nickname. If a list, must match the length of `lora_nickname`
- `target` (string or list of strings, optional): Which transformer(s) to apply the LoRA to. If a list, must match the length of `lora_nickname`. Valid values:
- `"all"` (default): Apply to all transformers
- `"transformer"`: Apply only to the primary transformer (high noise for Wan2.2)
- `"transformer_2"`: Apply only to transformer_2 (low noise for Wan2.2)
- `"critic"`: Apply only to the critic model
- `strength` (float or list of floats, optional): LoRA strength for merge, default 1.0. If a list, must match the length of `lora_nickname`. Values < 1.0 reduce the effect, values > 1.0 amplify the effect
**Curl Example:**
**Single LoRA Example:**
```bash
curl -X POST http://localhost:30010/v1/set_lora \
@@ -259,10 +263,43 @@ curl -X POST http://localhost:30010/v1/set_lora \
-d '{
"lora_nickname": "lora_name",
"lora_path": "/path/to/lora.safetensors",
"target": "all",
"strength": 0.8
}'
```
**Multiple LoRA Example:**
```bash
curl -X POST http://localhost:30010/v1/set_lora \
-H "Content-Type: application/json" \
-d '{
"lora_nickname": ["lora_1", "lora_2"],
"lora_path": ["/path/to/lora1.safetensors", "/path/to/lora2.safetensors"],
"target": ["transformer", "transformer_2"],
"strength": [0.8, 1.0]
}'
```
**Multiple LoRA with Same Target:**
```bash
curl -X POST http://localhost:30010/v1/set_lora \
-H "Content-Type: application/json" \
-d '{
"lora_nickname": ["style_lora", "character_lora"],
"lora_path": ["/path/to/style.safetensors", "/path/to/character.safetensors"],
"target": "all",
"strength": [0.7, 0.9]
}'
```
> [!NOTE]
> When using multiple LoRAs:
> - All list parameters (`lora_nickname`, `lora_path`, `target`, `strength`) must have the same length
> - If `target` or `strength` is a single value, it will be applied to all LoRAs
> - Multiple LoRAs applied to the same target will be merged in order
#### Merge LoRA Weights

View File

@@ -11,7 +11,7 @@ diffusion models.
import multiprocessing as mp
import os
import time
from typing import Any
from typing import Any, List, Union
import numpy as np
@@ -21,6 +21,7 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.utils import (
MergeLoraWeightsReq,
SetLoraReq,
UnmergeLoraWeightsReq,
format_lora_message,
)
from sglang.multimodal_gen.runtime.entrypoints.utils import (
post_process_sample,
@@ -317,23 +318,25 @@ class DiffGenerator:
def set_lora(
self,
lora_nickname: str,
lora_path: str | None = None,
target: str = "all",
strength: float = 1.0,
lora_nickname: Union[str, List[str]],
lora_path: Union[str, None, List[Union[str, None]]] = None,
target: Union[str, List[str]] = "all",
strength: Union[float, List[float]] = 1.0,
) -> None:
"""
Set a LoRA adapter for the specified transformer(s).
Set LoRA adapter(s) for the specified transformer(s).
Supports both single LoRA (backward compatible) and multiple LoRA adapters.
Args:
lora_nickname: The nickname of the adapter.
lora_path: Path to the LoRA adapter.
target: Which transformer(s) to apply the LoRA to. One of:
lora_nickname: The nickname(s) of the adapter(s). Can be a string or a list of strings.
lora_path: Path(s) to the LoRA adapter(s). Can be a string, None, or a list of strings/None.
target: Which transformer(s) to apply the LoRA to. Can be a string or a list of strings.
Valid values:
- "all": Apply to all transformers (default)
- "transformer": Apply only to the primary transformer (high noise for Wan2.2)
- "transformer_2": Apply only to transformer_2 (low noise for Wan2.2)
- "critic": Apply only to the critic model
strength: LoRA strength for merge, default 1.0.
strength: LoRA strength(s) for merge, default 1.0. Can be a float or a list of floats.
"""
req = SetLoraReq(
lora_nickname=lora_nickname,
@@ -341,9 +344,13 @@ class DiffGenerator:
target=target,
strength=strength,
)
nickname_str, target_str, strength_str = format_lora_message(
lora_nickname, target, strength
)
self._send_lora_request(
req,
f"Successfully set LoRA adapter: {lora_nickname} (target: {target}, strength: {strength})",
f"Successfully set LoRA adapter(s): {nickname_str} (target: {target_str}, strength: {strength_str})",
"Failed to set LoRA adapter",
)

View File

@@ -1,4 +1,4 @@
from typing import Any, Optional
from typing import Any, List, Optional, Union
from fastapi import APIRouter, Body, HTTPException
from fastapi.responses import ORJSONResponse
@@ -9,6 +9,7 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.utils import (
MergeLoraWeightsReq,
SetLoraReq,
UnmergeLoraWeightsReq,
format_lora_message,
)
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch
from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client
@@ -48,23 +49,27 @@ async def _handle_lora_request(req: Any, success_msg: str, failure_msg: str):
@router.post("/set_lora")
async def set_lora(
lora_nickname: str = Body(..., embed=True),
lora_path: Optional[str] = Body(None, embed=True),
target: str = Body("all", embed=True),
strength: float = Body(1.0, embed=True),
lora_nickname: Union[str, List[str]] = Body(..., embed=True),
lora_path: Optional[Union[str, List[Optional[str]]]] = Body(None, embed=True),
target: Union[str, List[str]] = Body("all", embed=True),
strength: Union[float, List[float]] = Body(1.0, embed=True),
):
"""
Set a LoRA adapter for the specified transformer(s).
Set LoRA adapter(s) for the specified transformer(s).
Supports both single LoRA (backward compatible) and multiple LoRA adapters.
Args:
lora_nickname: The nickname of the adapter.
lora_path: Path to the LoRA adapter (local path or HF repo id).
target: Which transformer(s) to apply the LoRA to. One of:
lora_nickname: The nickname(s) of the adapter(s). Can be a string or a list of strings.
lora_path: Path(s) to the LoRA adapter(s) (local path or HF repo id).
Can be a string, None, or a list of strings/None. Must match the length of lora_nickname.
target: Which transformer(s) to apply the LoRA to. Can be a string or a list of strings.
If a list, must match the length of lora_nickname. Valid values:
- "all": Apply to all transformers (default)
- "transformer": Apply only to the primary transformer (high noise for Wan2.2)
- "transformer_2": Apply only to transformer_2 (low noise for Wan2.2)
- "critic": Apply only to the critic model
strength: LoRA strength for merge, default 1.0. Values < 1.0 reduce the effect,
strength: LoRA strength(s) for merge, default 1.0. Can be a float or a list of floats.
If a list, must match the length of lora_nickname. Values < 1.0 reduce the effect,
values > 1.0 amplify the effect.
"""
req = SetLoraReq(
@@ -73,9 +78,13 @@ async def set_lora(
target=target,
strength=strength,
)
nickname_str, target_str, strength_str = format_lora_message(
lora_nickname, target, strength
)
return await _handle_lora_request(
req,
f"Successfully set LoRA adapter: {lora_nickname} (target: {target}, strength: {strength})",
f"Successfully set LoRA adapter(s): {nickname_str} (target: {target_str}, strength: {strength_str})",
"Failed to set LoRA adapter",
)

View File

@@ -24,10 +24,10 @@ logger = init_logger(__name__)
@dataclasses.dataclass
class SetLoraReq:
lora_nickname: str
lora_path: Optional[str] = None
target: str = "all" # "all", "transformer", "transformer_2", "critic"
strength: float = 1.0 # LoRA strength for merge, default 1.0
lora_nickname: Union[str, List[str]]
lora_path: Optional[Union[str, List[Optional[str]]]] = None
target: Union[str, List[str]] = "all"
strength: Union[float, List[float]] = 1.0 # LoRA strength for merge, default 1.0
@dataclasses.dataclass
@@ -47,6 +47,31 @@ class ListLorasReq:
pass
def format_lora_message(
lora_nickname: Union[str, List[str]],
target: Union[str, List[str]],
strength: Union[float, List[float]],
) -> tuple[str, str, str]:
"""Format success message for single or multiple LoRAs"""
if isinstance(lora_nickname, list):
nickname_str = ", ".join(lora_nickname)
target_str = ", ".join(target) if isinstance(target, list) else target
strength_str = (
", ".join(f"{s:.2f}" for s in strength)
if isinstance(strength, list)
else f"{strength:.2f}"
)
else:
nickname_str = lora_nickname
target_str = target if isinstance(target, str) else ", ".join(target)
strength_str = (
f"{strength:.2f}"
if isinstance(strength, (int, float))
else ", ".join(f"{s:.2f}" for s in strength)
)
return nickname_str, target_str, strength_str
def _parse_size(size: str) -> tuple[int, int] | tuple[None, None]:
try:
parts = size.lower().replace(" ", "").split("x")

View File

@@ -55,6 +55,9 @@ class BaseLayerWithLoRA(nn.Module):
self.disable_lora: bool = True
self.lora_rank = lora_rank
self.lora_alpha = lora_alpha
self.lora_weights_list: list[
tuple[torch.nn.Parameter, torch.nn.Parameter, str | None, float]
] = []
self.lora_path: str | None = None
self.strength: float = 1.0
@@ -77,6 +80,7 @@ class BaseLayerWithLoRA(nn.Module):
lora_B = self.lora_B.to_local()
lora_A = self.lora_A.to_local()
# TODO: Support multiple LoRA adapters when use not merged mode
if not self.merged and not self.disable_lora:
lora_A_sliced = self.slice_lora_a_weights(lora_A.to(x, non_blocking=True))
lora_B_sliced = self.slice_lora_b_weights(lora_B.to(x, non_blocking=True))
@@ -104,15 +108,70 @@ class BaseLayerWithLoRA(nn.Module):
B: torch.Tensor,
lora_path: str | None = None,
strength: float = 1.0,
clear_existing: bool = False,
) -> None:
self.lora_A = torch.nn.Parameter(
"""
Set LoRA weights. Supports multiple LoRA adapters.
Args:
A: LoRA A weight tensor
B: LoRA B weight tensor
lora_path: Path to the LoRA adapter (for logging)
strength: LoRA strength
clear_existing: If True, clear existing LoRA weights before adding new one.
If False, append to existing list (for multi-LoRA support).
"""
lora_A_param = torch.nn.Parameter(
A
) # share storage with weights in the pipeline
self.lora_B = torch.nn.Parameter(B)
self.disable_lora = False
self.strength = strength
self.merge_lora_weights()
lora_B_param = torch.nn.Parameter(B)
if clear_existing:
self.lora_weights_list.clear()
# Also clear backward compatibility attributes
self.lora_A = None
self.lora_B = None
self.lora_path = None
self.strength = 1.0
# Add to list for multi-LoRA support
self.lora_weights_list.append((lora_A_param, lora_B_param, lora_path, strength))
# Set backward compatibility attributes to point to the last LoRA (for single LoRA case)
# This ensures backward compatibility while supporting multiple LoRA
self.lora_A = lora_A_param
self.lora_B = lora_B_param
self.lora_path = lora_path
self.strength = strength
self.disable_lora = False
self.merge_lora_weights()
@torch.no_grad()
def _merge_lora_into_data(
self,
data: torch.Tensor,
lora_list: list[
tuple[torch.nn.Parameter, torch.nn.Parameter, str | None, float]
],
) -> None:
"""
Merge all LoRA adapters into the data tensor in-place.
Args:
data: The base weight tensor to merge LoRA into (modified in-place)
lora_list: List of (lora_A, lora_B, lora_path, lora_strength) tuples
"""
# Merge all LoRA adapters in order
for lora_A, lora_B, _, lora_strength in lora_list:
lora_delta = self.slice_lora_b_weights(
lora_B.to(data)
) @ self.slice_lora_a_weights(lora_A.to(data))
# Apply lora_alpha / lora_rank scaling for consistency with forward()
if self.lora_alpha is not None and self.lora_rank is not None:
if self.lora_alpha != self.lora_rank:
lora_delta = lora_delta * (self.lora_alpha / self.lora_rank)
data += lora_strength * lora_delta
@torch.no_grad()
def merge_lora_weights(self, strength: float | None = None) -> None:
@@ -124,9 +183,15 @@ class BaseLayerWithLoRA(nn.Module):
if self.merged:
self.unmerge_lora_weights()
assert (
self.lora_A is not None and self.lora_B is not None
), "LoRA weights not set. Please set them first."
# Use lora_weights_list if available, otherwise fall back to single LoRA for backward compatibility
lora_list = self.lora_weights_list if self.lora_weights_list else []
if not lora_list and self.lora_A is not None and self.lora_B is not None:
lora_list = [(self.lora_A, self.lora_B, self.lora_path, self.strength)]
if not lora_list:
raise ValueError("LoRA weights not set. Please set them first.")
if isinstance(self.base_layer.weight, DTensor):
mesh = self.base_layer.weight.data.device_mesh
unsharded_base_layer = ReplicatedLinear(
@@ -143,14 +208,9 @@ class BaseLayerWithLoRA(nn.Module):
data = self.base_layer.weight.data.to(
get_local_torch_device()
).full_tensor()
lora_delta = self.slice_lora_b_weights(self.lora_B).to(
data
) @ self.slice_lora_a_weights(self.lora_A).to(data)
# Apply lora_alpha / lora_rank scaling for consistency with forward()
if self.lora_alpha is not None and self.lora_rank is not None:
if self.lora_alpha != self.lora_rank:
lora_delta = lora_delta * (self.lora_alpha / self.lora_rank)
data += self.strength * lora_delta
self._merge_lora_into_data(data, lora_list)
unsharded_base_layer.weight = nn.Parameter(data.to(current_device))
if isinstance(getattr(self.base_layer, "bias", None), DTensor):
unsharded_base_layer.bias = nn.Parameter(
@@ -173,14 +233,9 @@ class BaseLayerWithLoRA(nn.Module):
else:
current_device = self.base_layer.weight.data.device
data = self.base_layer.weight.data.to(get_local_torch_device())
lora_delta = self.slice_lora_b_weights(
self.lora_B.to(data)
) @ self.slice_lora_a_weights(self.lora_A.to(data))
# Apply lora_alpha / lora_rank scaling for consistency with forward()
if self.lora_alpha is not None and self.lora_rank is not None:
if self.lora_alpha != self.lora_rank:
lora_delta = lora_delta * (self.lora_alpha / self.lora_rank)
data += self.strength * lora_delta
self._merge_lora_into_data(data, lora_list)
self.base_layer.weight.data = data.to(current_device, non_blocking=True)
self.merged = True
@@ -403,6 +458,7 @@ class LinearWithLoRA(BaseLayerWithLoRA):
lora_B = self.lora_B.to_local()
lora_A = self.lora_A.to_local()
# TODO: Support multiple LoRA adapters when use not merged mode
if not self.merged and not self.disable_lora:
lora_A_sliced = self.slice_lora_a_weights(lora_A.to(x, non_blocking=True))
lora_B_sliced = self.slice_lora_b_weights(lora_B.to(x, non_blocking=True))

View File

@@ -4,7 +4,7 @@
import multiprocessing as mp
import os
import time
from typing import List
from typing import List, Union
import torch
from setproctitle import setproctitle
@@ -206,19 +206,20 @@ class GPUWorker:
def set_lora(
self,
lora_nickname: str,
lora_path: str | None = None,
target: str = "all",
strength: float = 1.0,
lora_nickname: Union[str, List[str]],
lora_path: Union[str, None, List[Union[str, None]]] = None,
target: Union[str, List[str]] = "all",
strength: Union[float, List[float]] = 1.0,
) -> OutputBatch:
"""
Set the LoRA adapter for the pipeline.
Set the LoRA adapter(s) for the pipeline.
Supports both single LoRA (backward compatible) and multiple LoRA adapters.
Args:
lora_nickname: The nickname of the adapter.
lora_path: Path to the LoRA adapter.
target: Which transformer(s) to apply the LoRA to.
strength: LoRA strength for merge, default 1.0.
lora_nickname: The nickname(s) of the adapter(s). Can be a string or a list of strings.
lora_path: Path(s) to the LoRA adapter(s). Can be a string, None, or a list of strings/None.
target: Which transformer(s) to apply the LoRA to. Can be a string or a list of strings.
strength: LoRA strength(s) for merge, default 1.0. Can be a float or a list of floats.
"""
if not isinstance(self.pipeline, LoRAPipeline):
return OutputBatch(error="Lora is not enabled")

View File

@@ -48,6 +48,7 @@ class LoRAPipeline(ComposedPipelineBase):
cur_adapter_name: dict[str, str]
cur_adapter_path: dict[str, str]
cur_adapter_strength: dict[str, float] # Track current strength per module
cur_adapter_config: dict[str, tuple[list[str], list[float]]]
# [dit_layer_name] = wrapped_lora_layer
lora_layers: dict[str, BaseLayerWithLoRA]
lora_layers_critic: dict[str, BaseLayerWithLoRA]
@@ -74,6 +75,8 @@ class LoRAPipeline(ComposedPipelineBase):
self.cur_adapter_name = {}
self.cur_adapter_path = {}
self.cur_adapter_strength = {}
# Track full LoRA config: {module_name: (nickname_list, strength_list)}
self.cur_adapter_config = {}
self.lora_layers = {}
self.lora_layers_critic = {}
self.lora_layers_transformer_2 = {}
@@ -297,50 +300,167 @@ class LoRAPipeline(ComposedPipelineBase):
converted_count_critic,
)
def _normalize_lora_params(
self,
lora_nickname: str | list[str],
lora_path: str | None | list[str | None],
strength: float | list[float],
target: str | list[str],
) -> tuple[list[str], list[str | None], list[float], list[str]]:
"""
Normalize LoRA parameters to lists for multi-LoRA support.
Requirements:
- each nickname must have a corresponding lora_path (no implicit repeat)
- strength / target if scalar broadcast, else length must match nickname
"""
# nickname
if isinstance(lora_nickname, str):
lora_nicknames = [lora_nickname]
else:
lora_nicknames = lora_nickname
# lora_path: require 1:1 mapping with nickname (no implicit repeat)
if isinstance(lora_path, list):
lora_paths = lora_path
else:
lora_paths = [lora_path]
if len(lora_paths) != len(lora_nicknames):
raise ValueError(
f"Length mismatch: lora_nickname has {len(lora_nicknames)} items, "
f"but lora_path has {len(lora_paths)} items. "
"Provide one path per nickname."
)
# strength and target: allow scalar broadcast, else length must match
if isinstance(strength, (int, float)):
strengths = [float(strength)] * len(lora_nicknames)
else:
strengths = [float(s) for s in strength]
if len(strengths) != len(lora_nicknames):
raise ValueError(
f"Length mismatch: lora_nickname has {len(lora_nicknames)} items, "
f"but strength has {len(strengths)} items"
)
if isinstance(target, str):
targets = [target] * len(lora_nicknames)
else:
targets = target
if len(targets) != len(lora_nicknames):
raise ValueError(
f"Length mismatch: lora_nickname has {len(lora_nicknames)} items, "
f"but target has {len(targets)} items"
)
return lora_nicknames, lora_paths, strengths, targets
def _check_lora_config_matches(
self,
module_name: str,
target_nicknames: list[str],
target_strengths: list[float],
adapter_updated: bool,
) -> bool:
"""
Check if current LoRA configuration matches the target configuration.
Args:
module_name: The name of the module to check.
target_nicknames: List of LoRA nicknames to apply.
target_strengths: List of LoRA strengths to apply.
adapter_updated: Whether any adapter was updated/loaded.
Returns:
True if the configuration matches exactly (including order and strength), False otherwise.
"""
if not self.is_lora_merged.get(module_name, False):
return False
if adapter_updated:
return False # Adapter was updated, need to reapply
stored_config = self.cur_adapter_config.get(module_name)
if stored_config is None:
return False
stored_nicknames, stored_strengths = stored_config
# Compare: nickname list and strength list must match exactly (including order)
return (
stored_nicknames == target_nicknames
and stored_strengths == target_strengths
)
def _apply_lora_to_layers(
self,
lora_layers: dict[str, BaseLayerWithLoRA],
lora_nickname: str,
lora_path: str | None,
lora_nicknames: list[str],
lora_paths: list[str | None],
rank: int,
strength: float = 1.0,
strengths: list[float],
clear_existing: bool = False,
) -> int:
"""
Apply LoRA weights to the given lora_layers.
Apply LoRA weights to the given lora_layers. Supports multiple LoRA adapters.
Args:
lora_layers: The dictionary of LoRA layers to apply weights to.
lora_nickname: The nickname of the LoRA adapter.
lora_path: The path to the LoRA adapter.
lora_nicknames: The list of nicknames of the LoRA adapters.
lora_paths: The list of paths to the LoRA adapters. Must match length of lora_nicknames.
rank: The distributed rank (for logging).
strength: LoRA strength for merge, default 1.0.
strengths: The list of LoRA strengths for merge. Must match length of lora_nicknames.
clear_existing: If True, clear existing LoRA weights before adding new ones.
Returns:
The number of layers that had LoRA weights applied.
"""
if len(lora_paths) != len(lora_nicknames):
raise ValueError(
f"Length mismatch: lora_nicknames has {len(lora_nicknames)} items, "
f"but lora_paths has {len(lora_paths)} items"
)
if len(strengths) != len(lora_nicknames):
raise ValueError(
f"Length mismatch: lora_nicknames has {len(lora_nicknames)} items, "
f"but strengths has {len(strengths)} items"
)
adapted_count = 0
for name, layer in lora_layers.items():
lora_A_name = name + ".lora_A"
lora_B_name = name + ".lora_B"
if (
lora_A_name in self.lora_adapters[lora_nickname]
and lora_B_name in self.lora_adapters[lora_nickname]
# Apply all LoRA adapters in order
for idx, (nickname, path, lora_strength) in enumerate(
zip(lora_nicknames, lora_paths, strengths)
):
layer.set_lora_weights(
self.lora_adapters[lora_nickname][lora_A_name],
self.lora_adapters[lora_nickname][lora_B_name],
lora_path=lora_path,
strength=strength,
)
adapted_count += 1
else:
if rank == 0:
logger.warning(
"LoRA adapter %s does not contain the weights for layer '%s'. LoRA will not be applied to it.",
lora_path,
name,
lora_A_name = name + ".lora_A"
lora_B_name = name + ".lora_B"
if (
lora_A_name in self.lora_adapters[nickname]
and lora_B_name in self.lora_adapters[nickname]
):
layer.set_lora_weights(
self.lora_adapters[nickname][lora_A_name],
self.lora_adapters[nickname][lora_B_name],
lora_path=path,
strength=lora_strength,
clear_existing=(
clear_existing and idx == 0
), # Only clear on first LoRA
)
layer.disable_lora = True
adapted_count += 1
else:
if rank == 0 and idx == 0: # Only warn for first missing LoRA
logger.warning(
"LoRA adapter %s does not contain the weights for layer '%s'. LoRA will not be applied to it.",
path,
name,
)
# Only disable if no LoRA was applied at all
if idx == len(lora_nicknames) - 1:
has_any_lora = any(
name + ".lora_A" in self.lora_adapters[n]
and name + ".lora_B" in self.lora_adapters[n]
for n in lora_nicknames
)
if not has_any_lora:
layer.disable_lora = True
return adapted_count
def is_lora_effective(self, target: str = "all") -> bool:
@@ -438,46 +558,25 @@ class LoRAPipeline(ComposedPipelineBase):
def set_lora(
self,
lora_nickname: str,
lora_path: str | None = None,
target: str = "all",
strength: float = 1.0,
lora_nickname: str | list[str],
lora_path: str | None | list[str | None] = None,
target: str | list[str] = "all",
strength: float | list[float] = 1.0,
): # type: ignore
"""
Load a LoRA adapter into the pipeline and apply it to the specified transformer(s).
Args:
lora_nickname: The "nick name" of the adapter when referenced in the pipeline.
lora_path: The path to the adapter, either a local path or a Hugging Face repo id.
target: Which transformer(s) to apply the LoRA to. One of:
- "all": Apply to all transformers (default, backward compatible)
- "transformer": Apply only to the primary transformer (high noise for Wan2.2)
- "transformer_2": Apply only to transformer_2 (low noise for Wan2.2)
- "critic": Apply only to the critic model (fake_score_transformer)
strength: LoRA strength for merge, default 1.0.
Load LoRA adapter(s) into the pipeline and apply them to the specified transformer(s).
Supports both single LoRA (backward compatible) and multiple LoRA adapters.
"""
if target not in self.VALID_TARGETS:
raise ValueError(
f"Invalid target: {target}. Valid targets: {self.VALID_TARGETS}"
)
# Normalize inputs to lists for multi-LoRA support
lora_nicknames, lora_paths, strengths, targets = self._normalize_lora_params(
lora_nickname, lora_path, strength, target
)
# Check if any target module has a different LoRA merged
target_modules, error = self._get_target_lora_layers(target)
if error:
logger.warning("set_lora: %s", error)
for module_name, _ in target_modules:
if (
self.is_lora_merged.get(module_name, False)
and self.cur_adapter_name.get(module_name) != lora_nickname
):
raise ValueError(
f"LoRA '{self.cur_adapter_name.get(module_name)}' is currently merged in {module_name}. "
"Please call 'unmerge_lora_weights' before setting a new LoRA."
)
if lora_nickname not in self.lora_adapters and lora_path is None:
# Validate targets
invalid_targets = [t for t in targets if t not in self.VALID_TARGETS]
if invalid_targets:
raise ValueError(
f"Adapter {lora_nickname} not found in the pipeline. Please provide lora_path to load it."
f"Invalid target(s): {invalid_targets}. Valid targets: {self.VALID_TARGETS}"
)
# Disable layerwise offload before convert_to_lora_layers to ensure weights are accessible
@@ -489,64 +588,100 @@ class LoRAPipeline(ComposedPipelineBase):
):
self.convert_to_lora_layers()
# Re-fetch target_modules after convert_to_lora_layers() to get populated dicts
target_modules, error = self._get_target_lora_layers(target)
if error:
logger.warning("set_lora: %s", error)
# Check adapter presence and load missing adapters
adapter_updated = False
rank = dist.get_rank()
should_load = False
if lora_path is not None:
if lora_nickname not in self.loaded_adapter_paths:
should_load = True
elif self.loaded_adapter_paths[lora_nickname] != lora_path:
should_load = True
# load required adapters
for nickname, path in zip(lora_nicknames, lora_paths):
if nickname not in self.lora_adapters and path is None:
raise ValueError(
f"Adapter {nickname} not found in the pipeline. Please provide lora_path to load it."
)
# Check if adapter needs to be loaded
should_load = False
if path is not None:
if nickname not in self.loaded_adapter_paths:
should_load = True
elif self.loaded_adapter_paths[nickname] != path:
should_load = True
if should_load:
adapter_updated = True
self.load_lora_adapter(path, nickname, rank)
if should_load:
adapter_updated = True
self.load_lora_adapter(lora_path, lora_nickname, rank)
# Group by target to apply separately
target_to_indices = {}
for idx, tgt in enumerate(targets):
if tgt not in target_to_indices:
target_to_indices[tgt] = []
target_to_indices[tgt].append(idx)
# Check if we can skip (same adapter already applied to all target modules with same strength)
all_already_applied = all(
not adapter_updated
and self.cur_adapter_name.get(module_name) == lora_nickname
and self.is_lora_merged.get(module_name, False)
and self.cur_adapter_strength.get(module_name) == strength
for module_name, _ in target_modules
adapted_count = 0
for tgt, idx_list in target_to_indices.items():
target_modules, error = self._get_target_lora_layers(tgt)
if error:
logger.warning("set_lora: %s", error)
if not target_modules:
continue
# Disable layerwise offload if enabled: load all layers to GPU
# the LoRA weights merging process requires weights being on device
with self._temporarily_disable_offload(target_modules=target_modules):
tgt_nicknames = [lora_nicknames[i] for i in idx_list]
tgt_paths = [lora_paths[i] for i in idx_list]
tgt_strengths = [strengths[i] for i in idx_list]
merged_name = (
",".join(tgt_nicknames)
if len(tgt_nicknames) > 1
else tgt_nicknames[0]
)
# Skip if LoRA configuration matches exactly (including order and strength)
# Since all modules for the same target apply the same config, checking one is sufficient
first_module_name, _ = target_modules[0]
if self._check_lora_config_matches(
first_module_name, tgt_nicknames, tgt_strengths, adapter_updated
):
logger.info("LoRA configuration matches exactly, skipping")
continue
# Apply LoRA to modules for this target
for module_name, lora_layers_dict in target_modules:
count = self._apply_lora_to_layers(
lora_layers_dict,
tgt_nicknames,
tgt_paths,
rank,
tgt_strengths,
clear_existing=True,
)
adapted_count += count
self.cur_adapter_name[module_name] = merged_name
self.cur_adapter_path[module_name] = ",".join(
str(p or self.loaded_adapter_paths.get(n, ""))
for n, p in zip(tgt_nicknames, tgt_paths)
)
self.is_lora_merged[module_name] = True
self.cur_adapter_strength[module_name] = tgt_strengths[0]
# Store full configuration for multi-LoRA support (preserves order and all strengths)
self.cur_adapter_config[module_name] = (
tgt_nicknames.copy(),
tgt_strengths.copy(),
)
logger.info(
"Rank %d: LoRA adapter(s) %s applied to %d layers (targets: %s, strengths: %s)",
rank,
", ".join(map(str, lora_paths)) if lora_paths else None,
adapted_count,
", ".join(targets) if len(set(targets)) > 1 else targets[0],
(
", ".join(f"{s:.2f}" for s in strengths)
if len(strengths) > 1
else f"{strengths[0]:.2f}"
),
)
if all_already_applied:
return
# Disable layerwise offload if enabled: load all layers to GPU
with self._temporarily_disable_offload(target_modules=target_modules):
# Apply LoRA to target modules (now all layers are on GPU)
adapted_count = 0
for module_name, lora_layers_dict in target_modules:
count = self._apply_lora_to_layers(
lora_layers_dict,
lora_nickname,
lora_path,
rank,
strength,
)
adapted_count += count
self.cur_adapter_name[module_name] = lora_nickname
self.cur_adapter_path[module_name] = (
lora_path or self.loaded_adapter_paths.get(lora_nickname, "")
)
self.is_lora_merged[module_name] = True
self.cur_adapter_strength[module_name] = strength
logger.info(
"Rank %d: LoRA adapter %s applied to %d layers (target: %s, strength: %s)",
rank,
lora_path,
adapted_count,
target,
strength,
)
def merge_lora_weights(self, target: str = "all", strength: float = 1.0) -> None:
"""
@@ -649,7 +784,8 @@ class LoRAPipeline(ComposedPipelineBase):
continue
self.is_lora_merged[module_name] = False
self.cur_adapter_strength.pop(module_name, None)
logger.info("LoRA weights unmerged for %s", module_name)
self.cur_adapter_config.pop(module_name, None)
logger.info("LoRA weights unmerged for %s", module_name)
def get_lora_status(self) -> dict[str, Any]:
"""

View File

@@ -529,6 +529,31 @@
"expected_avg_denoise_ms": 94.15,
"expected_median_denoise_ms": 102.03
},
"zimage_image_t2i_multi_lora": {
"stages_ms": {
"InputValidationStage": 0.04,
"TextEncodingStage": 103.81,
"ConditioningStage": 0.01,
"TimestepPreparationStage": 1.3,
"LatentPreparationStage": 0.11,
"DenoisingStage": 813.7,
"DecodingStage": 34.51
},
"denoise_step_ms": {
"0": 30.35,
"1": 74.53,
"2": 99.34,
"3": 100.92,
"4": 99.46,
"5": 100.57,
"6": 99.72,
"7": 100.86,
"8": 103.87
},
"expected_e2e_ms": 955.14,
"expected_avg_denoise_ms": 89.96,
"expected_median_denoise_ms": 99.72
},
"zimage_image_t2i_2_gpus": {
"stages_ms": {
"InputValidationStage": 0.08,

View File

@@ -556,6 +556,76 @@ Consider updating perf_baselines.json with the snippets below:
assert resp.status_code == 200, f"Dynamic set_lora failed: {resp.text}"
logger.info("[Dynamic LoRA] set_lora succeeded for %s", case.id)
def _test_multi_lora_e2e(
self,
ctx: ServerContext,
case: DiffusionTestCase,
generate_fn: Callable[[str, openai.Client], str],
first_lora_path: str,
second_lora_path: str,
) -> None:
"""
Test multiple LoRA adapters with different set_lora input scenarios.
Tests: basic multi-LoRA, different strengths, cached adapters, switch back to single.
"""
base_url = f"http://localhost:{ctx.port}/v1"
client = OpenAI(base_url=base_url, api_key="dummy")
# Test 1: Basic multi-LoRA with list format
resp = requests.post(
f"{base_url}/set_lora",
json={
"lora_nickname": ["default", "lora2"],
"lora_path": [first_lora_path, second_lora_path],
"target": "all",
"strength": [1.0, 1.0],
},
)
assert (
resp.status_code == 200
), f"set_lora with multiple adapters failed: {resp.text}"
assert generate_fn(case.id, client) is not None
# Test 2: Different strengths
resp = requests.post(
f"{base_url}/set_lora",
json={
"lora_nickname": ["default", "lora2"],
"lora_path": [first_lora_path, second_lora_path],
"target": "all",
"strength": [0.8, 0.5],
},
)
assert (
resp.status_code == 200
), f"set_lora with different strengths failed: {resp.text}"
assert generate_fn(case.id, client) is not None
# Test 3: Different targets
requests.post(f"{base_url}/set_lora", json={"lora_nickname": "default"})
resp = requests.post(
f"{base_url}/set_lora",
json={
"lora_nickname": ["default", "lora2"],
"lora_path": [first_lora_path, second_lora_path],
"target": ["transformer", "transformer_2"],
"strength": [0.8, 0.5],
},
)
assert (
resp.status_code == 200
), f"set_lora with cached adapters failed: {resp.text}"
assert generate_fn(case.id, client) is not None
# Test 4: Switch back to single LoRA
resp = requests.post(f"{base_url}/set_lora", json={"lora_nickname": "default"})
assert (
resp.status_code == 200
), f"set_lora back to single adapter failed: {resp.text}"
assert generate_fn(case.id, client) is not None
logger.info("[Multi-LoRA] All multi-LoRA tests passed for %s", case.id)
def _test_v1_models_endpoint(
self, ctx: ServerContext, case: DiffusionTestCase
) -> None:
@@ -681,3 +751,21 @@ Consider updating perf_baselines.json with the snippets below:
# LoRA API functionality test with E2E validation (only for LoRA-enabled cases)
if case.server_args.lora_path or case.server_args.dynamic_lora_path:
self._test_lora_api_functionality(diffusion_server, case, generate_fn)
# Test dynamic LoRA switching (requires a second LoRA adapter)
if case.server_args.second_lora_path:
self._test_lora_dynamic_switch_e2e(
diffusion_server,
case,
generate_fn,
case.server_args.second_lora_path,
)
# Test multi-LoRA functionality
self._test_multi_lora_e2e(
diffusion_server,
case,
generate_fn,
case.server_args.lora_path,
case.server_args.second_lora_path,
)

View File

@@ -155,6 +155,9 @@ class DiffusionServerArgs:
dynamic_lora_path: str | None = (
None # LoRA path for dynamic loading test (loaded via set_lora after startup)
)
second_lora_path: str | None = (
None # Second LoRA adapter path for multi-LoRA testing
)
# misc
enable_warmup: bool = False
@@ -361,6 +364,17 @@ ONE_GPU_CASES_A: list[DiffusionTestCase] = [
),
T2I_sampling_params,
),
# Multi-LoRA test case for Z-Image-Turbo
DiffusionTestCase(
"zimage_image_t2i_multi_lora",
DiffusionServerArgs(
model_path="Tongyi-MAI/Z-Image-Turbo",
modality="image",
lora_path="reverentelusarca/elusarca-anime-style-lora-z-image-turbo",
second_lora_path="tarn59/pixel_art_style_lora_z_image_turbo",
),
T2I_sampling_params,
),
# === Text and Image to Image (TI2I) ===
DiffusionTestCase(
"qwen_image_edit_ti2i",