Compare commits

..
Author SHA1 Message Date
Hongsheng Liu 0db924bc54 [Misc] Keep serving changes out of example refactor
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-12 00:22:41 +08:00
Hongsheng Liu 39e1c70aea [Misc] Keep model-specific curl recipes
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-12 00:14:40 +08:00
Hongsheng Liu b9572edb81 [Misc] Apply shared prompt builder review feedback
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-11 23:59:14 +08:00
Hongsheng Liu a4516ea5db [Misc] Address prompt adapter review feedback
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-11 23:41:44 +08:00
Hongsheng Liu 280c4f56f6 [Misc] Build canonical image task prompts
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-11 17:42:06 +08:00
47 changed files with 351 additions and 1682 deletions

No files matched your search

+5 -5
View File
@@ -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.5.0
- export CACHE_DIT_VERSION=1.3.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.5.0
- export CACHE_DIT_VERSION=1.3.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.5.0
- export CACHE_DIT_VERSION=1.3.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.5.0
- export CACHE_DIT_VERSION=1.3.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.5.0
- export CACHE_DIT_VERSION=1.3.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
+1 -1
View File
@@ -115,7 +115,7 @@ steps:
DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export CACHE_DIT_VERSION=1.5.0
- export CACHE_DIT_VERSION=1.3.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
+7 -105
View File
@@ -1,18 +1,14 @@
# 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 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.
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.
## Overview
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 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 Adapter Format
### PEFT (Parameter-Efficient Fine-Tuning) format (default)
A typical PEFT-format LoRA adapter directory structure:
LoRA adapters must be in **PEFT (Parameter-Efficient Fine-Tuning)** format. A typical LoRA adapter directory structure:
```
lora_adapter/
@@ -29,20 +25,19 @@ The `adapter_config.json` file contains metadata about the LoRA adapter, includi
### Offline Inference
#### PEFT backend: pre-loaded LoRA
#### Pre-loaded LoRA
Load a PEFT-format LoRA adapter at initialization. The adapter is pre-loaded into the cache and can be activated per request:
Load a LoRA adapter at initialization. This adapter is pre-loaded into the cache and can be activated by requests:
```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_backend="peft", # default, can be omitted
lora_path=lora_path
)
lora_request = LoRARequest(
@@ -58,98 +53,6 @@ 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.
@@ -238,7 +141,6 @@ 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
@@ -88,9 +88,8 @@ 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 or checkpoint file |
| `--lora-path` | str | — | Path to PEFT LoRA adapter folder |
| `--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:**
@@ -125,9 +125,6 @@ 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
@@ -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,
build_image_to_video_prompt as build_model_image_to_video_prompt,
)
from vllm_omni.model_extras import (
get_extra_body_params,
get_model_class_name,
)
@@ -85,6 +85,22 @@ 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=(
@@ -296,30 +312,6 @@ 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"),
@@ -513,12 +505,6 @@ 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:
@@ -549,36 +535,22 @@ 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,
@@ -589,7 +561,7 @@ def main():
num_inference_steps=num_inference_steps,
num_frames=num_frames,
frame_rate=frame_rate,
extra_args=sampling_extra_args,
extra_args={"sample_solver": args.sample_solver},
)
if flow_shift is not None:
sampling_params.extra_args["flow_shift"] = flow_shift
@@ -57,8 +57,9 @@ python examples/offline_inference/text_to_image/text_to_image.py \
--output lingbot_t2i.png
```
The LingBot prompt builder selects the `image` output modality and requests one
frame. Passing an input image is not supported in LingBot T2I mode.
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.
## Quick Start
@@ -109,9 +110,8 @@ 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 or checkpoint file |
| `--lora-path` | str | — | Path to PEFT LoRA adapter folder |
| `--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). |
@@ -21,7 +21,9 @@ 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,
build_text_to_image_prompt as build_model_text_to_image_prompt,
)
from vllm_omni.model_extras import (
get_extra_body_params,
get_model_class_name,
should_init_extra_args_for_non_diffusion_stages,
@@ -56,6 +58,17 @@ 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 = []
@@ -260,7 +273,6 @@ 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.",
)
@@ -270,15 +282,6 @@ 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,
@@ -414,12 +417,8 @@ def main():
# Prepare LoRA kwargs for Omni initialization
lora_args: dict[str, Any] = {}
if 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}")
lora_args["lora_path"] = args.lora_path
print(f"Using LoRA from: {args.lora_path}")
# Build quantization kwargs: use quantization_config dict when
# ignored_layers is specified so the list flows through OmniDiffusionConfig
@@ -508,24 +507,23 @@ def main():
# Build LoRA request when --lora-path is set
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)
if args.lora_path:
lora_request_id = stable_lora_int_id(args.lora_path)
lora_request = LoRARequest(
lora_name=Path(lora_path).stem,
lora_name=Path(args.lora_path).stem,
lora_int_id=lora_request_id,
lora_path=lora_path,
lora_path=args.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,
)
@@ -218,9 +218,6 @@ 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
@@ -14,8 +14,6 @@ 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
@@ -336,29 +334,6 @@ 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",
@@ -465,12 +440,6 @@ 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).
@@ -500,19 +469,6 @@ 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
@@ -520,11 +476,6 @@ 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,
@@ -532,7 +483,6 @@ 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
@@ -1,6 +1,9 @@
# Image-To-Video
This example demonstrates how to deploy the Wan2.2 image-to-video model for online video generation using vLLM-Omni.
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.
## Start Server
@@ -20,9 +23,7 @@ bash run_server.sh
The script allows overriding:
- `MODEL` (default: `Wan-AI/Wan2.2-I2V-A14B-Diffusers`)
- `PORT` (default: `8091`)
- `BOUNDARY_RATIO` (default: `0.875`)
- `FLOW_SHIFT` (default: `12.0`)
- `PORT` (default: `8099`)
- `CACHE_BACKEND` (default: `none`)
- `ENABLE_CACHE_DIT_SUMMARY` (default: `0`)
@@ -1,12 +1,12 @@
#!/bin/bash
# Wan2.2 image-to-video server start script
# Shared 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 Wan2.2 I2V server..."
echo "Starting image-to-video server..."
echo "Model: $MODEL"
echo "Port: $PORT"
echo "Cache backend: $CACHE_BACKEND"
@@ -1,6 +1,9 @@
# Text-To-Image
This example demonstrates how to deploy Qwen-Image model for online image generation service using vLLM-Omni.
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.
## Start Server
@@ -115,23 +118,7 @@ python gradio_demo.py
## LoRA
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`.
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).
### Using Python Client with LoRA
@@ -242,8 +229,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 | 50 | Number of denoising steps |
| `true_cfg_scale` | float | 4.0 | Qwen-Image CFG scale |
| `num_inference_steps` | int | model | Number of denoising steps |
| `true_cfg_scale` | float | model | True CFG scale when supported |
| `seed` | int | None | Random seed (reproducible) |
| `negative_prompt` | str | None | Negative prompt |
| `num_outputs_per_prompt` | int | 1 | Number of images to generate |
@@ -1,6 +1,6 @@
#!/usr/bin/env python3
"""
Qwen-Image Gradio Demo for online serving.
Shared text-to-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="Qwen-Image Demo") as demo:
gr.Markdown("# Qwen-Image Online Generation")
gr.Markdown("Generate images using Qwen-Image model")
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.Row():
with gr.Column(scale=1):
@@ -115,7 +115,6 @@ 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,
)
@@ -173,7 +172,7 @@ def create_demo(server_url: str):
def main():
parser = argparse.ArgumentParser(description="Qwen-Image Gradio Demo")
parser = argparse.ArgumentParser(description="Shared text-to-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")
@@ -1,6 +1,6 @@
#!/usr/bin/env python3
"""
Qwen-Image OpenAI-compatible image generation client.
Shared OpenAI-compatible text-to-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: Qwen-Image CFG scale
true_cfg_scale: True CFG scale for models that support it
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:
if negative_prompt is not None:
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="Qwen-Image chat client")
parser = argparse.ArgumentParser(description="Shared text-to-image client")
parser.add_argument("--prompt", "-p", default="a cup of coffee on the table", help="Text prompt")
parser.add_argument("--output", "-o", default="qwen_image_output.png", help="Output file")
parser.add_argument("--output", "-o", default="text_to_image_output.png", help="Output file")
parser.add_argument("--server", "-s", default="http://localhost:8091", help="Server URL")
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("--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("--negative", help="Negative prompt")
parser.add_argument("--lora-path", default=None, help="Server-local LoRA adapter folder (PEFT format)")
@@ -1,10 +1,10 @@
#!/bin/bash
# Qwen-Image online serving startup script
# Shared text-to-image online serving startup script
MODEL="${MODEL:-Qwen/Qwen-Image}"
PORT="${PORT:-8091}"
echo "Starting Qwen-Image server..."
echo "Starting text-to-image server..."
echo "Model: $MODEL"
echo "Port: $PORT"
+1 -1
View File
@@ -82,7 +82,7 @@ example and the DiT sections below.
- Optional environment variables:
```bash
export CACHE_DIT_VERSION=1.5.0
export CACHE_DIT_VERSION=1.3.0
```
HunyuanImage-3.0 sets the diffusion attention backend to `TORCH_SDPA`
+1 -1
View File
@@ -16,7 +16,7 @@ diffusers==0.38.0
safetensors>=0.8.0
accelerate==1.12.0
soundfile>=0.13.1
cache-dit==1.5.0
cache-dit==1.3.0
tqdm>=4.66.0
torchsde>=0.2.6
openai-whisper>=20250625
-438
View File
@@ -1,438 +0,0 @@
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
@@ -19,11 +19,38 @@ 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"),
# Talker context is 4096 (tts_config.max_position_embeddings); KV sizing
# is left automatic so duplex matches the simplex deploy profiles.
# 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.
"stages": {
1: {"max_model_len": 4096},
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},
},
# 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"
@@ -44,6 +71,17 @@ 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:
@@ -74,11 +112,9 @@ def validated_soft_interrupt_wav() -> Path:
def resolve_ref_audio(model_prefix: str) -> Path:
# ``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():
if model_prefix:
model_root = Path(model_prefix) / MODEL
else:
model_root = Path(snapshot_download(MODEL, local_files_only=True))
ref_audio = model_root / REF_AUDIO_RELATIVE_PATH
if not ref_audio.is_file():
@@ -14,15 +14,47 @@ 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"
BASE_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"})
PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"}, num_cards=2)
def _get_diffusion_feature_cases(model: str):
return [
pytest.param(
OmniServerParams(model=model),
id="base",
marks=BASE_FEATURE_MARKS,
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",
),
]
+15 -23
View File
@@ -1,7 +1,10 @@
"""E2E online tests for MiniCPM-o 4.5 multimodal input and audio/text output.
"""
E2E Online tests for MiniCPM-o 4.5 model with multimodal input and audio / text output.
Exercises async chunk streaming (``--async-chunk``) across separate Thinker,
Talker, and Code2Wav stages.
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.
"""
import os
@@ -11,17 +14,12 @@ 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, modify_stage_config
from tests.helpers.stage_config import get_deploy_config_path
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
_MODEL = "openbmb/MiniCPM-o-4_5"
_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}}
},
)
_CI_DEPLOY = get_deploy_config_path("minicpmo_4_5.yaml")
test_params = [
pytest.param(
@@ -29,10 +27,12 @@ test_params = [
model=_MODEL,
stage_config_path=_CI_DEPLOY,
use_stage_cli=False,
server_args=["--trust-remote-code", "--async-chunk"],
server_args=[
"--trust-remote-code",
],
),
id="async_chunk",
),
id="default",
)
]
@@ -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 Talker TTS region detection and the Code2Wav stage.
This exercises the talker TTS region detection and in-process token2wav vocoder.
Deploy Setting: default single GPU
Input Modal: text
Output Modal: text + audio
@@ -132,21 +132,13 @@ 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(),
content_text=get_prompt("mix"),
)
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())
@@ -11,6 +11,7 @@ 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,
@@ -106,7 +107,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", SERVER_PARAMS, indirect=True)
@pytest.mark.parametrize("omni_server", CORE_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(
@@ -133,12 +134,8 @@ 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
assert result["ok"] is True, json.dumps(result, ensure_ascii=False, indent=2)
assert result["audio_delta_count"] > 0
assert result["done_count"] == 2
assert result["error_count"] == 0
@@ -163,7 +160,7 @@ def test_duplex_two_sessions_resume_and_takeover(omni_server, model_prefix: str,
)
)
)
assert result["ok"] is True
assert result["ok"] is True, json.dumps(result, ensure_ascii=False, indent=2)
assert result["session_count"] == 2
assert result["resume"]["ok"] is True
assert result["takeover"]["ok"] is True
@@ -2,9 +2,6 @@
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
@@ -23,27 +20,18 @@ 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", "--no-async-chunk"],
server_args=[
"--trust-remote-code",
],
),
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",
),
)
]
@@ -284,37 +284,6 @@ 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()
@@ -0,0 +1,47 @@
# 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",
}
+38 -107
View File
@@ -3,6 +3,8 @@
from __future__ import annotations
from typing import Any
import pytest
from PIL import Image
@@ -135,67 +137,6 @@ 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"])
@@ -342,8 +283,11 @@ 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.",
negative_prompt="ugly, blurry",
{
"prompt": "Please draw a cute cat.",
"modalities": ["image"],
"negative_prompt": "ugly, blurry",
},
height=512,
width=768,
) == {
@@ -360,8 +304,7 @@ 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.",
negative_prompt=None,
{"prompt": "Draw a poster.", "modalities": ["image"]},
) == {
"prompt": "Draw a poster.",
"modalities": ["image"],
@@ -391,28 +334,6 @@ 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:
@@ -505,8 +426,11 @@ 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",
negative_prompt="blurry",
{
"prompt": "a cat",
"modalities": ["image"],
"negative_prompt": "blurry",
},
height=512,
width=768,
) == {
@@ -543,14 +467,17 @@ def test_bagel_image_to_image_prompt_builder() -> None:
@pytest.mark.core_model
@pytest.mark.cpu
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"}
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
)
@pytest.mark.core_model
@@ -569,12 +496,14 @@ def test_unknown_pipeline_uses_default_image_to_image_prompt() -> None:
}
def _build_vace_prompt(media_inputs: dict[str, object], *, num_frames: int = 5) -> dict:
def _build_vace_prompt(media_inputs: dict[str, Any], *, num_frames: int = 5) -> dict:
return build_image_to_video_prompt(
"WanVACEPipeline",
prompt="a bird flying",
negative_prompt=None,
media_inputs=media_inputs,
{
"prompt": "a bird flying",
"modalities": ["video"],
"multi_modal_data": media_inputs,
},
height=16,
width=320,
num_frames=num_frames,
@@ -597,7 +526,7 @@ def _build_vace_prompt(media_inputs: dict[str, object], *, num_frames: int = 5)
],
ids=["i2v", "v2lf", "flf2v", "inpaint", "r2v"],
)
def test_vace_image_to_video_prompt_builder(media_inputs: dict[str, object]) -> None:
def test_vace_image_to_video_prompt_builder(media_inputs: dict[str, Any]) -> None:
result = _build_vace_prompt(media_inputs)
mmd = result["multi_modal_data"]
if "reference_images" in mmd:
@@ -616,7 +545,7 @@ def test_vace_image_to_video_prompt_builder(media_inputs: dict[str, object]) ->
({"control_image": Image.new("RGB", (320, 16))}, "Unsupported VACE media input"),
],
)
def test_vace_rejects_invalid_media_combinations(media_inputs: dict[str, object], message: str) -> None:
def test_vace_rejects_invalid_media_combinations(media_inputs: dict[str, Any], message: str) -> None:
with pytest.raises(ValueError, match=message):
_build_vace_prompt(media_inputs)
@@ -664,8 +593,7 @@ def test_mammothmoda2_extra_registry_declares_request_and_response_params() -> N
wrapper_prompt = build_text_to_image_prompt(
"MammothModa2ForConditionalGeneration",
prompt="a cat",
negative_prompt=None,
{"prompt": "a cat", "modalities": ["image"]},
height=256,
width=256,
)
@@ -679,8 +607,11 @@ 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",
negative_prompt="blurry",
{
"prompt": "a cat",
"modalities": ["image"],
"negative_prompt": "blurry",
},
height=512,
width=768,
) == {
@@ -37,19 +37,6 @@ 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
@@ -70,24 +57,23 @@ 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_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,
"text_encoder": shrink_clip_text_encoder,
"text_encoder_2": shrink_t5_text_encoder,
"transformer": partial(_shrink_dit_rope_config, num_single_layers=2, default_axes_dims_rope=[16, 56, 56]),
},
)
@@ -59,16 +59,6 @@ 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,
@@ -48,6 +48,7 @@ EXCLUDED_MODELS = [
"InternVLAA1Pipeline",
"LongCatImageEditPipeline",
"StableDiffusion3Pipeline",
"FluxKontextPipeline",
"HunyuanImage3ForCausalMM",
"ErnieImagePipeline",
"NextStep11Pipeline",
+1 -2
View File
@@ -563,9 +563,8 @@ 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 | list[str] | None = None
lora_path: 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
-3
View File
@@ -385,9 +385,6 @@ 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
+2 -4
View File
@@ -27,7 +27,6 @@ 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
@@ -683,9 +682,8 @@ class OmniDiffusionConfig:
# pipeline_config: PipelineConfig = field(default_factory=PipelineConfig, repr=False)
# LoRA parameters
lora_path: str | list[str] | None = None
lora_scale: float | list[float] = 1.0
lora_backend: LoRABackend = LoRABackend.PEFT # available choices: ["peft", "distill"]
lora_path: str | None = None
lora_scale: float = 1.0
max_cpu_loras: int | None = None
output_type: str = "pil"
-461
View File
@@ -1,461 +0,0 @@
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]
-6
View File
@@ -4,7 +4,6 @@
import math
import time
from collections import OrderedDict
from enum import Enum
from typing import get_args
import torch
@@ -34,11 +33,6 @@ 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.
@@ -25,7 +25,6 @@ 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
@@ -265,11 +264,7 @@ def apply_rotary_emb_qwen(
class QwenImagePipeline(
nn.Module,
QwenImageCFGParallelMixin,
DiffusionPipelineProfilerMixin,
SupportsComponentDiscovery,
QwenImageLoraLoaderMixin,
nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery
):
supports_request_batch = True
_dit_modules: ClassVar[list[str]] = ["transformer"]
@@ -25,7 +25,6 @@ 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
@@ -286,7 +285,6 @@ class Wan22Pipeline(
DenoiseProgressMixin,
DiffusionPipelineProfilerMixin,
SupportsComponentDiscovery,
WanLoraLoaderMixin,
):
_dit_modules: ClassVar[list[str]] = ["transformer", "transformer_2"]
_encoder_modules: ClassVar[list[str]] = ["text_encoder"]
@@ -25,7 +25,6 @@ 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
@@ -148,7 +147,6 @@ class Wan22I2VPipeline(
DiffusionPipelineProfilerMixin,
DenoiseProgressMixin,
SupportsComponentDiscovery,
WanLoraLoaderMixin,
):
"""
Wan2.2 Image-to-Video Pipeline.
+3 -4
View File
@@ -70,10 +70,9 @@ def find_module_with_attr(model, attr_name="transformer"):
return None
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)
def get_transformer_from_pipeline(pipeline: Any):
pipe = find_module_with_attr(pipeline, attr_name="transformer")
if pipe is not None:
return getattr(pipe, transformer_name)
return pipe.transformer
return None
+9 -26
View File
@@ -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, LoRABackend
from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
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,31 +388,14 @@ class DiffusionWorker:
"""Initialize the LoRA manager for this worker."""
if self.model_runner.pipeline is None:
return
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__}")
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,
)
def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None:
"""Start or stop profiling for this GPU worker.
-3
View File
@@ -471,9 +471,6 @@ 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
-6
View File
@@ -1048,9 +1048,6 @@ 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"]),
@@ -1254,9 +1251,6 @@ 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
-33
View File
@@ -164,13 +164,6 @@ 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:
@@ -434,32 +427,6 @@ 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"],
-26
View File
@@ -3,32 +3,6 @@
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",
-45
View File
@@ -3,11 +3,6 @@
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",
@@ -24,43 +19,3 @@ 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
+29 -55
View File
@@ -28,21 +28,12 @@ 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,
@@ -147,18 +138,6 @@ 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,
@@ -175,20 +154,6 @@ 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,
@@ -208,12 +173,11 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = {
"Cosmos3OmniDiffusersPipeline": {
"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,
# The shared T2I example already supplies modalities=["image"].
},
"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,
@@ -229,8 +193,8 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = {
},
"LingBotVideoPipeline": {
"extra_body_params": LINGBOT_VIDEO_EXTRA_BODY_PARAMS,
"text_to_image_prompt_builder": build_lingbot_text_to_image_prompt,
"image_to_video_prompt_builder": build_lingbot_image_to_video_prompt,
# Shared T2I/I2V envelopes select the output modality. LingBot's
# pipeline owns model-specific validation and normalization.
},
**{
model_class_name: {
@@ -263,7 +227,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 adapters.
# submodule, so both names must resolve to the same request builders.
_EXTRA_SPECS["MammothModa2ForConditionalGeneration"] = _EXTRA_SPECS["MammothModa2DiTPipeline"]
_EXTRA_SPECS["Mammothmoda2Model"] = _EXTRA_SPECS["MammothModa2DiTPipeline"]
@@ -305,18 +269,21 @@ 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: str,
negative_prompt: str | None,
prompt: dict[str, Any],
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 = (
spec.get("text_to_image_prompt_builder", default_text_to_image_prompt)
if spec is not None
else default_text_to_image_prompt
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,
)
return builder(prompt, negative_prompt, height, width)
def build_image_to_image_prompt(
@@ -338,17 +305,24 @@ def build_image_to_image_prompt(
def build_image_to_video_prompt(
model_class_name: str | None,
prompt: str,
negative_prompt: str | None,
media_inputs: Mapping[str, Any],
prompt: dict[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 = (
spec.get("image_to_video_prompt_builder", default_image_to_video_prompt)
if spec is not None
else default_image_to_video_prompt
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,
)
return builder(prompt, negative_prompt, media_inputs, height, width, num_frames)
+5 -30
View File
@@ -3,8 +3,6 @@
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
@@ -32,6 +30,10 @@ 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
@@ -113,10 +115,6 @@ 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."""
@@ -144,7 +142,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.musa.get_device_capability(device_id)
major, minor = torch.cuda.get_device_capability(device_id)
return DeviceCapability(major=major, minor=minor)
@classmethod
@@ -185,26 +183,3 @@ 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)