[Refactor]Base class for dit pipelines with unified parameter declaration[1/N] (#4225)
Signed-off-by: princepride <wangzhipeng628@gmail.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
26 files changed
+1026
-1586
No files matched your search
@@ -1,391 +0,0 @@
|
||||
# BAGEL-7B-MoT
|
||||
|
||||
## Setup
|
||||
|
||||
Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
|
||||
|
||||
## Architecture
|
||||
|
||||
BAGEL-7B-MoT is a Mixture-of-Transformers (MoT) model supporting both image generation and understanding. It offers two deployment topologies:
|
||||
|
||||
| Topology | Stages | Description |
|
||||
| :------- | :----- | :---------- |
|
||||
| **Two-stage** (default) | Stage 0 (Thinker, AR) + Stage 1 (DiT, Diffusion) | Thinker handles text/understanding via vLLM AR engine; DiT handles image generation. KV cache is transferred between stages. |
|
||||
| **Single-stage** | Stage 0 (DiT, Diffusion) only | The DiT stage contains a full LLM, ViT, VAE, and tokenizer internally. All modalities are handled within a single diffusion process. |
|
||||
|
||||
Both topologies support all four modalities: `text2img`, `img2img`, `img2text`, `text2text`.
|
||||
|
||||
## Quick Start
|
||||
|
||||
```bash
|
||||
cd examples/offline_inference/bagel
|
||||
|
||||
# Default two-stage mode (auto-detected)
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2img \
|
||||
--prompts "A cute cat"
|
||||
|
||||
# Single-stage mode
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2img \
|
||||
--prompts "A cute cat" \
|
||||
--deploy-config vllm_omni/deploy/bagel_single_stage.yaml
|
||||
```
|
||||
|
||||
> **Note**: These examples work with the default configuration on an **NVIDIA A100 (80GB)**. For dual-GPU setups, modify the deploy YAML to distribute stages across devices.
|
||||
|
||||
## Modality Control
|
||||
|
||||
Control the mode using the `--modality` argument:
|
||||
|
||||
| Modality | Input | Output | Description |
|
||||
| :------- | :---- | :----- | :---------- |
|
||||
| `text2img` | Text | Image | Generate images from text prompts |
|
||||
| `img2img` | Image + Text | Image | Transform images using text guidance |
|
||||
| `img2text` | Image + Text | Text | Generate text descriptions from images |
|
||||
| `text2text` | Text | Text | Pure text generation (language model mode) |
|
||||
|
||||
### Text to Image (text2img)
|
||||
|
||||
```bash
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2img \
|
||||
--prompts "A cute cat" \
|
||||
--steps 50
|
||||
```
|
||||
|
||||
### Image to Image (img2img)
|
||||
|
||||
```bash
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality img2img \
|
||||
--image-path /path/to/image.jpg \
|
||||
--prompts "Let the woman wear a blue dress"
|
||||
```
|
||||
|
||||
### Image to Text (img2text)
|
||||
|
||||
```bash
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality img2text \
|
||||
--image-path /path/to/image.jpg \
|
||||
--prompts "Describe this image in detail"
|
||||
```
|
||||
|
||||
### Text to Text (text2text)
|
||||
|
||||
```bash
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2text \
|
||||
--prompts "What is the capital of France?"
|
||||
|
||||
# Load prompts from a text file (one prompt per line):
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2text \
|
||||
--txt-prompts /path/to/prompts.txt
|
||||
```
|
||||
|
||||
## Think Mode
|
||||
|
||||
Think mode enables the model to generate `<think>...</think>` planning/reasoning tokens before producing the final output. This improves generation quality for complex prompts.
|
||||
|
||||
- **Two-stage**: The Thinker (AR) stage decodes think tokens, then transfers the augmented KV cache to the DiT stage for image generation.
|
||||
- **Single-stage**: The DiT's internal LLM generates think tokens in-place before proceeding to denoise.
|
||||
|
||||
```bash
|
||||
# Think + text2img: plan before generating
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2img \
|
||||
--prompts "A futuristic city with flying cars" \
|
||||
--think \
|
||||
--max-think-tokens 1000
|
||||
|
||||
# Think + img2img: reason about the edit
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality img2img \
|
||||
--image-path /path/to/image.jpg \
|
||||
--prompts "Make it look like a watercolor painting" \
|
||||
--think
|
||||
|
||||
# Think + img2text: reason before describing
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality img2text \
|
||||
--image-path /path/to/image.jpg \
|
||||
--prompts "What is happening in this image?" \
|
||||
--think
|
||||
|
||||
# Think + text2text: chain-of-thought reasoning
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2text \
|
||||
--prompts "Solve: 23 * 47" \
|
||||
--think
|
||||
```
|
||||
|
||||
Think mode parameters:
|
||||
|
||||
| Argument | Default | Description |
|
||||
| :------- | :------ | :---------- |
|
||||
| `--think` | `False` | Enable thinking mode |
|
||||
| `--max-think-tokens` | `1000` | Maximum tokens for think generation |
|
||||
| `--do-sample` | `False` | Enable sampling (vs. greedy) for text generation |
|
||||
| `--text-temperature` | `0.3` | Temperature for text generation sampling |
|
||||
|
||||
## Classifier-Free Guidance (CFG)
|
||||
|
||||
CFG controls the trade-off between prompt fidelity and diversity. These parameters apply to image generation modalities (`text2img`, `img2img`).
|
||||
|
||||
```bash
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2img \
|
||||
--prompts "A photorealistic portrait" \
|
||||
--cfg-text-scale 6.0 \
|
||||
--cfg-img-scale 2.0 \
|
||||
--negative-prompt "blurry, low quality, distorted" \
|
||||
--cfg-interval 0.4 1.0 \
|
||||
--cfg-renorm-type global \
|
||||
--cfg-renorm-min 0.0
|
||||
```
|
||||
|
||||
| Argument | Default | Description |
|
||||
| :------- | :------ | :---------- |
|
||||
| `--cfg-text-scale` | `4.0` | Text CFG scale (higher = more prompt-adherent) |
|
||||
| `--cfg-img-scale` | `1.5` | Image CFG scale (for img2img) |
|
||||
| `--negative-prompt` | `None` | Negative prompt for CFG conditioning |
|
||||
| `--cfg-interval` | pipeline default | CFG active interval `[start, end]` as fractions of total timesteps |
|
||||
| `--cfg-renorm-type` | `None` | Renormalization type: `global`, `text_channel`, `channel` |
|
||||
| `--cfg-renorm-min` | `None` | Minimum renormalization value |
|
||||
| `--cfg-parallel-size` | `1` | CFG parallel size: `1` = batched (single GPU), `2` = 2-branch parallel, `3` = full 3-GPU parallel |
|
||||
|
||||
## Deployment Topologies
|
||||
|
||||
### Two-Stage (Default)
|
||||
|
||||
The default topology auto-detected from the model. No extra flags needed.
|
||||
|
||||
```bash
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2img \
|
||||
--prompts "A cute cat"
|
||||
```
|
||||
|
||||
The pipeline is defined in [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml). Stage 0 (Thinker) and Stage 1 (DiT) share GPU 0 by default. For dual-GPU setups, customize the deploy YAML and set `devices: "1"` for stage 1.
|
||||
|
||||
### Single-Stage
|
||||
|
||||
Pass the single-stage deploy config via `--deploy-config`:
|
||||
|
||||
```bash
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2img \
|
||||
--prompts "A cute cat" \
|
||||
--deploy-config vllm_omni/deploy/bagel_single_stage.yaml
|
||||
```
|
||||
|
||||
See [`bagel_single_stage.yaml`](../../../vllm_omni/deploy/bagel_single_stage.yaml) for configuration details. The `pipeline: bagel_single_stage` field selects the single-stage topology from the pipeline registry.
|
||||
|
||||
### Tensor Parallelism (TP)
|
||||
|
||||
For larger models or multi-GPU environments, customize the deploy YAML (see [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml)) and set per-stage `tensor_parallel_size` and `devices`:
|
||||
|
||||
```yaml
|
||||
# Example: TP=2 on GPUs 0,1 for the Thinker stage
|
||||
stages:
|
||||
- stage_id: 0
|
||||
tensor_parallel_size: 2
|
||||
devices: "0,1"
|
||||
```
|
||||
|
||||
### VAE Patch Parallelism
|
||||
|
||||
[VAE Patch Parallelism](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism/vae_patch_parallel.html) splits Bagel VAE **decode/encode** tiles across multiple GPUs on the **DiT stage**, reducing **per-GPU peak memory during VAE decode**. Use it when high-resolution `text2img` or `img2img` hits VAE OOM or large decode spikes.
|
||||
|
||||
**Bagel-specific notes:**
|
||||
|
||||
- Implemented in `BagelPipeline` via `DistributedAutoEncoder` (DiT stage only).
|
||||
- **Single-stage** is the simplest path: one DiT process with TP + VAE patch parallel.
|
||||
- **Two-stage**: enable on **stage 1 (DiT)** only; stage 0 (Thinker) keeps encoder-only `VAEEncoder` and does not use VAE patch parallel.
|
||||
- You need a DiT `world_size` ≥ `vae_patch_parallel_size` (typically `tensor_parallel_size=2` on that stage). VAE PP reuses the DiT process group; it is not a standalone second-GPU VAE worker.
|
||||
|
||||
**Single-stage via deploy YAML** (recommended for `end2end.py`):
|
||||
|
||||
```yaml
|
||||
pipeline: bagel_single_stage
|
||||
async_chunk: false
|
||||
|
||||
stages:
|
||||
- stage_id: 0
|
||||
max_num_batched_tokens: 32768
|
||||
max_num_seqs: 1
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
enable_prefix_caching: false
|
||||
devices: "0,1"
|
||||
vae_use_tiling: true
|
||||
parallel_config:
|
||||
tensor_parallel_size: 2
|
||||
vae_patch_parallel_size: 2
|
||||
default_sampling_params:
|
||||
seed: 52
|
||||
```
|
||||
|
||||
```bash
|
||||
cd examples/offline_inference/bagel
|
||||
|
||||
CUDA_VISIBLE_DEVICES=0,1 python end2end.py \
|
||||
--model /path/to/BAGEL-7B-MoT \
|
||||
--deploy-config /path/to/bagel_single_stage_vae_pp.yaml \
|
||||
--modality text2img \
|
||||
--prompts "A cute cat" \
|
||||
--steps 10 \
|
||||
--output ./out_vae_pp
|
||||
```
|
||||
|
||||
**Single-stage via `Omni` kwargs** (same flags as online serving):
|
||||
|
||||
```python
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
omni = Omni(
|
||||
model="ByteDance-Seed/BAGEL-7B-MoT",
|
||||
deploy_config="vllm_omni/deploy/bagel_single_stage.yaml",
|
||||
tensor_parallel_size=2,
|
||||
vae_patch_parallel_size=2,
|
||||
vae_use_tiling=True,
|
||||
)
|
||||
# Then call omni.generate(...) as in end2end.py
|
||||
```
|
||||
|
||||
**Two-stage (VAE PP on DiT only):**
|
||||
|
||||
```yaml
|
||||
stages:
|
||||
- stage_id: 0
|
||||
devices: "0"
|
||||
# AR Thinker — no vae_patch_parallel here
|
||||
|
||||
- stage_id: 1
|
||||
devices: "0,1"
|
||||
vae_use_tiling: true
|
||||
parallel_config:
|
||||
tensor_parallel_size: 2
|
||||
vae_patch_parallel_size: 2
|
||||
```
|
||||
|
||||
```bash
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--deploy-config /path/to/bagel_vae_pp.yaml \
|
||||
--modality text2img \
|
||||
--prompts "A cute cat"
|
||||
```
|
||||
|
||||
**Startup log checks:**
|
||||
|
||||
```text
|
||||
INFO ... vae_patch_parallel_size=2 requires vae_use_tiling; automatically enabling it.
|
||||
```
|
||||
|
||||
| Setting | Role |
|
||||
| :------ | :--- |
|
||||
| `parallel_config.tensor_parallel_size` | DiT world size / TP (must be ≥ `vae_patch_parallel_size`) |
|
||||
| `parallel_config.vae_patch_parallel_size` | Number of ranks for distributed VAE tiles (`1` = off) |
|
||||
| `vae_use_tiling` | Enable spatial tiling (auto-enabled when `vae_patch_parallel_size > 1`) |
|
||||
|
||||
#### Hybrid Sharded Data Parallel (HSDP)
|
||||
|
||||
For larger Bagel deployments on multiple GPUs, you can enable HSDP (Hybrid Sharded Data Parallel) by modifying the stage configuration (for example, [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml)). HSDP shards transformer weights across GPUs to reduce per-GPU memory usage.
|
||||
|
||||
1. **Enable HSDP**: Set `use_hsdp: true`.
|
||||
2. **Set shard size**: Set `hsdp_shard_size` to the number of GPUs used for sharding (for example, `4`).
|
||||
3. **Set replicate size**: Usually keep `hsdp_replicate_size: 1` unless you want replicated HSDP groups.
|
||||
4. **Set devices**: Specify the comma-separated GPU IDs used by the diffusion stage (for example, `"0,1,2,3"`).
|
||||
|
||||
Example configuration for HSDP across 4 GPUs:
|
||||
```yaml
|
||||
- stage_id: 1
|
||||
devices: "0,1,2,3"
|
||||
parallel_config:
|
||||
use_hsdp: true
|
||||
hsdp_shard_size: 4
|
||||
hsdp_replicate_size: 1
|
||||
```
|
||||
|
||||
Then pass the custom deploy YAML:
|
||||
|
||||
```bash
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2img \
|
||||
--prompts "A cute cat" \
|
||||
--deploy-config /path/to/custom_bagel.yaml
|
||||
```
|
||||
|
||||
### FP8 Quantization
|
||||
|
||||
```bash
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2img \
|
||||
--prompts "A cute cat" \
|
||||
--quantization fp8
|
||||
```
|
||||
|
||||
## Command Line Reference
|
||||
|
||||
### Core Arguments
|
||||
|
||||
| Argument | Type | Default | Description |
|
||||
| :------- | :--- | :------ | :---------- |
|
||||
| `--model` | string | `ByteDance-Seed/BAGEL-7B-MoT` | Model path or HuggingFace name |
|
||||
| `--modality` | choice | `text2img` | `text2img`, `img2img`, `img2text`, `text2text` |
|
||||
| `--prompts` | list | `None` | Input text prompts |
|
||||
| `--txt-prompts` | string | `None` | Path to text file with one prompt per line |
|
||||
| `--image-path` | string | `None` | Input image path (required for `img2img`/`img2text`) |
|
||||
| `--output` | string | `.` | Output directory for saved images |
|
||||
| `--steps` | int | `50` | Number of diffusion inference steps |
|
||||
| `--seed` | int | `None` | Random seed for reproducibility |
|
||||
|
||||
### Think Mode Arguments
|
||||
|
||||
| Argument | Type | Default | Description |
|
||||
| :------- | :--- | :------ | :---------- |
|
||||
| `--think` | flag | `False` | Enable `<think>...</think>` planning/reasoning |
|
||||
| `--max-think-tokens` | int | `1000` | Maximum tokens for think generation |
|
||||
| `--do-sample` | flag | `False` | Use sampling instead of greedy decoding |
|
||||
| `--text-temperature` | float | `0.3` | Sampling temperature for text generation |
|
||||
|
||||
### CFG Arguments
|
||||
|
||||
| Argument | Type | Default | Description |
|
||||
| :------- | :--- | :------ | :---------- |
|
||||
| `--cfg-text-scale` | float | `4.0` | Text CFG guidance scale |
|
||||
| `--cfg-img-scale` | float | `1.5` | Image CFG guidance scale |
|
||||
| `--negative-prompt` | string | `None` | Negative prompt for CFG |
|
||||
| `--cfg-parallel-size` | int | `1` | CFG parallel GPU count (1, 2, or 3) |
|
||||
| `--cfg-interval` | float[2] | pipeline default | CFG active window `[start, end]` |
|
||||
| `--cfg-renorm-type` | string | `None` | `global`, `text_channel`, or `channel` |
|
||||
| `--cfg-renorm-min` | float | `None` | Minimum renormalization value |
|
||||
|
||||
### Engine Arguments
|
||||
|
||||
| Argument | Type | Default | Description |
|
||||
| :------- | :--- | :------ | :---------- |
|
||||
| `--deploy-config` | string | `None` | Path to deploy YAML (auto-detected if omitted) |
|
||||
| `--worker-backend` | choice | `process` | `process` or `ray` |
|
||||
| `--ray-address` | string | `None` | Ray cluster address |
|
||||
| `--quantization` | string | `None` | Quantization method (e.g. `fp8`) |
|
||||
| `--log-stats` | flag | `False` | Enable statistics logging |
|
||||
| `--init-timeout` | int | `300` | Initialization timeout (seconds) |
|
||||
| `--batch-timeout` | int | `5` | Batch timeout (seconds) |
|
||||
| `--enable-diffusion-pipeline-profiler` | flag | `False` | Profile diffusion stage durations |
|
||||
|
||||
## FAQ
|
||||
|
||||
- If you encounter OOM errors, try decreasing `max_model_len` or `gpu_memory_utilization` in the deploy YAML.
|
||||
|
||||
**Two-stage VRAM usage:**
|
||||
|
||||
| Stage | VRAM |
|
||||
| :---- | :--- |
|
||||
| Stage 0 (Thinker) | **15.04 GiB + KV Cache** |
|
||||
| Stage 1 (DiT) | **26.50 GiB** |
|
||||
| Total | **~42 GiB + KV Cache** |
|
||||
|
||||
**Single-stage VRAM usage:** The DiT loads the full model (~42 GiB) in one process.
|
||||
@@ -1,278 +0,0 @@
|
||||
import argparse
|
||||
import os
|
||||
|
||||
from vllm_omni.inputs.data import OmniPromptType
|
||||
from vllm_omni.model_executor.stage_input_processors.bagel import (
|
||||
GEN_THINK_SYSTEM_PROMPT,
|
||||
VLM_THINK_SYSTEM_PROMPT,
|
||||
)
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
default="ByteDance-Seed/BAGEL-7B-MoT",
|
||||
help="Path to merged model directory.",
|
||||
)
|
||||
parser.add_argument("--prompts", nargs="+", default=None, help="Input text prompts.")
|
||||
parser.add_argument(
|
||||
"--txt-prompts",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to a .txt file with one prompt per line (preferred).",
|
||||
)
|
||||
parser.add_argument("--prompt-type", default="text", choices=["text"])
|
||||
|
||||
parser.add_argument(
|
||||
"--modality",
|
||||
default="text2img",
|
||||
choices=["text2img", "img2img", "img2text", "text2text"],
|
||||
help="Modality mode to control stage execution.",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--image-path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to input image for img2img.",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
type=str,
|
||||
default=".",
|
||||
help="Output directory to save images.",
|
||||
)
|
||||
|
||||
# OmniLLM init args
|
||||
parser.add_argument("--log-stats", action="store_true", default=False)
|
||||
parser.add_argument("--init-sleep-seconds", type=int, default=20)
|
||||
parser.add_argument("--batch-timeout", type=int, default=5)
|
||||
parser.add_argument("--init-timeout", type=int, default=300)
|
||||
parser.add_argument("--shm-threshold-bytes", type=int, default=65536)
|
||||
parser.add_argument("--worker-backend", type=str, default="process", choices=["process", "ray"])
|
||||
parser.add_argument("--ray-address", type=str, default=None)
|
||||
parser.add_argument(
|
||||
"--deploy-config",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to deploy YAML. If unset, auto-loads vllm_omni/deploy/bagel.yaml based on the HF model_type.",
|
||||
)
|
||||
parser.add_argument("--steps", type=int, default=50, help="Number of inference steps.")
|
||||
|
||||
parser.add_argument("--cfg-text-scale", type=float, default=4.0, help="Text CFG scale (default: 4.0)")
|
||||
parser.add_argument("--cfg-img-scale", type=float, default=1.5, help="Image CFG scale (default: 1.5)")
|
||||
parser.add_argument(
|
||||
"--negative-prompt", type=str, default=None, help="Negative prompt for CFG (default: empty prompt)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--cfg-parallel-size",
|
||||
type=int,
|
||||
default=1,
|
||||
choices=[1, 2, 3],
|
||||
help="CFG parallel size: 1=batched (single GPU), 2=parallel with 2 branches (text CFG only), 3=parallel (3 GPUs).",
|
||||
)
|
||||
parser.add_argument("--seed", type=int, default=None, help="Random seed for generation.")
|
||||
parser.add_argument(
|
||||
"--cfg-interval",
|
||||
type=float,
|
||||
nargs=2,
|
||||
default=None,
|
||||
help="CFG interval [start, end] (default: pipeline default)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--cfg-renorm-type", type=str, default=None, help="CFG renorm type: global, text_channel, channel"
|
||||
)
|
||||
parser.add_argument("--cfg-renorm-min", type=float, default=None, help="CFG renorm min")
|
||||
parser.add_argument(
|
||||
"--enable-diffusion-pipeline-profiler",
|
||||
action="store_true",
|
||||
help="Enable diffusion pipeline profiler to display stage durations.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--quantization",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Quantization method (e.g. 'fp8').",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--think",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Enable thinking mode: AR stage decodes <think>...</think> planning tokens before image generation.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-think-tokens",
|
||||
type=int,
|
||||
default=1000,
|
||||
help="Maximum number of tokens for thinking text generation (default: 1000).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--do-sample",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Enable sampling for text generation (default: greedy).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--text-temperature",
|
||||
type=float,
|
||||
default=0.3,
|
||||
help="Temperature for text generation sampling (default: 0.3).",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
os.makedirs(args.output, exist_ok=True)
|
||||
model_name = args.model
|
||||
prompts: list[OmniPromptType] = []
|
||||
try:
|
||||
if getattr(args, "txt_prompts", None) and args.prompt_type == "text":
|
||||
with open(args.txt_prompts, encoding="utf-8") as f:
|
||||
lines = [ln.strip() for ln in f.readlines()]
|
||||
prompts = [ln for ln in lines if ln != ""]
|
||||
print(f"[Info] Loaded {len(prompts)} prompts from {args.txt_prompts}")
|
||||
else:
|
||||
prompts = args.prompts
|
||||
except Exception as e:
|
||||
print(f"[Error] Failed to load prompts: {e}")
|
||||
raise
|
||||
|
||||
if not prompts:
|
||||
prompts = ["A cute cat"]
|
||||
print(f"[Info] No prompts provided, using default: {prompts}")
|
||||
|
||||
from PIL import Image
|
||||
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
omni_kwargs = vars(args).copy()
|
||||
omni_kwargs.pop("model", None)
|
||||
deploy_config = args.deploy_config
|
||||
if args.think and deploy_config is None:
|
||||
deploy_config = "vllm_omni/deploy/bagel_think.yaml"
|
||||
print(f"[Info] Think mode enabled, using deploy config: {deploy_config}")
|
||||
if deploy_config:
|
||||
omni_kwargs["deploy_config"] = deploy_config
|
||||
|
||||
if args.quantization:
|
||||
omni_kwargs["quantization_config"] = args.quantization
|
||||
|
||||
omni = Omni(model=model_name, **omni_kwargs)
|
||||
|
||||
formatted_prompts = []
|
||||
for p in prompts:
|
||||
if args.modality == "img2img":
|
||||
if not args.image_path or not os.path.exists(args.image_path):
|
||||
raise ValueError(f"img2img requires --image-path pointing to an existing file, got: {args.image_path}")
|
||||
loaded_image = Image.open(args.image_path).convert("RGB")
|
||||
think_prefix = f"<|im_start|>{GEN_THINK_SYSTEM_PROMPT}<|im_end|>" if args.think else ""
|
||||
final_prompt_text = f"{think_prefix}<|fim_middle|><|im_start|>{p}<|im_end|>"
|
||||
prompt_dict = {
|
||||
"prompt": final_prompt_text,
|
||||
"multi_modal_data": {"img2img": loaded_image},
|
||||
"modalities": ["img2img"],
|
||||
}
|
||||
if args.negative_prompt is not None:
|
||||
prompt_dict["negative_prompt"] = args.negative_prompt
|
||||
formatted_prompts.append(prompt_dict)
|
||||
elif args.modality == "img2text":
|
||||
if args.image_path:
|
||||
loaded_image = Image.open(args.image_path).convert("RGB")
|
||||
think_prefix = f"<|im_start|>system\n{VLM_THINK_SYSTEM_PROMPT}<|im_end|>\n" if args.think else ""
|
||||
final_prompt_text = (
|
||||
f"{think_prefix}<|im_start|>user\n<|image_pad|>\n{p}<|im_end|>\n<|im_start|>assistant\n"
|
||||
)
|
||||
prompt_dict = {
|
||||
"prompt": final_prompt_text,
|
||||
"multi_modal_data": {"image": loaded_image},
|
||||
"modalities": ["text"],
|
||||
}
|
||||
formatted_prompts.append(prompt_dict)
|
||||
elif args.modality == "text2text":
|
||||
think_prefix = f"<|im_start|>{VLM_THINK_SYSTEM_PROMPT}<|im_end|>" if args.think else ""
|
||||
final_prompt_text = f"{think_prefix}<|im_start|>{p}<|im_end|><|im_start|>"
|
||||
prompt_dict = {"prompt": final_prompt_text, "modalities": ["text"]}
|
||||
formatted_prompts.append(prompt_dict)
|
||||
else:
|
||||
think_prefix = f"<|im_start|>{GEN_THINK_SYSTEM_PROMPT}<|im_end|>" if args.think else ""
|
||||
final_prompt_text = f"{think_prefix}<|im_start|>{p}<|im_end|>"
|
||||
prompt_dict = {"prompt": final_prompt_text, "modalities": ["image"]}
|
||||
if args.negative_prompt is not None:
|
||||
prompt_dict["negative_prompt"] = args.negative_prompt
|
||||
formatted_prompts.append(prompt_dict)
|
||||
|
||||
params_list = omni.default_sampling_params_list
|
||||
# Bagel exposes 1 sampling param set for single-stage (DiT-only) and
|
||||
# 2 for two-stage (Thinker + DiT). This heuristic may need updating
|
||||
# if future pipelines break that 1:1 mapping.
|
||||
is_single_stage = len(params_list) == 1
|
||||
|
||||
diffusion_params_idx = 0 if is_single_stage else (1 if len(params_list) > 1 else 0)
|
||||
diffusion_params = params_list[diffusion_params_idx]
|
||||
|
||||
if args.modality in ("text2img", "img2img"):
|
||||
diffusion_params.num_inference_steps = args.steps # type: ignore
|
||||
diffusion_params.cfg_parallel_size = args.cfg_parallel_size # type: ignore
|
||||
if args.seed is not None:
|
||||
diffusion_params.seed = args.seed # type: ignore
|
||||
|
||||
extra = getattr(diffusion_params, "extra_args", {}) or {}
|
||||
extra["cfg_text_scale"] = args.cfg_text_scale
|
||||
extra["cfg_img_scale"] = args.cfg_img_scale
|
||||
if args.cfg_interval is not None:
|
||||
extra["cfg_interval"] = tuple(args.cfg_interval)
|
||||
if args.cfg_renorm_type is not None:
|
||||
extra["cfg_renorm_type"] = args.cfg_renorm_type
|
||||
if args.cfg_renorm_min is not None:
|
||||
extra["cfg_renorm_min"] = args.cfg_renorm_min
|
||||
if args.negative_prompt is not None:
|
||||
extra["negative_prompt"] = args.negative_prompt
|
||||
|
||||
needs_text_gen = is_single_stage and (args.think or args.modality in ("text2text", "img2text"))
|
||||
if needs_text_gen:
|
||||
if args.think:
|
||||
extra["think"] = True
|
||||
extra["max_think_tokens"] = args.max_think_tokens
|
||||
extra["do_sample"] = args.do_sample
|
||||
extra["text_temperature"] = args.text_temperature
|
||||
diffusion_params.extra_args = extra # type: ignore
|
||||
|
||||
omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list))
|
||||
|
||||
img_idx = 0
|
||||
for req_output in omni_outputs:
|
||||
# 2-stage think mode: text output from thinker stage
|
||||
ro = getattr(req_output, "request_output", None)
|
||||
if ro and getattr(ro, "outputs", None):
|
||||
txt = "".join(getattr(o, "text", "") or "" for o in ro.outputs)
|
||||
if txt:
|
||||
if args.think:
|
||||
print(f"[Think]\n{txt}")
|
||||
else:
|
||||
print(f"[Output] Text:\n{txt}")
|
||||
|
||||
# Single-stage DiT: text from custom_output
|
||||
custom = getattr(req_output, "_custom_output", {}) or {}
|
||||
if custom.get("think_text"):
|
||||
print(f"[Think]\n{custom['think_text']}")
|
||||
if custom.get("text_output"):
|
||||
print(f"[Output] Text:\n{custom['text_output']}")
|
||||
|
||||
images = getattr(req_output, "images", None)
|
||||
if not images:
|
||||
continue
|
||||
|
||||
for j, img in enumerate(images):
|
||||
save_path = os.path.join(args.output, f"output_{img_idx}_{j}.png")
|
||||
img.save(save_path)
|
||||
print(f"[Output] Saved image to {save_path}")
|
||||
img_idx += 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -97,9 +97,16 @@ import torch
|
||||
from PIL import Image
|
||||
|
||||
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.entrypoints.openai.stage_params import clone_sampling_params
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
from vllm_omni.model_extras import (
|
||||
build_image_to_image_prompt,
|
||||
get_extra_body_params,
|
||||
get_model_class_name,
|
||||
should_init_extra_args_for_non_diffusion_stages,
|
||||
)
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
|
||||
@@ -187,6 +194,12 @@ def parse_args() -> argparse.Namespace:
|
||||
parser.add_argument(
|
||||
"--guidance-scale-2", type=float, default=None, help="image guidance scale for image-to-image generation."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--extra-args",
|
||||
type=parse_profiler_config,
|
||||
default=None,
|
||||
help="JSON object copied to OmniDiffusionSamplingParams.extra_args, e.g. '{\"cfg_text_scale\": 4.0}'.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
type=str,
|
||||
@@ -468,6 +481,8 @@ def main():
|
||||
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
|
||||
profiler_config=args.profiler_config,
|
||||
)
|
||||
model_class_name = get_model_class_name(omni)
|
||||
declared_extra_body_params = get_extra_body_params(model_class_name)
|
||||
print("Pipeline loaded")
|
||||
|
||||
profiler_enabled = args.profiler_config is not None
|
||||
@@ -478,6 +493,8 @@ def main():
|
||||
print(f" Model: {args.model}")
|
||||
print(f" Inference steps: {args.num_inference_steps}")
|
||||
print(f" Cache backend: {args.cache_backend if args.cache_backend else 'None (no acceleration)'}")
|
||||
if args.height is not None or args.width is not None:
|
||||
print(f" Output size: {args.width or 'auto'}x{args.height or 'auto'}")
|
||||
if isinstance(input_image, list):
|
||||
print(f" Number of input images: {len(input_image)}")
|
||||
for idx, img in enumerate(input_image):
|
||||
@@ -487,7 +504,6 @@ def main():
|
||||
print(
|
||||
f" Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size}, enable_expert_parallel: {args.enable_expert_parallel}"
|
||||
)
|
||||
print(f" Output image size: {args.width or 'default'}x{args.height or 'default'}")
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
generation_start = time.perf_counter()
|
||||
@@ -496,10 +512,20 @@ def main():
|
||||
print("[Profiler] Starting profiling...")
|
||||
omni.start_profile()
|
||||
|
||||
# Generate edited image
|
||||
sampling_params = OmniDiffusionSamplingParams(
|
||||
width=args.width,
|
||||
prompt_dict = build_image_to_image_prompt(
|
||||
model_class_name=model_class_name,
|
||||
prompt=args.prompt,
|
||||
negative_prompt=args.negative_prompt,
|
||||
input_image=input_image,
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
)
|
||||
|
||||
extra_args_from_cli = dict(args.extra_args or {})
|
||||
if args.negative_prompt is not None:
|
||||
extra_args_from_cli.setdefault("negative_prompt", args.negative_prompt)
|
||||
|
||||
diffusion_params = OmniDiffusionSamplingParams(
|
||||
generator=generator,
|
||||
true_cfg_scale=args.cfg_scale,
|
||||
guidance_scale=args.guidance_scale,
|
||||
@@ -507,17 +533,44 @@ def main():
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
num_outputs_per_prompt=args.num_outputs_per_prompt,
|
||||
layers=args.layers,
|
||||
resolution=args.resolution,
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
)
|
||||
if args.resolution:
|
||||
sampling_params.resolution = args.resolution
|
||||
outputs = omni.generate(
|
||||
{
|
||||
"prompt": args.prompt,
|
||||
"negative_prompt": args.negative_prompt,
|
||||
"multi_modal_data": {"image": input_image},
|
||||
},
|
||||
sampling_params,
|
||||
if declared_extra_body_params:
|
||||
apply_declared_extra_args(
|
||||
diffusion_params,
|
||||
declared_extra_body_params,
|
||||
extra_args_from_cli,
|
||||
)
|
||||
else:
|
||||
diffusion_params.extra_args.update({k: v for k, v in extra_args_from_cli.items() if v is not None})
|
||||
|
||||
# Build per-stage sampling params for multi-stage models
|
||||
init_non_diffusion = should_init_extra_args_for_non_diffusion_stages(
|
||||
model_class_name,
|
||||
)
|
||||
defaults = list(omni.default_sampling_params_list or [])
|
||||
sampling_params_list = [clone_sampling_params(p) for p in defaults]
|
||||
if not sampling_params_list:
|
||||
sampling_params_list = [diffusion_params]
|
||||
|
||||
diffusion_replaced = False
|
||||
for idx, params in enumerate(sampling_params_list):
|
||||
if isinstance(params, OmniDiffusionSamplingParams):
|
||||
merged_extra = dict(getattr(params, "extra_args", {}) or {})
|
||||
merged_extra.update(diffusion_params.extra_args)
|
||||
diffusion_params.extra_args = merged_extra
|
||||
sampling_params_list[idx] = diffusion_params
|
||||
diffusion_replaced = True
|
||||
elif init_non_diffusion and hasattr(params, "extra_args"):
|
||||
if params.extra_args is None:
|
||||
params.extra_args = {}
|
||||
|
||||
if not diffusion_replaced and len(sampling_params_list) == 1:
|
||||
sampling_params_list = [diffusion_params]
|
||||
|
||||
outputs = omni.generate(prompt_dict, sampling_params_list=sampling_params_list)
|
||||
generation_end = time.perf_counter()
|
||||
generation_time = generation_end - generation_start
|
||||
|
||||
@@ -544,17 +597,16 @@ def main():
|
||||
if not outputs:
|
||||
raise ValueError("No output generated from omni.generate()")
|
||||
|
||||
# Extract images from OmniRequestOutput
|
||||
# omni.generate() returns list[OmniRequestOutput], extract images from request_output.images
|
||||
first_output = outputs[0]
|
||||
if not hasattr(first_output, "request_output") or not first_output.request_output:
|
||||
raise ValueError("No request_output found in OmniRequestOutput")
|
||||
images = None
|
||||
for output in outputs:
|
||||
images = getattr(output, "images", None)
|
||||
if images:
|
||||
break
|
||||
req_out = getattr(output, "request_output", None)
|
||||
images = getattr(req_out, "images", None) if req_out is not None else None
|
||||
if images:
|
||||
break
|
||||
|
||||
req_out = first_output.request_output
|
||||
if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
|
||||
raise ValueError("Invalid request_output structure or missing 'images' key")
|
||||
|
||||
images = req_out.images
|
||||
if not images:
|
||||
raise ValueError("No images found in request_output")
|
||||
|
||||
|
||||
@@ -10,10 +10,18 @@ from typing import Any
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.data import logger
|
||||
from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.entrypoints.openai.stage_params import clone_sampling_params
|
||||
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,
|
||||
get_extra_body_params,
|
||||
get_model_class_name,
|
||||
should_init_extra_args_for_non_diffusion_stages,
|
||||
)
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
|
||||
@@ -259,6 +267,19 @@ def parse_args() -> argparse.Namespace:
|
||||
action="store_true",
|
||||
help="[NextStep-1.1 only] Apply layer normalization to sampled tokens.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--extra-body",
|
||||
type=parse_profiler_config,
|
||||
default=None,
|
||||
help=(
|
||||
"Model-specific generation params as a JSON object, e.g. "
|
||||
'\'{"timestep_shift": 3.0, "cfg_text_scale": 4.0, "cfg_interval": [0.4, 1.0]}\'. '
|
||||
"Each key is filtered against the model's declared extra_body_params "
|
||||
"(see vllm_omni/model_extras), so unknown keys for the chosen model are "
|
||||
"silently dropped. Values here take precedence over the equivalent "
|
||||
"model-specific flags above."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable-diffusion-pipeline-profiler",
|
||||
action="store_true",
|
||||
@@ -405,6 +426,8 @@ def main():
|
||||
# NextStep-1.1 requires explicit pipeline class
|
||||
omni_kwargs["model_class_name"] = "NextStep11Pipeline"
|
||||
omni = Omni(**omni_kwargs)
|
||||
model_class_name = get_model_class_name(omni)
|
||||
declared_extra_body_params = get_extra_body_params(model_class_name)
|
||||
|
||||
if profiler_enabled:
|
||||
print("[Profiler] Starting profiling...")
|
||||
@@ -445,34 +468,79 @@ def main():
|
||||
)
|
||||
|
||||
generation_start = time.perf_counter()
|
||||
extra_args = {
|
||||
|
||||
prompt_dict = build_text_to_image_prompt(
|
||||
model_class_name=model_class_name,
|
||||
prompt=args.prompt,
|
||||
negative_prompt=args.negative_prompt,
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
)
|
||||
|
||||
diffusion_params = OmniDiffusionSamplingParams(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
seed=args.seed,
|
||||
generator=generator,
|
||||
true_cfg_scale=args.cfg_scale,
|
||||
guidance_scale=args.guidance_scale,
|
||||
guidance_scale_2=args.guidance_scale_2,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
num_outputs_per_prompt=args.num_images_per_prompt,
|
||||
)
|
||||
|
||||
# Base layer: backward-compatible model-specific flags. New model params should
|
||||
# instead be declared in vllm_omni/model_extras and passed via --extra-body, so
|
||||
# this dict does not need to grow per-model.
|
||||
user_extra = {
|
||||
"cfg_scale": args.cfg_scale,
|
||||
"cfg_text_scale": args.cfg_scale,
|
||||
"negative_prompt": args.negative_prompt,
|
||||
"timestep_shift": args.timesteps_shift,
|
||||
"timesteps_shift": args.timesteps_shift,
|
||||
"cfg_schedule": args.cfg_schedule,
|
||||
"use_norm": args.use_norm,
|
||||
"use_system_prompt": args.use_system_prompt,
|
||||
"system_prompt": args.system_prompt,
|
||||
}
|
||||
if lora_request:
|
||||
extra_args["lora_request"] = lora_request
|
||||
extra_args["lora_scale"] = args.lora_scale
|
||||
# Override layer: generic JSON passthrough wins over the flags above. Keys are
|
||||
# still filtered against the model's declared extra_body_params downstream.
|
||||
if args.extra_body:
|
||||
user_extra.update(args.extra_body)
|
||||
if declared_extra_body_params:
|
||||
apply_declared_extra_args(diffusion_params, declared_extra_body_params, user_extra)
|
||||
else:
|
||||
diffusion_params.extra_args.update({k: v for k, v in user_extra.items() if v is not None})
|
||||
|
||||
outputs = omni.generate(
|
||||
{
|
||||
"prompt": args.prompt,
|
||||
"negative_prompt": args.negative_prompt,
|
||||
},
|
||||
OmniDiffusionSamplingParams(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
generator=generator,
|
||||
true_cfg_scale=args.cfg_scale,
|
||||
guidance_scale=args.guidance_scale,
|
||||
guidance_scale_2=args.guidance_scale_2,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
num_outputs_per_prompt=args.num_images_per_prompt,
|
||||
extra_args=extra_args,
|
||||
),
|
||||
if lora_request:
|
||||
diffusion_params.extra_args["lora_request"] = lora_request
|
||||
diffusion_params.extra_args["lora_scale"] = args.lora_scale
|
||||
|
||||
# Build per-stage sampling params for multi-stage models (e.g. BAGEL),
|
||||
# or wrap single diffusion params for single-stage models.
|
||||
init_non_diffusion = should_init_extra_args_for_non_diffusion_stages(
|
||||
model_class_name,
|
||||
)
|
||||
defaults = list(omni.default_sampling_params_list or [])
|
||||
sampling_params_list = [clone_sampling_params(p) for p in defaults]
|
||||
if not sampling_params_list:
|
||||
sampling_params_list = [diffusion_params]
|
||||
|
||||
diffusion_replaced = False
|
||||
for idx, params in enumerate(sampling_params_list):
|
||||
if isinstance(params, OmniDiffusionSamplingParams):
|
||||
sampling_params_list[idx] = diffusion_params
|
||||
diffusion_replaced = True
|
||||
elif init_non_diffusion and hasattr(params, "extra_args"):
|
||||
if params.extra_args is None:
|
||||
params.extra_args = {}
|
||||
if args.seed is not None and hasattr(params, "seed"):
|
||||
params.seed = args.seed
|
||||
|
||||
if not diffusion_replaced and len(sampling_params_list) == 1:
|
||||
sampling_params_list = [diffusion_params]
|
||||
|
||||
outputs = omni.generate(prompt_dict, sampling_params_list=sampling_params_list)
|
||||
|
||||
generation_end = time.perf_counter()
|
||||
generation_time = generation_end - generation_start
|
||||
@@ -502,15 +570,16 @@ def main():
|
||||
raise ValueError("No output generated from omni.generate()")
|
||||
logger.info(f"Outputs: {outputs}")
|
||||
|
||||
first_output = outputs[0]
|
||||
if not hasattr(first_output, "request_output") or not first_output.request_output:
|
||||
raise ValueError("No request_output found in OmniRequestOutput")
|
||||
images = None
|
||||
for output in outputs:
|
||||
images = getattr(output, "images", None)
|
||||
if images:
|
||||
break
|
||||
req_out = getattr(output, "request_output", None)
|
||||
images = getattr(req_out, "images", None) if req_out is not None else None
|
||||
if images:
|
||||
break
|
||||
|
||||
req_out = first_output.request_output
|
||||
if not hasattr(req_out, "images"):
|
||||
raise ValueError("Invalid request_output structure or missing 'images'.")
|
||||
|
||||
images = req_out.images
|
||||
if not images:
|
||||
raise ValueError("No images found in request_output")
|
||||
|
||||
|
||||
@@ -1,404 +0,0 @@
|
||||
# BAGEL-7B-MoT
|
||||
|
||||
## Installation
|
||||
|
||||
Please refer to [README.md](../../../README.md)
|
||||
|
||||
## Architecture
|
||||
|
||||
BAGEL-7B-MoT is a Mixture-of-Transformers (MoT) model supporting both image generation and understanding. It offers two deployment topologies:
|
||||
|
||||
| Topology | Stages | Description |
|
||||
| :------- | :----- | :---------- |
|
||||
| **Two-stage** (default) | Stage 0 (Thinker, AR) + Stage 1 (DiT, Diffusion) | Thinker handles text/understanding via vLLM AR engine; DiT handles image generation. KV cache is transferred between stages. |
|
||||
| **Single-stage** | Stage 0 (DiT, Diffusion) only | The DiT stage contains a full LLM, ViT, VAE, and tokenizer internally. All modalities are handled within a single diffusion process. |
|
||||
|
||||
Both topologies support all four modalities: `text2img`, `img2img`, `img2text`, `text2text`.
|
||||
|
||||
> **Note**: These examples work with the default configuration on an **NVIDIA A100 (80GB)**. We also tested on dual **NVIDIA RTX 5000 Ada (32GB each)**. For dual-GPU setups, modify the deploy YAML to distribute stages across devices.
|
||||
|
||||
## Launch the Server
|
||||
|
||||
### Two-Stage (Default)
|
||||
|
||||
The default pipeline is auto-detected from the model. No extra flags needed:
|
||||
|
||||
```bash
|
||||
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091
|
||||
```
|
||||
|
||||
Or use the convenience script:
|
||||
|
||||
```bash
|
||||
cd examples/online_serving/bagel
|
||||
bash run_server.sh
|
||||
|
||||
# Initialize each stage in a discrete isolated process terminal
|
||||
bash run_server_stage_cli.sh --stage 0
|
||||
bash run_server_stage_cli.sh --stage 1
|
||||
```
|
||||
|
||||
To use a custom deploy YAML, pass it via `--deploy-config`:
|
||||
|
||||
```bash
|
||||
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \
|
||||
--deploy-config /path/to/deploy_config.yaml
|
||||
```
|
||||
|
||||
See [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml) for the default two-stage deploy configuration.
|
||||
|
||||
### Single-Stage
|
||||
|
||||
The DiT stage contains a full LLM, ViT, VAE, and tokenizer, so it can handle all modalities (text2img, img2img, img2text, text2text, think) without a separate Thinker stage:
|
||||
|
||||
```bash
|
||||
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \
|
||||
--deploy-config vllm_omni/deploy/bagel_single_stage.yaml
|
||||
```
|
||||
|
||||
See [`bagel_single_stage.yaml`](../../../vllm_omni/deploy/bagel_single_stage.yaml) for configuration. The `pipeline: bagel_single_stage` field selects the single-stage topology from the pipeline registry.
|
||||
|
||||
### Tensor Parallelism (TP)
|
||||
|
||||
For larger models or multi-GPU environments, enable TP via CLI:
|
||||
|
||||
```bash
|
||||
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --tensor-parallel-size 2
|
||||
```
|
||||
|
||||
Or set `tensor_parallel_size` per stage in a custom deploy YAML.
|
||||
|
||||
### VAE Patch Parallelism
|
||||
|
||||
[VAE Patch Parallelism](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism/vae_patch_parallel.html) distributes Bagel VAE **decode/encode** across multiple GPUs by splitting latent tiles. It lowers **per-GPU peak memory during VAE decode**, which helps high-resolution `text2img` / `img2img` when VAE becomes a bottleneck.
|
||||
|
||||
**Scope for Bagel:**
|
||||
|
||||
| Topology | VAE patch parallel |
|
||||
| :------- | :----------------- |
|
||||
| **Single-stage** (DiT only) | Supported on stage 0 (`BagelPipeline` + `DistributedAutoEncoder`) |
|
||||
| **Two-stage** | Supported on **stage 1 (DiT)** only; stage 0 (Thinker) uses encoder-only VAE and is unrelated |
|
||||
|
||||
**Requirements:**
|
||||
|
||||
- `vae_patch_parallel_size > 1` and a distributed VAE (`DistributedAutoEncoder` on the DiT pipeline).
|
||||
- The DiT process group must have at least `vae_patch_parallel_size` ranks. In practice this means the diffusion stage `world_size` must be ≥ 2 (commonly `tensor_parallel_size=2` on that stage).
|
||||
- `vae_use_tiling` must be enabled. If you set `vae_patch_parallel_size > 1` and omit tiling, the registry auto-enables `vae_use_tiling` at startup.
|
||||
|
||||
VAE patch parallel **reuses the DiT process group** (`dit_group`); it does not create a separate VAE-only worker pool. It is not a substitute for single-GPU VAE tiling (`vae_pp=1`).
|
||||
|
||||
**Online serving (single-stage, 2 GPUs):**
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0,1 vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \
|
||||
--deploy-config vllm_omni/deploy/bagel_single_stage.yaml \
|
||||
--tensor-parallel-size 2 \
|
||||
--vae-patch-parallel-size 2 \
|
||||
--vae-use-tiling
|
||||
```
|
||||
|
||||
**Online serving (two-stage, VAE PP on DiT stage 1):** use a custom deploy YAML, for example:
|
||||
|
||||
```yaml
|
||||
stages:
|
||||
- stage_id: 0
|
||||
devices: "0"
|
||||
# Thinker (AR) — no VAE patch parallel here
|
||||
|
||||
- stage_id: 1
|
||||
devices: "0,1"
|
||||
vae_use_tiling: true
|
||||
parallel_config:
|
||||
tensor_parallel_size: 2
|
||||
vae_patch_parallel_size: 2
|
||||
```
|
||||
|
||||
```bash
|
||||
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \
|
||||
--deploy-config /path/to/bagel_vae_pp.yaml
|
||||
```
|
||||
|
||||
**Verify it is active** (check server logs at startup):
|
||||
|
||||
```text
|
||||
INFO ... vae_patch_parallel_size=2 requires vae_use_tiling; automatically enabling it.
|
||||
```
|
||||
|
||||
| CLI flag | Default | Description |
|
||||
| :------- | :------ | :---------- |
|
||||
| `--vae-patch-parallel-size` | `1` | Number of DiT ranks used for VAE tile parallelism. Set to `2` or higher to enable. Should be ≤ DiT process group size (typically match `--tensor-parallel-size` on the diffusion stage). |
|
||||
| `--vae-use-tiling` | off | Enable VAE spatial tiling. Required for VAE patch parallel (auto-enabled when `vae_patch_parallel_size > 1`). |
|
||||
|
||||
#### Hybrid Sharded Data Parallel (HSDP)
|
||||
|
||||
For larger Bagel deployments on multiple GPUs, you can enable HSDP (Hybrid Sharded Data Parallel) by modifying the stage configuration (for example, [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml)). HSDP shards transformer weights across GPUs to reduce per-GPU memory usage.
|
||||
|
||||
1. **Enable HSDP**: Set `use_hsdp: true`.
|
||||
2. **Set shard size**: Set `hsdp_shard_size` to the number of GPUs used for sharding (for example, `4`).
|
||||
3. **Set replicate size**: Usually keep `hsdp_replicate_size: 1` unless you want replicated HSDP groups.
|
||||
4. **Set devices**: Specify the comma-separated GPU IDs used by the diffusion stage (for example, `"0,1,2,3"`).
|
||||
|
||||
Example configuration for HSDP across 4 GPUs:
|
||||
```yaml
|
||||
- stage_id: 1
|
||||
devices: "0,1,2,3"
|
||||
parallel_config:
|
||||
use_hsdp: true
|
||||
hsdp_shard_size: 4
|
||||
hsdp_replicate_size: 1
|
||||
```
|
||||
|
||||
### Multi-Node Deployment
|
||||
|
||||
Deploy each stage on a **separate node** for better resource utilization. Replace `<ORCHESTRATOR_IP>` with the actual IP address of your orchestrator node.
|
||||
|
||||
**1. Launch Stage 0 (Thinker / Orchestrator)** on the orchestrator node:
|
||||
|
||||
```bash
|
||||
# API server port for client requests: 8000
|
||||
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni \
|
||||
--port 8000 \
|
||||
--stage-id 0 \
|
||||
--omni-master-address <ORCHESTRATOR_IP> \
|
||||
--omni-master-port 8091
|
||||
```
|
||||
|
||||
**2. Launch Stage 1 (DiT)** on the remote node in headless mode:
|
||||
|
||||
```bash
|
||||
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni \
|
||||
--stage-id 1 \
|
||||
--headless \
|
||||
--omni-master-address <ORCHESTRATOR_IP> \
|
||||
--omni-master-port 8091
|
||||
```
|
||||
|
||||
Or use the convenience script:
|
||||
|
||||
```bash
|
||||
# Terminal 1: Stage 0
|
||||
bash run_server_stage_cli.sh --stage 0
|
||||
|
||||
# Terminal 2: Stage 1
|
||||
bash run_server_stage_cli.sh --stage 1
|
||||
|
||||
# With extra args
|
||||
bash run_server_stage_cli.sh --stage 0 -- --tensor-parallel-size 2
|
||||
bash run_server_stage_cli.sh --stage 1 -- --gpu-memory-utilization 0.9
|
||||
```
|
||||
|
||||
**vllm serve arguments:**
|
||||
|
||||
| Argument | Description |
|
||||
| :------- | :---------- |
|
||||
| `--stage-id` | Which stage this process runs (0 = Thinker, 1 = DiT) |
|
||||
| `--headless` | Run without the API server (worker-only mode) |
|
||||
| `-oma` / `--omni-master-address` | Orchestrator master address |
|
||||
| `-omp` / `--omni-master-port` | Orchestrator master port |
|
||||
|
||||
> [!IMPORTANT]
|
||||
> **Startup Order**: Stage 0 (orchestrator) must be launched **before** Stage 1 (headless).
|
||||
> Stage 0 will appear to hang on startup until Stage 1 (worker) connects — this is expected behavior.
|
||||
|
||||
### Inter-Stage Connectors
|
||||
|
||||
When deploying stages across nodes, configure the connector type in the deploy YAML:
|
||||
|
||||
- **SharedMemoryConnector** (default): Used for single-node deployments. No explicit configuration needed.
|
||||
- **MooncakeTransferEngineConnector**: For multi-node setups with RDMA hardware. Defined in [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml) under `connectors.rdma_connector`.
|
||||
|
||||
To use Mooncake, create a custom deploy YAML that binds `output_connectors` / `input_connectors` on each stage to the `rdma_connector` defined in the `connectors` section.
|
||||
|
||||
## Send Requests
|
||||
|
||||
```bash
|
||||
cd examples/online_serving/bagel
|
||||
```
|
||||
|
||||
### Text to Image (text2img)
|
||||
|
||||
**Python client:**
|
||||
|
||||
```bash
|
||||
python openai_chat_client.py \
|
||||
--prompt "A beautiful sunset over mountains" \
|
||||
--modality text2img \
|
||||
--output sunset.png \
|
||||
--steps 50
|
||||
```
|
||||
|
||||
**curl:**
|
||||
|
||||
```bash
|
||||
curl http://localhost:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"messages": [{"role": "user", "content": [{"type": "text", "text": "<|im_start|>A beautiful sunset over mountains<|im_end|>"}]}],
|
||||
"modalities": ["image"],
|
||||
"height": 512,
|
||||
"width": 512,
|
||||
"num_inference_steps": 50,
|
||||
"seed": 42
|
||||
}'
|
||||
```
|
||||
|
||||
### Image to Image (img2img)
|
||||
|
||||
**Python client:**
|
||||
|
||||
```bash
|
||||
python openai_chat_client.py \
|
||||
--prompt "Make the cat stand up" \
|
||||
--modality img2img \
|
||||
--image-url /path/to/input.jpg \
|
||||
--output transformed.png
|
||||
```
|
||||
|
||||
**curl:**
|
||||
|
||||
```bash
|
||||
IMAGE_BASE64=$(base64 -w 0 cat.jpg)
|
||||
|
||||
cat <<EOF > payload.json
|
||||
{
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "<|im_start|>Make the cat stand up<|im_end|>"},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,${IMAGE_BASE64}"}}
|
||||
]
|
||||
}],
|
||||
"modalities": ["image"],
|
||||
"height": 512,
|
||||
"width": 512,
|
||||
"num_inference_steps": 50,
|
||||
"seed": 42
|
||||
}
|
||||
EOF
|
||||
|
||||
curl http://localhost:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d @payload.json
|
||||
```
|
||||
|
||||
### Image to Text (img2text)
|
||||
|
||||
**Python client:**
|
||||
|
||||
```bash
|
||||
python openai_chat_client.py \
|
||||
--prompt "Describe this image in detail" \
|
||||
--modality img2text \
|
||||
--image-url /path/to/image.jpg
|
||||
```
|
||||
|
||||
**curl:**
|
||||
|
||||
```bash
|
||||
IMAGE_BASE64=$(base64 -w 0 cat.jpg)
|
||||
|
||||
cat <<EOF > payload.json
|
||||
{
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "<|im_start|>user\n<|image_pad|>\nDescribe this image in detail<|im_end|>\n<|im_start|>assistant\n"},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,${IMAGE_BASE64}"}}
|
||||
]
|
||||
}],
|
||||
"modalities": ["text"]
|
||||
}
|
||||
EOF
|
||||
|
||||
curl http://localhost:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d @payload.json
|
||||
```
|
||||
|
||||
### Text to Text (text2text)
|
||||
|
||||
**Python client:**
|
||||
|
||||
```bash
|
||||
python openai_chat_client.py \
|
||||
--prompt "What is the capital of France?" \
|
||||
--modality text2text
|
||||
```
|
||||
|
||||
**curl:**
|
||||
|
||||
```bash
|
||||
curl http://localhost:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"messages": [{"role": "user", "content": [{"type": "text", "text": "<|im_start|>user\nWhat is the capital of France?<|im_end|>\n<|im_start|>assistant\n"}]}],
|
||||
"modalities": ["text"]
|
||||
}'
|
||||
```
|
||||
|
||||
### Python Client Arguments
|
||||
|
||||
| Argument | Default | Description |
|
||||
| :------- | :------ | :---------- |
|
||||
| `--prompt` / `-p` | `A cute cat` | Text prompt |
|
||||
| `--output` / `-o` | `bagel_output.png` | Output file path |
|
||||
| `--server` / `-s` | `http://localhost:8091` | Server URL |
|
||||
| `--image-url` / `-i` | `None` | Input image URL or local path (img2img/img2text) |
|
||||
| `--modality` / `-m` | `text2img` | `text2img`, `img2img`, `img2text`, `text2text` |
|
||||
| `--height` | `512` | Image height in pixels |
|
||||
| `--width` | `512` | Image width in pixels |
|
||||
| `--steps` | `25` | Number of inference steps |
|
||||
| `--seed` | `42` | Random seed |
|
||||
| `--negative` | `None` | Negative prompt for CFG |
|
||||
|
||||
Example with custom parameters:
|
||||
|
||||
```bash
|
||||
python openai_chat_client.py \
|
||||
--prompt "A futuristic city" \
|
||||
--modality text2img \
|
||||
--height 768 \
|
||||
--width 768 \
|
||||
--steps 50 \
|
||||
--seed 42 \
|
||||
--negative "blurry, low quality"
|
||||
```
|
||||
|
||||
## Configuration Reference
|
||||
|
||||
### Deploy YAML Files
|
||||
|
||||
| File | Description |
|
||||
| :--- | :---------- |
|
||||
| [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml) | Two-stage default (Thinker + DiT on GPU 0) |
|
||||
| [`bagel_single_stage.yaml`](../../../vllm_omni/deploy/bagel_single_stage.yaml) | Single-stage (DiT only) |
|
||||
|
||||
### Key Deploy YAML Fields
|
||||
|
||||
| Field | Scope | Description |
|
||||
| :---- | :---- | :---------- |
|
||||
| `pipeline` | top-level | Override auto-detected pipeline (e.g. `bagel_single_stage`) |
|
||||
| `stages[].stage_id` | per-stage | Stage identifier (0, 1, ...) |
|
||||
| `stages[].devices` | per-stage | GPU device IDs (e.g. `"0"`, `"0,1"`) |
|
||||
| `stages[].max_num_seqs` | per-stage | Maximum concurrent sequences |
|
||||
| `stages[].gpu_memory_utilization` | per-stage | Fraction of GPU memory to use |
|
||||
| `stages[].enforce_eager` | per-stage | Disable CUDA graphs |
|
||||
| `stages[].tensor_parallel_size` | per-stage | TP degree for this stage |
|
||||
| `stages[].parallel_config.vae_patch_parallel_size` | per-stage (DiT) | VAE tile parallelism degree (DiT stage only) |
|
||||
| `stages[].vae_use_tiling` | per-stage (DiT) | Enable VAE tiling (required for VAE patch parallel) |
|
||||
| `connectors` | top-level | Define available connector instances (SHM, Mooncake) |
|
||||
| `platforms` | top-level | Platform-specific overrides (e.g. `xpu`) |
|
||||
|
||||
## FAQ
|
||||
|
||||
- If you encounter OOM errors, try decreasing `max_model_len` or `gpu_memory_utilization` in the deploy YAML.
|
||||
|
||||
**Two-stage VRAM usage:**
|
||||
|
||||
| Stage | VRAM |
|
||||
| :---- | :--- |
|
||||
| Stage 0 (Thinker) | **15.04 GiB + KV Cache** |
|
||||
| Stage 1 (DiT) | **26.50 GiB** |
|
||||
| Total | **~42 GiB + KV Cache** |
|
||||
|
||||
**Single-stage VRAM usage:** The DiT loads the full model (~42 GiB) in one process.
|
||||
@@ -1,184 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Bagel OpenAI-compatible chat client for image generation and multimodal tasks.
|
||||
|
||||
Usage:
|
||||
python openai_chat_client.py --prompt "A cute cat" --output output.png
|
||||
python openai_chat_client.py --prompt "Describe this image" --image-url https://example.com/image.png
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
from pathlib import Path
|
||||
|
||||
import requests
|
||||
|
||||
|
||||
def generate_image(
|
||||
prompt: str,
|
||||
server_url: str = "http://localhost:8091",
|
||||
image_url: str | None = None,
|
||||
height: int | None = None,
|
||||
width: int | None = None,
|
||||
steps: int | None = None,
|
||||
seed: int | None = None,
|
||||
negative_prompt: str | None = None,
|
||||
modality: str = "text2img", # "text2img" (default), "img2img", "img2text", "text2text"
|
||||
) -> bytes | str | None:
|
||||
"""Generate an image or text using the chat completions API.
|
||||
|
||||
Args:
|
||||
prompt: Text description or prompt
|
||||
server_url: Server URL
|
||||
image_url: URL or path to input image (for img2img/img2text)
|
||||
height: Image height in pixels
|
||||
width: Image width in pixels
|
||||
steps: Number of inference steps
|
||||
seed: Random seed
|
||||
negative_prompt: Negative prompt
|
||||
modality: Task modality hint
|
||||
|
||||
Returns:
|
||||
Image bytes (for image outputs) or Text string (for text outputs) or None if failed
|
||||
"""
|
||||
|
||||
# Construct Message Content
|
||||
content = [{"type": "text", "text": f"<|im_start|>{prompt}<|im_end|>"}]
|
||||
|
||||
if image_url:
|
||||
# Check if local file
|
||||
if Path(image_url).exists():
|
||||
with open(image_url, "rb") as f:
|
||||
b64_data = base64.b64encode(f.read()).decode("utf-8")
|
||||
final_image_url = f"data:image/jpeg;base64,{b64_data}"
|
||||
else:
|
||||
final_image_url = image_url
|
||||
|
||||
content.append({"type": "image_url", "image_url": {"url": final_image_url}})
|
||||
|
||||
messages = [{"role": "user", "content": content}]
|
||||
|
||||
# Build request payload with all parameters at top level
|
||||
# Note: vLLM ignores "extra_body", so we put parameters directly in the payload
|
||||
payload = {"messages": messages}
|
||||
|
||||
# Set output modalities at top level
|
||||
if modality == "text2img" or modality == "img2img":
|
||||
payload["modalities"] = ["image"]
|
||||
elif modality == "img2text" or modality == "text2text":
|
||||
payload["modalities"] = ["text"]
|
||||
|
||||
# Add generation parameters directly to payload
|
||||
if height is not None:
|
||||
payload["height"] = height
|
||||
if width is not None:
|
||||
payload["width"] = width
|
||||
if steps is not None:
|
||||
payload["num_inference_steps"] = steps
|
||||
if seed is not None:
|
||||
payload["seed"] = seed
|
||||
if negative_prompt:
|
||||
payload["negative_prompt"] = negative_prompt
|
||||
|
||||
# Send request
|
||||
try:
|
||||
print(f"Sending request to {server_url} with modality {modality}...")
|
||||
response = requests.post(
|
||||
f"{server_url}/v1/chat/completions",
|
||||
headers={"Content-Type": "application/json"},
|
||||
json=payload,
|
||||
timeout=300,
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
|
||||
# Extract content - check ALL choices since server may return multiple
|
||||
# (e.g., text in choices[0], image in choices[1])
|
||||
choices = data.get("choices", [])
|
||||
|
||||
# First pass: look for image output in any choice
|
||||
for choice in choices:
|
||||
choice_content = choice.get("message", {}).get("content")
|
||||
|
||||
# Handle Image Output
|
||||
if isinstance(choice_content, list) and len(choice_content) > 0:
|
||||
first_item = choice_content[0]
|
||||
if isinstance(first_item, dict) and "image_url" in first_item:
|
||||
img_url_str = first_item["image_url"].get("url", "")
|
||||
if img_url_str.startswith("data:image"):
|
||||
_, b64_data = img_url_str.split(",", 1)
|
||||
return base64.b64decode(b64_data)
|
||||
|
||||
# Second pass: look for text output if no image found
|
||||
for choice in choices:
|
||||
choice_content = choice.get("message", {}).get("content")
|
||||
if isinstance(choice_content, str) and choice_content:
|
||||
return choice_content
|
||||
|
||||
print(f"Unexpected response format: {choices}")
|
||||
return None
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Bagel multimodal chat client")
|
||||
parser.add_argument("--prompt", "-p", default="A cute cat", help="Text prompt")
|
||||
parser.add_argument("--output", "-o", default="bagel_output.png", help="Output file (for image results)")
|
||||
parser.add_argument("--server", "-s", default="http://localhost:8091", help="Server URL")
|
||||
|
||||
# Modality Control
|
||||
parser.add_argument("--image-url", "-i", type=str, help="Input image URL or local path")
|
||||
parser.add_argument(
|
||||
"--modality",
|
||||
"-m",
|
||||
default="text2img",
|
||||
choices=["text2img", "img2img", "img2text", "text2text"],
|
||||
help="Task modality",
|
||||
)
|
||||
|
||||
# Generation Params
|
||||
parser.add_argument("--height", type=int, default=512, help="Image height")
|
||||
parser.add_argument("--width", type=int, default=512, help="Image width")
|
||||
parser.add_argument("--steps", type=int, default=25, help="Inference steps")
|
||||
parser.add_argument("--seed", type=int, default=42, help="Random seed")
|
||||
parser.add_argument("--negative", help="Negative prompt")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
print(f"Mode: {args.modality}")
|
||||
if args.image_url:
|
||||
print(f"Input Image: {args.image_url}")
|
||||
|
||||
result = generate_image(
|
||||
prompt=args.prompt,
|
||||
server_url=args.server,
|
||||
image_url=args.image_url,
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
steps=args.steps,
|
||||
seed=args.seed,
|
||||
negative_prompt=args.negative,
|
||||
modality=args.modality,
|
||||
)
|
||||
|
||||
if result:
|
||||
if isinstance(result, bytes):
|
||||
# It's an image
|
||||
output_path = Path(args.output)
|
||||
output_path.write_bytes(result)
|
||||
print(f"Image saved to: {output_path}")
|
||||
print(f"Size: {len(result) / 1024:.1f} KB")
|
||||
elif isinstance(result, str):
|
||||
# It's text
|
||||
print("Response:")
|
||||
print(result)
|
||||
else:
|
||||
print("Failed to generate response")
|
||||
exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,12 +0,0 @@
|
||||
#!/bin/bash
|
||||
# Bagel online serving startup script
|
||||
|
||||
MODEL="${MODEL:-ByteDance-Seed/BAGEL-7B-MoT}"
|
||||
PORT="${PORT:-8091}"
|
||||
|
||||
echo "Starting Bagel server..."
|
||||
echo "Model: $MODEL"
|
||||
echo "Port: $PORT"
|
||||
|
||||
vllm serve "$MODEL" --omni \
|
||||
--port "$PORT"
|
||||
@@ -1,164 +0,0 @@
|
||||
#!/bin/bash
|
||||
# Bagel multi-stage online serving startup script.
|
||||
#
|
||||
# Usage:
|
||||
# ./run_server_stage_cli.sh --stage 0
|
||||
# ./run_server_stage_cli.sh --stage 1
|
||||
# ./run_server_stage_cli.sh --stage 0 -- --tensor-parallel-size 2
|
||||
# ./run_server_stage_cli.sh --stage 1 -- --gpu-memory-utilization 0.9
|
||||
#
|
||||
# By default, `--stage all` keeps the old behavior and launches both stages in
|
||||
# one session. Use `--stage 0` / `--stage 1` to launch each stage separately in
|
||||
# different terminal sessions, with stage-specific extra CLI arguments passed
|
||||
# after `--`.
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
MODEL="${MODEL:-ByteDance-Seed/BAGEL-7B-MoT}"
|
||||
PORT="${PORT:-8091}"
|
||||
MASTER_ADDRESS="${MASTER_ADDRESS:-127.0.0.1}"
|
||||
MASTER_PORT="${MASTER_PORT:-8092}"
|
||||
STAGE="all"
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
DEPLOY_CONFIG="${DEPLOY_CONFIG:-$SCRIPT_DIR/../../../vllm_omni/deploy/bagel.yaml}"
|
||||
EXTRA_ARGS=()
|
||||
|
||||
usage() {
|
||||
cat <<EOF
|
||||
Usage: $0 [OPTIONS] [-- EXTRA_VLLM_ARGS...]
|
||||
|
||||
Options:
|
||||
--stage {0|1|all} Stage to launch (default: all)
|
||||
--model MODEL Model name/path (default: $MODEL)
|
||||
--port PORT API port for stage 0 (default: $PORT)
|
||||
--master-address ADDRESS Master/orchestrator address (default: $MASTER_ADDRESS)
|
||||
--master-port PORT Master/orchestrator port (default: $MASTER_PORT)
|
||||
--deploy-config PATH Deploy config YAML path (default: $DEPLOY_CONFIG)
|
||||
--help Show this help message
|
||||
|
||||
Examples:
|
||||
$0 --stage 0
|
||||
$0 --stage 1
|
||||
$0 --stage 0 -- --tensor-parallel-size 2
|
||||
$0 --stage 1 -- --gpu-memory-utilization 0.9
|
||||
|
||||
Notes:
|
||||
- Use different terminal sessions to launch stage 0 and stage 1 separately.
|
||||
- Extra args after '--' are forwarded only to the selected stage.
|
||||
- When using '--stage all', the extra args are forwarded to both stages.
|
||||
EOF
|
||||
}
|
||||
|
||||
while [[ $# -gt 0 ]]; do
|
||||
case "$1" in
|
||||
--stage)
|
||||
STAGE="$2"
|
||||
shift 2
|
||||
;;
|
||||
--model)
|
||||
MODEL="$2"
|
||||
shift 2
|
||||
;;
|
||||
--port)
|
||||
PORT="$2"
|
||||
shift 2
|
||||
;;
|
||||
--master-address)
|
||||
MASTER_ADDRESS="$2"
|
||||
shift 2
|
||||
;;
|
||||
--master-port)
|
||||
MASTER_PORT="$2"
|
||||
shift 2
|
||||
;;
|
||||
--deploy-config)
|
||||
DEPLOY_CONFIG="$2"
|
||||
shift 2
|
||||
;;
|
||||
--help|-h)
|
||||
usage
|
||||
exit 0
|
||||
;;
|
||||
--)
|
||||
shift
|
||||
EXTRA_ARGS=("$@")
|
||||
break
|
||||
;;
|
||||
*)
|
||||
echo "Unknown option: $1" >&2
|
||||
usage
|
||||
exit 1
|
||||
;;
|
||||
esac
|
||||
done
|
||||
|
||||
if [[ "$STAGE" != "0" && "$STAGE" != "1" && "$STAGE" != "all" ]]; then
|
||||
echo "Invalid --stage value: $STAGE" >&2
|
||||
usage
|
||||
exit 1
|
||||
fi
|
||||
|
||||
print_config() {
|
||||
echo "Model: $MODEL"
|
||||
echo "API Port: $PORT"
|
||||
echo "Master Address: $MASTER_ADDRESS"
|
||||
echo "Master Port: $MASTER_PORT"
|
||||
echo "Deploy Config: $DEPLOY_CONFIG"
|
||||
echo "Selected Stage: $STAGE"
|
||||
if [[ ${#EXTRA_ARGS[@]} -gt 0 ]]; then
|
||||
echo "Extra Args: ${EXTRA_ARGS[*]}"
|
||||
fi
|
||||
}
|
||||
|
||||
run_stage_0() {
|
||||
echo "Starting Stage 0 (Thinker) as master..."
|
||||
vllm serve "$MODEL" --omni \
|
||||
--port "$PORT" \
|
||||
--deploy-config "$DEPLOY_CONFIG" \
|
||||
--stage-id 0 \
|
||||
--omni-master-address "$MASTER_ADDRESS" \
|
||||
--omni-master-port "$MASTER_PORT" \
|
||||
"${EXTRA_ARGS[@]}"
|
||||
}
|
||||
|
||||
run_stage_1() {
|
||||
echo "Starting Stage 1 (DiT) in headless mode..."
|
||||
vllm serve "$MODEL" --omni \
|
||||
--deploy-config "$DEPLOY_CONFIG" \
|
||||
--stage-id 1 \
|
||||
--headless \
|
||||
--omni-master-address "$MASTER_ADDRESS" \
|
||||
--omni-master-port "$MASTER_PORT" \
|
||||
"${EXTRA_ARGS[@]}"
|
||||
}
|
||||
|
||||
echo "Starting Bagel multi-stage server..."
|
||||
print_config
|
||||
|
||||
case "$STAGE" in
|
||||
0)
|
||||
run_stage_0
|
||||
;;
|
||||
1)
|
||||
run_stage_1
|
||||
;;
|
||||
all)
|
||||
echo "Launching both stages in one session (legacy mode)..."
|
||||
echo "Starting Stage 0 (Thinker) in background first..."
|
||||
run_stage_0 &
|
||||
STAGE_0_PID=$!
|
||||
|
||||
cleanup() {
|
||||
if [[ -n "${STAGE_0_PID:-}" ]]; then
|
||||
kill "$STAGE_0_PID" 2>/dev/null || true
|
||||
wait "$STAGE_0_PID" 2>/dev/null || true
|
||||
fi
|
||||
}
|
||||
|
||||
trap cleanup EXIT INT TERM
|
||||
|
||||
echo "Waiting briefly for Stage 0 to initialize..."
|
||||
sleep 2
|
||||
run_stage_1
|
||||
;;
|
||||
esac
|
||||
@@ -1,6 +1,8 @@
|
||||
# Image-To-Image
|
||||
|
||||
This example demonstrates how to deploy Qwen-Image-Edit model for online image editing service using vLLM-Omni.
|
||||
This example demonstrates how to deploy image-to-image models for online image editing service using vLLM-Omni.
|
||||
|
||||
Supported models include Qwen-Image-Edit, BAGEL, and other image-to-image pipelines.
|
||||
|
||||
For **multi-image** input editing, use **Qwen-Image-Edit-2509** (QwenImageEditPlusPipeline) and send multiple images in the user message content.
|
||||
|
||||
@@ -21,6 +23,12 @@ vllm serve Qwen/Qwen-Image-Edit --omni --port 8092
|
||||
vllm serve Qwen/Qwen-Image-Edit-2509 --omni --port 8092
|
||||
```
|
||||
|
||||
### BAGEL
|
||||
|
||||
```bash
|
||||
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091
|
||||
```
|
||||
|
||||
### Start with Parameters
|
||||
|
||||
|
||||
@@ -119,6 +127,16 @@ python openai_chat_client.py --input input.png --prompt "Convert to oil painting
|
||||
python openai_chat_client.py --input input1.png input2.png --prompt "Combine these images into a single scene" --output output.png
|
||||
```
|
||||
|
||||
Pass model-specific parameters through `--extra-body` (e.g. for BAGEL):
|
||||
|
||||
```bash
|
||||
python openai_chat_client.py \
|
||||
--input input.png \
|
||||
--prompt "Make the scene look like a watercolor painting" \
|
||||
--server http://localhost:8091 \
|
||||
--extra-body '{"cfg_text_scale": 4.0, "cfg_img_scale": 1.5}'
|
||||
```
|
||||
|
||||
### Method 4: Using Gradio Demo
|
||||
|
||||
```bash
|
||||
@@ -318,6 +336,9 @@ count, use `size` and `n` rather than `height`, `width`, or
|
||||
| `layers` | int | 4 | Number of layers (Qwen-Image-Layered) |
|
||||
| `resolution` | int | 640 | Resolution, 640 or 1024 (Qwen-Image-Layered) |
|
||||
|
||||
Models like BAGEL accept additional parameters via `extra_body` (e.g.
|
||||
`cfg_text_scale`, `cfg_img_scale`). See the [BAGEL recipe](../../../recipes/Bagel/BAGEL-7B-MoT.md) for the full list.
|
||||
|
||||
## Response Format
|
||||
|
||||
```json
|
||||
|
||||
@@ -38,6 +38,7 @@ def edit_image(
|
||||
guidance_scale: float | None = None,
|
||||
seed: int | None = None,
|
||||
negative_prompt: str | None = None,
|
||||
extra_body: dict | None = None,
|
||||
) -> bytes | None:
|
||||
"""Edit an image using the chat completions API.
|
||||
|
||||
@@ -51,6 +52,7 @@ def edit_image(
|
||||
guidance_scale: CFG guidance scale
|
||||
seed: Random seed
|
||||
negative_prompt: Negative prompt
|
||||
extra_body: Additional model-specific params (e.g. cfg_text_scale for BAGEL)
|
||||
|
||||
Returns:
|
||||
Edited image bytes or None if failed
|
||||
@@ -75,24 +77,26 @@ def edit_image(
|
||||
]
|
||||
|
||||
# Build extra_body with generation parameters
|
||||
extra_body = {}
|
||||
merged_extra_body: dict[str, object] = {}
|
||||
if height is not None:
|
||||
extra_body["height"] = height
|
||||
merged_extra_body["height"] = height
|
||||
if width is not None:
|
||||
extra_body["width"] = width
|
||||
merged_extra_body["width"] = width
|
||||
if steps is not None:
|
||||
extra_body["num_inference_steps"] = steps
|
||||
merged_extra_body["num_inference_steps"] = steps
|
||||
if guidance_scale is not None:
|
||||
extra_body["guidance_scale"] = guidance_scale
|
||||
merged_extra_body["guidance_scale"] = guidance_scale
|
||||
if seed is not None:
|
||||
extra_body["seed"] = seed
|
||||
merged_extra_body["seed"] = seed
|
||||
if negative_prompt:
|
||||
extra_body["negative_prompt"] = negative_prompt
|
||||
merged_extra_body["negative_prompt"] = negative_prompt
|
||||
if extra_body:
|
||||
merged_extra_body.update(extra_body)
|
||||
|
||||
# Build request payload
|
||||
payload = {"messages": messages}
|
||||
if extra_body:
|
||||
payload["extra_body"] = extra_body
|
||||
payload: dict[str, object] = {"messages": messages}
|
||||
if merged_extra_body:
|
||||
payload["extra_body"] = merged_extra_body
|
||||
|
||||
# Send request
|
||||
try:
|
||||
@@ -121,6 +125,19 @@ def edit_image(
|
||||
return None
|
||||
|
||||
|
||||
def parse_extra_body(value: str) -> dict:
|
||||
"""Parse a JSON string into a dict for --extra-body."""
|
||||
import json
|
||||
|
||||
try:
|
||||
obj = json.loads(value)
|
||||
except json.JSONDecodeError as e:
|
||||
raise argparse.ArgumentTypeError(f"--extra-body must be valid JSON: {e}") from e
|
||||
if not isinstance(obj, dict):
|
||||
raise argparse.ArgumentTypeError("--extra-body must be a JSON object")
|
||||
return obj
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Qwen-Image-Edit chat client")
|
||||
parser.add_argument("--input", "-i", required=True, nargs="+", help="Input image path(s)")
|
||||
@@ -133,6 +150,12 @@ def main():
|
||||
parser.add_argument("--guidance", type=float, default=7.5, help="Guidance scale")
|
||||
parser.add_argument("--seed", type=int, default=0, help="Random seed")
|
||||
parser.add_argument("--negative", help="Negative prompt")
|
||||
parser.add_argument(
|
||||
"--extra-body",
|
||||
type=parse_extra_body,
|
||||
default=None,
|
||||
help="JSON object merged into request extra_body, e.g. '{\"cfg_text_scale\": 4.0}'.",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
@@ -152,6 +175,7 @@ def main():
|
||||
guidance_scale=args.guidance,
|
||||
seed=args.seed,
|
||||
negative_prompt=args.negative,
|
||||
extra_body=args.extra_body,
|
||||
)
|
||||
|
||||
if image_bytes:
|
||||
|
||||
@@ -48,7 +48,7 @@ def generate(
|
||||
forwarded as top-level fields in the request payload. The serving layer
|
||||
maps standard keys (height, width, seed, num_inference_steps) to sampling
|
||||
params and forwards model-specific keys (think, cfg_norm, img_cfg_scale,
|
||||
…) to ``extra_args`` based on the pipeline's ``EXTRA_BODY_PARAMS``.
|
||||
…) to ``extra_args`` based on the model extra registry.
|
||||
|
||||
Returns:
|
||||
Image bytes (for image outputs) or text string (for text outputs).
|
||||
@@ -131,7 +131,7 @@ def main():
|
||||
parser.add_argument("--width", type=int, default=2048, help="Image width")
|
||||
parser.add_argument("--num-steps", type=int, default=50, help="Inference steps")
|
||||
parser.add_argument("--seed", type=int, default=42, help="Random seed")
|
||||
# Model-specific parameters (forwarded via EXTRA_BODY_PARAMS)
|
||||
# Model-specific parameters forwarded via the model extra registry.
|
||||
parser.add_argument("--cfg-scale", type=float, default=4.0, help="CFG scale")
|
||||
parser.add_argument("--img-cfg-scale", type=float, help="Image CFG scale (img2img)")
|
||||
parser.add_argument("--cfg-norm", type=str, help="CFG normalization mode")
|
||||
|
||||
@@ -0,0 +1,227 @@
|
||||
# BAGEL-7B-MoT
|
||||
|
||||
> BAGEL image generation through the shared online and offline image examples
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: ByteDance Seed
|
||||
- Model: `ByteDance-Seed/BAGEL-7B-MoT`
|
||||
- Task: Text-to-image and image-to-image generation
|
||||
- Mode: Offline inference and OpenAI-compatible online serving
|
||||
- Maintainer: Community
|
||||
|
||||
## When to use this recipe
|
||||
|
||||
Use this recipe when you want to run BAGEL through the shared image examples
|
||||
instead of model-specific example clients. The generic examples can format
|
||||
BAGEL text-to-image and image-to-image prompts, select the image output
|
||||
modality, attach reference images, and forward BAGEL-specific generation
|
||||
parameters through the pipeline-declared `extra_args` contract.
|
||||
|
||||
## References
|
||||
|
||||
- Upstream model:
|
||||
[`ByteDance-Seed/BAGEL-7B-MoT`](https://huggingface.co/ByteDance-Seed/BAGEL-7B-MoT)
|
||||
- Related offline example:
|
||||
[`examples/offline_inference/text_to_image/text_to_image.py`](../../examples/offline_inference/text_to_image/text_to_image.py)
|
||||
- Related offline image-to-image example:
|
||||
[`examples/offline_inference/image_to_image/image_edit.py`](../../examples/offline_inference/image_to_image/image_edit.py)
|
||||
- Related online example:
|
||||
[`examples/online_serving/text_to_image/openai_chat_client.py`](../../examples/online_serving/text_to_image/openai_chat_client.py)
|
||||
- Related online image-to-image example:
|
||||
[`examples/online_serving/image_to_image/openai_chat_client.py`](../../examples/online_serving/image_to_image/openai_chat_client.py)
|
||||
- Default deploy configs:
|
||||
[`vllm_omni/deploy/bagel.yaml`](../../vllm_omni/deploy/bagel.yaml),
|
||||
[`vllm_omni/deploy/bagel_single_stage.yaml`](../../vllm_omni/deploy/bagel_single_stage.yaml)
|
||||
|
||||
## Hardware Support
|
||||
|
||||
This recipe documents the CUDA layouts used by the in-repo BAGEL deploy
|
||||
configs. The default two-stage config shares one 80 GB GPU; for more headroom,
|
||||
move the diffusion stage to a second GPU in a custom deploy config.
|
||||
|
||||
## GPU
|
||||
|
||||
### 1x A100 80GB
|
||||
|
||||
#### Environment
|
||||
|
||||
- OS: Linux
|
||||
- Python: Match the repository requirements for your checkout
|
||||
- Driver / runtime: NVIDIA CUDA environment with one A100 80 GB GPU
|
||||
- vLLM version: Match the repository requirements for your checkout
|
||||
- vLLM-Omni version or commit: Use the commit you are deploying from
|
||||
|
||||
#### Offline Commands
|
||||
|
||||
Run text-to-image with the shared offline example from the repository root.
|
||||
Pick the topology with `--stage-configs-path` and forward BAGEL-specific
|
||||
generation parameters as a JSON object through `--extra-body`:
|
||||
|
||||
```bash
|
||||
# Two-stage (Thinker + DiT), shares one GPU by default
|
||||
python examples/offline_inference/text_to_image/text_to_image.py \
|
||||
--model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--stage-configs-path vllm_omni/deploy/bagel.yaml \
|
||||
--prompt "A beautiful sunset over mountains" \
|
||||
--height 1024 \
|
||||
--width 1024 \
|
||||
--num-inference-steps 50 \
|
||||
--negative-prompt "blurry, low quality" \
|
||||
--seed 42 \
|
||||
--extra-body '{"timestep_shift": 3.0, "cfg_text_scale": 4.0, "cfg_img_scale": 1.5, "cfg_interval": [0.4, 1.0], "cfg_renorm_type": "global", "cfg_renorm_min": 0.0, "think": false}' \
|
||||
--output /tmp/bagel_text2img.png
|
||||
```
|
||||
|
||||
```bash
|
||||
# Single-stage (DiT only, with internal LLM/ViT/VAE)
|
||||
python examples/offline_inference/text_to_image/text_to_image.py \
|
||||
--model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--stage-configs-path vllm_omni/deploy/bagel_single_stage.yaml \
|
||||
--prompt "A beautiful sunset over mountains" \
|
||||
--height 1024 \
|
||||
--width 1024 \
|
||||
--num-inference-steps 50 \
|
||||
--seed 42 \
|
||||
--extra-body '{"timestep_shift": 3.0, "cfg_text_scale": 4.0, "cfg_interval": [0.4, 1.0], "cfg_renorm_type": "global", "cfg_renorm_min": 0.0}' \
|
||||
--output /tmp/bagel_text2img_single.png
|
||||
```
|
||||
|
||||
The `--extra-body` JSON forwards BAGEL-specific parameters into
|
||||
`OmniDiffusionSamplingParams.extra_args`. Keys are filtered against the model's
|
||||
declared `extra_body_params` (see
|
||||
[`vllm_omni/model_extras/bagel.py`](../../vllm_omni/model_extras/bagel.py)), so
|
||||
unknown keys for BAGEL are silently dropped. This is the preferred way to pass
|
||||
model-specific params — `timestep_shift` (use `3.0` for 1024×1024; the
|
||||
deprecated `--timesteps-shift` flag defaults to `1.0` and is NextStep-oriented),
|
||||
`cfg_text_scale`, `cfg_img_scale`, `cfg_interval`, `cfg_renorm_type`,
|
||||
`cfg_renorm_min`, and `think`. `--seed` makes generation reproducible.
|
||||
|
||||
> **Note**: the two-stage `bagel.yaml` runs both stages on GPU 0. If you hit
|
||||
> CUDA OOM during warmup, lower stage-0 `gpu_memory_utilization` in the YAML or
|
||||
> split the stages across two GPUs (see the `2x CUDA GPUs` section).
|
||||
|
||||
Run image-to-image with the shared offline image-to-image example:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/image_to_image/image_edit.py \
|
||||
--model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--prompt "Make the scene look like a watercolor painting" \
|
||||
--image /tmp/bagel_text2img.png \
|
||||
--height 512 \
|
||||
--width 512 \
|
||||
--extra-args '{"cfg_text_scale": 4.0, "cfg_img_scale": 1.5}' \
|
||||
--output /tmp/bagel_img2img.png
|
||||
```
|
||||
|
||||
The `--extra-args` JSON forwards BAGEL-specific parameters (e.g. `cfg_text_scale`,
|
||||
`cfg_img_scale`, `cfg_interval`, `cfg_renorm_type`) into
|
||||
`OmniDiffusionSamplingParams.extra_args` via the model-extras registry.
|
||||
|
||||
#### Online Commands
|
||||
|
||||
Start the OpenAI-compatible server:
|
||||
|
||||
```bash
|
||||
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091
|
||||
```
|
||||
|
||||
To use the single-stage topology online:
|
||||
|
||||
```bash
|
||||
vllm serve ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--deploy-config vllm_omni/deploy/bagel_single_stage.yaml
|
||||
```
|
||||
|
||||
Send a text-to-image request with BAGEL-specific generation parameters:
|
||||
|
||||
```bash
|
||||
curl http://localhost:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "<|im_start|>A beautiful sunset over mountains<|im_end|>"}
|
||||
]
|
||||
}
|
||||
],
|
||||
"modalities": ["image"],
|
||||
"extra_body": {
|
||||
"height": 512,
|
||||
"width": 512,
|
||||
"num_inference_steps": 50,
|
||||
"cfg_text_scale": 4.0,
|
||||
"cfg_img_scale": 1.5,
|
||||
"negative_prompt": "blurry, low quality",
|
||||
"seed": 42
|
||||
}
|
||||
}'
|
||||
```
|
||||
|
||||
The important part is that BAGEL-specific keys such as `cfg_text_scale`,
|
||||
`cfg_img_scale`, `cfg_interval`, `cfg_renorm_type`, and `cfg_renorm_min` belong
|
||||
in `extra_body`. The serving layer routes the declared keys into
|
||||
`OmniDiffusionSamplingParams.extra_args` for the BAGEL pipeline.
|
||||
|
||||
#### Verification
|
||||
|
||||
Decode the returned data URL into an image:
|
||||
|
||||
```bash
|
||||
curl http://localhost:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "<|im_start|>A ceramic teapot on a wooden table<|im_end|>"}
|
||||
]
|
||||
}
|
||||
],
|
||||
"modalities": ["image"],
|
||||
"extra_body": {
|
||||
"height": 512,
|
||||
"width": 512,
|
||||
"num_inference_steps": 25,
|
||||
"cfg_text_scale": 4.0,
|
||||
"seed": 42
|
||||
}
|
||||
}' | jq -r '.choices[]?.message.content[]? | select(.image_url.url) | .image_url.url' \
|
||||
| head -n 1 \
|
||||
| cut -d',' -f2- \
|
||||
| base64 -d > /tmp/bagel_online.png
|
||||
|
||||
ls -lh /tmp/bagel_online.png
|
||||
```
|
||||
|
||||
### 2x CUDA GPUs
|
||||
|
||||
Create a custom deploy config from `vllm_omni/deploy/bagel.yaml` and move the
|
||||
diffusion stage to GPU 1:
|
||||
|
||||
```yaml
|
||||
stages:
|
||||
- stage_id: 0
|
||||
devices: "0"
|
||||
# keep the remaining stage-0 settings from bagel.yaml
|
||||
- stage_id: 1
|
||||
devices: "1"
|
||||
# keep the remaining stage-1 settings from bagel.yaml
|
||||
```
|
||||
|
||||
Then start serving with that config:
|
||||
|
||||
```bash
|
||||
vllm serve ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--deploy-config /path/to/custom_bagel_2gpu.yaml
|
||||
```
|
||||
|
||||
Use the online curl request from the `1x A100 80GB` section to verify that the
|
||||
server returns an image.
|
||||
@@ -29,6 +29,7 @@ recipes/
|
||||
|--------|------|----------|
|
||||
| [`audiox/AudioX.md`](./audiox/AudioX.md) | Offline + online unified text/video→audio diffusion | 1x L4 24GB |
|
||||
| [`Baidu/ERNIE-Image.md`](./Baidu/ERNIE-Image.md) | Text-to-image online serving (ERNIE-Image 8B) | 1x or 2x RTX 4090 24GB |
|
||||
| [`Bagel/BAGEL-7B-MoT.md`](./Bagel/BAGEL-7B-MoT.md) | Text-to-image with shared online/offline examples | 1x A100 80GB / 2x CUDA GPUs |
|
||||
| [`ByteDance/Lance.md`](./ByteDance/Lance.md) | Unified AR+diffusion: text/img/video gen + understanding (Lance 3B) | 1x B300 / A100 80GB |
|
||||
| [`fishaudio/Fish-Speech-S2-Pro.md`](./fishaudio/Fish-Speech-S2-Pro.md) | Online serving for TTS | 1x A800 80GB |
|
||||
| [`Helios/Helios.md`](./Helios/Helios.md) | Text-to-video, image-to-video, and video-to-video generation | 1x NVIDIA H20 |
|
||||
|
||||
@@ -72,6 +72,31 @@ pytest -s -v tests/e2e/offline_inference/test_sensenova_u1_text2img.py \
|
||||
- No deploy YAML needed — the engine auto-generates a single-stage diffusion config.
|
||||
- Think mode (`--think`) is recommended for higher image quality.
|
||||
|
||||
#### Text-to-Image via the shared example
|
||||
|
||||
SenseNova-U1 is also registered with the shared text-to-image example, so you
|
||||
can run it through `text_to_image.py` instead of `end2end.py`. Forward the
|
||||
SenseNova-specific generation parameters as a JSON object through `--extra-body`:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_image/text_to_image.py \
|
||||
--model SenseNova/SenseNova-U1-8B-MoT \
|
||||
--prompt "A beautiful sunset over mountains" \
|
||||
--width 2048 --height 2048 \
|
||||
--num-inference-steps 50 \
|
||||
--seed 42 \
|
||||
--extra-body '{"think": true, "cfg_scale": 4.0, "cfg_norm": "none", "timestep_shift": 3.0, "t_eps": 0.02}' \
|
||||
--output sensenova_text2img.png
|
||||
```
|
||||
|
||||
The `--extra-body` keys are filtered against the model's declared
|
||||
`extra_body_params` (see
|
||||
[`vllm_omni/model_extras/sensenova_u1.py`](../../vllm_omni/model_extras/sensenova_u1.py)):
|
||||
`think`, `cfg_scale`, `cfg_norm`, `timestep_shift`, `t_eps`, `img_cfg_scale`,
|
||||
and `max_tokens`. No deploy YAML is needed — the engine auto-generates a
|
||||
single-stage diffusion config. For richer task modalities (img2text, text2text)
|
||||
and CPU offload, use the dedicated `end2end.py` script described above.
|
||||
|
||||
#### Image-to-Image Editing (img2img)
|
||||
|
||||
```bash
|
||||
|
||||
@@ -0,0 +1,157 @@
|
||||
# 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 vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.model_extras import (
|
||||
build_image_to_image_prompt,
|
||||
build_text_to_image_prompt,
|
||||
get_extra_body_params,
|
||||
get_extra_output_params,
|
||||
should_init_extra_args_for_non_diffusion_stages,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_bagel_extra_registry_declares_request_and_response_params() -> None:
|
||||
assert get_extra_body_params("BagelPipeline") == frozenset(
|
||||
{
|
||||
"cfg_text_scale",
|
||||
"cfg_img_scale",
|
||||
"cfg_interval",
|
||||
"cfg_renorm_type",
|
||||
"cfg_renorm_min",
|
||||
"negative_prompt",
|
||||
"think",
|
||||
"max_think_tokens",
|
||||
"do_sample",
|
||||
"text_temperature",
|
||||
"timestep_shift",
|
||||
}
|
||||
)
|
||||
assert get_extra_output_params("BagelPipeline") == frozenset({"text_output", "think_text"})
|
||||
assert should_init_extra_args_for_non_diffusion_stages("BagelPipeline") is True
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_sensenova_extra_registry_declares_request_and_response_params() -> None:
|
||||
assert get_extra_body_params("SenseNovaU1Pipeline") == frozenset(
|
||||
{
|
||||
"think",
|
||||
"cfg_scale",
|
||||
"cfg_norm",
|
||||
"timestep_shift",
|
||||
"t_eps",
|
||||
"img_cfg_scale",
|
||||
"max_tokens",
|
||||
}
|
||||
)
|
||||
assert get_extra_output_params("SenseNovaU1Pipeline") == frozenset({"think_text"})
|
||||
assert should_init_extra_args_for_non_diffusion_stages("SenseNovaU1Pipeline") is False
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_unknown_pipeline_has_empty_extra_registry() -> None:
|
||||
assert get_extra_body_params("UnknownPipeline") == frozenset()
|
||||
assert get_extra_output_params("UnknownPipeline") == frozenset()
|
||||
assert should_init_extra_args_for_non_diffusion_stages("UnknownPipeline") is False
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_bagel_text_to_image_prompt_builder() -> None:
|
||||
assert build_text_to_image_prompt(
|
||||
"BagelPipeline",
|
||||
prompt="a cat",
|
||||
negative_prompt="blurry",
|
||||
height=512,
|
||||
width=768,
|
||||
) == {
|
||||
"prompt": "<|im_start|>a cat<|im_end|>",
|
||||
"modalities": ["image"],
|
||||
"mm_processor_kwargs": {
|
||||
"target_h": 512,
|
||||
"target_w": 768,
|
||||
"modalities": ["image"],
|
||||
},
|
||||
"negative_prompt": "blurry",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_bagel_image_to_image_prompt_builder() -> None:
|
||||
dummy_image = Image.new("RGB", (64, 64))
|
||||
result = build_image_to_image_prompt(
|
||||
"BagelPipeline",
|
||||
prompt="paint it",
|
||||
negative_prompt="ugly",
|
||||
input_image=dummy_image,
|
||||
height=256,
|
||||
width=256,
|
||||
)
|
||||
assert result["prompt"] == "<|fim_middle|><|im_start|>paint it<|im_end|>"
|
||||
assert result["modalities"] == ["img2img"]
|
||||
assert result["multi_modal_data"]["img2img"] is dummy_image
|
||||
assert result["mm_processor_kwargs"]["target_h"] == 256
|
||||
assert result["mm_processor_kwargs"]["target_w"] == 256
|
||||
assert result["negative_prompt"] == "ugly"
|
||||
|
||||
|
||||
@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"}
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_unknown_pipeline_uses_default_image_to_image_prompt() -> None:
|
||||
dummy_image = Image.new("RGB", (64, 64))
|
||||
result = build_image_to_image_prompt(
|
||||
"UnknownPipeline",
|
||||
prompt="edit",
|
||||
negative_prompt=None,
|
||||
input_image=dummy_image,
|
||||
)
|
||||
assert result == {
|
||||
"prompt": "edit",
|
||||
"multi_modal_data": {"image": dummy_image},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_declared_extra_args_apply_to_existing_sampling_params() -> None:
|
||||
params = OmniDiffusionSamplingParams(extra_args={"existing": 1})
|
||||
|
||||
declared_extra_params: frozenset[str] = frozenset({"cfg_text_scale", "think"})
|
||||
apply_declared_extra_args(
|
||||
params,
|
||||
declared_extra_params,
|
||||
{
|
||||
"cfg_text_scale": 4.0,
|
||||
"think": False,
|
||||
"unknown": "ignored",
|
||||
},
|
||||
)
|
||||
|
||||
assert params.extra_args == {
|
||||
"existing": 1,
|
||||
"cfg_text_scale": 4.0,
|
||||
"think": False,
|
||||
}
|
||||
@@ -423,6 +423,67 @@ class DiffusionCacheConfig:
|
||||
raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{item}'")
|
||||
|
||||
|
||||
def resolve_model_class_name(model: str | None, diffusion_load_format: str = "default") -> str | None:
|
||||
"""Resolve the diffusion pipeline class name from the model config.
|
||||
|
||||
Read-only counterpart of ``OmniDiffusionConfig.enrich_config``, safe to call
|
||||
client-side. Returns ``None`` if the pipeline can't be determined.
|
||||
"""
|
||||
from vllm.transformers_utils.config import get_hf_file_to_dict
|
||||
|
||||
if not model:
|
||||
return None
|
||||
|
||||
is_lance_subfolder = os.path.basename(str(model).rstrip("/")) in {"Lance_3B", "Lance_3B_Video"}
|
||||
|
||||
# Diffusers models: read _class_name from model_index.json.
|
||||
try:
|
||||
model_index = get_hf_file_to_dict("model_index.json", model)
|
||||
except Exception:
|
||||
model_index = None
|
||||
if model_index is not None:
|
||||
return model_index.get("_class_name")
|
||||
if diffusion_load_format == "diffusers":
|
||||
return "DiffusersAdapterPipeline"
|
||||
|
||||
# Other models: map model_type / architecture from config.json.
|
||||
try:
|
||||
cfg = get_hf_file_to_dict("config.json", model) or {}
|
||||
except Exception:
|
||||
cfg = {}
|
||||
model_type = cfg.get("model_type")
|
||||
architectures = cfg.get("architectures") or []
|
||||
|
||||
if model_type == "bagel" or "BagelForConditionalGeneration" in architectures:
|
||||
return "BagelPipeline"
|
||||
if (
|
||||
model_type == "lance"
|
||||
or "LancePipeline" in architectures
|
||||
or cfg.get("model_name") == "Lance"
|
||||
or is_lance_subfolder
|
||||
):
|
||||
return "LancePipeline"
|
||||
if model_type == "neo_chat":
|
||||
return "SenseNovaU1Pipeline"
|
||||
if "BailingMM2NativeForConditionalGeneration" in architectures or model_type in (
|
||||
"bailingmm_moe_v2_lite",
|
||||
"ming_flash_omni",
|
||||
"ming_flash_omni_thinker",
|
||||
):
|
||||
return "MingImagePipeline"
|
||||
if model_type == "nextstep":
|
||||
return "NextStep11Pipeline"
|
||||
if model_type == "s2v":
|
||||
return "WanS2VPipeline"
|
||||
if model_type == "vla":
|
||||
from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero
|
||||
|
||||
return "DreamZeroPipeline" if _looks_like_dreamzero(model) else None
|
||||
if len(architectures) == 1:
|
||||
return architectures[0]
|
||||
return None
|
||||
|
||||
|
||||
@dataclass
|
||||
class OmniDiffusionConfig:
|
||||
# Model and path configuration (for convenience)
|
||||
|
||||
@@ -27,8 +27,6 @@ from vllm_omni.diffusion.data import (
|
||||
from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
|
||||
from vllm_omni.diffusion.io_support import (
|
||||
get_dummy_run_num_frames,
|
||||
get_extra_body_params,
|
||||
get_extra_output_params,
|
||||
image_color_format,
|
||||
supports_audio_output,
|
||||
supports_multimodal_input,
|
||||
@@ -61,8 +59,6 @@ __all__ = [
|
||||
"_RpcTask",
|
||||
"_move_tensor_tree_to_cpu",
|
||||
"get_dummy_run_num_frames",
|
||||
"get_extra_body_params",
|
||||
"get_extra_output_params",
|
||||
"image_color_format",
|
||||
"supports_audio_output",
|
||||
"supports_multimodal_input",
|
||||
|
||||
@@ -47,31 +47,3 @@ def get_dummy_run_num_frames(model_class_name: str, supports_audio_input: bool)
|
||||
if model_cls is not None and hasattr(model_cls, "dummy_run_num_frames"):
|
||||
return int(getattr(model_cls, "dummy_run_num_frames"))
|
||||
return 2 if supports_audio_input or supports_audio_output(model_class_name) else 1
|
||||
|
||||
|
||||
def get_extra_body_params(model_class_name: str) -> frozenset[str]:
|
||||
"""Return the set of extra_body keys accepted by a pipeline.
|
||||
|
||||
Each pipeline can declare ``EXTRA_BODY_PARAMS: ClassVar[frozenset[str]]``
|
||||
to advertise which request-level parameters should be forwarded from
|
||||
``extra_body`` to ``OmniDiffusionSamplingParams.extra_args``.
|
||||
Returns an empty frozenset when the pipeline does not declare any.
|
||||
"""
|
||||
model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name)
|
||||
if model_cls is None:
|
||||
return frozenset()
|
||||
return frozenset(getattr(model_cls, "EXTRA_BODY_PARAMS", frozenset()))
|
||||
|
||||
|
||||
def get_extra_output_params(model_class_name: str) -> frozenset[str]:
|
||||
"""Return the set of custom_output keys to expose in API response metrics.
|
||||
|
||||
Each pipeline can declare ``EXTRA_OUTPUT_PARAMS: ClassVar[frozenset[str]]``
|
||||
to advertise which ``DiffusionOutput.custom_output`` keys should be
|
||||
copied into the response ``metrics`` dict.
|
||||
Returns an empty frozenset when the pipeline does not declare any.
|
||||
"""
|
||||
model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name)
|
||||
if model_cls is None:
|
||||
return frozenset()
|
||||
return frozenset(getattr(model_cls, "EXTRA_OUTPUT_PARAMS", frozenset()))
|
||||
@@ -504,30 +504,6 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
|
||||
|
||||
support_image_input = True
|
||||
|
||||
# Model-specific parameters accepted via ``extra_body`` in the online
|
||||
# serving API. The serving layer reads this set and forwards matching
|
||||
# keys from the request payload to ``OmniDiffusionSamplingParams.extra_args``.
|
||||
EXTRA_BODY_PARAMS: ClassVar[frozenset[str]] = frozenset(
|
||||
{
|
||||
"think",
|
||||
"cfg_scale",
|
||||
"cfg_norm",
|
||||
"timestep_shift",
|
||||
"t_eps",
|
||||
"img_cfg_scale",
|
||||
"max_tokens",
|
||||
}
|
||||
)
|
||||
|
||||
# Keys from ``DiffusionOutput.custom_output`` that should be surfaced in
|
||||
# the API response ``metrics`` dict. The serving layer reads this set and
|
||||
# copies matching entries from ``custom_output`` into the response.
|
||||
EXTRA_OUTPUT_PARAMS: ClassVar[frozenset[str]] = frozenset(
|
||||
{
|
||||
"think_text",
|
||||
}
|
||||
)
|
||||
|
||||
# CPU-offload protocol: language_model carries the denoising blocks; the
|
||||
# vision and FM modules are lightweight encoders pinned on GPU during the
|
||||
# diffusion loop. There is no separate VAE.
|
||||
|
||||
@@ -0,0 +1,27 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
|
||||
def apply_declared_extra_args(
|
||||
sampling_params: OmniDiffusionSamplingParams,
|
||||
declared_params: frozenset[str],
|
||||
user_kwargs: dict[str, object],
|
||||
) -> None:
|
||||
"""Route pipeline-declared request params into ``sampling_params.extra_args``.
|
||||
|
||||
Both online serving and offline examples call this so that model-specific
|
||||
keys (e.g. ``cfg_text_scale`` for BAGEL) end up in ``extra_args`` instead
|
||||
of being silently dropped.
|
||||
|
||||
This is a no-op when no declared params are present in ``user_kwargs``, so
|
||||
it is safe to call on non-diffusion (e.g. AR) sampling params whose
|
||||
``extra_args`` defaults to ``None``.
|
||||
"""
|
||||
declared = {key: user_kwargs[key] for key in declared_params if user_kwargs.get(key) is not None}
|
||||
if not declared:
|
||||
return
|
||||
sampling_params.extra_args = {**(sampling_params.extra_args or {}), **declared}
|
||||
@@ -220,6 +220,8 @@ class AsyncOmniEngine:
|
||||
self.model = model
|
||||
self.tokenizer = tokenizer
|
||||
self.diffusion_batch_size = diffusion_batch_size
|
||||
# Cached by get_diffusion_od_config().
|
||||
self._diffusion_od_config_view: Any = None
|
||||
startup_timeout = int(init_timeout)
|
||||
# Forwarded into Orchestrator so its _forward_to_next_stage path can
|
||||
# emit per-edge transfer_tx_s / transfer_size_bytes histograms.
|
||||
@@ -323,6 +325,20 @@ class AsyncOmniEngine:
|
||||
|
||||
logger.info(f"[AsyncOmniEngine] Orchestrator ready with {self.num_stages} stages")
|
||||
|
||||
def get_diffusion_od_config(self) -> Any:
|
||||
"""Expose the diffusion ``model_class_name`` to client-side model-extras.
|
||||
|
||||
The worker holds the full config; here we just resolve the pipeline class
|
||||
name from the model config (cached). ``model_class_name`` may be ``None``.
|
||||
"""
|
||||
if self._diffusion_od_config_view is None:
|
||||
from types import SimpleNamespace
|
||||
|
||||
from vllm_omni.diffusion.data import resolve_model_class_name
|
||||
|
||||
self._diffusion_od_config_view = SimpleNamespace(model_class_name=resolve_model_class_name(self.model))
|
||||
return self._diffusion_od_config_view
|
||||
|
||||
def _initialize_stages(self, stage_init_timeout: int) -> None:
|
||||
"""Initialize stage clients/processors via StageRuntime and assign to self."""
|
||||
self._runtime = create_stage_runtime(
|
||||
|
||||
@@ -23,7 +23,7 @@ from vllm.entrypoints.chat_utils import (
|
||||
make_tool_call_id,
|
||||
)
|
||||
|
||||
from vllm_omni.diffusion.io_support import get_extra_body_params, get_extra_output_params
|
||||
from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
|
||||
from vllm_omni.entrypoints.async_omni import AsyncOmni
|
||||
from vllm_omni.entrypoints.openai.protocol.chat_completion import OmniChatCompletionResponse
|
||||
from vllm_omni.entrypoints.utils import coerce_param_message_types
|
||||
@@ -33,6 +33,7 @@ from vllm_omni.metrics.modality import (
|
||||
observe_audio_first_packet,
|
||||
observe_audio_streaming_finalize,
|
||||
)
|
||||
from vllm_omni.model_extras import get_extra_body_params, get_extra_output_params
|
||||
|
||||
try:
|
||||
import soundfile
|
||||
@@ -162,12 +163,16 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
instance._diffusion_mode = True
|
||||
instance._diffusion_engine = diffusion_engine
|
||||
instance._diffusion_model_name = model_name
|
||||
instance._diffusion_extra_body_params = None
|
||||
instance._diffusion_extra_output_params = None
|
||||
instance.engine_client = None
|
||||
instance.has_kv_connector = False
|
||||
# Extra body/output params are resolved lazily on first use; see
|
||||
# _get_diffusion_extra_body_params / _get_diffusion_extra_output_params.
|
||||
return instance
|
||||
|
||||
def _get_diffusion_extra_body_params(self) -> frozenset[str]:
|
||||
"""Return the pipeline-declared ``EXTRA_BODY_PARAMS`` (cached)."""
|
||||
"""Return model-specific extra_body params from the extra registry."""
|
||||
if self._diffusion_extra_body_params is not None:
|
||||
return self._diffusion_extra_body_params
|
||||
|
||||
@@ -177,7 +182,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
if od_config is not None and getattr(od_config, "model_class_name", None):
|
||||
params = get_extra_body_params(od_config.model_class_name)
|
||||
except Exception as e:
|
||||
logger.warning("Failed to read EXTRA_BODY_PARAMS from pipeline: %s", e)
|
||||
logger.warning("Failed to read model extra_body params: %s", e)
|
||||
|
||||
self._diffusion_extra_body_params = params
|
||||
return params
|
||||
@@ -186,7 +191,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
self,
|
||||
custom_output: dict[str, Any] | None,
|
||||
) -> dict[str, Any] | None:
|
||||
"""Pick pipeline-declared ``EXTRA_OUTPUT_PARAMS`` keys from *custom_output*."""
|
||||
"""Pick model-specific extra output keys from *custom_output*."""
|
||||
if not custom_output:
|
||||
return None
|
||||
|
||||
@@ -197,7 +202,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
if od_config is not None and getattr(od_config, "model_class_name", None):
|
||||
params = get_extra_output_params(od_config.model_class_name)
|
||||
except Exception as e:
|
||||
logger.warning("Failed to read EXTRA_OUTPUT_PARAMS from pipeline: %s", e)
|
||||
logger.warning("Failed to read model extra output params: %s", e)
|
||||
self._diffusion_extra_output_params = params
|
||||
|
||||
if not self._diffusion_extra_output_params:
|
||||
@@ -397,8 +402,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
request.modalities = output_modalities if output_modalities is not None else engine_output_modalities
|
||||
|
||||
num_inference_steps = None
|
||||
cfg_text_scale = None
|
||||
cfg_img_scale = None
|
||||
extra_body: dict[str, Any] = {}
|
||||
# Omni multistage image generation: Stage-0 (AR) should receive a clean
|
||||
# text prompt (and optional conditioning image/size) so the model's own
|
||||
# processor can construct the correct inputs.
|
||||
@@ -416,9 +420,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
# `extra_body` is flattented and merged into the payload's root.
|
||||
# These extra fields are accessible via `model_extra` property (from Pydantic base class).
|
||||
# When sending raw request with curl, no flattening happens. Directly read the `extra_body` dict.
|
||||
extra_body = getattr(request, "extra_body", None)
|
||||
if not extra_body:
|
||||
extra_body = request.model_extra or {}
|
||||
extra_body = getattr(request, "extra_body", None) or request.model_extra or {}
|
||||
|
||||
height, width = self._resolve_height_width_from_extra_body(extra_body)
|
||||
|
||||
@@ -430,9 +432,6 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
num_inference_steps = None
|
||||
|
||||
negative_prompt = extra_body.get("negative_prompt")
|
||||
cfg_text_scale = extra_body.get("cfg_text_scale")
|
||||
cfg_img_scale = extra_body.get("cfg_img_scale")
|
||||
|
||||
engine_prompt_image: dict[str, Any] | None = None
|
||||
if reference_images:
|
||||
# Best-effort decode first reference image for i2i.
|
||||
@@ -559,15 +558,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
sp.width = _image_gen_width
|
||||
if hasattr(sp, "num_inference_steps") and num_inference_steps is not None:
|
||||
sp.num_inference_steps = num_inference_steps
|
||||
if hasattr(sp, "extra_args") and sp.extra_args is not None:
|
||||
if cfg_text_scale is not None:
|
||||
sp.extra_args["cfg_text_scale"] = cfg_text_scale
|
||||
if cfg_img_scale is not None:
|
||||
sp.extra_args["cfg_img_scale"] = cfg_img_scale
|
||||
for _key in self._get_diffusion_extra_body_params():
|
||||
_val = extra_body.get(_key)
|
||||
if _val is not None:
|
||||
sp.extra_args[_key] = _val
|
||||
apply_declared_extra_args(sp, self._get_diffusion_extra_body_params(), extra_body)
|
||||
|
||||
self._log_inputs(
|
||||
request_id,
|
||||
@@ -2596,6 +2587,11 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
layers=layers,
|
||||
resolution=resolution,
|
||||
)
|
||||
apply_declared_extra_args(
|
||||
default_stage_params,
|
||||
self._get_diffusion_extra_body_params(),
|
||||
extra_body,
|
||||
)
|
||||
if lora_body and isinstance(lora_body, dict):
|
||||
try:
|
||||
lora_req, lora_scale = parse_lora_request(lora_body)
|
||||
@@ -2990,8 +2986,6 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
num_inference_steps = extra_body.get("num_inference_steps")
|
||||
guidance_scale = extra_body.get("guidance_scale")
|
||||
true_cfg_scale = extra_body.get("true_cfg_scale") or extra_body.get("cfg_scale")
|
||||
cfg_text_scale = extra_body.get("cfg_text_scale")
|
||||
cfg_img_scale = extra_body.get("cfg_img_scale")
|
||||
seed = extra_body.get("seed")
|
||||
if seed is None:
|
||||
seed = getattr(request, "seed", None)
|
||||
@@ -3049,14 +3043,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
|
||||
gen_params.guidance_scale = guidance_scale
|
||||
if true_cfg_scale is not None:
|
||||
gen_params.true_cfg_scale = true_cfg_scale
|
||||
if cfg_text_scale is not None:
|
||||
gen_params.extra_args["cfg_text_scale"] = cfg_text_scale
|
||||
if cfg_img_scale is not None:
|
||||
gen_params.extra_args["cfg_img_scale"] = cfg_img_scale
|
||||
for _key in self._get_diffusion_extra_body_params():
|
||||
_val = extra_body.get(_key)
|
||||
if _val is not None:
|
||||
gen_params.extra_args[_key] = _val
|
||||
apply_declared_extra_args(gen_params, self._get_diffusion_extra_body_params(), extra_body)
|
||||
if num_frames is not None:
|
||||
gen_params.num_frames = num_frames
|
||||
if guidance_scale_2 is not None:
|
||||
|
||||
@@ -0,0 +1,20 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from vllm_omni.model_extras.registry import (
|
||||
build_image_to_image_prompt,
|
||||
build_text_to_image_prompt,
|
||||
get_extra_body_params,
|
||||
get_extra_output_params,
|
||||
get_model_class_name,
|
||||
should_init_extra_args_for_non_diffusion_stages,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"build_image_to_image_prompt",
|
||||
"build_text_to_image_prompt",
|
||||
"get_extra_body_params",
|
||||
"get_extra_output_params",
|
||||
"get_model_class_name",
|
||||
"should_init_extra_args_for_non_diffusion_stages",
|
||||
]
|
||||
@@ -0,0 +1,74 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from PIL import Image
|
||||
|
||||
|
||||
def build_text_to_image_prompt(
|
||||
prompt: str,
|
||||
negative_prompt: str | None,
|
||||
height: int | None = None,
|
||||
width: int | None = None,
|
||||
) -> dict[str, Any]:
|
||||
text_prompt: dict[str, Any] = {
|
||||
"prompt": f"<|im_start|>{prompt}<|im_end|>",
|
||||
"modalities": ["image"],
|
||||
"mm_processor_kwargs": {
|
||||
"target_h": height,
|
||||
"target_w": width,
|
||||
"modalities": ["image"],
|
||||
},
|
||||
}
|
||||
if negative_prompt is not None:
|
||||
text_prompt["negative_prompt"] = negative_prompt
|
||||
return text_prompt
|
||||
|
||||
|
||||
def build_image_to_image_prompt(
|
||||
prompt: str,
|
||||
negative_prompt: str | None,
|
||||
input_image: Image.Image | list[Image.Image],
|
||||
height: int | None = None,
|
||||
width: int | None = None,
|
||||
) -> dict[str, Any]:
|
||||
img_prompt: dict[str, Any] = {
|
||||
"prompt": f"<|fim_middle|><|im_start|>{prompt}<|im_end|>",
|
||||
"modalities": ["img2img"],
|
||||
"multi_modal_data": {"img2img": input_image},
|
||||
"mm_processor_kwargs": {"modalities": ["img2img"]},
|
||||
}
|
||||
if height is not None:
|
||||
img_prompt["mm_processor_kwargs"]["target_h"] = height
|
||||
if width is not None:
|
||||
img_prompt["mm_processor_kwargs"]["target_w"] = width
|
||||
if negative_prompt is not None:
|
||||
img_prompt["negative_prompt"] = negative_prompt
|
||||
return img_prompt
|
||||
|
||||
|
||||
BAGEL_EXTRA_BODY_PARAMS = frozenset(
|
||||
{
|
||||
"cfg_text_scale",
|
||||
"cfg_img_scale",
|
||||
"cfg_interval",
|
||||
"cfg_renorm_type",
|
||||
"cfg_renorm_min",
|
||||
"negative_prompt",
|
||||
"think",
|
||||
"max_think_tokens",
|
||||
"do_sample",
|
||||
"text_temperature",
|
||||
"timestep_shift",
|
||||
}
|
||||
)
|
||||
BAGEL_EXTRA_OUTPUT_PARAMS = frozenset(
|
||||
{
|
||||
"text_output",
|
||||
"think_text",
|
||||
}
|
||||
)
|
||||
BAGEL_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES = True
|
||||
@@ -0,0 +1,147 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
from PIL import Image
|
||||
|
||||
from vllm_omni.model_extras.bagel import (
|
||||
BAGEL_EXTRA_BODY_PARAMS,
|
||||
BAGEL_EXTRA_OUTPUT_PARAMS,
|
||||
BAGEL_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES,
|
||||
)
|
||||
from vllm_omni.model_extras.bagel import (
|
||||
build_image_to_image_prompt as build_bagel_image_to_image_prompt,
|
||||
)
|
||||
from vllm_omni.model_extras.bagel import (
|
||||
build_text_to_image_prompt as build_bagel_text_to_image_prompt,
|
||||
)
|
||||
from vllm_omni.model_extras.sensenova_u1 import (
|
||||
SENSENOVA_U1_EXTRA_BODY_PARAMS,
|
||||
SENSENOVA_U1_EXTRA_OUTPUT_PARAMS,
|
||||
)
|
||||
|
||||
# TODO: extend with TextToVideoPromptBuilder, ImageToVideoPromptBuilder
|
||||
# once a video model needs model-specific prompt construction (follow-up PR).
|
||||
TextToImagePromptBuilder = Callable[
|
||||
[str, str | None, int | None, int | None],
|
||||
dict[str, Any],
|
||||
]
|
||||
ImageToImagePromptBuilder = Callable[
|
||||
[str, str | None, "Image.Image | list[Image.Image]", int | None, int | None],
|
||||
dict[str, Any],
|
||||
]
|
||||
|
||||
|
||||
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,
|
||||
input_image: Image.Image | list[Image.Image],
|
||||
height: int | None = None,
|
||||
width: int | None = None,
|
||||
) -> dict[str, Any]:
|
||||
result: dict[str, Any] = {
|
||||
"prompt": prompt,
|
||||
"multi_modal_data": {"image": input_image},
|
||||
}
|
||||
if negative_prompt is not None:
|
||||
result["negative_prompt"] = negative_prompt
|
||||
return result
|
||||
|
||||
|
||||
_EXTRA_SPECS: dict[str, dict[str, Any]] = {
|
||||
"BagelPipeline": {
|
||||
"extra_body_params": BAGEL_EXTRA_BODY_PARAMS,
|
||||
"extra_output_params": BAGEL_EXTRA_OUTPUT_PARAMS,
|
||||
"init_extra_args_for_non_diffusion_stages": BAGEL_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES,
|
||||
"text_to_image_prompt_builder": build_bagel_text_to_image_prompt,
|
||||
"image_to_image_prompt_builder": build_bagel_image_to_image_prompt,
|
||||
},
|
||||
"SenseNovaU1Pipeline": {
|
||||
"extra_body_params": SENSENOVA_U1_EXTRA_BODY_PARAMS,
|
||||
"extra_output_params": SENSENOVA_U1_EXTRA_OUTPUT_PARAMS,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _get_spec(model_class_name: str | None) -> dict[str, Any] | None:
|
||||
if not model_class_name:
|
||||
return None
|
||||
return _EXTRA_SPECS.get(model_class_name)
|
||||
|
||||
|
||||
def get_model_class_name(omni: Any) -> str | None:
|
||||
"""Extract model_class_name from an Omni/AsyncOmni instance.
|
||||
|
||||
This hides the internal ODConfig plumbing from example scripts.
|
||||
"""
|
||||
engine = getattr(omni, "engine", None)
|
||||
if engine is None:
|
||||
return None
|
||||
od_config = getattr(engine, "od_config", None)
|
||||
if od_config is None and hasattr(engine, "get_diffusion_od_config"):
|
||||
od_config = engine.get_diffusion_od_config()
|
||||
return getattr(od_config, "model_class_name", None) if od_config else None
|
||||
|
||||
|
||||
def get_extra_body_params(model_class_name: str | None) -> frozenset[str]:
|
||||
spec = _get_spec(model_class_name)
|
||||
return spec.get("extra_body_params", frozenset()) if spec is not None else frozenset()
|
||||
|
||||
|
||||
def get_extra_output_params(model_class_name: str | None) -> frozenset[str]:
|
||||
spec = _get_spec(model_class_name)
|
||||
return spec.get("extra_output_params", frozenset()) if spec is not None else frozenset()
|
||||
|
||||
|
||||
def should_init_extra_args_for_non_diffusion_stages(model_class_name: str | None) -> bool:
|
||||
spec = _get_spec(model_class_name)
|
||||
return bool(spec and spec.get("init_extra_args_for_non_diffusion_stages", False))
|
||||
|
||||
|
||||
def build_text_to_image_prompt(
|
||||
model_class_name: str | None,
|
||||
prompt: str,
|
||||
negative_prompt: str | None,
|
||||
height: int | None = None,
|
||||
width: int | None = None,
|
||||
) -> dict[str, Any]:
|
||||
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
|
||||
)
|
||||
return builder(prompt, negative_prompt, height, width)
|
||||
|
||||
|
||||
def build_image_to_image_prompt(
|
||||
model_class_name: str | None,
|
||||
prompt: str,
|
||||
negative_prompt: str | None,
|
||||
input_image: Image.Image | list[Image.Image],
|
||||
height: int | None = None,
|
||||
width: int | None = None,
|
||||
) -> dict[str, Any]:
|
||||
spec = _get_spec(model_class_name)
|
||||
builder: ImageToImagePromptBuilder = (
|
||||
spec.get("image_to_image_prompt_builder", default_image_to_image_prompt)
|
||||
if spec is not None
|
||||
else default_image_to_image_prompt
|
||||
)
|
||||
return builder(prompt, negative_prompt, input_image, height, width)
|
||||
@@ -0,0 +1,21 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
SENSENOVA_U1_EXTRA_BODY_PARAMS = frozenset(
|
||||
{
|
||||
"think",
|
||||
"cfg_scale",
|
||||
"cfg_norm",
|
||||
"timestep_shift",
|
||||
"t_eps",
|
||||
"img_cfg_scale",
|
||||
"max_tokens",
|
||||
}
|
||||
)
|
||||
SENSENOVA_U1_EXTRA_OUTPUT_PARAMS = frozenset(
|
||||
{
|
||||
"think_text",
|
||||
}
|
||||
)
|
||||
Reference in new issue
Block a user