[diffusion] feat: support distilled vae generic (#14195)

Co-authored-by: BBuf <1182563586@qq.com>
Co-authored-by: Mick <mickjagger19@icloud.com>

python/sglang/multimodal_gen/runtime/server_args.py

@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # Inspired by SGLang: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/server_args.py
 """The arguments of sglang-diffusion Inference."""
+
 import argparse
 import dataclasses
 import inspect
@@ -252,6 +253,9 @@ class ServerArgs:
     # (Wenxuan) prefer to keep it here instead of in pipeline config to not make it complicated.
     lora_path: str | None = None
     lora_nickname: str = "default"  # for swapping adapters in the pipeline
+
+    # VAE parameters
+    vae_path: str | None = None  # Custom VAE path (e.g., for distilled autoencoder)
     # can restrict layers to adapt, e.g. ["q_proj"]
     # Will adapt only q, k, v, o by default.
     lora_target_modules: list[str] | None = None
@@ -374,6 +378,12 @@ class ServerArgs:
             type=str,
             help="Directory containing StepVideo model",
         )
+        parser.add_argument(
+            "--vae-path",
+            type=str,
+            default=ServerArgs.vae_path,
+            help="Custom path to VAE model (e.g., for distilled autoencoder). If not specified, VAE will be loaded from the main model path.",
+        )
 
         # attention
         parser.add_argument(
@@ -825,17 +835,15 @@ class ServerArgs:
         if self.ulysses_degree is None:
             self.ulysses_degree = 1
             logger.info(
-                f"Ulysses degree not set, " f"using default value {self.ulysses_degree}"
+                f"Ulysses degree not set, using default value {self.ulysses_degree}"
             )
 
         if self.ring_degree is None:
             self.ring_degree = 1
-            logger.info(
-                f"Ring degree not set, " f"using default value {self.ring_degree}"
-            )
+            logger.info(f"Ring degree not set, using default value {self.ring_degree}")
 
         if self.ring_degree > 1:
-            if self.attention_backend != None and self.attention_backend != "fa":
+            if self.attention_backend is not None and self.attention_backend != "fa":
                 raise ValueError(
                     "Ring Attention is only supported for flash attention backend for now"
                 )