[Dependency] Upgrade Cache-DiT to 1.5.0

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>

.buildkite/cuda/test-nightly.yml

@@ -345,7 +345,7 @@ steps:
         commands:
           - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
           - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
-          - export CACHE_DIT_VERSION=1.3.0
+          - export CACHE_DIT_VERSION=1.5.0
           - |
             set +e
             pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
@@ -361,7 +361,7 @@ steps:
         commands:
           - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
           - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
-          - export CACHE_DIT_VERSION=1.3.0
+          - export CACHE_DIT_VERSION=1.5.0
           - |
             set +e
             pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json
@@ -377,7 +377,7 @@ steps:
         commands:
           - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
           - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
-          - export CACHE_DIT_VERSION=1.3.0
+          - export CACHE_DIT_VERSION=1.5.0
           - |
             set +e
             pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
@@ -393,7 +393,7 @@ steps:
         commands:
           - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
           - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
-          - export CACHE_DIT_VERSION=1.3.0
+          - export CACHE_DIT_VERSION=1.5.0
           - |
             set +e
             pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json
@@ -409,7 +409,7 @@ steps:
         commands:
           - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
           - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
-          - export CACHE_DIT_VERSION=1.3.0
+          - export CACHE_DIT_VERSION=1.5.0
           - |
             set +e
             pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_bagel_vllm_omni.json

.buildkite/npu/test-npu-nightly.yml

@@ -115,7 +115,7 @@ steps:
           DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
         commands:
           - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
-          - export CACHE_DIT_VERSION=1.3.0
+          - export CACHE_DIT_VERSION=1.5.0
           - |
             set +e
             pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m npu --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json

docs/user_guide/diffusion/lora.md

@@ -1,14 +1,18 @@
 # LoRA (Low-Rank Adaptation) Guide
 
-LoRA (Low-Rank Adaptation) enables fine-tuning diffusion models by adding trainable low-rank matrices to existing model weights. vLLM-Omni currently supports PEFT-style LoRA adapters, allowing you to customize model behavior without modifying the base model weights.
+LoRA (Low-Rank Adaptation) enables fine-tuning diffusion models by adding trainable low-rank matrices to existing model weights. vLLM-Omni supports two LoRA backends: **PEFT** for PEFT-style adapters, and **Distill** usually for few-steps inference. **PEFT** backend allows you to customize model behavior without modifying the base model weights. **Distill** backend fuses LoRA weights into the base model at initialization.
 
 ## Overview
 
-LoRA adapters are lightweight, model-specific fine-tuning weights that can be dynamically loaded and applied to diffusion models. vLLM-Omni uses a unified LoRA handling mechanism similar to vLLM with LRU cache management.
+LoRA adapters are lightweight, model-specific fine-tuning weights that can be applied to diffusion models in two ways:
+
+- **PEFT backend** (`--lora-backend peft`, default): Loads a PEFT-format adapter folder via `DiffusionLoRAManager`. Adapters are cached (LRU) and activated per request via `LoRARequest`. It uses a unified LoRA handling mechanisms similar to vLLM with LRU cache management.
+- **Distill backend** (`--lora-backend distill`): Calls `pipeline.load_lora_weights` once at initialization to fuse one or more concrete checkpoint files directly into the base weights. Typically used for distilled few-step LoRAs (e.g. Lightning, LightX2V).
 
 ## LoRA Adapter Format
 
-LoRA adapters must be in **PEFT (Parameter-Efficient Fine-Tuning)** format. A typical LoRA adapter directory structure:
+### PEFT (Parameter-Efficient Fine-Tuning) format (default)
+A typical PEFT-format LoRA adapter directory structure:
 
 ```
 lora_adapter/
@@ -25,19 +29,20 @@ The `adapter_config.json` file contains metadata about the LoRA adapter, includi
 
 ### Offline Inference
 
-#### Pre-loaded LoRA
+#### PEFT backend: pre-loaded LoRA
 
-Load a LoRA adapter at initialization. This adapter is pre-loaded into the cache and can be activated by requests:
+Load a PEFT-format LoRA adapter at initialization. The adapter is pre-loaded into the cache and can be activated per request:
 
 ```python
 from vllm_omni import Omni
 from vllm_omni.lora.request import LoRARequest
 
-lora_path="/path/to/lora_adapter"
+lora_path = "/path/to/lora_adapter"
 
 omni = Omni(
     model="stabilityai/stable-diffusion-3.5-medium",
-    lora_path=lora_path
+    lora_path=lora_path,
+    lora_backend="peft",  # default, can be omitted
 )
 
 lora_request = LoRARequest(
@@ -53,6 +58,98 @@ outputs = omni.generate(
 )
 ```
 
+#### Distill backend: fuse distilled LoRA at init
+
+For distilled few-step LoRAs, pass `lora_backend="distill"` together with one or more concrete `.safetensors` files. The weights are fused into the base model once at init; subsequent `generate()` calls do not need a `LoRARequest`.
+
+##### Supported pipelines
+
+For Qwen-Image and Wan pipelines, the distill backend calls `pipeline.load_lora_weights(...)` during worker initialization.
+
+| Pipeline | Supported distilled LoRA repo | Notes |
+|----------|--------------------------------|-------|
+| `QwenImagePipeline` | `lightx2v/Qwen-Image-2512-Lightning` | Used with Qwen-Image-2512 Lightning-style few-step inference. |
+| `Wan22Pipeline` | `lightx2v/Wan2.1-Distill-Loras`, `lightx2v/Wan2.2-Distill-Loras` | Wan2.1 uses one LoRA file. For dual-transformer Wan2.2 MoE, pass high-noise then low-noise LoRA files. |
+| `Wan22I2VPipeline` | `lightx2v/Wan2.2-Distill-Loras` | For dual-transformer Wan2.2 MoE, pass high-noise then low-noise LoRA files. |
+
+Other diffusion pipelines are not currently listed as supporting distilled LoRA. Use the PEFT backend for request-time adapters, or bake converted weights into a local Diffusers directory before serving.
+
+Single-file example (Qwen-Image-Lightning):
+
+```python
+from vllm_omni import Omni
+
+omni = Omni(
+    model="Qwen/Qwen-Image-2512",
+    lora_path="/path/to/Qwen-Image-2512-Lightning.safetensors",
+    lora_backend="distill",
+)
+
+outputs = omni.generate(prompt="A piece of cheesecake")
+```
+
+Multi-file example (Wan2.2 MoE, high + low noise):
+
+```python
+from vllm_omni import Omni
+
+omni = Omni(
+    model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+    lora_path=[
+        "/path/to/wan2.2_high_noise_lora.safetensors",   # -> transformer
+        "/path/to/wan2.2_low_noise_lora.safetensors",    # -> transformer_2
+    ],
+    lora_backend="distill",
+)
+```
+
+The CLI examples under `examples/offline_inference/` accept the same flags, e.g.:
+
+```bash
+python examples/offline_inference/text_to_video/text_to_video.py \
+  --model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
+  --lora-backend distill \
+  --lora-path /path/to/high.safetensors /path/to/low.safetensors \
+  --prompt "A cat playing with yarn"
+```
+
+##### Online serving
+
+Distilled LoRAs can also be fused when an online diffusion server starts. The
+adapter remains active for the lifetime of that server, so requests do not pass
+a per-request `lora` field.
+
+```bash
+vllm serve Qwen/Qwen-Image-2512 \
+  --omni \
+  --port 8091 \
+  --lora-backend distill \
+  --lora-path /path/to/Qwen-Image-2512-Lightning-4steps.safetensors
+```
+
+Then send a normal generation request with the sampling settings expected by
+the distilled checkpoint:
+
+```bash
+curl -s http://localhost:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "messages": [
+      {"role": "user", "content": "A Chinese female college student, around 20 years old, with a very short haircut that conveys a gentle, artistic vibe. Her hair naturally falls to partially cover her cheeks, projecting a tomboyish yet charming demeanor. She has cool-toned fair skin and delicate features, with a slightly shy yet subtly confident expression—her mouth crooked in a playful, youthful smirk. She wears an off-shoulder top, revealing one shoulder, with a well-proportioned figure. The image is framed as a close-up selfie: she dominates the foreground, while the background clearly shows her dormitory—a neatly made bed with white linens on the top bunk, a tidy study desk with organized stationery, and wooden cabinets and drawers. The photo is captured on a smartphone under soft, even ambient lighting, with natural tones, high clarity, and a bright, lively atmosphere full of youthful, everyday energy."}
+    ],
+    "extra_body": {
+      "height": 1024,
+      "width": 1024,
+      "num_inference_steps": 4,
+      "true_cfg_scale": 1.0,
+      "seed": 42
+    }
+  }' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2 | base64 -d > qwen_image_distill_lora_output.png
+```
+
+For Wan2.2 MoE serving, pass the high-noise checkpoint first and the low-noise
+checkpoint second to `--lora-path`.
+
 !!! note "Server-side Path Requirement"
     The LoRA adapter path (`local_path`) must be readable on the **server** machine. If your client and server are on different machines, ensure the LoRA adapter is accessible via a shared mount or copied to the server.
 
@@ -141,6 +238,7 @@ Notes:
 
 - This route avoids runtime LoRA loading changes in vLLM-Omni when you choose to bake converted weights into a local Diffusers directory.
 - Output quality and speed depend on the replacement checkpoints and sampling params you choose.
+- If you only need to fuse distilled LoRAs into a Wan2.2 checkpoint at load time (without the full LightX2V convert + assemble pipeline), you can instead pass them directly via `--lora-backend distill --lora-path <high>.safetensors <low>.safetensors`. See the [Distill backend](#distill-backend-fuse-distilled-lora-at-init) section above.
 
 
 ## See Also

docs/user_guide/examples/offline_inference/text_to_image.md

@@ -88,8 +88,9 @@ python text_to_image.py \
 | `--vae-use-tiling` | flag | off | Enable VAE tiling for memory optimization |
 | `--cfg-parallel-size` | int | `1` | Set to `2` to enable CFG Parallel |
 | `--enable-cpu-offload` | flag | off | Enable CPU offloading for diffusion models |
-| `--lora-path` | str | — | Path to PEFT LoRA adapter folder |
+| `--lora-path` | str | — | Path to PEFT LoRA adapter folder or checkpoint file |
 | `--lora-scale` | float | `1.0` | Scale factor for LoRA weights |
+| `--lora-backend` | str |`"peft"`| LoRA backend for loading LoRA adapters. Default: peft. Available options: peft, distill |
 
 **NextStep-1.1 specific arguments:**
 

docs/user_guide/examples/offline_inference/text_to_video.md

@@ -125,6 +125,9 @@ python text_to_video.py \
 - `--audio-sample-rate`: fallback audio sample rate when the pipeline returns audio.
 - `--quantization`: quantization method (`fp8` for FP8, `gguf` for GGUF).
 - `--flow-shift`: scheduler flow_shift parameter.
+- `--lora-path`: path to PEFT LoRA adapter folder or checkpoint file.
+- `--lora-scale`: scale factor for LoRA weights.
+- `--lora-backend`: backend for loading LoRA adapters. Default: peft. Available options: peft, distill.
 
 ### Wan2.2-specific
 

examples/offline_inference/image_to_video/image_to_video.py

@@ -60,10 +60,10 @@ from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.lora.request import LoRARequest
+from vllm_omni.lora.utils import stable_lora_int_id
 from vllm_omni.model_extras import (
-    build_image_to_video_prompt as build_model_image_to_video_prompt,
-)
-from vllm_omni.model_extras import (
+    build_image_to_video_prompt,
     get_extra_body_params,
     get_model_class_name,
 )
@@ -85,22 +85,6 @@ def parse_json_object(value: str, flag_name: str = "argument") -> dict[str, Any]
 parse_profiler_config = functools.partial(parse_json_object, flag_name="--profiler-config")
 
 
-def build_image_to_video_prompt(
-    prompt: str,
-    negative_prompt: str | None,
-    media_inputs: dict[str, Any],
-) -> dict[str, Any]:
-    """Build the canonical request envelope for the shared I2V example."""
-    result: dict[str, Any] = {
-        "prompt": prompt,
-        "modalities": ["video"],
-        "multi_modal_data": media_inputs,
-    }
-    if negative_prompt is not None:
-        result["negative_prompt"] = negative_prompt
-    return result
-
-
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(
         description=(
@@ -312,6 +296,30 @@ def parse_args() -> argparse.Namespace:
         default=None,
         help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
     )
+    parser.add_argument(
+        "--lora-path",
+        type=str,
+        nargs="+",
+        default=None,
+        help="Path to a LoRA adapter folder (PEFT format) or concrete LoRA checkpoint files. "
+        "For Wan2.2 MoE models, pass the high-noise checkpoint first and the low-noise checkpoint second.",
+    )
+    parser.add_argument(
+        "--lora-scale",
+        type=float,
+        default=1.0,
+        help="Scale factor for PEFT LoRA weights (default: 1.0).",
+    )
+    parser.add_argument(
+        "--lora-backend",
+        type=str,
+        default="peft",
+        choices=["peft", "distill"],
+        help=(
+            "LoRA backend. 'peft' loads a PEFT-format adapter for request-time activation; "
+            "'distill' fuses one or more concrete LoRA checkpoint files at initialization."
+        ),
+    )
     parser.add_argument(
         "--extra-body",
         type=functools.partial(parse_json_object, flag_name="--extra-body"),
@@ -505,6 +513,12 @@ def main():
         omni_kwargs["flow_shift"] = flow_shift
     if args.quantization is not None:
         omni_kwargs["quantization"] = args.quantization
+    if args.lora_path is not None:
+        lora_path = args.lora_path
+        if len(lora_path) == 1:
+            lora_path = lora_path[0]
+        omni_kwargs["lora_path"] = lora_path
+        omni_kwargs["lora_backend"] = args.lora_backend
     # Cosmos3 loads its (gated) guardrail models at build time, so the guardrails
     # gate is an engine-level config (offline analog of the server's --no-guardrails).
     if args.extra_body and "guardrails" in args.extra_body:
@@ -535,22 +549,36 @@ def main():
     print(f"  Video size: {width}x{height}")
     print(f"{'=' * 60}\n")
 
+    lora_request = None
+    if args.lora_path and args.lora_backend == "peft":
+        if len(args.lora_path) != 1:
+            raise ValueError("Only one LoRA path is expected for PEFT backend.")
+
+        lora_path = args.lora_path[0]
+        lora_request = LoRARequest(
+            lora_name=Path(lora_path).stem,
+            lora_int_id=stable_lora_int_id(lora_path),
+            lora_path=lora_path,
+        )
+
     negative_prompt = args.negative_prompt
     if negative_prompt is None and not is_ltx2:
         # Preserve the historical empty-prompt behavior for non-LTX examples.
         negative_prompt = ""
     prompt_dict = build_image_to_video_prompt(
+        model_class_name=model_class_name,
         prompt=args.prompt,
         negative_prompt=negative_prompt,
         media_inputs=media_inputs,
-    )
-    prompt_dict = build_model_image_to_video_prompt(
-        model_class_name=model_class_name,
-        prompt=prompt_dict,
         height=height,
         width=width,
         num_frames=num_frames,
     )
+    sampling_extra_args = {"sample_solver": args.sample_solver}
+    if lora_request is not None:
+        sampling_extra_args["lora_request"] = lora_request
+        sampling_extra_args["lora_scale"] = args.lora_scale
+
     sampling_params = OmniDiffusionSamplingParams(
         height=height,
         width=width,
@@ -561,7 +589,7 @@ def main():
         num_inference_steps=num_inference_steps,
         num_frames=num_frames,
         frame_rate=frame_rate,
-        extra_args={"sample_solver": args.sample_solver},
+        extra_args=sampling_extra_args,
     )
     if flow_shift is not None:
         sampling_params.extra_args["flow_shift"] = flow_shift

examples/offline_inference/text_to_image/README.md

@@ -57,9 +57,8 @@ python examples/offline_inference/text_to_image/text_to_image.py \
   --output lingbot_t2i.png
 ```
 
-The shared T2I entry point selects the `image` output modality. The LingBot
-pipeline validates that T2I requests produce one frame and do not include an
-input image.
+The LingBot prompt builder selects the `image` output modality and requests one
+frame. Passing an input image is not supported in LingBot T2I mode.
 
 ## Quick Start
 
@@ -110,8 +109,9 @@ python text_to_image.py \
 | `--ring-degree` | int | `1` | Ring sequence parallel degree for hybrid Ulysses + Ring inference |
 | `--ulysses-mode` | str | `"strict"` | Ulysses SP mode: `"strict"` or `"advanced_uaa"` |
 | `--enable-cpu-offload` | flag | off | Enable CPU offloading for diffusion models |
-| `--lora-path` | str | — | Path to PEFT LoRA adapter folder |
+| `--lora-path` | str | — | Path to PEFT LoRA adapter folder or checkpoint file |
 | `--lora-scale` | float | `1.0` | Scale factor for LoRA weights |
+| `--lora-backend` | str |`"peft"`| LoRA backend for loading LoRA adapters. Default: peft. Available options: peft, distill |
 | `--use-system-prompt` | str | `None` | System prompt preset: `en_unified`, `en_vanilla`, `en_recaption`, `en_think_recaption`, `dynamic`, `None`, or custom text. Recommended: `en_unified`. Only for HunyuanImage-3.0.|
 | `--system-prompt` | str | `None` | Custom system prompt text. Only used when `--use-system-prompt` is set to `custom`. Only for HunyuanImage-3.0.|
 | `--auxiliary-text-encoder` | str | `None` | Supplementary auxiliary text encoder parameters model name or path (especially for Hidream-l1-full). |

examples/offline_inference/text_to_image/text_to_image.py

@@ -21,9 +21,7 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.lora.request import LoRARequest
 from vllm_omni.lora.utils import stable_lora_int_id
 from vllm_omni.model_extras import (
-    build_text_to_image_prompt as build_model_text_to_image_prompt,
-)
-from vllm_omni.model_extras import (
+    build_text_to_image_prompt,
     get_extra_body_params,
     get_model_class_name,
     should_init_extra_args_for_non_diffusion_stages,
@@ -58,17 +56,6 @@ def parse_json_object(value: str, flag_name: str = "argument") -> dict[str, Any]
 parse_profiler_config = functools.partial(parse_json_object, flag_name="--profiler-config")
 
 
-def build_text_to_image_prompt(prompt: str, negative_prompt: str | None) -> dict[str, Any]:
-    """Build the canonical request envelope for the shared T2I example."""
-    result: dict[str, Any] = {
-        "prompt": prompt,
-        "modalities": ["image"],
-    }
-    if negative_prompt is not None:
-        result["negative_prompt"] = negative_prompt
-    return result
-
-
 def _normalize_images_for_save(images: list[Any]) -> list[Any]:
     """Convert NumPy diffusion outputs to PIL images before saving."""
     normalized = []
@@ -273,6 +260,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--lora-path",
         type=str,
+        nargs="+",
         default=None,
         help="Path to LoRA adapter folder (PEFT format). Loaded at initialization and used for generation.",
     )
@@ -282,6 +270,15 @@ def parse_args() -> argparse.Namespace:
         default=1.0,
         help="Scale factor for LoRA weights (default: 1.0).",
     )
+    parser.add_argument(
+        "--lora-backend",
+        type=str,
+        default="peft",
+        choices=["peft", "distill"],
+        help="LoRA backend for loading LoRA adapters. Default: peft"
+        "'peft' loads a PEFT-format adapter folder, used e.g. for RL"
+        "'distill' fuses one or more concrete LoRA checkpoint files, used e.g. for distilled few-step LoRAs",
+    )
     parser.add_argument(
         "--vae-patch-parallel-size",
         type=int,
@@ -417,8 +414,12 @@ def main():
     # Prepare LoRA kwargs for Omni initialization
     lora_args: dict[str, Any] = {}
     if args.lora_path:
-        lora_args["lora_path"] = args.lora_path
-        print(f"Using LoRA from: {args.lora_path}")
+        lora_path = args.lora_path
+        if len(lora_path) == 1:
+            lora_path = lora_path[0]
+        lora_args["lora_path"] = lora_path
+        lora_args["lora_backend"] = args.lora_backend
+        print(f"Using LoRA from: {lora_path} with backend: {args.lora_backend}")
 
     # Build quantization kwargs: use quantization_config dict when
     # ignored_layers is specified so the list flows through OmniDiffusionConfig
@@ -507,23 +508,24 @@ def main():
 
     # Build LoRA request when --lora-path is set
     lora_request = None
-    if args.lora_path:
-        lora_request_id = stable_lora_int_id(args.lora_path)
+    if args.lora_path and args.lora_backend == "peft":
+        if len(args.lora_path) != 1:
+            raise ValueError("Only one LoRA path is expected for PEFT backend.")
+
+        lora_path = args.lora_path[0]
+        lora_request_id = stable_lora_int_id(lora_path)
         lora_request = LoRARequest(
-            lora_name=Path(args.lora_path).stem,
+            lora_name=Path(lora_path).stem,
             lora_int_id=lora_request_id,
-            lora_path=args.lora_path,
+            lora_path=lora_path,
         )
 
     generation_start = time.perf_counter()
 
     prompt_dict = build_text_to_image_prompt(
+        model_class_name=model_class_name,
         prompt=args.prompt,
         negative_prompt=args.negative_prompt,
-    )
-    prompt_dict = build_model_text_to_image_prompt(
-        model_class_name=model_class_name,
-        prompt=prompt_dict,
         height=args.height,
         width=args.width,
     )

examples/offline_inference/text_to_video/text_to_video.md

@@ -218,6 +218,9 @@ python text_to_video.py \
 - `--audio-sample-rate`: fallback audio sample rate when the pipeline returns audio.
 - `--quantization`: quantization method (such as `fp8` for FP8).
 - `--flow-shift`: scheduler flow_shift parameter.
+- `--lora-path`: path to PEFT LoRA adapter folder or checkpoint file.
+- `--lora-scale`: scale factor for LoRA weights.
+- `--lora-backend`: backend for loading LoRA adapters. Default: peft. Available options: peft, distill.
 - `--extra-body`: JSON object of model-specific generation params, filtered against the model's declared `extra_body_params` (see [`vllm_omni/model_extras`](../../../vllm_omni/model_extras)). Used by Cosmos3 (see above).
 
 ### Wan2.2-specific

examples/offline_inference/text_to_video/text_to_video.py

@@ -14,6 +14,8 @@ from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.lora.request import LoRARequest
+from vllm_omni.lora.utils import stable_lora_int_id
 from vllm_omni.model_extras import get_extra_body_params, get_model_class_name
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.platforms import current_omni_platform
@@ -334,6 +336,29 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable expert parallelism for MoE layers.",
     )
+    parser.add_argument(
+        "--lora-path",
+        type=str,
+        nargs="+",
+        default=None,
+        help="Path to LoRA adapter folder (PEFT format) or concrete LoRA checkpoint files. Loaded at initialization and used for generation."
+        "Note: for Wan2.2 MoE models, two checkpoints positionally mapped to high-noise and low-noise modules.",
+    )
+    parser.add_argument(
+        "--lora-scale",
+        type=float,
+        default=1.0,
+        help="Scale factor for LoRA weights (default: 1.0).",
+    )
+    parser.add_argument(
+        "--lora-backend",
+        type=str,
+        default="peft",
+        choices=["peft", "distill"],
+        help="LoRA backend for loading LoRA adapters. Default: peft"
+        "'peft' loads a PEFT-format adapter folder, used e.g. for RL"
+        "'distill' fuses one or more concrete LoRA checkpoint files, used e.g. for distilled few-step LoRAs",
+    )
     parser.add_argument(
         "--use-hsdp",
         action="store_true",
@@ -440,6 +465,12 @@ def main():
         omni_kwargs["cache_backend"] = args.cache_backend
         omni_kwargs["cache_config"] = cache_config
         omni_kwargs["enable_cache_dit_summary"] = args.enable_cache_dit_summary
+    if args.lora_path is not None:
+        lora_path = args.lora_path
+        if len(lora_path) == 1:
+            lora_path = lora_path[0]
+        omni_kwargs["lora_path"] = lora_path
+        omni_kwargs["lora_backend"] = args.lora_backend
 
     # Cosmos3 loads its (gated) guardrail models at build time, so the guardrails
     # gate is an engine-level config (offline analog of the server's --no-guardrails).
@@ -469,6 +500,19 @@ def main():
     print(f"  Video size: {args.width}x{args.height}")
     print(f"{'=' * 60}\n")
 
+    lora_request = None
+    if args.lora_path and args.lora_backend == "peft":
+        if len(args.lora_path) != 1:
+            raise ValueError("Only one LoRA path is expected for PEFT backend.")
+
+        lora_path = args.lora_path[0]
+        lora_request_id = stable_lora_int_id(lora_path)
+        lora_request = LoRARequest(
+            lora_name=Path(lora_path).stem,
+            lora_int_id=lora_request_id,
+            lora_path=lora_path,
+        )
+
     prompt_dict = {"prompt": args.prompt}
     if args.negative_prompt is not None:
         prompt_dict["negative_prompt"] = args.negative_prompt
@@ -476,6 +520,11 @@ def main():
         # Preserve the historical empty-prompt behavior for non-LTX examples.
         prompt_dict["negative_prompt"] = ""
 
+    extra_args = {}
+    if lora_request:
+        extra_args["lora_request"] = lora_request
+        extra_args["lora_scale"] = args.lora_scale
+
     sampling_kwargs = dict(
         height=args.height,
         width=args.width,
@@ -483,6 +532,7 @@ def main():
         guidance_scale=args.guidance_scale,
         num_inference_steps=args.num_inference_steps,
         num_frames=args.num_frames,
+        extra_args=extra_args,
     )
     if args.guidance_scale_high is not None:
         sampling_kwargs["guidance_scale_2"] = args.guidance_scale_high

examples/online_serving/image_to_video/README.md

@@ -1,9 +1,6 @@
 # Image-To-Video
 
-This example demonstrates online image-to-video generation with vLLM-Omni.
-The startup script defaults to Wan2.2, but `MODEL` can select any supported
-image-to-video model. The existing curl helper remains a concrete Wan2.2
-request through the standard Videos API.
+This example demonstrates how to deploy the Wan2.2 image-to-video model for online video generation using vLLM-Omni.
 
 ## Start Server
 
@@ -23,7 +20,9 @@ bash run_server.sh
 
 The script allows overriding:
 - `MODEL` (default: `Wan-AI/Wan2.2-I2V-A14B-Diffusers`)
-- `PORT` (default: `8099`)
+- `PORT` (default: `8091`)
+- `BOUNDARY_RATIO` (default: `0.875`)
+- `FLOW_SHIFT` (default: `12.0`)
 - `CACHE_BACKEND` (default: `none`)
 - `ENABLE_CACHE_DIT_SUMMARY` (default: `0`)
 

examples/online_serving/image_to_video/run_server.sh

@@ -1,12 +1,12 @@
 #!/bin/bash
-# Shared image-to-video server start script
+# Wan2.2 image-to-video server start script
 
 MODEL="${MODEL:-Wan-AI/Wan2.2-I2V-A14B-Diffusers}"
 PORT="${PORT:-8099}"
 CACHE_BACKEND="${CACHE_BACKEND:-none}"
 ENABLE_CACHE_DIT_SUMMARY="${ENABLE_CACHE_DIT_SUMMARY:-0}"
 
-echo "Starting image-to-video server..."
+echo "Starting Wan2.2 I2V server..."
 echo "Model: $MODEL"
 echo "Port: $PORT"
 echo "Cache backend: $CACHE_BACKEND"

examples/online_serving/text_to_image/README.md

@@ -1,9 +1,6 @@
 # Text-To-Image
 
-This example demonstrates online text-to-image generation with vLLM-Omni. The
-startup script and Python clients default to Qwen-Image, but `MODEL` can select
-any supported text-to-image model. The existing curl helper remains a concrete
-Qwen-Image request through the standard Images API.
+This example demonstrates how to deploy Qwen-Image model for online image generation service using vLLM-Omni.
 
 ## Start Server
 
@@ -118,7 +115,23 @@ python gradio_demo.py
 
 ## LoRA
 
-This example supports Peft-compatible LoRA (Low-Rank Adaptation) adapters for diffusion models. The LoRA adapter path must be readable on the **server** machine (usually a local path or a mounted directory).
+This example supports both startup-time distilled LoRAs and request-time
+PEFT-compatible adapters. The LoRA path must be readable on the **server**
+machine (usually a local path or a mounted directory).
+
+### Start Server with a Distilled LoRA
+
+```bash
+vllm serve Qwen/Qwen-Image-2512 \
+  --omni \
+  --port 8091 \
+  --lora-backend distill \
+  --lora-path /path/to/Qwen-Image-2512-Lightning-4steps.safetensors
+```
+
+The distilled LoRA is fused once during server initialization and applies to
+every request. Use the checkpoint's few-step sampling settings, for example
+`num_inference_steps=4` and `true_cfg_scale=1.0`.
 
 ### Using Python Client with LoRA
 
@@ -229,8 +242,8 @@ count, use `size` and `n` rather than `height`, `width`, or
 | `height`                 | int   | None    | Image height in pixels         |
 | `width`                  | int   | None    | Image width in pixels          |
 | `size`                   | str   | None    | Image size (e.g., "1024x1024") |
-| `num_inference_steps`    | int   | model   | Number of denoising steps      |
-| `true_cfg_scale`         | float | model   | True CFG scale when supported  |
+| `num_inference_steps`    | int   | 50      | Number of denoising steps      |
+| `true_cfg_scale`         | float | 4.0     | Qwen-Image CFG scale           |
 | `seed`                   | int   | None    | Random seed (reproducible)     |
 | `negative_prompt`        | str   | None    | Negative prompt                |
 | `num_outputs_per_prompt` | int   | 1       | Number of images to generate   |

examples/online_serving/text_to_image/gradio_demo.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 """
-Shared text-to-image Gradio demo for online serving.
+Qwen-Image Gradio Demo for online serving.
 
 Usage:
     python gradio_demo.py [--server http://localhost:8091] [--port 7860]
@@ -77,9 +77,9 @@ def generate_image(
 def create_demo(server_url: str):
     """Create Gradio demo interface."""
 
-    with gr.Blocks(title="Text-to-Image Demo") as demo:
-        gr.Markdown("# Text-to-Image Online Generation")
-        gr.Markdown("Generate images using the model configured on the server")
+    with gr.Blocks(title="Qwen-Image Demo") as demo:
+        gr.Markdown("# Qwen-Image Online Generation")
+        gr.Markdown("Generate images using Qwen-Image model")
 
         with gr.Row():
             with gr.Column(scale=1):
@@ -115,6 +115,7 @@ def create_demo(server_url: str):
                         label="Inference Steps",
                         minimum=10,
                         maximum=100,
+                        # Default steps aligned with run_curl_text_to_image.sh to 100
                         value=100,
                         step=5,
                     )
@@ -172,7 +173,7 @@ def create_demo(server_url: str):
 
 
 def main():
-    parser = argparse.ArgumentParser(description="Shared text-to-image Gradio demo")
+    parser = argparse.ArgumentParser(description="Qwen-Image Gradio Demo")
     parser.add_argument("--server", default="http://localhost:8091", help="Server URL")
     parser.add_argument("--port", type=int, default=7860, help="Gradio port")
     parser.add_argument("--share", action="store_true", help="Create public link")

examples/online_serving/text_to_image/openai_chat_client.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 """
-Shared OpenAI-compatible text-to-image generation client.
+Qwen-Image OpenAI-compatible image generation client.
 
 Usage:
     python openai_chat_client.py --prompt "A beautiful landscape" --output output.png
@@ -39,7 +39,7 @@ def generate_image(
         height: Image height in pixels
         width: Image width in pixels
         steps: Number of diffusion steps
-        true_cfg_scale: True CFG scale for models that support it
+        true_cfg_scale: Qwen-Image CFG scale
         seed: Random seed
         negative_prompt: Negative prompt
         num_outputs_per_prompt: Number of images to generate
@@ -70,7 +70,7 @@ def generate_image(
         payload["num_inference_steps"] = steps
     if true_cfg_scale is not None:
         payload["true_cfg_scale"] = true_cfg_scale
-    if negative_prompt is not None:
+    if negative_prompt:
         payload["negative_prompt"] = negative_prompt
     if seed is not None:
         payload["seed"] = seed
@@ -115,15 +115,15 @@ def generate_image(
 
 
 def main():
-    parser = argparse.ArgumentParser(description="Shared text-to-image client")
+    parser = argparse.ArgumentParser(description="Qwen-Image chat client")
     parser.add_argument("--prompt", "-p", default="a cup of coffee on the table", help="Text prompt")
-    parser.add_argument("--output", "-o", default="text_to_image_output.png", help="Output file")
+    parser.add_argument("--output", "-o", default="qwen_image_output.png", help="Output file")
     parser.add_argument("--server", "-s", default="http://localhost:8091", help="Server URL")
-    parser.add_argument("--height", type=int, help="Image height; defaults to the model setting")
-    parser.add_argument("--width", type=int, help="Image width; defaults to the model setting")
-    parser.add_argument("--steps", type=int, help="Inference steps; defaults to the model setting")
-    parser.add_argument("--cfg-scale", type=float, help="True CFG scale when supported")
-    parser.add_argument("--seed", type=int, help="Random seed")
+    parser.add_argument("--height", type=int, default=1024, help="Image height")
+    parser.add_argument("--width", type=int, default=1024, help="Image width")
+    parser.add_argument("--steps", type=int, default=50, help="Inference steps")
+    parser.add_argument("--cfg-scale", type=float, default=4.0, help="True CFG scale")
+    parser.add_argument("--seed", type=int, default=0, help="Random seed")
     parser.add_argument("--negative", help="Negative prompt")
 
     parser.add_argument("--lora-path", default=None, help="Server-local LoRA adapter folder (PEFT format)")

examples/online_serving/text_to_image/run_server.sh

@@ -1,10 +1,10 @@
 #!/bin/bash
-# Shared text-to-image online serving startup script
+# Qwen-Image online serving startup script
 
 MODEL="${MODEL:-Qwen/Qwen-Image}"
 PORT="${PORT:-8091}"
 
-echo "Starting text-to-image server..."
+echo "Starting Qwen-Image server..."
 echo "Model: $MODEL"
 echo "Port: $PORT"
 

recipes/Tencent/HunyuanImage-3.0-Instruct.md

@@ -82,7 +82,7 @@ example and the DiT sections below.
 - Optional environment variables:
 
 ```bash
-export CACHE_DIT_VERSION=1.3.0
+export CACHE_DIT_VERSION=1.5.0
 ```
 
 HunyuanImage-3.0 sets the diffusion attention backend to `TORCH_SDPA`

requirements/common.txt

@@ -16,7 +16,7 @@ diffusers==0.38.0
 safetensors>=0.8.0
 accelerate==1.12.0
 soundfile>=0.13.1
-cache-dit==1.3.0
+cache-dit==1.5.0
 tqdm>=4.66.0
 torchsde>=0.2.6
 openai-whisper>=20250625

tests/diffusion/lora/test_loader.py

@@ -0,0 +1,438 @@
+from collections import defaultdict
+
+import pytest
+import torch
+import torch.nn as nn
+from diffusers.loaders.lora_conversion_utils import (
+    _convert_non_diffusers_qwen_lora_to_diffusers,
+)
+from pytest_mock import MockerFixture
+from torch.testing import assert_close
+
+from vllm_omni.diffusion.lora.loader import (
+    LoraLoaderMixin,
+    QwenImageLoraLoaderMixin,
+    WanLoraLoaderMixin,
+    _prepare_lora_delta,
+    _remap_state_dict_keys,
+)
+
+HEAD_DIM = 24
+RANK = 2
+NUM_LAYERS = 10
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class DummyTransformerBlock(nn.Module):
+    def __init__(
+        self,
+        d_model: int,
+        bias: bool = False,
+        stacked_params: bool = False,
+    ):
+        super().__init__()
+        self.attn = nn.Module()
+        if stacked_params:
+            self.attn.to_qkv = nn.Linear(d_model * 3, d_model, bias=bias)
+        else:
+            self.attn.to_q = nn.Linear(d_model, d_model, bias=bias)
+            self.attn.to_k = nn.Linear(d_model, d_model, bias=bias)
+            self.attn.to_v = nn.Linear(d_model, d_model, bias=bias)
+        self.attn.to_out = nn.Linear(d_model, d_model, bias=bias)
+
+
+class DummyTransformer(nn.Module):
+    def __init__(self, num_layers: int, d_model: int, bias: bool = False, stacked_params: bool = False):
+        super().__init__()
+
+        self.blocks = nn.ModuleList(
+            [DummyTransformerBlock(d_model, bias=bias, stacked_params=stacked_params) for _ in range(num_layers)]
+        )
+        if stacked_params:
+            self.stacked_params_mapping = [
+                (".attn.to_qkv", ".attn.to_q", "q"),
+                (".attn.to_qkv", ".attn.to_k", "k"),
+                (".attn.to_qkv", ".attn.to_v", "v"),
+            ]
+
+
+class DummyPipeline(nn.Module, LoraLoaderMixin):
+    def __init__(self, num_layers: int, d_model: int, bias: bool = False, stacked_params: bool = False):
+        super().__init__()
+        self.transformer = DummyTransformer(num_layers, d_model, bias=bias, stacked_params=stacked_params)
+
+
+# ======================================================
+# helper functions for test
+# ======================================================
+def _filter_parameters(named_params, mode: str = "all"):
+    if mode == "all":
+        return named_params
+
+    filtered = []
+    for name, param in named_params:
+        if name.endswith(mode):
+            filtered.append((name, param))
+
+    return filtered
+
+
+def assert_params_equal(module, original_params, mode: str = "all"):
+    for name, param in _filter_parameters(module.named_parameters(), mode):
+        assert_close(param, original_params[name])
+
+
+def assert_params_not_equal(module, original_params, mode: str = "all"):
+    for name, param in _filter_parameters(module.named_parameters(), mode):
+        assert not torch.allclose(param, original_params[name])
+
+
+def make_dummy_lora_state_dict(
+    base_key,
+    weight_a: torch.Tensor | None = None,
+    weight_b: torch.Tensor | None = None,
+    weight_bias: torch.Tensor | None = None,
+    bias: bool = False,
+    lora_a_suffix: str = "lora_A.weight",
+    lora_b_suffix: str = "lora_B.weight",
+    lora_bias_suffix: str = "bias",
+):
+    """
+    Create a dummy LoRA state dict.
+
+    Returns:
+        tuple: (state_dict, expected_delta)
+    """
+    key_a = f"{base_key}.{lora_a_suffix}"
+    key_b = f"{base_key}.{lora_b_suffix}"
+    key_bias = f"{base_key}.{lora_bias_suffix}" if bias else None
+
+    if weight_a is None:
+        weight_a = torch.randn(RANK, HEAD_DIM)
+    if weight_b is None:
+        weight_b = torch.randn(HEAD_DIM, RANK)
+
+    state_dict = {}
+    expected_delta = None
+
+    if bias:
+        state_dict[key_a] = weight_a
+        state_dict[key_b] = weight_b
+        weight_bias = torch.randn(HEAD_DIM)
+        state_dict[key_bias] = weight_bias
+        expected_delta = weight_bias
+    else:
+        state_dict[key_a] = weight_a
+        state_dict[key_b] = weight_b
+        expected_delta = weight_b @ weight_a
+
+    return state_dict, expected_delta
+
+
+def make_lora_state_dict_for_module(
+    module,
+    lora_rank=RANK,
+    bias=False,
+    lora_a_suffix: str = "lora_A.weight",
+    lora_b_suffix: str = "lora_B.weight",
+    lora_bias_suffix: str = "bias",
+    prefix: str = "transformer",
+    remap_proj_out: bool = True,
+):
+    param_to_weight_names = defaultdict(list)
+    if hasattr(module, "stacked_params_mapping"):
+        for param_name, weight_name, _ in module.stacked_params_mapping:
+            param_to_weight_names[param_name].append(weight_name)
+
+    state_dict = {}
+    for name, param in module.named_parameters(prefix=prefix):
+        if name.endswith(".weight"):
+            base_key = name[: -len(".weight")]
+            if remap_proj_out and base_key.endswith(".to_out"):
+                base_key = base_key.replace(".to_out", ".to_out.0")
+
+            d_in, d_out = param.shape
+            is_stacked_param = False
+            for param_name, weight_names in param_to_weight_names.items():
+                if param_name not in name:
+                    continue
+                is_stacked_param = True
+                target_in = d_in // len(weight_names)
+                for weight_name in weight_names:
+                    lora_key = base_key.replace(param_name, weight_name)
+                    weight_a = torch.randn(lora_rank, d_out)
+                    weight_b = torch.randn(target_in, lora_rank)
+                    state_dict[f"{lora_key}.{lora_a_suffix}"] = weight_a
+                    state_dict[f"{lora_key}.{lora_b_suffix}"] = weight_b
+
+            if is_stacked_param:
+                continue
+
+            weight_a = torch.randn(lora_rank, d_out)
+            weight_b = torch.randn(d_in, lora_rank)
+            key_a = f"{base_key}.{lora_a_suffix}"
+            key_b = f"{base_key}.{lora_b_suffix}"
+            state_dict[key_a] = weight_a
+            state_dict[key_b] = weight_b
+        elif name.endswith(".bias") and bias:
+            base_key = name[: -len(".bias")]
+            if remap_proj_out and base_key.endswith(".to_out"):
+                base_key = base_key.replace(".to_out", ".to_out.0")
+
+            d_bias = param.shape[0]
+            is_stacked_param = False
+            for param_name, weight_names in param_to_weight_names.items():
+                if param_name not in name:
+                    continue
+                is_stacked_param = True
+                target_bias = d_bias // len(weight_names)
+                for weight_name in weight_names:
+                    lora_key = base_key.replace(param_name, weight_name)
+                    weight_bias = torch.randn(target_bias)
+                    state_dict[f"{lora_key}.{lora_bias_suffix}"] = weight_bias
+
+            if is_stacked_param:
+                continue
+
+            key_bias = f"{base_key}.{lora_bias_suffix}"
+            state_dict[key_bias] = torch.randn_like(param)
+
+    return state_dict
+
+
+# ======================================================
+# Test loader helper functions
+# ======================================================
+def test_prepare_lora_delta_with_stacked_params():
+    base_key = "transformer.attn.to_qkv"
+
+    sd0, delta0 = make_dummy_lora_state_dict("transformer.attn.to_q")
+    sd1, delta1 = make_dummy_lora_state_dict("transformer.attn.to_k")
+    sd2, delta2 = make_dummy_lora_state_dict("transformer.attn.to_v")
+
+    lora_state_dict = {}
+    lora_state_dict.update(sd0)
+    lora_state_dict.update(sd1)
+    lora_state_dict.update(sd2)
+    expected_delta = torch.concat([delta0, delta1, delta2])
+
+    stacked_mapping = {".attn.to_qkv": [".attn.to_q", ".attn.to_k", ".attn.to_v"]}
+
+    actual_delta, used_keys = _prepare_lora_delta(
+        lora_state_dict,
+        base_key,
+        stacked_mapping,
+    )
+
+    # validate the delta
+    assert actual_delta is not None
+    assert_close(actual_delta, expected_delta)
+
+    # validate the used keys
+    key_q_a = "transformer.attn.to_q.lora_A.weight"
+    key_q_b = "transformer.attn.to_q.lora_B.weight"
+    key_k_a = "transformer.attn.to_k.lora_A.weight"
+    key_k_b = "transformer.attn.to_k.lora_B.weight"
+    key_v_a = "transformer.attn.to_v.lora_A.weight"
+    key_v_b = "transformer.attn.to_v.lora_B.weight"
+
+    for key in [key_q_a, key_q_b, key_k_a, key_k_b, key_v_a, key_v_b]:
+        assert key in used_keys
+
+
+def test_remap_state_dict_keys_matched_two():
+    lora_state_dict, _ = make_dummy_lora_state_dict(
+        "transformer.attn.to_out.0", bias=True, lora_bias_suffix="lora_B.bias"
+    )
+
+    lora_state_dict = _remap_state_dict_keys(
+        lora_state_dict,
+        [(".to_out.0.", ".to_out."), (".lora_B.bias", ".bias")],
+    )
+
+    assert len(lora_state_dict) == 3
+    assert "transformer.attn.to_out.bias" in lora_state_dict
+
+
+# ======================================================
+# Test LoraLoaderMixin
+# ======================================================
+class TestLoraLoaderMixin:
+    @pytest.mark.parametrize(
+        "bias, stacked_params",
+        [
+            (False, False),
+            (True, False),
+            (False, True),
+            (True, True),
+        ],
+        ids=[
+            "basic",
+            "bias",
+            "stacked_params",
+            "bias_stacked_params",
+        ],
+    )
+    def test_load_lora_into_module(self, bias: bool, stacked_params: bool):
+        pipeline = DummyPipeline(NUM_LAYERS, HEAD_DIM, bias=bias, stacked_params=stacked_params)
+        original_params = {name: param.clone() for name, param in pipeline.transformer.named_parameters()}
+
+        lora_state_dict = make_lora_state_dict_for_module(pipeline.transformer, bias=bias)
+        assert len(lora_state_dict) > 0
+
+        # remap the keys to vllm-omni format
+        lora_state_dict = _remap_state_dict_keys(lora_state_dict, [(".to_out.0.", ".to_out.")])
+
+        used_keys = pipeline.load_lora_into_module(lora_state_dict, pipeline.transformer)
+        # validate the parameters are updated
+        assert_params_not_equal(pipeline.transformer, original_params)
+
+        # validate all keys are used
+        missing_keys = set(lora_state_dict.keys()) - used_keys
+        assert len(missing_keys) == 0
+
+    def test_unload_module_lora(self):
+        pipeline = DummyPipeline(NUM_LAYERS, HEAD_DIM)
+        original_params = {name: param.clone() for name, param in pipeline.transformer.named_parameters()}
+
+        lora_state_dict = make_lora_state_dict_for_module(pipeline.transformer)
+        # validate the state dict is valid
+        assert len(lora_state_dict) > 0
+
+        # remap the keys to vllm-omni format
+        lora_state_dict = _remap_state_dict_keys(lora_state_dict, [(".to_out.0.", ".to_out.")])
+
+        used_keys = pipeline.load_lora_into_module(lora_state_dict, pipeline.transformer)
+
+        # validate the weights are updated
+        assert_params_not_equal(pipeline.transformer, original_params)
+
+        # validate all keys are used
+        missing_keys = set(lora_state_dict.keys()) - used_keys
+        assert len(missing_keys) == 0
+
+        unload_used_keys = pipeline.unload_module_lora(lora_state_dict, pipeline.transformer)
+        missing_keys = set(lora_state_dict.keys()) - unload_used_keys
+        assert len(missing_keys) == 0
+        assert len(unload_used_keys - used_keys) == 0
+
+        # validate the weights are restored
+        assert_params_equal(pipeline.transformer, original_params)
+
+
+# ======================================================
+# Test QwenImageLoraLoaderMixin
+# ======================================================
+class DummyQwenImagePipeline(nn.Module, QwenImageLoraLoaderMixin):
+    def __init__(self, num_layers: int, d_model: int, bias: bool = False):
+        super().__init__()
+        self.transformer = DummyTransformer(num_layers, d_model, bias=bias, stacked_params=True)
+
+
+class TestQwenImageLoraLoaderMixin:
+    def test_load_lora_weights(self, mocker: MockerFixture):
+        pipeline = DummyQwenImagePipeline(NUM_LAYERS, HEAD_DIM)
+        original_params = {name: param.clone() for name, param in pipeline.transformer.named_parameters()}
+
+        lora_state_dict = make_lora_state_dict_for_module(pipeline.transformer)
+        assert len(lora_state_dict) > 0
+
+        mocker.patch(
+            "vllm_omni.diffusion.lora.loader.get_converter_by_pipeline",
+            return_value=_convert_non_diffusers_qwen_lora_to_diffusers,
+        )
+        pipeline.load_lora_weights(lora_state_dict, "adapter0")
+
+        # validate the weights are updated after lora loaded
+        assert_params_not_equal(pipeline.transformer, original_params)
+
+        # validate lora_loaded map is updated after lora loaded
+        assert "adapter0" in pipeline.lora_loaded
+
+    def test_unload_lora_weights(self, mocker: MockerFixture):
+        pipeline = DummyQwenImagePipeline(NUM_LAYERS, HEAD_DIM)
+        original_params = {name: param.clone() for name, param in pipeline.transformer.named_parameters()}
+
+        lora_state_dict = make_lora_state_dict_for_module(pipeline.transformer)
+        assert len(lora_state_dict) > 0
+
+        mocker.patch(
+            "vllm_omni.diffusion.lora.loader.get_converter_by_pipeline",
+            return_value=_convert_non_diffusers_qwen_lora_to_diffusers,
+        )
+        pipeline.load_lora_weights(lora_state_dict, "adapter0")
+        # validate the weights are updated after lora loaded
+        assert_params_not_equal(pipeline.transformer, original_params)
+
+        # validate lora_loaded map is updated after lora loaded
+        assert "adapter0" in pipeline.lora_loaded
+
+        pipeline.unload_lora_weights("adapter0")
+        # validate the weights are restored after lora unloaded
+        assert_params_equal(pipeline.transformer, original_params)
+
+        # validate lora_loaded map is updated
+        assert "adapter0" not in pipeline.lora_loaded
+
+    def test_load_diffusers_lora_applies_non_unit_alpha(self, mocker: MockerFixture):
+        pipeline = DummyQwenImagePipeline(1, HEAD_DIM)
+        to_out = pipeline.transformer.blocks[0].attn.to_out
+        original_weight = to_out.weight.detach().clone()
+        base_key = "blocks.0.attn.to_out.0"
+        lora_a = torch.randn(RANK, HEAD_DIM)
+        lora_b = torch.randn(HEAD_DIM, RANK)
+        alpha = torch.tensor(6.0)
+        state_dict = {
+            f"{base_key}.lora_A.weight": lora_a,
+            f"{base_key}.lora_B.weight": lora_b,
+            f"{base_key}.alpha": alpha,
+        }
+        mocker.patch(
+            "vllm_omni.diffusion.lora.loader.get_converter_by_pipeline",
+            return_value=_convert_non_diffusers_qwen_lora_to_diffusers,
+        )
+
+        pipeline.load_lora_weights(state_dict, "adapter0")
+
+        expected_weight = original_weight + (lora_b @ lora_a) * (alpha.item() / RANK)
+        assert_close(to_out.weight, expected_weight)
+
+        pipeline.unload_lora_weights("adapter0")
+        assert_close(to_out.weight, original_weight)
+
+
+# ======================================================
+# Test WanLoraLoaderMixin
+# ======================================================
+class DummyWanPipeline(nn.Module, WanLoraLoaderMixin):
+    def __init__(self, num_layers: int, d_model: int):
+        super().__init__()
+        self.has_transformer_2 = True
+        self.transformer = DummyTransformer(num_layers, d_model, stacked_params=True)
+        self.transformer_2 = DummyTransformer(num_layers, d_model, stacked_params=True)
+
+
+def test_wan_load_lora_weights_targets_both_transformers(mocker: MockerFixture):
+    pipeline = DummyWanPipeline(NUM_LAYERS, HEAD_DIM)
+    original_high = {name: param.clone() for name, param in pipeline.transformer.named_parameters()}
+    original_low = {name: param.clone() for name, param in pipeline.transformer_2.named_parameters()}
+    high_state_dict = make_lora_state_dict_for_module(pipeline.transformer)
+    low_state_dict = make_lora_state_dict_for_module(pipeline.transformer_2)
+
+    mocker.patch(
+        "vllm_omni.diffusion.lora.loader._load_lora_state_dict",
+        side_effect=[high_state_dict, low_state_dict],
+    )
+
+    pipeline.load_lora_weights(["high.safetensors", "low.safetensors"], "adapter0")
+
+    assert_params_not_equal(pipeline.transformer, original_high)
+    assert_params_not_equal(pipeline.transformer_2, original_low)
+    assert set(pipeline.lora_loaded["adapter0"]) == {"transformer", "transformer_2"}
+
+    pipeline.unload_lora_weights("adapter0")
+
+    assert_params_equal(pipeline.transformer, original_high)
+    assert_params_equal(pipeline.transformer_2, original_low)
+    assert "adapter0" not in pipeline.lora_loaded

tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py

@@ -19,38 +19,11 @@ DEPLOY_CONFIG = modify_stage_config(
     get_deploy_config_path("minicpmo_4_5_duplex.yaml"),
     updates={
         "base_config": get_deploy_config_path("minicpmo_4_5.yaml"),
-        # Cap per-stage KV so Thinker/Talker/Code2Wav share one GPU with
-        # max_sessions=2. Talker must stay within its 4096 context so the
-        # 0.5 GiB budget passes vLLM's min-KV check at init.
+        # Talker context is 4096 (tts_config.max_position_embeddings); KV sizing
+        # is left automatic so duplex matches the simplex deploy profiles.
         "stages": {
-            0: {"kv_cache_memory_bytes": 6 * 1024 * 1024 * 1024},
-            1: {
-                "max_model_len": 4096,
-                "kv_cache_memory_bytes": 512 * 1024 * 1024,
-            },
-            2: {"kv_cache_memory_bytes": 256 * 1024 * 1024},
+            1: {"max_model_len": 4096},
         },
-        # Platform overrides apply after ordinary stage settings and would
-        # otherwise reinstate the base CUDA 2 GiB Talker default.
-        "platforms": {
-            "cuda": {
-                "stages": [
-                    {
-                        "stage_id": 1,
-                        "kv_cache_memory_bytes": 512 * 1024 * 1024,
-                    }
-                ]
-            }
-        },
-    },
-)
-CORE_DEPLOY_CONFIG = modify_stage_config(
-    DEPLOY_CONFIG,
-    updates={
-        "stages": {
-            0: {"enforce_eager": True},
-            1: {"enforce_eager": True},
-        }
     },
 )
 ASSET_DIR = Path(__file__).resolve().parents[3] / "assets" / "minicpmo_4_5"
@@ -71,17 +44,6 @@ SERVER_PARAMS = [
         id="three-stage-single-gpu",
     )
 ]
-CORE_SERVER_PARAMS = [
-    pytest.param(
-        OmniServerParams(
-            model=MODEL,
-            stage_config_path=CORE_DEPLOY_CONFIG,
-            use_stage_cli=False,
-            server_args=["--trust-remote-code"],
-        ),
-        id="three-stage-single-gpu",
-    )
-]
 
 
 def validated_wav(path: Path, expected_sha256: str) -> Path:
@@ -112,9 +74,11 @@ def validated_soft_interrupt_wav() -> Path:
 
 
 def resolve_ref_audio(model_prefix: str) -> Path:
-    if model_prefix:
-        model_root = Path(model_prefix) / MODEL
-    else:
+    # ``MODEL`` may be a local checkpoint directory instead of a Hub repo id
+    # (``Path("/prefix") / "/abs/path"`` collapses to the absolute path), so only
+    # fall back to the Hub cache when no such directory exists on disk.
+    model_root = Path(model_prefix) / MODEL if model_prefix else Path(MODEL)
+    if not model_root.is_dir():
         model_root = Path(snapshot_download(MODEL, local_files_only=True))
     ref_audio = model_root / REF_AUDIO_RELATIVE_PATH
     if not ref_audio.is_file():

tests/e2e/online_serving/test_flux_kontext_expansion.py

@@ -14,47 +14,15 @@ pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
 EDIT_PROMPT = "Transform this modern, geometrist image into a Vincent van Gogh style impressionist painting."
 NEGATIVE_PROMPT = "blurry, low quality, modern, geometrist"
 MODEL = "black-forest-labs/FLUX.1-Kontext-dev"
-PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"}, num_cards=2)
+BASE_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"})
 
 
 def _get_diffusion_feature_cases(model: str):
     return [
         pytest.param(
-            OmniServerParams(
-                model=model,
-                server_args=[
-                    "--tensor-parallel-size",
-                    "2",
-                    "--enable-cpu-offload",
-                ],
-            ),
-            id="parallel_tp_2",
-            marks=PARALLEL_FEATURE_MARKS,
-        ),
-        pytest.param(
-            OmniServerParams(
-                model=model,
-                server_args=[
-                    "--enable-cpu-offload",
-                    "--cfg-parallel-size",
-                    "2",
-                ],
-            ),
-            id="parallel_cfg_2",
-            marks=PARALLEL_FEATURE_MARKS,
-        ),
-        pytest.param(
-            OmniServerParams(
-                model=model,
-                server_args=[
-                    "--tensor-parallel-size",
-                    "2",
-                    "--enable-cpu-offload",
-                    "--cache-backend",
-                    "cache_dit",
-                ],
-            ),
-            id="parallel_002",
+            OmniServerParams(model=model),
+            id="base",
+            marks=BASE_FEATURE_MARKS,
         ),
     ]
 

tests/e2e/online_serving/test_minicpmo_4_5.py

@@ -1,10 +1,7 @@
-"""
-E2E Online tests for MiniCPM-o 4.5 model with multimodal input and audio / text output.
+"""E2E online tests for MiniCPM-o 4.5 multimodal input and audio/text output.
 
-MiniCPM-o 4.5 is a three-stage pipeline: Stage 0 Thinker (text), Stage 1 Talker
-(TTS codec), Stage 2 Code2Wav (vocoder). The talker -> Code2Wav bridge supports
-both async_chunk (streaming) and full-payload (``--no-async-chunk``) modes, so the
-tests below exercise the sync path as well.
+Exercises async chunk streaming (``--async-chunk``) across separate Thinker,
+Talker, and Code2Wav stages.
 """
 
 import os
@@ -14,12 +11,17 @@ import pytest
 from tests.helpers.mark import hardware_test
 from tests.helpers.media import generate_synthetic_audio, generate_synthetic_image, generate_synthetic_video
 from tests.helpers.runtime import OmniServerParams, dummy_messages_from_mix_data
-from tests.helpers.stage_config import get_deploy_config_path
+from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
 _MODEL = "openbmb/MiniCPM-o-4_5"
-_CI_DEPLOY = get_deploy_config_path("minicpmo_4_5.yaml")
+_CI_DEPLOY = modify_stage_config(
+    get_deploy_config_path("minicpmo_4_5.yaml"),
+    updates={
+        "stages": {0: {"default_sampling_params.max_tokens": 64}, 1: {"default_sampling_params.max_tokens": 1024}}
+    },
+)
 
 test_params = [
     pytest.param(
@@ -27,12 +29,10 @@ test_params = [
             model=_MODEL,
             stage_config_path=_CI_DEPLOY,
             use_stage_cli=False,
-            server_args=[
-                "--trust-remote-code",
-            ],
+            server_args=["--trust-remote-code", "--async-chunk"],
         ),
-        id="default",
-    )
+        id="async_chunk",
+    ),
 ]
 
 
@@ -98,7 +98,7 @@ def test_text_to_text_001(omni_server, openai_client) -> None:
 def test_text_to_audio_001(omni_server, openai_client) -> None:
     """
     Test text-only input generating text + audio output via OpenAI API.
-    This exercises the talker TTS region detection and in-process token2wav vocoder.
+    This exercises Talker TTS region detection and the Code2Wav stage.
     Deploy Setting: default single GPU
     Input Modal: text
     Output Modal: text + audio
@@ -132,13 +132,21 @@ def test_audio_to_text_audio_001(omni_server, openai_client) -> None:
     messages = dummy_messages_from_mix_data(
         system_prompt=get_system_prompt(),
         audio_data_url=audio_data_url,
-        content_text=get_prompt("mix"),
+        content_text=get_prompt(),
     )
 
     request_config = {
         "model": omni_server.model,
         "messages": messages,
         "stream": True,
+        "key_words": {"text": ["Beijing"]},
+        "modalities": ["text", "audio"],
+        "extra_body": {
+            "chat_template_kwargs": {
+                "use_tts_template": True,
+                "enable_thinking": False,
+            }
+        },
     }
 
     openai_client.send_omni_request(request_config, request_num=get_max_batch_size())

tests/e2e/online_serving/test_minicpmo_4_5_duplex.py

@@ -11,7 +11,6 @@ import pytest
 import websockets
 
 from tests.e2e.online_serving.helpers.minicpmo_4_5_duplex import (
-    CORE_SERVER_PARAMS,
     SERVER_PARAMS,
     demo_args,
     multi_session_args,
@@ -107,7 +106,7 @@ async def _run_protocol_smoke(*, url: str, model: str, ref_audio: Path) -> list[
 
 @pytest.mark.core_model
 @hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
-@pytest.mark.parametrize("omni_server", CORE_SERVER_PARAMS, indirect=True)
+@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True)
 def test_duplex_websocket_protocol_smoke(omni_server, model_prefix: str) -> None:
     ref_audio = resolve_ref_audio(model_prefix)
     events = asyncio.run(
@@ -134,8 +133,12 @@ def test_duplex_single_session_response_required(omni_server, model_prefix: str,
         output_dir=tmp_path / "single_session",
     )
     args.turns = 2
+    # Every turn replays the same active-speech window as the first one. The default
+    # shorter follow-up window is a different mid-utterance slice, which the native
+    # duplex model may legitimately answer with "listen" instead of a response.
+    args.turn_duration_ms = [args.first_turn_ms] * args.turns
     result = asyncio.run(run_demo(args))
-    assert result["ok"] is True, json.dumps(result, ensure_ascii=False, indent=2)
+    assert result["ok"] is True
     assert result["audio_delta_count"] > 0
     assert result["done_count"] == 2
     assert result["error_count"] == 0
@@ -160,7 +163,7 @@ def test_duplex_two_sessions_resume_and_takeover(omni_server, model_prefix: str,
             )
         )
     )
-    assert result["ok"] is True, json.dumps(result, ensure_ascii=False, indent=2)
+    assert result["ok"] is True
     assert result["session_count"] == 2
     assert result["resume"]["ok"] is True
     assert result["takeover"]["ok"] is True

tests/e2e/online_serving/test_minicpmo_4_5_expansion.py

@@ -2,6 +2,9 @@
 
 These cover modality combinations, separate Code2Wav behavior, sequential
 request isolation, longer audio output, and Chinese speech.
+
+Like Qwen3-Omni expansion, ``default`` exercises ``--no-async-chunk`` and
+``async_chunk`` exercises ``--async-chunk``.
 """
 
 import os
@@ -20,18 +23,27 @@ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 _MODEL = "openbmb/MiniCPM-o-4_5"
 _CI_DEPLOY = get_deploy_config_path("minicpmo_4_5.yaml")
 
+# Deploy YAML defaults to ``async_chunk: true``; keep sync and async as
+# separate parametrizations so both handoff paths are exercised.
 test_params = [
     pytest.param(
         OmniServerParams(
             model=_MODEL,
             stage_config_path=_CI_DEPLOY,
             use_stage_cli=False,
-            server_args=[
-                "--trust-remote-code",
-            ],
+            server_args=["--trust-remote-code", "--no-async-chunk"],
         ),
         id="default",
-    )
+    ),
+    pytest.param(
+        OmniServerParams(
+            model=_MODEL,
+            stage_config_path=_CI_DEPLOY,
+            use_stage_cli=False,
+            server_args=["--trust-remote-code", "--async-chunk"],
+        ),
+        id="async_chunk",
+    ),
 ]
 
 

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -284,6 +284,37 @@ def test_serve_cli_accepts_diffusion_pipeline_profiler_flag():
     assert stage_cfg["engine_args"]["enable_diffusion_pipeline_profiler"] is True
 
 
+def test_serve_cli_forwards_distilled_lora_to_diffusion_stage():
+    """Ensure startup distilled LoRA options reach the online diffusion stage."""
+    parser = TrackingArgumentParser()
+    subparsers = parser.add_subparsers(dest="command")
+    OmniServeCommand().subparser_init(subparsers)
+
+    args = parser.parse_args(
+        [
+            "serve",
+            "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+            "--omni",
+            "--lora-backend",
+            "distill",
+            "--lora-path",
+            "/models/high.safetensors",
+            "/models/low.safetensors",
+        ]
+    )
+
+    explicit_kwargs = args.get_explicit_kwargs_dict()
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0]
+    engine_args = stage_cfg["engine_args"]
+
+    assert explicit_kwargs["lora_backend"] == "distill"
+    assert engine_args["lora_backend"] == "distill"
+    assert engine_args["lora_path"] == [
+        "/models/high.safetensors",
+        "/models/low.safetensors",
+    ]
+
+
 def test_serve_cli_forwards_distributed_offload_residency():
     """Ensure the two-GPU DLO placement controls reach the diffusion stage."""
     parser = TrackingArgumentParser()

tests/examples/offline_inference/test_image_task_prompts.py

@@ -1,47 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from __future__ import annotations
-
-import pytest
-from PIL import Image
-
-from examples.offline_inference.image_to_video.image_to_video import build_image_to_video_prompt
-from examples.offline_inference.text_to_image.text_to_image import build_text_to_image_prompt
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-@pytest.mark.parametrize(
-    ("negative_prompt", "expected_negative_prompt"),
-    [(None, None), ("", ""), ("blurry", "blurry")],
-)
-def test_text_to_image_builds_canonical_prompt(
-    negative_prompt: str | None,
-    expected_negative_prompt: str | None,
-) -> None:
-    result = build_text_to_image_prompt("a red fox", negative_prompt)
-
-    assert result["prompt"] == "a red fox"
-    assert result["modalities"] == ["image"]
-    if expected_negative_prompt is None:
-        assert "negative_prompt" not in result
-    else:
-        assert result["negative_prompt"] == expected_negative_prompt
-
-
-def test_image_to_video_builds_canonical_prompt() -> None:
-    image = Image.new("RGB", (32, 16), "red")
-
-    result = build_image_to_video_prompt(
-        prompt="the fox turns toward the camera",
-        negative_prompt="flicker",
-        media_inputs={"image": image},
-    )
-
-    assert result == {
-        "prompt": "the fox turns toward the camera",
-        "modalities": ["video"],
-        "multi_modal_data": {"image": image},
-        "negative_prompt": "flicker",
-    }

tests/model_extras/test_model_extras.py

@@ -3,8 +3,6 @@
 
 from __future__ import annotations
 
-from typing import Any
-
 import pytest
 from PIL import Image
 
@@ -137,6 +135,67 @@ def test_lingbot_extra_registry_declares_request_params() -> None:
     assert should_init_extra_args_for_non_diffusion_stages("LingBotVideoPipeline") is False
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_lingbot_text_to_image_prompt_builder_preserves_empty_negative_prompt() -> None:
+    assert build_text_to_image_prompt(
+        "LingBotVideoPipeline",
+        prompt="a red fox in fresh snow",
+        negative_prompt="",
+        height=192,
+        width=320,
+    ) == {
+        "prompt": "a red fox in fresh snow",
+        "modalities": ["image"],
+        "negative_prompt": "",
+    }
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_lingbot_image_to_video_prompt_builder() -> None:
+    image = Image.new("RGB", (320, 192), "red")
+    result = build_image_to_video_prompt(
+        "LingBotVideoPipeline",
+        prompt="the fox looks toward the camera",
+        negative_prompt=None,
+        media_inputs={"image": image},
+        height=192,
+        width=320,
+        num_frames=9,
+    )
+    assert result == {
+        "prompt": "the fox looks toward the camera",
+        "modalities": ["video"],
+        "multi_modal_data": {"image": image},
+        "num_frames": 9,
+    }
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+@pytest.mark.parametrize(
+    "media_inputs",
+    [
+        {},
+        {"image": "input.png"},
+        {"image": [Image.new("RGB", (16, 16))]},
+        {"image": Image.new("RGB", (16, 16)), "video": b"video"},
+    ],
+    ids=["missing", "path", "multiple", "extra-modality"],
+)
+def test_lingbot_image_to_video_prompt_builder_rejects_invalid_media(
+    media_inputs: dict[str, object],
+) -> None:
+    with pytest.raises(ValueError, match="exactly one PIL image"):
+        build_image_to_video_prompt(
+            "LingBotVideoPipeline",
+            prompt="move",
+            negative_prompt=None,
+            media_inputs=media_inputs,
+        )
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 @pytest.mark.parametrize("pipeline_name", ["Cosmos3OmniDiffusersPipeline", "Cosmos3OmniPipeline"])
@@ -283,11 +342,8 @@ def test_ming_flash_omni_declared_extra_args_route_into_sampling_params() -> Non
 def test_ming_flash_omni_text_to_image_prompt_builder() -> None:
     assert build_text_to_image_prompt(
         "MingImagePipeline",
-        {
-            "prompt": "Please draw a cute cat.",
-            "modalities": ["image"],
-            "negative_prompt": "ugly, blurry",
-        },
+        prompt="Please draw a cute cat.",
+        negative_prompt="ugly, blurry",
         height=512,
         width=768,
     ) == {
@@ -304,7 +360,8 @@ def test_ming_flash_omni_text_to_image_prompt_builder() -> None:
     # pipeline's 1024x1024 default still applies.
     assert build_text_to_image_prompt(
         "MingImagePipeline",
-        {"prompt": "Draw a poster.", "modalities": ["image"]},
+        prompt="Draw a poster.",
+        negative_prompt=None,
     ) == {
         "prompt": "Draw a poster.",
         "modalities": ["image"],
@@ -334,6 +391,28 @@ def test_ming_flash_omni_image_to_image_prompt_builder() -> None:
     assert "negative_prompt" not in result
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+@pytest.mark.parametrize("pipeline_name", ["Cosmos3OmniDiffusersPipeline", "Cosmos3OmniPipeline"])
+def test_cosmos3_text_to_image_prompt_builder_selects_image_modality(pipeline_name: str) -> None:
+    assert build_text_to_image_prompt(
+        pipeline_name,
+        prompt="a red sports car at golden hour",
+        negative_prompt="blurry, distorted",
+        height=1024,
+        width=1024,
+    ) == {
+        "prompt": "a red sports car at golden hour",
+        "modalities": ["image"],
+        "negative_prompt": "blurry, distorted",
+    }
+    assert build_text_to_image_prompt(
+        pipeline_name,
+        prompt="a red sports car",
+        negative_prompt=None,
+    ) == {"prompt": "a red sports car", "modalities": ["image"]}
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_audiox_extra_registry_declares_request_and_response_params() -> None:
@@ -426,11 +505,8 @@ def test_unknown_pipeline_has_empty_extra_registry() -> None:
 def test_bagel_text_to_image_prompt_builder() -> None:
     assert build_text_to_image_prompt(
         "BagelPipeline",
-        {
-            "prompt": "a cat",
-            "modalities": ["image"],
-            "negative_prompt": "blurry",
-        },
+        prompt="a cat",
+        negative_prompt="blurry",
         height=512,
         width=768,
     ) == {
@@ -467,17 +543,14 @@ def test_bagel_image_to_image_prompt_builder() -> None:
 
 @pytest.mark.core_model
 @pytest.mark.cpu
-def test_unknown_pipeline_preserves_canonical_text_to_image_prompt() -> None:
-    canonical_prompt = {"prompt": "a cat", "modalities": ["image"]}
-    assert (
-        build_text_to_image_prompt(
-            "UnknownPipeline",
-            canonical_prompt,
-            height=512,
-            width=512,
-        )
-        is canonical_prompt
-    )
+def test_unknown_pipeline_uses_default_text_to_image_prompt() -> None:
+    assert build_text_to_image_prompt(
+        "UnknownPipeline",
+        prompt="a cat",
+        negative_prompt=None,
+        height=512,
+        width=512,
+    ) == {"prompt": "a cat"}
 
 
 @pytest.mark.core_model
@@ -496,14 +569,12 @@ def test_unknown_pipeline_uses_default_image_to_image_prompt() -> None:
     }
 
 
-def _build_vace_prompt(media_inputs: dict[str, Any], *, num_frames: int = 5) -> dict:
+def _build_vace_prompt(media_inputs: dict[str, object], *, num_frames: int = 5) -> dict:
     return build_image_to_video_prompt(
         "WanVACEPipeline",
-        {
-            "prompt": "a bird flying",
-            "modalities": ["video"],
-            "multi_modal_data": media_inputs,
-        },
+        prompt="a bird flying",
+        negative_prompt=None,
+        media_inputs=media_inputs,
         height=16,
         width=320,
         num_frames=num_frames,
@@ -526,7 +597,7 @@ def _build_vace_prompt(media_inputs: dict[str, Any], *, num_frames: int = 5) ->
     ],
     ids=["i2v", "v2lf", "flf2v", "inpaint", "r2v"],
 )
-def test_vace_image_to_video_prompt_builder(media_inputs: dict[str, Any]) -> None:
+def test_vace_image_to_video_prompt_builder(media_inputs: dict[str, object]) -> None:
     result = _build_vace_prompt(media_inputs)
     mmd = result["multi_modal_data"]
     if "reference_images" in mmd:
@@ -545,7 +616,7 @@ def test_vace_image_to_video_prompt_builder(media_inputs: dict[str, Any]) -> Non
         ({"control_image": Image.new("RGB", (320, 16))}, "Unsupported VACE media input"),
     ],
 )
-def test_vace_rejects_invalid_media_combinations(media_inputs: dict[str, Any], message: str) -> None:
+def test_vace_rejects_invalid_media_combinations(media_inputs: dict[str, object], message: str) -> None:
     with pytest.raises(ValueError, match=message):
         _build_vace_prompt(media_inputs)
 
@@ -593,7 +664,8 @@ def test_mammothmoda2_extra_registry_declares_request_and_response_params() -> N
 
     wrapper_prompt = build_text_to_image_prompt(
         "MammothModa2ForConditionalGeneration",
-        {"prompt": "a cat", "modalities": ["image"]},
+        prompt="a cat",
+        negative_prompt=None,
         height=256,
         width=256,
     )
@@ -607,11 +679,8 @@ def test_mammothmoda2_text_to_image_prompt_builder() -> None:
     # prompt is ignored (MammothModa2 t2i uses CFG, not an explicit negative path).
     assert build_text_to_image_prompt(
         "MammothModa2DiTPipeline",
-        {
-            "prompt": "a cat",
-            "modalities": ["image"],
-            "negative_prompt": "blurry",
-        },
+        prompt="a cat",
+        negative_prompt="blurry",
         height=512,
         width=768,
     ) == {

tests/model_tests/diffusion/diff_model_builders.py

@@ -37,6 +37,19 @@ def tiny_ltx2_builder() -> str:
     return build_tiny_from_configs("LTX2Pipeline", "Lightricks/LTX-2", TINY_CONFIGS_DIR / "LTX2Pipeline")
 
 
+def _shrink_flux_clip_text_encoder(config: dict) -> dict:
+    config["num_hidden_layers"] = 2
+    return config
+
+
+def _shrink_flux_t5_text_encoder(config: dict) -> dict:
+    config["num_layers"] = 2
+    config["num_decoder_layers"] = 2
+    config["d_ff"] = 64
+    config["num_heads"] = 4
+    return config
+
+
 def tiny_qwen_image_builder() -> str:
     def shrink_text_encoder(config: dict) -> dict:
         config["num_hidden_layers"] = 2
@@ -57,23 +70,24 @@ def tiny_qwen_image_builder() -> str:
 
 
 def tiny_flux_builder() -> str:
-    def shrink_clip_text_encoder(config: dict) -> dict:
-        config["num_hidden_layers"] = 2
-        return config
-
-    def shrink_t5_text_encoder(config: dict) -> dict:
-        config["num_layers"] = 2
-        config["num_decoder_layers"] = 2
-        config["d_ff"] = 64
-        config["num_heads"] = 4
-        return config
-
     return build_tiny_from_configs(
         "FluxPipeline",
         "black-forest-labs/FLUX.1-schnell",
         transform={
-            "text_encoder": shrink_clip_text_encoder,
-            "text_encoder_2": shrink_t5_text_encoder,
+            "text_encoder": _shrink_flux_clip_text_encoder,
+            "text_encoder_2": _shrink_flux_t5_text_encoder,
+            "transformer": partial(_shrink_dit_rope_config, num_single_layers=2, default_axes_dims_rope=[16, 56, 56]),
+        },
+    )
+
+
+def tiny_flux_kontext_builder() -> str:
+    return build_tiny_from_configs(
+        "FluxKontextPipeline",
+        "black-forest-labs/FLUX.1-Kontext-dev",
+        transform={
+            "text_encoder": _shrink_flux_clip_text_encoder,
+            "text_encoder_2": _shrink_flux_t5_text_encoder,
             "transformer": partial(_shrink_dit_rope_config, num_single_layers=2, default_axes_dims_rope=[16, 56, 56]),
         },
     )

tests/model_tests/diffusion/model_settings.py

@@ -59,6 +59,16 @@ DIFFUSION_TEST_SETTINGS = {
         builder=diff_model_builders.tiny_flux_builder,
         supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE],
     ),
+    "FluxKontextPipeline": DiffusionModelTestOpts(
+        model="black-forest-labs/FLUX.1-Kontext-dev",
+        builder=diff_model_builders.tiny_flux_kontext_builder,
+        supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE, DiffusionTasks.IMAGE_TO_IMAGE],
+        extra_test_groups=[
+            [DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CPU_OFFLOAD],
+            [DiffusionAccs.CFG_PARALLEL, DiffusionAccs.CPU_OFFLOAD],
+            [DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CPU_OFFLOAD, DiffusionAccs.CACHE_DIT],
+        ],
+    ),
     "Flux2Pipeline": DiffusionModelTestOpts(
         model="black-forest-labs/FLUX.2-dev",
         builder=diff_model_builders.tiny_flux2_builder,

tests/model_tests/diffusion/test_alignment.py

@@ -48,7 +48,6 @@ EXCLUDED_MODELS = [
     "InternVLAA1Pipeline",
     "LongCatImageEditPipeline",
     "StableDiffusion3Pipeline",
-    "FluxKontextPipeline",
     "HunyuanImage3ForCausalMM",
     "ErnieImagePipeline",
     "NextStep11Pipeline",

vllm_omni/config/omni_config.py

@@ -563,8 +563,9 @@ class _DiffusionConfigProjection:
     diffusers_load_kwargs: dict[str, Any] = field(default_factory=dict)
     diffusers_call_kwargs: dict[str, Any] = field(default_factory=dict)
     diffusers_pipeline_cls: Any = None
-    lora_path: str | None = None
+    lora_path: str | list[str] | None = None
     lora_scale: float = 1.0
+    lora_backend: str = "peft"
     max_cpu_loras: int | None = None
     output_type: str = "pil"
     enable_cpu_offload: bool = False

vllm_omni/config/stage_config.py

@@ -385,6 +385,9 @@ class StageDeployConfig:
     # Diffusion model loading and adapter construction.
     model_class_name: str | None = None
     diffusion_load_format: str | None = None
+    lora_path: str | list[str] | None = None
+    lora_backend: str | None = None
+    lora_scale: float | None = None
     diffusers_load_kwargs: dict[str, Any] | None = None
     diffusers_call_kwargs: dict[str, Any] | None = None
     diffusion_quantization_config: str | None = None

vllm_omni/diffusion/data.py

@@ -27,6 +27,7 @@ from vllm_omni.diffusion.diffusion_kv.config import (
     DiffusionKVCacheMode,
     parse_diffusion_kv_cache_mode,
 )
+from vllm_omni.diffusion.lora.manager import LoRABackend
 from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata
 from vllm_omni.diffusion.utils.network_utils import is_port_available
 from vllm_omni.errors import client_error_metadata
@@ -682,8 +683,9 @@ class OmniDiffusionConfig:
     # pipeline_config: PipelineConfig = field(default_factory=PipelineConfig, repr=False)
 
     # LoRA parameters
-    lora_path: str | None = None
-    lora_scale: float = 1.0
+    lora_path: str | list[str] | None = None
+    lora_scale: float | list[float] = 1.0
+    lora_backend: LoRABackend = LoRABackend.PEFT  # available choices: ["peft", "distill"]
     max_cpu_loras: int | None = None
 
     output_type: str = "pil"

vllm_omni/diffusion/lora/loader.py

@@ -0,0 +1,461 @@
+import os
+from collections import OrderedDict, defaultdict
+from collections.abc import Callable
+
+import torch
+from diffusers.loaders.lora_conversion_utils import (
+    _convert_non_diffusers_qwen_lora_to_diffusers,
+    _convert_non_diffusers_wan_lora_to_diffusers,
+)
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.utils.tf_utils import get_transformer_from_pipeline
+
+logger = init_logger(__name__)
+
+lora_convert_mapping: dict[str, Callable] = {
+    "QwenImagePipeline": _convert_non_diffusers_qwen_lora_to_diffusers,
+    "QwenImageEditPipeline": _convert_non_diffusers_qwen_lora_to_diffusers,
+    "QwenImageEditPlusPipeline": _convert_non_diffusers_qwen_lora_to_diffusers,
+    "Wan22Pipeline": _convert_non_diffusers_wan_lora_to_diffusers,
+    "Wan22I2VPipeline": _convert_non_diffusers_wan_lora_to_diffusers,
+}
+
+
+def get_converter_by_pipeline(pipeline):
+    for pipeline_cls in pipeline.__class__.__mro__:
+        converter = lora_convert_mapping.get(pipeline_cls.__name__)
+        if converter is not None:
+            return converter
+    return None
+
+
+def _prepare_lora_delta(
+    lora_state_dict,
+    base_key: str,
+    param_to_weight_names: dict[str, list[str]] | None = None,
+    is_bias: bool = False,
+    lora_a_suffix: str = "lora_A.weight",
+    lora_b_suffix: str = "lora_B.weight",
+    lora_bias_suffix: str = "bias",
+):
+    used_keys = set()
+    # stacked_params_mapping                       param_to_weight_names
+    # [(".to_qkv", ".to_q.", "q")
+    # (".to_qkv", ".to_k.", "k")  ========> {".to_qkv": [".to_q", ".to_k", ".to_v"]}
+    # (".to_qkv", ".to_v.", "v")]
+
+    # stacked_sd is to store packed parameter deltas (format: [str, list[torch.Tensor]])
+    # example: {".to_qkv": [delta_q, delta_k, delta_v]}
+    stacked_deltas = []
+    is_stacked_param = False
+    if param_to_weight_names is None:
+        param_to_weight_names = defaultdict(list)
+    for param_name, weight_names in param_to_weight_names.items():
+        if param_name not in base_key:
+            continue
+        is_stacked_param = True
+        # handle lora_a_key and lora_b_key together
+        if is_bias:
+            for weight_name in weight_names:
+                lora_bias_key = f"{base_key.replace(param_name, weight_name)}.{lora_bias_suffix}"
+                if lora_bias_key not in lora_state_dict:
+                    return None, used_keys
+                delta = lora_state_dict[lora_bias_key]
+                stacked_deltas.append(delta)
+                used_keys.add(lora_bias_key)
+        else:
+            for weight_name in weight_names:
+                lora_a_key = f"{base_key.replace(param_name, weight_name)}.{lora_a_suffix}"
+                lora_b_key = lora_a_key.replace(lora_a_suffix, lora_b_suffix)
+                if lora_a_key not in lora_state_dict or lora_b_key not in lora_state_dict:
+                    return None, used_keys
+                a = lora_state_dict[lora_a_key]
+                b = lora_state_dict[lora_b_key]
+                delta = torch.matmul(b, a)
+                stacked_deltas.append(delta)
+                used_keys.add(lora_a_key)
+                used_keys.add(lora_b_key)
+        continue
+
+    if is_stacked_param:
+        return torch.concat(stacked_deltas), used_keys
+
+    if is_bias:
+        lora_bias_key = f"{base_key}.{lora_bias_suffix}"
+        if lora_bias_key not in lora_state_dict:
+            return None, used_keys
+        used_keys.add(lora_bias_key)
+        return lora_state_dict[lora_bias_key], used_keys
+
+    lora_a_key = f"{base_key}.{lora_a_suffix}"
+    lora_b_key = f"{base_key}.{lora_b_suffix}"
+    if lora_a_key not in lora_state_dict or lora_b_key not in lora_state_dict:
+        return None, used_keys
+    a = lora_state_dict[lora_a_key]
+    b = lora_state_dict[lora_b_key]
+    used_keys.add(lora_a_key)
+    used_keys.add(lora_b_key)
+    return torch.matmul(b, a), used_keys
+
+
+def _load_lora_state_dict(
+    pretrained_model_name_or_path: str,
+    weights_name: str | None = None,
+    use_safetensors: bool = True,
+    subfolder: str | None = None,
+    cache_dir: str | None = None,
+    force_download: bool = False,
+    local_files_only: bool | None = None,
+):
+    # first we try to load it from local disk
+    if use_safetensors and pretrained_model_name_or_path.endswith(".safetensors"):
+        return load_file(pretrained_model_name_or_path)
+
+    if os.path.isdir(pretrained_model_name_or_path):
+        model_file = pretrained_model_name_or_path
+        if subfolder:
+            model_file = os.path.join(model_file, subfolder)
+        if weights_name is None:
+            raise ValueError("weights_name is required when loading from a directory")
+        model_file = os.path.join(model_file, weights_name)
+        return load_file(model_file)
+
+    # finally, we try to load it from the internet
+    try:
+        model_file = hf_hub_download(
+            pretrained_model_name_or_path,
+            filename=weights_name,
+            subfolder=subfolder,
+            cache_dir=cache_dir,
+            force_download=force_download,
+            local_files_only=local_files_only,
+        )
+        return load_file(model_file)
+    except Exception as e:
+        logger.error(f"Failed to download {weights_name} from {pretrained_model_name_or_path}: {e}")
+        raise e
+
+
+def _remap_state_dict_keys(sd, rules):
+    """
+    Remap keys in a state_dict by sequentially applying a list of substring substitution rules.
+
+    This utility function transforms parameter names by applying a sequence of
+    (old_substring, new_substring) pairs. For each key in the input dictionary,
+    **all rules are evaluated in the given order**, and every rule whose
+    `old_substring` is found within the current key triggers a replacement.
+    Substitutions are cumulative: the key is updated after each match, and
+    subsequent rules operate on the result of previous replacements.
+
+    The order of rules matters; users are responsible for arranging them to
+    achieve the desired final key. Keys that do not match any rule are copied
+    unchanged.
+
+    Args:
+        sd (dict[str, Any]): Original state_dict mapping parameter names to tensors.
+        rules (List[Tuple[str, str]]): A list of (old_substring, new_substring) pairs.
+            Rules are applied sequentially, and **all matching rules are executed**.
+
+    Returns:
+        dict[str, Any]: A new state_dict with remapped keys. Values are unchanged.
+    """
+    new_sd = OrderedDict()
+    for k, v in sd.items():
+        new_key = k
+        matched = False
+        for p1, p2 in rules:
+            if p1 not in k:
+                continue
+            new_key = new_key.replace(p1, p2)
+            matched = True
+        if matched:
+            new_sd[new_key] = v
+        else:
+            new_sd[k] = v
+    return new_sd
+
+
+def _apply_diffusers_lora_alpha_scaling(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+    """Fold Diffusers-format LoRA alpha values into their B matrices."""
+    scaled_state_dict = dict(state_dict)
+    for alpha_key in (key for key in state_dict if key.endswith(".alpha")):
+        base_key = alpha_key[: -len(".alpha")]
+        lora_a_key = f"{base_key}.lora_A.weight"
+        lora_b_key = f"{base_key}.lora_B.weight"
+        if lora_a_key not in state_dict or lora_b_key not in state_dict:
+            raise ValueError(f"LoRA alpha key {alpha_key!r} does not have matching lora_A and lora_B weights.")
+
+        rank = state_dict[lora_a_key].shape[0]
+        alpha = state_dict[alpha_key]
+
+        scaled_state_dict[lora_b_key] = state_dict[lora_b_key] * (alpha.item() / rank)
+        scaled_state_dict.pop(alpha_key)
+
+    return scaled_state_dict
+
+
+class LoraLoaderMixin:
+    transformer_name = "transformer"
+
+    # Lazy initialization to avoid MRO issues: __init__ may not be called
+    # when mixin with nn.Module
+    @property
+    def lora_loaded(self):
+        if not hasattr(self, "_lora_loaded"):
+            self._lora_loaded = {}
+        return self._lora_loaded
+
+    @lora_loaded.setter
+    def lora_loaded(self, value):
+        self._lora_loaded = value
+
+    @classmethod
+    def load_lora_into_module(
+        cls,
+        state_dict,
+        module,
+        prefix: str = "transformer",
+        lora_a_suffix: str = "lora_A.weight",
+        lora_b_suffix: str = "lora_B.weight",
+        lora_bias_suffix: str = "bias",
+    ):
+        lora_loaded_keys = set()
+
+        # prepare stacked parameters mapping for later use
+        # stacked_params_mapping                       param_to_weight_names
+        # [(".to_qkv", ".to_q.", "q")
+        # (".to_qkv", ".to_k.", "k")  ========> {".to_qkv": [".to_q", ".to_k", ".to_v"]}
+        # (".to_qkv", ".to_v.", "v")]
+        param_to_weight_names = defaultdict(list)
+        if hasattr(module, "stacked_params_mapping"):
+            for param_name, weight_name, _ in module.stacked_params_mapping:
+                param_to_weight_names[param_name].append(weight_name)
+
+        for name, params in module.named_parameters(prefix):
+            is_bias = False
+            if name.endswith(".bias"):
+                base_key = name[: -len(".bias")]
+                is_bias = True
+            elif name.endswith(".weight"):
+                base_key = name[: -len(".weight")]
+            else:
+                continue
+
+            delta, used_keys = _prepare_lora_delta(
+                state_dict,
+                base_key,
+                param_to_weight_names,
+                is_bias,
+                lora_a_suffix,
+                lora_b_suffix,
+                lora_bias_suffix,
+            )
+            if delta is None:
+                continue
+            lora_loaded_keys.update(used_keys)
+
+            delta = delta.to(device=params.device, dtype=params.dtype)
+            with torch.no_grad():
+                params.add_(delta)
+            del delta
+
+        missing_keys = set(state_dict.keys()) - lora_loaded_keys
+        for k in missing_keys:
+            logger.warning(f"Missing loading lora key: {k}")
+
+        logger.info(f"{len(lora_loaded_keys)} lora keys loaded into {module.__class__.__name__}.")
+        return lora_loaded_keys
+
+    @classmethod
+    def unload_module_lora(
+        cls,
+        state_dict,
+        module,
+        prefix: str = "transformer",
+        lora_a_suffix: str = "lora_A.weight",
+        lora_b_suffix: str = "lora_B.weight",
+        lora_bias_suffix: str = "bias",
+    ):
+        lora_unloaded_keys = set()
+
+        param_to_weight_names = defaultdict(list)
+        if hasattr(module, "stacked_params_mapping"):
+            for param_name, weight_name, _ in module.stacked_params_mapping:
+                param_to_weight_names[param_name].append(weight_name)
+
+        for name, param in module.named_parameters(prefix):
+            is_bias = False
+            if name.endswith(".bias"):
+                base_key = name[: -len(".bias")]
+                is_bias = True
+            elif name.endswith(".weight"):
+                base_key = name[: -len(".weight")]
+            else:
+                continue
+
+            delta, used_keys = _prepare_lora_delta(
+                state_dict,
+                base_key,
+                param_to_weight_names,
+                is_bias,
+                lora_a_suffix,
+                lora_b_suffix,
+                lora_bias_suffix,
+            )
+            if delta is None:
+                continue
+            lora_unloaded_keys.update(used_keys)
+
+            delta = delta.to(device=param.device, dtype=param.dtype)
+            with torch.no_grad():
+                param.sub_(delta)
+            del delta
+
+        missing_keys = set(state_dict.keys()) - lora_unloaded_keys
+        for k in missing_keys:
+            logger.warning(f"Missing unloading lora key: {k}")
+
+        logger.info(f"{len(lora_unloaded_keys)} lora keys unloaded from {module.__class__.__name__}.")
+        return lora_unloaded_keys
+
+
+class QwenImageLoraLoaderMixin(LoraLoaderMixin):
+    def load_lora_weights(
+        self,
+        pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor],
+        adapter_name: str | None = None,
+    ):
+        if adapter_name in self.lora_loaded:
+            return
+
+        if isinstance(pretrained_model_name_or_path_or_dict, dict):
+            state_dict = pretrained_model_name_or_path_or_dict
+        else:
+            state_dict = _load_lora_state_dict(pretrained_model_name_or_path_or_dict)
+
+        has_alpha = any(k.endswith(".alpha") for k in state_dict)
+        is_non_diffusers_format = any(k.startswith("diffusion_model.") for k in state_dict)
+        is_diffusers_lora_format = any(k.endswith(".lora_A.weight") for k in state_dict)
+        if has_alpha and is_diffusers_lora_format:
+            state_dict = _apply_diffusers_lora_alpha_scaling(state_dict)
+
+        if has_alpha or is_non_diffusers_format:
+            converter = get_converter_by_pipeline(self)
+            if converter is None:
+                raise ValueError(f"Converter for Lora weights not found for {self.__class__.__name__}")
+
+            state_dict = converter(state_dict)
+
+        state_dict = _remap_state_dict_keys(state_dict, [(".to_out.0.", ".to_out.")])
+
+        self.load_lora_into_module(
+            state_dict,
+            self.transformer,
+            prefix=self.transformer_name,
+        )
+
+        self.lora_loaded[adapter_name] = state_dict
+
+    def unload_lora_weights(self, adapter_name: str):
+        if adapter_name not in self.lora_loaded:
+            return
+
+        state_dict = self.lora_loaded[adapter_name]
+
+        transformer = get_transformer_from_pipeline(self)
+        self.unload_module_lora(state_dict, transformer, prefix=self.transformer_name)
+
+        del self.lora_loaded[adapter_name]
+
+
+class WanLoraLoaderMixin(LoraLoaderMixin):
+    transformer_2_name = "transformer_2"
+
+    def load_lora_weights(
+        self,
+        pretrained_model_name_or_path: str | list[str],
+        adapter_name: str | None = None,
+    ):
+        """Load LoRA weights into the pipeline's transformer(s).
+
+        Args:
+            pretrained_model_name_or_path: Path(s) to concrete .safetensors LoRA files.
+                - str: a single .safetensors file. Only valid for Wan2.1
+                  (single-transformer pipelines).
+                - list[str]: one file per transformer, matched **by position** to
+                  [transformer, transformer_2]. Required for Wan2.2 MoE.
+            adapter_name: Name to register the adapter under.
+        """
+        if adapter_name in self.lora_loaded:
+            return
+
+        # Normalize to list[(lora_path, target_module_name)].
+        lora_paths = None
+        if isinstance(pretrained_model_name_or_path, str):
+            lora_paths = [pretrained_model_name_or_path]
+        else:
+            lora_paths = list(pretrained_model_name_or_path)
+
+        target_modules = (
+            [self.transformer_name, self.transformer_2_name] if self.has_transformer_2 else [self.transformer_name]
+        )
+        if len(lora_paths) != len(target_modules):
+            raise ValueError(
+                f"{self.__class__.__name__} expects {len(target_modules)} LoRA file(s) "
+                f"(target_modules={target_modules}, paths={lora_paths}), got {len(lora_paths)}. Pass a list of concrete "  # noqa
+                f".safetensors files, one per transformer."
+            )
+
+        module_to_lora_sd = dict()
+        for lora_path, target_module_name in zip(lora_paths, target_modules):
+            state_dict = _load_lora_state_dict(lora_path)
+            is_non_diffusers_format = any(k.startswith("diffusion_model.") for k in state_dict)
+            if is_non_diffusers_format:
+                converter = get_converter_by_pipeline(self)
+                if converter is None:
+                    raise ValueError(f"Converter for Lora weights not found for {self.__class__.__name__}")
+
+                state_dict = converter(state_dict)
+
+            state_dict = _remap_state_dict_keys(
+                state_dict,
+                [
+                    (".ffn.net.0.", ".ffn.net_0."),
+                    (".ffn.net.2.", ".ffn.net_2."),
+                    (
+                        ".to_out.0.",
+                        ".to_out.",
+                    ),
+                    # '.bias_B.bias' is a horrible name in diffuser's.
+                    # So we remap it to '.bias'
+                    (".lora_B.bias", ".bias"),
+                ],
+            )
+
+            module = get_transformer_from_pipeline(self, transformer_name=target_module_name)
+            if module is None:
+                logger.warning(
+                    f"Skip LoRA {lora_path}: target module '{target_module_name}' "
+                    f"is not loaded in pipeline {self.__class__.__name__} "
+                    f"(e.g. disabled by boundary_ratio)."
+                )
+                continue
+
+            self.load_lora_into_module(state_dict, module, prefix=self.transformer_name, lora_bias_suffix="bias")
+            module_to_lora_sd[target_module_name] = state_dict
+
+        self.lora_loaded[adapter_name] = module_to_lora_sd
+
+    def unload_lora_weights(self, adapter_name: str):
+        if adapter_name not in self.lora_loaded:
+            return
+
+        module_to_lora_sd = self.lora_loaded[adapter_name]
+        for module_name, lora_sd in module_to_lora_sd.items():
+            module = get_transformer_from_pipeline(self, transformer_name=module_name)
+            self.unload_module_lora(lora_sd, module, prefix=self.transformer_name)
+
+        del self.lora_loaded[adapter_name]

vllm_omni/diffusion/lora/manager.py

@@ -4,6 +4,7 @@
 import math
 import time
 from collections import OrderedDict
+from enum import Enum
 from typing import get_args
 
 import torch
@@ -33,6 +34,11 @@ from vllm_omni.lora.utils import stable_lora_int_id
 logger = init_logger(__name__)
 
 
+class LoRABackend(str, Enum):
+    PEFT = "peft"
+    DISTILL = "distill"
+
+
 class DiffusionLoRAManager:
     """Manager for LoRA adapters in diffusion models.
 

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py

@@ -25,6 +25,7 @@ from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_qwenimage import DistributedAutoencoderKLQwenImage
 from vllm_omni.diffusion.distributed.utils import get_local_device
+from vllm_omni.diffusion.lora.loader import QwenImageLoraLoaderMixin
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
 from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
@@ -264,7 +265,11 @@ def apply_rotary_emb_qwen(
 
 
 class QwenImagePipeline(
-    nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery
+    nn.Module,
+    QwenImageCFGParallelMixin,
+    DiffusionPipelineProfilerMixin,
+    SupportsComponentDiscovery,
+    QwenImageLoraLoaderMixin,
 ):
     supports_request_batch = True
     _dit_modules: ClassVar[list[str]] = ["transformer"]

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py

@@ -25,6 +25,7 @@ from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.pipeline_parallel import AsyncLatents, PipelineParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.forward_context import DenoiseProgressMixin
+from vllm_omni.diffusion.lora.loader import WanLoraLoaderMixin
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
 from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
@@ -285,6 +286,7 @@ class Wan22Pipeline(
     DenoiseProgressMixin,
     DiffusionPipelineProfilerMixin,
     SupportsComponentDiscovery,
+    WanLoraLoaderMixin,
 ):
     _dit_modules: ClassVar[list[str]] = ["transformer", "transformer_2"]
     _encoder_modules: ClassVar[list[str]] = ["text_encoder"]

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py

@@ -25,6 +25,7 @@ from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.pipeline_parallel import AsyncLatents, PipelineParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.forward_context import DenoiseProgressMixin
+from vllm_omni.diffusion.lora.loader import WanLoraLoaderMixin
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
 from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
@@ -147,6 +148,7 @@ class Wan22I2VPipeline(
     DiffusionPipelineProfilerMixin,
     DenoiseProgressMixin,
     SupportsComponentDiscovery,
+    WanLoraLoaderMixin,
 ):
     """
     Wan2.2 Image-to-Video Pipeline.

vllm_omni/diffusion/utils/tf_utils.py

@@ -70,9 +70,10 @@ def find_module_with_attr(model, attr_name="transformer"):
     return None
 
 
-def get_transformer_from_pipeline(pipeline: Any):
-    pipe = find_module_with_attr(pipeline, attr_name="transformer")
+def get_transformer_from_pipeline(pipeline: Any, transformer_name: str = "transformer"):
+    """Return a named transformer from a pipeline or pipeline wrapper."""
+    pipe = find_module_with_attr(pipeline, attr_name=transformer_name)
 
     if pipe is not None:
-        return pipe.transformer
+        return getattr(pipe, transformer_name)
     return None

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -49,7 +49,7 @@ from vllm_omni.diffusion.distributed.parallel_state import (
 )
 from vllm_omni.diffusion.forward_context import set_forward_context
 from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE, pack_diffusion_output_shm
-from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
+from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager, LoRABackend
 from vllm_omni.diffusion.registry import get_diffusion_ir_op_priority_func
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput, KVPrefetchJob
@@ -388,14 +388,31 @@ class DiffusionWorker:
         """Initialize the LoRA manager for this worker."""
         if self.model_runner.pipeline is None:
             return
-        self.lora_manager = DiffusionLoRAManager(
-            pipeline=self.model_runner.pipeline,
-            device=self.device,
-            dtype=self.od_config.dtype,
-            max_cached_adapters=self.od_config.max_cpu_loras,
-            lora_path=self.od_config.lora_path,
-            lora_scale=self.od_config.lora_scale,
-        )
+
+        lora_path = self.od_config.lora_path
+        if isinstance(lora_path, list) and len(lora_path) == 1:
+            lora_path = lora_path[0]
+
+        lora_backend = self.od_config.lora_backend
+        if lora_backend == LoRABackend.PEFT:
+            self.lora_manager = DiffusionLoRAManager(
+                pipeline=self.model_runner.pipeline,
+                device=self.device,
+                dtype=self.od_config.dtype,
+                max_cached_adapters=self.od_config.max_cpu_loras,
+                lora_path=lora_path,
+                lora_scale=self.od_config.lora_scale,
+            )
+        elif lora_backend == LoRABackend.DISTILL:
+            pipeline = self.model_runner.pipeline
+            if hasattr(pipeline, "load_lora_weights"):
+                if self.od_config.lora_scale > 1.0:
+                    logger.warning("lora_scale > 1.0 may not take any effect when using distilled LoRA backend.")
+                pipeline.load_lora_weights(lora_path)
+            else:
+                logger.warning("Pipeline does not support loading distilled LoRA weights for now.")
+        else:
+            raise ValueError(f"Unknown LoRA backend: {lora_backend}. Available choices: {LoRABackend.__members__}")
 
     def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None:
         """Start or stop profiling for this GPU worker.

vllm_omni/engine/arg_utils.py

@@ -471,6 +471,9 @@ class OrchestratorArgs:
     num_gpus: int | None = None
     model_class_name: str | None = None
     diffusion_load_format: str | None = None
+    lora_path: list[str] | None = None
+    lora_backend: str | None = None
+    lora_scale: float | None = None
     diffusers_load_kwargs: str = "{}"
     diffusers_call_kwargs: str = "{}"
     ulysses_degree: int | None = None

vllm_omni/engine/async_omni_engine.py

@@ -1048,6 +1048,9 @@ class AsyncOmniEngine:
             "boundary_ratio": kwargs.get("boundary_ratio", None),
             "flow_shift": kwargs.get("flow_shift", None),
             "diffusion_load_format": kwargs.get("diffusion_load_format", "default"),
+            "lora_path": kwargs.get("lora_path", None),
+            "lora_scale": kwargs.get("lora_scale", 1.0),
+            "lora_backend": kwargs.get("lora_backend", "peft"),
             "custom_pipeline_args": kwargs.get("custom_pipeline_args", None),
             "worker_extension_cls": kwargs.get("worker_extension_cls", None),
             "trust_remote_code": (False if kwargs.get("trust_remote_code") is None else kwargs["trust_remote_code"]),
@@ -1251,6 +1254,9 @@ class AsyncOmniEngine:
                 if lora_scale is not None:
                     if not hasattr(cfg.engine_args, "lora_scale") or cfg.engine_args.lora_scale is None:
                         cfg.engine_args.lora_scale = lora_scale
+                if kwargs.get("lora_backend") is not None:
+                    if not hasattr(cfg.engine_args, "lora_backend") or cfg.engine_args.lora_backend is None:
+                        cfg.engine_args.lora_backend = kwargs["lora_backend"]
                 if (
                     kwargs.get("diffusion_attention_config") is not None
                     or kwargs.get("diffusion_attention_backend") is not None

vllm_omni/entrypoints/cli/serve.py

@@ -164,6 +164,13 @@ class OmniServeCommand(CLISubcommand):
                     "(single-runtime) or `--omni-dp-size-local` (headless / multi-runtime)."
                 )
 
+        lora_backend = getattr(args, "lora_backend", None)
+        lora_path = getattr(args, "lora_path", None)
+        if lora_backend == "distill" and not lora_path:
+            raise ValueError("--lora-backend distill requires --lora-path")
+        if (lora_backend or "peft") == "peft" and isinstance(lora_path, list) and len(lora_path) > 1:
+            raise ValueError("--lora-backend peft accepts only one startup LoRA path")
+
         # --omni-lb-policy is validated against the LoadBalancingPolicy enum.
         omni_lb_policy = getattr(args, "omni_lb_policy", None)
         if omni_lb_policy is not None:
@@ -427,6 +434,32 @@ class OmniServeCommand(CLISubcommand):
                 "custom_pipeline, dummy, or diffusers for the HF diffusers adapter."
             ),
         )
+        omni_config_group.add_argument(
+            "--lora-path",
+            type=str,
+            nargs="+",
+            default=None,
+            help=(
+                "LoRA checkpoint path(s) loaded when the diffusion server starts. "
+                "The distill backend accepts one file per pipeline transformer."
+            ),
+        )
+        omni_config_group.add_argument(
+            "--lora-backend",
+            type=str,
+            choices=["peft", "distill"],
+            default=None,
+            help=(
+                "Diffusion LoRA loading backend. 'distill' fuses checkpoint files "
+                "into the base model at server startup; 'peft' uses the adapter manager."
+            ),
+        )
+        omni_config_group.add_argument(
+            "--lora-scale",
+            type=float,
+            default=None,
+            help="Scale for a startup PEFT LoRA. Distilled LoRAs are fused at their checkpoint scale.",
+        )
         omni_config_group.add_argument(
             "--diffusion-compile-granularity",
             choices=["regional", "full"],

vllm_omni/model_extras/cosmos3.py

@@ -3,6 +3,32 @@
 
 from __future__ import annotations
 
+from typing import Any
+
+
+def build_text_to_image_prompt(
+    prompt: str,
+    negative_prompt: str | None,
+    height: int | None = None,
+    width: int | None = None,
+) -> dict[str, Any]:
+    """Build a Cosmos3 T2I prompt by selecting image output modality.
+
+    Cosmos3 uses the same pipeline class for T2I and video modes. The OpenAI
+    image endpoint selects T2I by adding ``modalities=["image"]``. ``height``
+    and ``width`` are accepted for registry compatibility; sizing travels
+    through sampling params instead of the prompt payload.
+    """
+    del height, width
+    text_prompt: dict[str, Any] = {
+        "prompt": prompt,
+        "modalities": ["image"],
+    }
+    if negative_prompt is not None:
+        text_prompt["negative_prompt"] = negative_prompt
+    return text_prompt
+
+
 COSMOS3_EXTRA_BODY_PARAMS = frozenset(
     {
         "flow_shift",

vllm_omni/model_extras/lingbot_video.py

@@ -3,6 +3,11 @@
 
 from __future__ import annotations
 
+from collections.abc import Mapping
+from typing import Any
+
+from PIL import Image
+
 LINGBOT_VIDEO_EXTRA_BODY_PARAMS = frozenset(
     {
         "batch_cfg",
@@ -19,3 +24,43 @@ LINGBOT_VIDEO_EXTRA_BODY_PARAMS = frozenset(
         "t_thresh",
     }
 )
+
+
+def build_text_to_image_prompt(
+    prompt: str,
+    negative_prompt: str | None,
+    height: int | None = None,
+    width: int | None = None,
+) -> dict[str, Any]:
+    del height, width
+    result: dict[str, Any] = {
+        "prompt": prompt,
+        "modalities": ["image"],
+    }
+    if negative_prompt is not None:
+        result["negative_prompt"] = negative_prompt
+    return result
+
+
+def build_image_to_video_prompt(
+    prompt: str,
+    negative_prompt: str | None,
+    media_inputs: Mapping[str, Any],
+    height: int | None = None,
+    width: int | None = None,
+    num_frames: int | None = None,
+) -> dict[str, Any]:
+    del height, width
+    if set(media_inputs) != {"image"} or not isinstance(media_inputs["image"], Image.Image):
+        raise ValueError("LingBot text-image-to-video requires exactly one PIL image input.")
+
+    result: dict[str, Any] = {
+        "prompt": prompt,
+        "modalities": ["video"],
+        "multi_modal_data": {"image": media_inputs["image"]},
+    }
+    if num_frames is not None:
+        result["num_frames"] = num_frames
+    if negative_prompt is not None:
+        result["negative_prompt"] = negative_prompt
+    return result

vllm_omni/model_extras/registry.py

@@ -28,12 +28,21 @@ from vllm_omni.model_extras.cosmos3 import (
     COSMOS3_EXTRA_BODY_PARAMS,
     COSMOS3_EXTRA_OUTPUT_PARAMS,
 )
+from vllm_omni.model_extras.cosmos3 import (
+    build_text_to_image_prompt as build_cosmos3_text_to_image_prompt,
+)
 from vllm_omni.model_extras.helios import (
     HELIOS_EXTRA_BODY_PARAMS,
     HELIOS_EXTRA_OUTPUT_PARAMS,
 )
 from vllm_omni.model_extras.hunyuan_image3 import build_x_to_text_prompt as build_hunyuan_x_to_text_prompt
 from vllm_omni.model_extras.lingbot_video import LINGBOT_VIDEO_EXTRA_BODY_PARAMS
+from vllm_omni.model_extras.lingbot_video import (
+    build_image_to_video_prompt as build_lingbot_image_to_video_prompt,
+)
+from vllm_omni.model_extras.lingbot_video import (
+    build_text_to_image_prompt as build_lingbot_text_to_image_prompt,
+)
 from vllm_omni.model_extras.ltx2 import LTX_EXTRA_BODY_PARAMS, LTX_EXTRA_OUTPUT_PARAMS
 from vllm_omni.model_extras.magi_human import (
     MAGI_HUMAN_EXTRA_BODY_PARAMS,
@@ -138,6 +147,18 @@ def build_x_to_text_prompt(
     return builder(model, prompt, has_image)
 
 
+def default_text_to_image_prompt(
+    prompt: str,
+    negative_prompt: str | None,
+    height: int | None = None,
+    width: int | None = None,
+) -> dict[str, Any]:
+    result: dict[str, Any] = {"prompt": prompt}
+    if negative_prompt is not None:
+        result["negative_prompt"] = negative_prompt
+    return result
+
+
 def default_image_to_image_prompt(
     prompt: str,
     negative_prompt: str | None,
@@ -154,6 +175,20 @@ def default_image_to_image_prompt(
     return result
 
 
+def default_image_to_video_prompt(
+    prompt: str,
+    negative_prompt: str | None,
+    media_inputs: Mapping[str, Any],
+    height: int | None = None,
+    width: int | None = None,
+    num_frames: int | None = None,
+) -> dict[str, Any]:
+    del height, width, num_frames
+    if set(media_inputs) != {"image"} or not isinstance(media_inputs["image"], Image.Image):
+        raise ValueError("This model only supports a single --image input in the shared image-to-video example.")
+    return default_image_to_image_prompt(prompt, negative_prompt, media_inputs["image"])
+
+
 _EXTRA_SPECS: dict[str, dict[str, Any]] = {
     "AudioXPipeline": {
         "extra_body_params": AUDIOX_EXTRA_BODY_PARAMS,
@@ -173,11 +208,12 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = {
     "Cosmos3OmniDiffusersPipeline": {
         "extra_body_params": COSMOS3_EXTRA_BODY_PARAMS,
         "extra_output_params": COSMOS3_EXTRA_OUTPUT_PARAMS,
-        # The shared T2I example already supplies modalities=["image"].
+        "text_to_image_prompt_builder": build_cosmos3_text_to_image_prompt,
     },
     "Cosmos3OmniPipeline": {
         "extra_body_params": COSMOS3_EXTRA_BODY_PARAMS,
         "extra_output_params": COSMOS3_EXTRA_OUTPUT_PARAMS,
+        "text_to_image_prompt_builder": build_cosmos3_text_to_image_prompt,
     },
     "MagiHumanPipeline": {
         "extra_body_params": MAGI_HUMAN_EXTRA_BODY_PARAMS,
@@ -193,8 +229,8 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = {
     },
     "LingBotVideoPipeline": {
         "extra_body_params": LINGBOT_VIDEO_EXTRA_BODY_PARAMS,
-        # Shared T2I/I2V envelopes select the output modality. LingBot's
-        # pipeline owns model-specific validation and normalization.
+        "text_to_image_prompt_builder": build_lingbot_text_to_image_prompt,
+        "image_to_video_prompt_builder": build_lingbot_image_to_video_prompt,
     },
     **{
         model_class_name: {
@@ -227,7 +263,7 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = {
 }
 
 # Multi-stage discovery reports the top-level wrapper rather than its DiT
-# submodule, so both names must resolve to the same request builders.
+# submodule, so both names must resolve to the same request adapters.
 _EXTRA_SPECS["MammothModa2ForConditionalGeneration"] = _EXTRA_SPECS["MammothModa2DiTPipeline"]
 _EXTRA_SPECS["Mammothmoda2Model"] = _EXTRA_SPECS["MammothModa2DiTPipeline"]
 
@@ -269,21 +305,18 @@ def should_init_extra_args_for_non_diffusion_stages(model_class_name: str | None
 
 def build_text_to_image_prompt(
     model_class_name: str | None,
-    prompt: dict[str, Any],
+    prompt: str,
+    negative_prompt: str | None,
     height: int | None = None,
     width: int | None = None,
 ) -> dict[str, Any]:
-    """Build a model-specific T2I prompt from an example-owned envelope."""
     spec = _get_spec(model_class_name)
-    builder: TextToImagePromptBuilder | None = spec.get("text_to_image_prompt_builder") if spec else None
-    if builder is None:
-        return prompt
-    return builder(
-        prompt=str(prompt["prompt"]),
-        negative_prompt=prompt.get("negative_prompt"),
-        height=height,
-        width=width,
+    builder: TextToImagePromptBuilder = (
+        spec.get("text_to_image_prompt_builder", default_text_to_image_prompt)
+        if spec is not None
+        else default_text_to_image_prompt
     )
+    return builder(prompt, negative_prompt, height, width)
 
 
 def build_image_to_image_prompt(
@@ -305,24 +338,17 @@ def build_image_to_image_prompt(
 
 def build_image_to_video_prompt(
     model_class_name: str | None,
-    prompt: dict[str, Any],
+    prompt: str,
+    negative_prompt: str | None,
+    media_inputs: Mapping[str, Any],
     height: int | None = None,
     width: int | None = None,
     num_frames: int | None = None,
 ) -> dict[str, Any]:
-    """Build a model-specific I2V prompt from an example-owned envelope."""
     spec = _get_spec(model_class_name)
-    builder: ImageToVideoPromptBuilder | None = spec.get("image_to_video_prompt_builder") if spec else None
-    if builder is None:
-        return prompt
-    media_inputs = prompt.get("multi_modal_data") or {}
-    if not isinstance(media_inputs, Mapping):
-        raise TypeError("Canonical I2V prompt multi_modal_data must be a mapping.")
-    return builder(
-        prompt=str(prompt["prompt"]),
-        negative_prompt=prompt.get("negative_prompt"),
-        media_inputs=media_inputs,
-        height=height,
-        width=width,
-        num_frames=num_frames,
+    builder: ImageToVideoPromptBuilder = (
+        spec.get("image_to_video_prompt_builder", default_image_to_video_prompt)
+        if spec is not None
+        else default_image_to_video_prompt
     )
+    return builder(prompt, negative_prompt, media_inputs, height, width, num_frames)

vllm_omni/platforms/musa/platform.py

@@ -3,6 +3,8 @@
 
 
 import torch
+from vllm.config import VllmConfig
+from vllm.config.kernel import IrOpPriorityConfig
 from vllm.logger import init_logger
 from vllm.platforms.interface import DeviceCapability
 from vllm_musa.platform import MUSAPlatformBase
@@ -30,10 +32,6 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase):
     def get_omni_generation_worker_cls(cls) -> str:
         return "vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker"
 
-    @classmethod
-    def get_default_stage_config_path(cls) -> str:
-        return "vllm_omni/deploy"
-
     @classmethod
     def has_flash_attn_package(cls) -> bool:
         from vllm_omni.diffusion.attention.backends.utils.fa import is_flash_attn_installed
@@ -115,6 +113,10 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase):
         """MUSA supports torch.compile with inductor backend."""
         return True
 
+    @classmethod
+    def get_default_stage_config_path(cls) -> str:
+        return "vllm_omni/deploy"
+
     @classmethod
     def supports_talker_mtp_graph_capture(cls) -> bool:
         """MUSA keeps Qwen3 talker MTP outside its dedicated FULL graph."""
@@ -142,7 +144,7 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase):
     @classmethod
     def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:
         """Get the compute capability of the MUSA device."""
-        major, minor = torch.cuda.get_device_capability(device_id)
+        major, minor = torch.musa.get_device_capability(device_id)
         return DeviceCapability(major=major, minor=minor)
 
     @classmethod
@@ -183,3 +185,26 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase):
     @classmethod
     def get_device_name(cls, device_id: int = 0) -> str:
         return torch.musa.get_device_name(device_id)
+
+    @classmethod
+    def set_device_control_env_var(cls, devices: str | int | None) -> None:
+        import os
+
+        os.environ["MUSA_VISIBLE_DEVICES"] = devices
+
+    @classmethod
+    def unset_device_control_env_var(cls) -> None:
+        import os
+
+        os.environ.pop("MUSA_VISIBLE_DEVICES", None)
+
+    @classmethod
+    def get_default_ir_op_priority(cls, vllm_config: VllmConfig) -> IrOpPriorityConfig:
+        """Prefer native ops while compiling, otherwise prefer vLLM kernels."""
+        from vllm.config.compilation import CompilationMode
+
+        cc = vllm_config.compilation_config
+        using_inductor = cc.backend == "inductor" and cc.mode != CompilationMode.NONE
+        default = ["native"] if using_inductor else ["vllm_c", "native"]
+
+        return IrOpPriorityConfig.with_default(default)