chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] feat: support LoRA (#13859)

Browse Source
This commit is contained in:
Mick
2025-11-26 00:21:33 +08:00
committed by GitHub
parent 46673b4224
commit dfd7ab9682
29 changed files with 463 additions and 278 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
python/sglang/multimodal_gen/configs/models/dits/base.py
View File

@@ -13,9 +13,16 @@ from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
class DiTArchConfig(ArchConfig):
_fsdp_shard_conditions: list = field(default_factory=list)
_compile_conditions: list = field(default_factory=list)
# convert weights name from HF-format to SGLang-dit-format
param_names_mapping: dict = field(default_factory=dict)
reverse_param_names_mapping: dict = field(default_factory=dict)
# convert weights name from misc-format to HF-format
# usually applicable if the LoRA is trained with official repo implementation
lora_param_names_mapping: dict = field(default_factory=dict)
# Reverse mapping for saving checkpoints: custom -> hf
reverse_param_names_mapping: dict = field(default_factory=dict)
_supported_attention_backends: set[AttentionBackendEnum] = field(
default_factory=lambda: {
AttentionBackendEnum.SLIDING_TILE_ATTN,

6
python/sglang/multimodal_gen/configs/models/dits/flux.py
View File

@@ -21,6 +21,12 @@ class FluxArchConfig(DiTArchConfig):
guidance_embeds: bool = False
axes_dims_rope: Tuple[int, int, int] = (16, 56, 56)
param_names_mapping: dict = field(
default_factory=lambda: {
r"transformer\.(\w*)\.(.*)$": r"\1.\2",
}
)
def __post_init__(self):
super().__post_init__()
self.out_channels = self.out_channels or self.in_channels

1
python/sglang/multimodal_gen/configs/models/dits/hunyuanvideo.py
View File

@@ -148,7 +148,6 @@ class HunyuanVideoArchConfig(DiTArchConfig):
}
)
# Reverse mapping for saving checkpoints: custom -> hf
reverse_param_names_mapping: dict = field(default_factory=lambda: {})
patch_size: int = 2

7
python/sglang/multimodal_gen/configs/models/dits/qwenimage.py
View File

@@ -21,6 +21,12 @@ class QwenImageArchConfig(DiTArchConfig):
guidance_embeds: bool = False
axes_dims_rope: Tuple[int, int, int] = (16, 56, 56)
param_names_mapping: dict = field(
default_factory=lambda: {
r"^(transformer_blocks\.\d+\.attn\..*\.lora_[AB])\.default$": r"\1",
}
)
def __post_init__(self):
super().__post_init__()
self.out_channels = self.out_channels or self.in_channels
@@ -30,7 +36,6 @@ class QwenImageArchConfig(DiTArchConfig):
@dataclass
class QwenImageDitConfig(DiTConfig):
arch_config: DiTArchConfig = field(default_factory=QwenImageArchConfig)
prefix: str = "qwenimage"

1
python/sglang/multimodal_gen/configs/models/dits/wanvideo.py
View File

@@ -37,7 +37,6 @@ class WanVideoArchConfig(DiTArchConfig):
}
)
# Reverse mapping for saving checkpoints: custom -> hf
reverse_param_names_mapping: dict = field(default_factory=lambda: {})
# Some LoRA adapters use the original official layer names instead of hf layer names,

1
python/sglang/multimodal_gen/configs/models/encoders/qwen_image.py
View File

@@ -46,6 +46,7 @@ class QwenImageArchConfig(TextEncoderArchConfig):
head_dim: int | None = None
hidden_state_skip_layer: int = 2
text_len: int = 256
stacked_params_mapping: list[tuple[str, str, str]] = field(
default_factory=lambda: [
# (param_name, shard_name, shard_id)

111
python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py
View File

@@ -21,18 +21,25 @@ import torchvision
from einops import rearrange
from sglang.multimodal_gen.configs.sample.base import DataType, SamplingParams
from sglang.multimodal_gen.runtime.entrypoints.openai.utils import (
MergeLoraWeightsReq,
SetLoraReq,
UnmergeLoraWeightsReq,
)
from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request
from sglang.multimodal_gen.runtime.launch_server import launch_server
from sglang.multimodal_gen.runtime.managers.schedulerbase import SchedulerBase
from sglang.multimodal_gen.runtime.pipelines_core import Req
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch
from sglang.multimodal_gen.runtime.server_args import PortArgs, ServerArgs
from sglang.multimodal_gen.runtime.sync_scheduler_client import sync_scheduler_client
from sglang.multimodal_gen.runtime.utils.logging_utils import (
init_logger,
suppress_loggers,
suppress_other_loggers,
)
suppress_loggers(["imageio", "imageio_ffmpeg", "PIL", "PIL_Image"])
logger = init_logger(__name__)
# TODO: move to somewhere appropriate
@@ -71,6 +78,9 @@ class DiffGenerator:
# The executor is now a client to the Scheduler service
self.local_scheduler_process: list[mp.Process] | None = None
self.owns_scheduler_client: bool = False
self._current_lora_path: str | None = None
self._current_lora_nickname: str | None = None
self._is_lora_merged: bool = False
@classmethod
def from_pretrained(
@@ -111,7 +121,6 @@ class DiffGenerator:
Returns:
The created DiffGenerator
"""
executor_class = SchedulerBase.get_class(server_args)
instance = cls(
server_args=server_args,
)
@@ -361,32 +370,98 @@ class DiffGenerator:
"""
return sync_scheduler_client.forward(batch)
def set_lora_adapter(
self, lora_nickname: str, lora_path: str | None = None
) -> None:
# self.scheduler.set_lora_adapter(lora_nickname, lora_path)
pass # Removed as per edit hint
# LoRA
def _send_lora_request(self, req: Any, success_msg: str, failure_msg: str):
response = sync_scheduler_client.forward(req)
if isinstance(response, dict) and response.get("status") == "ok":
logger.info(success_msg)
else:
error_msg = (
response.get("message", "Unknown error")
if isinstance(response, dict)
else "Unknown response format"
)
raise RuntimeError(f"{failure_msg}: {error_msg}")
def set_lora(self, lora_nickname: str, lora_path: str | None = None) -> None:
req = SetLoraReq(lora_nickname=lora_nickname, lora_path=lora_path)
self._send_lora_request(
req,
f"Successfully set LoRA adapter: {lora_nickname}",
"Failed to set LoRA adapter",
)
def unmerge_lora_weights(self) -> None:
"""
Use unmerged weights for inference to produce outputs that align with
validation outputs generated during training.
"""
# self.scheduler.unmerge_lora_weights()
pass # Removed as per edit hint
req = UnmergeLoraWeightsReq()
self._send_lora_request(
req,
"Successfully unmerged LoRA weights",
"Failed to unmerge LoRA weights",
)
self._is_lora_merged = False
def merge_lora_weights(self) -> None:
# self.scheduler.merge_lora_weights()
pass # Removed as per edit hint
req = MergeLoraWeightsReq()
self._send_lora_request(
req, "Successfully merged LoRA weights", "Failed to merge LoRA weights"
)
self._is_lora_merged = True
def _ensure_lora_state(
self,
lora_path: str | None,
lora_nickname: str | None = None,
merge_lora: bool = True,
) -> None:
if lora_path is None:
if self._is_lora_merged:
self.unmerge_lora_weights()
self._current_lora_path = None
self._current_lora_nickname = None
self._is_lora_merged = False
return
lora_nickname = lora_nickname or self.server_args.lora_nickname
if self._current_lora_path != lora_path:
if self._is_lora_merged:
self.unmerge_lora_weights()
self._is_lora_merged = False
self.set_lora(lora_nickname, lora_path)
self._current_lora_path = lora_path
self._current_lora_nickname = lora_nickname
self._is_lora_merged = False
if merge_lora and not self._is_lora_merged:
self.merge_lora_weights()
elif not merge_lora:
self._is_lora_merged = False
def generate_with_lora(
self,
prompt: str | list[str] | None = None,
sampling_params: SamplingParams | None = None,
*,
lora_path: str | None = None,
lora_nickname: str | None = None,
merge_lora: bool = True,
**kwargs,
):
self._ensure_lora_state(
lora_path=lora_path, lora_nickname=lora_nickname, merge_lora=merge_lora
)
return self.generate(
prompt=prompt,
sampling_params=sampling_params,
**kwargs,
)
def shutdown(self):
"""
Shutdown the generator.
If in local mode, it also shuts down the scheduler server.
"""
# This sends the shutdown command to the server
# self.scheduler.shutdown()
# sends the shutdown command to the server
if self.local_scheduler_process:
logger.info("Waiting for local worker processes to terminate...")
for process in self.local_scheduler_process:

3
python/sglang/multimodal_gen/runtime/entrypoints/http_server.py
View File

@@ -55,6 +55,9 @@ def create_app(server_args: ServerArgs):
app.include_router(health_router)
from sglang.multimodal_gen.runtime.entrypoints.openai import common_api
app.include_router(common_api.router)
app.include_router(image_api.router)
app.include_router(video_api.router)

62
python/sglang/multimodal_gen/runtime/entrypoints/openai/common_api.py Normal file
View File

@@ -0,0 +1,62 @@
from typing import Any, Optional
from fastapi import APIRouter, Body, HTTPException
from sglang.multimodal_gen.runtime.entrypoints.openai.utils import (
MergeLoraWeightsReq,
SetLoraReq,
UnmergeLoraWeightsReq,
)
from sglang.multimodal_gen.runtime.scheduler_client import scheduler_client
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
router = APIRouter(prefix="/v1")
logger = init_logger(__name__)
async def _handle_lora_request(req: Any, success_msg: str, failure_msg: str):
try:
response = await scheduler_client.forward(req)
if isinstance(response, dict) and response.get("status") == "ok":
return {"status": "ok", "message": success_msg}
else:
error_msg = (
response.get("message", "Unknown error")
if isinstance(response, dict)
else "Unknown response format"
)
raise HTTPException(status_code=500, detail=f"{failure_msg}: {error_msg}")
except Exception as e:
if isinstance(e, HTTPException):
raise
logger.error(f"Error during '{failure_msg}': {e}", exc_info=True)
raise HTTPException(status_code=500, detail=str(e))
@router.post("/set_lora")
async def set_lora(
lora_nickname: str = Body(..., embed=True),
lora_path: Optional[str] = Body(None, embed=True),
):
req = SetLoraReq(lora_nickname=lora_nickname, lora_path=lora_path)
return await _handle_lora_request(
req,
f"Successfully set LoRA adapter: {lora_nickname}",
"Failed to set LoRA adapter",
)
@router.post("/merge_lora_weights")
async def merge_lora_weights():
req = MergeLoraWeightsReq()
return await _handle_lora_request(
req, "Successfully merged LoRA weights", "Failed to merge LoRA weights"
)
@router.post("/unmerge_lora_weights")
async def unmerge_lora_weights():
req = UnmergeLoraWeightsReq()
return await _handle_lora_request(
req, "Successfully unmerged LoRA weights", "Failed to unmerge LoRA weights"
)

22
python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py
View File

@@ -21,7 +21,7 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.stores import IMAGE_STORE
from sglang.multimodal_gen.runtime.entrypoints.openai.utils import (
_parse_size,
_save_upload_to_path,
post_process_sample,
process_generation_batch,
)
from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
@@ -110,15 +110,7 @@ async def generations(
sampling_params=sampling,
)
# Run synchronously for images and save to disk
result = await scheduler_client.forward([batch])
save_file_path = os.path.join(batch.output_path, batch.output_file_name)
post_process_sample(
result.output[0],
batch.data_type,
1,
batch.save_output,
save_file_path,
)
save_file_path = await process_generation_batch(scheduler_client, batch)
await IMAGE_STORE.upsert(
request_id,
@@ -186,15 +178,7 @@ async def edits(
)
batch = _build_req_from_sampling(sampling)
result = await scheduler_client.forward([batch])
save_file_path = os.path.join(batch.output_path, batch.output_file_name)
post_process_sample(
result.output[0],
batch.data_type,
1,
batch.save_output,
save_file_path,
)
save_file_path = await process_generation_batch(scheduler_client, batch)
await IMAGE_STORE.upsert(
request_id,

38
python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py
View File

@@ -1,6 +1,7 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
import dataclasses
import os
from typing import Optional
import imageio
import numpy as np
@@ -15,6 +16,22 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__)
@dataclasses.dataclass
class SetLoraReq:
lora_nickname: str
lora_path: Optional[str] = None
@dataclasses.dataclass
class MergeLoraWeightsReq:
pass
@dataclasses.dataclass
class UnmergeLoraWeightsReq:
pass
def post_process_sample(
sample: torch.Tensor,
data_type: DataType,
@@ -75,3 +92,22 @@ async def _save_upload_to_path(upload: UploadFile, target_path: str) -> str:
with open(target_path, "wb") as f:
f.write(content)
return target_path
async def process_generation_batch(
scheduler_client,
batch,
):
result = await scheduler_client.forward([batch])
if result.output is None:
raise RuntimeError("Model generation returned no output.")
save_file_path = str(os.path.join(batch.output_path, batch.output_file_name))
post_process_sample(
result.output[0],
batch.data_type,
batch.fps,
batch.save_output,
save_file_path,
)
return save_file_path

11
python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py
View File

@@ -31,7 +31,7 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.stores import VIDEO_STORE
from sglang.multimodal_gen.runtime.entrypoints.openai.utils import (
_parse_size,
_save_upload_to_path,
post_process_sample,
process_generation_batch,
)
from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
@@ -99,14 +99,7 @@ async def _dispatch_job_async(job_id: str, batch: Req) -> None:
from sglang.multimodal_gen.runtime.scheduler_client import scheduler_client
try:
result = await scheduler_client.forward([batch])
post_process_sample(
result.output[0],
batch.data_type,
batch.fps,
batch.save_output,
os.path.join(batch.output_path, batch.output_file_name),
)
await process_generation_batch(scheduler_client, batch)
await VIDEO_STORE.update_fields(
job_id,
{"status": "completed", "progress": 100, "completed_at": int(time.time())},

5
python/sglang/multimodal_gen/runtime/layers/lora/linear.py
View File

@@ -351,11 +351,14 @@ class RowParallelLinearWithLoRA(BaseLayerWithLoRA):
return B
def get_lora_layer(
def wrap_with_lora_layer(
layer: nn.Module,
lora_rank: int | None = None,
lora_alpha: int | None = None,
) -> BaseLayerWithLoRA | None:
"""
transform the given layer to its corresponding LoRA layer
"""
supported_layer_types: dict[type[LinearBase], type[BaseLayerWithLoRA]] = {
# the order matters
# VocabParallelEmbedding: VocabParallelEmbeddingWithLoRA,

6
python/sglang/multimodal_gen/runtime/loader/component_loader.py
View File

@@ -55,7 +55,7 @@ class skip_init_modules:
cls.reset_parameters = lambda self: None # skip init
def __exit__(self, exc_type, exc_value, traceback):
# Restore originals
# restore originals
for cls, orig in self._orig_reset.items():
cls.reset_parameters = orig
@@ -199,7 +199,7 @@ class TextEncoderLoader(ComponentLoader):
def _get_weights_iterator(
self, source: "Source", to_cpu: bool
) -> Generator[tuple[str, torch.Tensor], None, None]:
"""Get an iterator for the model weights based on the load format."""
"""get an iterator for the model weights based on the load format."""
hf_folder, hf_weights_files, use_safetensors = self._prepare_weights(
source.model_or_path,
source.fall_back_to_pt,
@@ -214,7 +214,7 @@ class TextEncoderLoader(ComponentLoader):
if self.counter_before_loading_weights == 0.0:
self.counter_before_loading_weights = time.perf_counter()
# Apply the prefix.
# apply the prefix.
return ((source.prefix + name, tensor) for (name, tensor) in weights_iterator)
def _get_all_weights(

7
python/sglang/multimodal_gen/runtime/loader/utils.py
View File

@@ -32,7 +32,6 @@ def get_param_names_mapping(
Args:
mapping_dict (Dict[str, str]): Dictionary mapping regex patterns to replacement patterns
param_name (str): The parameter name to be transformed
Returns:
Callable[[str], str]: A function that maps parameter names from source to target format
@@ -44,13 +43,13 @@ def get_param_names_mapping(
match = re.match(pattern, name)
if match:
merge_index = None
total_splitted_params = None
total_split_params = None
if isinstance(replacement, tuple):
merge_index = replacement[1]
total_splitted_params = replacement[2]
total_split_params = replacement[2]
replacement = replacement[0]
name = re.sub(pattern, replacement, name)
return name, merge_index, total_splitted_params
return name, merge_index, total_split_params
# If no pattern matches, return the original name
return name, None, None

6
python/sglang/multimodal_gen/runtime/managers/gpu_worker.py
View File

@@ -121,14 +121,12 @@ class GPUWorker:
finally:
return output_batch
def set_lora_adapter(
self, lora_nickname: str, lora_path: str | None = None
) -> None:
def set_lora(self, lora_nickname: str, lora_path: str | None = None) -> None:
"""
Set the LoRA adapter for the pipeline.
"""
assert self.pipeline is not None
self.pipeline.set_lora_adapter(lora_nickname, lora_path)
self.pipeline.set_lora(lora_nickname, lora_path)
def merge_lora_weights(self) -> None:
"""

67
python/sglang/multimodal_gen/runtime/managers/scheduler.py
View File

@@ -1,11 +1,17 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
# SPDX-License-Identifier: Apache-2.0
from typing import Any
from typing import Any, List
import zmq
from sglang.multimodal_gen.runtime.entrypoints.openai.utils import (
MergeLoraWeightsReq,
SetLoraReq,
UnmergeLoraWeightsReq,
)
from sglang.multimodal_gen.runtime.managers.gpu_worker import GPUWorker
from sglang.multimodal_gen.runtime.pipelines_core import Req
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch
from sglang.multimodal_gen.runtime.server_args import (
PortArgs,
@@ -62,6 +68,30 @@ class Scheduler:
self.gpu_id = gpu_id
self._running = True
self.request_handlers = {
SetLoraReq: self._handle_set_lora,
MergeLoraWeightsReq: self._handle_merge_lora,
UnmergeLoraWeightsReq: self._handle_unmerge_lora,
Req: self._handle_generation,
}
def _handle_set_lora(self, reqs: List[Any]):
# TODO: return set status
req = reqs[0]
self.worker.set_lora(req.lora_nickname, req.lora_path)
return {"status": "ok"}
def _handle_merge_lora(self, _reqs: List[Any]):
self.worker.merge_lora_weights()
return {"status": "ok"}
def _handle_unmerge_lora(self, _reqs: List[Any]):
self.worker.unmerge_lora_weights()
return {"status": "ok"}
def _handle_generation(self, reqs: List[Req]):
return self.worker.execute_forward(reqs)
def return_result(self, output_batch: OutputBatch):
"""
replies to client, only on rank 0
@@ -69,13 +99,20 @@ class Scheduler:
if self.receiver is not None:
self.receiver.send_pyobj(output_batch)
def recv_reqs(self):
def recv_reqs(self) -> List[Any]:
"""
For non-main schedulers, reqs are broadcasted from main using broadcast_pyobj
"""
if self.receiver is not None:
recv_reqs = self.receiver.recv_pyobj()
assert isinstance(recv_reqs, list)
try:
recv_reqs = self.receiver.recv_pyobj()
except zmq.ZMQError:
# re-raise or handle appropriately to let the outer loop continue
raise
# Ensure recv_reqs is a list
if not isinstance(recv_reqs, list):
recv_reqs = [recv_reqs]
else:
recv_reqs = None
@@ -133,13 +170,27 @@ class Scheduler:
# 2: execute, make sure a reply is always sent
try:
output_batch = self.worker.execute_forward(reqs)
first_req = reqs[0] if reqs else None
handler = self.request_handlers.get(type(first_req))
if handler:
output_batch = handler(reqs)
else:
output_batch = {
"status": "error",
"message": f"Unknown request type: {type(first_req)}",
}
except Exception as e:
logger.error(
f"Error executing forward in scheduler event loop: {e}",
f"Error executing request in scheduler event loop: {e}",
exc_info=True,
)
output_batch = OutputBatch(error=str(e))
# Determine appropriate error response format
output_batch = (
OutputBatch(error=str(e))
if reqs and isinstance(reqs[0], Req)
else {"status": "error", "message": str(e)}
)
try:
self.return_result(output_batch)
@@ -171,7 +222,7 @@ class Scheduler:
return {"status": "ok", "result": result}
return {"status": "error", "error": f"Unknown method: {method}"}
def _collect_slave_results(self) -> list[dict[str, Any]]:
def _collect_slave_results(self) -> List[dict[str, Any]]:
"""Collect results from all slave worker processes."""
results = []
for pipe in self.result_pipes_from_slaves:

104
python/sglang/multimodal_gen/runtime/managers/schedulerbase.py
View File

@@ -1,104 +0,0 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
# SPDX-License-Identifier: Apache-2.0
from abc import ABC
from typing import TypeVar
import zmq
from sglang.multimodal_gen.runtime.pipelines_core import Req
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch
from sglang.multimodal_gen.runtime.server_args import ServerArgs
from sglang.multimodal_gen.utils import init_logger
logger = init_logger(__name__)
_R = TypeVar("_R")
class SchedulerBase(ABC):
"""
Abstract base class for all schedulers.
"""
def __init__(self, server_args: "ServerArgs"):
"""
Initialize the scheduler.
Args:
server_args: The inference arguments
"""
self.server_args = server_args
self.context = zmq.Context()
self.socket = self.context.socket(zmq.REQ)
self.socket.connect(self.server_args.scheduler_endpoint())
@classmethod
def get_class(cls, server_args: "ServerArgs") -> type["SchedulerBase"]:
"""
Get the scheduler class based on the server arguments.
"""
if server_args.distributed_executor_backend == "mp":
from sglang.multimodal_gen.runtime.managers.scheduler import Scheduler
# For now, always return the new Scheduler
return Scheduler
else:
raise ValueError(
f"Unsupported distributed executor backend: {server_args.distributed_executor_backend}"
)
# @abstractmethod
def start(self) -> None:
"""
Start the scheduler service.
"""
raise NotImplementedError
def execute_forward(self, batch: Req, server_args: "ServerArgs") -> OutputBatch:
"""
Execute a forward pass. This method now sends a request over ZMQ.
"""
payload = {"method": "execute_forward", "batch": batch}
self.socket.send_pyobj(payload)
output_batch = self.socket.recv_pyobj()
return output_batch
def set_lora_adapter(
self, lora_nickname: str, lora_path: str | None = None
) -> None:
"""
Set the LoRA adapter.
"""
payload = {
"method": "set_lora_adapter",
"lora_nickname": lora_nickname,
"lora_path": lora_path,
}
self.socket.send_pyobj(payload)
self.socket.recv_pyobj() # Wait for confirmation
# @abstractmethod
def unmerge_lora_weights(self) -> None:
"""
Unmerge the LoRA weights for the workers.
"""
raise NotImplementedError
# @abstractmethod
def merge_lora_weights(self) -> None:
"""
Merge the LoRA weights for the workers.
"""
raise NotImplementedError
def shutdown(self) -> None:
"""
Shutdown the scheduler.
"""
logger.info("Shutting down scheduler client.")
payload = {"method": "shutdown"}
self.socket.send_pyobj(payload)
self.socket.recv_pyobj() # Wait for shutdown confirmation
self.socket.close()
self.context.term()

4
python/sglang/multimodal_gen/runtime/pipelines/flux.py
View File

@@ -1,6 +1,8 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
# SPDX-License-Identifier: Apache-2.0
from sglang.multimodal_gen.runtime.pipelines_core import LoRAPipeline
from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
ComposedPipelineBase,
)
@@ -54,7 +56,7 @@ def prepare_mu(batch: Req, server_args: ServerArgs):
return "mu", mu
class FluxPipeline(ComposedPipelineBase):
class FluxPipeline(LoRAPipeline, ComposedPipelineBase):
pipeline_name = "FluxPipeline"
_required_config_modules = [

5
python/sglang/multimodal_gen/runtime/pipelines/qwen_image.py
View File

@@ -3,6 +3,7 @@
# SPDX-License-Identifier: Apache-2.0
from diffusers.image_processor import VaeImageProcessor
from sglang.multimodal_gen.runtime.pipelines_core import LoRAPipeline
from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
ComposedPipelineBase,
)
@@ -58,7 +59,7 @@ def prepare_mu(batch: Req, server_args: ServerArgs):
return "mu", mu
class QwenImagePipeline(ComposedPipelineBase):
class QwenImagePipeline(LoRAPipeline, ComposedPipelineBase):
pipeline_name = "QwenImagePipeline"
_required_config_modules = [
@@ -119,7 +120,7 @@ class QwenImagePipeline(ComposedPipelineBase):
)
class QwenImageEditPipeline(ComposedPipelineBase):
class QwenImageEditPipeline(LoRAPipeline, ComposedPipelineBase):
pipeline_name = "QwenImageEditPipeline"
_required_config_modules = [

18
python/sglang/multimodal_gen/runtime/pipelines_core/README.md
View File

@@ -1,18 +0,0 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
# Adding a New Custom Pipeline
Please see documentation [here](https://hao-ai-lab.github.io/sgl-diffusion/contributing/add_pipeline.html)
# PipelineStages
Basic components in a pipeline, which can be used by customed pipelines of different models.
The stages form a partial order
# PipelineExecutors
Runs the stages in a pipeline in various way. Supported ways:
1. sync
2. async

11
python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py
View File

@@ -55,6 +55,12 @@ class ComposedPipelineBase(ABC):
# the name of the pipeline it associated with, in diffusers
pipeline_name: str
def is_lora_effective(self):
return False
def is_lora_set(self):
return False
def __init__(
self,
model_path: str,
@@ -345,6 +351,11 @@ class ComposedPipelineBase(ABC):
if not self.post_init_called:
self.post_init()
if self.is_lora_set() and not self.is_lora_effective():
logger.warning(
"LoRA adapter is set, but not effective. Please make sure the LoRA weights are merged"
)
# Execute each stage
logger.info(
"Running pipeline stages: %s",

174
python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py
View File

@@ -1,6 +1,7 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
# SPDX-License-Identifier: Apache-2.0
import os
from collections import defaultdict
from collections.abc import Hashable
from typing import Any
@@ -12,8 +13,8 @@ from safetensors.torch import load_file
from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
from sglang.multimodal_gen.runtime.layers.lora.linear import (
BaseLayerWithLoRA,
get_lora_layer,
replace_submodule,
wrap_with_lora_layer,
)
from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping
from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
@@ -23,6 +24,9 @@ from sglang.multimodal_gen.runtime.server_args import ServerArgs
from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import maybe_download_lora
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
# to avoid deadlocks when forking
os.environ["TOKENIZERS_PARALLELISM"] = "false"
logger = init_logger(__name__)
@@ -31,11 +35,15 @@ class LoRAPipeline(ComposedPipelineBase):
Pipeline that supports injecting LoRA adapters into the diffusion transformer.
"""
# [lora_nickname][target_LoRA_weight_name_in_SGLang_dit] = weight
# e.g., [jinx][transformer_blocks.0.attn.to_v.lora_A]
lora_adapters: dict[str, dict[str, torch.Tensor]] = defaultdict(
dict
) # state dicts of loaded lora adapters
loaded_adapter_paths: dict[str, str] = {} # nickname -> lora_path
cur_adapter_name: str = ""
cur_adapter_path: str = ""
# [dit_layer_name] = wrapped_lora_layer
lora_layers: dict[str, BaseLayerWithLoRA] = {}
lora_layers_critic: dict[str, BaseLayerWithLoRA] = {}
server_args: ServerArgs
@@ -47,19 +55,20 @@ class LoRAPipeline(ComposedPipelineBase):
lora_rank: int | None = None
lora_alpha: int | None = None
lora_initialized: bool = False
is_lora_merged: bool = False
def __init__(self, *args, **kwargs) -> None:
super().__init__(*args, **kwargs)
self.device = get_local_torch_device()
self.exclude_lora_layers = self.modules[
"transformer"
].config.arch_config.exclude_lora_layers
self.exclude_lora_layers = (
self.server_args.pipeline_config.dit_config.arch_config.exclude_lora_layers
)
self.lora_target_modules = self.server_args.lora_target_modules
self.lora_path = self.server_args.lora_path
self.lora_nickname = self.server_args.lora_nickname
if self.lora_path is not None:
self.convert_to_lora_layers()
self.set_lora_adapter(
self.set_lora(
self.lora_nickname, self.lora_path # type: ignore
) # type: ignore
@@ -82,22 +91,20 @@ class LoRAPipeline(ComposedPipelineBase):
if not self.is_target_layer(name):
continue
excluded = False
for exclude_layer in self.exclude_lora_layers:
if exclude_layer in name:
excluded = True
break
excluded = any(
exclude_layer in name for exclude_layer in self.exclude_lora_layers
)
if excluded:
continue
layer = get_lora_layer(
lora_layer = wrap_with_lora_layer(
layer,
lora_rank=self.lora_rank,
lora_alpha=self.lora_alpha,
)
if layer is not None:
self.lora_layers[name] = layer
replace_submodule(self.modules["transformer"], name, layer)
if lora_layer is not None:
self.lora_layers[name] = lora_layer
replace_submodule(self.modules["transformer"], name, lora_layer)
converted_count += 1
logger.info("Converted %d layers to LoRA layers", converted_count)
@@ -105,7 +112,7 @@ class LoRAPipeline(ComposedPipelineBase):
for name, layer in self.modules["fake_score_transformer"].named_modules():
if not self.is_target_layer(name):
continue
layer = get_lora_layer(
layer = wrap_with_lora_layer(
layer,
lora_rank=self.lora_rank,
lora_alpha=self.lora_alpha,
@@ -121,7 +128,68 @@ class LoRAPipeline(ComposedPipelineBase):
converted_count,
)
def set_lora_adapter(
def is_lora_effective(self):
return self.is_lora_merged
def is_lora_set(self):
return self.lora_initialized and self.cur_adapter_name is not None
def load_lora_adapter(self, lora_path: str, lora_nickname: str, rank: int):
"""
Load the LoRA, and setup the lora_adapters for later weight replacement
"""
assert lora_path is not None
lora_local_path = maybe_download_lora(lora_path)
lora_state_dict = load_file(lora_local_path)
if lora_nickname in self.lora_adapters:
self.lora_adapters[lora_nickname].clear()
config = self.server_args.pipeline_config.dit_config.arch_config
param_names_mapping_fn = get_param_names_mapping(
config.param_names_mapping
or self.modules["transformer"].param_names_mapping
)
lora_param_names_mapping_fn = get_param_names_mapping(
config.lora_param_names_mapping
or self.modules["transformer"].lora_param_names_mapping
)
to_merge_params: defaultdict[Hashable, dict[Any, Any]] = defaultdict(dict)
for name, weight in lora_state_dict.items():
name = name.replace("diffusion_model.", "")
name = name.replace(".weight", "")
# misc-format -> HF-format
name, _, _ = lora_param_names_mapping_fn(name)
# HF-format (LoRA) -> SGLang-dit-format
target_name, merge_index, num_params_to_merge = param_names_mapping_fn(name)
# for (in_dim, r) @ (r, out_dim), we only merge (r, out_dim * n) where n is the number of linear layers to fuse
# see param mapping in HunyuanVideoArchConfig
if merge_index is not None and "lora_B" in name:
to_merge_params[target_name][merge_index] = weight
if len(to_merge_params[target_name]) == num_params_to_merge:
# cat at output dim according to the merge_index order
sorted_tensors = [
to_merge_params[target_name][i]
for i in range(num_params_to_merge)
]
weight = torch.cat(sorted_tensors, dim=1)
del to_merge_params[target_name]
else:
continue
print(f"{name} -> {target_name}")
if target_name in self.lora_adapters[lora_nickname]:
raise ValueError(
f"Dit target weight name {target_name} already exists in lora_adapters[{lora_nickname}]"
)
self.lora_adapters[lora_nickname][target_name] = weight.to(self.device)
self.cur_adapter_path = lora_path
self.loaded_adapter_paths[lora_nickname] = lora_path
logger.info("Rank %d: loaded LoRA adapter %s", rank, lora_path)
def set_lora(
self, lora_nickname: str, lora_path: str | None = None
): # type: ignore
"""
@@ -130,6 +198,11 @@ class LoRAPipeline(ComposedPipelineBase):
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.
"""
if self.is_lora_merged and self.cur_adapter_name != lora_nickname:
raise ValueError(
f"LoRA '{self.cur_adapter_name}' is currently merged. "
"Please call 'unmerge_lora_weights' before setting a new LoRA."
)
if lora_nickname not in self.lora_adapters and lora_path is None:
raise ValueError(
@@ -137,53 +210,26 @@ class LoRAPipeline(ComposedPipelineBase):
)
if not self.lora_initialized:
self.convert_to_lora_layers()
adapter_updated = False
rank = dist.get_rank()
if lora_path is not None and lora_path != self.cur_adapter_path:
lora_local_path = maybe_download_lora(lora_path)
lora_state_dict = load_file(lora_local_path)
# Map the hf layer names to our custom layer names
param_names_mapping_fn = get_param_names_mapping(
self.modules["transformer"].param_names_mapping
)
lora_param_names_mapping_fn = get_param_names_mapping(
self.modules["transformer"].lora_param_names_mapping
)
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
to_merge_params: defaultdict[Hashable, dict[Any, Any]] = defaultdict(dict)
for name, weight in lora_state_dict.items():
name = name.replace("diffusion_model.", "")
name = name.replace(".weight", "")
name, _, _ = lora_param_names_mapping_fn(name)
target_name, merge_index, num_params_to_merge = param_names_mapping_fn(
name
)
# for (in_dim, r) @ (r, out_dim), we only merge (r, out_dim * n) where n is the number of linear layers to fuse
# see param mapping in HunyuanVideoArchConfig
if merge_index is not None and "lora_B" in name:
to_merge_params[target_name][merge_index] = weight
if len(to_merge_params[target_name]) == num_params_to_merge:
# cat at output dim according to the merge_index order
sorted_tensors = [
to_merge_params[target_name][i]
for i in range(num_params_to_merge)
]
weight = torch.cat(sorted_tensors, dim=1)
del to_merge_params[target_name]
else:
continue
if target_name in self.lora_adapters[lora_nickname]:
raise ValueError(
f"Target name {target_name} already exists in lora_adapters[{lora_nickname}]"
)
self.lora_adapters[lora_nickname][target_name] = weight.to(self.device)
if should_load:
adapter_updated = True
self.cur_adapter_path = lora_path
logger.info("Rank %d: loaded LoRA adapter %s", rank, lora_path)
self.load_lora_adapter(lora_path, lora_nickname, rank)
if not adapter_updated and self.cur_adapter_name == lora_nickname:
if (
not adapter_updated
and self.cur_adapter_name == lora_nickname
and self.is_lora_merged
):
return
self.cur_adapter_name = lora_nickname
@@ -205,11 +251,12 @@ class LoRAPipeline(ComposedPipelineBase):
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 adapter %s does not contain the weights for layer '%s'. LoRA will not be applied to it.",
lora_path,
name,
)
layer.disable_lora = True
self.is_lora_merged = True
logger.info(
"Rank %d: LoRA adapter %s applied to %d layers",
rank,
@@ -218,9 +265,20 @@ class LoRAPipeline(ComposedPipelineBase):
)
def merge_lora_weights(self) -> None:
if self.is_lora_merged:
logger.warning("LoRA weights are already merged")
return
for name, layer in self.lora_layers.items():
layer.merge_lora_weights()
logger.info("LoRA weights merged")
self.is_lora_merged = True
def unmerge_lora_weights(self) -> None:
if not self.is_lora_merged:
logger.warning("LoRA weights are not merged.")
return
for name, layer in self.lora_layers.items():
layer.unmerge_lora_weights()
self.is_lora_merged = False

4
python/sglang/multimodal_gen/runtime/scheduler_client.py
View File

@@ -1,11 +1,11 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
import asyncio
from typing import Any
import zmq
import zmq.asyncio
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
from sglang.multimodal_gen.runtime.server_args import ServerArgs
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
@@ -47,7 +47,7 @@ class SchedulerClient:
)
# Worker will be lazily started on the first forward call to ensure a running loop exists
async def forward(self, batch: Req) -> Req:
async def forward(self, batch: Any) -> Any:
"""Enqueue a request to the backend Scheduler and await the reply."""
if self._closing:
raise RuntimeError(

13
python/sglang/multimodal_gen/runtime/server_args.py
View File

@@ -621,6 +621,19 @@ class ServerArgs:
default=ServerArgs.override_transformer_cls_name,
help="Override transformer cls name",
)
# LoRA
parser.add_argument(
"--lora-path",
type=str,
default=ServerArgs.lora_path,
help="The path to the LoRA adapter weights (can be local file path or HF hub id) to launch with",
)
parser.add_argument(
"--lora-nickname",
type=str,
default=ServerArgs.lora_nickname,
help="The nickname for the LoRA adapter to launch with",
)
# Add pipeline configuration arguments
PipelineConfig.add_cli_args(parser)

6
python/sglang/multimodal_gen/runtime/sync_scheduler_client.py
View File

@@ -1,8 +1,8 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
from typing import Any
import zmq
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
from sglang.multimodal_gen.runtime.server_args import ServerArgs
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
@@ -45,8 +45,8 @@ class SyncSchedulerClient:
f"SyncSchedulerClient connected to backend scheduler at {scheduler_endpoint}"
)
def forward(self, batch: Req) -> Req:
"""Sends a batch to the scheduler and waits for the response."""
def forward(self, batch: Any) -> Any:
"""Sends a batch or request to the scheduler and waits for the response."""
try:
self.scheduler_socket.send_pyobj(batch)
output_batch = self.scheduler_socket.recv_pyobj()

4
python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py
View File

@@ -214,6 +214,10 @@ def maybe_download_lora(
"""
local_path = maybe_download_model(model_name_or_path, local_dir, download)
# return directly if local_path is a file
if os.path.isfile(local_path):
return local_path
weight_name = _best_guess_weight_name(
model_name_or_path, file_extension=".safetensors"
)

28
python/sglang/multimodal_gen/runtime/utils/logging_utils.py
View File

@@ -381,6 +381,17 @@ def configure_logger(server_args, prefix: str = ""):
set_uvicorn_logging_configs()
def suppress_loggers(loggers_to_suppress: list[str]):
original_levels = {}
for logger_name in loggers_to_suppress:
logger = logging.getLogger(logger_name)
original_levels[logger_name] = logger.level
logger.setLevel(logging.WARNING)
return original_levels
@contextmanager
def suppress_other_loggers(not_suppress_on_main_rank: bool = False):
"""
@@ -400,21 +411,8 @@ def suppress_other_loggers(not_suppress_on_main_rank: bool = False):
if get_is_main_process() == 0:
should_suppress = False
loggers_to_suppress = ["urllib3"]
original_levels = {}
if should_suppress:
for logger_name in loggers_to_suppress:
logger = logging.getLogger(logger_name)
original_levels[logger_name] = logger.level
logger.setLevel(logging.WARNING)
# Suppress verbose logging from imageio, which is triggered when saving images.
logging.getLogger("imageio").setLevel(logging.WARNING)
logging.getLogger("imageio_ffmpeg").setLevel(logging.WARNING)
# Suppress Pillow plugin import logs when app log level is DEBUG
logging.getLogger("PIL").setLevel(logging.WARNING)
logging.getLogger("PIL.Image").setLevel(logging.WARNING)
loggers_to_suppress = ["urllib3", "imageio", "imageio_ffmpeg", "PIL", "PIL_Image"]
original_levels = suppress_loggers(loggers_to_suppress)
try:
yield

7
python/sglang/multimodal_gen/test/server/test_server_common.py
View File

@@ -73,7 +73,7 @@ def diffusion_server(case: DiffusionTestCase) -> ServerContext:
prompt=sampling_params.prompt or "A colorful raccoon icon",
output_size=output_size,
)
warmup.run_text_warmups(case.server_args.warmup_text)
warmup.run_text_warmups(server_args.warmup_text)
if (
case.server_args.warmup_edit > 0
@@ -81,15 +81,14 @@ def diffusion_server(case: DiffusionTestCase) -> ServerContext:
and sampling_params.image_path
):
# Handle URL or local path
image_path = sampling_params.image_path
if is_image_url(sampling_params.image_path):
image_path = download_image_from_url(str(sampling_params.image_path))
else:
image_path = Path(sampling_params.image_path)
warmup.run_edit_warmups(
count=case.server_args.warmup_edit,
edit_prompt=case.sampling_params.prompt,
count=server_args.warmup_edit,
edit_prompt=sampling_params.prompt,
image_path=image_path,
)
except Exception as exc: