diff --git a/examples/offline_inference/bagel/README.md b/examples/offline_inference/bagel/README.md deleted file mode 100644 index f2685d56f..000000000 --- a/examples/offline_inference/bagel/README.md +++ /dev/null @@ -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 `...` 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 `...` 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. diff --git a/examples/offline_inference/bagel/end2end.py b/examples/offline_inference/bagel/end2end.py deleted file mode 100644 index 088687253..000000000 --- a/examples/offline_inference/bagel/end2end.py +++ /dev/null @@ -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 ... 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() diff --git a/examples/offline_inference/image_to_image/image_edit.py b/examples/offline_inference/image_to_image/image_edit.py index 0f2445673..299efb33b 100644 --- a/examples/offline_inference/image_to_image/image_edit.py +++ b/examples/offline_inference/image_to_image/image_edit.py @@ -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") diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index b65ffa0da..06e5dc181 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -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") diff --git a/examples/online_serving/bagel/README.md b/examples/online_serving/bagel/README.md deleted file mode 100644 index b7a54da35..000000000 --- a/examples/online_serving/bagel/README.md +++ /dev/null @@ -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 `` 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 \ - --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 \ - --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 < 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 < 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. diff --git a/examples/online_serving/bagel/openai_chat_client.py b/examples/online_serving/bagel/openai_chat_client.py deleted file mode 100755 index cc9ec32db..000000000 --- a/examples/online_serving/bagel/openai_chat_client.py +++ /dev/null @@ -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() diff --git a/examples/online_serving/bagel/run_server.sh b/examples/online_serving/bagel/run_server.sh deleted file mode 100755 index a64057ef0..000000000 --- a/examples/online_serving/bagel/run_server.sh +++ /dev/null @@ -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" diff --git a/examples/online_serving/bagel/run_server_stage_cli.sh b/examples/online_serving/bagel/run_server_stage_cli.sh deleted file mode 100644 index 912e212f9..000000000 --- a/examples/online_serving/bagel/run_server_stage_cli.sh +++ /dev/null @@ -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 <&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 diff --git a/examples/online_serving/image_to_image/README.md b/examples/online_serving/image_to_image/README.md index 59b1f0e2c..8a6c6a0e6 100644 --- a/examples/online_serving/image_to_image/README.md +++ b/examples/online_serving/image_to_image/README.md @@ -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 diff --git a/examples/online_serving/image_to_image/openai_chat_client.py b/examples/online_serving/image_to_image/openai_chat_client.py index 7aac2b214..46f2ad5a7 100644 --- a/examples/online_serving/image_to_image/openai_chat_client.py +++ b/examples/online_serving/image_to_image/openai_chat_client.py @@ -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: diff --git a/examples/online_serving/sensenova_u1/openai_chat_client.py b/examples/online_serving/sensenova_u1/openai_chat_client.py index df9afa83d..50ff202d4 100644 --- a/examples/online_serving/sensenova_u1/openai_chat_client.py +++ b/examples/online_serving/sensenova_u1/openai_chat_client.py @@ -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") diff --git a/recipes/Bagel/BAGEL-7B-MoT.md b/recipes/Bagel/BAGEL-7B-MoT.md new file mode 100644 index 000000000..a7b0524e4 --- /dev/null +++ b/recipes/Bagel/BAGEL-7B-MoT.md @@ -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. diff --git a/recipes/README.md b/recipes/README.md index a1c6973d5..1f5b5be57 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -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 | diff --git a/recipes/SenseNova/SenseNova-U1.md b/recipes/SenseNova/SenseNova-U1.md index b1a14eaa2..5cb30a21d 100644 --- a/recipes/SenseNova/SenseNova-U1.md +++ b/recipes/SenseNova/SenseNova-U1.md @@ -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 diff --git a/tests/model_extras/test_model_extras.py b/tests/model_extras/test_model_extras.py new file mode 100644 index 000000000..09c732d07 --- /dev/null +++ b/tests/model_extras/test_model_extras.py @@ -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, + } diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index a2dfb5a15..c4cca97bb 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -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) diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 3b4843af0..a7869b140 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -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", diff --git a/vllm_omni/diffusion/io_support.py b/vllm_omni/diffusion/io_support.py index 9c7b2b027..8d1405a6c 100644 --- a/vllm_omni/diffusion/io_support.py +++ b/vllm_omni/diffusion/io_support.py @@ -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())) diff --git a/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py b/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py index fa7f95c66..44446ae83 100644 --- a/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py +++ b/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py @@ -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. diff --git a/vllm_omni/diffusion/utils/param_utils.py b/vllm_omni/diffusion/utils/param_utils.py new file mode 100644 index 000000000..e56fcd072 --- /dev/null +++ b/vllm_omni/diffusion/utils/param_utils.py @@ -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} diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index dc1385e38..721e58c7b 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -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( diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 24998436e..9a0a9ed1c 100755 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -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: diff --git a/vllm_omni/model_extras/__init__.py b/vllm_omni/model_extras/__init__.py new file mode 100644 index 000000000..e2f98cd41 --- /dev/null +++ b/vllm_omni/model_extras/__init__.py @@ -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", +] diff --git a/vllm_omni/model_extras/bagel.py b/vllm_omni/model_extras/bagel.py new file mode 100644 index 000000000..e8cbcbf11 --- /dev/null +++ b/vllm_omni/model_extras/bagel.py @@ -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 diff --git a/vllm_omni/model_extras/registry.py b/vllm_omni/model_extras/registry.py new file mode 100644 index 000000000..639a4b67d --- /dev/null +++ b/vllm_omni/model_extras/registry.py @@ -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) diff --git a/vllm_omni/model_extras/sensenova_u1.py b/vllm_omni/model_extras/sensenova_u1.py new file mode 100644 index 000000000..29fcc52c4 --- /dev/null +++ b/vllm_omni/model_extras/sensenova_u1.py @@ -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", + } +)