From dfd7ab9682594d143937236b42f23c5303d81a4d Mon Sep 17 00:00:00 2001 From: Mick Date: Wed, 26 Nov 2025 00:21:33 +0800 Subject: [PATCH] [diffusion] feat: support LoRA (#13859) --- .../configs/models/dits/base.py | 9 +- .../configs/models/dits/flux.py | 6 + .../configs/models/dits/hunyuanvideo.py | 1 - .../configs/models/dits/qwenimage.py | 7 +- .../configs/models/dits/wanvideo.py | 1 - .../configs/models/encoders/qwen_image.py | 1 + .../entrypoints/diffusion_generator.py | 111 +++++++++-- .../runtime/entrypoints/http_server.py | 3 + .../runtime/entrypoints/openai/common_api.py | 62 +++++++ .../runtime/entrypoints/openai/image_api.py | 22 +-- .../runtime/entrypoints/openai/utils.py | 38 +++- .../runtime/entrypoints/openai/video_api.py | 11 +- .../runtime/layers/lora/linear.py | 5 +- .../runtime/loader/component_loader.py | 6 +- .../multimodal_gen/runtime/loader/utils.py | 7 +- .../runtime/managers/gpu_worker.py | 6 +- .../runtime/managers/scheduler.py | 67 ++++++- .../runtime/managers/schedulerbase.py | 104 ----------- .../multimodal_gen/runtime/pipelines/flux.py | 4 +- .../runtime/pipelines/qwen_image.py | 5 +- .../runtime/pipelines_core/README.md | 18 -- .../pipelines_core/composed_pipeline_base.py | 11 ++ .../runtime/pipelines_core/lora_pipeline.py | 174 ++++++++++++------ .../runtime/scheduler_client.py | 4 +- .../multimodal_gen/runtime/server_args.py | 13 ++ .../runtime/sync_scheduler_client.py | 6 +- .../runtime/utils/hf_diffusers_utils.py | 4 + .../runtime/utils/logging_utils.py | 28 ++- .../test/server/test_server_common.py | 7 +- 29 files changed, 463 insertions(+), 278 deletions(-) create mode 100644 python/sglang/multimodal_gen/runtime/entrypoints/openai/common_api.py delete mode 100644 python/sglang/multimodal_gen/runtime/managers/schedulerbase.py delete mode 100644 python/sglang/multimodal_gen/runtime/pipelines_core/README.md diff --git a/python/sglang/multimodal_gen/configs/models/dits/base.py b/python/sglang/multimodal_gen/configs/models/dits/base.py index 22da409a1..77fecfb4d 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/base.py +++ b/python/sglang/multimodal_gen/configs/models/dits/base.py @@ -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, diff --git a/python/sglang/multimodal_gen/configs/models/dits/flux.py b/python/sglang/multimodal_gen/configs/models/dits/flux.py index 285acecc0..f50c9b1e6 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/flux.py +++ b/python/sglang/multimodal_gen/configs/models/dits/flux.py @@ -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 diff --git a/python/sglang/multimodal_gen/configs/models/dits/hunyuanvideo.py b/python/sglang/multimodal_gen/configs/models/dits/hunyuanvideo.py index 23a6c715b..1cae921ff 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/hunyuanvideo.py +++ b/python/sglang/multimodal_gen/configs/models/dits/hunyuanvideo.py @@ -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 diff --git a/python/sglang/multimodal_gen/configs/models/dits/qwenimage.py b/python/sglang/multimodal_gen/configs/models/dits/qwenimage.py index 4cf46a089..f9003264b 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/qwenimage.py +++ b/python/sglang/multimodal_gen/configs/models/dits/qwenimage.py @@ -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" diff --git a/python/sglang/multimodal_gen/configs/models/dits/wanvideo.py b/python/sglang/multimodal_gen/configs/models/dits/wanvideo.py index 68e6801d7..3a9939c4d 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/wanvideo.py +++ b/python/sglang/multimodal_gen/configs/models/dits/wanvideo.py @@ -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, diff --git a/python/sglang/multimodal_gen/configs/models/encoders/qwen_image.py b/python/sglang/multimodal_gen/configs/models/encoders/qwen_image.py index 0a5f245f4..02f57ca65 100644 --- a/python/sglang/multimodal_gen/configs/models/encoders/qwen_image.py +++ b/python/sglang/multimodal_gen/configs/models/encoders/qwen_image.py @@ -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) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py index 96862c36e..ab408e95e 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py @@ -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: diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py b/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py index 25d5e8fc4..489a08e20 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py @@ -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) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/common_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/common_api.py new file mode 100644 index 000000000..2dcbeb77c --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/common_api.py @@ -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" + ) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py index 1ba388023..17ce57cbd 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py @@ -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, diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py index 42bda15e0..89c08c222 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py @@ -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 diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py index 734dce04d..4304a636f 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py @@ -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())}, diff --git a/python/sglang/multimodal_gen/runtime/layers/lora/linear.py b/python/sglang/multimodal_gen/runtime/layers/lora/linear.py index fbe6a4495..f5c7d7ea4 100644 --- a/python/sglang/multimodal_gen/runtime/layers/lora/linear.py +++ b/python/sglang/multimodal_gen/runtime/layers/lora/linear.py @@ -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, diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loader.py index 9cf0c1b92..0975f8741 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loader.py @@ -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( diff --git a/python/sglang/multimodal_gen/runtime/loader/utils.py b/python/sglang/multimodal_gen/runtime/loader/utils.py index fe3c2de69..9f375b9a7 100644 --- a/python/sglang/multimodal_gen/runtime/loader/utils.py +++ b/python/sglang/multimodal_gen/runtime/loader/utils.py @@ -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 diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index aeca02b0e..7fa2014f5 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -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: """ diff --git a/python/sglang/multimodal_gen/runtime/managers/scheduler.py b/python/sglang/multimodal_gen/runtime/managers/scheduler.py index 8b2e33f58..35da11dcf 100644 --- a/python/sglang/multimodal_gen/runtime/managers/scheduler.py +++ b/python/sglang/multimodal_gen/runtime/managers/scheduler.py @@ -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: diff --git a/python/sglang/multimodal_gen/runtime/managers/schedulerbase.py b/python/sglang/multimodal_gen/runtime/managers/schedulerbase.py deleted file mode 100644 index a2da3cc75..000000000 --- a/python/sglang/multimodal_gen/runtime/managers/schedulerbase.py +++ /dev/null @@ -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() diff --git a/python/sglang/multimodal_gen/runtime/pipelines/flux.py b/python/sglang/multimodal_gen/runtime/pipelines/flux.py index 7ede4b388..d738c8ba6 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/flux.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/flux.py @@ -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 = [ diff --git a/python/sglang/multimodal_gen/runtime/pipelines/qwen_image.py b/python/sglang/multimodal_gen/runtime/pipelines/qwen_image.py index e74ebb8ae..458c00987 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/qwen_image.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/qwen_image.py @@ -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 = [ diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/README.md b/python/sglang/multimodal_gen/runtime/pipelines_core/README.md deleted file mode 100644 index 14a9531a4..000000000 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/README.md +++ /dev/null @@ -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 diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py b/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py index b9d22fc46..175c14fe8 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py @@ -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", diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py index 3af42dee4..33f8ad6c7 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py @@ -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 diff --git a/python/sglang/multimodal_gen/runtime/scheduler_client.py b/python/sglang/multimodal_gen/runtime/scheduler_client.py index 4b70de3fb..c0c0a00a1 100644 --- a/python/sglang/multimodal_gen/runtime/scheduler_client.py +++ b/python/sglang/multimodal_gen/runtime/scheduler_client.py @@ -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( diff --git a/python/sglang/multimodal_gen/runtime/server_args.py b/python/sglang/multimodal_gen/runtime/server_args.py index e23dba802..9d1e274eb 100644 --- a/python/sglang/multimodal_gen/runtime/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args.py @@ -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) diff --git a/python/sglang/multimodal_gen/runtime/sync_scheduler_client.py b/python/sglang/multimodal_gen/runtime/sync_scheduler_client.py index 8e89728f6..3d32576fb 100644 --- a/python/sglang/multimodal_gen/runtime/sync_scheduler_client.py +++ b/python/sglang/multimodal_gen/runtime/sync_scheduler_client.py @@ -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() diff --git a/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py b/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py index 56877d159..f9979151b 100644 --- a/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py +++ b/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py @@ -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" ) diff --git a/python/sglang/multimodal_gen/runtime/utils/logging_utils.py b/python/sglang/multimodal_gen/runtime/utils/logging_utils.py index fc7c41c61..b6f54c3b6 100644 --- a/python/sglang/multimodal_gen/runtime/utils/logging_utils.py +++ b/python/sglang/multimodal_gen/runtime/utils/logging_utils.py @@ -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 diff --git a/python/sglang/multimodal_gen/test/server/test_server_common.py b/python/sglang/multimodal_gen/test/server/test_server_common.py index f3f1e8ef2..18b56f5d7 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_common.py +++ b/python/sglang/multimodal_gen/test/server/test_server_common.py @@ -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: