vllm-omni/recipes/Bagel/BAGEL-7B-MoT.md

# 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, image-to-image, text-to-text, and image-to-text
- 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 offline understanding example:
  [`examples/offline_inference/x_to_text/x_to_text.py`](../../examples/offline_inference/x_to_text/x_to_text.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_think.yaml`](../../vllm_omni/deploy/bagel_think.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 `--deploy-config` 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 \
  --deploy-config 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 \
  --deploy-config 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.
For BAGEL img2img, an explicit `--height`/`--width` (or `size` on
`/v1/images/edits`) sets the output canvas, in both the single-stage and the
two-stage deployment; the input image is still resized to a stride-aligned size
for the VAE/ViT prefill. The latent grid needs multiples of 16, so a requested
side that is not one is floored to the next multiple (1000x700 generates
992x688) with a server-side warning, and the response reports the generated
size. A side above the checkpoint limit is rejected with HTTP 400. When no size
is requested
(`size=auto`, or no `height`/`width`), the pipeline derives the canvas from the
input image: it preserves the input aspect ratio, aligns dimensions to the
latent stride, and applies the checkpoint size limit. Step-mode admission
resolves this effective size before batching so differently shaped requests are
scheduled separately.

Run text-to-text with the shared understanding example. BAGEL's default
`bagel.yaml` deploy config is discovered from the checkpoint, so no model-specific
example or explicit deploy config is required:

```bash
python examples/offline_inference/x_to_text/x_to_text.py \
  --model ByteDance-Seed/BAGEL-7B-MoT \
  --prompt "Where is the capital of France?"
```

Add `--image` to run image-to-text with the same example:

```bash
python examples/offline_inference/x_to_text/x_to_text.py \
  --model ByteDance-Seed/BAGEL-7B-MoT \
  --image /path/to/input.jpg \
  --prompt "Please describe this image in detail."
```

#### 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
```

#### Step Execution and Continuous Batching

BAGEL supports step execution for image generation with all three in-repo
deploy configs:

- `bagel.yaml`: the flag applies to the Stage 1 diffusion engine; Stage 0
  remains on the normal autoregressive path.
- `bagel_think.yaml`: the flag also applies only to the Stage 1 diffusion
  engine; Thinker decoding in Stage 0 is unchanged.
- `bagel_single_stage.yaml`: image requests use step execution in Stage 0.
  Explicit text-output requests keep the existing complete-request path, so
  `text2text` and `img2text` behavior is unchanged.

Start with one active request when validating correctness. For example, the
following command uses the single-stage topology:

```bash
vllm serve ByteDance-Seed/BAGEL-7B-MoT \
  --omni \
  --port 8091 \
  --deploy-config vllm_omni/deploy/bagel_single_stage.yaml \
  --step-execution
```

To let compatible requests share denoising waves, override the diffusion-stage
capacity. The deploy files set `max_num_seqs: 1` for their diffusion stages, so
the per-stage value takes precedence over the global `--max-num-seqs` flag.
Use Stage 0 for `bagel_single_stage.yaml`:

```bash
vllm serve ByteDance-Seed/BAGEL-7B-MoT \
  --omni \
  --port 8091 \
  --deploy-config vllm_omni/deploy/bagel_single_stage.yaml \
  --step-execution \
  --stage-overrides '{"0":{"max_num_seqs":4}}'
```

For either two-stage config, use Stage 1 instead:

```bash
vllm serve ByteDance-Seed/BAGEL-7B-MoT \
  --omni \
  --port 8091 \
  --deploy-config vllm_omni/deploy/bagel.yaml \
  --step-execution \
  --stage-overrides '{"1":{"max_num_seqs":4}}'
```

Use at least two inference steps for image requests. BAGEL builds
`num_inference_steps - 1` scheduler updates, so a one-step image request has no
denoising iteration to schedule. The scheduler automatically separates
requests with incompatible image shapes or BAGEL CFG/renormalization settings.

Step execution currently does not support BAGEL sequence parallelism or a
diffusion cache backend. Do not combine it with BAGEL SP, TeaCache, or
Cache-DiT. Tensor parallelism and CFG parallelism are separate features. For
best utilization, match the CFG parallel world size to the number of active
BAGEL CFG branches. A larger world size is supported, but extra ranks run dummy
forward passes to participate in collectives and their outputs are ignored.
BAGEL currently does not support two CFG ranks when three CFG branches are
active.

For a matched throughput comparison, start the first server with the default
stage capacity of one and run:

```bash
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
  --base-url http://localhost:8091 \
  --endpoint /v1/chat/completions \
  --model ByteDance-Seed/BAGEL-7B-MoT \
  --task t2i \
  --dataset random \
  --num-prompts 20 \
  --max-concurrency 4 \
  --width 512 \
  --height 512 \
  --num-inference-steps 20 \
  --seed 42 \
  --warmup-requests 4 \
  --warmup-concurrency 4 \
  --warmup-num-inference-steps 2 \
  --output-file /tmp/bagel_step_seq1.json \
  --disable-tqdm
```

Restart the single-stage server with `--stage-overrides
'{"0":{"max_num_seqs":4}}'` (or use Stage 1 for either two-stage config),
rerun the same benchmark with
`--output-file /tmp/bagel_step_seq4.json`, and compare completed requests,
duration, throughput, and latency percentiles. Keep the model, device topology,
request count, client concurrency, resolution, steps, seed, and warmup settings
fixed between runs.

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.