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