[diffusion] feat: support LoRA (#13859)

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,

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

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

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"

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,

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)

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(
+    # LoRA
-        self, lora_nickname: str, lora_path: str | None = None
+    def _send_lora_request(self, req: Any, success_msg: str, failure_msg: str):
-    ) -> None:
+        response = sync_scheduler_client.forward(req)
-        # self.scheduler.set_lora_adapter(lora_nickname, lora_path)
+        if isinstance(response, dict) and response.get("status") == "ok":
-        pass  # Removed as per edit hint
+            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:
-        """
+        req = UnmergeLoraWeightsReq()
-        Use unmerged weights for inference to produce outputs that align with
+        self._send_lora_request(
-        validation outputs generated during training.
+            req,
-        """
+            "Successfully unmerged LoRA weights",
-        # self.scheduler.unmerge_lora_weights()
+            "Failed to unmerge LoRA weights",
-        pass  # Removed as per edit hint
+        )
+        self._is_lora_merged = False
 
     def merge_lora_weights(self) -> None:
-        # self.scheduler.merge_lora_weights()
+        req = MergeLoraWeightsReq()
-        pass  # Removed as per edit hint
+        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
+        # sends the shutdown command to the server
-        # self.scheduler.shutdown()
-
         if self.local_scheduler_process:
             logger.info("Waiting for local worker processes to terminate...")
             for process in self.local_scheduler_process:

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)
 

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"
+    )

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 = await process_generation_batch(scheduler_client, 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,
-    )
 
     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 = await process_generation_batch(scheduler_client, 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,
-    )
 
     await IMAGE_STORE.upsert(
         request_id,

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

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])
+        await process_generation_batch(scheduler_client, 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 VIDEO_STORE.update_fields(
             job_id,
             {"status": "completed", "progress": 100, "completed_at": int(time.time())},

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,

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(

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

python/sglang/multimodal_gen/runtime/managers/gpu_worker.py

@@ -121,14 +121,12 @@ class GPUWorker:
         finally:
             return output_batch
 
-    def set_lora_adapter(
+    def set_lora(self, lora_nickname: str, lora_path: str | None = None) -> None:
-        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:
         """

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:
+            try:
                 recv_reqs = self.receiver.recv_pyobj()
-            assert isinstance(recv_reqs, list)
+            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:

python/sglang/multimodal_gen/runtime/managers/schedulerbase.py

@@ -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()

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 = [

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 = [

python/sglang/multimodal_gen/runtime/pipelines_core/README.md

@@ -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

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",

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[
+        self.exclude_lora_layers = (
-            "transformer"
+            self.server_args.pipeline_config.dit_config.arch_config.exclude_lora_layers
-        ].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
+            excluded = any(
-            for exclude_layer in self.exclude_lora_layers:
+                exclude_layer in name for exclude_layer in self.exclude_lora_layers
-                if exclude_layer in name:
+            )
-                    excluded = True
-                    break
             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:
+            if lora_layer is not None:
-                self.lora_layers[name] = layer
+                self.lora_layers[name] = lora_layer
-                replace_submodule(self.modules["transformer"], name, 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,44 +128,42 @@ class LoRAPipeline(ComposedPipelineBase):
                 converted_count,
             )
 
-    def set_lora_adapter(
+    def is_lora_effective(self):
-        self, lora_nickname: str, lora_path: str | None = None
+        return self.is_lora_merged
-    ):  # type: ignore
-        """
-        Load a LoRA adapter into the pipeline and merge it into the transformer.
-        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.
-        """
 
-        if lora_nickname not in self.lora_adapters and lora_path is None:
+    def is_lora_set(self):
-            raise ValueError(
+        return self.lora_initialized and self.cur_adapter_name is not None
-                f"Adapter {lora_nickname} not found in the pipeline. Please provide lora_path to load it."
+
-            )
+    def load_lora_adapter(self, lora_path: str, lora_nickname: str, rank: int):
-        if not self.lora_initialized:
+        """
-            self.convert_to_lora_layers()
+        Load the LoRA, and setup the lora_adapters for later weight replacement
-        adapter_updated = False
+        """
-        rank = dist.get_rank()
+        assert lora_path is not None
-        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
+        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(
-                self.modules["transformer"].param_names_mapping
+            config.param_names_mapping
+            or self.modules["transformer"].param_names_mapping
         )
         lora_param_names_mapping_fn = get_param_names_mapping(
-                self.modules["transformer"].lora_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)
-                target_name, merge_index, num_params_to_merge = param_names_mapping_fn(
+            # HF-format (LoRA) -> SGLang-dit-format
-                    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:
@@ -174,16 +179,57 @@ class LoRAPipeline(ComposedPipelineBase):
                 else:
                     continue
 
+            print(f"{name} -> {target_name}")
             if target_name in self.lora_adapters[lora_nickname]:
                 raise ValueError(
-                        f"Target name {target_name} already exists in lora_adapters[{lora_nickname}]"
+                    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)
-            adapter_updated = True
         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)
 
-        if not adapter_updated and self.cur_adapter_name == lora_nickname:
+    def set_lora(
+        self, lora_nickname: str, lora_path: str | None = None
+    ):  # type: ignore
+        """
+        Load a LoRA adapter into the pipeline and merge it into the transformer.
+        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.
+        """
+        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(
+                f"Adapter {lora_nickname} not found in the pipeline. Please provide lora_path to load it."
+            )
+        if not self.lora_initialized:
+            self.convert_to_lora_layers()
+
+        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
+
+        if should_load:
+            adapter_updated = True
+            self.load_lora_adapter(lora_path, lora_nickname, rank)
+
+        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

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(

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)
 

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:
+    def forward(self, batch: Any) -> Any:
-        """Sends a batch to the scheduler and waits for the response."""
+        """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()

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"
     )

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"]
+    loggers_to_suppress = ["urllib3", "imageio", "imageio_ffmpeg", "PIL", "PIL_Image"]
-    original_levels = {}
+    original_levels = suppress_loggers(loggers_to_suppress)
-
-    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)
 
     try:
         yield

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,
+                count=server_args.warmup_edit,
-                edit_prompt=case.sampling_params.prompt,
+                edit_prompt=sampling_params.prompt,
                 image_path=image_path,
             )
     except Exception as exc: